MindsDB 1.2.9__py3-none-any.whl → 24.12.4.0__py3-none-any.whl

mindsdb/interfaces/skills/retrieval_tool.py

@@ -0,0 +1,110 @@
+from mindsdb.integrations.utilities.rag.rag_pipeline_builder import RAG
+from mindsdb.integrations.utilities.rag.config_loader import load_rag_config
+
+from mindsdb.interfaces.agents.constants import DEFAULT_EMBEDDINGS_MODEL_CLASS
+from mindsdb.interfaces.skills.skill_tool import skill_tool
+from mindsdb.interfaces.storage import db
+from mindsdb.interfaces.storage.db import KnowledgeBase
+from mindsdb.utilities import log
+from langchain_core.tools import Tool
+from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embedding_handler import construct_model_from_args
+
+logger = log.getLogger(__name__)
+
+
+def build_retrieval_tool(tool: dict, pred_args: dict, skill: db.Skills):
+    """
+    Builds a retrieval tool i.e RAG
+
+    Args:
+        tool: Tool configuration dictionary
+        pred_args: Predictor arguments dictionary
+        skill: Skills database object
+
+    Returns:
+        Tool: Configured retrieval tool
+
+    Raises:
+        ValueError: If knowledge base is not found or configuration is invalid
+    """
+    # build RAG config
+    tools_config = tool['config']
+    tools_config.update(pred_args)
+
+    kb_params = {}
+    embeddings_model = None
+
+    if 'source' in tool:
+        kb_name = tool['source']
+        executor = skill_tool.get_command_executor()
+        kb = _get_knowledge_base(kb_name, skill.project_id, executor)
+
+        if not kb:
+            raise ValueError(f"Knowledge base not found: {kb_name}")
+
+        kb_table = executor.session.kb_controller.get_table(kb.name, kb.project_id)
+        kb_params = {
+            'vector_store_config': {
+                'kb_table': kb_table
+            }
+        }
+
+        # Get embedding model from knowledge base table
+        if kb_table._kb.embedding_model:
+            # Extract embedding model args from knowledge base table
+            embedding_args = kb_table._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}")
+        else:
+            embeddings_model = DEFAULT_EMBEDDINGS_MODEL_CLASS()
+            logger.debug("Using default embedding model as knowledge base has no embedding model")
+    elif 'embedding_model' not in tools_config:
+        embeddings_model = DEFAULT_EMBEDDINGS_MODEL_CLASS()
+        logger.debug("Using default embedding model as no knowledge base provided")
+
+    # Load and validate config
+    try:
+        rag_config = load_rag_config(tools_config, kb_params, embeddings_model)
+        # build retriever
+        rag_pipeline = RAG(rag_config)
+        logger.debug(f"RAG pipeline created with config: {rag_config}")
+
+        def rag_wrapper(query: str) -> str:
+            try:
+                result = rag_pipeline(query)
+                logger.debug(f"RAG pipeline result: {result}")
+                return result['answer']
+            except Exception as e:
+                logger.error(f"Error in RAG pipeline: {str(e)}")
+                return f"Error in retrieval: {str(e)}"
+
+        # Create RAG tool
+        return Tool(
+            func=rag_wrapper,
+            name=tool['name'],
+            description=tool['description'],
+            response_format='content',
+            # Return directly by default since we already use an LLM against retrieved context to generate a response.
+            return_direct=tools_config.get('return_direct', True)
+        )
+
+    except Exception as e:
+        logger.error(f"Error building RAG pipeline: {str(e)}")
+        raise ValueError(f"Failed to build RAG pipeline: {str(e)}")
+
+
+def _get_knowledge_base(knowledge_base_name: str, project_id, executor) -> KnowledgeBase:
+    """
+    Get knowledge base by name and project ID
+
+    Args:
+        knowledge_base_name: Name of the knowledge base
+        project_id: Project ID
+        executor: Command executor instance
+
+    Returns:
+        KnowledgeBase: Knowledge base instance if found, None otherwise
+    """
+    kb = executor.session.kb_controller.get(knowledge_base_name, project_id)
+    return kb

mindsdb/interfaces/skills/skill_tool.py

