""" Model Context Protocol (MCP) Integration Example with Real OpenAI Call and Dynamic Dispatch. Demonstrates: 1. Registering tool schemas and calling OpenAI via Infralo Gateway. 2. Dynamically dispatching tool calls based on the LLM's response. 3. Connecting to a real local MCP server spawned in a subprocess. 4. Executing the tool calls dynamically on the MCP session and wrapping them with Infralo. """ import asyncio import json import logging import os import sys from openai import OpenAI from infralo import Infralo # Configure standard logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger("infralo.mcp_example") # Load environment configuration INFRALO_API_KEY = os.environ.get( "INFRALO_API_KEY", "your_workspace_api_key" ) INFRALO_ENDPOINT = os.environ.get("INFRALO_ENDPOINT", "http://localhost:8000") MODEL_DEPLOYMENT = os.environ.get("INFRALO_MODEL_DEPLOYMENT", "your_model_name") # Initialize Infralo Client infralo = Infralo( api_key=INFRALO_API_KEY, endpoint=INFRALO_ENDPOINT, flush_interval=1.0, ) openai_client = OpenAI( base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1", api_key=INFRALO_API_KEY, ) # ── 1. MCP SERVER DEFINITION ──────────────────────────────────────────────── # If this file is run with the "--server" flag, it behaves as a stdio MCP server. from mcp.server.fastmcp import FastMCP mcp = FastMCP("Infralo Example Server") @mcp.tool() def get_factorial(n: int) -> int: """Calculate the factorial of a given integer n.""" import math if n < 0: raise ValueError("Factorial is not defined for negative numbers.") return math.factorial(n) @mcp.tool() def fetch_title(url: str) -> str: """Fetch the text content from a URL (mocked to avoid network dependencies).""" if "error" in url: raise RuntimeError("Network timeout connecting to server.") return f"Title: Welcome to {url}" if __name__ == "__main__" and len(sys.argv) > 1 and sys.argv[1] == "--server": # Run the server on stdio transport and exit mcp.run(transport="stdio") sys.exit(0) # ── 2. MCP CLIENT & TRACING INTEGRATION ────────────────────────────────────── import contextvars # ContextVar to hold the active MCP ClientSession so decorated tool functions # can access it without taking ClientSession as a parameter (which fails JSON serialization) mcp_session_var = contextvars.ContextVar("mcp_session_var", default=None) @infralo.tool(name="get_factorial") async def call_get_factorial(n: int) -> list[dict | str]: logger.info("Calling MCP tool 'get_factorial' with arguments n=%s", n) session = mcp_session_var.get() if session is None: raise RuntimeError("No active MCP session found.") res = await session.call_tool(name="get_factorial", arguments={"n": n}) previews = [block.text for block in res.content if getattr(block, "type", None) == "text"] output_preview = "\n".join(previews) if getattr(res, "isError", False): logger.warning("MCP tool 'get_factorial' returned an error: %s", output_preview) raise RuntimeError(output_preview) logger.info("MCP tool 'get_factorial' succeeded: %s", output_preview) return [ block.model_dump() if hasattr(block, "model_dump") else str(block) for block in res.content ] @infralo.tool(name="fetch_title") async def call_fetch_title(url: str) -> list[dict | str]: logger.info("Calling MCP tool 'fetch_title' with arguments url=%r", url) session = mcp_session_var.get() if session is None: raise RuntimeError("No active MCP session found.") res = await session.call_tool(name="fetch_title", arguments={"url": url}) previews = [block.text for block in res.content if getattr(block, "type", None) == "text"] output_preview = "\n".join(previews) if getattr(res, "isError", False): logger.warning("MCP tool 'fetch_title' returned an error: %s", output_preview) raise RuntimeError(output_preview) logger.info("MCP tool 'fetch_title' succeeded: %s", output_preview) return [ block.model_dump() if hasattr(block, "model_dump") else str(block) for block in res.content ] # OpenAI tools schema for LLM matching MCP server capabilities tools_schema = [ { "type": "function", "function": { "name": "get_factorial", "description": "Calculate the factorial of a given integer n.", "parameters": { "type": "object", "properties": { "n": {"type": "integer"}, }, "required": ["n"], }, }, }, { "type": "function", "function": { "name": "fetch_title", "description": "Fetch text content from a URL.", "parameters": { "type": "object", "properties": { "url": {"type": "string"}, }, "required": ["url"], }, }, }, ] async def run_client(): from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client print("=== Running Real MCP Server/Client Integration Example with Real LLM ===") # Define connection parameters to spawn ourselves in "--server" mode server_params = StdioServerParameters( command=sys.executable, args=[__file__, "--server"], ) # Start trace context with infralo.start_trace(session_id="mcp_live_session") as trace: print(f"Started trace: {trace.trace_id}") # Call LLM via Gateway with tools registered prompt = ( "Please call the get_factorial tool for n=5, call fetch_title for url 'example.com', " "and call fetch_title for url 'error-page.com' (which will fail)." ) try: raw_response = openai_client.chat.completions.with_raw_response.create( model=MODEL_DEPLOYMENT, messages=[{"role": "user", "content": prompt}], tools=tools_schema, extra_headers=trace.headers(), ) except Exception as exc: logger.error("Failed to connect to LLM gateway. Details: %s", exc) return # Read LLM span context from header trace.update_from_response(raw_response) response = raw_response.parse() # Connect to the MCP server async with stdio_client(server_params) as (read, write): async with ClientSession(read, write) as session: # Initialize the connection and set the ContextVar await session.initialize() token = mcp_session_var.set(session) try: # List tools to verify server is connected and working tools = await session.list_tools() logger.info( "Connected to MCP server. Available tools: %s", [t.name for t in tools.tools], ) # Map of MCP tool names to client-side decorated wrapper functions mcp_tools_map = { "get_factorial": call_get_factorial, "fetch_title": call_fetch_title, } # Process the tool calls returned by the LLM dynamically tool_calls = getattr(response.choices[0].message, "tool_calls", None) or [] for tool_call in tool_calls: name = tool_call.function.name arguments = json.loads(tool_call.function.arguments) call_id = tool_call.id tool_func = mcp_tools_map.get(name) if tool_func: try: # Invoke the decorated function with target tool call ID and MCP metadata await tool_func( _infralo_tool_call_id=call_id, _infralo_metadata={"protocol": "mcp", "transport": "stdio"}, **arguments, ) except Exception: pass # Trace records the failure automatically finally: mcp_session_var.reset(token) if __name__ == "__main__": asyncio.run(run_client()) infralo.shutdown()