MindsDB 25.2.3.0__py3-none-any.whl → 25.3.1.0__py3-none-any.whl

mindsdb/interfaces/skills/retrieval_tool.py

@@ -1,33 +1,24 @@
+import traceback
+
 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.rag.settings import RAGPipelineModel
+from mindsdb.integrations.utilities.sql_utils import FilterCondition, FilterOperator
 
 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.documents import Document
 from langchain_core.tools import Tool
+from mindsdb.integrations.libs.response import RESPONSE_TYPE
 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
+def _load_rag_config(tool: dict, pred_args: dict, skill: db.Skills) -> RAGPipelineModel:
     tools_config = tool['config']
     tools_config.update(pred_args)
 
@@ -71,34 +62,132 @@ def build_retrieval_tool(tool: dict, pred_args: dict, skill: db.Skills):
         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)
+    return load_rag_config(tools_config, kb_params, embeddings_model)
+
+
+def _build_rag_pipeline_tool(tool: dict, pred_args: dict, skill: db.Skills):
+    rag_config = _load_rag_config(tool, pred_args, skill)
+    # 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)}")
+            logger.error(traceback.format_exc())
+            return f"Error in retrieval: {str(e)}"
+
+    # Create RAG tool
+    tools_config = tool['config']
+    tools_config.update(pred_args)
+    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)
+    )
+
+
+def _build_name_lookup_tool(tool: dict, pred_args: dict, skill: db.Skills):
+    if 'source' not in tool:
+        raise ValueError("Knowledge base for tool not found")
+    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)
+    vector_db_handler = kb_table.get_vector_db()
+
+    rag_config = _load_rag_config(tool, pred_args, skill)
+    metadata_config = rag_config.metadata_config
+
+    def _get_document_by_name(name: str):
+        if metadata_config.name_column_index is not None:
+            tsquery_str = ' & '.join(name.split(' '))
+            documents_response = vector_db_handler.native_query(
+                f'SELECT * FROM {metadata_config.table} WHERE {metadata_config.name_column_index} @@ to_tsquery(\'{tsquery_str}\') LIMIT 1;'
+            )
+        else:
+            documents_response = vector_db_handler.native_query(
+                f'SELECT * FROM {metadata_config.table} WHERE "{metadata_config.name_column}" ILIKE \'%{name}%\' LIMIT 1;'
+            )
+        if documents_response.resp_type == RESPONSE_TYPE.ERROR:
+            raise RuntimeError(f'There was an error looking up documents: {documents_response.error_message}')
+        if documents_response.data_frame.empty:
+            return None
+        document_row = documents_response.data_frame.head(1)
+        # Restore document from chunks, keeping in mind max context.
+        id_filter_condition = FilterCondition(
+            f"{metadata_config.embeddings_metadata_column}->>'{metadata_config.doc_id_key}'",
+            FilterOperator.EQUAL,
+            str(document_row.get(metadata_config.id_column).item())
+        )
+        document_chunks_df = vector_db_handler.select(
+            metadata_config.embeddings_table,
+            conditions=[id_filter_condition]
         )
+        if document_chunks_df.empty:
+            return None
+        sort_col = 'chunk_id' if 'chunk_id' in document_chunks_df.columns else 'id'
+        document_chunks_df.sort_values(by=sort_col)
+        content = ''
+        for _, chunk in document_chunks_df.iterrows():
+            if len(content) > metadata_config.max_document_context:
+                break
+            content += chunk.get(metadata_config.content_column, '')
+
+        return Document(
+            page_content=content,
+            metadata=document_row.to_dict(orient='records')[0]
+        )
+
+    def _lookup_document_by_name(name: str):
+        found_document = _get_document_by_name(name)
+        if found_document is None:
+            return f'I could not find any document with name {name}. Please make sure the document name matches exactly.'
+        return f"I found document {found_document.metadata.get(metadata_config.id_column)} with name {found_document.metadata.get(metadata_config.name_column)}. Here is the full document to use as context:\n\n{found_document.page_content}"
+
+    return Tool(
+        func=_lookup_document_by_name,
+        name=tool.get('name', '') + '_name_lookup',
+        description='You must use this tool ONLY when the user is asking about a specific document by name or title. The input should be the exact name of the document the user is looking for.',
+        return_direct=False
+    )
 
