fast-agent-mcp 0.4.7__py3-none-any.whl

fast_agent/mcp/sampling.py

@@ -0,0 +1,200 @@
+"""
+This simplified implementation directly converts between MCP types and PromptMessageExtended.
+"""
+
+from typing import TYPE_CHECKING
+
+from mcp import ClientSession
+from mcp.types import CreateMessageRequestParams, CreateMessageResult, TextContent
+
+from fast_agent.agents.agent_types import AgentConfig
+from fast_agent.agents.llm_agent import LlmAgent
+from fast_agent.core.logging.logger import get_logger
+from fast_agent.interfaces import FastAgentLLMProtocol
+from fast_agent.llm.sampling_converter import SamplingConverter
+from fast_agent.mcp.helpers.server_config_helpers import get_server_config
+from fast_agent.types.llm_stop_reason import LlmStopReason
+
+if TYPE_CHECKING:
+    from fast_agent.types import PromptMessageExtended
+
+logger = get_logger(__name__)
+
+
+def create_sampling_llm(
+    params: CreateMessageRequestParams, model_string: str, api_key: str | None
+) -> FastAgentLLMProtocol:
+    """
+    Create an LLM instance for sampling without tools support.
+    This utility function creates a minimal LLM instance based on the model string.
+
+    Args:
+        mcp_ctx: The MCP ClientSession
+        model_string: The model to use (e.g. "passthrough", "claude-3-5-sonnet-latest")
+
+    Returns:
+        An initialized LLM instance ready to use
+    """
+    from fast_agent.llm.model_factory import ModelFactory
+
+    app_context = None
+    try:
+        from fast_agent.context import get_current_context
+
+        app_context = get_current_context()
+    except Exception:
+        logger.warning("App context not available for sampling call")
+
+    agent = LlmAgent(
+        config=sampling_agent_config(params),
+        context=app_context,
+    )
+
+    # Create the LLM using the factory
+    factory = ModelFactory.create_factory(model_string)
+    llm = factory(agent=agent, api_key=api_key)
+
+    # Attach the LLM to the agent
+    agent._llm = llm
+
+    return llm
+
+
+async def sample(mcp_ctx: ClientSession, params: CreateMessageRequestParams) -> CreateMessageResult:
+    """
+    Handle sampling requests from the MCP protocol using SamplingConverter.
+
+    This function:
+    1. Extracts the model from the request
+    2. Uses SamplingConverter to convert types
+    3. Calls the LLM's generate_prompt method
+    4. Returns the result as a CreateMessageResult
+
+    Args:
+        mcp_ctx: The MCP ClientSession
+        params: The sampling request parameters
+
+    Returns:
+        A CreateMessageResult containing the LLM's response
+    """
+    # Get server name for notification tracking
+    server_name = "unknown"
+    if hasattr(mcp_ctx, "session") and hasattr(mcp_ctx.session, "session_server_name"):
+        server_name = mcp_ctx.session.session_server_name or "unknown"
+
+    # Start tracking sampling operation
+    try:
+        from fast_agent.ui import notification_tracker
+        notification_tracker.start_sampling(server_name)
+    except Exception:
+        # Don't let notification tracking break sampling
+        pass
+
+    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)
+
+        # First priority: explicitly configured sampling model
+        if server_config and server_config.sampling:
+            model = server_config.sampling.model
+
+        # Second priority: auto_sampling fallback (if enabled at application level)
+        if model is None:
+            # Check if auto_sampling is enabled
+            auto_sampling_enabled = False
+            try:
+                from fast_agent.context import get_current_context
+
+                app_context = get_current_context()
+                if app_context and app_context.config:
+                    auto_sampling_enabled = getattr(app_context.config, "auto_sampling", True)
+            except Exception as e:
+                logger.debug(f"Could not get application config: {e}")
+                auto_sampling_enabled = True  # Default to enabled
+
+            if auto_sampling_enabled:
+                # Import here to avoid circular import
+                from fast_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
+                ):
+                    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
+                            logger.debug(f"Using system default model for sampling: {model}")
+                    except Exception as e:
+                        logger.debug(f"Could not get system default model: {e}")
+
+        if model is None:
+            raise ValueError(
+                "No model configured for sampling (server config, agent model, or system default)"
+            )
+
+        # Create an LLM instance
+        llm = create_sampling_llm(params, model, api_key)
+
+        # Extract all messages from the request params
+        if not params.messages:
+            raise ValueError("No messages provided")
+
+        # Convert all SamplingMessages to PromptMessageExtended objects
+        conversation = SamplingConverter.convert_messages(params.messages)
+
+        # Extract request parameters using our converter
+        request_params = SamplingConverter.extract_request_params(params)
+
+        llm_response: PromptMessageExtended = await llm.generate(conversation, request_params)
+        logger.info(f"Complete sampling request : {llm_response.first_text()[:50]}...")
+
+        return CreateMessageResult(
+            role=llm_response.role,
+            content=TextContent(type="text", text=llm_response.first_text()),
+            model=model,
+            stopReason=LlmStopReason.END_TURN.value,
+        )
+    except Exception as e:
+        logger.error(f"Error in sampling: {str(e)}")
+        return SamplingConverter.error_result(
+            error_message=f"Error in sampling: {str(e)}", model=model
+        )
+    finally:
+        # End tracking sampling operation
+        try:
+            from fast_agent.ui import notification_tracker
+            notification_tracker.end_sampling(server_name)
+        except Exception:
+            # Don't let notification tracking break sampling
+            pass
+
+
+def sampling_agent_config(
+    params: CreateMessageRequestParams | None = None,
+) -> AgentConfig:
+    """
+    Build a sampling AgentConfig based on request parameters.
+
+    Args:
+        params: Optional CreateMessageRequestParams that may contain a system prompt
+
+    Returns:
+        An initialized AgentConfig for use in sampling
+    """
+    # Use systemPrompt from params if available, otherwise use default
+    instruction = "You are a helpful AI Agent."
+    if params and params.systemPrompt is not None:
+        instruction = params.systemPrompt
+
+    return AgentConfig(name="sampling_agent", instruction=instruction)

fast_agent/mcp/server/__init__.py

@@ -0,0 +1,4 @@
+# Import and re-export AgentMCPServer to avoid circular imports
+from fast_agent.mcp.server.agent_server import AgentMCPServer
+
+__all__ = ["AgentMCPServer"]