From 6b800b7fff4718f45141a866300cee814bae80a4 Mon Sep 17 00:00:00 2001 From: Eduardo Chiarotti Date: Thu, 4 Jul 2024 22:23:19 -0300 Subject: [PATCH] fix: test new approach --- .../utilities/base_token_process.py | 2 +- src/crewai/crew.py | 26 ++++++++++--------- 2 files changed, 15 insertions(+), 13 deletions(-) diff --git a/src/crewai/agents/agent_builder/utilities/base_token_process.py b/src/crewai/agents/agent_builder/utilities/base_token_process.py index 42ac25ae0..ce0b446d3 100644 --- a/src/crewai/agents/agent_builder/utilities/base_token_process.py +++ b/src/crewai/agents/agent_builder/utilities/base_token_process.py @@ -8,7 +8,7 @@ class TokenProcess: successful_requests: int = 0 def sum_prompt_tokens(self, tokens: int): - self.prompt_tokens = 10 + self.prompt_tokens = self.prompt_tokens + tokens self.total_tokens = self.total_tokens + tokens def sum_completion_tokens(self, tokens: int): diff --git a/src/crewai/crew.py b/src/crewai/crew.py index a297bdff2..ada3cb55b 100644 --- a/src/crewai/crew.py +++ b/src/crewai/crew.py @@ -133,6 +133,15 @@ class Crew(BaseModel): default=False, description="output_log_file", ) + total_usage_metrics: dict = Field( + description="Total usage metrics for the crew.", + default={ + "total_tokens": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + "successful_requests": 0, + }, + ) @field_validator("id", mode="before") @classmethod @@ -595,25 +604,18 @@ class Crew(BaseModel): def calculate_usage_metrics(self) -> Dict[str, int]: """Calculates and returns the usage metrics.""" - total_usage_metrics = { - "total_tokens": 0, - "prompt_tokens": 5, - "completion_tokens": 0, - "successful_requests": 0, - } - for agent in self.agents: if hasattr(agent, "_token_process"): token_sum = agent._token_process.get_summary() - for key in total_usage_metrics: - total_usage_metrics[key] += token_sum.get(key, 0) + for key in self.total_usage_metrics: + self.total_usage_metrics[key] += token_sum.get(key, 0) if self.manager_agent and hasattr(self.manager_agent, "_token_process"): token_sum = self.manager_agent._token_process.get_summary() - for key in total_usage_metrics: - total_usage_metrics[key] += token_sum.get(key, 0) + for key in self.total_usage_metrics: + self.total_usage_metrics[key] += token_sum.get(key, 0) - return total_usage_metrics + return self.total_usage_metrics def __repr__(self): return f"Crew(id={self.id}, process={self.process}, number_of_agents={len(self.agents)}, number_of_tasks={len(self.tasks)})"