MCP Implementation

Author: kostisd

mcp_config.json

@@ -0,0 +1,9 @@
+{
+  "mcpServers": {
+    "sql_generator": {
+      "command": "python",
+      "args": ["-m", "datu.mcp.tools.sql_generator"],
+      "env": { "PYTHONPATH": "." }
+    }
+  }
+}

pyproject.toml

@@ -38,6 +38,7 @@ dependencies = [
     "openai>=1.30.1",
     "fastmcp>=2.10.5",
     "mcp-use[search]>=1.3.7",
+    "onnxruntime==1.19.2 ; sys_platform == 'darwin' and platform_machine == 'x86_64'",
 ]
 
 [project.urls]
@@ -307,6 +308,12 @@ ignore_missing_imports = true
 [[tool.mypy.overrides]]
 module = "langchain_mcp_adapters.tools"
 ignore_missing_imports = true
+[[tool.mypy.overrides]]
+module = "mcp_use"
+ignore_missing_imports = true
+[[tool.mypy.overrides]]
+module = "mcp_use.*"
+ignore_missing_imports = true
 
 [tool.towncrier]
 directory = "changelog.d"

src/datu/app_config.py

@@ -15,6 +15,7 @@
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from datu.integrations.config import IntegrationConfigs
+from datu.mcp.config import MCPConfig
 from datu.services.config import SchemaRAGConfig
 
 
@@ -72,7 +73,11 @@ class DatuConfig(BaseSettings):
         simulate_llm_response (str): Whether to simulate LLM responses.
         schema_sample_limit (int): The maximum number of rows to sample from the schema.
         schema_categorical_threshold (int): The threshold for categorical columns in the schema.
+        enable_mcp (bool): Whether to enable MCP integration.
+        mcp (MCPConfig | None): Configuration settings for MCP integration.
         enable_schema_rag (bool): Enable RAG for schema extraction.
+        schema_rag (SchemaRAGConfig | None): Configuration settings for schema RAG.
+
 
     """
 
@@ -94,6 +99,11 @@ class DatuConfig(BaseSettings):
     schema_categorical_detection: bool = True
     schema_sample_limit: int = 1000
     schema_categorical_threshold: int = 10
+    enable_mcp: bool = True
+    mcp: MCPConfig | None = Field(
+        default_factory=MCPConfig,
+        description="Configuration settings for MCP integration.",
+    )
     enable_schema_rag: bool = False
     schema_rag: SchemaRAGConfig | None = Field(
         default_factory=SchemaRAGConfig,

src/datu/base/chat_schema.py

@@ -0,0 +1,27 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class ChatMessage(BaseModel):
+    """Represents a single message in a chat conversation.
+
+    Attributes:
+        role (str): The role of the message sender (e.g., "user", "assistant").
+        content (str): The content of the message.
+    """
+
+    role: str
+    content: str
+
+
+class ChatRequest(BaseModel):
+    """Represents a chat request containing a list of messages and an optional system prompt.
+
+    Attributes:
+        messages (List[ChatMessage]): A list of messages in the chat conversation.
+        system_prompt (Optional[str]): An optional system prompt to provide context for the conversation.
+    """
+
+    messages: List[ChatMessage]
+    system_prompt: Optional[str] = None

src/datu/base/llm_client.py

@@ -2,15 +2,34 @@
 
 from abc import ABC, abstractmethod
 
+from mcp_use import MCPClient
+
+from datu.app_config import settings
+
 
 class BaseLLMClient(ABC):
     """BaseLLMClient class to provide a common interface for LLM clients.
     This class serves as an abstract base class for all LLM clients,
     providing a common interface and shared functionality.
     """
 
+    def __init__(self):
+        """Initializes the BaseLLMClient.
+            Sets up the client and MCP client if enabled in the settings.
+
+        Attributes:
+            client: The LLM client instance.
+            mcp_client: The MCP client instance if MCP is enabled in the settings.
+            agent: The agent instance if applicable.
+        """
+        self.client = None
+        self.mcp_client = None
+        if settings.enable_mcp:
+            self.mcp_client = MCPClient.from_config_file(settings.mcp.config_file)
+        self.agent = None
+
     @abstractmethod
-    def chat_completion(self, messages: list, system_prompt: str | None = None) -> str:
+    async def chat_completion(self, messages: list, system_prompt: str | None = None) -> str:
         """Given a conversation (and an optional system prompt), returns the assistant's text response."""
 
     @abstractmethod

