feat(llm): add token usage tracking and statistics for LLM providers and resolve the conflicts and modify files in src/llm/providers/archived to be consistent with the current main branch. (#92)

* support token usage stats

---------

Co-authored-by: Einsiedler <[email protected]>

Author: mmltt-star

src/core/orchestrator.py

@@ -391,6 +391,9 @@ async def run_sub_agent(
         # Start new sub-agent session
         self.task_log.start_sub_agent_session(sub_agent_name, task_description)
 
+        # Reset sub-agent usage stats for independent tracking
+        self.sub_agent_llm_client.reset_usage_stats()
+
         # Simplified initial user content (no file attachments)
         initial_user_content = [{"type": "text", "text": task_description}]
         message_history = [{"role": "user", "content": initial_user_content}]
@@ -667,6 +670,14 @@ async def run_sub_agent(
         ] = {"system_prompt": system_prompt, "message_history": message_history}  # type: ignore
         self.task_log.save()
 
+        # Record sub-agent cumulative usage
+        usage_log = self.sub_agent_llm_client.get_usage_log()
+        self.task_log.log_step(
+            "usage_calculation",
+            usage_log,
+            metadata={"session_id": self.task_log.current_sub_agent_session_id},
+        )
+
         self.task_log.end_sub_agent_session(sub_agent_name)
         self.task_log.log_step(
             "sub_agent_completed", f"Sub agent {sub_agent_name} completed", "info"
@@ -688,6 +699,9 @@ async def run_main_agent(
         if task_file_name:
             logger.debug(f"Associated File: {task_file_name}")
 
+        # Reset main agent usage stats for independent tracking
+        self.llm_client.reset_usage_stats()
+
         # 1. Process input
         initial_user_content, task_description = process_input(
             task_description, task_file_name
@@ -1103,6 +1117,14 @@ async def run_main_agent(
             "task_completed", f"Main agent task {task_id} completed successfully"
         )
 
+        # Record main agent cumulative usage
+        usage_log = self.llm_client.get_usage_log()
+        self.task_log.log_step(
+            "usage_calculation",
+            usage_log,
+            metadata={"session_id": "main_agent"},
+        )
+
         if "browsecomp-zh" in self.cfg.benchmark.name:
             return final_summary, final_summary
         else:

src/llm/provider_client_base.py

@@ -34,6 +34,11 @@ class LLMProviderClientBase(ABC):
 
     # post_init
     client: Any = dataclasses.field(init=False)
+    # Usage tracking - cumulative for each agent session
+    total_input_tokens: int = dataclasses.field(init=False, default=0)
+    total_input_cached_tokens: int = dataclasses.field(init=False, default=0)
+    total_output_tokens: int = dataclasses.field(init=False, default=0)
+    total_output_reasoning_tokens: int = dataclasses.field(init=False, default=0)
 
     def __post_init__(self):
         # Explicitly assign from cfg object
@@ -196,6 +201,21 @@ async def create_message(
             tool_definitions,
             keep_tool_result=keep_tool_result,
         )
+
+        # Accumulate usage for agent session
+        if response:
+            try:
+                usage = self._extract_usage_from_response(response)
+                if usage:
+                    self.total_input_tokens += usage.get("input_tokens", 0)
+                    self.total_input_cached_tokens += usage.get("cached_tokens", 0)
+                    self.total_output_tokens += usage.get("output_tokens", 0)
+                    self.total_output_reasoning_tokens += usage.get(
+                        "reasoning_tokens", 0
+                    )
+            except Exception as e:
+                logger.warning(f"Failed to accumulate usage: {e}")
+
         return response
 
     @staticmethod
@@ -315,3 +335,54 @@ def handle_max_turns_reached_summary_prompt(
         self, message_history: list[dict[str, Any]], summary_prompt: str
     ):
         raise NotImplementedError("must implement in subclass")
+
+    def _extract_usage_from_response(self, response):
+        """Default Extract usage - OpenAI Chat Completions format"""
+        if not hasattr(response, "usage"):
+            return {
+                "input_tokens": 0,
+                "cached_tokens": 0,
+                "output_tokens": 0,
+                "reasoning_tokens": 0,
+            }
+
+        usage = response.usage
+        prompt_tokens_details = getattr(usage, "prompt_tokens_details", {}) or {}
+        if hasattr(prompt_tokens_details, "to_dict"):
+            prompt_tokens_details = prompt_tokens_details.to_dict()
+        completion_tokens_details = (
+            getattr(usage, "completion_tokens_details", {}) or {}
+        )
+        if hasattr(completion_tokens_details, "to_dict"):
+            completion_tokens_details = completion_tokens_details.to_dict()
+
+        usage_dict = {
+            "input_tokens": getattr(usage, "prompt_tokens", 0),
+            "cached_tokens": prompt_tokens_details.get("cached_tokens", 0),
+            "output_tokens": getattr(usage, "completion_tokens", 0),
+            "reasoning_tokens": completion_tokens_details.get("reasoning_tokens", 0),
+        }
+
+        return usage_dict
+
+    def get_usage_log(self) -> str:
+        """Get cumulative usage for current agent session as formatted string"""
+        # Format: [Provider | Model] Total Input: X, Cache Input: Y, Output: Z, ...
+        provider_model = f"[{self.provider_class} | {self.model_name}]"
+        input_uncached = self.total_input_tokens - self.total_input_cached_tokens
+        output_response = self.total_output_tokens - self.total_output_reasoning_tokens
+        total_tokens = self.total_input_tokens + self.total_output_tokens
+
+        return (
+            f"Usage log: {provider_model}, "
+            f"Total Input: {self.total_input_tokens} (Cached: {self.total_input_cached_tokens}, Uncached: {input_uncached}), "
+            f"Total Output: {self.total_output_tokens} (Reasoning: {self.total_output_reasoning_tokens}, Response: {output_response}), "
+            f"Total Tokens: {total_tokens}"
+        )
+
+    def reset_usage_stats(self):
+        """Reset usage stats for new agent session"""
+        self.total_input_tokens = 0
+        self.total_input_cached_tokens = 0
+        self.total_output_tokens = 0
+        self.total_output_reasoning_tokens = 0

src/llm/providers/archived/gpt_openai_response_client.py

@@ -93,7 +93,7 @@ async def _create_message(
             response = self._convert_response_to_serializable(response)
 
             # Update token count
-            self._update_token_usage(response.get("usage", None))
+            # self._update_token_usage(response.get("usage", None))
             logger.debug(
                 f"LLM Response API call status: {response.get('error', 'N/A')}"
             )
@@ -269,3 +269,30 @@ def _convert_response_to_serializable(self, response):
         }
 
         return serializable_response
+
+    def _extract_usage_from_response(self, response):
+        """Extract usage - OpenAI Responses API format"""
+        if not response or not response.get("usage"):
+            return {
+                "input_tokens": 0,
+                "cached_tokens": 0,
+                "output_tokens": 0,
+                "reasoning_tokens": 0,
+            }
+
+        usage = response.get("usage", {}) or {}
+        input_tokens_details = usage.get("input_tokens_details", {}) or {}
+        if hasattr(input_tokens_details, "to_dict"):
+            input_tokens_details = input_tokens_details.to_dict()
+        output_tokens_details = usage.get("output_tokens_details", {}) or {}
+        if hasattr(output_tokens_details, "to_dict"):
+            output_tokens_details = output_tokens_details.to_dict()
+
+        usage_dict = {
+            "input_tokens": usage.get("input_tokens", 0),
+            "cached_tokens": input_tokens_details.get("cached_tokens", 0),
+            "output_tokens": usage.get("output_tokens", 0),
+            "reasoning_tokens": output_tokens_details.get("reasoning_tokens", 0),
+        }
+
+        return usage_dict

src/llm/providers/claude_anthropic_client.py

@@ -185,6 +185,33 @@ def handle_max_turns_reached_summary_prompt(self, message_history, summary_promp
         else:
             return summary_prompt
 
+    def _extract_usage_from_response(self, response):
+        """Extract usage - Anthropic format"""
+        if not hasattr(response, "usage"):
+            return {
+                "input_tokens": 0,
+                "cached_tokens": 0,
+                "output_tokens": 0,
+                "reasoning_tokens": 0,
+            }
+
+        usage = response.usage
+        cache_creation_input_tokens = getattr(usage, "cache_creation_input_tokens", 0)
+        cache_read_input_tokens = getattr(usage, "cache_read_input_tokens", 0)
+        input_tokens = getattr(usage, "input_tokens", 0)
+        output_tokens = getattr(usage, "output_tokens", 0)
+
+        usage_dict = {
+            "input_tokens": cache_creation_input_tokens
+            + cache_read_input_tokens
+            + input_tokens,
+            "cached_tokens": cache_read_input_tokens,
+            "output_tokens": output_tokens,
+            "reasoning_tokens": 0,
+        }
+
+        return usage_dict
+
     def _apply_cache_control(self, messages):
         """Apply cache control to the last user message and system message (if applicable)"""
         cached_messages = []

src/llm/providers/claude_openrouter_client.py

@@ -133,6 +133,8 @@ async def _create_message(
             if self.repetition_penalty != 1.0:
                 extra_body["repetition_penalty"] = self.repetition_penalty
 
+            extra_body["usage"] = {"include": True}
+
             params = {
                 "model": self.model_name,
                 "temperature": temperature,