fast-agent-mcp 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

mcp_agent/logging/transport.py

@@ -6,12 +6,14 @@ Transports for the Logger module for MCP Agent, including:
 
 import asyncio
 import json
+import traceback
 from abc import ABC, abstractmethod
-from typing import Dict, List, Protocol
 from pathlib import Path
+from typing import Dict, List, Protocol
 
 import aiohttp
 from opentelemetry import trace
+from rich import print
 from rich.json import JSON
 from rich.text import Text
 
@@ -20,8 +22,6 @@ from mcp_agent.console import console
 from mcp_agent.logging.events import Event, EventFilter
 from mcp_agent.logging.json_serializer import JSONSerializer
 from mcp_agent.logging.listeners import EventListener, LifecycleAwareListener
-from rich import print
-import traceback
 
 
 class EventTransport(Protocol):
@@ -30,7 +30,7 @@ class EventTransport(Protocol):
     (Kafka, RabbitMQ, REST, etc.).
     """
 
-    async def send_event(self, event: Event):
+    async def send_event(self, event: Event) -> None:
         """
         Send an event to the external system.
         Args:
@@ -44,10 +44,10 @@ class FilteredEventTransport(EventTransport, ABC):
     Event transport that filters events based on a filter before sending.
     """
 
-    def __init__(self, event_filter: EventFilter | None = None):
+    def __init__(self, event_filter: EventFilter | None = None) -> None:
         self.filter = event_filter
 
-    async def send_event(self, event: Event):
+    async def send_event(self, event: Event) -> None:
         if not self.filter or self.filter.matches(event):
             await self.send_matched_event(event)
 
@@ -59,7 +59,7 @@ class FilteredEventTransport(EventTransport, ABC):
 class NoOpTransport(FilteredEventTransport):
     """Default transport that does nothing (purely local)."""
 
-    async def send_matched_event(self, event):
+    async def send_matched_event(self, event) -> None:
         """Do nothing."""
         pass
 
@@ -67,7 +67,7 @@ class NoOpTransport(FilteredEventTransport):
 class ConsoleTransport(FilteredEventTransport):
     """Simple transport that prints events to console."""
 
-    def __init__(self, event_filter: EventFilter | None = None):
+    def __init__(self, event_filter: EventFilter | None = None) -> None:
         super().__init__(event_filter=event_filter)
         # Use shared console instances
         self._serializer = JSONSerializer()
@@ -78,7 +78,7 @@ class ConsoleTransport(FilteredEventTransport):
             "error": "bold red",
         }
 
-    async def send_matched_event(self, event: Event):
+    async def send_matched_event(self, event: Event) -> None:
         # Map log levels to styles
         style = self.log_level_styles.get(event.type, "white")
 
@@ -114,7 +114,7 @@ class FileTransport(FilteredEventTransport):
         event_filter: EventFilter | None = None,
         mode: str = "a",
         encoding: str = "utf-8",
-    ):
+    ) -> None:
         """Initialize FileTransport.
 
         Args:
@@ -186,7 +186,7 @@ class HTTPTransport(FilteredEventTransport):
         batch_size: int = 100,
         timeout: float = 5.0,
         event_filter: EventFilter | None = None,
-    ):
+    ) -> None:
         super().__init__(event_filter=event_filter)
         self.endpoint = endpoint
         self.headers = headers or {}
@@ -198,14 +198,14 @@ class HTTPTransport(FilteredEventTransport):
         self._session: aiohttp.ClientSession | None = None
         self._serializer = JSONSerializer()
 
-    async def start(self):
+    async def start(self) -> None:
         """Initialize HTTP session."""
         if not self._session:
             self._session = aiohttp.ClientSession(
                 headers=self.headers, timeout=aiohttp.ClientTimeout(total=self.timeout)
             )
 
-    async def stop(self):
+    async def stop(self) -> None:
         """Close HTTP session and flush any remaining events."""
         if self.batch:
             await self._flush()
@@ -213,14 +213,14 @@ class HTTPTransport(FilteredEventTransport):
             await self._session.close()
             self._session = None
 