src/datu/llm_clients/openai_client.py

@@ -8,6 +8,7 @@
 from langchain_community.chat_message_histories import ChatMessageHistory
 from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
 from langchain_openai import ChatOpenAI
+from mcp_use import MCPAgent
 
 from datu.app_config import get_logger, settings
 from datu.base.llm_client import BaseLLMClient
@@ -65,15 +66,37 @@ class OpenAIClient(BaseLLMClient):
     """
 
     def __init__(self):
+        super().__init__()
         self.model = getattr(settings, "openai_model", "gpt-4o-mini")
         self.client = ChatOpenAI(
             api_key=settings.openai_api_key,
             model=self.model,
             temperature=settings.llm_temperature,
         )
         self.history = ChatMessageHistory()
+        if settings.enable_mcp:
+            if not self.mcp_client:
+                raise RuntimeError("MCP is enabled but mcp_client was not initialized. ")
+            try:
+                self.agent = MCPAgent(
+                    llm=self.client,
+                    client=self.mcp_client,
+                    max_steps=settings.mcp.max_steps,
+                    use_server_manager=settings.mcp.use_server_manager,
+                )
+            except Exception:
+                # Prefer failing early so misconfig doesn’t silently degrade behavior
+                logger.exception("Failed to construct MCPAgent with provided MCP settings.")
+                raise
+
+    async def chat(self, input_text: str) -> str:
+        response = await self.agent.run(
+            input_text,
+            max_steps=30,
+        )
+        return response
 
-    def chat_completion(self, messages: list[BaseMessage], system_prompt: str | None = None) -> str:
+    async def chat_completion(self, messages: list[BaseMessage], system_prompt: str | None = None) -> str:
         if settings.simulate_llm_response:
             return create_simulated_llm_response()
         if not messages:
@@ -114,9 +137,20 @@ def chat_completion(self, messages: list[BaseMessage], system_prompt: str | None
             )
 
         self.history.add_message(HumanMessage(content=last_user_message))
-        response = self.client.invoke(self.history.messages)
-        self.history.add_message(response)
-        return response.content if response else ""
+        # Convert entire history messages to plain text for chat()
+        # Adjust this if your llm_with_tools expects different format
+        input_text = "\n".join(msg.content for msg in self.history.messages if hasattr(msg, "content"))
+
+        response = await self.chat(input_text)
+
+        # Assuming response is a BaseMessage or similar with 'content'
+        if hasattr(response, "content"):
+            self.history.add_message(response)
+            return response.content
+        else:
+            # If response is plain text string
+            self.history.add_message(HumanMessage(content=response))
+            return response
 
     def fix_sql_error(self, sql_code: str, error_msg: str, loop_count: int) -> str:
         """Generates a corrected SQL query based on the provided SQL code and error message.

src/datu/main.py

@@ -11,18 +11,14 @@
 from fastapi.staticfiles import StaticFiles
 
 from datu.app_config import get_logger, settings
-from datu.mcp.launcher import launch_mcp_server
 from datu.routers import chat, metadata, transformations
 from datu.schema_extractor.schema_cache import load_schema_cache
 
 logger = get_logger(__name__)
 
 # Optionally load schema and graph-rag in cache for use in prompts or logging
 if settings.app_environment != "test":
-    if settings.enable_schema_rag:
-        launch_mcp_server("schema_rag_server")
-    else:
-        schema_data = load_schema_cache()
+    schema_data = load_schema_cache()
 
 
 # Create the FastAPI application instance.

src/datu/mcp/client.py

@@ -0,0 +1,23 @@
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+
+class MCPConfig(BaseSettings):
+    """Configuration settings for MCP integration.
+    This class uses Pydantic's BaseSettings to manage environment variables
+    and provides a structured way to access MCP configuration settings.
+
+    Args:
+        max_steps (int): The maximum number of steps for MCP operations.
+        use_server_manager (bool): Whether to use the server manager for MCP operations.
+        config_file (str): The file path for the MCP configuration file.
+
+    Attributes:
+        max_steps (int): The maximum number of steps for MCP operations.
+        use_server_manager (bool): Whether to use the server manager for MCP operations.
+        config_file (str): The file path for the MCP configuration file.
+    """
+
+    max_steps: int = Field(default=30, description="Maximum number of steps MCPAgent can run.")
+    use_server_manager: bool = Field(default=True, description="Whether to enable the Server Manager.")
+    config_file: str = Field(default="mcp_config.json", description="Path to the MCP configuration file.")