MindsDB 25.9.3rc1__py3-none-any.whl → 25.10.0rc1__py3-none-any.whl

mindsdb/interfaces/knowledge_base/controller.py

@@ -10,7 +10,6 @@ from pydantic import BaseModel, ValidationError
 from sqlalchemy.orm.attributes import flag_modified
 
 from mindsdb_sql_parser.ast import BinaryOperation, Constant, Identifier, Select, Update, Delete, Star
-from mindsdb_sql_parser.ast.mindsdb import CreatePredictor
 from mindsdb_sql_parser import parse_sql
 
 from mindsdb.integrations.libs.keyword_search_base import KeywordSearchBase
@@ -23,6 +22,7 @@ from mindsdb.integrations.libs.vectordatabase_handler import (
     VectorStoreHandler,
 )
 from mindsdb.integrations.utilities.handler_utils import get_api_key
+from mindsdb.integrations.utilities.handlers.auth_utilities.snowflake import get_validated_jwt
 
 from mindsdb.interfaces.agents.constants import DEFAULT_EMBEDDINGS_MODEL_CLASS, MAX_INSERT_BATCH_SIZE
 from mindsdb.interfaces.agents.langchain_agent import create_chat_model, get_llm_provider
@@ -42,6 +42,7 @@ 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
+from mindsdb.interfaces.knowledge_base.llm_client import LLMClient
 
 logger = log.getLogger(__name__)
 
@@ -72,6 +73,10 @@ def get_model_params(model_params: dict, default_config_key: str):
         if not isinstance(model_params, dict):
             raise ValueError("Model parameters must be passed as a JSON object")
 
+        # if provider mismatches - don't use default values
+        if "provider" in model_params and model_params["provider"] != combined_model_params.get("provider"):
+            return model_params
+
         combined_model_params.update(model_params)
 
     combined_model_params.pop("use_default_llm", None)
@@ -142,6 +147,28 @@ def to_json(obj):
         return obj
 
 
+def rotate_provider_api_key(params):
+    """
+    Check api key for specific providers. At the moment it checks and updated jwt token of snowflake provider
+    :param params: input params, can be modified by this function
+    :return: a new api key if it is refreshed
+    """
+    provider = params.get("provider").lower()
+
+    if provider == "snowflake":
+        api_key = params.get("api_key")
+        api_key2 = get_validated_jwt(
+            api_key,
+            account=params.get("snowflake_account_id"),
+            user=params.get("user"),
+            private_key=params.get("private_key"),
+        )
+        if api_key2 != api_key:
+            # update keys
+            params["api_key"] = api_key2
+            return api_key2
+
+
 class KnowledgeBaseTable:
     """
     Knowledge base table interface
@@ -194,6 +221,22 @@ class KnowledgeBaseTable:
         executor = KnowledgeBaseQueryExecutor(self)
         df = executor.run(query)
 
+        # copy metadata to columns
+        if "metadata" in df.columns:
+            meta_columns = self._get_allowed_metadata_columns()
+            if meta_columns:
+                meta_data = pd.json_normalize(df["metadata"])
+                # exclude absent columns and used colunns
+                df_columns = list(df.columns)
+                meta_columns = list(set(meta_columns).intersection(meta_data.columns).difference(df_columns))
+
+                # add columns
+                df = df.join(meta_data[meta_columns])
+
+                # put metadata in the end
+                df_columns.remove("metadata")
+                df = df[df_columns + meta_columns + ["metadata"]]
+
         if (
             query_copy.group_by is not None
             or query_copy.order_by is not None
@@ -384,7 +427,7 @@ class KnowledgeBaseTable:
         # if relevance filtering method is strictly GREATER THAN we filter the df
         if gt_filtering:
             relevance_scores = TableField.RELEVANCE.value
-            df = df[relevance_scores > relevance_threshold]
+            df = df[df[relevance_scores] > relevance_threshold]
 
         return df
 
@@ -402,6 +445,7 @@ class KnowledgeBaseTable:
         return [col.lower() for col in columns]
 
     def score_documents(self, query_text, documents, reranking_model_params):
+        rotate_provider_api_key(reranking_model_params)
         reranker = get_reranking_model_from_params(reranking_model_params)
         return reranker.get_scores(query_text, documents)
 
@@ -412,6 +456,15 @@ class KnowledgeBaseTable:
         if reranking_model_params and query_text and len(df) > 0 and not disable_reranking:
             # Use reranker for relevance score
 
+            new_api_key = rotate_provider_api_key(reranking_model_params)
+            if new_api_key:
+                # update key
+                if "reranking_model" not in self._kb.params:
+                    self._kb.params["reranking_model"] = {}
+                self._kb.params["reranking_model"]["api_key"] = new_api_key
+                flag_modified(self._kb, "params")
+                db.session.commit()
+
             # Apply custom filtering threshold if provided
             if relevance_threshold is not None:
                 reranking_model_params["filtering_threshold"] = relevance_threshold
@@ -864,10 +917,12 @@ class KnowledgeBaseTable:
         model_id = self._kb.embedding_model_id
 
         if model_id is None:
-            # call litellm handler
             messages = list(df[TableField.CONTENT.value])
             embedding_params = get_model_params(self._kb.params.get("embedding_model", {}), "default_embedding_model")
-            results = self.call_litellm_embedding(self.session, embedding_params, messages)
+
+            llm_client = LLMClient(embedding_params, session=self.session)
+            results = llm_client.embeddings(messages)
+
             results = [[val] for val in results]
             return pd.DataFrame(results, columns=[TableField.EMBEDDINGS.value])
 
@@ -1053,6 +1108,26 @@ class KnowledgeBaseController:
     def __init__(self, session) -> None:
         self.session = session
 
+    def _check_kb_input_params(self, params):
+        # check names and types KB params
+        try:
+            KnowledgeBaseInputParams.model_validate(params)
+        except ValidationError as e:
+            problems = []
+            for error in e.errors():
+                parameter = ".".join([str(i) for i in error["loc"]])
+                param_type = error["type"]
+                if param_type == "extra_forbidden":
+                    msg = f"Parameter '{parameter}' is not allowed"
+                else:
+                    msg = f"Error in '{parameter}' (type: {param_type}): {error['msg']}. Input: {repr(error['input'])}"
+                problems.append(msg)
+
+            msg = "\n".join(problems)
+            if len(problems) > 1:
+                msg = "\n" + msg
+            raise ValueError(f"Problem with knowledge base parameters: {msg}") from e
+
     def add(
         self,
         name: str,
@@ -1070,36 +1145,18 @@ 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
         """
