diff --git a/lib/crewai/src/crewai/crews/crew_output.py b/lib/crewai/src/crewai/crews/crew_output.py index 4541ae02a..13b431feb 100644 --- a/lib/crewai/src/crewai/crews/crew_output.py +++ b/lib/crewai/src/crewai/crews/crew_output.py @@ -24,9 +24,23 @@ class CrewOutput(BaseModel): description="Output of each task", default_factory=list ) token_usage: UsageMetrics = Field( - description="Processed token summary", default_factory=UsageMetrics + description=( + "Processed token summary; ``usage_metrics`` exposes the same " + "data as a plain dict" + ), + default_factory=UsageMetrics, ) + @property + def usage_metrics(self) -> dict[str, Any]: + """Token usage as a plain dict. + + Same attribute name and shape as ``LiteAgentOutput.usage_metrics`` + (the ``Agent.kickoff()`` result), so a usage accessor written for one + result type works on both. + """ + return self.token_usage.model_dump() + @property def json(self) -> str | None: # type: ignore[override] if self.tasks_output[-1].output_format != OutputFormat.JSON: diff --git a/lib/crewai/src/crewai/lite_agent_output.py b/lib/crewai/src/crewai/lite_agent_output.py index 1ac79d422..1ba3680bb 100644 --- a/lib/crewai/src/crewai/lite_agent_output.py +++ b/lib/crewai/src/crewai/lite_agent_output.py @@ -6,6 +6,7 @@ from typing import Any from pydantic import BaseModel, Field +from crewai.types.usage_metrics import UsageMetrics from crewai.utilities.planning_types import TodoItem from crewai.utilities.types import LLMMessage @@ -38,7 +39,11 @@ class LiteAgentOutput(BaseModel): ) agent_role: str = Field(description="Role of the agent that produced this output") usage_metrics: dict[str, Any] | None = Field( - description="Token usage metrics for this execution", default=None + description=( + "Token usage metrics for this execution as a plain dict; " + "``token_usage`` exposes the same data as a UsageMetrics object" + ), + default=None, ) messages: list[LLMMessage] = Field( description="Messages of the agent", default_factory=list @@ -86,6 +91,19 @@ class LiteAgentOutput(BaseModel): return self.pydantic.model_dump() return {} + @property + def token_usage(self) -> UsageMetrics: + """Token usage as a ``UsageMetrics`` object. + + Same attribute name and type as ``CrewOutput.token_usage``, so a + usage accessor written for one result type works on both. Returns + zeroed metrics when no usage was captured (``usage_metrics`` is + ``None``). + """ + if not self.usage_metrics: + return UsageMetrics() + return UsageMetrics.model_validate(self.usage_metrics) + @property def completed_todos(self) -> list[TodoExecutionResult]: """Get only the completed todos.""" diff --git a/lib/crewai/tests/test_usage_shape_parity.py b/lib/crewai/tests/test_usage_shape_parity.py new file mode 100644 index 000000000..1b7234a53 --- /dev/null +++ b/lib/crewai/tests/test_usage_shape_parity.py @@ -0,0 +1,143 @@ +# mypy: ignore-errors +"""Regression tests for EPD-178: token usage was exposed in different shapes +and attribute names per code path — ``Agent.kickoff()`` results carried a +plain dict at ``.usage_metrics`` (no ``token_usage`` attribute at all), while +``Crew.kickoff()`` results carried a ``UsageMetrics`` object at +``.token_usage`` (no ``usage_metrics`` attribute), so any single accessor +written for one path raised ``AttributeError`` on the other. + +Both result types now expose both surfaces: ``.token_usage`` as a +``UsageMetrics`` object and ``.usage_metrics`` as a plain dict. +""" + +from crewai import Agent, Crew, Task +from crewai.crews.crew_output import CrewOutput +from crewai.lite_agent_output import LiteAgentOutput +from crewai.llms.base_llm import BaseLLM +from crewai.types.usage_metrics import UsageMetrics + + +class _FixedUsageLLM(BaseLLM): + """Offline BaseLLM that records fixed usage (100/10 tokens) per call.""" + + def __init__(self): + super().__init__(model="fixed-usage-model") + + def call( + self, + messages, + tools=None, + callbacks=None, + available_functions=None, + from_task=None, + from_agent=None, + response_model=None, + ) -> str: + self._track_token_usage_internal( + {"prompt_tokens": 100, "completion_tokens": 10, "total_tokens": 110} + ) + return "Thought: I know the answer.\nFinal Answer: fake answer" + + def supports_function_calling(self) -> bool: + return False + + def supports_stop_words(self) -> bool: + return False + + def get_context_window_size(self) -> int: + return 4096 + + +class TestUsageShapeUnitParity: + def test_lite_agent_output_exposes_token_usage_object(self): + metrics = UsageMetrics( + total_tokens=110, + prompt_tokens=100, + completion_tokens=10, + successful_requests=1, + ) + output = LiteAgentOutput( + agent_role="analyst", usage_metrics=metrics.model_dump() + ) + + assert output.token_usage == metrics + assert isinstance(output.token_usage, UsageMetrics) + + def test_lite_agent_output_token_usage_zeroed_when_absent(self): + output = LiteAgentOutput(agent_role="analyst") + + assert output.usage_metrics is None + assert output.token_usage == UsageMetrics() + + def test_crew_output_exposes_usage_metrics_dict(self): + metrics = UsageMetrics( + total_tokens=110, + prompt_tokens=100, + completion_tokens=10, + successful_requests=1, + ) + output = CrewOutput(token_usage=metrics) + + assert output.usage_metrics == metrics.model_dump() + assert isinstance(output.usage_metrics, dict) + + def test_both_shapes_carry_identical_keys(self): + """The dict shape has exactly the UsageMetrics fields on both types.""" + crew_dict = CrewOutput(token_usage=UsageMetrics()).usage_metrics + lite = LiteAgentOutput( + agent_role="analyst", usage_metrics=UsageMetrics().model_dump() + ) + + assert set(crew_dict) == set(UsageMetrics.model_fields) + assert set(lite.usage_metrics) == set(UsageMetrics.model_fields) + + +class TestUsageShapeEndToEnd: + """Mirror of the EPD-178 clean-room repro, offline via a fake BaseLLM.""" + + @staticmethod + def _read_via_object(result) -> int: + """Single accessor written against the CrewOutput shape.""" + return result.token_usage.prompt_tokens + + @staticmethod + def _read_via_dict(result) -> int: + """Single accessor written against the LiteAgentOutput shape.""" + return result.usage_metrics["prompt_tokens"] + + def test_single_accessor_works_on_both_kickoff_paths(self): + agent_a = Agent( + role="analyst", + goal="Answer questions.", + backstory="Test agent.", + llm=_FixedUsageLLM(), + verbose=False, + ) + result_agent = agent_a.kickoff("a question") + + agent_b = Agent( + role="analyst", + goal="Answer questions.", + backstory="Test agent.", + llm=_FixedUsageLLM(), + verbose=False, + ) + task = Task( + description="Answer: a question", + expected_output="A short answer.", + agent=agent_b, + ) + crew = Crew(agents=[agent_b], tasks=[task], verbose=False) + result_crew = crew.kickoff() + + assert isinstance(result_agent, LiteAgentOutput) + assert isinstance(result_crew, CrewOutput) + + # Both accessors work on both result types and agree with each other. + for result in (result_agent, result_crew): + object_read = self._read_via_object(result) + dict_read = self._read_via_dict(result) + assert object_read == dict_read + assert object_read > 0 + assert isinstance(result.token_usage, UsageMetrics) + assert isinstance(result.usage_metrics, dict)