+
+def build_retrieval_tools(tool: dict, pred_args: dict, skill: db.Skills):
+    """
+    Builds a list of tools for retrieval i.e RAG
+
+    Args:
+        tool: Tool configuration dictionary
+        pred_args: Predictor arguments dictionary
+        skill: Skills database object
+
+    Returns:
+        Tool: Configured list of retrieval tools
+
+    Raises:
+        ValueError: If knowledge base is not found or configuration is invalid
+    """
+    # Catch configuration errors before creating tools.
+    try:
+        rag_config = _load_rag_config(tool, pred_args, skill)
     except Exception as e:
         logger.error(f"Error building RAG pipeline: {str(e)}")
         raise ValueError(f"Failed to build RAG pipeline: {str(e)}")
+    tools = [_build_rag_pipeline_tool(tool, pred_args, skill)]
+    if rag_config.metadata_config is None:
+        return tools
+    tools.append(_build_name_lookup_tool(tool, pred_args, skill))
+    return tools
 
 
 def _get_knowledge_base(knowledge_base_name: str, project_id, executor) -> KnowledgeBase:

mindsdb/interfaces/skills/skill_tool.py

@@ -10,6 +10,7 @@ from mindsdb_sql_parser.ast import Select, BinaryOperation, Identifier, Constant
 
 from mindsdb.utilities import log
 from mindsdb.utilities.cache import get_cache
+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
@@ -106,7 +107,7 @@ class SkillToolController:
             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'
+            sql_session.database = config.get('default_project')
 
             self.command_executor = ExecuteCommands(sql_session)
         return self.command_executor
@@ -222,8 +223,8 @@ class SkillToolController:
         pred_args = {}
         pred_args['llm'] = llm
 
-        from .retrieval_tool import build_retrieval_tool
-        return build_retrieval_tool(tool, pred_args, skill)
+        from .retrieval_tool import build_retrieval_tools
+        return build_retrieval_tools(tool, pred_args, skill)
 
     def _get_rag_query_function(self, skill: db.Skills):
         session_controller = self.get_command_executor().session
@@ -295,10 +296,9 @@ class SkillToolController:
                     for skill in skills
                 ]
             elif skill_type == SkillType.RETRIEVAL:
-                tools[skill_type] = [
-                    self._make_retrieval_tools(skill, llm, embedding_model)
-                    for skill in skills
-                ]
+                tools[skill_type] = []
+                for skill in skills:
+                    tools[skill_type] += self._make_retrieval_tools(skill, llm, embedding_model)
         return tools
 
 

mindsdb/interfaces/skills/skills_controller.py

@@ -1,11 +1,15 @@
 import datetime
 from typing import Dict, List, Optional
 
-from sqlalchemy import null
+from sqlalchemy import null, func
 from sqlalchemy.orm.attributes import flag_modified
 
 from mindsdb.interfaces.storage import db
 from mindsdb.interfaces.database.projects import ProjectController
+from mindsdb.utilities.config import config
+
+
+default_project = config.get('default_project')
 
 
 class SkillsController:
@@ -16,7 +20,7 @@ class SkillsController:
             project_controller = ProjectController()
         self.project_controller = project_controller
 
-    def get_skill(self, skill_name: str, project_name: str = 'mindsdb') -> Optional[db.Skills]:
+    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.
 
@@ -33,7 +37,7 @@ class SkillsController:
 
         project = self.project_controller.get(name=project_name)
         return db.Skills.query.filter(
-            db.Skills.name == skill_name,
+            func.lower(db.Skills.name) == func.lower(skill_name),
             db.Skills.project_id == project.id,
             db.Skills.deleted_at == null()
         ).first()
@@ -90,7 +94,7 @@ class SkillsController:
             ValueError: If `project_name` does not exist or skill already exists
         '''
         if project_name is None:
-            project_name = 'mindsdb'
+            project_name = default_project
         project = self.project_controller.get(name=project_name)
 
         skill = self.get_skill(name, project_name)
@@ -113,7 +117,7 @@ class SkillsController:
             self,
             skill_name: str,
             new_name: str = None,
-            project_name: str = 'mindsdb',
+            project_name: str = default_project,
             type: str = None,
             params: Dict[str, str] = None):
         '''
@@ -158,7 +162,7 @@ class SkillsController:
 
         return existing_skill
 