-        if not name.islower():
-            raise ValueError(f"The name must be in lower case: {name}")
 
         # fill variables
         params = variables_controller.fill_parameters(params)
 
-        try:
-            KnowledgeBaseInputParams.model_validate(params)
-        except ValidationError as e:
-            problems = []
-            for error in e.errors():
-                parameter = ".".join([str(i) for i in error["loc"]])
-                param_type = error["type"]
-                if param_type == "extra_forbidden":
-                    msg = f"Parameter '{parameter}' is not allowed"
-                else:
-                    msg = f"Error in '{parameter}' (type: {param_type}): {error['msg']}. Input: {repr(error['input'])}"
-                problems.append(msg)
-
-            msg = "\n".join(problems)
-            if len(problems) > 1:
-                msg = "\n" + msg
-            raise ValueError(f"Problem with knowledge base parameters: {msg}") from e
-
         # Validate preprocessing config first if provided
         if preprocessing_config is not None:
             PreprocessingConfig(**preprocessing_config)  # Validate before storing
             params = params or {}
             params["preprocessing"] = preprocessing_config
 
+        self._check_kb_input_params(params)
+
         # Check if vector_size is provided when using sparse vectors
         is_sparse = params.get("is_sparse")
         vector_size = params.get("vector_size")
@@ -1110,8 +1167,6 @@ class KnowledgeBaseController:
         project = self.session.database_controller.get_project(project_name)
         project_id = project.id
 
-        # not difference between cases in sql
-        name = name.lower()
         # check if knowledge base already exists
         kb = self.get(name, project_id)
         if kb is not None:
@@ -1123,42 +1178,25 @@ class KnowledgeBaseController:
         params["embedding_model"] = embedding_params
 
         # if model_name is None:  # Legacy
