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

mcp_agent/mcp/interfaces.py

@@ -0,0 +1,152 @@
+"""
+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 mcp import ClientSession
+from mcp.types import CreateMessageRequestParams
+from pydantic import Field
+
+
+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(
+        self,
+        server_name: str,
+        client_session_factory=None,
+        init_hook=None,
+    ) -> AsyncGenerator[ClientSession, None]:
+        """Initialize a server and yield a client session."""
+        ...
+
+    @property
+    def connection_manager(self) -> "ConnectionManagerProtocol":
+        """Get the connection manager."""
+        ...
+
+
+class ConnectionManagerProtocol(Protocol):
+    """
+    Protocol defining the minimal interface of ConnectionManager needed.
+    """
+
+    async def get_server(
+        self,
+        server_name: str,
+        client_session_factory=None,
+    ):
+        """Get a server connection."""
+        ...
+
+    async def disconnect_server(self, server_name: str) -> None:
+        """Disconnect from a server."""
+        ...
+
+    async def disconnect_all_servers(self) -> None:
+        """Disconnect from all servers."""
+        ...
+
+
+# Type variables for generic protocols
+MessageParamT = TypeVar("MessageParamT")
+"""A type representing an input message to an LLM."""
+
+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.
+    """
+
+    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]):
+    """Protocol defining the interface for augmented LLMs"""
+
+    async def generate(
+        self,
+        message: str | MessageParamT | List[MessageParamT],
+        request_params: RequestParams | None = None,
+    ) -> List[MessageT]:
+        """Request an LLM generation, which may run multiple iterations, and return the result"""
+
+    async def generate_str(
+        self,
+        message: str | MessageParamT | List[MessageParamT],
+        request_params: RequestParams | None = None,
+    ) -> str:
+        """Request an LLM generation and return the string representation of the result"""
+
+    async def generate_structured(
+        self,
+        message: str | MessageParamT | List[MessageParamT],
+        response_model: Type[ModelT],
+        request_params: RequestParams | None = None,
+    ) -> ModelT:
+        """Request a structured LLM generation and return the result as a Pydantic model."""
+
+
+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
+    ) -> Callable[..., AugmentedLLMProtocol[Any, Any]]:
+        """
+        Creates a factory function that can be used to construct an LLM instance.
+
+        Args:
+            model_string: The model specification string
+            request_params: Optional parameters to configure LLM behavior
+
+        Returns:
+            A factory function that can create an LLM instance
+        """
+        ...

mcp_agent/mcp/mcp_agent_client_session.py

