MindsDB 25.5.4.2__py3-none-any.whl → 25.6.2.0__py3-none-any.whl

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

@@ -6,9 +6,7 @@ from langchain_core.tools import BaseTool
 
 
 class KnowledgeBaseListToolInput(BaseModel):
-    tool_input: str = Field(
-        "", description="An empty string to list all knowledge bases."
-    )
+    tool_input: str = Field("", description="An empty string to list all knowledge bases.")
 
 
 class KnowledgeBaseListTool(BaseTool):
@@ -21,7 +19,11 @@ class KnowledgeBaseListTool(BaseTool):
 
     def _run(self, tool_input: str) -> str:
         """List all knowledge bases."""
-        return self.db.get_usable_knowledge_base_names()
+        kb_names = self.db.get_usable_knowledge_base_names()
+        # Convert list to a formatted string for better readability
+        if not kb_names:
+            return "No knowledge bases found."
+        return json.dumps(kb_names)
 
 
 class KnowledgeBaseInfoToolInput(BaseModel):
@@ -41,8 +43,27 @@ class KnowledgeBaseInfoTool(BaseTool):
 
     def _extract_kb_names(self, tool_input: str) -> List[str]:
         """Extract knowledge base names from the tool input."""
+        # First, check if the input is already a list (passed directly from include_knowledge_bases)
+        if isinstance(tool_input, list):
+            return tool_input
+
+        # Next, try to parse it as JSON in case it was serialized as a JSON string
+        try:
+            parsed_input = json.loads(tool_input)
+            if isinstance(parsed_input, list):
+                return parsed_input
+        except (json.JSONDecodeError, TypeError):
+            pass
+
+        # Finally, try the original regex pattern for $START$ and $STOP$ markers
         match = re.search(r"\$START\$(.*?)\$STOP\$", tool_input, re.DOTALL)
         if not match:
+            # If no markers found, check if it's a simple comma-separated string
+            if "," in tool_input:
+                return [kb.strip() for kb in tool_input.split(",")]
+            # If it's just a single string without formatting, return it as a single item
+            if tool_input.strip():
+                return [tool_input.strip()]
             return []
 
         # Extract and clean the knowledge base names
@@ -55,66 +76,84 @@ class KnowledgeBaseInfoTool(BaseTool):
         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$."
+            return "No valid knowledge base names provided. Please provide knowledge base names as a list, comma-separated string, or 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}`;"
-                )
+                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."
-                    )
+                    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"
+
+                # Handle different return types for schema_result
+                if isinstance(schema_result, str):
+                    kb_info += f"{schema_result}\n"
+                elif isinstance(schema_result, list):
+                    for row in schema_result:
+                        if isinstance(row, dict):
+                            kb_info += f"{json.dumps(row, indent=2)}\n"
+                        else:
+                            kb_info += f"{str(row)}\n"
+                else:
+                    kb_info += f"{str(schema_result)}\n"
+
                 kb_info += "```\n\n"
 
+                # Get sample data
+                sample_data = self.db.run_no_throw(f"SELECT * FROM `{kb_name}` LIMIT 10;")
+
                 # 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:
+
+                # Handle different return types for sample_data
+                if not sample_data:
                     kb_info += "No sample data available.\n"
+                elif isinstance(sample_data, str):
+                    kb_info += f"```\n{sample_data}\n```\n"
+                elif isinstance(sample_data, list) and len(sample_data) > 0:
+                    # Only try to extract columns if we have a list of dictionaries
+                    if isinstance(sample_data[0], dict):
+                        # 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:
+                        # If it's a list but not of dictionaries, just format as text
+                        kb_info += "```\n"
+                        for item in sample_data:
+                            kb_info += f"{str(item)}\n"
+                        kb_info += "```\n"
+                else:
+                    # For any other type, just convert to string
+                    kb_info += f"```\n{str(sample_data)}\n```\n"
 
                 results.append(kb_info)
 
             except Exception as e:
