MindsDB 25.4.3.2__py3-none-any.whl → 25.4.5.0__py3-none-any.whl

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/functions/controller.py

@@ -3,6 +3,7 @@ import os
 from duckdb.typing import BIGINT, DOUBLE, VARCHAR, BLOB, BOOLEAN
 from mindsdb.interfaces.functions.to_markdown import ToMarkdown
 from mindsdb.interfaces.storage.model_fs import HandlerStorage
+from mindsdb.utilities.config import config
 
 
 def python_to_duckdb_type(py_type):
@@ -164,7 +165,7 @@ class FunctionController(BYOMFunctionsController):
             return self.callbacks[name]
 
         def callback(file_path_or_url, use_llm):
-            chat_model_params = self._parse_chat_model_params()
+            chat_model_params = self._parse_chat_model_params('TO_MARKDOWN_FUNCTION_')
 
             llm_client = None
             llm_model = None
@@ -192,7 +193,7 @@ class FunctionController(BYOMFunctionsController):
         """
         Parses the environment variables for chat model parameters.
         """
-        chat_model_params = {}
+        chat_model_params = config.get("default_llm") or {}
         for k, v in os.environ.items():
             if k.startswith(param_prefix):
                 param_name = k[len(param_prefix):]

mindsdb/interfaces/knowledge_base/controller.py

