fast-agent-mcp 0.2.33__py3-none-any.whl → 0.2.35__py3-none-any.whl

mcp_agent/llm/providers/augmented_llm_openai.py

@@ -31,6 +31,7 @@ from mcp_agent.llm.providers.multipart_converter_openai import OpenAIConverter,
 from mcp_agent.llm.providers.sampling_converter_openai import (
     OpenAISamplingConverter,
 )
+from mcp_agent.llm.usage_tracking import TurnUsage
 from mcp_agent.logging.logger import get_logger
 from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
 
@@ -90,15 +91,14 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
 
     def _initialize_default_params(self, kwargs: dict) -> RequestParams:
         """Initialize OpenAI-specific default parameters"""
+        # Get base defaults from parent (includes ModelDatabase lookup)
+        base_params = super()._initialize_default_params(kwargs)
+
+        # Override with OpenAI-specific settings
         chosen_model = kwargs.get("model", DEFAULT_OPENAI_MODEL)
+        base_params.model = chosen_model
 
-        return RequestParams(
-            model=chosen_model,
-            systemPrompt=self.instruction,
-            parallel_tool_calls=True,
-            max_iterations=20,
-            use_history=True,
-        )
+        return base_params
 
     def _base_url(self) -> str:
         return self.context.config.openai.base_url if self.context.config.openai else None
@@ -166,6 +166,19 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
 
             response = executor_result[0]
 
