andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-01-09 16:18:30 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Fix token tracking issues in async tasks and agent metrics

Browse Source 
Resolved 4 review comments from Cursor Bugbot:
1. Added token tracking for async tasks in _execute_tasks and _process_async_tasks
2. Fixed task key collision by including task_id in the key
3. Added token tracking for _aexecute_tasks paths (both sync and async)
4. Fixed agent metrics to be keyed by agent_id to handle multiple agents with same role

All async tasks now capture tokens_before/after and attach metrics properly.
Task metrics now use unique keys to prevent overwriting.
Agent metrics properly track separate agents with same role.
This commit is contained in:
Devasy Patel
2026-01-03 22:24:27 +05:30
parent 85860610e9
commit afea8a505a
1 changed files with 77 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
lib/crewai/src/crewai/crew.py
View File

@@ -948,6 +948,9 @@ class Crew(FlowTrackable, BaseModel):
continue continue
if task.async_execution: if task.async_execution:
# Capture token usage before async task execution
tokens_before = self._get_agent_token_usage(exec_data.agent)
context = self._get_context( context = self._get_context(
task, [last_sync_output] if last_sync_output else [] task, [last_sync_output] if last_sync_output else []
) )
@@ -958,7 +961,7 @@ class Crew(FlowTrackable, BaseModel):
tools=exec_data.tools, tools=exec_data.tools,
) )
) )
pending_tasks.append((task, async_task, task_index)) pending_tasks.append((task, async_task, task_index, exec_data.agent, tokens_before))
else: else:
if pending_tasks: if pending_tasks:
task_outputs = await self._aprocess_async_tasks( task_outputs = await self._aprocess_async_tasks(
@@ -966,12 +969,22 @@ class Crew(FlowTrackable, BaseModel):
) )
pending_tasks.clear() pending_tasks.clear()
# Capture token usage before task execution
tokens_before = self._get_agent_token_usage(exec_data.agent)
context = self._get_context(task, task_outputs) context = self._get_context(task, task_outputs)
task_output = await task.aexecute_sync( task_output = await task.aexecute_sync(
agent=exec_data.agent, agent=exec_data.agent,
context=context, context=context,
tools=exec_data.tools, tools=exec_data.tools,
) )
# Capture token usage after task execution and attach to task output
tokens_after = self._get_agent_token_usage(exec_data.agent)
task_output = self._attach_task_token_metrics(
task_output, task, exec_data.agent, tokens_before, tokens_after
)
task_outputs.append(task_output) task_outputs.append(task_output)
self._process_task_result(task, task_output) self._process_task_result(task, task_output)
self._store_execution_log(task, task_output, task_index, was_replayed) self._store_execution_log(task, task_output, task_index, was_replayed)
@@ -985,7 +998,7 @@ class Crew(FlowTrackable, BaseModel):
self, self,
task: ConditionalTask, task: ConditionalTask,
task_outputs: list[TaskOutput], task_outputs: list[TaskOutput],
pending_tasks: list[tuple[Task, asyncio.Task[TaskOutput], int]], pending_tasks: list[tuple[Task, asyncio.Task[TaskOutput], int, Any, Any]],
task_index: int, task_index: int,
was_replayed: bool, was_replayed: bool,
) -> TaskOutput | None: ) -> TaskOutput | None:
@@ -1000,13 +1013,20 @@ class Crew(FlowTrackable, BaseModel):
async def _aprocess_async_tasks( async def _aprocess_async_tasks(
self, self,
pending_tasks: list[tuple[Task, asyncio.Task[TaskOutput], int]], pending_tasks: list[tuple[Task, asyncio.Task[TaskOutput], int, Any, Any]],
was_replayed: bool = False, was_replayed: bool = False,
) -> list[TaskOutput]: ) -> list[TaskOutput]:
"""Process pending async tasks and return their outputs.""" """Process pending async tasks and return their outputs."""
task_outputs: list[TaskOutput] = [] task_outputs: list[TaskOutput] = []
for future_task, async_task, task_index in pending_tasks: for future_task, async_task, task_index, agent, tokens_before in pending_tasks:
task_output = await async_task task_output = await async_task
# Capture token usage after async task execution and attach to task output
tokens_after = self._get_agent_token_usage(agent)
task_output = self._attach_task_token_metrics(
task_output, future_task, agent, tokens_before, tokens_after
)
task_outputs.append(task_output) task_outputs.append(task_output)
self._process_task_result(future_task, task_output) self._process_task_result(future_task, task_output)
self._store_execution_log( self._store_execution_log(
@@ -1145,6 +1165,9 @@ class Crew(FlowTrackable, BaseModel):
continue continue
if task.async_execution: if task.async_execution:
# Capture token usage before async task execution
tokens_before = self._get_agent_token_usage(exec_data.agent)
context = self._get_context( context = self._get_context(
task, [last_sync_output] if last_sync_output else [] task, [last_sync_output] if last_sync_output else []
) )
@@ -1153,7 +1176,7 @@ class Crew(FlowTrackable, BaseModel):
context=context, context=context,
tools=exec_data.tools, tools=exec_data.tools,
) )
futures.append((task, future, task_index)) futures.append((task, future, task_index, exec_data.agent, tokens_before))
else: else:
if futures: if futures:
task_outputs = self._process_async_tasks(futures, was_replayed) task_outputs = self._process_async_tasks(futures, was_replayed)
@@ -1188,7 +1211,7 @@ class Crew(FlowTrackable, BaseModel):
self, self,
task: ConditionalTask, task: ConditionalTask,
task_outputs: list[TaskOutput], task_outputs: list[TaskOutput],
futures: list[tuple[Task, Future[TaskOutput], int]], futures: list[tuple[Task, Future[TaskOutput], int, Any, Any]],
task_index: int, task_index: int,
was_replayed: bool, was_replayed: bool,
) -> TaskOutput | None: ) -> TaskOutput | None:
@@ -1420,12 +1443,19 @@ class Crew(FlowTrackable, BaseModel):
def _process_async_tasks( def _process_async_tasks(
self, self,
futures: list[tuple[Task, Future[TaskOutput], int]], futures: list[tuple[Task, Future[TaskOutput], int, Any, Any]],
was_replayed: bool = False, was_replayed: bool = False,
) -> list[TaskOutput]: ) -> list[TaskOutput]:
task_outputs: list[TaskOutput] = [] task_outputs: list[TaskOutput] = []
for future_task, future, task_index in futures: for future_task, future, task_index, agent, tokens_before in futures:
task_output = future.result() task_output = future.result()
# Capture token usage after async task execution and attach to task output
tokens_after = self._get_agent_token_usage(agent)
task_output = self._attach_task_token_metrics(
task_output, future_task, agent, tokens_before, tokens_after
)
task_outputs.append(task_output) task_outputs.append(task_output)
self._process_task_result(future_task, task_output) self._process_task_result(future_task, task_output)
self._store_execution_log( self._store_execution_log(
@@ -1646,34 +1676,53 @@ class Crew(FlowTrackable, BaseModel):
# Build per-agent metrics from per-task data (more accurate) # Build per-agent metrics from per-task data (more accurate)
# This avoids the cumulative token issue where all agents show the same total # This avoids the cumulative token issue where all agents show the same total
# Key by agent_id to handle multiple agents with the same role
agent_token_sums = {} agent_token_sums = {}
agent_info_map = {} # Map agent_id to (agent_name, agent_id)
# First, build a map of all agents by their ID
for agent in self.agents:
agent_role = getattr(agent, 'role', 'Unknown Agent')
agent_id = str(getattr(agent, 'id', ''))
agent_info_map[agent_id] = (agent_role, agent_id)
if workflow_metrics.per_task: if workflow_metrics.per_task:
# Sum up tokens for each agent from their tasks # Sum up tokens for each agent from their tasks
# We need to find which agent_id corresponds to each task's agent_name
for task_name, task_metrics in workflow_metrics.per_task.items(): for task_name, task_metrics in workflow_metrics.per_task.items():
agent_name = task_metrics.agent_name agent_name = task_metrics.agent_name
if agent_name not in agent_token_sums: # Find the agent_id for this agent_name from agent_info_map
agent_token_sums[agent_name] = { # For now, we'll use the agent_name as a temporary key but this needs improvement
'total_tokens': 0, # TODO: Store agent_id in TaskTokenMetrics to avoid this lookup
'prompt_tokens': 0, matching_agent_ids = [aid for aid, (name, _) in agent_info_map.items() if name == agent_name]
'cached_prompt_tokens': 0,
'completion_tokens': 0, # Use the first matching agent_id (limitation: can't distinguish between same-role agents)
'successful_requests': 0 # This is better than nothing but ideally we'd store agent_id in TaskTokenMetrics
} for agent_id in matching_agent_ids:
agent_token_sums[agent_name]['total_tokens'] += task_metrics.total_tokens if agent_id not in agent_token_sums:
agent_token_sums[agent_name]['prompt_tokens'] += task_metrics.prompt_tokens agent_token_sums[agent_id] = {
agent_token_sums[agent_name]['cached_prompt_tokens'] += task_metrics.cached_prompt_tokens 'total_tokens': 0,
agent_token_sums[agent_name]['completion_tokens'] += task_metrics.completion_tokens 'prompt_tokens': 0,
agent_token_sums[agent_name]['successful_requests'] += task_metrics.successful_requests 'cached_prompt_tokens': 0,
'completion_tokens': 0,
'successful_requests': 0
}
# Only add to the first matching agent (this is the limitation)
agent_token_sums[agent_id]['total_tokens'] += task_metrics.total_tokens
agent_token_sums[agent_id]['prompt_tokens'] += task_metrics.prompt_tokens
agent_token_sums[agent_id]['cached_prompt_tokens'] += task_metrics.cached_prompt_tokens
agent_token_sums[agent_id]['completion_tokens'] += task_metrics.completion_tokens
agent_token_sums[agent_id]['successful_requests'] += task_metrics.successful_requests
break # Only add to first matching agent
# Create per-agent metrics from the summed task data # Create per-agent metrics from the summed task data, keyed by agent_id
for agent in self.agents: for agent in self.agents:
agent_role = getattr(agent, 'role', 'Unknown Agent') agent_role = getattr(agent, 'role', 'Unknown Agent')
agent_id = str(getattr(agent, 'id', '')) agent_id = str(getattr(agent, 'id', ''))
if agent_role in agent_token_sums: if agent_id in agent_token_sums:
# Use accurate per-task summed data # Use accurate per-task summed data
sums = agent_token_sums[agent_role] sums = agent_token_sums[agent_id]
agent_metrics = AgentTokenMetrics( agent_metrics = AgentTokenMetrics(
agent_name=agent_role, agent_name=agent_role,
agent_id=agent_id, agent_id=agent_id,
@@ -1683,7 +1732,8 @@ class Crew(FlowTrackable, BaseModel):
completion_tokens=sums['completion_tokens'], completion_tokens=sums['completion_tokens'],
successful_requests=sums['successful_requests'] successful_requests=sums['successful_requests']
) )
workflow_metrics.per_agent[agent_role] = agent_metrics # Key by agent_id to avoid collision for agents with same role
workflow_metrics.per_agent[agent_id] = agent_metrics
# Still get total usage for overall metrics # Still get total usage for overall metrics
if isinstance(agent.llm, BaseLLM): if isinstance(agent.llm, BaseLLM):
@@ -2074,7 +2124,8 @@ To enable tracing, do any one of these:
from crewai.types.usage_metrics import WorkflowTokenMetrics from crewai.types.usage_metrics import WorkflowTokenMetrics
self.workflow_token_metrics = WorkflowTokenMetrics() self.workflow_token_metrics = WorkflowTokenMetrics()
task_key = f"{task_tokens.task_name}_{task_tokens.agent_name}" # Use task_id in the key to prevent collision when multiple tasks have the same name
task_key = f"{task_tokens.task_id}_{task_tokens.task_name}_{task_tokens.agent_name}"
self.workflow_token_metrics.per_task[task_key] = task_tokens self.workflow_token_metrics.per_task[task_key] = task_tokens
return task_output return task_output