From b65c8487d2ff0de8ac76c5d8ea0cbd57901dfd24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Moura?= Date: Fri, 10 Jul 2026 20:29:34 -0300 Subject: [PATCH] fix(output): expose token usage under both names on agent and crew results (#6507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Agent.kickoff() returned LiteAgentOutput with a plain dict at .usage_metrics and no token_usage attribute, while Crew.kickoff() returned CrewOutput with a UsageMetrics object at .token_usage and no usage_metrics attribute — so a usage accessor written for one path raised AttributeError on the other, and every consumer had to duck-type both shapes. Give both result types both surfaces, each name with one consistent shape everywhere: .token_usage is a UsageMetrics object and .usage_metrics is a plain dict, on both LiteAgentOutput and CrewOutput. Added as read-only properties, so existing fields, serialization, and constructors are unchanged. Fixes EPD-178. Co-authored-by: Claude Fable 5 --- lib/crewai/src/crewai/crews/crew_output.py | 16 ++- lib/crewai/src/crewai/lite_agent_output.py | 20 ++- lib/crewai/tests/test_usage_shape_parity.py | 143 ++++++++++++++++++++ 3 files changed, 176 insertions(+), 3 deletions(-) create mode 100644 lib/crewai/tests/test_usage_shape_parity.py 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 bd17f5855..4b4bec446 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 @@ -39,8 +40,10 @@ 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 kickoff call only (guardrail " - "retries included), not the LLM instance's cumulative totals" + "Token usage for this kickoff call only (guardrail retries " + "included), not the LLM instance's cumulative totals, as a " + "plain dict; ``token_usage`` exposes the same data as a " + "UsageMetrics object" ), default=None, ) @@ -90,6 +93,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)