@@ -7,7 +7,6 @@ from typing import Optional
 
 from mcp import ClientSession
 from mcp.shared.session import (
-    RequestResponder,
     ReceiveResultT,
     ReceiveNotificationT,
     RequestId,
@@ -16,26 +15,43 @@ from mcp.shared.session import (
     SendResultT,
 )
 from mcp.types import (
-    ClientResult,
-    CreateMessageRequest,
-    CreateMessageResult,
     ErrorData,
-    JSONRPCNotification,
-    JSONRPCRequest,
-    ServerRequest,
-    TextContent,
-    ListRootsRequest,
     ListRootsResult,
     Root,
 )
+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
 
 logger = get_logger(__name__)
 
 
+async def list_roots(ctx: ClientSession) -> ListRootsResult:
+    """List roots callback that will be called by the MCP library."""
+
+    roots = []
+    if (
+        hasattr(ctx, "session")
+        and hasattr(ctx.session, "server_config")
+        and ctx.session.server_config
+        and hasattr(ctx.session.server_config, "roots")
+        and ctx.session.server_config.roots
+    ):
+        roots = [
+            Root(
+                uri=AnyUrl(
+                    root.server_uri_alias or root.uri,
+                ),
+                name=root.name,
+            )
+            for root in ctx.session.server_config.roots
+        ]
+    return ListRootsResult(roots=roots or [])
+
+
 class MCPAgentClientSession(ClientSession, ContextDependent):
     """
     MCP Agent framework acts as a client to the servers providing tools/resources/prompts for the agent workloads.
@@ -48,36 +64,11 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
     """
 
     def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+        super().__init__(
+            *args, **kwargs, list_roots_callback=list_roots, sampling_callback=sample
+        )
         self.server_config: Optional[MCPServerSettings] = None
 
-    async def _received_request(
-        self, responder: RequestResponder[ServerRequest, ClientResult]
-    ) -> None:
-        logger.debug("Received request:", data=responder.request.model_dump())
-        request = responder.request.root
-
-        if isinstance(request, CreateMessageRequest):
-            return await self.handle_sampling_request(request, responder)
-        elif isinstance(request, ListRootsRequest):
-            # Handle list_roots request by returning configured roots
-            if hasattr(self, "server_config") and self.server_config.roots:
-                roots = [
-                    Root(
-                        uri=root.server_uri_alias or root.uri,
-                        name=root.name,
-                    )
-                    for root in self.server_config.roots
-                ]
-
-                await responder.respond(ListRootsResult(roots=roots))
-            else:
-                await responder.respond(ListRootsResult(roots=[]))
-            return
-
-        # Handle other requests as usual
-        await super()._received_request(responder)
-
     async def send_request(
         self,
         request: SendRequestT,
@@ -89,7 +80,7 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
             logger.debug("send_request: response=", data=result.model_dump())
             return result
         except Exception as e:
-            logger.error(f"send_request failed: {e}")
+            logger.error(f"send_request failed: {str(e)}")
             raise
 
     async def send_notification(self, notification: SendNotificationT) -> None:
@@ -132,111 +123,4 @@ class MCPAgentClientSession(ClientSession, ContextDependent):
         )
         return await super().send_progress_notification(
             progress_token=progress_token, progress=progress, total=total
-        )
-
-    async def _receive_loop(self) -> None:
-        async with (
-            self._read_stream,
-            self._write_stream,
-            self._incoming_message_stream_writer,
-        ):
-            async for message in self._read_stream:
-                if isinstance(message, Exception):
-                    await self._incoming_message_stream_writer.send(message)
-                elif isinstance(message.root, JSONRPCRequest):
-                    validated_request = self._receive_request_type.model_validate(
-                        message.root.model_dump(
-                            by_alias=True, mode="json", exclude_none=True
-                        )
-                    )
-                    responder = RequestResponder(
-                        request_id=message.root.id,
-                        request_meta=validated_request.root.params.meta
-                        if validated_request.root.params
-                        else None,
-                        request=validated_request,
-                        session=self,
-                    )
-
-                    await self._received_request(responder)
-                    if not responder._responded:
-                        await self._incoming_message_stream_writer.send(responder)
-                elif isinstance(message.root, JSONRPCNotification):
-                    notification = self._receive_notification_type.model_validate(
-                        message.root.model_dump(
-                            by_alias=True, mode="json", exclude_none=True
-                        )
-                    )
-
-                    await self._received_notification(notification)
-                    await self._incoming_message_stream_writer.send(notification)
-                else:  # Response or error
-                    stream = self._response_streams.pop(message.root.id, None)
-                    if stream:
-                        await stream.send(message.root)
-                    else:
-                        await self._incoming_message_stream_writer.send(
-                            RuntimeError(
-                                "Received response with an unknown "
-                                f"request ID: {message}"
-                            )
-                        )
-
-    async def handle_sampling_request(
-        self,
-        request: CreateMessageRequest,
-        responder: RequestResponder[ServerRequest, ClientResult],
-    ):
-        logger.info("Handling sampling request: %s", request)
-        config = self.context.config
-        session = self.context.upstream_session
-        if session is None:
-            # TODO: saqadri - consider whether we should be handling the sampling request here as a client
-            logger.warning(
-                "Error: No upstream client available for sampling requests. Request:",
-                data=request,
-            )
-            try:
-                from anthropic import AsyncAnthropic
-
-                client = AsyncAnthropic(api_key=config.anthropic.api_key)
-
-                params = request.params
-                response = await client.messages.create(
-                    model="claude-3-sonnet-20240229",
-                    max_tokens=params.maxTokens,
-                    messages=[
-                        {
-                            "role": m.role,
-                            "content": m.content.text
-                            if hasattr(m.content, "text")
-                            else m.content.data,
-                        }
-                        for m in params.messages
-                    ],
-                    system=getattr(params, "systemPrompt", None),
-                    temperature=getattr(params, "temperature", 0.7),
-                    stop_sequences=getattr(params, "stopSequences", None),
-                )
-
-                await responder.respond(
-                    CreateMessageResult(
-                        model="claude-3-sonnet-20240229",
-                        role="assistant",
-                        content=TextContent(type="text", text=response.content[0].text),
-                    )
-                )
-            except Exception as e:
-                logger.error(f"Error handling sampling request: {e}")
-                await responder.respond(ErrorData(code=-32603, message=str(e)))
-        else:
-            try:
-                # If a session is available, we'll pass-through the sampling request to the upstream client
-                result = await session.send_request(
-                    request=ServerRequest(request), result_type=CreateMessageResult
-                )
-
-                # Pass the result from the upstream client back to the server. We just act as a pass-through client here.
-                await responder.send_result(result)
-            except Exception as e:
-                await responder.send_error(code=-32603, message=str(e))
+        )

