fast-agent-mcp 0.2.40__py3-none-any.whl → 0.2.42__py3-none-any.whl

mcp_agent/llm/providers/augmented_llm_google_native.py

@@ -295,7 +295,7 @@ class GoogleNativeAugmentedLLM(AugmentedLLM[types.Content, types.Content]):
                         turn_usage = TurnUsage.from_google(
                             api_response.usage_metadata, request_params.model
                         )
-                        self.usage_accumulator.add_turn(turn_usage)
+                        self._finalize_turn_usage(turn_usage)
 
                     except Exception as e:
                         self.logger.warning(f"Failed to track usage: {e}")
@@ -439,6 +439,9 @@ class GoogleNativeAugmentedLLM(AugmentedLLM[types.Content, types.Content]):
         """
         Applies the prompt messages and potentially calls the LLM for completion.
         """
+        # Reset tool call counter for new turn
+        self._reset_turn_tool_calls()
+
         request_params = self.get_request_params(
             request_params=request_params
         )  # Get request params

mcp_agent/llm/providers/augmented_llm_openai.py

@@ -84,7 +84,9 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
         # TODO -- move this to model capabilities, add o4.
         chosen_model = self.default_request_params.model if self.default_request_params else None
         self._reasoning = chosen_model and (
-            chosen_model.startswith("o3") or chosen_model.startswith("o1")
+            chosen_model.startswith("o3")
+            or chosen_model.startswith("o1")
+            or chosen_model.startswith("o4")
         )
         if self._reasoning:
             self.logger.info(
@@ -108,6 +110,7 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
     def _openai_client(self) -> AsyncOpenAI:
         try:
             return AsyncOpenAI(api_key=self._api_key(), base_url=self._base_url())
+
         except AuthenticationError as e:
             raise ProviderKeyError(
                 "Invalid OpenAI API key",
@@ -355,7 +358,7 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
                 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)
+                    self._finalize_turn_usage(turn_usage)
                 except Exception as e:
                     self.logger.warning(f"Failed to track usage: {e}")
 
@@ -389,7 +392,7 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
             messages.append(message)
 
             message_text = message.content
-            if choice.finish_reason in ["tool_calls", "function_call"] and message.tool_calls:
+            if await self._is_tool_stop_reason(choice.finish_reason) and message.tool_calls:
                 if message_text:
                     await self.show_assistant_message(
                         message_text,
@@ -477,12 +480,18 @@ class OpenAIAugmentedLLM(AugmentedLLM[ChatCompletionMessageParam, ChatCompletion
 
         return responses
 
+    async def _is_tool_stop_reason(self, finish_reason: str) -> bool:
+        return True
+
     async def _apply_prompt_provider_specific(
         self,
         multipart_messages: List["PromptMessageMultipart"],
         request_params: RequestParams | None = None,
         is_template: bool = False,
     ) -> PromptMessageMultipart:
+        # Reset tool call counter for new turn
+        self._reset_turn_tool_calls()
+
         last_message = multipart_messages[-1]
 
         # Add all previous messages to history (or all messages if last is from assistant)

mcp_agent/llm/providers/augmented_llm_xai.py

@@ -0,0 +1,38 @@
+import os
+
+from mcp_agent.core.request_params import RequestParams
+from mcp_agent.llm.provider_types import Provider
+from mcp_agent.llm.providers.augmented_llm_openai import OpenAIAugmentedLLM
+
+XAI_BASE_URL = "https://api.x.ai/v1"
+DEFAULT_XAI_MODEL = "grok-3"
+
+
+class XAIAugmentedLLM(OpenAIAugmentedLLM):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(
+            *args, provider=Provider.XAI, **kwargs
+        )  # Properly pass args and kwargs to parent
+
+    def _initialize_default_params(self, kwargs: dict) -> RequestParams:
+        """Initialize xAI parameters"""
+        chosen_model = kwargs.get("model", DEFAULT_XAI_MODEL)
+
+        return RequestParams(
+            model=chosen_model,
+            systemPrompt=self.instruction,
+            parallel_tool_calls=False,
+            max_iterations=10,
+            use_history=True,
+        )
+
+    def _base_url(self) -> str:
+        base_url = os.getenv("XAI_BASE_URL", XAI_BASE_URL)
+        if self.context.config and self.context.config.xai:
+            base_url = self.context.config.xai.base_url
+
+        return base_url
+
+    async def _is_tool_stop_reason(self, finish_reason: str) -> bool:
+        # grok uses Null as the finish reason for tool calls?
+        return await super()._is_tool_stop_reason(finish_reason) or finish_reason is None

mcp_agent/llm/usage_tracking.py

@@ -78,6 +78,9 @@ class TurnUsage(BaseModel):
     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")
 
+    # Tool call count for this turn
+    tool_calls: int = Field(default=0, description="Number of tool calls made in this turn")
+
     # Raw usage data from provider (preserves all original data)
     raw_usage: ProviderUsage
 
@@ -86,7 +89,11 @@ class TurnUsage(BaseModel):
     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
+        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
@@ -106,11 +113,20 @@ class TurnUsage(BaseModel):
         """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
+            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
 
+    def set_tool_calls(self, count: int) -> None:
+        """Set the number of tool calls made in this turn"""
+        # Use object.__setattr__ since this is a Pydantic model
+        object.__setattr__(self, "tool_calls", count)
+
     @classmethod
     def from_anthropic(cls, usage: AnthropicUsage, model: str) -> "TurnUsage":
         # Extract cache tokens with proper null handling
@@ -219,7 +235,9 @@ class UsageAccumulator(BaseModel):
     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
+            turn.input_tokens
+            + turn.cache_usage.cache_read_tokens
+            + turn.cache_usage.cache_write_tokens
             for turn in self.turns
         )
 
@@ -247,6 +265,12 @@ class UsageAccumulator(BaseModel):
         """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_tool_calls(self) -> int:
+        """Total tool calls made across all turns"""
+        return sum(turn.tool_calls for turn in self.turns)
+
     @computed_field
     @property
     def cumulative_cache_hit_tokens(self) -> int:
@@ -333,6 +357,7 @@ class UsageAccumulator(BaseModel):
             "cumulative_billing_tokens": self.cumulative_billing_tokens,
             "cumulative_tool_use_tokens": self.cumulative_tool_use_tokens,
             "cumulative_reasoning_tokens": self.cumulative_reasoning_tokens,
+            "cumulative_tool_calls": self.cumulative_tool_calls,
             "current_context_tokens": self.current_context_tokens,
             "context_window_size": self.context_window_size,
             "context_usage_percentage": self.context_usage_percentage,

mcp_agent/logging/logger.py

@@ -8,6 +8,7 @@ Logger module for the MCP Agent, which provides:
 """
 
 import asyncio
+import logging
 import threading
 import time
 from contextlib import asynccontextmanager, contextmanager
@@ -206,6 +207,12 @@ class LoggingConfig:
         if cls._initialized:
             return
 
+        # Suppress boto3/botocore logging to prevent flooding
+        logging.getLogger('boto3').setLevel(logging.WARNING)
+        logging.getLogger('botocore').setLevel(logging.WARNING)
+        logging.getLogger('urllib3').setLevel(logging.WARNING)
+        logging.getLogger('s3transfer').setLevel(logging.WARNING)
+
         bus = AsyncEventBus.get(transport=transport)
 
         # Add standard listeners

mcp_agent/mcp/hf_auth.py

@@ -69,14 +69,29 @@ def should_add_hf_auth(url: str, existing_headers: Optional[Dict[str, str]]) ->
     """
     # Only add HF auth if:
     # 1. URL is a HuggingFace URL
-    # 2. No existing Authorization header is set
+    # 2. No existing Authorization/X-HF-Authorization header is set
     # 3. HF_TOKEN environment variable is available
     
     if not is_huggingface_url(url):
         return False
         
-    if existing_headers and "Authorization" in existing_headers:
-        return False
+    if existing_headers:
+        # Check if this is a .hf.space domain
+        try:
+            parsed = urlparse(url)
+            hostname = parsed.hostname
+            if hostname and hostname.endswith(".hf.space"):
+                # For .hf.space, check for X-HF-Authorization header
+                if "X-HF-Authorization" in existing_headers:
+                    return False
+            else:
+                # For other HF domains, check for Authorization header
+                if "Authorization" in existing_headers:
+                    return False
+        except Exception:
+            # Fallback to checking Authorization header
+            if "Authorization" in existing_headers:
+                return False
         
     return get_hf_token_from_env() is not None
 
@@ -101,6 +116,19 @@ def add_hf_auth_header(url: str, headers: Optional[Dict[str, str]]) -> Optional[
         
     # Create new headers dict or copy existing one
     result_headers = dict(headers) if headers else {}
-    result_headers["Authorization"] = f"Bearer {hf_token}"
+    
+    # Check if this is a .hf.space domain
+    try:
+        parsed = urlparse(url)
+        hostname = parsed.hostname
+        if hostname and hostname.endswith(".hf.space"):
+            # Use X-HF-Authorization for .hf.space domains
+            result_headers["X-HF-Authorization"] = f"Bearer {hf_token}"
+        else:
+            # Use standard Authorization header for other HF domains
+            result_headers["Authorization"] = f"Bearer {hf_token}"
+    except Exception:
+        # Fallback to standard Authorization header
+        result_headers["Authorization"] = f"Bearer {hf_token}"
     
     return result_headers

mcp_agent/mcp/mcp_agent_client_session.py

@@ -78,6 +78,8 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
         self.agent_model: str | None = kwargs.pop("agent_model", None)
         # Extract agent_name if provided
         self.agent_name: str | None = kwargs.pop("agent_name", None)
+        # Extract api_key if provided
+        self.api_key: str | None = kwargs.pop("api_key", None)
         # Extract custom elicitation handler if provided
         custom_elicitation_handler = kwargs.pop("elicitation_handler", None)
 

mcp_agent/mcp/mcp_aggregator.py

@@ -221,6 +221,7 @@ class MCPAggregator(ContextDependent):
                     agent_model: str | None = None
                     agent_name: str | None = None
                     elicitation_handler = None
+                    api_key: str | None = None
 
                     # Check if this aggregator is part of an Agent (which has config)
                     # Import here to avoid circular dependency
@@ -230,6 +231,7 @@ class MCPAggregator(ContextDependent):
                         agent_model = self.config.model
                         agent_name = self.config.name
                         elicitation_handler = self.config.elicitation_handler
+                        api_key = self.config.api_key
 
                     return MCPAgentClientSession(
                         read_stream,
@@ -238,6 +240,7 @@ class MCPAggregator(ContextDependent):
                         server_name=server_name,
                         agent_model=agent_model,
                         agent_name=agent_name,
+                        api_key=api_key,
                         elicitation_handler=elicitation_handler,
                         tool_list_changed_callback=self._handle_tool_list_changed,
                         **kwargs,  # Pass through any additional kwargs like server_config
@@ -292,6 +295,8 @@ class MCPAggregator(ContextDependent):
                     # Get agent's model and name if this aggregator is part of an agent
                     agent_model: str | None = None
                     agent_name: str | None = None
+                    elicitation_handler = None
+                    api_key: str | None = None
 
                     # Check if this aggregator is part of an Agent (which has config)
                     # Import here to avoid circular dependency
@@ -300,6 +305,8 @@ class MCPAggregator(ContextDependent):
                     if isinstance(self, BaseAgent):
                         agent_model = self.config.model
                         agent_name = self.config.name
+                        elicitation_handler = self.config.elicitation_handler
+                        api_key = self.config.api_key
 
                     return MCPAgentClientSession(
                         read_stream,
@@ -308,6 +315,8 @@ class MCPAggregator(ContextDependent):
                         server_name=server_name,
                         agent_model=agent_model,
                         agent_name=agent_name,
+                        api_key=api_key,
+                        elicitation_handler=elicitation_handler,
                         tool_list_changed_callback=self._handle_tool_list_changed,
                         **kwargs,  # Pass through any additional kwargs like server_config
                     )
@@ -957,58 +966,43 @@ class MCPAggregator(ContextDependent):
 
         async with self._refresh_lock:
             try:
+                # Create a factory function that will include our parameters
+                def create_session(read_stream, write_stream, read_timeout):
+                    # Get agent name if available
+                    agent_model: str | None = None
+                    agent_name: str | None = None
+                    elicitation_handler = None
+                    api_key: str | None = None
+
+                    # Import here to avoid circular dependency
+                    from mcp_agent.agents.base_agent import BaseAgent
+
+                    if isinstance(self, BaseAgent):
+                        agent_model = self.config.model
+                        agent_name = self.config.name
+                        elicitation_handler = self.config.elicitation_handler
+                        api_key = self.config.api_key
+
+                    return MCPAgentClientSession(
+                        read_stream,
+                        write_stream,
+                        read_timeout,
+                        server_name=server_name,
+                        agent_model=agent_model,
+                        agent_name=agent_name,
+                        api_key=api_key,
+                        elicitation_handler=elicitation_handler,
+                        tool_list_changed_callback=self._handle_tool_list_changed,
+                    )
+
                 # Fetch new tools from the server
                 if self.connection_persistence:
-                    # Create a factory function that will include our parameters
-                    def create_session(read_stream, write_stream, read_timeout):
-                        # Get agent name if available
-                        agent_name: str | None = None
-
-                        # Import here to avoid circular dependency
-                        from mcp_agent.agents.base_agent import BaseAgent
-
-                        if isinstance(self, BaseAgent):
-                            agent_name = self.config.name
-                            elicitation_handler = self.config.elicitation_handler
-
-                        return MCPAgentClientSession(
-                            read_stream,
-                            write_stream,
-                            read_timeout,
-                            server_name=server_name,
-                            agent_name=agent_name,
-                            elicitation_handler=elicitation_handler,
-                            tool_list_changed_callback=self._handle_tool_list_changed,
-                        )
-
                     server_connection = await self._persistent_connection_manager.get_server(
                         server_name, client_session_factory=create_session
                     )
                     tools_result = await server_connection.session.list_tools()
                     new_tools = tools_result.tools or []
                 else:
-                    # Create a factory function for the client session
-                    def create_session(read_stream, write_stream, read_timeout):
-                        # Get agent name if available
-                        agent_name: str | None = None
-
-                        # Import here to avoid circular dependency
-                        from mcp_agent.agents.base_agent import BaseAgent
-
-                        if isinstance(self, BaseAgent):
-                            agent_name = self.config.name
-                            elicitation_handler = self.config.elicitation_handler
-
-                        return MCPAgentClientSession(
-                            read_stream,
-                            write_stream,
-                            read_timeout,
-                            server_name=server_name,
-                            agent_name=agent_name,
-                            elicitation_handler=elicitation_handler,
-                            tool_list_changed_callback=self._handle_tool_list_changed,
-                        )
-
                     async with gen_client(
                         server_name,
                         server_registry=self.context.server_registry,

mcp_agent/mcp/sampling.py

@@ -20,7 +20,7 @@ logger = get_logger(__name__)
 
 
 def create_sampling_llm(
-    params: CreateMessageRequestParams, model_string: str
+    params: CreateMessageRequestParams, model_string: str, api_key: str | None
 ) -> AugmentedLLMProtocol:
     """
     Create an LLM instance for sampling without tools support.
@@ -52,7 +52,7 @@ def create_sampling_llm(
 
     # Create the LLM using the factory
     factory = ModelFactory.create_factory(model_string)
-    llm = factory(agent=agent)
+    llm = factory(agent=agent, api_key=api_key)
 
     # Attach the LLM to the agent
     agent._llm = llm
@@ -77,7 +77,8 @@ async def sample(mcp_ctx: ClientSession, params: CreateMessageRequestParams) ->
     Returns:
         A CreateMessageResult containing the LLM's response
     """
-    model = None
+    model: str | None = None
+    api_key: str | None = None
     try:
         # Extract model from server config using type-safe helper
         server_config = get_server_config(mcp_ctx)
@@ -104,13 +105,16 @@ async def sample(mcp_ctx: ClientSession, params: CreateMessageRequestParams) ->
                 from mcp_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
                 
                 # Try agent's model first (from the session)
-                if (hasattr(mcp_ctx, 'session') and 
-                    isinstance(mcp_ctx.session, MCPAgentClientSession) and
-                    mcp_ctx.session.agent_model):
-                    model = mcp_ctx.session.agent_model
-                    logger.debug(f"Using agent's model for sampling: {model}")
-                else:
-                    # Fall back to system default model
+                if hasattr(mcp_ctx, "session") and isinstance(mcp_ctx.session, MCPAgentClientSession):
+                    if mcp_ctx.session.agent_model:
+                        model = mcp_ctx.session.agent_model
+                        logger.debug(f"Using agent's model for sampling: {model}")
+                    if mcp_ctx.session.api_key:
+                        api_key = mcp_ctx.session.api_key
+                        logger.debug(f"Using agent's API KEY for sampling: {api_key}")
+
+                # Fall back to system default model
+                if model is None:
                     try:
                         if app_context and app_context.config and app_context.config.default_model:
                             model = app_context.config.default_model
@@ -122,7 +126,7 @@ async def sample(mcp_ctx: ClientSession, params: CreateMessageRequestParams) ->
             raise ValueError("No model configured for sampling (server config, agent model, or system default)")
 
         # Create an LLM instance
-        llm = create_sampling_llm(params, model)
+        llm = create_sampling_llm(params, model, api_key)
 
         # Extract all messages from the request params
         if not params.messages:

mcp_agent/resources/examples/mcp/elicitations/forms_demo.py

@@ -51,8 +51,6 @@ async def main():
         else:
             console.print("[red]No registration data received[/red]")
 
-        console.print("\n" + "─" * 50 + "\n")
-
         # Example 2: Product Review
         console.print("[bold yellow]Example 2: Product Review Form[/bold yellow]")
         console.print(
@@ -66,8 +64,6 @@ async def main():
             )
             console.print(review_panel)
 
-        console.print("\n" + "─" * 50 + "\n")
-
         # Example 3: Account Settings
         console.print("[bold yellow]Example 3: Account Settings Form[/bold yellow]")
         console.print(
@@ -81,8 +77,6 @@ async def main():
             )
             console.print(settings_panel)
 
-        console.print("\n" + "─" * 50 + "\n")
-
         # Example 4: Service Appointment
         console.print("[bold yellow]Example 4: Service Appointment Booking[/bold yellow]")
         console.print(

mcp_agent/resources/examples/workflows/router.py

@@ -7,6 +7,8 @@ Demonstrates router's ability to either:
 
 import asyncio
 
+from rich.console import Console
+
 from mcp_agent.core.fastagent import FastAgent
 
 # Create the application
@@ -45,7 +47,14 @@ SAMPLE_REQUESTS = [
     agents=["code_expert", "general_assistant", "fetcher"],
 )
 async def main() -> None:
+    console = Console()
+    console.print(
+        "\n[bright_red]Router Workflow Demo[/bright_red]\n\n"
+        "Enter a request to route it to the appropriate agent.\nEnter [bright_red]STOP[/bright_red] to run the demo, [bright_red]EXIT[/bright_red] to leave"
+    )
+
     async with fast.run() as agent:
+        await agent.interactive(agent_name="route")
         for request in SAMPLE_REQUESTS:
             await agent.route(request)