Mcp connectivity (#6)

Author: kostisd

CHANGELOG.md

@@ -0,0 +1,8 @@
+## [0.0.1](https://github.com/Datuanalytics/datu-core/tree/0.0.1) - 2025-08-18
+
+### Added
+
+- Basic LLM and postgres, msql integrations
+- MCP Connectivity
+- Prompt updated for welcoming and tool listing dialogue.
+- SQL generator is activated and connected by default.

changelog.d/+ab395eb1.added.md

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

mcp_config.json

@@ -27,14 +27,18 @@ dependencies = [
     "types-pyyaml>=6.0.12.20250402",
     "sql-metadata>=2.17.0",
     "sentence-transformers>=2.5.1",
+    "safetensors>=0.6.2,<0.7",
+    "transformers<4.44.0",
     "torch==2.2.2",
     "numpy<2.0.0",
     "networkx>=3.0",
     "langgraph>=0.0.37",
     "langchain-core>=0.1.41",
     "langchain-mcp-adapters>=0.1.8",
     "openai>=1.30.1",
-    "fastmcp>=2.10.5" 
+    "fastmcp>=2.10.5",
+    "mcp-use[search]>=1.3.7",
+    "onnxruntime==1.19.2 ; sys_platform == 'darwin' and platform_machine == 'x86_64'",
 ]
 
 [project.urls]
@@ -298,11 +302,17 @@ 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"
 filename = "CHANGELOG.md"
-package = "datu-core"
+package = "datu"
 package_dir = "src"
 start_string = "<!-- towncrier release notes start -->\n"
 underlines = ["", "", ""]

pyproject.toml

@@ -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 = False
+    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/app_config.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/chat_schema.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/base/llm_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,54 @@ class OpenAIClient(BaseLLMClient):
     """
 
     def __init__(self):
+        """Initializes the OpenAIClient with the configured model and API key."""
+        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()
+        self.agent = None
+        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:
+        """Sends a chat message to the MCP agent and returns the response.
+        Args:
+            input_text (str): The input text to send to the agent.
+        Returns:
+            str: The response from the agent."""
 
-    def chat_completion(self, messages: list[BaseMessage], system_prompt: str | None = None) -> str:
+        if not settings.enable_mcp or self.agent is None:
+            raise RuntimeError("chat() requires MCP enabled and an initialized agent.")
+        response = await self.agent.run(
+            input_text,
+            max_steps=30,
+        )
+        return response
+
+    async def chat_completion(self, messages: list[BaseMessage], system_prompt: str | None = None) -> str:
+        """Generates a chat completion response based on the provided messages and system prompt.
+        Args:
+            messages (list[BaseMessage]): A list of messages to send to the LLM.
+            system_prompt (str | None): An optional system prompt to guide the LLM's response.
+        Returns:
+            str: The generated response from the LLM.
+        """
         if settings.simulate_llm_response:
             return create_simulated_llm_response()
         if not messages:
@@ -114,9 +154,25 @@ 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"))
+
+        if settings.enable_mcp:
+            # uses MCP agent
+            response = await self.chat(input_text)
+        else:
+            # direct LLM call without MCP
+            response = await self.client.ainvoke(self.history.messages)
+
+        # 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/llm_clients/openai_client.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.