MindsDB 25.4.2.0__py3-none-any.whl → 25.4.3.0__py3-none-any.whl

mindsdb/interfaces/knowledge_base/controller.py

@@ -3,7 +3,6 @@ import copy
 from typing import Dict, List, Optional
 
 import pandas as pd
-import hashlib
 import numpy as np
 
 from mindsdb_sql_parser.ast import (
@@ -27,6 +26,8 @@ from mindsdb.integrations.libs.vectordatabase_handler import (
 )
 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.handler_utils import get_api_key
+from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embedding_handler import construct_model_from_args, row_to_document
 
 from mindsdb.interfaces.agents.constants import DEFAULT_EMBEDDINGS_MODEL_CLASS
 from mindsdb.interfaces.agents.langchain_agent import create_chat_model, get_llm_provider
@@ -36,6 +37,7 @@ from mindsdb.interfaces.knowledge_base.preprocessing.document_preprocessor impor
 from mindsdb.interfaces.model.functions import PredictorRecordNotFound
 from mindsdb.utilities.exception import EntityExistsError, EntityNotExistsError
 from mindsdb.integrations.utilities.sql_utils import FilterCondition, FilterOperator
+from mindsdb.utilities.context import context as ctx
 
 from mindsdb.api.executor.command_executor import ExecuteCommands
 from mindsdb.utilities import log
@@ -50,6 +52,42 @@ KB_TO_VECTORDB_COLUMNS = {
 }
 
 
+def get_embedding_model_from_params(embedding_model_params: dict):
+    """
+    Create embedding model from parameters.
+    """
+    params_copy = copy.deepcopy(embedding_model_params)
+    provider = params_copy.pop('provider', None).lower()
+    api_key = get_api_key(provider, params_copy, strict=False) or params_copy.get('api_key')
+    # Underscores are replaced because the provider name ultimately gets mapped to a class name.
+    # This is mostly to support Azure OpenAI (azure_openai); the mapped class name is 'AzureOpenAIEmbeddings'.
+    params_copy['class'] = provider.replace('_', '')
+    if provider == 'azure_openai':
+        # Azure OpenAI expects the api_key to be passed as 'openai_api_key'.
+        params_copy['openai_api_key'] = api_key
+    else:
+        params_copy[f"{provider}_api_key"] = api_key
+    params_copy.pop('api_key', None)
+    params_copy['model'] = params_copy.pop('model_name', None)
+
+    return construct_model_from_args(params_copy)
+
+
+def get_reranking_model_from_params(reranking_model_params: dict):
+    """
+    Create reranking model from parameters.
+    """
+    params_copy = copy.deepcopy(reranking_model_params)
+    provider = params_copy.pop('provider', "openai").lower()
+    if provider != 'openai':
+        raise ValueError("Only OpenAI provider is supported for the reranking model.")
+    params_copy[f"{provider}_api_key"] = get_api_key(provider, params_copy, strict=False) or params_copy.get('api_key')
+    params_copy.pop('api_key', None)
+    params_copy['model'] = params_copy.pop('model_name', None)
+
+    return LLMReranker(**params_copy)
+
+
 class KnowledgeBaseTable:
     """
     Knowledge base table interface
@@ -116,19 +154,19 @@ class KnowledgeBaseTable:
         # extract values from conditions and prepare for vectordb
         conditions = []
         query_text = None
-        reranking_threshold = None
+        relevance_threshold = None
         query_conditions = db_handler.extract_conditions(query.where)
         if query_conditions is not None:
             for item in query_conditions:
-                if item.column == "reranking_threshold" and item.op.value == "=":
+                if item.column == "relevance_threshold" and item.op.value == "=":
                     try:
-                        reranking_threshold = float(item.value)
+                        relevance_threshold = float(item.value)
                         # Validate range: must be between 0 and 1
-                        if not (0 <= reranking_threshold <= 1):
-                            raise ValueError(f"reranking_threshold must be between 0 and 1, got: {reranking_threshold}")
-                        logger.debug(f"Found reranking_threshold in query: {reranking_threshold}")
+                        if not (0 <= relevance_threshold <= 1):
+                            raise ValueError(f"relevance_threshold must be between 0 and 1, got: {relevance_threshold}")
+                        logger.debug(f"Found relevance_threshold in query: {relevance_threshold}")
                     except (ValueError, TypeError) as e:
-                        error_msg = f"Invalid reranking_threshold value: {item.value}. {str(e)}"
+                        error_msg = f"Invalid relevance_threshold value: {item.value}. {str(e)}"
                         logger.error(error_msg)
                         raise ValueError(error_msg)
                 elif item.column == TableField.CONTENT.value:
@@ -146,6 +184,16 @@ class KnowledgeBaseTable:
         logger.debug(f"Extracted query text: {query_text}")
 
         self.addapt_conditions_columns(conditions)
+
+        # Set default limit if query is present
+        if query_text is not None:
+            limit = query.limit.value if query.limit is not None else None
+            if limit is None:
+                limit = 10
+            elif limit > 100:
+                limit = 100
+            query.limit = Constant(limit)
+
         df = db_handler.dispatch_select(query, conditions)
         df = self.addapt_result_columns(df)
 
@@ -153,28 +201,27 @@ class KnowledgeBaseTable:
         logger.debug(f"Columns in response: {df.columns.tolist()}")
         # Check if we have a rerank_model configured in KB params
 
-        df = self.add_relevance(df, query_text, reranking_threshold)
+        df = self.add_relevance(df, query_text, relevance_threshold)
 
         # filter by targets
         if requested_kb_columns is not None:
             df = df[requested_kb_columns]
         return df
 
-    def add_relevance(self, df, query_text, reranking_threshold=None):
+    def add_relevance(self, df, query_text, relevance_threshold=None):
         relevance_column = TableField.RELEVANCE.value
 
-        rerank_model = self._kb.params.get("rerank_model")
-        if rerank_model and query_text and len(df) > 0:
+        reranking_model_params = self._kb.params.get("reranking_model")
+        if reranking_model_params and query_text and len(df) > 0:
             # Use reranker for relevance score
             try:
-                logger.info(f"Using reranker model {rerank_model} for relevance calculation")
-                reranker_params = {"model": rerank_model}
+                logger.info(f"Using knowledge reranking model from params: {reranking_model_params}")
                 # Apply custom filtering threshold if provided
-                if reranking_threshold is not None:
-                    reranker_params["filtering_threshold"] = reranking_threshold
-                    logger.info(f"Using custom filtering threshold: {reranking_threshold}")
+                if relevance_threshold is not None:
+                    reranking_model_params["filtering_threshold"] = relevance_threshold
+                    logger.info(f"Using custom filtering threshold: {relevance_threshold}")
 
-                reranker = LLMReranker(**reranker_params)
+                reranker = get_reranking_model_from_params(reranking_model_params)
                 # Get documents to rerank
                 documents = df['chunk_content'].tolist()
                 # Use the get_scores method with disable_events=True
@@ -185,7 +232,7 @@ class KnowledgeBaseTable:
                 # Filter by threshold
                 scores_array = np.array(scores)
                 df = df[scores_array > reranker.filtering_threshold]
-                logger.debug(f"Applied reranking with model {rerank_model}, threshold: {reranker.filtering_threshold}")
+                logger.debug(f"Applied reranking with params: {reranking_model_params}")
             except Exception as e:
                 logger.error(f"Error during reranking: {str(e)}")
                 # Fallback to distance-based relevance
@@ -198,6 +245,8 @@ class KnowledgeBaseTable:
             # Calculate relevance from distance
             logger.info("Calculating relevance from vector distance")
             df[relevance_column] = 1 / (1 + df['distance'])
+            if relevance_threshold is not None:
+                df = df[df[relevance_column] > relevance_threshold]
 
         else:
             df[relevance_column] = None
@@ -293,12 +342,21 @@ class KnowledgeBaseTable:
 
         emb_col = TableField.EMBEDDINGS.value
         cont_col = TableField.CONTENT.value
+
+        db_handler = self.get_vector_db()
+        conditions = db_handler.extract_conditions(query.where)
+        doc_id = None
+        for condition in conditions:
+            if condition.column == 'chunk_id' and condition.op == FilterOperator.EQUAL:
+                doc_id = condition.value
+
         if cont_col in query.update_columns:
             content = query.update_columns[cont_col]
 
             # Apply preprocessing to content if configured
             if self.document_preprocessor:
                 doc = Document(
+                    id=doc_id,
                     content=content.value,
                     metadata={}  # Empty metadata for content-only updates
                 )
@@ -314,8 +372,6 @@ class KnowledgeBaseTable:
         query.table = Identifier(parts=[self._kb.vector_database_table])
 
         # send to vectordb
-        db_handler = self.get_vector_db()
-        conditions = db_handler.extract_conditions(query.where)
         self.addapt_conditions_columns(conditions)
         db_handler.dispatch_update(query, conditions)
 
@@ -369,10 +425,24 @@ class KnowledgeBaseTable:
         db_handler.delete(self._kb.vector_database_table)
 
     def insert(self, df: pd.DataFrame):
-        """Insert dataframe to KB table."""
+        """Insert dataframe to KB table.
+
+        Args:
+            df: DataFrame to insert
+        """
         if df.empty:
             return
 
+        try:
+            run_query_id = ctx.run_query_id
+            # Link current KB to running query (where KB is used to insert data)
+            if run_query_id is not None:
+                self._kb.query_id = run_query_id
+                db.session.commit()
+
+        except AttributeError:
+            ...
+
         # First adapt column names to identify content and metadata columns
         adapted_df = self._adapt_column_names(df)
         content_columns = self._kb.params.get('content_columns', [TableField.CONTENT.value])
@@ -577,36 +647,48 @@ class KnowledgeBaseTable:
         if df.empty:
             return pd.DataFrame([], columns=[TableField.EMBEDDINGS.value])
 
+        # keep only content
+        df = df[[TableField.CONTENT.value]]
+
         model_id = self._kb.embedding_model_id
-        # get the input columns
-        model_rec = db.session.query(db.Predictor).filter_by(id=model_id).first()
+        if model_id:
+            # get the input columns
+            model_rec = db.session.query(db.Predictor).filter_by(id=model_id).first()
 
-        assert model_rec is not None, f"Model not found: {model_id}"
-        model_project = db.session.query(db.Project).filter_by(id=model_rec.project_id).first()
+            assert model_rec is not None, f"Model not found: {model_id}"
+            model_project = db.session.query(db.Project).filter_by(id=model_rec.project_id).first()
 
-        project_datanode = self.session.datahub.get(model_project.name)
+            project_datanode = self.session.datahub.get(model_project.name)
 
-        # keep only content
-        df = df[[TableField.CONTENT.value]]
+            model_using = model_rec.learn_args.get('using', {})
+            input_col = model_using.get('question_column')
+            if input_col is None:
+                input_col = model_using.get('input_column')
+
+            if input_col is not None and input_col != TableField.CONTENT.value:
+                df = df.rename(columns={TableField.CONTENT.value: input_col})
 
-        model_using = model_rec.learn_args.get('using', {})
-        input_col = model_using.get('question_column')
-        if input_col is None:
-            input_col = model_using.get('input_column')
+            df_out = project_datanode.predict(
+                model_name=model_rec.name,
+                df=df,
+                params=self.model_params
+            )
 
-        if input_col is not None and input_col != TableField.CONTENT.value:
-            df = df.rename(columns={TableField.CONTENT.value: input_col})
+            target = model_rec.to_predict[0]
+            if target != TableField.EMBEDDINGS.value:
+                # adapt output for vectordb
+                df_out = df_out.rename(columns={target: TableField.EMBEDDINGS.value})
 
-        df_out = project_datanode.predict(
-            model_name=model_rec.name,
-            df=df,
-            params=self.model_params
-        )
+        elif self._kb.params.get('embedding_model'):
+            embedding_model = get_embedding_model_from_params(self._kb.params.get('embedding_model'))
+
+            df_texts = df.apply(row_to_document, axis=1)
+            embeddings = embedding_model.embed_documents(df_texts.tolist())
+            df_out = df.copy().assign(**{TableField.EMBEDDINGS.value: embeddings})
+
+        else:
+            raise ValueError("No embedding model found for the knowledge base.")
 
-        target = model_rec.to_predict[0]
-        if target != TableField.EMBEDDINGS.value:
-            # adapt output for vectordb
-            df_out = df_out.rename(columns={target: TableField.EMBEDDINGS.value})
         df_out = df_out[[TableField.EMBEDDINGS.value]]
 
         return df_out
@@ -640,9 +722,11 @@ class KnowledgeBaseTable:
             # Extract embedding model args from knowledge base table
             embedding_args = self._kb.embedding_model.learn_args.get('using', {})
             # Construct the embedding model directly
-            from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embedding_handler import construct_model_from_args
             embeddings_model = construct_model_from_args(embedding_args)
             logger.debug(f"Using knowledge base embedding model with args: {embedding_args}")
+        elif self._kb.params.get('embedding_model'):
+            embeddings_model = get_embedding_model_from_params(self._kb.params['embedding_model'])
+            logger.debug(f"Using knowledge base embedding model from params: {self._kb.params['embedding_model']}")
         else:
             embeddings_model = DEFAULT_EMBEDDINGS_MODEL_CLASS()
             logger.debug("Using default embedding model as knowledge base has no embedding model")
@@ -690,22 +774,9 @@ class KnowledgeBaseTable:
         return {}
 
     def _generate_document_id(self, content: str, content_column: str, provided_id: str = None) -> str:
-        """
-        Generate a deterministic document ID from content and column name.
-        If provided_id exists, combines it with content_column.
-
-        Args:
-            content: The content string
-            content_column: Name of the content column
-            provided_id: Optional user-provided ID
-        Returns:
-            Deterministic document ID
-        """
-        if provided_id is not None:
-            return f"{provided_id}_{content_column}"
-
-        id_string = f"content={content}_column={content_column}"
-        return hashlib.sha256(id_string.encode()).hexdigest()
+        """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)
 
     def _convert_metadata_value(self, value):
         """
@@ -788,26 +859,46 @@ class KnowledgeBaseController:
                 return kb
             raise EntityExistsError("Knowledge base already exists", name)
 
-        if embedding_model is None:
-            # create default embedding model
-            model_name = self._get_default_embedding_model(project.name, params=params)
-            params['default_embedding_model'] = model_name
-        else:
-            # get embedding model from input
+        embedding_model_params = params.get('embedding_model', None)
+        reranking_model_params = params.get('reranking_model', None)
+
+        if embedding_model:
             model_name = embedding_model.parts[-1]
 
+        elif embedding_model_params:
+            # Get embedding model from params.
+            # This is called here to check validaity of the parameters.
+            get_embedding_model_from_params(
+                embedding_model_params
+            )
+
+        else:
+            model_name = self._get_default_embedding_model(
+                project.name,
+                params=params
+            )
+            params['default_embedding_model'] = model_name
+
+        model_project = None
         if embedding_model is not None and len(embedding_model.parts) > 1:
             # model project is set
             model_project = self.session.database_controller.get_project(embedding_model.parts[-2])
-        else:
+        elif not embedding_model_params:
             model_project = project
 
-        model = self.session.model_controller.get_model(
-            name=model_name,
-            project_name=model_project.name
-        )
-        model_record = db.Predictor.query.get(model['id'])
-        embedding_model_id = model_record.id
+        embedding_model_id = None
+        if model_project:
+            model = self.session.model_controller.get_model(
+                name=model_name,
+                project_name=model_project.name
+            )
+            model_record = db.Predictor.query.get(model['id'])
+            embedding_model_id = model_record.id
+
+        if reranking_model_params:
+            # Get reranking model from params.
+            # This is called here to check validaity of the parameters.
+            get_reranking_model_from_params(reranking_model_params)
 
         # search for the vector database table
         if storage is None:
@@ -1029,6 +1120,7 @@ class KnowledgeBaseController:
                 'embedding_model': embedding_model.name if embedding_model is not None else None,
                 'vector_database': None if vector_database is None else vector_database.name,
                 'vector_database_table': record.vector_database_table,
+                'query_id': record.query_id,
                 'params': record.params
             })
 

mindsdb/interfaces/knowledge_base/preprocessing/document_preprocessor.py

@@ -1,7 +1,6 @@
 from typing import List, Dict, Optional, Any
 import pandas as pd
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-import hashlib
 import asyncio
 
 
@@ -43,7 +42,11 @@ class DocumentPreprocessor:
         self.splitter = None  # Will be set by child classes
 
     def process_documents(self, documents: List[Document]) -> List[ProcessedChunk]:
-        """Base implementation - should be overridden by child classes"""
+        """Base implementation - should be overridden by child classes
+
+        Args:
+            documents: List of documents to process
+        """
         raise NotImplementedError("Subclasses must implement process_documents")
 
     def _split_document(self, doc: Document) -> List[Document]:
@@ -80,27 +83,22 @@ class DocumentPreprocessor:
             metadata=data.get("metadata", {}),
         )
 
-    def _generate_deterministic_id(
-        self, content: str, content_column: str = None, provided_id: str = None
-    ) -> str:
-        """Generate a deterministic ID based on content and column"""
-        if provided_id is not None:
-            return f"{provided_id}_{content_column}"
-
-        id_string = f"content={content}_column={content_column}"
-        return hashlib.sha256(id_string.encode()).hexdigest()
-
     def _generate_chunk_id(
         self,
         chunk_index: Optional[int] = None,
+        total_chunks: Optional[int] = None,
+        start_char: Optional[int] = None,
+        end_char: Optional[int] = None,
         provided_id: str = None,
     ) -> str:
-        """Generate deterministic ID for a chunk"""
-        base_id = provided_id
-        chunk_id = (
-            f"{base_id}_chunk_{chunk_index}" if chunk_index is not None else base_id
-        )
-        logger.debug(f"Generated chunk ID: {chunk_id} for content hash: {base_id}")
+        """Generate human-readable deterministic ID for a chunk
+        Format: <doc_id>:<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}"
+        logger.debug(f"Generated chunk ID: {chunk_id}")
         return chunk_id
 
     def _prepare_chunk_metadata(
@@ -207,14 +205,10 @@ Please give a short succinct context to situate this chunk within the overall do
         processed_chunks = []
 
         for doc_index, doc in enumerate(documents):
-            # Get content_column from metadata if available
-            content_column = (
-                doc.metadata.get("content_column") if doc.metadata else None
-            )
 
-            # Ensure document has an ID
+            # Document ID must be provided by this point
             if doc.id is None:
-                doc.id = self._generate_deterministic_id(doc.content, content_column)
+                raise ValueError("Document ID must be provided before preprocessing")
 
             # Skip empty or whitespace-only content
             if not doc.content or not doc.content.strip():
@@ -298,68 +292,55 @@ class TextChunkingPreprocessor(DocumentPreprocessor):
         processed_chunks = []
 
         for doc in documents:
-            # Get content_column from metadata if available
-            content_column = (
-                doc.metadata.get("content_column") if doc.metadata else None
-            )
 
-            # Ensure document has an ID
+            # Document ID must be provided by this point
             if doc.id is None:
-                doc.id = self._generate_deterministic_id(doc.content, content_column)
+                raise ValueError("Document ID must be provided before preprocessing")
 
             # Skip empty or whitespace-only content
             if not doc.content or not doc.content.strip():
                 continue
 
             chunk_docs = self._split_document(doc)
+            total_chunks = len(chunk_docs)
 
-            # Single chunk case
-            if len(chunk_docs) == 1:
-                chunk_doc = chunk_docs[0]
+            # Track character positions
+            current_pos = 0
+            for i, chunk_doc in enumerate(chunk_docs):
                 if not chunk_doc.content or not chunk_doc.content.strip():
                     continue
 
+                # Calculate chunk positions
+                start_char = current_pos
+                end_char = start_char + len(chunk_doc.content)
+                current_pos = end_char + 1  # +1 for separator
+
                 # Initialize metadata
                 metadata = {}
                 if doc.metadata:
                     metadata.update(doc.metadata)
 
-                # Pass through doc.id and content_column
-                id = self._generate_chunk_id(
-                    chunk_index=0, provided_id=doc.id
+                # Add position metadata
+                metadata["start_char"] = start_char
+                metadata["end_char"] = end_char
+
+                # Generate chunk ID with total chunks
+                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
                 )
+
                 processed_chunks.append(
                     ProcessedChunk(
-                        id=id,
+                        id=chunk_id,
                         content=chunk_doc.content,
                         embeddings=doc.embeddings,
-                        metadata=self._prepare_chunk_metadata(doc.id, None, metadata),
+                        metadata=self._prepare_chunk_metadata(doc.id, i, metadata),
                     )
                 )
-            else:
-                # Multiple chunks case
-                for i, chunk_doc in enumerate(chunk_docs):
-                    if not chunk_doc.content or not chunk_doc.content.strip():
-                        continue
-
-                    # Initialize metadata
-                    metadata = {}
-                    if doc.metadata:
-                        metadata.update(doc.metadata)
-
-                    # Pass through doc.id and content_column
-                    chunk_id = self._generate_chunk_id(
-                        chunk_index=i,
-                        provided_id=doc.id,
-                    )
-                    processed_chunks.append(
-                        ProcessedChunk(
-                            id=chunk_id,
-                            content=chunk_doc.content,
-                            embeddings=doc.embeddings,
-                            metadata=self._prepare_chunk_metadata(doc.id, i, metadata),
-                        )
-                    )
 
         return processed_chunks
 

mindsdb/interfaces/knowledge_base/utils.py

@@ -0,0 +1,28 @@
+"""Utilities for knowledge base operations."""
+import hashlib
+
+
+def generate_document_id(content: str, content_column: str, provided_id: str = None) -> str:
+    """
+    Generate a deterministic document ID from content and column name.
+    If provided_id exists, combines it with content_column.
+    For generated IDs, uses a short hash of just the content to ensure
+    same content gets same base ID across different columns.
+
+    Args:
+        content: The content string
+        content_column: Name of the content column
+        provided_id: Optional user-provided ID
+    Returns:
+        Deterministic document ID in format: <base_id>_<column>
+        where base_id is either the provided_id or a 16-char hash of content
+    """
+    if provided_id is not None:
+        base_id = provided_id
+    else:
+        # Generate a shorter 16-character hash based only on content
+        hash_obj = hashlib.md5(content.encode())
+        base_id = hash_obj.hexdigest()[:16]
+
+    # Append column name to maintain uniqueness across columns
+    return f"{base_id}_{content_column}"