-        model_name = self._create_embedding_model(
+        self._check_embedding_model(
             project.name,
             params=embedding_params,
             kb_name=name,
         )
-        if model_name is not None:
-            params["created_embedding_model"] = model_name
-
-        embedding_model_id = None
-        if model_name is not None:
-            model = self.session.model_controller.get_model(name=model_name, project_name=project.name)
-            model_record = db.Predictor.query.get(model["id"])
-            embedding_model_id = model_record.id
-
-            if model_record.learn_args.get("using", {}).get("sparse"):
-                is_sparse = True
 
         # if params.get("reranking_model", {}) is bool and False we evaluate it to empty dictionary
         reranking_model_params = params.get("reranking_model", {})
 
         if isinstance(reranking_model_params, bool) and not reranking_model_params:
             params["reranking_model"] = {}
-        # if params.get("reranking_model", {}) is string and false in any case we evaluate it to empty dictionary
-        if isinstance(reranking_model_params, str) and reranking_model_params.lower() == "false":
-            params["reranking_model"] = {}
 
         reranking_model_params = get_model_params(reranking_model_params, "default_reranking_model")
         params["reranking_model"] = reranking_model_params
         if reranking_model_params:
             # Get reranking model from params.
             # This is called here to check validaity of the parameters.
-            try:
-                reranker = get_reranking_model_from_params(reranking_model_params)
-                reranker.get_scores("test", ["test"])
-            except (ValueError, RuntimeError) as e:
-                raise RuntimeError(f"Problem with reranker config: {e}") from e
+            rotate_provider_api_key(reranking_model_params)
+            self._test_reranking(reranking_model_params)
 
         # search for the vector database table
         if storage is None:
@@ -1211,13 +1249,115 @@ class KnowledgeBaseController:
             project_id=project_id,
             vector_database_id=vector_database_id,
             vector_database_table=vector_table_name,
-            embedding_model_id=embedding_model_id,
+            embedding_model_id=None,
             params=params,
         )
         db.session.add(kb)
         db.session.commit()
         return kb
 
+    def update(
+        self,
+        name: str,
+        project_name: str,
+        params: dict,
+        preprocessing_config: Optional[dict] = None,
+    ) -> db.KnowledgeBase:
+        """
+        Update the knowledge base
+        :param name: The name of the knowledge base
+        :param project_name: Current project name
+        :param params: The parameters to update
+        :param preprocessing_config: Optional preprocessing configuration to validate and store
+        """
+
+        # 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
+            params = params or {}
+            params["preprocessing"] = preprocessing_config
+
+        self._check_kb_input_params(params)
+
+        # get project id
+        project = self.session.database_controller.get_project(project_name)
+        project_id = project.id
+
+        # get existed KB
+        kb = self.get(name.lower(), project_id)
+        if kb is None:
+            raise EntityNotExistsError("Knowledge base doesn't exists", name)
+
+        if "embedding_model" in params:
+            new_config = params["embedding_model"]
+            # update embedding
+            embed_params = kb.params.get("embedding_model", {})
+            if not embed_params:
+                # maybe old version of KB
+                raise ValueError("No embedding config to update")
+
+            # some parameters are not allowed to update
+            for key in ("provider", "model_name"):
+                if key in new_config and new_config[key] != embed_params.get(key):
+                    raise ValueError(f"You can't update '{key}' setting")
+
+            embed_params.update(new_config)
+
+            self._check_embedding_model(
+                project.name,
+                params=embed_params,
+                kb_name=name,
+            )
+            kb.params["embedding_model"] = embed_params
+
+        if "reranking_model" in params:
+            new_config = params["reranking_model"]
+            # update embedding
+            rerank_params = kb.params.get("reranking_model", {})
+
+            if new_config is False:
+                # disable reranking
+                rerank_params = {}
+            elif "provider" in new_config and new_config["provider"] != rerank_params.get("provider"):
+                # use new config (and include default config)
+                rerank_params = get_model_params(new_config, "default_reranking_model")
+            else:
+                # update current config
+                rerank_params.update(new_config)
+
+            if rerank_params:
+                self._test_reranking(rerank_params)
+
+            kb.params["reranking_model"] = rerank_params
+
+        # update other keys
+        for key in ["id_column", "metadata_columns", "content_columns", "preprocessing"]:
+            if key in params:
+                kb.params[key] = params[key]
+
+        flag_modified(kb, "params")
+        db.session.commit()
+
+        return self.get(name.lower(), project_id)
+
+    def _test_reranking(self, params):
+        try:
+            reranker = get_reranking_model_from_params(params)
+            reranker.get_scores("test", ["test"])
+        except (ValueError, RuntimeError) as e:
+            if params["provider"] in ("azure_openai", "openai") and params.get("method") != "no-logprobs":
+                # check with no-logprobs
+                params["method"] = "no-logprobs"
+                self._test_reranking(params)
+                logger.warning(
+                    f"logprobs is not supported for this model: {params.get('model_name')}. using no-logprobs mode"
+                )
+            else:
+                raise RuntimeError(f"Problem with reranker config: {e}") from e
+
     def _create_persistent_pgvector(self, params=None):
         """Create default vector database for knowledge base, if not specified"""
         vector_store_name = "kb_pgvector_store"