-    async def send_matched_event(self, event: Event):
+    async def send_matched_event(self, event: Event) -> None:
         """Add event to batch, flush if batch is full."""
         async with self.lock:
             self.batch.append(event)
             if len(self.batch) >= self.batch_size:
                 await self._flush()
 
-    async def _flush(self):
+    async def _flush(self) -> None:
         """Send batch of events to HTTP endpoint."""
         if not self.batch:
             return
@@ -266,7 +266,7 @@ class AsyncEventBus:
 
     _instance = None
 
-    def __init__(self, transport: EventTransport | None = None):
+    def __init__(self, transport: EventTransport | None = None) -> None:
         self.transport: EventTransport = transport or NoOpTransport()
         self.listeners: Dict[str, EventListener] = {}
         self._queue = asyncio.Queue()
@@ -290,7 +290,7 @@ class AsyncEventBus:
             # Update transport if provided
             cls._instance.transport = transport
         return cls._instance
-        
+
     @classmethod
     def reset(cls) -> None:
         """
@@ -302,11 +302,11 @@ class AsyncEventBus:
             # Signal shutdown
             cls._instance._running = False
             cls._instance._stop_event.set()
-            
+
             # Clear the singleton instance
             cls._instance = None
 
-    async def start(self):
+    async def start(self) -> None:
         """Start the event bus and all lifecycle-aware listeners."""
         if self._running:
             return
@@ -321,7 +321,7 @@ class AsyncEventBus:
         self._running = True
         self._task = asyncio.create_task(self._process_events())
 
-    async def stop(self):
+    async def stop(self) -> None:
         """Stop the event bus and all lifecycle-aware listeners."""
         if not self._running:
             return
@@ -369,7 +369,7 @@ class AsyncEventBus:
                 except Exception as e:
                     print(f"Error stopping listener: {e}")
 
-    async def emit(self, event: Event):
+    async def emit(self, event: Event) -> None:
         """Emit an event to all listeners and transport."""
         # Inject current tracing info if available
         span = trace.get_current_span()
@@ -387,15 +387,15 @@ class AsyncEventBus:
         # Then queue for listeners
         await self._queue.put(event)
 
-    def add_listener(self, name: str, listener: EventListener):
+    def add_listener(self, name: str, listener: EventListener) -> None:
         """Add a listener to the event bus."""
         self.listeners[name] = listener
 
-    def remove_listener(self, name: str):
+    def remove_listener(self, name: str) -> None:
         """Remove a listener from the event bus."""
         self.listeners.pop(name, None)
 
-    async def _process_events(self):
+    async def _process_events(self) -> None:
         """Process events from the queue until stopped."""
         while self._running:
             event = None

mcp_agent/mcp/gen_client.py

@@ -28,9 +28,7 @@ async def gen_client(
     For persistent connections, use connect() or MCPConnectionManager instead.
     """
     if not server_registry:
