MindsDB 25.5.3.0__py3-none-any.whl → 25.5.4.0__py3-none-any.whl

mindsdb/interfaces/skills/custom/text2sql/mindsdb_kb_tools.py

@@ -0,0 +1,190 @@
+from typing import Type, List, Any
+import re
+import json
+from pydantic import BaseModel, Field
+from langchain_core.tools import BaseTool
+
+
+class KnowledgeBaseListToolInput(BaseModel):
+    tool_input: str = Field(
+        "", description="An empty string to list all knowledge bases."
+    )
+
+
+class KnowledgeBaseListTool(BaseTool):
+    """Tool for listing knowledge bases in MindsDB."""
+
+    name: str = "kb_list_tool"
+    description: str = "List all knowledge bases in MindsDB."
+    args_schema: Type[BaseModel] = KnowledgeBaseListToolInput
+    db: Any = None
+
+    def _run(self, tool_input: str) -> str:
+        """List all knowledge bases."""
+        return self.db.get_usable_knowledge_base_names()
+
+
+class KnowledgeBaseInfoToolInput(BaseModel):
+    tool_input: str = Field(
+        ...,
+        description="A comma-separated list of knowledge base names enclosed between $START$ and $STOP$.",
+    )
+
+
+class KnowledgeBaseInfoTool(BaseTool):
+    """Tool for getting information about knowledge bases in MindsDB."""
+
+    name: str = "kb_info_tool"
+    description: str = "Get information about knowledge bases in MindsDB."
+    args_schema: Type[BaseModel] = KnowledgeBaseInfoToolInput
+    db: Any = None
+
+    def _extract_kb_names(self, tool_input: str) -> List[str]:
+        """Extract knowledge base names from the tool input."""
+        match = re.search(r"\$START\$(.*?)\$STOP\$", tool_input, re.DOTALL)
+        if not match:
+            return []
+
+        # Extract and clean the knowledge base names
+        kb_names_str = match.group(1).strip()
+        kb_names = re.findall(r"`([^`]+)`", kb_names_str)
+        return kb_names
+
+    def _run(self, tool_input: str) -> str:
+        """Get information about specified knowledge bases."""
+        kb_names = self._extract_kb_names(tool_input)
+
+        if not kb_names:
+            return "No valid knowledge base names provided. Please provide names enclosed in backticks between $START$ and $STOP$."
+
+        results = []
+
+        for kb_name in kb_names:
+            try:
+                # Get knowledge base schema
+                schema_result = self.db.run_no_throw(
+                    f"DESCRIBE KNOWLEDGE_BASE `{kb_name}`;"
+                )
+
+                if not schema_result:
+                    results.append(
+                        f"Knowledge base `{kb_name}` not found or has no schema information."
+                    )
+                    continue
+
+                # Get sample data
+                sample_data = self.db.run_no_throw(
+                    f"SELECT * FROM `{kb_name}` LIMIT 10;"
+                )
+
+                # Format the results
+                kb_info = f"## Knowledge Base: `{kb_name}`\n\n"
+
+                # Schema information
+                kb_info += "### Schema Information:\n"
+                kb_info += "```\n"
+                for row in schema_result:
+                    kb_info += f"{json.dumps(row, indent=2)}\n"
+                kb_info += "```\n\n"
+
+                # Sample data
+                kb_info += "### Sample Data:\n"
+                if sample_data:
+                    # Extract column names
+                    columns = list(sample_data[0].keys())
+
+                    # Create markdown table header
+                    kb_info += "| " + " | ".join(columns) + " |\n"
+                    kb_info += "| " + " | ".join(["---" for _ in columns]) + " |\n"
+
+                    # Add rows
+                    for row in sample_data:
+                        formatted_row = []
+                        for col in columns:
+                            cell_value = row[col]
+                            if isinstance(cell_value, dict):
+                                cell_value = json.dumps(cell_value, ensure_ascii=False)
+                            formatted_row.append(str(cell_value).replace("|", "\\|"))
+                        kb_info += "| " + " | ".join(formatted_row) + " |\n"
+                else:
+                    kb_info += "No sample data available.\n"
+
+                results.append(kb_info)
+
+            except Exception as e:
+                results.append(
+                    f"Error getting information for knowledge base `{kb_name}`: {str(e)}"
+                )
+
+        return "\n\n".join(results)
+
+
+class KnowledgeBaseQueryToolInput(BaseModel):
+    tool_input: str = Field(
+        ...,
+        description="A SQL query for knowledge bases. Can be provided directly or enclosed between $START$ and $STOP$.",
+    )
+
+
+class KnowledgeBaseQueryTool(BaseTool):
+    """Tool for querying knowledge bases in MindsDB."""
+
+    name: str = "kb_query_tool"
+    description: str = "Query knowledge bases in MindsDB."
+    args_schema: Type[BaseModel] = KnowledgeBaseQueryToolInput
+    db: Any = None
+
+    def _extract_query(self, tool_input: str) -> str:
+        """Extract the SQL query from the tool input."""
+        # First check if the input is wrapped in $START$ and $STOP$
+        match = re.search(r"\$START\$(.*?)\$STOP\$", tool_input, re.DOTALL)
+        if match:
+            return match.group(1).strip()
+
+        # If not wrapped in delimiters, use the input directly
+        # Check for SQL keywords to validate it's likely a query
+        if re.search(
+            r"\b(SELECT|FROM|WHERE|LIMIT|ORDER BY)\b", tool_input, re.IGNORECASE
+        ):
+            return tool_input.strip()
+
+        return ""
+
+    def _run(self, tool_input: str) -> str:
+        """Execute a knowledge base query."""
+        query = self._extract_query(tool_input)
+
+        if not query:
+            return "No valid SQL query provided. Please provide a SQL query that includes SELECT, FROM, or other SQL keywords."
+
+        try:
+            # Execute the query
+            result = self.db.run_no_throw(query)
+
+            if not result:
+                return "Query executed successfully, but no results were returned."
+
+            # Format the results as a markdown table
+            if isinstance(result, list) and len(result) > 0:
+                # Extract column names
+                columns = list(result[0].keys())
+
+                # Create markdown table header
+                table = "| " + " | ".join(columns) + " |\n"
+                table += "| " + " | ".join(["---" for _ in columns]) + " |\n"
+
+                # Add rows
+                for row in result:
+                    formatted_row = []
+                    for col in columns:
+                        cell_value = row[col]
+                        if isinstance(cell_value, dict):
+                            cell_value = json.dumps(cell_value, ensure_ascii=False)
+                        formatted_row.append(str(cell_value).replace("|", "\\|"))
+                    table += "| " + " | ".join(formatted_row) + " |\n"
+
+                return table
+
+            return result
+        except Exception as e:
+            return f"Error executing query: {str(e)}"

