Add support for sampling

examples/mcp_primitives/mcp_basic_agent/demo_server.py

@@ -1,54 +1,14 @@
 from mcp.server.fastmcp import FastMCP
-import datetime
+from mcp.types import ModelPreferences, ModelHint, SamplingMessage, TextContent
 import json
 
-# Store server start time
-SERVER_START_TIME = datetime.datetime.utcnow()
-
 mcp = FastMCP("Resource Demo MCP Server")
 
-# Define some static resources
-STATIC_RESOURCES = {
-    "demo://docs/readme": {
-        "name": "README",
-        "description": "A sample README file.",
-        "content_type": "text/markdown",
-        "content": "# Demo Resource Server\n\nThis is a sample README resource provided by the demo MCP server.",
-    },
-    "demo://data/users": {
-        "name": "User Data",
-        "description": "Sample user data in JSON format.",
-        "content_type": "application/json",
-        "content": json.dumps(
-            [
-                {"id": 1, "name": "Alice"},
-                {"id": 2, "name": "Bob"},
-                {"id": 3, "name": "Charlie"},
-            ],
-            indent=2,
-        ),
-    },
-}
-
 
 @mcp.resource("demo://docs/readme")
 def get_readme():
     """Provide the README file content."""
-    meta = STATIC_RESOURCES["demo://docs/readme"]
-    return meta["content"]
-
-
-@mcp.resource("demo://data/users")
-def get_users():
-    """Provide user data."""
-    meta = STATIC_RESOURCES["demo://data/users"]
-    return meta["content"]
-
-
-@mcp.resource("demo://{city}/weather")
-def get_weather(city: str) -> str:
-    """Provide a simple weather report for a given city."""
-    return f"It is sunny in {city} today!"
+    return "# Demo Resource Server\n\nThis is a sample README resource provided by the demo MCP server."
 
 
 @mcp.prompt()
@@ -60,6 +20,53 @@ def echo(message: str) -> str:
     return f"Prompt: {message}"
 
 
+@mcp.resource("demo://data/friends")
+def get_users():
+    """Provide my friend list."""
+    return (
+        json.dumps(
+            [
+                {"id": 1, "friend": "Alice"},
+            ],
+        ),
+    )
+
+
+@mcp.prompt()
+def get_haiku_prompt(topic: str) -> str:
+    """Get a haiku prompt about a given topic."""
+    return f"I am fascinated about {topic}. Can you generate a haiku combining {topic} + my friend name?"
+
+
+@mcp.tool()
+async def get_haiku(topic: str) -> str:
+    """Get a haiku about a given topic."""
+    haiku = await mcp.get_context().session.create_message(
+        messages=[
+            SamplingMessage(
+                role="user",
+                content=TextContent(
+                    type="text", text=f"Generate a haiku about {topic}."
+                ),
+            )
+        ],
+        system_prompt="You are a poet.",
+        max_tokens=100,
+        temperature=0.7,
+        model_preferences=ModelPreferences(
+            hints=[ModelHint(name="gpt-4o-mini")],
+            costPriority=0.1,
+            speedPriority=0.8,
+            intelligencePriority=0.1,
+        ),
+    )
+
+    if isinstance(haiku.content, TextContent):
+        return haiku.content.text
+    else:
+        return "Haiku generation failed, unexpected content type."
+
+
 def main():
     """Main entry point for the MCP server."""
     mcp.run()

examples/mcp_primitives/mcp_basic_agent/main.py

@@ -11,6 +11,8 @@
 )
 from mcp_agent.agents.agent import Agent
 from mcp_agent.workflows.llm.augmented_llm_openai import OpenAIAugmentedLLM
+from mcp_agent.human_input.handler import console_input_callback
+
 
 settings = Settings(
     execution_engine="asyncio",
@@ -30,7 +32,9 @@
 
 # Settings can either be specified programmatically,
 # or loaded from mcp_agent.config.yaml/mcp_agent.secrets.yaml
-app = MCPApp(name="mcp_basic_agent")  # settings=settings)
+app = MCPApp(
+    name="mcp_basic_agent", human_input_callback=console_input_callback
+)  # settings=settings)
 
 
 async def example_usage():
@@ -71,13 +75,16 @@ async def example_usage():
             )
 
             llm = await agent.attach_llm(OpenAIAugmentedLLM)
