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

mindsdb/interfaces/knowledge_base/controller.py

@@ -32,6 +32,7 @@ from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embeddi
 from mindsdb.interfaces.agents.constants import DEFAULT_EMBEDDINGS_MODEL_CLASS
 from mindsdb.interfaces.agents.langchain_agent import create_chat_model, get_llm_provider
 from mindsdb.interfaces.database.projects import ProjectController
+from mindsdb.interfaces.variables.variables_controller import variables_controller
 from mindsdb.interfaces.knowledge_base.preprocessing.models import PreprocessingConfig, Document
 from mindsdb.interfaces.knowledge_base.preprocessing.document_preprocessor import PreprocessorFactory
 from mindsdb.interfaces.model.functions import PredictorRecordNotFound
@@ -41,13 +42,14 @@ from mindsdb.utilities.config import config
 from mindsdb.utilities.context import context as ctx
 
 from mindsdb.api.executor.command_executor import ExecuteCommands
+from mindsdb.api.executor.utilities.sql import query_df
 from mindsdb.utilities import log
 from mindsdb.integrations.utilities.rag.rerankers.base_reranker import BaseLLMReranker
 
 logger = log.getLogger(__name__)
 
 KB_TO_VECTORDB_COLUMNS = {
-    'id': 'original_row_id',
+    'id': 'original_doc_id',
     'chunk_id': 'id',
     'chunk_content': 'content'
 }
@@ -124,6 +126,11 @@ class KnowledgeBaseTable:
         logger.debug(f"Configuring preprocessing with config: {config}")
         self.document_preprocessor = None  # Reset existing preprocessor
         if config is not None:
+            # Ensure content_column is set for JSON chunking if not already specified
+            if config.get('type') == 'json_chunking' and config.get('json_chunking_config'):
+                if 'content_column' not in config['json_chunking_config']:
+                    config['json_chunking_config']['content_column'] = 'content'
+
             preprocessing_config = PreprocessingConfig(**config)
             self.document_preprocessor = PreprocessorFactory.create_preprocessor(preprocessing_config)
             logger.debug(f"Created preprocessor of type: {type(self.document_preprocessor)}")
@@ -150,13 +157,8 @@ class KnowledgeBaseTable:
         query.from_table = Identifier(parts=[self._kb.vector_database_table])
         logger.debug(f"Set table name to: {self._kb.vector_database_table}")
 