-    def delete_skill(self, skill_name: str, project_name: str = 'mindsdb'):
+    def delete_skill(self, skill_name: str, project_name: str = default_project):
         '''
         Deletes a skill by name.
 

mindsdb/interfaces/skills/sql_agent.py

@@ -287,6 +287,7 @@ class SQLAgent:
         return info
 
     def _get_sample_rows(self, table: str, fields: List[str]) -> str:
+        logger.info(f'_get_sample_rows: table={table} fields={fields}')
         command = f"select {', '.join(fields)} from {table} limit {self._sample_rows_in_table_info};"
         try:
             ret = self._call_engine(command)
@@ -300,7 +301,7 @@ class SQLAgent:
                 map(lambda row: [truncate_value(value) for value in row], sample_rows))
             sample_rows_str = "\n" + list_to_csv_str([fields] + sample_rows)
         except Exception as e:
-            logger.warning(e)
+            logger.info(f'_get_sample_rows error: {e}')
             sample_rows_str = "\n" + "\t [error] Couldn't retrieve sample rows!"
 
         return sample_rows_str
@@ -347,14 +348,18 @@ class SQLAgent:
 
     def get_table_info_safe(self, table_names: Optional[List[str]] = None) -> str:
         try:
+            logger.info(f'get_table_info_safe: {table_names}')
             return self.get_table_info(table_names)
         except Exception as e:
+            logger.info(f'get_table_info_safe error: {e}')
             return f"Error: {e}"
 
     def query_safe(self, command: str, fetch: str = "all") -> str:
         try:
+            logger.info(f'query_safe (fetch={fetch}): {command}')
             return self.query(command, fetch)
         except Exception as e:
+            logger.info(f'query_safe error: {e}')
             msg = f"Error: {e}"
             if 'does not exist' in msg and ' relation ' in msg:
                 msg += '\nAvailable tables: ' + ', '.join(self.get_usable_table_names())

mindsdb/interfaces/storage/db.py

@@ -10,6 +10,7 @@ from sqlalchemy import (
     DateTime,
     Index,
     Integer,
+    LargeBinary,
     Numeric,
     String,
     UniqueConstraint,
@@ -213,23 +214,12 @@ class Project(Base):
     deleted_at = Column(DateTime)
     name = Column(String, nullable=False)
     company_id = Column(Integer, default=0)
+    metadata_: dict = Column("metadata", JSON, nullable=True)
     __table_args__ = (
         UniqueConstraint("name", "company_id", name="unique_project_name_company_id"),
     )
 
 
-class Log(Base):
-    __tablename__ = "log"
-
-    id = Column(Integer, primary_key=True)
-    created_at = Column(DateTime, default=datetime.datetime.now)
-    log_type = Column(String)  # log, info, warning, traceback etc
-    source = Column(String)  # file + line
-    company_id = Column(Integer)
-    payload = Column(String)
-    created_at_index = Index("some_index", "created_at_index")
-
-
 class Integration(Base):
     __tablename__ = "integration"
     id = Column(Integer, primary_key=True)
@@ -288,8 +278,20 @@ class JsonStorage(Base):
     resource_id = Column(Integer)
     name = Column(String)
     content = Column(JSON)
+    encrypted_content = Column(LargeBinary, nullable=True)
     company_id = Column(Integer)
 
+    def to_dict(self) -> Dict:
+        return {
+            "id": self.id,
+            "resource_group": self.resource_group,
+            "resource_id": self.resource_id,
+            "name": self.name,
+            "content": self.content,
+            "encrypted_content": self.encrypted_content,
+            "company_id": self.company_id,
+        }
+
 
 class Jobs(Base):
     __tablename__ = "jobs"

mindsdb/interfaces/storage/json.py

@@ -1,3 +1,5 @@
+from mindsdb.utilities.functions import decrypt_json, encrypt_json
+from mindsdb.utilities.config import config
 from mindsdb.interfaces.storage import db
 from mindsdb.interfaces.storage.fs import RESOURCE_GROUP
 from mindsdb.utilities.context import context as ctx
@@ -90,8 +92,65 @@ class JsonStorage:
             logger.error('cant delete records from JSON storage')
 
 
+class EncryptedJsonStorage(JsonStorage):
+    def __init__(self, resource_group: str, resource_id: int):
+        super().__init__(resource_group, resource_id)
+        self.secret_key = config.get('secret_key', 'dummy-key')
+
+    def __setitem__(self, key: str, value: dict) -> None:
+        if isinstance(value, dict) is False:
+            raise TypeError(f"got {type(value)} instead of dict")
+
+        encrypted_value = encrypt_json(value, self.secret_key)
+
+        existing_record = self.get_record(key)
+        if existing_record is None:
+            record = db.JsonStorage(
+                name=key,
+                resource_group=self.resource_group,
+                resource_id=self.resource_id,
+                company_id=ctx.company_id,
+                encrypted_content=encrypted_value
+            )
+            db.session.add(record)
+        else:
+            existing_record.encrypted_content = encrypted_value
+        db.session.commit()
+
+    def set_bytes(self, key: str, encrypted_value: bytes):
+        existing_record = self.get_record(key)
+        if existing_record is None:
+            record = db.JsonStorage(
+                name=key,
+                resource_group=self.resource_group,
+                resource_id=self.resource_id,
+                company_id=ctx.company_id,
+                encrypted_content=encrypted_value
+            )
+            db.session.add(record)
+        else:
+            existing_record.encrypted_content = encrypted_value
+        db.session.commit()
+
+    def set_str(self, key: str, encrypted_value: str):
+        self.set_bytes(key, encrypted_value.encode())
+
+    def __getitem__(self, key: str) -> dict:
+        record = self.get_record(key)
+        if record is None:
+            return None
+        return decrypt_json(record.encrypted_content, self.secret_key)
+
+
 def get_json_storage(resource_id: int, resource_group: str = RESOURCE_GROUP.PREDICTOR):
     return JsonStorage(
         resource_group=resource_group,
         resource_id=resource_id,
     )
+
+
+def get_encrypted_json_storage(resource_id: int, resource_group: str = RESOURCE_GROUP.PREDICTOR):
+    return EncryptedJsonStorage(
+        resource_group=resource_group,
+        resource_id=resource_id,
+    )

mindsdb/interfaces/storage/model_fs.py

@@ -1,5 +1,6 @@
 import os
 import re
+import json
 import io
 import zipfile
 from typing import Union
@@ -7,7 +8,10 @@ from typing import Union
 import mindsdb.interfaces.storage.db as db
 
 from .fs import RESOURCE_GROUP, FileStorageFactory, SERVICE_FILES_NAMES
-from .json import get_json_storage
+from .json import get_json_storage, get_encrypted_json_storage
+
+
+JSON_STORAGE_FILE = 'json_storage.json'
 
 
 class ModelStorage:
@@ -119,6 +123,13 @@ class ModelStorage:
         )
         return json_storage.set(name, data)
 
+    def encrypted_json_set(self, name: str, data: dict) -> None:
+        json_storage = get_encrypted_json_storage(
+            resource_id=self.predictor_id,
+            resource_group=RESOURCE_GROUP.PREDICTOR
+        )
+        return json_storage.set(name, data)
+
     def json_get(self, name):
         json_storage = get_json_storage(
             resource_id=self.predictor_id,
@@ -126,6 +137,13 @@ class ModelStorage:
         )
         return json_storage.get(name)
 
+    def encrypted_json_get(self, name: str) -> dict:
+        json_storage = get_encrypted_json_storage(
+            resource_id=self.predictor_id,
+            resource_group=RESOURCE_GROUP.PREDICTOR
+        )
+        return json_storage.get(name)
+
     def json_list(self):
         ...
 
@@ -237,6 +255,13 @@ class HandlerStorage:
         )
         return json_storage.set(name, content)
 
+    def encrypted_json_set(self, name: str, content: dict) -> None:
+        json_storage = get_encrypted_json_storage(
+            resource_id=self.integration_id,
+            resource_group=RESOURCE_GROUP.INTEGRATION
+        )
+        return json_storage.set(name, content)
+
     def json_get(self, name):
         json_storage = get_json_storage(
             resource_id=self.integration_id,
@@ -244,6 +269,13 @@ class HandlerStorage:
         )
         return json_storage.get(name)
 
+    def encrypted_json_get(self, name: str) -> dict:
+        json_storage = get_encrypted_json_storage(
+            resource_id=self.integration_id,
+            resource_group=RESOURCE_GROUP.INTEGRATION
+        )
+        return json_storage.get(name)
+
     def json_list(self):
         ...
 
@@ -251,8 +283,11 @@ class HandlerStorage:
         ...
 
     def export_files(self) -> bytes:
-        if self.is_empty():
+        json_storage = self.export_json_storage()
+
+        if self.is_empty() and not json_storage:
             return None
+
         folder_path = self.folder_get('')
 
         zip_fd = io.BytesIO()
@@ -265,6 +300,11 @@ class HandlerStorage:
                     abs_path = os.path.join(root, file_name)
                     zipf.write(abs_path, os.path.relpath(abs_path, folder_path))
 
+            # If JSON storage is not empty, add it to the zip file.
+            if json_storage:
+                json_str = json.dumps(json_storage)
+                zipf.writestr(JSON_STORAGE_FILE, json_str)
+
         zip_fd.seek(0)
         return zip_fd.read()
 
@@ -277,6 +317,48 @@ class HandlerStorage:
         zip_fd.seek(0)
 
         with zipfile.ZipFile(zip_fd, 'r') as zip_ref:
-            zip_ref.extractall(folder_path)
+            for name in zip_ref.namelist():
+                # If JSON storage file is in the zip file, import the content to the JSON storage.
+                # Thereafter, remove the file from the folder.
+                if name == JSON_STORAGE_FILE:
+                    json_storage = zip_ref.read(JSON_STORAGE_FILE)
+                    self.import_json_storage(json_storage)
+
+                else:
+                    zip_ref.extract(name, folder_path)
 
         self.folder_sync('')
+
+    def export_json_storage(self) -> list[dict]:
+        json_storage = get_json_storage(
+            resource_id=self.integration_id,
+            resource_group=RESOURCE_GROUP.INTEGRATION
+        )
+
+        records = []
+        for record in json_storage.get_all_records():
+            record_dict = record.to_dict()
+            if record_dict.get('encrypted_content'):
+                record_dict['encrypted_content'] = record_dict['encrypted_content'].decode()
+            records.append(record_dict)
+
+        return records
+
+    def import_json_storage(self, records: bytes) -> None:
+        json_storage = get_json_storage(
+            resource_id=self.integration_id,
+            resource_group=RESOURCE_GROUP.INTEGRATION
+        )
+
+        encrypted_json_storage = get_encrypted_json_storage(
+            resource_id=self.integration_id,
+            resource_group=RESOURCE_GROUP.INTEGRATION
+        )
+
+        records = json.loads(records.decode())
+
+        for record in records:
+            if record['encrypted_content']:
+                encrypted_json_storage.set_str(record['name'], record['encrypted_content'])
+            else:
+                json_storage.set(record['name'], record['content'])

mindsdb/interfaces/triggers/triggers_controller.py

@@ -5,6 +5,7 @@ from mindsdb_sql_parser import parse_sql, ParsingException
 from mindsdb.interfaces.storage import db
 from mindsdb.interfaces.database.projects import ProjectController
 from mindsdb.utilities.context import context as ctx
+from mindsdb.utilities.config import config
 
 from mindsdb.api.executor.controllers.session_controller import SessionController
 
@@ -16,7 +17,7 @@ class TriggersController:
         name = name.lower()
 
         if project_name is None:
-            project_name = 'mindsdb'
+            project_name = config.get('default_project')
         project_controller = ProjectController()
         project = project_controller.get(name=project_name)
 

mindsdb/migrations/versions/2022-10-14_43c52d23845a_projects.py

@@ -32,12 +32,26 @@ def upgrade():
         sa.UniqueConstraint('name', 'company_id', name='unique_project_name_company_id')
     )
 
+    project_table = sa.Table(
+        'project',
+        sa.MetaData(),
+        sa.Column('id', sa.Integer()),
+        sa.Column('name', sa.String()),
+        sa.Column('company_id', sa.Integer()),
+    )
+
     conn = op.get_bind()
     session = sa.orm.Session(bind=conn)
 
-    project_record = db.Project(name='mindsdb')
-    session.add(project_record)
-    session.commit()
+    conn.execute(
+        project_table.insert().values(
+            name='mindsdb'
+        )
+    )
+
+    project_record = conn.execute(
+        project_table.select().where(project_table.c.name == 'mindsdb')
+    ).fetchone()
 
     with op.batch_alter_table('predictor', schema=None) as batch_op:
         batch_op.add_column(sa.Column('project_id', sa.Integer()))