How to get the prompt token after each and every task? → I need this to raise the userdeifned exception if the token limit exceeded.
TaskOutput doesn’t have the attribute of usage metrics. Help me on this.
3 Likes
I was wondering the same issue and it was an dead end. I was trying to use from langchain.callbacks import get_openai_callback or
from langchain.callbacks.openai_info import OpenAICallbackHandler but it yields no result. If you have something in your mind, please share. Thank you!
Same issue, but I was able to estimate per-agent tokens based on output proportions and then distributed the total_usage accordingly. I know this might not be the cleanest approach but this is how I got it working.
"""
Standalone test script for CrewAI token tracking - Version 3 (Final).
This version uses CrewAI's result.token_usage for total tokens and
estimates per-agent usage based on task count and output lengths.
Since CrewAI doesn't expose per-task usage_metrics (as of v1.3.0),
we use a reasonable estimation approach.
Strategy:
1. Get total tokens from result.token_usage (crew-level)
2. Get task list with agents from result.tasks_output
3. Estimate per-agent usage based on:
- Equal division for simple cases
- Output length proportions for better estimation
Requirements:
- crewai
- GEMINI_API_KEY in .env file
Usage:
python test_crewai_token_tracking_v3.py
"""
print("Starting token tracking test v3...")
import os
import sys
import logging
# Windows compatibility
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src', 'backend'))
try:
import win_compat
print("✓ Windows compatibility patches loaded")
except ImportError:
print("⚠ Windows compatibility module not found")
from dotenv import load_dotenv
from crewai import Agent, Task, Crew, Process
from crewai.llm import LLM
from pathlib import Path
# Load environment variables
env_paths = [
Path(__file__).parent / '.env',
Path(__file__).parent / 'src' / 'backend' / '.env',
]
for env_path in env_paths:
if env_path.exists():
load_dotenv(env_path)
print(f"✓ Loaded environment from: {env_path}")
break
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
def estimate_per_agent_tokens(result):
"""
Estimate per-agent token usage from CrewAI result.
Args:
result: CrewOutput from crew.kickoff()
Returns:
dict: {agent_name: {total_tokens, prompt_tokens, completion_tokens, task_count}}
"""
per_agent = {}
# Get total usage
total_tokens = 0
total_prompt = 0
total_completion = 0
if hasattr(result, 'token_usage') and result.token_usage:
usage_obj = result.token_usage
total_tokens = getattr(usage_obj, 'total_tokens', 0)
total_prompt = getattr(usage_obj, 'prompt_tokens', 0)
total_completion = getattr(usage_obj, 'completion_tokens', 0)
# Get task list
if not hasattr(result, 'tasks_output') or not result.tasks_output:
logger.warning("No tasks_output found")
return per_agent
tasks = result.tasks_output
num_tasks = len(tasks)
if num_tasks == 0:
return per_agent
# Method 1: Equal division (simple but fair for similar tasks)
# For better accuracy, we could use output length proportions
# Calculate output lengths for proportional distribution
task_data = []
total_output_length = 0
for task_output in tasks:
# Extract agent
agent_name = "Unknown"
if hasattr(task_output, 'agent'):
agent_obj = task_output.agent
if hasattr(agent_obj, 'role'):
agent_name = str(agent_obj.role)
elif isinstance(agent_obj, str):
agent_name = agent_obj
# Extract output length
output_length = 0
if hasattr(task_output, 'raw'):
output_length = len(str(task_output.raw))
elif hasattr(task_output, 'description'):
output_length = len(str(task_output.description))
total_output_length += output_length
task_data.append({
'agent': agent_name,
'output_length': output_length
})
# Distribute tokens proportionally based on output length
if total_output_length > 0:
for task_info in task_data:
agent_name = task_info['agent']
proportion = task_info['output_length'] / total_output_length
if agent_name not in per_agent:
per_agent[agent_name] = {
'total_tokens': 0,
'prompt_tokens': 0,
'completion_tokens': 0,
'task_count': 0
}
per_agent[agent_name]['total_tokens'] += int(total_tokens * proportion)
per_agent[agent_name]['prompt_tokens'] += int(total_prompt * proportion)
per_agent[agent_name]['completion_tokens'] += int(total_completion * proportion)
per_agent[agent_name]['task_count'] += 1
else:
# Fallback: equal division
tokens_per_task = total_tokens / num_tasks
prompt_per_task = total_prompt / num_tasks
completion_per_task = total_completion / num_tasks
for task_info in task_data:
agent_name = task_info['agent']
if agent_name not in per_agent:
per_agent[agent_name] = {
'total_tokens': 0,
'prompt_tokens': 0,
'completion_tokens': 0,
'task_count': 0
}
per_agent[agent_name]['total_tokens'] += int(tokens_per_task)
per_agent[agent_name]['prompt_tokens'] += int(prompt_per_task)
per_agent[agent_name]['completion_tokens'] += int(completion_per_task)
per_agent[agent_name]['task_count'] += 1
return per_agent
def create_test_crew():
"""Create a simple test crew with 2 agents and 2 tasks."""
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
raise ValueError("GEMINI_API_KEY not found in environment!")
logger.info("🔧 Creating LLM instance...")
llm = LLM(
model="gemini/gemini-2.0-flash-exp",
api_key=api_key
)
# Create agents
logger.info("🤖 Creating agents...")
researcher = Agent(
role="Research Analyst",
goal="Research and summarize information about a given topic",
backstory="You are an expert research analyst who excels at finding and summarizing information.",
llm=llm,
verbose=False, # Reduce verbosity for cleaner output
allow_delegation=False
)
writer = Agent(
role="Content Writer",
goal="Write engaging content based on research",
backstory="You are a skilled content writer who creates compelling narratives from research data.",
llm=llm,
verbose=False,
allow_delegation=False
)
# Create tasks
logger.info("📝 Creating tasks...")
research_task = Task(
description="Research the history of artificial intelligence and list 3 key milestones.",
expected_output="A list of 3 key milestones in AI history with brief descriptions.",
agent=researcher
)
writing_task = Task(
description="Write a brief 2-paragraph summary of AI history based on the research.",
expected_output="A 2-paragraph summary of AI history.",
agent=writer
)
# Create crew
logger.info("🚀 Creating crew...")
crew = Crew(
agents=[researcher, writer],
tasks=[research_task, writing_task],
process=Process.sequential,
verbose=False # Reduce verbosity
)
return crew
def main():
"""Run the token tracking test."""
print("\n" + "="*70)
print("CrewAI Token Tracking Test - Version 3 (Production Ready)")
print("="*70 + "\n")
# Create crew
print("Creating crew...")
crew = create_test_crew()
print("✓ Crew created\n")
# Run crew
print("-"*70)
print("Running CrewAI Workflow (this takes 1-2 minutes)...")
print("-"*70 + "\n")
try:
result = crew.kickoff()
logger.info("✅ Crew execution completed")
except Exception as e:
logger.error(f"❌ Crew execution failed: {e}", exc_info=True)
return
# Extract and display token usage
print("\n" + "="*70)
print("📊 Token Usage Report")
print("="*70 + "\n")
# 1. Total usage
if hasattr(result, 'token_usage') and result.token_usage:
usage_obj = result.token_usage
total_tokens = getattr(usage_obj, 'total_tokens', 0)
total_prompt = getattr(usage_obj, 'prompt_tokens', 0)
total_completion = getattr(usage_obj, 'completion_tokens', 0)
print("Total Crew Usage:")
print(f" Total Tokens: {total_tokens:,}")
print(f" Prompt Tokens: {total_prompt:,}")
print(f" Completion Tokens: {total_completion:,}")
# 2. Per-agent estimation
per_agent = estimate_per_agent_tokens(result)
if per_agent:
print("\nPer-Agent Breakdown (Estimated):")
print("─" * 70)
for agent_name, metrics in per_agent.items():
print(f"\n{agent_name}:")
print(f" Tasks Completed: {metrics['task_count']}")
print(f" Total Tokens: {metrics['total_tokens']:,}")
print(f" Prompt Tokens: {metrics['prompt_tokens']:,}")
print(f" Completion Tokens: {metrics['completion_tokens']:,}")
else:
print("\n⚠️ Could not estimate per-agent breakdown")
print("\n" + "="*70)
print("✅ Test Complete!")
print("="*70 + "\n")
print("Note: Per-agent tokens are estimated based on output length proportions.")
print("CrewAI v1.3.0 does not expose per-task usage_metrics.")
print("\n")
if __name__ == "__main__":
main()
1 Like