-        raise ValueError(
-            "Server registry not found in the context. Please specify one either on this method, or in the context."
-        )
+        raise ValueError("Server registry not found in the context. Please specify one either on this method, or in the context.")
 
     async with server_registry.initialize_server(
         server_name=server_name,
@@ -53,9 +51,7 @@ async def connect(
     If required, callers can specify their own message receive loop and ClientSession class constructor to customize further.
     """
     if not server_registry:
-        raise ValueError(
-            "Server registry not found in the context. Please specify one either on this method, or in the context."
-        )
+        raise ValueError("Server registry not found in the context. Please specify one either on this method, or in the context.")
 
     server_connection = await server_registry.connection_manager.get_server(
         server_name=server_name,
@@ -73,13 +69,9 @@ async def disconnect(
     Disconnect from the specified server. If server_name is None, disconnect from all servers.
     """
     if not server_registry:
-        raise ValueError(
-            "Server registry not found in the context. Please specify one either on this method, or in the context."
-        )
+        raise ValueError("Server registry not found in the context. Please specify one either on this method, or in the context.")
 
     if server_name:
-        await server_registry.connection_manager.disconnect_server(
-            server_name=server_name
-        )
+        await server_registry.connection_manager.disconnect_server(server_name=server_name)
     else:
         await server_registry.connection_manager.disconnect_all_servers()

mcp_agent/mcp/interfaces.py

@@ -3,130 +3,149 @@ Interface definitions to prevent circular imports.
 This module defines protocols (interfaces) that can be used to break circular dependencies.
 """
 
-from contextlib import asynccontextmanager
-from typing import Any, AsyncGenerator, Callable, Generic, List, Optional, Protocol, Type, TypeVar
-
+from datetime import timedelta
+from typing import (
+    Any,
+    AsyncContextManager,
+    Callable,
+    Generic,
+    List,
+    Optional,
+    Protocol,
+    Type,
+    TypeVar,
+    Union,
+    runtime_checkable,
+)
+
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from mcp import ClientSession
-from mcp.types import CreateMessageRequestParams
-from pydantic import Field
+from mcp.types import PromptMessage
 
+from mcp_agent.core.request_params import RequestParams
+from mcp_agent.mcp.prompt_message_multipart import PromptMessageMultipart
 
-class ServerRegistryProtocol(Protocol):
-    """
-    Protocol defining the minimal interface of ServerRegistry needed by gen_client.
-    This allows gen_client to depend on this protocol rather than the full ServerRegistry class.
-    """
 
-    @asynccontextmanager
-    async def initialize_server(
+@runtime_checkable
+class MCPConnectionManagerProtocol(Protocol):
+    """Protocol for MCPConnectionManager functionality needed by ServerRegistry."""
+
+    async def get_server(
         self,
         server_name: str,
-        client_session_factory=None,
-        init_hook=None,
-    ) -> AsyncGenerator[ClientSession, None]:
-        """Initialize a server and yield a client session."""
-        ...
+        client_session_factory: Optional[
+            Callable[
+                [
+                    MemoryObjectReceiveStream,
+                    MemoryObjectSendStream,
+                    Optional[timedelta],
+                ],
+                ClientSession,
+            ]
+        ] = None,
+    ) -> "ServerConnection": ...
 
-    @property
-    def connection_manager(self) -> "ConnectionManagerProtocol":
-        """Get the connection manager."""
-        ...
+    async def disconnect_server(self, server_name: str) -> None: ...
 
+    async def disconnect_all_servers(self) -> None: ...
 
-class ConnectionManagerProtocol(Protocol):
-    """
-    Protocol defining the minimal interface of ConnectionManager needed.
-    """
 
-    async def get_server(
+@runtime_checkable
+class ServerRegistryProtocol(Protocol):
+    """Protocol defining the minimal interface of ServerRegistry needed by gen_client."""
+
+    @property
+    def connection_manager(self) -> MCPConnectionManagerProtocol: ...
+
+    def initialize_server(
         self,
         server_name: str,
-        client_session_factory=None,
-    ):
-        """Get a server connection."""
+        client_session_factory: Optional[
+            Callable[
+                [
+                    MemoryObjectReceiveStream,
+                    MemoryObjectSendStream,
+                    Optional[timedelta],
+                ],
+                ClientSession,
+            ]
+        ] = None,
+        init_hook: Optional[Callable] = None,
+    ) -> AsyncContextManager[ClientSession]:
+        """Initialize a server and yield a client session."""
         ...
 
-    async def disconnect_server(self, server_name: str) -> None:
-        """Disconnect from a server."""
-        ...
 
-    async def disconnect_all_servers(self) -> None:
-        """Disconnect from all servers."""
-        ...
+class ServerConnection(Protocol):
+    """Protocol for server connection objects returned by MCPConnectionManager."""
 
+    @property
+    def session(self) -> ClientSession: ...
 
-# Type variables for generic protocols
-MessageParamT = TypeVar("MessageParamT")
-"""A type representing an input message to an LLM."""
 
+# Regular invariant type variables
+MessageParamT = TypeVar("MessageParamT")
 MessageT = TypeVar("MessageT")
-"""A type representing an output message from an LLM."""
-
 ModelT = TypeVar("ModelT")
-"""A type representing a structured output message from an LLM."""
-
-
-class RequestParams(CreateMessageRequestParams):
-    """
-    Parameters to configure the AugmentedLLM 'generate' requests.
-    """
-
-    messages: None = Field(exclude=True, default=None)
-    """
-    Ignored. 'messages' are removed from CreateMessageRequestParams 
-    to avoid confusion with the 'message' parameter on 'generate' method.
-    """
-
-    maxTokens: int = 2048
-    """The maximum number of tokens to sample, as requested by the server."""
 
-    model: str | None = None
-    """
-    The model to use for the LLM generation.
-    If specified, this overrides the 'modelPreferences' selection criteria.
-    """
-
-    use_history: bool = True
-    """
-    Include the message history in the generate request.
-    """
-
-    max_iterations: int = 10
-    """
-    The maximum number of iterations to run the LLM for.
-    """
+# Variance-annotated type variables
+MessageParamT_co = TypeVar("MessageParamT_co", contravariant=True)
+MessageT_co = TypeVar("MessageT_co")
 
-    parallel_tool_calls: bool = True
-    """
-    Whether to allow multiple tool calls per iteration.
-    Also known as multi-step tool use.
-    """
 
-
-class AugmentedLLMProtocol(Protocol, Generic[MessageParamT, MessageT]):
+class AugmentedLLMProtocol(Protocol, Generic[MessageParamT_co, MessageT_co]):
     """Protocol defining the interface for augmented LLMs"""
 
     async def generate(
         self,
-        message: str | MessageParamT | List[MessageParamT],
+        message: Union[str, MessageParamT_co, List[MessageParamT_co]],
         request_params: RequestParams | None = None,
-    ) -> List[MessageT]:
+    ) -> List[MessageT_co]:
         """Request an LLM generation, which may run multiple iterations, and return the result"""
+        ...
 
     async def generate_str(
         self,
-        message: str | MessageParamT | List[MessageParamT],
+        message: Union[str, MessageParamT_co, List[MessageParamT_co]],
         request_params: RequestParams | None = None,
     ) -> str:
         """Request an LLM generation and return the string representation of the result"""
+        ...
 
-    async def generate_structured(
+    async def structured(
         self,
-        message: str | MessageParamT | List[MessageParamT],
-        response_model: Type[ModelT],
-        request_params: RequestParams | None = None,
+        prompt: Union[str, PromptMessage, PromptMessageMultipart, List[str]],
+        model: Type[ModelT],
+        request_params: RequestParams | None,
     ) -> ModelT:
-        """Request a structured LLM generation and return the result as a Pydantic model."""
+        """Apply the prompt and return the result as a Pydantic model, or None if coercion fails"""
+        ...
+
+    async def generate_prompt(
+        self,
+        prompt: Union[str, PromptMessage, PromptMessageMultipart, List[str]],
+        request_params: RequestParams | None,
+    ) -> str:
+        """Request an LLM generation and return a string representation of the result"""
+        ...
+
+    async def apply_prompt(
+        self,
+        multipart_messages: List["PromptMessageMultipart"],
+        request_params: RequestParams | None = None,
+    ) -> str:
+        """
+        Apply a list of PromptMessageMultipart messages directly to the LLM.
+        This is a cleaner interface to _apply_prompt_template_provider_specific.
+
+        Args:
+            multipart_messages: List of PromptMessageMultipart objects
+            request_params: Optional parameters to configure the LLM request
+
+        Returns:
+            String representation of the assistant's response
+        """
+        ...
 
 
 class ModelFactoryClassProtocol(Protocol):
@@ -134,7 +153,7 @@ class ModelFactoryClassProtocol(Protocol):
     Protocol defining the minimal interface of the ModelFactory class needed by sampling.
     This allows sampling.py to depend on this protocol rather than the concrete ModelFactory class.
     """
-    
+
     @classmethod
     def create_factory(
         cls, model_string: str, request_params: Optional[RequestParams] = None
@@ -149,4 +168,4 @@ class ModelFactoryClassProtocol(Protocol):
         Returns:
             A factory function that can create an LLM instance
         """
-        ...
+        ...

mcp_agent/mcp/mcp_agent_client_session.py

@@ -3,12 +3,12 @@ A derived client session for the MCP Agent framework.
 It adds logging and supports sampling requests.
 """
 
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 from mcp import ClientSession
 from mcp.shared.session import (
-    ReceiveResultT,
     ReceiveNotificationT,
+    ReceiveResultT,
     RequestId,
     SendNotificationT,
     SendRequestT,
@@ -21,11 +21,13 @@ from mcp.types import (
 )
 from pydantic import AnyUrl
 
-from mcp_agent.config import MCPServerSettings
 from mcp_agent.context_dependent import ContextDependent
 from mcp_agent.logging.logger import get_logger
 from mcp_agent.mcp.sampling import sample
 
+if TYPE_CHECKING:
+    from mcp_agent.config import MCPServerSettings
+
 logger = get_logger(__name__)
 
 
@@ -63,10 +65,8 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
     Developers can extend this class to add more custom functionality as needed
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(
-            *args, **kwargs, list_roots_callback=list_roots, sampling_callback=sample
-        )
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs, list_roots_callback=list_roots, sampling_callback=sample)
         self.server_config: Optional[MCPServerSettings] = None
 
     async def send_request(
@@ -91,9 +91,7 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
             logger.error("send_notification failed", data=e)
             raise
 
-    async def _send_response(
-        self, request_id: RequestId, response: SendResultT | ErrorData
-    ) -> None:
+    async def _send_response(self, request_id: RequestId, response: SendResultT | ErrorData) -> None:
         logger.debug(
             f"send_response: request_id={request_id}, response=",
             data=response.model_dump(),
@@ -111,16 +109,10 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
         )
         return await super()._received_notification(notification)
 
-    async def send_progress_notification(
-        self, progress_token: str | int, progress: float, total: float | None = None
-    ) -> None:
+    async def send_progress_notification(self, progress_token: str | int, progress: float, total: float | None = None) -> None:
         """
         Sends a progress notification for a request that is currently being
         processed.
         """
-        logger.debug(
-            "send_progress_notification: progress_token={progress_token}, progress={progress}, total={total}"
-        )
-        return await super().send_progress_notification(
-            progress_token=progress_token, progress=progress, total=total
-        )
+        logger.debug("send_progress_notification: progress_token={progress_token}, progress={progress}, total={total}")
+        return await super().send_progress_notification(progress_token=progress_token, progress=progress, total=total)

mcp_agent/mcp/mcp_agent_server.py

@@ -1,7 +1,9 @@
 import asyncio
+
 from mcp.server import NotificationOptions
 from mcp.server.fastmcp import FastMCP
 from mcp.server.stdio import stdio_server
+
 from mcp_agent.executor.temporal import get_temporal_client
 from mcp_agent.telemetry.tracing import setup_tracing
 
@@ -10,27 +12,23 @@ app = FastMCP("mcp-agent-server")
 setup_tracing("mcp-agent-server")
 
 
-async def run():
+async def run() -> None:
     async with stdio_server() as (read_stream, write_stream):
         await app._mcp_server.run(
             read_stream,
             write_stream,
-            app._mcp_server.create_initialization_options(
-                notification_options=NotificationOptions(
-                    tools_changed=True, resources_changed=True
-                )
-            ),
+            app._mcp_server.create_initialization_options(notification_options=NotificationOptions(tools_changed=True, resources_changed=True)),
         )
 
 
 @app.tool
-async def run_workflow(query: str):
+async def run_workflow(query: str) -> None:
     """Run the workflow given its name or id"""
     pass
 
 
 @app.tool
-async def pause_workflow(workflow_id: str):
+async def pause_workflow(workflow_id: str) -> None:
     """Pause a running workflow."""
     temporal_client = await get_temporal_client()
     handle = temporal_client.get_workflow_handle(workflow_id)
@@ -38,14 +36,14 @@ async def pause_workflow(workflow_id: str):
 
 
 @app.tool
-async def resume_workflow(workflow_id: str):
+async def resume_workflow(workflow_id: str) -> None:
     """Resume a paused workflow."""
     temporal_client = await get_temporal_client()
     handle = temporal_client.get_workflow_handle(workflow_id)
     await handle.signal("resume")
 
 
-async def provide_user_input(workflow_id: str, input_data: str):
+async def provide_user_input(workflow_id: str, input_data: str) -> None:
     """Provide user/human input to a waiting workflow step."""
     temporal_client = await get_temporal_client()
     handle = temporal_client.get_workflow_handle(workflow_id)