added sampling support

- add sampling
- vendor ClientSession from mcp
- improve stdio client
- add vscode launch profile

Author: SecretiveShell

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python Debugger: Module",
+            "type": "debugpy",
+            "request": "launch",
+            "django": true,
+            "module": "mcp_bridge.main",
+        }
+    ]
+}

README.md

@@ -23,6 +23,7 @@ working features:
 - non streaming completions without MCP
 
 - MCP tools
+- MCP sampling
 
 - SSE Bridge for external clients
 
@@ -129,6 +130,23 @@ an example config.json file with most of the options explicitly set:
         "base_url": "http://localhost:8000/v1",
         "api_key": "None"
     },
+    "sampling": {
+        "timeout": 10,
+        "models": [
+            {
+                "model": "gpt-4o",
+                "intelligence": 0.8,
+                "cost": 0.9,
+                "speed": 0.3
+            },
+            {
+                "model": "gpt-4o-mini",
+                "intelligence": 0.4,
+                "cost": 0.1,
+                "speed": 0.7
+            }
+        ]
+    },
     "mcp_servers": {
         "fetch": {
             "command": "uvx",

mcp_bridge/config/final.py

@@ -20,6 +20,18 @@ class Logging(BaseModel):
     log_server_pings: bool = Field(False, description="log server pings")
 
 
+class SamplingModel(BaseModel):
+    model: Annotated[str, Field(description="Name of the sampling model")]
+
+    intelligence: Annotated[float, Field(description="Intelligence of the sampling model")] = 0.5
+    cost: Annotated[float, Field(description="Cost of the sampling model")] = 0.5
+    speed: Annotated[float, Field(description="Speed of the sampling model")] = 0.5
+
+
+class Sampling(BaseModel):
+    timeout: Annotated[int, Field(description="Timeout for sampling requests")] = 10
+    models: Annotated[list[SamplingModel], Field(description="List of sampling models")] = []
+
 class SSEMCPServer(BaseModel):
     # TODO: expand this once I find a good definition for this
     url: str = Field(description="URL of the MCP server")
@@ -52,6 +64,11 @@ class Settings(BaseSettings):
         default_factory=dict, description="MCP servers configuration"
     )
 
+    sampling: Sampling = Field(
+        default_factory=lambda: Sampling.model_construct(),
+        description="sampling config",
+    )
+
     logging: Logging = Field(
         default_factory=lambda: Logging.model_construct(),
         description="logging config",

mcp_bridge/mcp_clients/AbstractClient.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, Optional
 from fastapi import HTTPException
-from mcp import ClientSession, McpError
+from mcp import McpError
 from mcp.types import (
     CallToolResult,
     ListToolsResult,
@@ -15,14 +15,15 @@
 )
 from loguru import logger
 from pydantic import AnyUrl
+from mcp_bridge.mcp_clients.session import McpClientSession
 from mcp_bridge.models.mcpServerStatus import McpServerStatus
 
 
 class GenericMcpClient(ABC):
     name: str
     config: Any
     client: Any
-    session: ClientSession | None = None
+    session: McpClientSession | None = None
 
     def __init__(self, name: str) -> None:
         super().__init__()
@@ -39,8 +40,10 @@ async def _session_maintainer(self):
         while True:
             try:
                 await self._maintain_session()
+            except FileNotFoundError as e:
+                logger.error(f"failed to maintain session for {self.name}: file {e.filename} not found.")
             except Exception as e:
-                logger.trace(f"failed to maintain session for {self.name}: {e}")
+                logger.error(f"failed to maintain session for {self.name}: {type(e)} {e.args}")
 
             logger.debug(f"restarting session for {self.name}")
             await asyncio.sleep(0.5)

mcp_bridge/mcp_clients/DockerClient.py

@@ -1,5 +1,7 @@
 import asyncio
 from mcp import ClientSession
+
+from mcp_bridge.mcp_clients.session import McpClientSession
 from .transports.docker import docker_client
 from mcp_bridge.config import config
 from mcp_bridge.config.final import DockerMCPServer
@@ -9,7 +11,6 @@
 
 class DockerClient(GenericMcpClient):
     config: DockerMCPServer
-    session: ClientSession | None = None
 
     def __init__(self, name: str, config: DockerMCPServer) -> None:
         super().__init__(name=name)
@@ -19,7 +20,7 @@ def __init__(self, name: str, config: DockerMCPServer) -> None:
     async def _maintain_session(self):
         async with docker_client(self.config) as client:
             logger.debug(f"made instance of docker client for {self.name}")
-            async with ClientSession(*client) as session:
+            async with McpClientSession(*client) as session:
                 await session.initialize()
                 logger.debug(f"finished initialise session for {self.name}")
                 self.session = session

mcp_bridge/mcp_clients/SseClient.py

@@ -1,8 +1,8 @@
 import asyncio
-from mcp import ClientSession
 from mcp.client.sse import sse_client
 from mcp_bridge.config import config
 from mcp_bridge.config.final import SSEMCPServer
+from mcp_bridge.mcp_clients.session import McpClientSession
 from .AbstractClient import GenericMcpClient
 from loguru import logger
 
@@ -17,7 +17,7 @@ def __init__(self, name: str, config: SSEMCPServer) -> None:
 
     async def _maintain_session(self):
         async with sse_client(self.config.url) as client:
-            async with ClientSession(*client) as session:
+            async with McpClientSession(*client) as session:
                 await session.initialize()
                 logger.debug(f"finished initialise session for {self.name}")
                 self.session = session

mcp_bridge/mcp_clients/StdioClient.py

@@ -1,7 +1,8 @@
 import asyncio
-from mcp import ClientSession, StdioServerParameters, stdio_client
+from mcp import StdioServerParameters, stdio_client
 
 from mcp_bridge.config import config
+from mcp_bridge.mcp_clients.session import McpClientSession
 from .AbstractClient import GenericMcpClient
 from loguru import logger
 import shutil
@@ -17,24 +18,28 @@ class StdioClient(GenericMcpClient):
     def __init__(self, name: str, config: StdioServerParameters) -> None:
         super().__init__(name=name)
 
+        # logger.debug(f"initializing settings for {name}: {config.command} {" ".join(config.args)}")
+
+        own_config = config.model_copy(deep=True)
+
         env = dict(os.environ.copy())
 
         env = {
             key: value for key, value in env.items()
             if not any(key.startswith(keyword) for keyword in venv_keywords)
         }
 
-        # logger.debug(f"env: {env}")
-
         if config.env is not None:
             env.update(config.env)
 
+        own_config.env = env
+
         command = shutil.which(config.command)
         if command is None:
             logger.error(f"could not find command {config.command}")
             exit(1)
 
-        own_config = config.model_copy(deep=True)
+        own_config.command = command
 
         # this changes the default to ignore
         if "encoding_error_handler" not in config.model_fields_set:
@@ -48,7 +53,7 @@ async def _maintain_session(self):
             logger.debug(f"entered stdio_client context manager for {self.name}")
             assert client[0] is not None, f"missing read stream for {self.name}"
             assert client[1] is not None, f"missing write stream for {self.name}"
-            async with ClientSession(*client) as session:
+            async with McpClientSession(*client) as session:
                 logger.debug(f"entered client session context manager for {self.name}")
                 await session.initialize()
                 logger.debug(f"finished initialise session for {self.name}")