mindsdb/interfaces/skills/custom/text2sql/mindsdb_sql_toolkit.py

@@ -7,6 +7,11 @@ from langchain_community.tools import ListSQLDatabaseTool, InfoSQLDatabaseTool,
 from langchain_core.tools import BaseTool
 
 from mindsdb.interfaces.skills.custom.text2sql.mindsdb_sql_tool import MindsDBSQLParserTool
+from mindsdb.interfaces.skills.custom.text2sql.mindsdb_kb_tools import (
+    KnowledgeBaseListTool,
+    KnowledgeBaseInfoTool,
+    KnowledgeBaseQueryTool
+)
 
 
 class MindsDBSQLToolkit(SQLDatabaseToolkit):
@@ -107,9 +112,96 @@ class MindsDBSQLToolkit(SQLDatabaseToolkit):
             description=mindsdb_sql_parser_tool_description
         )
 
+        # Knowledge base tools
+        kb_list_tool = KnowledgeBaseListTool(
+            name=f'kb_list_tool{prefix}',
+            db=self.db,
+            description=dedent("""\
+                Lists all available knowledge bases that can be queried.
+                Input: No input required, just call the tool directly.
+                Output: A table of all available knowledge bases with their names and creation dates.
+
+                Use this tool first when answering factual questions to see what knowledge bases are available.
+                Each knowledge base name is escaped using backticks.
+
+                Example usage: kb_list_tool()
+            """)
+        )
+
+        kb_info_tool = KnowledgeBaseInfoTool(
+            name=f'kb_info_tool{prefix}',
+            db=self.db,
+            description=dedent(f"""\
+                Gets detailed information about specific knowledge bases including their structure and metadata fields.
+
+                Input: A knowledge base name as a simple string.
+                Output: Schema, metadata columns, and sample rows for the specified knowledge base.
+
+                Use this after kb_list_tool to understand what information is contained in the knowledge base
+                and what metadata fields are available for filtering.
+
+                Example usage: kb_info_tool("kb_name")
+
+                Make sure the knowledge base exists by calling {kb_list_tool.name} first.
+            """)
+        )
+
+        kb_query_tool = KnowledgeBaseQueryTool(
+            name=f'kb_query_tool{prefix}',
+            db=self.db,
+            description=dedent(f"""\
+                Queries knowledge bases using SQL syntax to retrieve relevant information.
+
+                Input: A SQL query string that targets a knowledge base.
+                Output: Knowledge base search results or error message.
+
+                This tool is designed for semantic search and metadata filtering in MindsDB knowledge bases.
+
+                Query Types and Examples:
+                1. Basic semantic search:
+                   kb_query_tool("SELECT * FROM kb_name WHERE content = 'your search query';")
+
+                2. Metadata filtering:
+                   kb_query_tool("SELECT * FROM kb_name WHERE metadata_field = 'value';")
+
+                3. Combined search:
+                   kb_query_tool("SELECT * FROM kb_name WHERE content = 'query' AND metadata_field = 'value';")
+
+                4. Setting relevance threshold:
+                   kb_query_tool("SELECT * FROM kb_name WHERE content = 'query' AND relevance_threshold = 0.7;")
+
+                5. Limiting results:
+                   kb_query_tool("SELECT * FROM kb_name WHERE content = 'query' LIMIT 5;")
+
+                6. Getting sample data:
+                   kb_query_tool("SELECT * FROM kb_name LIMIT 3;")
+
+                7. Don't use LIKE operator on content filter ie semantic search:
+                SELECT * FROM `test_kb` WHERE content LIKE '%population of New York%' $STOP$
+
+                Like is not supported, use the following instead:
+                SELECT * FROM `test_kb` WHERE content = 'population of New York'
+
+                Result Format:
+                - Results include: id, chunk_id, chunk_content, metadata, distance, and relevance columns
+                - The metadata column contains a JSON object with all metadata fields
+
+                Best Practices:
+                - Always check available knowledge bases with {kb_list_tool.name} first
+                - Use {kb_info_tool.name} to understand the structure and metadata fields
+                - Always include a semicolon at the end of your SQL query
+
+                For factual questions, use this tool to retrieve information rather than relying on the model's knowledge.
+            """)
+        )
+
+        # Return standard SQL tools and knowledge base tools
         return [
             query_sql_database_tool,
             info_sql_database_tool,
             list_sql_database_tool,
             mindsdb_sql_parser_tool,
+            kb_list_tool,
+            kb_info_tool,
+            kb_query_tool,
         ]