-            res = await llm.generate_str(
+            summary = await llm.generate_str(
                 [
                     "Summarise what are my prompts and resources?",
                     *combined_messages,
                 ]
             )
-            logger.info(f"Summary: {res}")
+            logger.info(f"Summary: {summary}")
+
+            haiku = await llm.generate_str("Write me a haiku")
+            logger.info(f"Haiku: {haiku}")
 
 
 if __name__ == "__main__":

src/mcp_agent/agents/agent.py

@@ -199,6 +199,8 @@ async def attach_llm(
                 value = getattr(self.llm, attr, None)
                 if value is not None:
                     span.set_attribute(f"llm.{attr}", value)
+
+            self.context.active_llm = self.llm
             return self.llm
 
     async def initialize(self, force: bool = False):

src/mcp_agent/core/context.py

@@ -73,6 +73,9 @@ class Context(BaseModel):
     # Use this flag to conditionally serialize expensive data for tracing
     tracing_enabled: bool = False
 
+    # Store the currently active LLM instance for MCP sampling callbacks
+    active_llm: Optional[Any] = None  # AugmentedLLM instance
+
     model_config = ConfigDict(
         extra="allow",
         arbitrary_types_allowed=True,  # Tell Pydantic to defer type evaluation

src/mcp_agent/mcp/mcp_agent_client_session.py

@@ -39,7 +39,6 @@
     Implementation,
     JSONRPCMessage,
     ServerRequest,
-    TextContent,
     ListRootsResult,
     NotificationParams,
     RequestParams,
@@ -58,6 +57,7 @@
     MCP_TOOL_NAME,
 )
 from mcp_agent.tracing.telemetry import get_tracer, record_attributes
+from mcp_agent.mcp.sampling_handler import SamplingHandler
 
 if TYPE_CHECKING:
     from mcp_agent.core.context import Context
@@ -108,6 +108,7 @@ def __init__(
         )
 
         self.server_config: Optional[MCPServerSettings] = None
+        self._sampling_handler = SamplingHandler(context=self.context)
 
         # Session ID handling for Streamable HTTP transport
         self._get_session_id_callback: Optional[Callable[[], str | None]] = None
@@ -320,45 +321,11 @@ async def _handle_sampling_callback(
         context: RequestContext["ClientSession", Any],
         params: CreateMessageRequestParams,
     ) -> CreateMessageResult | ErrorData:
-        logger.info("Handling sampling request: %s", params)
-        config = self.context.config
+        logger.info(f"Handling sampling request: {params}")
         server_session = self.context.upstream_session
         if server_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=params,
-            )
-            try:
-                from anthropic import AsyncAnthropic
-
-                client = AsyncAnthropic(api_key=config.anthropic.api_key)
-
-                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),
-                )
-
-                return 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}")
-                return ErrorData(code=-32603, message=str(e))
+            # Enhanced sampling with human approval workflow
+            return await self._sampling_handler.handle_sampling_with_human_approval(params)
         else:
             try:
                 # If a server_session is available, we'll pass-through the sampling request to the upstream client
@@ -376,6 +343,7 @@ async def _handle_sampling_callback(
             except Exception as e:
                 return ErrorData(code=-32603, message=str(e))
 
+
     async def _handle_list_roots_callback(
         self,
         context: RequestContext["ClientSession", Any],

src/mcp_agent/mcp/mcp_aggregator.py

@@ -1209,6 +1209,29 @@ def getter(item: NamespacedResource):
         # No match found
         return None, None
 
+    def _find_server_name_from_uri(self, uri: str) -> str:
+        """
+        Find the server name from a resource URI.
+
+        Args:
+            uri: The URI of the resource.
+
+        Returns:
+            Server name if found, None otherwise
+        """
+        capability_map = self._server_to_resource_map
+
+        def getter(item: NamespacedResource):
+            return str(item.resource.uri)
+
+        for server_name, resources in capability_map.items():
+            for resource in resources:
+                if uri == getter(resource):
+                    return server_name
+
+        # No match found
+        return None
+
     async def _start_server(self, server_name: str):
         if self.connection_persistence:
             logger.info(