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

mindsdb/interfaces/query_context/context_controller.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional, Iterable
 import pickle
 import datetime as dt
 
@@ -32,7 +32,41 @@ class RunningQuery:
         self.sql = record.sql
         self.database = record.database or config.get('default_project')
 
-    def get_partition_query(self, step_num: int, query: Select) -> Select:
+    def get_partitions(self, dn, step_call, query: Select) -> Iterable:
+        """
+        Gets chunks of data from data handler for executing them in next steps of the planner
+        Check if datanode supports fetch with stream
+        :param dn: datanode to execute query
+        :param step_call: instance of StepCall to get some parameters from it
+        :param query: AST query to execute
+        :return: generator with query results
+        """
+        if dn.has_support_stream():
+            query2 = self.get_partition_query(step_call.current_step_num, query, stream=True)
+
+            for df in dn.query_stream(query2, fetch_size=self.batch_size):
+                max_track_value = self.get_max_track_value(df)
+                yield df
+                self.set_progress(df, max_track_value)
+
+        else:
+            while True:
+                query2 = self.get_partition_query(step_call.current_step_num, query, stream=False)
+
+                response = dn.query(
+                    query=query2,
+                    session=step_call.session
+                )
+                df = response.data_frame
+
+                if df is None or len(df) == 0:
+                    break
+
+                max_track_value = self.get_max_track_value(df)
+                yield df
+                self.set_progress(df, max_track_value)
+
+    def get_partition_query(self, step_num: int, query: Select, stream=False) -> Select:
         """
            Generate query for fetching the next partition
            It wraps query to
@@ -41,16 +75,34 @@ class RunningQuery:
               order by track_column
               limit size {batch_size}
            And fill track_column, previous_value, batch_size
+
+           If steam is true:
+             - if track_column is defined:
+                - don't add limit
+             - else:
+                - return user query without modifications
         """
 
-        track_column = self.record.parameters['track_column']
+        track_column = self.record.parameters.get('track_column')
+        if track_column is None and stream:
+            # if no track column for stream fetching: it is not resumable query, execute original query
+
+            # check if it is first run of the query
+            if self.record.processed_rows > 0:
+                raise RuntimeError("Can't resume query without track_column")
+            return query
+
+        if not stream and track_column is None:
+            raise ValueError('Track column is not defined')
 
         query = Select(
             targets=[Star()],
             from_table=query,
             order_by=[OrderBy(Identifier(track_column))],
-            limit=Constant(self.batch_size)
+
         )
+        if not stream:
+            query.limit = Constant(self.batch_size)
 
         track_value = self.record.context.get('track_value')
         # is it different step?
@@ -114,8 +166,6 @@ class RunningQuery:
             Store parameters of the step which is about to be split into partitions
         """
 
-        if 'track_column' not in params:
-            raise ValueError('Track column is not defined')
         if 'batch_size' not in params:
             params['batch_size'] = 1000
 
@@ -123,15 +173,18 @@ class RunningQuery:
         self.batch_size = self.record.parameters['batch_size']
         db.session.commit()
 
-    def get_max_track_value(self, df: pd.DataFrame) -> pd.DataFrame:
+    def get_max_track_value(self, df: pd.DataFrame) -> Optional[pd.DataFrame]:
         """
             return max value to use in `set_progress`.
             this function is called before execution substeps,
              `set_progress` function - after
         """
-
-        track_column = self.record.parameters['track_column']
-        return df[track_column].max()
+        if 'track_column' in self.record.parameters:
+            track_column = self.record.parameters['track_column']
+            return df[track_column].max()
+        else:
+            # stream mode
+            return None
 
     def set_progress(self, df: pd.DataFrame, max_track_value: int):
         """
@@ -508,6 +561,7 @@ class QueryContextController:
             db.Queries.finished_at < (dt.datetime.now() - dt.timedelta(days=1))
         )
         for rec in remove_query.all():
+            self.get_query(rec.id).remove_from_task()
             db.session.delete(rec)
 
         rec = db.Queries(
@@ -544,6 +598,8 @@ class QueryContextController:
         if rec is None:
             raise RuntimeError(f'Query not found: {query_id}')
 
+        self.get_query(rec.id).remove_from_task()
+
         # the query in progress will fail when it tries to update status
         db.session.delete(rec)
         db.session.commit()

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,
         ]