mindsdb/interfaces/skills/skill_tool.py

@@ -1,5 +1,6 @@
 import enum
 import inspect
+
 from dataclasses import dataclass
 from collections import defaultdict
 from typing import List, Dict, Optional
@@ -14,7 +15,7 @@ from mindsdb.utilities.config import config
 from mindsdb.interfaces.storage import db
 from mindsdb.interfaces.skills.sql_agent import SQLAgent
 from mindsdb.integrations.libs.vectordatabase_handler import TableField
-
+from mindsdb.interfaces.agents.constants import DEFAULT_TEXT2SQL_DATABASE
 
 _DEFAULT_TOP_K_SIMILARITY_SEARCH = 5
 _MAX_CACHE_SIZE = 1000
@@ -120,7 +121,6 @@ class SkillToolController:
         try:
             from mindsdb.interfaces.agents.mindsdb_database_agent import MindsDBSQL
             from mindsdb.interfaces.skills.custom.text2sql.mindsdb_sql_toolkit import MindsDBSQLToolkit
-            from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
         except ImportError:
             raise ImportError(
                 'To use the text-to-SQL skill, please install langchain with `pip install mindsdb[langchain]`')
@@ -132,75 +132,220 @@ class SkillToolController:
             return f'`{name}`'
 
         tables_list = []