@@ -0,0 +1,255 @@
+import enum
+from collections import defaultdict
+from typing import List, Optional
+from dataclasses import dataclass
+
+from langchain_core.embeddings import Embeddings
+from langchain_core.language_models import BaseChatModel
+from mindsdb_sql_parser.ast import Select, BinaryOperation, Identifier, Constant, Star
+
+from mindsdb.integrations.libs.vectordatabase_handler import TableField
+from mindsdb.interfaces.skills.sql_agent import SQLAgent
+from mindsdb.interfaces.storage import db
+from mindsdb.utilities import log
+from mindsdb.utilities.cache import get_cache
+
+
+_DEFAULT_TOP_K_SIMILARITY_SEARCH = 5
+_MAX_CACHE_SIZE = 1000
+
+logger = log.getLogger(__name__)
+
+
+class SkillType(enum.Enum):
+    TEXT2SQL_LEGACY = 'text2sql'
+    TEXT2SQL = 'sql'
+    KNOWLEDGE_BASE = 'knowledge_base'
+    RETRIEVAL = 'retrieval'
+
+
+@dataclass
+class SkillData:
+    """Storage for skill's data
+
+    Attributes:
+        name (str): name of the skill
+        type (str): skill's type (SkillType)
+        params (dict): skill's attributes
+        project_id (int): id of the project
+        agent_tables_list (Optional[List[str]]): the restriction on available tables for an agent using the skill
+    """
+    name: str
+    type: str
+    params: dict
+    project_id: int
+    agent_tables_list: Optional[List[str]]
+
+    @property
+    def tables_list(self) -> List[str]:
+        """List of tables which may use this skill. If the list is empty, there are no restrictions.
+        The result list is a combination of skill's and agent's tables lists.
+
+        Returns:
+            List[str]: List of tables.
+
+        Raises:
+            ValueError: if there is no intersection between skill's and agent's list.
+                This means that all tables restricted for use.
+        """
+        agent_tables_list = self.agent_tables_list or []
+        skill_tables_list = self.params.get('tables', [])
+        if len(skill_tables_list) > 0 and len(agent_tables_list) > 0:
+            diff = set(skill_tables_list) & set(agent_tables_list)
+            if len(diff) == 0:
+                raise ValueError("There are no tables allowed for use.")
+            return list(diff)
+        if len(skill_tables_list) > 0:
+            return skill_tables_list
+        return agent_tables_list
+
+
+class SkillToolController:
+    def __init__(self):
+        self.command_executor = None
+
+    def get_command_executor(self):
+        if self.command_executor is None:
+            from mindsdb.api.executor.command_executor import ExecuteCommands
+            from mindsdb.api.executor.controllers import SessionController  # Top-level import produces circular import in some cases TODO: figure out a fix without losing runtime improvements (context: see #9304)  # noqa
+
+            sql_session = SessionController()
+            sql_session.database = 'mindsdb'
+
+            self.command_executor = ExecuteCommands(sql_session)
+        return self.command_executor
+
+    def get_sql_agent(
+            self,
+            database: str,
+            include_tables: Optional[List[str]] = None,
+            ignore_tables: Optional[List[str]] = None,
+            sample_rows_in_table_info: int = 3,
+    ):
+        return SQLAgent(
+            self.get_command_executor(),
+            database,
+            include_tables,
+            ignore_tables,
+            sample_rows_in_table_info,
+            cache=get_cache('agent', max_size=_MAX_CACHE_SIZE)
+        )
+
+    def _make_text_to_sql_tools(self, skills: List[db.Skills], llm) -> List:
+        '''
+           Uses SQLAgent to execute tool
+        '''
+        # To prevent dependency on Langchain unless an actual tool uses it.
+        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]`')
+
+        tables_list = []
+        for skill in skills:
+            database = skill.params['database']
+            for table in skill.tables_list:
+                tables_list.append(f'{database}.{table}')
+
+        # use list databases
+        database = ','.join(set(s.params['database'] for s in skills))
+        db = MindsDBSQL(
+            engine=self.get_command_executor(),
+            database=database,
+            metadata=self.get_command_executor().session.integration_controller,
+            include_tables=tables_list
+        )
+
+        # 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.
+                tool.description = (
+                    f'Use this tool if you need data about {" OR ".join(descriptions)}. '
+                    'Use the conversation context to decide which table to query. '
+                    f'These are the available tables: {",".join(tables_list)}.\n' if len(tables_list) > 0 else '\n'
+                    f'ALWAYS consider these special cases:\n'
+                    f'- 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")'
+                    f'Here are the rest of the instructions:\n'
+                    f'{tool.description}'
+                )
+                sql_database_tools[i] = tool
+        return sql_database_tools
+
+    def _make_retrieval_tools(self, skill: db.Skills, llm, embedding_model):
+        """
+        creates advanced retrieval tool i.e. RAG
+        """
+        params = skill.params
+        config = params.get('config', {})
+        if 'llm' not in config:
+            # Set LLM if not explicitly provided in configs.
+            config['llm'] = llm
+        tool = dict(
+            name=params.get('name', skill.name),
+            source=params.get('source', None),
+            config=config,
+            description=f'You must use this tool to get more context or information '
+                        f'to answer a question about {params["description"]}. '
+                        f'The input should be the exact question the user is asking.',
+            type=skill.type
+        )
+        pred_args = {}
+        pred_args['llm'] = llm
+
+        from .retrieval_tool import build_retrieval_tool
+        return build_retrieval_tool(tool, pred_args, skill)
+
+    def _get_rag_query_function(self, skill: db.Skills):
+
+        session_controller = self.get_command_executor().session
+
+        def _answer_question(question: str) -> str:
+            knowledge_base_name = skill.params['source']
+
+            # make select in KB table
+            query = Select(
+                targets=[Star()],
+                where=BinaryOperation(op='=', args=[
+                    Identifier(TableField.CONTENT.value), Constant(question)
+                ]),
+                limit=Constant(_DEFAULT_TOP_K_SIMILARITY_SEARCH),
+            )
+            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)
+
+        return _answer_question
+
+    def _make_knowledge_base_tools(self, skill: db.Skills) -> dict:
+        # To prevent dependency on Langchain unless an actual tool uses it.
+        description = skill.params.get('description', '')
+
+        logger.warning("This skill is deprecated and will be removed in the future. "
+                       "Please use `retrieval` skill instead ")
+
+        return dict(
+            name='Knowledge Base Retrieval',
+            func=self._get_rag_query_function(skill),
+            description=f'Use this tool to get more context or information to answer a question about {description}. The input should be the exact question the user is asking.',
+            type=skill.type
+        )
+
+    def get_tools_from_skills(self, skills_data: List[SkillData], llm: BaseChatModel, embedding_model: Embeddings) -> dict:
+        """Creates function for skill and metadata (name, description)
+
+        Args:
+            skills_data (List[SkillData]): Skills to make a tool from
+            llm (BaseChatModel): LLM which will be used by skills
+            embedding_model (Embeddings): this model is used by retrieval skill
+
+        Returns:
+            dict: with keys: name, description, func
+        """
+
+        # group skills by type
+        skills_group = defaultdict(list)
+        for skill in skills_data:
+            try:
+                skill_type = SkillType(skill.type)
+            except ValueError:
+                raise NotImplementedError(
+                    f'skill of type {skill.type} is not supported as a tool, supported types are: {list(SkillType._member_names_)}')
+
+            if skill_type == SkillType.TEXT2SQL_LEGACY:
+                skill_type = SkillType.TEXT2SQL
+            skills_group[skill_type].append(skill)
+
+        tools = {}
+        for skill_type, skills in skills_group.items():
+            if skill_type == SkillType.TEXT2SQL:
+                tools[skill_type] = self._make_text_to_sql_tools(skills, llm)
+            elif skill_type == SkillType.KNOWLEDGE_BASE:
+                tools[skill_type] = [
+                    self._make_knowledge_base_tools(skill)
+                    for skill in skills
+                ]
+            elif skill_type == SkillType.RETRIEVAL:
+                tools[skill_type] = [
+                    self._make_retrieval_tools(skill, llm, embedding_model)
+                    for skill in skills
+                ]
+        return tools
+
+
+skill_tool = SkillToolController()

mindsdb/interfaces/skills/skills_controller.py

@@ -0,0 +1,175 @@
+import datetime
+from typing import Dict, List
+
+from sqlalchemy import null
+from sqlalchemy.orm.attributes import flag_modified
+
+from mindsdb.interfaces.storage import db
+from mindsdb.interfaces.database.projects import ProjectController
+
+
+class SkillsController:
+    '''Handles CRUD operations at the database level for Skills'''
+
+    def __init__(self, project_controller: ProjectController = None):
+        if project_controller is None:
+            project_controller = ProjectController()
+        self.project_controller = project_controller
+
+    def get_skill(self, skill_name: str, project_name: str = 'mindsdb') -> db.Skills:
+        '''
+        Gets a skill by name. Skills are expected to have unique names.
+
+        Parameters:
+            skill_name (str): The name of the skill
+            project_name (str): The name of the containing project
+
+        Returns:
+            skill (db.Skills): The database skill object
+
+        Raises:
+            ValueError: If `project_name` does not exist
+        '''
+
+        project = self.project_controller.get(name=project_name)
+        return db.Skills.query.filter(
+            db.Skills.name == skill_name,
+            db.Skills.project_id == project.id,
+            db.Skills.deleted_at == null()
+        ).first()
+
+    def get_skills(self, project_name: str) -> List[dict]:
+        '''
+        Gets all skills in a project.
+
+        Parameters:
+            project_name (str): The name of the containing project
+
+        Returns:
+            all_skills (List[db.Skills]): List of database skill object
+
+        Raises:
+            ValueError: If `project_name` does not exist
+        '''
+
+        project_controller = ProjectController()
+        projects = project_controller.get_list()
+        if project_name is not None:
+            projects = list([p for p in projects if p.name == project_name])
+        project_ids = list([p.id for p in projects])
+
+        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:
+        '''
+        Adds a skill to the database.
+
+        Parameters:
+            name (str): The name of the new skill
+            project_name (str): The containing project
+            type (str): The type of the skill (e.g. Knowledge Base)
+            params: (Dict[str, str]): Parameters associated with the skill
+
+        Returns:
+            bot (db.Skills): The created skill
+
+        Raises:
+            ValueError: If `project_name` does not exist or skill already exists
+        '''
+        if project_name is None:
+            project_name = 'mindsdb'
+        project = self.project_controller.get(name=project_name)
+
+        skill = self.get_skill(name, project_name)
+
+        if skill is not None:
+            raise ValueError(f'Skill with name already exists: {name}')
+
+        new_skill = db.Skills(
+            name=name,
+            project_id=project.id,
+            type=type,
+            params=params,
+        )
+        db.session.add(new_skill)
+        db.session.commit()
+
+        return new_skill
+
+    def update_skill(
+            self,
+            skill_name: str,
+            new_name: str = None,
+            project_name: str = 'mindsdb',
+            type: str = None,
+            params: Dict[str, str] = None):
+        '''
+        Updates an existing skill in the database.
+
+        Parameters:
+            skill_name (str): The name of the new skill, or existing skill to update
+            new_name (str): Updated name of the skill
+            project_name (str): The containing project
+            type (str): The type of the skill (e.g. Knowledge Base)
+            params: (Dict[str, str]): Parameters associated with the skill
+
+        Returns:
+            bot (db.Skills): The updated skill
+
+        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 new_name is not None:
+            existing_skill.name = new_name
+        if type is not None:
+            existing_skill.type = type
+        if params is not None:
+            # Merge params on update
+            existing_params = {} if not existing_skill.params else existing_skill.params
+            existing_params.update(params)
+            # Remove None values entirely.
+            params = {k: v for k, v in existing_params.items() if v is not None}
+            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')
+
+        db.session.commit()
+
+        return existing_skill
+
+    def delete_skill(self, skill_name: str, project_name: str = 'mindsdb'):
+        '''
+        Deletes a skill by name.
+
+        Parameters:
+            skill_name (str): The name of the skill to delete
+            project_name (str): The name of the containing project
+
+        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}")
+        skill.deleted_at = datetime.datetime.now()
+        db.session.commit()