@@ -27,7 +27,7 @@ from mindsdb.integrations.libs.vectordatabase_handler import (
 from mindsdb.integrations.utilities.rag.rag_pipeline_builder import RAG
 from mindsdb.integrations.utilities.rag.config_loader import load_rag_config
 from mindsdb.integrations.utilities.handler_utils import get_api_key
-from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embedding_handler import construct_model_from_args, row_to_document
+from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embedding_handler import construct_model_from_args
 
 from mindsdb.interfaces.agents.constants import DEFAULT_EMBEDDINGS_MODEL_CLASS
 from mindsdb.interfaces.agents.langchain_agent import create_chat_model, get_llm_provider
@@ -37,11 +37,12 @@ from mindsdb.interfaces.knowledge_base.preprocessing.document_preprocessor impor
 from mindsdb.interfaces.model.functions import PredictorRecordNotFound
 from mindsdb.utilities.exception import EntityExistsError, EntityNotExistsError
 from mindsdb.integrations.utilities.sql_utils import FilterCondition, FilterOperator
+from mindsdb.utilities.config import config
 from mindsdb.utilities.context import context as ctx
 
 from mindsdb.api.executor.command_executor import ExecuteCommands
 from mindsdb.utilities import log
-from mindsdb.integrations.utilities.rag.rerankers.reranker_compressor import LLMReranker
+from mindsdb.integrations.utilities.rag.rerankers.base_reranker import BaseLLMReranker
 
 logger = log.getLogger(__name__)
 
@@ -52,6 +53,18 @@ KB_TO_VECTORDB_COLUMNS = {
 }
 
 
+def get_model_params(model_params: dict, default_config_key: str):
+    """
+    Get model parameters by combining default config with user provided parameters.
+    """
+    combined_model_params = copy.deepcopy(config.get(default_config_key, {}))
+
+    if model_params:
+        combined_model_params.update(model_params)
+
+    return combined_model_params
+
+
 def get_embedding_model_from_params(embedding_model_params: dict):
     """
     Create embedding model from parameters.
@@ -65,6 +78,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,14 +96,13 @@ 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)
+    return BaseLLMReranker(**params_copy)
 
 
 class KnowledgeBaseTable:
@@ -211,7 +228,7 @@ class KnowledgeBaseTable:
     def add_relevance(self, df, query_text, relevance_threshold=None):
         relevance_column = TableField.RELEVANCE.value
 
-        reranking_model_params = self._kb.params.get("reranking_model")
+        reranking_model_params = get_model_params(self._kb.params.get("reranking_model"), "default_llm")
         if reranking_model_params and query_text and len(df) > 0:
             # Use reranker for relevance score
             try:
@@ -424,11 +441,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 +515,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:
         '''
@@ -647,47 +670,34 @@ class KnowledgeBaseTable:
         if df.empty:
             return pd.DataFrame([], columns=[TableField.EMBEDDINGS.value])
 
-        # keep only content
-        df = df[[TableField.CONTENT.value]]
-
         model_id = self._kb.embedding_model_id
-        if model_id:
-            # get the input columns
-            model_rec = db.session.query(db.Predictor).filter_by(id=model_id).first()
 
-            assert model_rec is not None, f"Model not found: {model_id}"
-            model_project = db.session.query(db.Project).filter_by(id=model_rec.project_id).first()
+        # get the input columns
+        model_rec = db.session.query(db.Predictor).filter_by(id=model_id).first()
 
-            project_datanode = self.session.datahub.get(model_project.name)
+        assert model_rec is not None, f"Model not found: {model_id}"
+        model_project = db.session.query(db.Project).filter_by(id=model_rec.project_id).first()
 
-            model_using = model_rec.learn_args.get('using', {})
-            input_col = model_using.get('question_column')
-            if input_col is None:
-                input_col = model_using.get('input_column')
-
-            if input_col is not None and input_col != TableField.CONTENT.value:
-                df = df.rename(columns={TableField.CONTENT.value: input_col})
-
-            df_out = project_datanode.predict(
-                model_name=model_rec.name,
-                df=df,
-                params=self.model_params
-            )
+        project_datanode = self.session.datahub.get(model_project.name)
 
-            target = model_rec.to_predict[0]
-            if target != TableField.EMBEDDINGS.value:
-                # adapt output for vectordb
-                df_out = df_out.rename(columns={target: TableField.EMBEDDINGS.value})
+        model_using = model_rec.learn_args.get('using', {})
+        input_col = model_using.get('question_column')
+        if input_col is None:
+            input_col = model_using.get('input_column')
 
-        elif self._kb.params.get('embedding_model'):
-            embedding_model = get_embedding_model_from_params(self._kb.params.get('embedding_model'))
+        if input_col is not None and input_col != TableField.CONTENT.value:
+            df = df.rename(columns={TableField.CONTENT.value: input_col})
 
-            df_texts = df.apply(row_to_document, axis=1)
-            embeddings = embedding_model.embed_documents(df_texts.tolist())
-            df_out = df.copy().assign(**{TableField.EMBEDDINGS.value: embeddings})
+        df_out = project_datanode.predict(
+            model_name=model_rec.name,
+            df=df,
+            params=self.model_params
+        )
 
-        else:
-            raise ValueError("No embedding model found for the knowledge base.")
+        target = model_rec.to_predict[0]
+        if target != TableField.EMBEDDINGS.value:
+            # adapt output for vectordb
+            df_out = df_out.rename(columns={target: TableField.EMBEDDINGS.value})
 
         df_out = df_out[[TableField.EMBEDDINGS.value]]
 
@@ -718,14 +728,15 @@ class KnowledgeBaseTable:
         """
         # Get embedding model from knowledge base
         embeddings_model = None
+        embedding_model_params = get_model_params(self._kb.params.get('embedding_model', {}), 'default_embedding_model')
         if self._kb.embedding_model:
             # Extract embedding model args from knowledge base table
             embedding_args = self._kb.embedding_model.learn_args.get('using', {})
             # Construct the embedding model directly
             embeddings_model = construct_model_from_args(embedding_args)
             logger.debug(f"Using knowledge base embedding model with args: {embedding_args}")
-        elif self._kb.params.get('embedding_model'):
-            embeddings_model = get_embedding_model_from_params(self._kb.params['embedding_model'])
+        elif embedding_model_params:
+            embeddings_model = get_embedding_model_from_params(embedding_model_params)
             logger.debug(f"Using knowledge base embedding model from params: {self._kb.params['embedding_model']}")
         else:
             embeddings_model = DEFAULT_EMBEDDINGS_MODEL_CLASS()
@@ -859,35 +870,33 @@ class KnowledgeBaseController:
                 return kb
             raise EntityExistsError("Knowledge base already exists", name)
 
-        embedding_model_params = params.get('embedding_model', None)
-        reranking_model_params = params.get('reranking_model', None)
+        embedding_params = copy.deepcopy(config.get('default_embedding_model', {}))
 
+        model_name = None
+        model_project = project
         if embedding_model:
             model_name = embedding_model.parts[-1]
+            if len(embedding_model.parts) > 1:
+                model_project = self.session.database_controller.get_project(embedding_model.parts[-2])
 
-        elif embedding_model_params:
-            # Get embedding model from params.
-            # This is called here to check validaity of the parameters.
-            get_embedding_model_from_params(
-                embedding_model_params
-            )
+        elif 'embedding_model' in params:
+            if isinstance(params['embedding_model'], str):
+                # it is model name
+                model_name = params['embedding_model']
+            else:
+                # it is params for model
+                embedding_params.update(params['embedding_model'])
 
-        else:
-            model_name = self._get_default_embedding_model(
+        if model_name is None:
+            model_name = self._create_embedding_model(
                 project.name,
-                params=params
+                params=embedding_params,
+                kb_name=name,
             )
-            params['default_embedding_model'] = model_name
-
-        model_project = None
-        if embedding_model is not None and len(embedding_model.parts) > 1:
-            # model project is set
-            model_project = self.session.database_controller.get_project(embedding_model.parts[-2])
-        elif not embedding_model_params:
-            model_project = project
+            params['created_embedding_model'] = model_name
 
         embedding_model_id = None
-        if model_project:
+        if model_name is not None:
             model = self.session.model_controller.get_model(
                 name=model_name,
                 project_name=model_project.name
@@ -895,6 +904,7 @@ class KnowledgeBaseController:
             model_record = db.Predictor.query.get(model['id'])
             embedding_model_id = model_record.id
 
+        reranking_model_params = get_model_params(params.get('reranking_model', {}), 'default_llm')
         if reranking_model_params:
             # Get reranking model from params.
             # This is called here to check validaity of the parameters.
@@ -979,38 +989,52 @@ class KnowledgeBaseController:
         self.session.integration_controller.add(vector_store_name, engine, connection_args)
         return vector_store_name
 
-    def _get_default_embedding_model(self, project_name, engine="langchain_embedding", params: dict = None):
+    def _create_embedding_model(self, project_name, engine="openai", params: dict = None, kb_name=''):
         """create a default embedding model for knowledge base, if not specified"""
-        model_name = "kb_default_embedding_model"
+        model_name = f"kb_embedding_{kb_name}"
 
-        # check exists
+        # drop if exists - parameters can be different
         try:
             model = self.session.model_controller.get_model(model_name, project_name=project_name)
             if model is not None:
-                return model_name
+                self.session.model_controller.delete_model(model_name, project_name)
         except PredictorRecordNotFound:
             pass
 
-        using_args = {
-            'engine': engine
-        }
-        if engine == 'langchain_embedding':
-            # Use default embeddings.
-            using_args['class'] = 'openai'
+        if 'provider' in params:
+            engine = params.pop('provider').lower()
+
+        if engine == 'azure_openai':
+            engine = 'openai'
+            params['provider'] = 'azure'
+
+        if engine == 'openai':
+            if 'question_column' not in params:
+                params['question_column'] = 'content'
+            if 'api_key' in params:
+                params[f"{engine}_api_key"] = params.pop('api_key')
+            if 'base_url' in params:
+                params['api_base'] = params.pop('base_url')
+
+        params['engine'] = engine
+        params['join_learn_process'] = True
+        params['mode'] = 'embedding'
 
         # Include API key if provided.
-        using_args.update({k: v for k, v in params.items() if 'api_key' in k})
         statement = CreatePredictor(
             name=Identifier(parts=[project_name, model_name]),
-            using=using_args,
+            using=params,
             targets=[
                 Identifier(parts=[TableField.EMBEDDINGS.value])
             ]
         )
 
         command_executor = ExecuteCommands(self.session)
-        command_executor.answer_create_predictor(statement, project_name)
-
+        resp = command_executor.answer_create_predictor(statement, project_name)
+        # check model status
+        record = resp.data.records[0]
+        if record['STATUS'] == 'error':
+            raise ValueError('Embedding model error:' + record['ERROR'])
         return model_name
 
     def delete(self, name: str, project_name: int, if_exists: bool = False) -> None:
@@ -1044,9 +1068,9 @@ class KnowledgeBaseController:
                 self.session.integration_controller.delete(kb.params['default_vector_storage'])
             except EntityNotExistsError:
                 pass
-        if 'default_embedding_model' in kb.params:
+        if 'created_embedding_model' in kb.params:
             try:
-                self.session.model_controller.delete_model(kb.params['default_embedding_model'], project_name)
+                self.session.model_controller.delete_model(kb.params['created_embedding_model'], project_name)
             except EntityNotExistsError:
                 pass