+        knowledge_bases_list = []
+        ignore_knowledge_bases_list = []
+
+        # Track databases extracted from dot notation
+        extracted_databases = set()
+
+        # Initialize knowledge_base_database with default value
+        knowledge_base_database = DEFAULT_TEXT2SQL_DATABASE  # Default to mindsdb project
+
+        # First pass: collect all database and knowledge base parameters
         for skill in skills:
-            database = skill.params['database']
-            restriction_on_tables = skill.restriction_on_tables
-            if restriction_on_tables is None:
-                handler = command_executor.session.integration_controller.get_data_handler(database)
-                if 'all' in inspect.signature(handler.get_tables).parameters:
-                    response = handler.get_tables(all=True)
-                else:
-                    response = handler.get_tables()
-                # no restrictions
-                columns = [c.lower() for c in response.data_frame.columns]
-                name_idx = columns.index('table_name') if 'table_name' in columns else 0
-
-                if 'table_schema' in response.data_frame.columns:
-                    for _, row in response.data_frame.iterrows():
-                        tables_list.append(f"{database}.{row['table_schema']}.{escape_table_name(row[name_idx])}")
-                else:
-                    for table_name in response.data_frame.iloc[:, name_idx]:
-                        tables_list.append(f"{database}.{escape_table_name(table_name)}")
+            # Update knowledge_base_database if specified in any skill
+            if skill.params.get('knowledge_base_database'):
+                knowledge_base_database = skill.params.get('knowledge_base_database')
+
+            # Extract databases from include_tables with dot notation
+            if skill.params.get('include_tables'):
+                include_tables = skill.params.get('include_tables')
+                if isinstance(include_tables, str):
+                    include_tables = [t.strip() for t in include_tables.split(',')]
+
+                # Extract database names from dot notation
+                for table in include_tables:
+                    if '.' in table:
+                        db_name = table.split('.')[0]
+                        extracted_databases.add(db_name)
+
+            # Extract databases from include_knowledge_bases with dot notation
+            if skill.params.get('include_knowledge_bases'):
+                include_kbs = skill.params.get('include_knowledge_bases')
+                if isinstance(include_kbs, str):
+                    include_kbs = [kb.strip() for kb in include_kbs.split(',')]
+
+                # Extract database names from dot notation
+                for kb in include_kbs:
+                    if '.' in kb:
+                        db_name = kb.split('.')[0]
+                        if db_name != knowledge_base_database:
+                            # Only update if it's different from the default
+                            knowledge_base_database = db_name
+
+        # Second pass: collect all tables and knowledge base restrictions
+        for skill in skills:
+            # Get database for tables (this is an actual database connection)
+            database = skill.params.get('database', DEFAULT_TEXT2SQL_DATABASE)
+
+            # Add databases extracted from dot notation if no explicit database is provided
+            if not database and extracted_databases:
+                # Use the first extracted database if no explicit database is provided
+                database = next(iter(extracted_databases))
+                # Update the skill params with the extracted database
+                skill.params['database'] = database
+
+            # Extract knowledge base restrictions if they exist in the skill params
+            if skill.params.get('include_knowledge_bases'):
+                # Convert to list if it's a string
+                include_kbs = skill.params.get('include_knowledge_bases')
+                if isinstance(include_kbs, str):
+                    include_kbs = [kb.strip() for kb in include_kbs.split(',')]
+
+                # Process each knowledge base name
+                for kb in include_kbs:
+                    # If it doesn't have a dot, prefix it with the knowledge_base_database
+                    if '.' not in kb:
+                        knowledge_bases_list.append(f"{knowledge_base_database}.{kb}")
+                    else:
+                        knowledge_bases_list.append(kb)
+
+            # Collect ignore_knowledge_bases
+            if skill.params.get('ignore_knowledge_bases'):
+                # Convert to list if it's a string
+                ignore_kbs = skill.params.get('ignore_knowledge_bases')
+                if isinstance(ignore_kbs, str):
+                    ignore_kbs = [kb.strip() for kb in ignore_kbs.split(',')]
+
+                # Process each knowledge base name to ignore
+                for kb in ignore_kbs:
+                    # If it doesn't have a dot, prefix it with the knowledge_base_database
+                    if '.' not in kb:
+                        ignore_knowledge_bases_list.append(f"{knowledge_base_database}.{kb}")
+                    else:
+                        ignore_knowledge_bases_list.append(kb)
+
+            # Skip if no database specified
+            if not database:
                 continue