-        requested_kb_columns = []
-        for target in query.targets:
-            if isinstance(target, Star):
-                requested_kb_columns = None
-                break
-            else:
-                requested_kb_columns.append(target.parts[-1].lower())
+        # Copy query for complex execution via DuckDB: DISTINCT, GROUP BY etc.
+        query_copy = copy.deepcopy(query)
 
         query.targets = [
             Identifier(TableField.ID.value),
@@ -220,9 +222,17 @@ class KnowledgeBaseTable:
 
         df = self.add_relevance(df, query_text, relevance_threshold)
 
-        # filter by targets
-        if requested_kb_columns is not None:
-            df = df[requested_kb_columns]
+        if (
+            query.group_by is not None
+            or query.order_by is not None
+            or query.having is not None
+            or query.distinct is True
+            or len(query.targets) != 1
+            or not isinstance(query.targets[0], Star)
+        ):
+            query_copy.where = None
+            df = query_df(df, query_copy, session=self.session)
+
         return df
 
     def add_relevance(self, df, query_text, relevance_threshold=None):
@@ -290,7 +300,7 @@ class KnowledgeBaseTable:
         columns = list(df.columns)
         # update id, get from metadata
         df[TableField.ID.value] = df[TableField.METADATA.value].apply(
-            lambda m: None if m is None else m.get('original_row_id')
+            lambda m: None if m is None else m.get('original_doc_id')
         )
 
         # id on first place
@@ -479,12 +489,9 @@ class KnowledgeBaseTable:
                     # Use provided_id directly if it exists, otherwise generate one
                     doc_id = self._generate_document_id(content_str, col, provided_id)
 
-                    # Need provided ID to link chunks back to original source (e.g. database row).
-                    row_id = provided_id if provided_id else idx
-
                     metadata = {
                         **base_metadata,
-                        'original_row_id': str(row_id),
+                        'original_row_index': str(idx),  # provide link to original row index
                         'content_column': col,
                     }
 
@@ -787,7 +794,7 @@ class KnowledgeBaseTable:
     def _generate_document_id(self, content: str, content_column: str, provided_id: str = None) -> str:
         """Generate a deterministic document ID using the utility function."""
         from mindsdb.interfaces.knowledge_base.utils import generate_document_id
-        return generate_document_id(content, content_column, provided_id)
+        return generate_document_id(content=content, provided_id=provided_id)
 
     def _convert_metadata_value(self, value):
         """
@@ -819,6 +826,15 @@ class KnowledgeBaseTable:
         # Convert everything else to string
         return str(value)
 
+    def create_index(self):
+        """
+        Create an index on the knowledge base table
+        :param index_name: name of the index
+        :param params: parameters for the index
+        """
+        db_handler = self.get_vector_db()
+        db_handler.create_index(self._kb.vector_database_table)
+
 
 class KnowledgeBaseController:
     """
@@ -845,6 +861,9 @@ class KnowledgeBaseController:
         :param is_sparse: Whether to use sparse vectors for embeddings
         :param vector_size: Optional size specification for vectors, required when is_sparse=True
         """
+        # fill variables
+        params = variables_controller.fill_parameters(params)
+
         # Validate preprocessing config first if provided
         if preprocessing_config is not None:
             PreprocessingConfig(**preprocessing_config)  # Validate before storing
@@ -1004,6 +1023,8 @@ class KnowledgeBaseController:
         if 'provider' in params:
             engine = params.pop('provider').lower()
 
+        api_key = get_api_key(engine, params, strict=False) or params.pop('api_key')
+
         if engine == 'azure_openai':
             engine = 'openai'
             params['provider'] = 'azure'
@@ -1011,8 +1032,10 @@ class KnowledgeBaseController:
         if engine == 'openai':
             if 'question_column' not in params:
                 params['question_column'] = 'content'
-            if 'api_key' in params:
-                params[f"{engine}_api_key"] = params.pop('api_key')
+            if api_key:
+                params[f"{engine}_api_key"] = api_key
+                if 'api_key' in params:
+                    params.pop('api_key')
             if 'base_url' in params:
                 params['api_base'] = params.pop('base_url')
 
@@ -1150,6 +1173,11 @@ class KnowledgeBaseController:
 
         return data
 
+    def create_index(self, table_name, project_name):
+        project_id = self.session.database_controller.get_project(project_name).id
+        kb_table = self.get_table(table_name, project_id)
+        kb_table.create_index()
+
     def update(self, name: str, project_id: int, **kwargs) -> db.KnowledgeBase:
         """
         Update a knowledge base record

mindsdb/interfaces/knowledge_base/preprocessing/document_preprocessor.py

@@ -25,6 +25,8 @@ from langchain_core.documents import Document as LangchainDocument
 
 logger = log.getLogger(__name__)
 
+_DEFAULT_CONTENT_COLUMN_NAME = "content"
+
 
 class DocumentPreprocessor:
     """Base class for document preprocessing"""
@@ -90,14 +92,18 @@ class DocumentPreprocessor:
         start_char: Optional[int] = None,
         end_char: Optional[int] = None,
         provided_id: str = None,
+        content_column: str = None,
     ) -> str:
         """Generate human-readable deterministic ID for a chunk
-        Format: <doc_id>:<chunk_number>of<total_chunks>:<start_char>to<end_char>
+        Format: <doc_id>:<content_column>:<chunk_number>of<total_chunks>:<start_char>to<end_char>
         """
         if provided_id is None:
             raise ValueError("Document ID must be provided for chunk ID generation")
 