@@ -1244,11 +1384,11 @@ class KnowledgeBaseController:
         self.session.integration_controller.add(vector_store_name, engine, connection_args)
         return vector_store_name
 
-    def _create_embedding_model(self, project_name, engine="openai", params: dict = None, kb_name=""):
-        """create a default embedding model for knowledge base, if not specified"""
-        model_name = f"kb_embedding_{kb_name}"
+    def _check_embedding_model(self, project_name, params: dict = None, kb_name=""):
+        """check embedding model for knowledge base"""
 
-        # drop if exists - parameters can be different
+        # if mindsdb model from old KB exists - drop it
+        model_name = f"kb_embedding_{kb_name}"
         try:
             model = self.session.model_controller.get_model(model_name, project_name=project_name)
             if model is not None:
@@ -1260,63 +1400,18 @@ class KnowledgeBaseController:
             raise ValueError("'provider' parameter is required for embedding model")
 
         # check available providers
-        avail_providers = ("openai", "azure_openai", "bedrock", "gemini", "google")
+        avail_providers = ("openai", "azure_openai", "bedrock", "gemini", "google", "ollama")
         if params["provider"] not in avail_providers:
             raise ValueError(
                 f"Wrong embedding provider: {params['provider']}. Available providers: {', '.join(avail_providers)}"
             )
 
-        if params["provider"] not in ("openai", "azure_openai"):
-            # try use litellm
-            try:
-                KnowledgeBaseTable.call_litellm_embedding(self.session, params, ["test"])
-            except Exception as e:
-                raise RuntimeError(f"Problem with embedding model config: {e}") from e
-            return
-
-        params = copy.deepcopy(params)
-        if "provider" in params:
-            engine = params.pop("provider").lower()
-
-        api_key = get_api_key(engine, params, strict=False)
-        if api_key is None:
-            if "api_key" in params:
-                params.pop("api_key")
-            else:
-                raise ValueError("'api_key' parameter is required for embedding model")
-
-        if engine == "azure_openai":
-            engine = "openai"
-            params["provider"] = "azure"
-
-        if engine == "openai":
-            if "question_column" not in params:
-                params["question_column"] = "content"
-            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")
-
-        params["engine"] = engine
-        params["join_learn_process"] = True
-        params["mode"] = "embedding"
-
-        # Include API key if provided.
-        statement = CreatePredictor(
-            name=Identifier(parts=[project_name, model_name]),
-            using=params,
-            targets=[Identifier(parts=[TableField.EMBEDDINGS.value])],
-        )
+        llm_client = LLMClient(params, session=self.session)
 
-        command_executor = ExecuteCommands(self.session)
-        resp = command_executor.answer_create_predictor(statement, project_name)
-        # check model status
-        record = resp.data.records[0]
-        if record["STATUS"] == "error":
-            raise ValueError("Embedding model error:" + record["ERROR"])
-        return model_name
+        try:
+            llm_client.embeddings(["test"])
+        except Exception as e:
+            raise RuntimeError(f"Problem with embedding model config: {e}") from e
 
     def delete(self, name: str, project_name: int, if_exists: bool = False) -> None:
         """
@@ -1422,12 +1517,6 @@ class KnowledgeBaseController:
         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
-        """
-        raise NotImplementedError()
-
     def evaluate(self, table_name: str, project_name: str, params: dict = None) -> pd.DataFrame:
         """
         Run evaluate and/or create test data for evaluation

mindsdb/interfaces/knowledge_base/evaluate.py

@@ -2,6 +2,7 @@ import json
 import math
 import re
 import time
+import copy
 from typing import List
 
 import pandas as pd
@@ -10,6 +11,7 @@ import datetime as dt
 from mindsdb.api.executor.sql_query.result_set import ResultSet
 from mindsdb_sql_parser import Identifier, Select, Constant, Star, parse_sql, BinaryOperation
 from mindsdb.utilities import log
+from mindsdb.utilities.config import config
 
 from mindsdb.interfaces.knowledge_base.llm_client import LLMClient
 
@@ -105,7 +107,12 @@ class EvaluateBase:
         if llm_params is None:
             llm_params = self.kb._kb.params.get("reranking_model")
 
-        self.llm_client = LLMClient(llm_params)
+        params = copy.deepcopy(config.get("default_llm", {}))
+
+        if llm_params:
+            params.update(llm_params)
+
+        self.llm_client = LLMClient(params)
 
     def generate_test_data(self, gen_params: dict) -> pd.DataFrame:
         # Extract source data (from users query or from KB itself) and call `generate` to get test data
@@ -241,6 +248,26 @@ class EvaluateBase:
 
         return cls(session, kb_table).run_evaluate(params)
 
+    def generate_question_answer(self, text: str) -> (str, str):
+        messages = [
+            {"role": "system", "content": GENERATE_QA_SYSTEM_PROMPT},
+            {"role": "user", "content": f"\n\nText:\n{text}\n\n"},
+        ]
+        answer = self.llm_client.completion(messages, json_output=True)[0]
+
+        # Sanitize the response by removing markdown code block formatting like ```json
+        sanitized_answer = sanitize_json_response(answer)
+
+        try:
+            output = json.loads(sanitized_answer)
+        except json.JSONDecodeError:
+            raise ValueError(f"Could not parse response from LLM: {answer}")
+
+        if "query" not in output or "reference_answer" not in output:
+            raise ValueError("Cant find question/answer in LLM response")
+
+        return output.get("query"), output.get("reference_answer")
+
 
 class EvaluateRerank(EvaluateBase):
     """
