mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-07-10 17:35:08 +00:00
fix(output): expose token usage under both names on agent and crew results
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 <noreply@anthropic.com>
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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."""
|
||||
|
||||
143
lib/crewai/tests/test_usage_shape_parity.py
Normal file
143
lib/crewai/tests/test_usage_shape_parity.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user