-                results.append(
-                    f"Error getting information for knowledge base `{kb_name}`: {str(e)}"
-                )
+                results.append(f"Error getting information for knowledge base `{kb_name}`: {str(e)}")
 
         return "\n\n".join(results)
 
@@ -143,9 +182,7 @@ class KnowledgeBaseQueryTool(BaseTool):
 
         # 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
-        ):
+        if re.search(r"\b(SELECT|FROM|WHERE|LIMIT|ORDER BY)\b", tool_input, re.IGNORECASE):
             return tool_input.strip()
 
         return ""
@@ -185,6 +222,9 @@ class KnowledgeBaseQueryTool(BaseTool):
 
                 return table
 
-            return result
+            # Ensure we always return a string
+            if isinstance(result, (list, dict)):
+                return json.dumps(result, indent=2)
+            return str(result)
         except Exception as e:
             return f"Error executing query: {str(e)}"

mindsdb/interfaces/skills/skills_controller.py

@@ -6,14 +6,19 @@ from sqlalchemy.orm.attributes import flag_modified
 
 from mindsdb.interfaces.storage import db
 from mindsdb.interfaces.database.projects import ProjectController
+from mindsdb.interfaces.data_catalog.data_catalog_loader import DataCatalogLoader
+from mindsdb.interfaces.skills.skill_tool import SkillType
 from mindsdb.utilities.config import config
+from mindsdb.utilities import log
 
 
-default_project = config.get('default_project')
+logger = log.getLogger(__name__)
+
+default_project = config.get("default_project")
 
 
 class SkillsController:
-    '''Handles CRUD operations at the database level for Skills'''
+    """Handles CRUD operations at the database level for Skills"""
 
     def __init__(self, project_controller: ProjectController = None):
         if project_controller is None:
@@ -21,7 +26,7 @@ class SkillsController:
         self.project_controller = project_controller
 
     def get_skill(self, skill_name: str, project_name: str = default_project) -> Optional[db.Skills]:
-        '''
+        """
         Gets a skill by name. Skills are expected to have unique names.
 
         Parameters:
@@ -33,17 +38,17 @@ class SkillsController:
 
         Raises:
             ValueError: If `project_name` does not exist
-        '''
+        """
 
         project = self.project_controller.get(name=project_name)
         return db.Skills.query.filter(
             func.lower(db.Skills.name) == func.lower(skill_name),
             db.Skills.project_id == project.id,
-            db.Skills.deleted_at == null()
+            db.Skills.deleted_at == null(),
         ).first()
 
     def get_skills(self, project_name: Optional[str]) -> List[dict]:
-        '''
+        """
         Gets all skills in a project.
 
         Parameters:
@@ -54,7 +59,7 @@ class SkillsController:
 
         Raises:
             ValueError: If `project_name` does not exist
-        '''
+        """
 
         if project_name is None:
             projects = self.project_controller.get_list()
@@ -63,23 +68,14 @@ class SkillsController:
             project = self.project_controller.get(name=project_name)
             project_ids = [project.id]
 
-        query = (
-            db.session.query(db.Skills)
-            .filter(
-                db.Skills.project_id.in_(project_ids),
-                db.Skills.deleted_at == null()
-            )
+        query = db.session.query(db.Skills).filter(
+            db.Skills.project_id.in_(project_ids), db.Skills.deleted_at == null()
         )
 
         return query.all()
 
-    def add_skill(
-            self,
-            name: str,
-            project_name: str,
-            type: str,
-            params: Dict[str, str] = {}) -> db.Skills:
-        '''
+    def add_skill(self, name: str, project_name: str, type: str, params: Dict[str, str] = {}) -> db.Skills:
+        """
         Adds a skill to the database.
 
         Parameters:
@@ -93,7 +89,7 @@ class SkillsController:
 
         Raises:
             ValueError: If `project_name` does not exist or skill already exists
-        '''
+        """
         if project_name is None:
             project_name = default_project
         project = self.project_controller.get(name=project_name)