mcp_agent/mcp/mcp_aggregator.py

@@ -8,14 +8,15 @@ from typing import (
     Callable,
     TypeVar,
 )
-from mcp import GetPromptResult
-from pydantic import BaseModel, ConfigDict
+from mcp import GetPromptResult, ReadResourceResult
+from pydantic import AnyUrl, BaseModel, ConfigDict
 from mcp.client.session import ClientSession
 from mcp.server.lowlevel.server import Server
 from mcp.server.stdio import stdio_server
 from mcp.types import (
     CallToolResult,
     ListToolsResult,
+    TextContent,
     Tool,
     Prompt,
 )
@@ -210,6 +211,7 @@ class MCPAggregator(ContextDependent):
                         "agent_name": self.agent_name,
                     },
                 )
+
                 await self._persistent_connection_manager.get_server(
                     server_name, client_session_factory=MCPAgentClientSession
                 )
@@ -458,7 +460,10 @@ class MCPAggregator(ContextDependent):
 
         if server_name is None or local_tool_name is None:
             logger.error(f"Error: Tool '{name}' not found")
-            return CallToolResult(isError=True, message=f"Tool '{name}' not found")
+            return CallToolResult(
+                isError=True, 
+                content=[TextContent(type="text", text=f"Tool '{name}' not found")]
+            )
 
         logger.info(
             "Requesting tool call",
@@ -476,7 +481,10 @@ class MCPAggregator(ContextDependent):
             operation_name=local_tool_name,
             method_name="call_tool",
             method_args={"name": local_tool_name, "arguments": arguments},
-            error_factory=lambda msg: CallToolResult(isError=True, message=msg),
+            error_factory=lambda msg: CallToolResult(
+                isError=True, 
+                content=[TextContent(type="text", text=msg)]
+            ),
         )
 
     async def get_prompt(
@@ -821,6 +829,53 @@ class MCPAggregator(ContextDependent):
         logger.debug(f"Available prompts across servers: {results}")
         return results
 
+    async def get_resource(
+        self, server_name: str, resource_uri: str
+    ) -> ReadResourceResult:
+        """
+        Get a resource directly from an MCP server by URI.
+
+        Args:
+            server_name: Name of the MCP server to retrieve the resource from
+            resource_uri: URI of the resource to retrieve
+
+        Returns:
+            ReadResourceResult object containing the resource content
+
+        Raises:
+            ValueError: If the server doesn't exist or the resource couldn't be found
+        """
+        if not self.initialized:
+            await self.load_servers()
+
+        if server_name not in self.server_names:
+            raise ValueError(f"Server '{server_name}' not found")
+
+        logger.info(
+            "Requesting resource",
+            data={
+                "progress_action": ProgressAction.CALLING_TOOL,
+                "resource_uri": resource_uri,
+                "server_name": server_name,
+                "agent_name": self.agent_name,
+            },
+        )
+
+        try:
+            uri = AnyUrl(resource_uri)
+        except Exception as e:
+            raise ValueError(f"Invalid resource URI: {resource_uri}. Error: {e}")
+
+        # Use the _execute_on_server method to call read_resource on the server
+        return await self._execute_on_server(
+            server_name=server_name,
+            operation_type="resource",
+            operation_name=resource_uri,
+            method_name="read_resource",
+            method_args={"uri": uri},
+            error_factory=lambda msg: ValueError(f"Failed to retrieve resource: {msg}"),
+        )
+
 
 class MCPCompoundServer(Server):
     """
@@ -850,7 +905,10 @@ class MCPCompoundServer(Server):
             result = await self.aggregator.call_tool(name=name, arguments=arguments)
             return result.content
         except Exception as e:
-            return CallToolResult(isError=True, message=f"Error calling tool: {e}")
+            return CallToolResult(
+                isError=True, 
+                content=[TextContent(type="text", text=f"Error calling tool: {e}")]
+            )
 
     async def _get_prompt(
         self, name: str = None, arguments: dict[str, str] = None

mcp_agent/mcp/mcp_connection_manager.py

@@ -163,7 +163,6 @@ async def _server_lifecycle_task(server_conn: ServerConnection) -> None:
         async with transport_context as (read_stream, write_stream):
             #      try:
             server_conn.create_session(read_stream, write_stream)
-            #       except FileNotFoundError as e:
 
             async with server_conn.session:
                 await server_conn.initialize_session()

mcp_agent/mcp/prompts/prompt_server.py

@@ -12,6 +12,8 @@ import logging
 import sys
 from pathlib import Path
 from typing import List, Dict, Optional, Callable, Awaitable, Literal, Any
+from mcp.server.fastmcp.resources import FileResource
+from pydantic import AnyUrl
 
 from mcp_agent.mcp import mime_utils, resource_utils
 
@@ -185,19 +187,19 @@ def create_prompt_handler(
 
 
 # Type for resource handler
-ResourceHandler = Callable[[], Awaitable[str]]
+ResourceHandler = Callable[[], Awaitable[str | bytes]]
 
 
 def create_resource_handler(resource_path: Path, mime_type: str) -> ResourceHandler:
     """Create a resource handler function for the given resource"""
 
-    async def get_resource() -> str:
+    async def get_resource() -> str | bytes:
         is_binary = mime_utils.is_binary_content(mime_type)
 
         if is_binary:
             # For binary files, read in binary mode and base64 encode
             with open(resource_path, "rb") as f:
-                return base64.b64encode(f.read()).decode("utf-8")
+                return f.read()
         else:
             # For text files, read as utf-8 text
             with open(resource_path, "r", encoding="utf-8") as f:
@@ -284,15 +286,14 @@ def register_prompt(file_path: Path):
                         exposed_resources[resource_id] = resource_file
                         mime_type = mime_utils.guess_mime_type(str(resource_file))
 
-                        # Register with the correct resource ID directly with MCP
-                        resource_handler = create_resource_handler(
-                            resource_file, mime_type
+                        mcp.add_resource(
+                            FileResource(
+                                uri=AnyUrl(resource_id),
+                                path=resource_file,
+                                mime_type=mime_type,
+                                is_binary=mime_utils.is_binary_content(mime_type),
+                            )
                         )
-                        mcp.resource(
-                            resource_id,
-                            description=f"Resource from {file_path.name}",
-                            mime_type=mime_type,
-                        )(resource_handler)
 
                         logger.info(
                             f"Registered resource: {resource_id} ({resource_file})"

mcp_agent/mcp/resource_utils.py

@@ -25,44 +25,6 @@ def find_resource_file(resource_path: str, prompt_files: List[Path]) -> Optional
     return None
 
 
-# TODO -- decide how to deal with this. Both Anthropic and OpenAI allow sending URLs in
-# input message
-# TODO -- used?
-# async def fetch_remote_resource(
-#     url: str, timeout: int = HTTP_TIMEOUT
-# ) -> ResourceContent:
-#     """
-#     Fetch a remote resource from a URL
-
-#     Returns:
-#         Tuple of (content, mime_type, is_binary)
-#         - content: Text content or base64-encoded binary content
-#         - mime_type: The MIME type of the resource
-#         - is_binary: Whether the content is binary (and base64-encoded)
-#     """
-
-#     async with httpx.AsyncClient(timeout=timeout) as client:
-#         response = await client.get(url)
-#         response.raise_for_status()
-
-#         # Get the content type or guess from URL
-#         mime_type = response.headers.get("content-type", "").split(";")[0]
-#         if not mime_type:
-#             mime_type = mime_utils.guess_mime_type(url)
-
-#         # Check if this is binary content
-#         is_binary = mime_utils.is_binary_content(mime_type)
-
-#         if is_binary:
-#             # For binary responses, get the binary content and base64 encode it
-#             content = base64.b64encode(response.content).decode("utf-8")
-#         else:
-#             # For text responses, just get the text
-#             content = response.text
-
-#         return content, mime_type, is_binary
-
-
 def load_resource_content(
     resource_path: str, prompt_files: List[Path]
 ) -> ResourceContent:
@@ -109,6 +71,36 @@ def create_resource_uri(path: str) -> str:
     return f"resource://fast-agent/{Path(path).name}"
 
 
+# Add this to your resource_utils.py module
+
+
+def create_resource_reference(uri: str, mime_type: str) -> "EmbeddedResource":
+    """
+    Create a reference to a resource without embedding its content directly.
+
+    This creates an EmbeddedResource that references another resource URI.
+    When the client receives this, it will make a separate request to fetch
+    the resource content using the provided URI.
+
+    Args:
+        uri: URI for the resource
+        mime_type: MIME type of the resource
+
+    Returns:
+        An EmbeddedResource object
+    """
+    from mcp.types import EmbeddedResource, TextResourceContents
+
+    # Create a resource reference
+    resource_contents = TextResourceContents(
+        uri=uri,
+        mimeType=mime_type,
+        text="",  # Empty text as we're just referencing
+    )
+
+    return EmbeddedResource(type="resource", resource=resource_contents)
+
+
 def create_embedded_resource(
     resource_path: str, content: str, mime_type: str, is_binary: bool = False
 ) -> EmbeddedResource:
@@ -149,6 +141,34 @@ def create_image_content(data: str, mime_type: str) -> ImageContent:
     )
 
 
+def create_blob_resource(
+    resource_path: str, content: str, mime_type: str
+) -> EmbeddedResource:
+    """Create an embedded resource for binary data"""
+    return EmbeddedResource(
+        type="resource",
+        resource=BlobResourceContents(
+            uri=resource_path,
+            mimeType=mime_type,
+            blob=content,  # Content should already be base64 encoded
+        ),
+    )
+
+
+def create_text_resource(
+    resource_path: str, content: str, mime_type: str
+) -> EmbeddedResource:
+    """Create an embedded resource for text data"""
+    return EmbeddedResource(
+        type="resource",
+        resource=TextResourceContents(
+            uri=resource_path,
+            mimeType=mime_type,
+            text=content,
+        ),
+    )
+
+
 def normalize_uri(uri_or_filename: str) -> str:
     """
     Normalize a URI or filename to ensure it's a valid URI.