@@ -268,28 +295,12 @@ class EvaluateRerank(EvaluateBase):
         df["id"] = df.index
         return df
 
-    def generate_question_answer(self, text: str) -> (str, str):
-        messages = [
-            {"role": "system", "content": GENERATE_QA_SYSTEM_PROMPT},
-            {"role": "user", "content": f"\n\nText:\n{text}\n\n"},
-        ]
-        answer = self.llm_client.completion(messages, json_output=True)
-
-        # Sanitize the response by removing markdown code block formatting like ```json
-        sanitized_answer = sanitize_json_response(answer)
-
-        try:
-            output = json.loads(sanitized_answer)
-        except json.JSONDecodeError:
-            raise ValueError(f"Could not parse response from LLM: {answer}")
-
-        if "query" not in output or "reference_answer" not in output:
-            raise ValueError("Cant find question/answer in LLM response")
-
-        return output.get("query"), output.get("reference_answer")
-
     def evaluate(self, test_data: pd.DataFrame) -> pd.DataFrame:
         json_to_log_list = []
+        if {"question", "answer"} - set(test_data.columns):
+            raise KeyError(
+                f'Test data must contain "question" and "answer" columns. Columns in the provided test data: {list(test_data.columns)}'
+            )
         questions = test_data.to_dict("records")
 
         for i, item in enumerate(questions):
@@ -483,28 +494,12 @@ class EvaluateDocID(EvaluateBase):
         df = pd.DataFrame(qa_data)
         return df
 
-    def generate_question_answer(self, text: str) -> (str, str):
-        messages = [
-            {"role": "system", "content": GENERATE_QA_SYSTEM_PROMPT},
-            {"role": "user", "content": f"\n\nText:\n{text}\n\n"},
-        ]
-        answer = self.llm_client.completion(messages, json_output=True)
-
-        # Sanitize the response by removing markdown code block formatting like ```json
-        sanitized_answer = sanitize_json_response(answer)
-
-        try:
-            output = json.loads(sanitized_answer)
-        except json.JSONDecodeError:
-            raise ValueError(f"Could not parse response from LLM: {answer}")
-
-        if "query" not in output or "reference_answer" not in output:
-            raise ValueError("Cant find question/answer in LLM response")
-
-        return output.get("query"), output.get("reference_answer")
-
     def evaluate(self, test_data: pd.DataFrame) -> pd.DataFrame:
         stats = []
+        if {"question", "doc_id"} - set(test_data.columns):
+            raise KeyError(
+                f'Test data must contain "question" and "doc_id" columns. Columns in the provided test data: {list(test_data.columns)}'
+            )
         questions = test_data.to_dict("records")
 
         for i, item in enumerate(questions):

mindsdb/interfaces/knowledge_base/executor.py

@@ -43,7 +43,18 @@ class KnowledgeBaseQueryExecutor:
         if isinstance(node, BinaryOperation):
             if isinstance(node.args[0], Identifier):
                 parts = node.args[0].parts
+
+                # map chunk_content to content
+                if parts[0].lower() == "chunk_content":
+                    parts[0] = self.content_column
+
                 if len(parts) == 1 and parts[0].lower() == self.content_column:
+                    if "LIKE" in node.op.upper():
+                        # remove '%'
+                        arg = node.args[1]
+                        if isinstance(arg, Constant) and isinstance(arg.value, str):
+                            arg.value = arg.value.strip(" %")
+
                     return True
         return False