diff --git a/lib/crewai/src/crewai/crew.py b/lib/crewai/src/crewai/crew.py index 0f77b2d224..381931c7bc 100644 --- a/lib/crewai/src/crewai/crew.py +++ b/lib/crewai/src/crewai/crew.py @@ -226,6 +226,7 @@ class Crew(FlowTrackable, BaseModel): ) _kickoff_event_id: str | None = PrivateAttr(default=None) _execution_start_dispatched: bool = PrivateAttr(default=False) + _usage_baselines: dict[int, UsageMetrics] = PrivateAttr(default_factory=dict) _execution_end_dispatched: bool = PrivateAttr(default=False) name: str | None = Field(default="crew") @@ -1047,6 +1048,7 @@ def run_crew() -> None: execution_token = begin_execution() + self._snapshot_usage_baselines() runtime_scope = crewai_event_bus._enter_runtime_scope() try: inputs = prepare_kickoff(self, inputs, input_files) @@ -1264,6 +1266,7 @@ async def run_crew() -> None: execution_token = begin_execution() + self._snapshot_usage_baselines() runtime_scope = crewai_event_bus._enter_runtime_scope() try: inputs = prepare_kickoff(self, inputs, input_files) @@ -2201,28 +2204,50 @@ def _finish_execution(self, final_string_output: str) -> None: if self.max_rpm: self._rpm_controller.stop_rpm_counter() - def calculate_usage_metrics(self) -> UsageMetrics: - """Calculates and returns the usage metrics.""" - total_usage_metrics = UsageMetrics() + def _llm_instances(self) -> list[BaseLLM]: + """Return the distinct LLM instances this crew runs on. - for agent in self.agents: - if isinstance(agent.llm, BaseLLM): - llm_usage = agent.llm.get_token_usage_summary() + De-duplicated by object identity: an instance shared by several agents + holds one set of counters, so it must be measured once. + """ + instances: list[BaseLLM] = [] + seen: set[int] = set() + for agent in (*self.agents, self.manager_agent): + llm = getattr(agent, "llm", None) + if isinstance(llm, BaseLLM) and id(llm) not in seen: + seen.add(id(llm)) + instances.append(llm) + return instances + + def _snapshot_usage_baselines(self) -> None: + """Record each LLM instance's counters at the start of a kickoff. + + An instance's counters are cumulative for its lifetime, so usage for + one run is the difference between these baselines and the counters at + the end. Snapshotting per instance rather than per agent keeps a shared + instance from being counted once per agent, and needs no per-agent + window, so concurrent tasks on one instance stay correct. + """ + self._usage_baselines = { + id(llm): llm.get_token_usage_summary() for llm in self._llm_instances() + } - total_usage_metrics.add_usage_metrics(llm_usage) - else: - if hasattr(agent, "_token_process"): - token_sum = agent._token_process.get_summary() - total_usage_metrics.add_usage_metrics(token_sum) + def calculate_usage_metrics(self) -> UsageMetrics: + """Return the token usage accrued during the most recent kickoff.""" + total_usage_metrics = UsageMetrics() - if self.manager_agent and hasattr(self.manager_agent, "_token_process"): - token_sum = self.manager_agent._token_process.get_summary() - total_usage_metrics.add_usage_metrics(token_sum) + for llm in self._llm_instances(): + current = llm.get_token_usage_summary() + baseline = self._usage_baselines.get(id(llm)) + usage = current.delta_since(baseline) if baseline is not None else current + total_usage_metrics.add_usage_metrics(usage) - if self.manager_agent: - if isinstance(self.manager_agent.llm, BaseLLM): - llm_usage = self.manager_agent.llm.get_token_usage_summary() - total_usage_metrics.add_usage_metrics(llm_usage) + for agent in (*self.agents, self.manager_agent): + if agent is None or isinstance(getattr(agent, "llm", None), BaseLLM): + continue + token_process = getattr(agent, "_token_process", None) + if token_process is not None: + total_usage_metrics.add_usage_metrics(token_process.get_summary()) self.usage_metrics = total_usage_metrics return total_usage_metrics diff --git a/lib/crewai/tests/test_crew.py b/lib/crewai/tests/test_crew.py index 0195112cb9..42fdfb172f 100644 --- a/lib/crewai/tests/test_crew.py +++ b/lib/crewai/tests/test_crew.py @@ -1834,13 +1834,22 @@ def test_hierarchical_kickoff_usage_metrics_include_manager(researcher): total_tokens=30, prompt_tokens=20, completion_tokens=10, successful_requests=1 ) - researcher.llm.get_token_usage_summary = MagicMock(return_value=researcher_metrics) - - # Mock the manager's _token_process since it uses the fallback path - manager._token_process = MagicMock( - get_summary=MagicMock(return_value=manager_metrics) + # Usage for a run is the growth of each LLM's counters across it, so the + # summaries read empty until the task runs and report totals afterwards. + consumed = {"done": False} + researcher.llm.get_token_usage_summary = ( + lambda: researcher_metrics if consumed["done"] else UsageMetrics() + ) + manager.llm.get_token_usage_summary = ( + lambda: manager_metrics if consumed["done"] else UsageMetrics() ) + def _execute(*_args, **_kwargs) -> TaskOutput: + consumed["done"] = True + return TaskOutput( + description="dummy", raw="Hello", agent=researcher.role, messages=[] + ) + crew = Crew( agents=[researcher], manager_agent=manager, # manager to be included @@ -1849,13 +1858,7 @@ def test_hierarchical_kickoff_usage_metrics_include_manager(researcher): ) # We don't care about LLM output here; patch execute_sync to avoid network - with patch.object( - Task, - "execute_sync", - return_value=TaskOutput( - description="dummy", raw="Hello", agent=researcher.role, messages=[] - ), - ): + with patch.object(Task, "execute_sync", side_effect=_execute): crew.kickoff() assert ( diff --git a/lib/crewai/tests/test_crew_usage_metrics.py b/lib/crewai/tests/test_crew_usage_metrics.py new file mode 100644 index 0000000000..cd51555150 --- /dev/null +++ b/lib/crewai/tests/test_crew_usage_metrics.py @@ -0,0 +1,99 @@ +"""Tests for crew-level token usage aggregation. + +``crew.usage_metrics`` reports what a kickoff consumed, measured as the growth +of each LLM instance's counters across the run. Reading those counters as +absolute totals instead both multiplied usage across agents sharing an +instance and carried earlier runs into later ones. + +Baselines are taken per distinct instance rather than per agent, so a shared +instance is measured once and no per-agent window exists to overlap when tasks +run concurrently. +""" + +from __future__ import annotations + +from crewai import Agent, Crew, Task, LLM +from crewai.types.usage_metrics import UsageMetrics + + +class _Counter: + """Stands in for an LLM instance's cumulative lifetime counters.""" + + def __init__(self) -> None: + self.total = 0 + + def bind(self, llm: LLM) -> LLM: + llm.get_token_usage_summary = lambda: UsageMetrics( # type: ignore[method-assign] + total_tokens=self.total, + prompt_tokens=self.total, + successful_requests=1 if self.total else 0, + ) + return llm + + def consume(self, tokens: int) -> None: + self.total += tokens + + +def _crew(*llms: LLM) -> Crew: + agents = [ + Agent(role=f"Role {i}", goal="goal", backstory="backstory", llm=llm) + for i, llm in enumerate(llms) + ] + tasks = [ + Task(description=f"task {i}", expected_output="out", agent=agent) + for i, agent in enumerate(agents) + ] + return Crew(agents=agents, tasks=tasks) + + +def test_agents_sharing_one_llm_are_counted_once() -> None: + """Three agents on one instance report that instance's usage, not 3x it.""" + counter = _Counter() + llm = counter.bind(LLM(model="gpt-4o")) + crew = _crew(llm, llm, llm) + + crew._snapshot_usage_baselines() + counter.consume(100) + + assert crew.calculate_usage_metrics().total_tokens == 100 + + +def test_distinct_llm_instances_are_summed() -> None: + """Separate instances still add up, even for the same model.""" + first, second = _Counter(), _Counter() + crew = _crew(first.bind(LLM(model="gpt-4o")), second.bind(LLM(model="gpt-4o"))) + + crew._snapshot_usage_baselines() + first.consume(100) + second.consume(50) + + assert crew.calculate_usage_metrics().total_tokens == 150 + + +def test_a_later_run_excludes_the_previous_one() -> None: + """Usage is the growth across this run, not the instance's lifetime.""" + counter = _Counter() + crew = _crew(counter.bind(LLM(model="gpt-4o"))) + + crew._snapshot_usage_baselines() + counter.consume(100) + assert crew.calculate_usage_metrics().total_tokens == 100 + + crew._snapshot_usage_baselines() + counter.consume(50) + assert crew.calculate_usage_metrics().total_tokens == 50 + + +def test_manager_sharing_an_agent_llm_is_counted_once() -> None: + """A manager on the same instance as its agents adds no extra usage.""" + counter = _Counter() + llm = counter.bind(LLM(model="gpt-4o")) + crew = _crew(llm) + crew.manager_agent = Agent( + role="Manager", goal="coordinate", backstory="backstory", llm=llm + ) + + crew._snapshot_usage_baselines() + counter.consume(100) + + assert crew.calculate_usage_metrics().total_tokens == 100