MindsDB 25.4.3.1__py3-none-any.whl → 25.4.4.0__py3-none-any.whl

mindsdb/interfaces/agents/mcp_client_agent.py

@@ -0,0 +1,252 @@
+import json
+import asyncio
+from typing import Dict, List, Any, Iterator, ClassVar
+from contextlib import AsyncExitStack
+
+import pandas as pd
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from mindsdb.utilities import log
+from mindsdb.interfaces.agents.langchain_agent import LangchainAgent
+from mindsdb.interfaces.storage import db
+from langchain_core.tools import BaseTool
+
+logger = log.getLogger(__name__)
+
+
+class MCPQueryTool(BaseTool):
+    """Tool that executes queries via MCP server"""
+
+    name: ClassVar[str] = "mcp_query"
+    description: ClassVar[str] = "Execute SQL queries against the MindsDB server via MCP protocol"
+
+    def __init__(self, session: ClientSession):
+        super().__init__()
+        self.session = session
+
+    async def _arun(self, query: str) -> str:
+        """Execute a query via MCP asynchronously"""
+        try:
+            logger.info(f"Executing MCP query: {query}")
+            # Find the appropriate tool for SQL queries
+            tools_response = await self.session.list_tools()
+            query_tool = None
+
+            for tool in tools_response.tools:
+                if tool.name == "query":
+                    query_tool = tool
+                    break
+
+            if not query_tool:
+                return "Error: No 'query' tool found in the MCP server"
+
+            # Call the query tool
+            result = await self.session.call_tool("query", {"query": query})
+
+            # Process the results
+            if isinstance(result.content, dict) and "data" in result.content and "column_names" in result.content:
+                # Create a DataFrame from the results
+                df = pd.DataFrame(result.content["data"], columns=result.content["column_names"])
+                return df.to_string()
+
+            # Return raw result for other types
+            return f"Query executed successfully: {json.dumps(result.content)}"
+
+        except Exception as e:
+            logger.error(f"Error executing MCP query: {str(e)}")
+            return f"Error executing query: {str(e)}"
+
+    def _run(self, query: str) -> str:
+        """Synchronous wrapper for async query function"""
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(self._arun(query))
+
+
+class MCPLangchainAgent(LangchainAgent):
+    """Extension of LangchainAgent that delegates to MCP server"""
+
+    def __init__(self, agent: db.Agents, model: dict = None, mcp_host: str = "127.0.0.1", mcp_port: int = 47337):
+        super().__init__(agent, model)
+        self.mcp_host = mcp_host
+        self.mcp_port = mcp_port
+        self.exit_stack = AsyncExitStack()
+        self.session = None
+        self.stdio = None
+        self.write = None
+
+    async def connect_to_mcp(self):
+        """Connect to the MCP server using stdio transport"""
+        if self.session is None:
+            logger.info(f"Connecting to MCP server at {self.mcp_host}:{self.mcp_port}")
+            try:
+                # For connecting to an already running MCP server
+                # Set up server parameters to connect to existing process
+                server_params = StdioServerParameters(
+                    command="python",
+                    args=["-m", "mindsdb", "--api=mcp"],
+                    env={"MCP_HOST": self.mcp_host, "MCP_PORT": str(self.mcp_port)}
+                )
+
+                logger.info(f"Connecting to MCP server at {self.mcp_host}:{self.mcp_port}")
+
+                # Connect to the server
+                stdio_transport = await self.exit_stack.enter_async_context(stdio_client(server_params))
+                self.stdio, self.write = stdio_transport
+                self.session = await self.exit_stack.enter_async_context(ClientSession(self.stdio, self.write))
+
+                await self.session.initialize()
+
+                # Test the connection by listing tools
+                tools_response = await self.session.list_tools()
+                logger.info(f"Successfully connected to MCP server. Available tools: {[tool.name for tool in tools_response.tools]}")
+
+            except Exception as e:
+                logger.error(f"Failed to connect to MCP server: {str(e)}")
+                raise ConnectionError(f"Failed to connect to MCP server: {str(e)}")
+
+    def _langchain_tools_from_skills(self, llm):
+        """Override to add MCP query tool along with other tools"""
+        # Get tools from parent implementation
+        tools = super()._langchain_tools_from_skills(llm)
+
+        # Initialize MCP connection
+        try:
+            # Using the event loop directly instead of asyncio.run()
+            loop = asyncio.get_event_loop()
+            if self.session is None:
+                loop.run_until_complete(self.connect_to_mcp())
+
+            # Add MCP query tool if session is established
+            if self.session:
+                tools.append(MCPQueryTool(self.session))
+                logger.info("Added MCP query tool to agent tools")
+        except Exception as e:
+            logger.error(f"Failed to add MCP query tool: {str(e)}")
+
+        return tools
+
+    def get_completion(self, messages, stream: bool = False):
+        """Override to ensure MCP connection is established before getting completion"""
+        try:
+            # Ensure connection to MCP is established
+            if self.session is None:
+                # Using the event loop directly instead of asyncio.run()
+                loop = asyncio.get_event_loop()
+                loop.run_until_complete(self.connect_to_mcp())
+        except Exception as e:
+            logger.error(f"Failed to connect to MCP server: {str(e)}")
+
+        # Call parent implementation to get completion
+        response = super().get_completion(messages, stream)
+
+        # Ensure response is a string (not a DataFrame)
+        if hasattr(response, 'to_string'):  # It's a DataFrame
+            return response.to_string()
+
+        return response
+
+    async def cleanup(self):
+        """Clean up resources"""
+        if self.exit_stack:
+            await self.exit_stack.aclose()
+            self.session = None
+            self.stdio = None
+            self.write = None
+
+
+class LiteLLMAgentWrapper:
+    """Wrapper for MCPLangchainAgent that provides LiteLLM-compatible interface"""
+
+    def __init__(self, agent: MCPLangchainAgent):
+        self.agent = agent
+
+    async def acompletion(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
+        """Async completion interface compatible with LiteLLM"""
+        # Convert messages to format expected by agent
+        formatted_messages = [
+            {
+                "question": msg["content"] if msg["role"] == "user" else "",
+                "answer": msg["content"] if msg["role"] == "assistant" else ""
+            }
+            for msg in messages
+        ]
+
+        # Get completion from agent
+        response = self.agent.get_completion(formatted_messages)
+
+        # Ensure response is a string
+        if not isinstance(response, str):
+            if hasattr(response, 'to_string'):  # It's a DataFrame
+                response = response.to_string()
+            else:
+                response = str(response)
+
+        # Format response in LiteLLM expected format
+        return {
+            "choices": [
+                {
+                    "message": {
+                        "role": "assistant",
+                        "content": response
+                    }
+                }
+            ],
+            "model": self.agent.args["model_name"],
+            "object": "chat.completion"
+        }
+
+    async def acompletion_stream(self, messages: List[Dict[str, str]], **kwargs) -> Iterator[Dict[str, Any]]:
+        """Async streaming completion interface compatible with LiteLLM"""
+        # Convert messages to format expected by agent
+        formatted_messages = [
+            {
+                "question": msg["content"] if msg["role"] == "user" else "",
+                "answer": msg["content"] if msg["role"] == "assistant" else ""
+            }
+            for msg in messages
+        ]
+
+        # Stream completion from agent
+        model_name = kwargs.get("model", self.agent.args.get("model_name", "mcp-agent"))
+        try:
+            # Handle synchronous generator from _get_completion_stream
+            for chunk in self.agent._get_completion_stream(formatted_messages):
+                content = chunk.get("output", "")
+                if content and isinstance(content, str):
+                    yield {
+                        "choices": [{"delta": {"role": "assistant", "content": content}}],
+                        "model": model_name,
+                        "object": "chat.completion.chunk"
+                    }
+                # Allow async context switch
+                await asyncio.sleep(0)
+        except Exception as e:
+            logger.error(f"Streaming error: {str(e)}")
+            raise
+
+    async def cleanup(self):
+        """Clean up resources"""
+        await self.agent.cleanup()
+
+
+def create_mcp_agent(agent_name: str, project_name: str, mcp_host: str = "127.0.0.1", mcp_port: int = 47337) -> LiteLLMAgentWrapper:
+    """Create an MCP agent and wrap it for LiteLLM compatibility"""
+    from mindsdb.interfaces.agents.agents_controller import AgentsController
+    from mindsdb.interfaces.storage import db
+
+    # Initialize database
+    db.init()
+
+    # Get the agent from database
+    agent_controller = AgentsController()
+    agent_db = agent_controller.get_agent(agent_name, project_name)
+
+    if agent_db is None:
+        raise ValueError(f"Agent {agent_name} not found in project {project_name}")
+
+    # Create MCP agent
+    mcp_agent = MCPLangchainAgent(agent_db, mcp_host=mcp_host, mcp_port=mcp_port)
+
+    # Wrap for LiteLLM compatibility
+    return LiteLLMAgentWrapper(mcp_agent)

mindsdb/interfaces/agents/run_mcp_agent.py

@@ -0,0 +1,205 @@
+import sys
+import argparse
+import asyncio
+from typing import List, Dict
+from contextlib import AsyncExitStack
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from mindsdb.utilities import log
+from mindsdb.interfaces.agents.mcp_client_agent import create_mcp_agent
+
+logger = log.getLogger(__name__)
+
+
+async def run_conversation(agent_wrapper, messages: List[Dict[str, str]], stream: bool = False):
+    """Run a conversation with the agent and print responses"""
+    try:
+        if stream:
+            logger.info("Streaming response:")
+            async for chunk in agent_wrapper.acompletion_stream(messages):
+                content = chunk["choices"][0]["delta"].get("content", "")
+                if content:
+                    # We still need to print content for streaming display
+                    # but we'll log it as debug as well
+                    logger.debug(f"Stream content: {content}")
+                    sys.stdout.write(content)
+                    sys.stdout.flush()
+            logger.debug("End of stream")
+            sys.stdout.write("\n\n")
+            sys.stdout.flush()
+        else:
+            logger.info("Getting response...")
+            response = await agent_wrapper.acompletion(messages)
+            content = response["choices"][0]["message"]["content"]
+            logger.info(f"Response: {content}")
+            # We still need to display the response to the user
+            sys.stdout.write(f"{content}\n")
+            sys.stdout.flush()
+    except Exception as e:
+        logger.error(f"Error during agent conversation: {str(e)}")
+
+
+async def execute_direct_query(query):
+    """Execute a direct SQL query using MCP"""
+    logger.info(f"Executing direct SQL query: {query}")
+
+    # Set up MCP client to connect to the running server
+    async with AsyncExitStack() as stack:
+        # Connect to MCP server
+        server_params = StdioServerParameters(
+            command="python",
+            args=["-m", "mindsdb", "--api=mcp"],
+            env=None
+        )
+
+        try:
+            stdio_transport = await stack.enter_async_context(stdio_client(server_params))
+            stdio, write = stdio_transport
+            session = await stack.enter_async_context(ClientSession(stdio, write))
+
+            await session.initialize()
+
+            # List available tools
+            tools_response = await session.list_tools()
+            tool_names = [tool.name for tool in tools_response.tools]
+            logger.info(f"Available tools: {tool_names}")
+
+            # Find query tool
+            query_tool = None
+            for tool in tools_response.tools:
+                if tool.name == "query":
+                    query_tool = tool
+                    break
+
+            if not query_tool:
+                logger.error("No 'query' tool found in MCP server")
+                return
+
+            # Execute query
+            result = await session.call_tool("query", {"query": query})
+            logger.info(f"Query result: {result.content}")
+        except Exception as e:
+            logger.error(f"Error executing query: {str(e)}")
+            logger.info("Make sure the MindsDB server is running with MCP enabled: python -m mindsdb --api=mysql,mcp,http")
+
+
+async def main():
+    parser = argparse.ArgumentParser(description="Run an agent as an MCP client")
+    parser.add_argument("--agent", type=str, help="Name of the agent to use")
+    parser.add_argument("--project", type=str, default="mindsdb", help="Project containing the agent")
+    parser.add_argument("--host", type=str, default="127.0.0.1", help="MCP server host")
+    parser.add_argument("--port", type=int, default=47337, help="MCP server port")
+    parser.add_argument("--query", type=str, help="Query to send to the agent")
+    parser.add_argument("--stream", action="store_true", help="Stream the response")
+    parser.add_argument("--execute-direct", type=str, help="Execute a direct SQL query via MCP (for testing)")
+
+    args = parser.parse_args()
+
+    try:
+        # Initialize database connection
+        from mindsdb.interfaces.storage import db
+        db.init()
+
+        # Direct SQL execution mode (for testing MCP connection)
+        if args.execute_direct:
+            await execute_direct_query(args.execute_direct)
+            return 0
+
+        # Make sure agent name is provided
+        if not args.agent:
+            parser.error("the --agent argument is required unless --execute-direct is used")
+
+        # Create the agent
+        logger.info(f"Creating MCP client agent for '{args.agent}' in project '{args.project}'")
+        logger.info(f"Connecting to MCP server at {args.host}:{args.port}")
+        logger.info("Make sure MindsDB server is running with MCP enabled: python -m mindsdb --api=mysql,mcp,http")
+
+        agent_wrapper = create_mcp_agent(
+            agent_name=args.agent,
+            project_name=args.project,
+            mcp_host=args.host,
+            mcp_port=args.port
+        )
+
+        # Run an example query if provided
+        if args.query:
+            messages = [{"role": "user", "content": args.query}]
+            await run_conversation(agent_wrapper, messages, args.stream)
+        else:
+            # Interactive mode
+            logger.info("Entering interactive mode. Type 'exit' to quit.")
+            logger.info("Available commands: exit/quit, clear, sql:")
+
+            # We still need to show these instructions to the user
+            sys.stdout.write("\nEntering interactive mode. Type 'exit' to quit.\n")
+            sys.stdout.write("\nAvailable commands:\n")
+            sys.stdout.write("  exit, quit - Exit the program\n")
+            sys.stdout.write("  clear - Clear conversation history\n")
+            sys.stdout.write("  sql: <query> - Execute a direct SQL query via MCP\n")
+            sys.stdout.flush()
+
+            messages = []
+
+            while True:
+                # We need to keep input for user interaction
+                user_input = input("\nYou: ")
+
+                # Check for special commands
+                if user_input.lower() in ["exit", "quit"]:
+                    logger.info("Exiting interactive mode")
+                    break
+                elif user_input.lower() == "clear":
+                    messages = []
+                    logger.info("Conversation history cleared")
+                    sys.stdout.write("Conversation history cleared\n")
+                    sys.stdout.flush()
+                    continue
+                elif user_input.lower().startswith("sql:"):
+                    # Direct SQL execution using the agent's session
+                    sql_query = user_input[4:].strip()
+                    logger.info(f"Executing SQL: {sql_query}")
+                    try:
+                        # Use the tool from the agent
+                        if hasattr(agent_wrapper.agent, "session") and agent_wrapper.agent.session:
+                            result = await agent_wrapper.agent.session.call_tool("query", {"query": sql_query})
+                            logger.info(f"SQL result: {result.content}")
+                            # We need to show the result to the user
+                            sys.stdout.write(f"Result: {result.content}\n")
+                            sys.stdout.flush()
+                        else:
+                            logger.error("No active MCP session")
+                            sys.stdout.write("Error: No active MCP session\n")
+                            sys.stdout.flush()
+                    except Exception as e:
+                        logger.error(f"SQL Error: {str(e)}")
+                        sys.stdout.write(f"SQL Error: {str(e)}\n")
+                        sys.stdout.flush()
+                    continue
+
+                messages.append({"role": "user", "content": user_input})
+                await run_conversation(agent_wrapper, messages, args.stream)
+
+                # Add assistant's response to the conversation history
+                if not args.stream:
+                    response = await agent_wrapper.acompletion(messages)
+                    messages.append({
+                        "role": "assistant",
+                        "content": response["choices"][0]["message"]["content"]
+                    })
+
+        # Clean up resources
+        logger.info("Cleaning up resources")
+        await agent_wrapper.cleanup()
+
+    except Exception as e:
+        logger.error(f"Error running MCP agent: {str(e)}")
+        logger.info("Make sure the MindsDB server is running with MCP enabled: python -m mindsdb --api=mysql,mcp,http")
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(asyncio.run(main()))

mindsdb/interfaces/knowledge_base/controller.py

@@ -65,6 +65,11 @@ def get_embedding_model_from_params(embedding_model_params: dict):
     if provider == 'azure_openai':
         # Azure OpenAI expects the api_key to be passed as 'openai_api_key'.
         params_copy['openai_api_key'] = api_key
+        params_copy['azure_endpoint'] = params_copy.pop('base_url')
+        if 'chunk_size' not in params_copy:
+            params_copy['chunk_size'] = 2048
+        if 'api_version' in params_copy:
+            params_copy['openai_api_version'] = params_copy['api_version']
     else:
         params_copy[f"{provider}_api_key"] = api_key
     params_copy.pop('api_key', None)
@@ -78,11 +83,10 @@ def get_reranking_model_from_params(reranking_model_params: dict):
     Create reranking model from parameters.
     """
     params_copy = copy.deepcopy(reranking_model_params)
-    provider = params_copy.pop('provider', "openai").lower()
-    if provider != 'openai':
-        raise ValueError("Only OpenAI provider is supported for the reranking model.")
-    params_copy[f"{provider}_api_key"] = get_api_key(provider, params_copy, strict=False) or params_copy.get('api_key')
-    params_copy.pop('api_key', None)
+    provider = params_copy.get('provider', "openai").lower()
+
+    if "api_key" not in params_copy:
+        params_copy["api_key"] = get_api_key(provider, params_copy, strict=False)
     params_copy['model'] = params_copy.pop('model_name', None)
 
     return LLMReranker(**params_copy)
@@ -424,11 +428,12 @@ class KnowledgeBaseTable:
         db_handler = self.get_vector_db()
         db_handler.delete(self._kb.vector_database_table)
 
-    def insert(self, df: pd.DataFrame):
+    def insert(self, df: pd.DataFrame, params: dict = None):
         """Insert dataframe to KB table.
 
         Args:
             df: DataFrame to insert
+            params: User parameters of insert
         """
         if df.empty:
             return
@@ -497,7 +502,12 @@ class KnowledgeBaseTable:
         df_emb = self._df_to_embeddings(df)
         df = pd.concat([df, df_emb], axis=1)
         db_handler = self.get_vector_db()
-        db_handler.do_upsert(self._kb.vector_database_table, df)
+
+        if params is not None and params.get('kb_no_upsert', False):
+            # speed up inserting by disable checking existing records
+            db_handler.insert(self._kb.vector_database_table, df)
+        else:
+            db_handler.do_upsert(self._kb.vector_database_table, df)
 
     def _adapt_column_names(self, df: pd.DataFrame) -> pd.DataFrame:
         '''

mindsdb/interfaces/skills/skill_tool.py

@@ -243,7 +243,13 @@ class SkillToolController:
             kb_table = session_controller.kb_controller.get_table(knowledge_base_name, skill.project_id)
 
             res = kb_table.select_query(query)
-            return '\n'.join(res.content)
+            # Handle both chunk_content and content column names
+            if hasattr(res, 'chunk_content'):
+                return '\n'.join(res.chunk_content)
+            elif hasattr(res, 'content'):
+                return '\n'.join(res.content)
+            else:
+                return "No content or chunk_content found in knowledge base response"
 
         return _answer_question
 

mindsdb/interfaces/skills/sql_agent.py

@@ -12,6 +12,7 @@ from mindsdb_sql_parser.ast import Select, Show, Describe, Explain, Identifier
 from mindsdb.utilities import log
 from mindsdb.utilities.context import context as ctx
 from mindsdb.integrations.utilities.query_traversal import query_traversal
+from mindsdb.integrations.libs.response import INF_SCHEMA_COLUMNS_NAMES
 
 logger = log.getLogger(__name__)
 
@@ -275,9 +276,13 @@ class SQLAgent:
         dn = self._command_executor.session.datahub.get(integration)
 
         fields, dtypes = [], []
-        for column in dn.get_table_columns(table_name, schema_name):
-            fields.append(column['name'])
-            dtypes.append(column.get('type', ''))
+        for df in dn.get_table_columns_df(table_name, schema_name):
+            df_records = df.to_dict(orient='records')
+            fields.append(df_records[INF_SCHEMA_COLUMNS_NAMES.COLUMN_NAME])
+            if df_records[INF_SCHEMA_COLUMNS_NAMES.MYSQL_DATA_TYPE] is not None:
+                dtypes.append(df_records[INF_SCHEMA_COLUMNS_NAMES.MYSQL_DATA_TYPE].value)
+            else:
+                dtypes.append(df_records[INF_SCHEMA_COLUMNS_NAMES.DATA_TYPE])
 
         info = f'Table named `{table_str}`:\n'
         info += f"\nSample with first {self._sample_rows_in_table_info} rows from table {table_str} in CSV format (dialect is 'excel'):\n"

mindsdb/utilities/config.py

@@ -209,6 +209,10 @@ class Config:
                     "restart_on_failure": True,
                     "max_restart_count": 1,
                     "max_restart_interval_seconds": 60
+                },
+                "litellm": {
+                    "host": "0.0.0.0",  # API server binds to all interfaces by default
+                    "port": "8000"
                 }
             },
             "cache": {
@@ -378,7 +382,8 @@ class Config:
         ):
             self._cmd_args = argparse.Namespace(
                 api=None, config=None, install_handlers=None, verbose=False,
-                no_studio=False, version=False, ml_task_queue_consumer=None
+                no_studio=False, version=False, ml_task_queue_consumer=None,
+                agent=None, project=None
             )
             return
 
@@ -390,6 +395,8 @@ class Config:
         parser.add_argument('--no_studio', action='store_true')
         parser.add_argument('-v', '--version', action='store_true')
         parser.add_argument('--ml_task_queue_consumer', action='store_true', default=None)
+        parser.add_argument('--agent', type=str, default=None, help='Name of the agent to use with litellm APIs')
+        parser.add_argument('--project', type=str, default=None, help='Project containing the agent (default: mindsdb)')
         self._cmd_args = parser.parse_args()
 
     def fetch_auto_config(self) -> bool:

mindsdb/utilities/starters.py

@@ -38,3 +38,10 @@ def start_mcp(*args, **kwargs):
     from mindsdb.api.mcp.start import start
 
     start(*args, **kwargs)
+
+
+def start_litellm(*args, **kwargs):
+    """Start the LiteLLM server"""
+    from mindsdb.api.litellm.start import start
+
+    start(*args, **kwargs)