+            # Track usage if response is valid and has usage data
+            if (
+                hasattr(response, "usage")
+                and response.usage
+                and not isinstance(response, BaseException)
+            ):
+                try:
+                    model_name = self.default_request_params.model or DEFAULT_OPENAI_MODEL
+                    turn_usage = TurnUsage.from_openai(response.usage, model_name)
+                    self.usage_accumulator.add_turn(turn_usage)
+                except Exception as e:
+                    self.logger.warning(f"Failed to track usage: {e}")
+
             self.logger.debug(
                 "OpenAI completion response:",
                 data=response,

mcp_agent/llm/usage_tracking.py

@@ -0,0 +1,402 @@
+"""
+Usage tracking system for LLM providers with comprehensive cache support.
+
+This module provides unified usage tracking across Anthropic, OpenAI, and Google providers,
+including detailed cache metrics and context window management.
+"""
+
+import time
+from typing import List, Optional, Union
+
+# Proper type imports for each provider
+from anthropic.types import Usage as AnthropicUsage
+from google.genai.types import GenerateContentResponseUsageMetadata as GoogleUsage
+from openai.types.completion_usage import CompletionUsage as OpenAIUsage
+from pydantic import BaseModel, Field, computed_field
+
+from mcp_agent.llm.model_database import ModelDatabase
+from mcp_agent.llm.provider_types import Provider
+
+
+# Fast-agent specific usage type for synthetic providers
+class FastAgentUsage(BaseModel):
+    """Usage data for fast-agent providers (passthrough, playback, slow)"""
+
+    input_chars: int = Field(description="Characters in input messages")
+    output_chars: int = Field(description="Characters in output messages")
+    model_type: str = Field(description="Type of fast-agent model (passthrough/playbook/slow)")
+    tool_calls: int = Field(default=0, description="Number of tool calls made")
+    delay_seconds: float = Field(default=0.0, description="Artificial delays added")
+
+
+# Union type for raw usage data from any provider
+ProviderUsage = Union[AnthropicUsage, OpenAIUsage, GoogleUsage, FastAgentUsage]
+
+
+class ModelContextWindows:
+    """Context window sizes and cache configurations for various models"""
+
+    @classmethod
+    def get_context_window(cls, model: str) -> Optional[int]:
+        return ModelDatabase.get_context_window(model)
+
+
+class CacheUsage(BaseModel):
+    """Cache-specific usage metrics"""
+
+    cache_read_tokens: int = Field(default=0, description="Tokens read from cache")
+    cache_write_tokens: int = Field(default=0, description="Tokens written to cache")
+    cache_hit_tokens: int = Field(default=0, description="Total tokens served from cache")
+
+    @computed_field
+    @property
+    def total_cache_tokens(self) -> int:
+        """Total cache-related tokens"""
+        return self.cache_read_tokens + self.cache_write_tokens + self.cache_hit_tokens
+
+    @computed_field
+    @property
+    def has_cache_activity(self) -> bool:
+        """Whether any cache activity occurred"""
+        return self.total_cache_tokens > 0
+
+
+class TurnUsage(BaseModel):
+    """Usage data for a single turn/completion with cache support"""
+
+    provider: Provider
+    model: str
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+    timestamp: float = Field(default_factory=time.time)
+
+    # Cache-specific metrics
+    cache_usage: CacheUsage = Field(default_factory=CacheUsage)
+
+    # Provider-specific token types
+    tool_use_tokens: int = Field(default=0, description="Tokens used for tool calling prompts")
+    reasoning_tokens: int = Field(default=0, description="Tokens used for reasoning/thinking")
+
+    # Raw usage data from provider (preserves all original data)
+    raw_usage: ProviderUsage
+
+    @computed_field
+    @property
+    def current_context_tokens(self) -> int:
+        """Current context size after this turn (total input including cache + output)"""
+        # For Anthropic: input_tokens + cache_read_tokens represents total input context
+        total_input = self.input_tokens + self.cache_usage.cache_read_tokens + self.cache_usage.cache_write_tokens
+        return total_input + self.output_tokens
+
+    @computed_field
+    @property
+    def effective_input_tokens(self) -> int:
+        """Input tokens actually processed (new tokens, not from cache)"""
+        # For Anthropic: input_tokens already excludes cached content
+        # For other providers: subtract cache hits from input_tokens
+        if self.provider == Provider.ANTHROPIC:
+            return self.input_tokens
+        else:
+            return max(0, self.input_tokens - self.cache_usage.cache_hit_tokens)
+
+    @computed_field
+    @property
+    def display_input_tokens(self) -> int:
+        """Input tokens to display for 'Last turn' (total submitted tokens)"""
+        # For Anthropic: input_tokens excludes cache, so add cache tokens
+        if self.provider == Provider.ANTHROPIC:
+            return self.input_tokens + self.cache_usage.cache_read_tokens + self.cache_usage.cache_write_tokens
+        else:
+            # For OpenAI/Google: input_tokens already includes cached tokens
+            return self.input_tokens
+
+    @classmethod
+    def from_anthropic(cls, usage: AnthropicUsage, model: str) -> "TurnUsage":
+        # Extract cache tokens with proper null handling
+        cache_creation_tokens = getattr(usage, "cache_creation_input_tokens", 0) or 0
+        cache_read_tokens = getattr(usage, "cache_read_input_tokens", 0) or 0
+
+        cache_usage = CacheUsage(
+            cache_read_tokens=cache_read_tokens,  # Tokens read from cache (90% discount)
+            cache_write_tokens=cache_creation_tokens,  # Tokens written to cache (25% surcharge)
+        )
+
+        return cls(
+            provider=Provider.ANTHROPIC,
+            model=model,
+            input_tokens=usage.input_tokens,
+            output_tokens=usage.output_tokens,
+            total_tokens=usage.input_tokens + usage.output_tokens,
+            cache_usage=cache_usage,
+            raw_usage=usage,  # Store the original Anthropic usage object
+        )
+
+    @classmethod
+    def from_openai(cls, usage: OpenAIUsage, model: str) -> "TurnUsage":
+        # Extract cache tokens with proper null handling
+        cached_tokens = 0
+        if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
+            cached_tokens = getattr(usage.prompt_tokens_details, "cached_tokens", 0) or 0
+
+        cache_usage = CacheUsage(
+            cache_hit_tokens=cached_tokens  # These are tokens served from cache (50% discount)
+        )
+
+        return cls(
+            provider=Provider.OPENAI,
+            model=model,
+            input_tokens=usage.prompt_tokens,
+            output_tokens=usage.completion_tokens,
+            total_tokens=usage.total_tokens,
+            cache_usage=cache_usage,
+            raw_usage=usage,  # Store the original OpenAI usage object
+        )
+
+    @classmethod
+    def from_google(cls, usage: GoogleUsage, model: str) -> "TurnUsage":
+        # Extract token counts with proper null handling
+        prompt_tokens = getattr(usage, "prompt_token_count", 0) or 0
+        candidates_tokens = getattr(usage, "candidates_token_count", 0) or 0
+        total_tokens = getattr(usage, "total_token_count", 0) or 0
+        cached_content_tokens = getattr(usage, "cached_content_token_count", 0) or 0
+
+        # Extract additional Google-specific token types
+        tool_use_tokens = getattr(usage, "tool_use_prompt_token_count", 0) or 0
+        thinking_tokens = getattr(usage, "thoughts_token_count", 0) or 0
+
+        # Google cache tokens are read hits (75% discount on Gemini 2.5)
+        cache_usage = CacheUsage(cache_hit_tokens=cached_content_tokens)
+
+        return cls(
+            provider=Provider.GOOGLE,
+            model=model,
+            input_tokens=prompt_tokens,
+            output_tokens=candidates_tokens,
+            total_tokens=total_tokens,
+            cache_usage=cache_usage,
+            tool_use_tokens=tool_use_tokens,
+            reasoning_tokens=thinking_tokens,
+            raw_usage=usage,  # Store the original Google usage object
+        )
+
+    @classmethod
+    def from_fast_agent(cls, usage: FastAgentUsage, model: str) -> "TurnUsage":
+        # For fast-agent providers, we use characters as "tokens"
+        # This provides a consistent unit of measurement across all providers
+        input_tokens = usage.input_chars
+        output_tokens = usage.output_chars
+        total_tokens = input_tokens + output_tokens
+
+        # Fast-agent providers don't have cache functionality
+        cache_usage = CacheUsage()
+
+        return cls(
+            provider=Provider.FAST_AGENT,
+            model=model,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=total_tokens,
+            cache_usage=cache_usage,
+            raw_usage=usage,  # Store the original FastAgentUsage object
+        )
+
+
+class UsageAccumulator(BaseModel):
+    """Accumulates usage data across multiple turns with cache analytics"""
+
+    turns: List[TurnUsage] = Field(default_factory=list)
+    model: Optional[str] = None
+
+    def add_turn(self, turn: TurnUsage) -> None:
+        """Add a new turn to the accumulator"""
+        self.turns.append(turn)
+        if self.model is None:
+            self.model = turn.model
+
+    @computed_field
+    @property
+    def cumulative_input_tokens(self) -> int:
+        """Total input tokens charged across all turns (including cache tokens)"""
+        return sum(
+            turn.input_tokens + turn.cache_usage.cache_read_tokens + turn.cache_usage.cache_write_tokens
+            for turn in self.turns
+        )
+
+    @computed_field
+    @property
+    def cumulative_output_tokens(self) -> int:
+        """Total output tokens charged across all turns"""
+        return sum(turn.output_tokens for turn in self.turns)
+
+    @computed_field
+    @property
+    def cumulative_billing_tokens(self) -> int:
+        """Total tokens charged across all turns (including cache tokens)"""
+        return self.cumulative_input_tokens + self.cumulative_output_tokens
+
+    @computed_field
+    @property
+    def cumulative_cache_read_tokens(self) -> int:
+        """Total tokens read from cache across all turns"""
+        return sum(turn.cache_usage.cache_read_tokens for turn in self.turns)
+
+    @computed_field
+    @property
+    def cumulative_cache_write_tokens(self) -> int:
+        """Total tokens written to cache across all turns"""
+        return sum(turn.cache_usage.cache_write_tokens for turn in self.turns)
+
+    @computed_field
+    @property
+    def cumulative_cache_hit_tokens(self) -> int:
+        """Total tokens served from cache across all turns"""
+        return sum(turn.cache_usage.cache_hit_tokens for turn in self.turns)
+
+    @computed_field
+    @property
+    def cumulative_effective_input_tokens(self) -> int:
+        """Total input tokens excluding cache reads across all turns"""
+        return sum(turn.effective_input_tokens for turn in self.turns)
+
+    @computed_field
+    @property
+    def cumulative_tool_use_tokens(self) -> int:
+        """Total tokens used for tool calling prompts across all turns"""
+        return sum(turn.tool_use_tokens for turn in self.turns)
+
+    @computed_field
+    @property
+    def cumulative_reasoning_tokens(self) -> int:
+        """Total tokens used for reasoning/thinking across all turns"""
+        return sum(turn.reasoning_tokens for turn in self.turns)
+
+    @computed_field
+    @property
+    def cache_hit_rate(self) -> Optional[float]:
+        """Percentage of total input context served from cache"""
+        cache_tokens = self.cumulative_cache_read_tokens + self.cumulative_cache_hit_tokens
+        total_input_context = self.cumulative_input_tokens + cache_tokens
+        if total_input_context == 0:
+            return None
+        return (cache_tokens / total_input_context) * 100
+
+    @computed_field
+    @property
+    def current_context_tokens(self) -> int:
+        """Current context usage (last turn's context tokens)"""
+        if not self.turns:
+            return 0
+        return self.turns[-1].current_context_tokens
+
+    @computed_field
+    @property
+    def context_window_size(self) -> Optional[int]:
+        """Get context window size for current model"""
+        if self.model:
+            return ModelContextWindows.get_context_window(self.model)
+        return None
+
+    @computed_field
+    @property
+    def context_usage_percentage(self) -> Optional[float]:
+        """Percentage of context window used"""
+        window_size = self.context_window_size
+        if window_size and window_size > 0:
+            return (self.current_context_tokens / window_size) * 100
+        return None
+
+    @computed_field
+    @property
+    def turn_count(self) -> int:
+        """Number of turns accumulated"""
+        return len(self.turns)
+
+    def get_cache_summary(self) -> dict[str, Union[int, float, None]]:
+        """Get cache-specific metrics summary"""
+        return {
+            "cumulative_cache_read_tokens": self.cumulative_cache_read_tokens,
+            "cumulative_cache_write_tokens": self.cumulative_cache_write_tokens,
+            "cumulative_cache_hit_tokens": self.cumulative_cache_hit_tokens,
+            "cache_hit_rate_percent": self.cache_hit_rate,
+            "cumulative_effective_input_tokens": self.cumulative_effective_input_tokens,
+        }
+
+    def get_summary(self) -> dict[str, Union[int, float, str, None]]:
+        """Get comprehensive usage statistics"""
+        cache_summary = self.get_cache_summary()
+        return {
+            "model": self.model,
+            "turn_count": self.turn_count,
+            "cumulative_input_tokens": self.cumulative_input_tokens,
+            "cumulative_output_tokens": self.cumulative_output_tokens,
+            "cumulative_billing_tokens": self.cumulative_billing_tokens,
+            "cumulative_tool_use_tokens": self.cumulative_tool_use_tokens,
+            "cumulative_reasoning_tokens": self.cumulative_reasoning_tokens,
+            "current_context_tokens": self.current_context_tokens,
+            "context_window_size": self.context_window_size,
+            "context_usage_percentage": self.context_usage_percentage,
+            **cache_summary,
+        }
+
+
+# Utility functions for fast-agent integration
+def create_fast_agent_usage(
+    input_content: str,
+    output_content: str,
+    model_type: str,
+    tool_calls: int = 0,
+    delay_seconds: float = 0.0,
+) -> FastAgentUsage:
+    """
+    Create FastAgentUsage from message content.
+
+    Args:
+        input_content: Input message content
+        output_content: Output message content
+        model_type: Type of fast-agent model (passthrough/playback/slow)
+        tool_calls: Number of tool calls made
+        delay_seconds: Artificial delays added
+
+    Returns:
+        FastAgentUsage object with character counts
+    """
+    return FastAgentUsage(
+        input_chars=len(input_content),
+        output_chars=len(output_content),
+        model_type=model_type,
+        tool_calls=tool_calls,
+        delay_seconds=delay_seconds,
+    )
+
+
+def create_turn_usage_from_messages(
+    input_content: str,
+    output_content: str,
+    model: str,
+    model_type: str,
+    tool_calls: int = 0,
+    delay_seconds: float = 0.0,
+) -> TurnUsage:
+    """
+    Create TurnUsage directly from message content for fast-agent providers.
+
+    Args:
+        input_content: Input message content
+        output_content: Output message content
+        model: Model name (e.g., "passthrough", "playback", "slow")
+        model_type: Type for internal tracking
+        tool_calls: Number of tool calls made
+        delay_seconds: Artificial delays added
+
+    Returns:
+        TurnUsage object ready for accumulation
+    """
+    usage = create_fast_agent_usage(
+        input_content=input_content,
+        output_content=output_content,
+        model_type=model_type,
+        tool_calls=tool_calls,
+        delay_seconds=delay_seconds,
+    )
+    return TurnUsage.from_fast_agent(usage, model)

mcp_agent/logging/events.py

@@ -117,3 +117,27 @@ class SamplingFilter(EventFilter):
         if not super().matches(event):
             return False
         return random.random() < self.sample_rate
+
+
+class StreamingExclusionFilter(EventFilter):
+    """
+    Event filter that excludes streaming progress events from logs.
+    This prevents token count updates from flooding the logs when info level is enabled.
+    """
+
+    def matches(self, event: Event) -> bool:
+        # First check if it passes the base filter
+        if not super().matches(event):
+            return False
+        
+        # Exclude events with "Streaming progress" message
+        if event.message == "Streaming progress":
+            return False
+        
+        # Also check for events with progress_action = STREAMING in data
+        if event.data and isinstance(event.data.get("data"), dict):
+            event_data = event.data["data"]
+            if event_data.get("progress_action") == "Streaming":
+                return False
+        
+        return True

mcp_agent/logging/rich_progress.py

@@ -73,6 +73,7 @@ class RichProgressDisplay:
             ProgressAction.LOADED: "dim green",
             ProgressAction.INITIALIZED: "dim green",
             ProgressAction.CHATTING: "bold blue",
+            ProgressAction.STREAMING: "bold blue",  # Same color as chatting
             ProgressAction.ROUTING: "bold blue",
             ProgressAction.PLANNING: "bold blue",
             ProgressAction.READY: "dim green",
@@ -100,9 +101,16 @@ class RichProgressDisplay:
             task_id = self._taskmap[task_name]
 
         # Ensure no None values in the update
+        # For streaming, use custom description immediately to avoid flashing
+        if event.action == ProgressAction.STREAMING and event.streaming_tokens:
+            formatted_tokens = f"↓ {event.streaming_tokens.strip()}".ljust(15)
+            description = f"[{self._get_action_style(event.action)}]{formatted_tokens}"
+        else:
+            description = f"[{self._get_action_style(event.action)}]{event.action.value:<15}"
+            
         self._progress.update(
             task_id,
-            description=f"[{self._get_action_style(event.action)}]{event.action.value:<15}",
+            description=description,
             target=event.target or task_name,  # Use task_name as fallback for target
             details=event.details or "",
             task_name=task_name,

mcp_agent/mcp/interfaces.py

@@ -5,6 +5,7 @@ This module defines protocols (interfaces) that can be used to break circular de
 
 from datetime import timedelta
 from typing import (
+    TYPE_CHECKING,
     Any,
     AsyncContextManager,
     Callable,
@@ -31,6 +32,9 @@ from mcp_agent.core.agent_types import AgentType
 from mcp_agent.core.request_params import RequestParams
 from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
 
+if TYPE_CHECKING:
+    from mcp_agent.llm.usage_tracking import UsageAccumulator
+
 
 @runtime_checkable
 class MCPConnectionManagerProtocol(Protocol):
@@ -132,6 +136,8 @@ class AugmentedLLMProtocol(Protocol):
         """
         ...
 
+    usage_accumulator: "UsageAccumulator"
+
 
 class AgentProtocol(AugmentedLLMProtocol, Protocol):
     """Protocol defining the standard agent interface"""