-            for schema_name, tables in restriction_on_tables.items():
-                for table in tables:
-                    if schema_name is None:
-                        tables_list.append(f'{database}.{escape_table_name(table)}')
+
+            # Process include_tables with dot notation
+            if skill.params.get('include_tables'):
+                include_tables = skill.params.get('include_tables')
+                if isinstance(include_tables, str):
+                    include_tables = [t.strip() for t in include_tables.split(',')]
+
+                for table in include_tables:
+                    # If table already has a database prefix, use it as is
+                    if '.' in table:
+                        # Check if the table already has backticks
+                        if '`' in table:
+                            tables_list.append(table)
+                        else:
+                            # Apply escape_table_name only to the table part
+                            parts = table.split('.')
+                            if len(parts) == 2:
+                                # Format: database.table
+                                tables_list.append(f"{parts[0]}.{escape_table_name(parts[1])}")
+                            elif len(parts) == 3:
+                                # Format: database.schema.table
+                                tables_list.append(f"{parts[0]}.{parts[1]}.{escape_table_name(parts[2])}")
+                            else:
+                                # Unusual format, escape the whole thing
+                                tables_list.append(escape_table_name(table))
                     else:
-                        tables_list.append(f'{database}.{schema_name}.{escape_table_name(table)}')
+                        # Otherwise, prefix with the database
+                        tables_list.append(f"{database}.{escape_table_name(table)}")
+
+                # Skip further table processing if include_tables is specified
+                continue
+
+            restriction_on_tables = skill.restriction_on_tables
+
+            if restriction_on_tables is None and database:
+                try:
+                    handler = command_executor.session.integration_controller.get_data_handler(database)
+                    if 'all' in inspect.signature(handler.get_tables).parameters:
+                        response = handler.get_tables(all=True)
+                    else:
+                        response = handler.get_tables()
+                    # no restrictions
+                    columns = [c.lower() for c in response.data_frame.columns]
+                    name_idx = columns.index('table_name') if 'table_name' in columns else 0
+
+                    if 'table_schema' in response.data_frame.columns:
+                        for _, row in response.data_frame.iterrows():
+                            tables_list.append(f"{database}.{row['table_schema']}.{escape_table_name(row[name_idx])}")
+                    else:
+                        for table_name in response.data_frame.iloc[:, name_idx]:
+                            tables_list.append(f"{database}.{escape_table_name(table_name)}")
+                except Exception as e:
+                    logger.warning(f"Could not get tables from database {database}: {str(e)}")
+                continue
+
+            # Handle table restrictions
+            if restriction_on_tables and database:
+                for schema_name, tables in restriction_on_tables.items():
+                    for table in tables:
+                        # Check if the table already has dot notation (e.g., 'postgresql_conn.home_rentals')
+                        if '.' in table:
+                            # Table already has database prefix, add it directly
+                            tables_list.append(escape_table_name(table))
+                        else:
+                            # No dot notation, apply schema and database as needed
+                            if schema_name is None:
+                                tables_list.append(f'{database}.{escape_table_name(table)}')
+                            else:
+                                tables_list.append(f'{database}.{schema_name}.{escape_table_name(table)}')
+                continue
+
+        # Remove duplicates from lists
+        tables_list = list(set(tables_list))
+        knowledge_bases_list = list(set(knowledge_bases_list))
+        ignore_knowledge_bases_list = list(set(ignore_knowledge_bases_list))
+
+        # Determine knowledge base parameters to pass to SQLAgent
+        include_knowledge_bases = knowledge_bases_list if knowledge_bases_list else None
+        ignore_knowledge_bases = ignore_knowledge_bases_list if ignore_knowledge_bases_list else None
+
+        # If both include and ignore lists exist, include takes precedence
+        if include_knowledge_bases:
+            ignore_knowledge_bases = None
+
+        # # Get all databases from skills and extracted databases
+        # all_databases = list(set([s.params.get('database', DEFAULT_TEXT2SQL_DATABASE) for s in skills if s.params.get('database')] + list(extracted_databases)))
+        #
+        #
+        # # If no databases were specified or extracted, use 'mindsdb' as a default
+        # if not all_databases:
+        #     all_databases = [DEFAULT_TEXT2SQL_DATABASE]
+        #
+
+        all_databases = []
+        # Filter out None values
+        all_databases = [db for db in all_databases if db is not None]
+
+        # Create a databases_struct dictionary that includes all extracted databases
+        databases_struct = {}
+
+        # First, add databases from skills with explicit database parameters
+        for skill in skills:
+            if skill.params.get('database'):
+                databases_struct[skill.params['database']] = skill.restriction_on_tables
+
+        # Then, add all extracted databases with no restrictions
+        for db_name in extracted_databases:
+            if db_name not in databases_struct:
+                databases_struct[db_name] = None
 
         sql_agent = SQLAgent(
             command_executor=command_executor,
-            databases=list(set(s.params['database'] for s in skills)),
-            databases_struct={
-                skill.params['database']: skill.restriction_on_tables
-                for skill in skills
-            },
+            databases=all_databases,
+            databases_struct=databases_struct,
             include_tables=tables_list,
             ignore_tables=None,
+            include_knowledge_bases=include_knowledge_bases,
+            ignore_knowledge_bases=ignore_knowledge_bases,
+            knowledge_base_database=knowledge_base_database,
             sample_rows_in_table_info=3,
+
             cache=get_cache('agent', max_size=_MAX_CACHE_SIZE)
         )
         db = MindsDBSQL.custom_init(
             sql_agent=sql_agent
         )