@@ -101,7 +97,28 @@ class SkillsController:
         skill = self.get_skill(name, project_name)
 
         if skill is not None:
-            raise ValueError(f'Skill with name already exists: {name}')
+            raise ValueError(f"Skill with name already exists: {name}")
+
+        # Load metadata to data catalog (if enabled) if the skill is Text-to-SQL.
+        if config.get("data_catalog", {}).get("enabled", False):
+            if type == SkillType.TEXT2SQL.value and "include_tables" in params:
+                # TODO: Is it possible to create a skill with complete access to the database with the new agent syntax?
+                # TODO: Handle the case where `ignore_tables` is provided. Is this a valid parameter?
+                # TODO: Knowledge Bases?
+                database_table_map = {}
+                for table in params["include_tables"]:
+                    parts = table.split(".", 1)
+                    database_table_map[parts[0]] = database_table_map.get(parts[0], []) + [parts[1]]
+
+                for database_name, table_names in database_table_map.items():
+                    data_catalog_loader = DataCatalogLoader(database_name=database_name, table_names=table_names)
+                    data_catalog_loader.load_metadata()
+
+            elif type in [SkillType.TEXT2SQL.value, SkillType.TEXT2SQL_LEGACY.value] and "database" in params:
+                data_catalog_loader = DataCatalogLoader(
+                    database_name=params["database"], table_names=params["tables"] if "tables" in params else None
+                )
+                data_catalog_loader.load_metadata()
 
         new_skill = db.Skills(
             name=name,
@@ -115,13 +132,14 @@ class SkillsController:
         return new_skill
 
     def update_skill(
-            self,
-            skill_name: str,
-            new_name: str = None,
-            project_name: str = default_project,
-            type: str = None,
-            params: Dict[str, str] = None):
-        '''
+        self,
+        skill_name: str,
+        new_name: str = None,
+        project_name: str = default_project,
+        type: str = None,
+        params: Dict[str, str] = None,
+    ):
+        """
         Updates an existing skill in the database.
 
         Parameters:
@@ -136,12 +154,12 @@ class SkillsController:
 
         Raises:
             ValueError: If `project_name` does not exist or skill doesn't exist
-        '''
+        """
 
         existing_skill = self.get_skill(skill_name, project_name)
         if existing_skill is None:
-            raise ValueError(f'Skill with name not found: {skill_name}')
-        if isinstance(existing_skill.params, dict) and existing_skill.params.get('is_demo') is True:
+            raise ValueError(f"Skill with name not found: {skill_name}")
+        if isinstance(existing_skill.params, dict) and existing_skill.params.get("is_demo") is True:
             raise ValueError("It is forbidden to change properties of the demo object")
 
         if new_name is not None:
@@ -157,14 +175,14 @@ class SkillsController:
             existing_skill.params = params
             # Some versions of SQL Alchemy won't handle JSON updates correctly without this.
             # See: https://docs.sqlalchemy.org/en/20/orm/session_api.html#sqlalchemy.orm.attributes.flag_modified
-            flag_modified(existing_skill, 'params')
+            flag_modified(existing_skill, "params")
 
         db.session.commit()
 
         return existing_skill
 
     def delete_skill(self, skill_name: str, project_name: str = default_project):
-        '''
+        """
         Deletes a skill by name.
 
         Parameters:
@@ -173,12 +191,12 @@ class SkillsController:
 
         Raises:
             ValueError: If `project_name` does not exist or skill doesn't exist
-        '''
+        """
 
         skill = self.get_skill(skill_name, project_name)
         if skill is None:
             raise ValueError(f"Skill with name doesn't exist: {skill_name}")
-        if isinstance(skill.params, dict) and skill.params.get('is_demo') is True:
+        if isinstance(skill.params, dict) and skill.params.get("is_demo") is True:
             raise ValueError("Unable to delete demo object")
         skill.deleted_at = datetime.datetime.now()
         db.session.commit()