-        chunk_id = f"{provided_id}:{chunk_index + 1}of{total_chunks}:{start_char}to{end_char}"
+        if content_column is None:
+            raise ValueError("Content column must be provided for chunk ID generation")
+
+        chunk_id = f"{provided_id}:{content_column}:{chunk_index + 1}of{total_chunks}:{start_char}to{end_char}"
         logger.debug(f"Generated chunk ID: {chunk_id}")
         return chunk_id
 
@@ -254,8 +260,15 @@ Please give a short succinct context to situate this chunk within the overall do
             if doc.metadata:
                 metadata.update(doc.metadata)
 
+            # Get content_column from metadata or use default
+            content_column = metadata.get('content_column')
+            if content_column is None:
+                # If content_column is not in metadata, use the default column name
+                content_column = _DEFAULT_CONTENT_COLUMN_NAME
+                logger.debug(f"No content_column found in metadata, using default: {_DEFAULT_CONTENT_COLUMN_NAME}")
+
             chunk_id = self._generate_chunk_id(
-                chunk_index=chunk_index, provided_id=doc.id
+                chunk_index=chunk_index, provided_id=doc.id, content_column=content_column
             )
             processed_chunks.append(
                 ProcessedChunk(
@@ -324,13 +337,23 @@ class TextChunkingPreprocessor(DocumentPreprocessor):
                 metadata["start_char"] = start_char
                 metadata["end_char"] = end_char
 
-                # Generate chunk ID with total chunks
+                # Get content_column from metadata or use default
+                content_column = None
+                if doc.metadata:
+                    content_column = doc.metadata.get('content_column')
+
+                if content_column is None:
+                    # If content_column is not in metadata, use the default column name
+                    content_column = _DEFAULT_CONTENT_COLUMN_NAME
+                    logger.debug(f"No content_column found in metadata, using default: {_DEFAULT_CONTENT_COLUMN_NAME}")
+
                 chunk_id = self._generate_chunk_id(
                     chunk_index=i,
                     total_chunks=total_chunks,
                     start_char=start_char,
                     end_char=end_char,
-                    provided_id=doc.id
+                    provided_id=doc.id,
+                    content_column=content_column
                 )
 
                 processed_chunks.append(
@@ -354,18 +377,21 @@ class PreprocessorFactory:
     ) -> DocumentPreprocessor:
         """
         Create appropriate preprocessor based on configuration
-        : param config: Preprocessing configuration
-        : return: Configured preprocessor instance
-        : raises ValueError: If unknown preprocessor type specified
+        :param config: Preprocessing configuration
+        :return: Configured preprocessor instance
+        :raises ValueError: If unknown preprocessor type specified
         """
         if config is None:
+            # Default to text chunking if no config provided
             return TextChunkingPreprocessor()
 
-        if config.type == PreprocessorType.CONTEXTUAL:
-            return ContextualPreprocessor(
-                config.contextual_config or ContextualConfig()
-            )
-        elif config.type == PreprocessorType.TEXT_CHUNKING:
+        if config.type == PreprocessorType.TEXT_CHUNKING:
             return TextChunkingPreprocessor(config.text_chunking_config)
-
-        raise ValueError(f"Unknown preprocessor type: {config.type}")
+        elif config.type == PreprocessorType.CONTEXTUAL:
+            return ContextualPreprocessor(config.contextual_config)
+        elif config.type == PreprocessorType.JSON_CHUNKING:
+            # Import here to avoid circular imports
+            from mindsdb.interfaces.knowledge_base.preprocessing.json_chunker import JSONChunkingPreprocessor
+            return JSONChunkingPreprocessor(config.json_chunking_config)
+        else:
+            raise ValueError(f"Unknown preprocessor type: {config.type}")