-
-        # Users probably don't need to configure this for now.
-        sql_database_tools = MindsDBSQLToolkit(db=db, llm=llm).get_tools()
-        descriptions = []
-        for skill in skills:
-            description = skill.params.get('description', '')
-            if description:
-                descriptions.append(description)
-
-        for i, tool in enumerate(sql_database_tools):
-            if isinstance(tool, QuerySQLDataBaseTool):
-                # Add our own custom description so our agent knows when to query this table.
-                original_description = tool.description
-                tool.description = ''
-                if len(descriptions) > 0:
-                    tool.description += f'Use this tool if you need data about {" OR ".join(descriptions)}.\n'
-                tool.description += 'Use the conversation context to decide which table to query.\n'
-                if len(tables_list) > 0:
-                    f'These are the available tables: {",".join(tables_list)}.\n'
-                tool.description += (
-                    'ALWAYS consider these special cases:\n'
-                    ' - For TIMESTAMP type columns, make sure you include the time portion in your query (e.g. WHERE date_column = "2020-01-01 12:00:00")\n'
-                    'Here are the rest of the instructions:\n'
-                    f'{original_description}'
-                )
-                sql_database_tools[i] = tool
-        return sql_database_tools
+        toolkit = MindsDBSQLToolkit(db=db, llm=llm)
+        return toolkit.get_tools()
 
     def _make_retrieval_tools(self, skill: db.Skills, llm, embedding_model):
         """