MindsDB 25.5.3.0__py3-none-any.whl → 25.5.4.1__py3-none-any.whl

mindsdb/interfaces/knowledge_base/controller.py

@@ -5,15 +5,7 @@ from typing import Dict, List, Optional
 import pandas as pd
 import numpy as np
 
-from mindsdb_sql_parser.ast import (
-    BinaryOperation,
-    Constant,
-    Identifier,
-    Select,
-    Update,
-    Delete,
-    Star
-)
+from mindsdb_sql_parser.ast import BinaryOperation, Constant, Identifier, Select, Update, Delete, Star
 from mindsdb_sql_parser.ast.mindsdb import CreatePredictor
 
 from mindsdb.integrations.utilities.query_traversal import query_traversal
@@ -27,11 +19,14 @@ 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
+from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embedding_handler import (
+    construct_model_from_args,
+)
 
 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
@@ -47,11 +42,7 @@ from mindsdb.integrations.utilities.rag.rerankers.base_reranker import BaseLLMRe
 
 logger = log.getLogger(__name__)
 
-KB_TO_VECTORDB_COLUMNS = {
-    'id': 'original_doc_id',
-    'chunk_id': 'id',
-    'chunk_content': 'content'
-}
+KB_TO_VECTORDB_COLUMNS = {"id": "original_doc_id", "chunk_id": "id", "chunk_content": "content"}
 
 
 def get_model_params(model_params: dict, default_config_key: str):
@@ -71,23 +62,23 @@ 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')
+    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':
+    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
-        params_copy['azure_endpoint'] = params_copy.pop('base_url')
-        if 'chunk_size' not in params_copy:
-            params_copy['chunk_size'] = 2048
-        if 'api_version' in params_copy:
-            params_copy['openai_api_version'] = params_copy['api_version']
+        params_copy["openai_api_key"] = api_key
+        params_copy["azure_endpoint"] = params_copy.pop("base_url")
+        if "chunk_size" not in params_copy:
+            params_copy["chunk_size"] = 2048
+        if "api_version" in params_copy:
+            params_copy["openai_api_version"] = params_copy["api_version"]
     else:
         params_copy[f"{provider}_api_key"] = api_key
-    params_copy.pop('api_key', None)
-    params_copy['model'] = params_copy.pop('model_name', None)
+    params_copy.pop("api_key", None)
+    params_copy["model"] = params_copy.pop("model_name", None)
 
     return construct_model_from_args(params_copy)
 
@@ -97,15 +88,26 @@ 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.get('provider', "openai").lower()
+    provider = params_copy.get("provider", "openai").lower()
 
     if "api_key" not in params_copy:
         params_copy["api_key"] = get_api_key(provider, params_copy, strict=False)
-    params_copy['model'] = params_copy.pop('model_name', None)
+    params_copy["model"] = params_copy.pop("model_name", None)
 
     return BaseLLMReranker(**params_copy)
 
 
+def safe_pandas_is_datetime(value: str) -> bool:
+    """
+    Check if the value can be parsed as a datetime.
+    """
+    try:
+        result = pd.api.types.is_datetime64_any_dtype(value)
+        return result
+    except ValueError:
+        return False
+
+
 class KnowledgeBaseTable:
     """
     Knowledge base table interface
@@ -125,6 +127,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)}")
@@ -186,11 +193,13 @@ class KnowledgeBaseTable:
                     query_text = item.value
 
                     # replace content with embeddings
-                    conditions.append(FilterCondition(
-                        column=TableField.EMBEDDINGS.value,
-                        value=self._content_to_embeddings(item.value),
-                        op=FilterOperator.EQUAL,
-                    ))
+                    conditions.append(
+                        FilterCondition(
+                            column=TableField.EMBEDDINGS.value,
+                            value=self._content_to_embeddings(item.value),
+                            op=FilterOperator.EQUAL,
+                        )
+                    )
                 else:
                     conditions.append(item)
 
@@ -232,7 +241,7 @@ class KnowledgeBaseTable:
     def add_relevance(self, df, query_text, relevance_threshold=None):
         relevance_column = TableField.RELEVANCE.value
 
-        reranking_model_params = get_model_params(self._kb.params.get("reranking_model"), "default_llm")
+        reranking_model_params = get_model_params(self._kb.params.get("reranking_model"), "default_reranking_model")
         if reranking_model_params and query_text and len(df) > 0:
             # Use reranker for relevance score
             try:
@@ -244,7 +253,7 @@ class KnowledgeBaseTable:
 
                 reranker = get_reranking_model_from_params(reranking_model_params)
                 # Get documents to rerank
-                documents = df['chunk_content'].tolist()
+                documents = df["chunk_content"].tolist()
                 # Use the get_scores method with disable_events=True
                 scores = reranker.get_scores(query_text, documents)
                 # Add scores as the relevance column
@@ -257,21 +266,21 @@ class KnowledgeBaseTable:
             except Exception as e:
                 logger.error(f"Error during reranking: {str(e)}")
                 # Fallback to distance-based relevance
-                if 'distance' in df.columns:
-                    df[relevance_column] = 1 / (1 + df['distance'])
+                if "distance" in df.columns:
+                    df[relevance_column] = 1 / (1 + df["distance"])
                 else:
                     logger.info("No distance or reranker available")
 
-        elif 'distance' in df.columns:
+        elif "distance" in df.columns:
             # Calculate relevance from distance
             logger.info("Calculating relevance from vector distance")
-            df[relevance_column] = 1 / (1 + df['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
-            df['distance'] = None
+            df["distance"] = None
         # Sort by relevance
         df = df.sort_values(by=relevance_column, ascending=False)
         return df
@@ -294,7 +303,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_doc_id')
+            lambda m: None if m is None else m.get("original_doc_id")
         )
 
         # id on first place
@@ -309,23 +318,14 @@ class KnowledgeBaseTable:
         if documents:
             self.insert_documents(documents)
 
-    def insert_web_pages(
-            self,
-            urls: List[str],
-            crawl_depth: int,
-            limit: int,
-            filters: List[str] = None
-    ):
+    def insert_web_pages(self, urls: List[str], crawl_depth: int, limit: int, filters: List[str] = None):
         """Process and insert web pages"""
         if not self.document_loader:
             raise ValueError("Document loader not configured")
 
-        documents = list(self.document_loader.load_web_pages(
-            urls,
-            limit=limit,
-            crawl_depth=crawl_depth,
-            filters=filters
-        ))
+        documents = list(
+            self.document_loader.load_web_pages(urls, limit=limit, crawl_depth=crawl_depth, filters=filters)
+        )
         if documents:
             self.insert_documents(documents)
 
@@ -343,11 +343,9 @@ class KnowledgeBaseTable:
         if not rows:
             return
 
-        documents = [Document(
-            content=row.get('content', ''),
-            id=row.get('id'),
-            metadata=row.get('metadata', {})
-        ) for row in rows]
+        documents = [
+            Document(content=row.get("content", ""), id=row.get("id"), metadata=row.get("metadata", {})) for row in rows
+        ]
 
         self.insert_documents(documents)
 
@@ -368,7 +366,7 @@ class KnowledgeBaseTable:
         conditions = db_handler.extract_conditions(query.where)
         doc_id = None
         for condition in conditions:
-            if condition.column == 'chunk_id' and condition.op == FilterOperator.EQUAL:
+            if condition.column == "chunk_id" and condition.op == FilterOperator.EQUAL:
                 doc_id = condition.value
 
         if cont_col in query.update_columns:
@@ -379,7 +377,7 @@ class KnowledgeBaseTable:
                 doc = Document(
                     id=doc_id,
                     content=content.value,
-                    metadata={}  # Empty metadata for content-only updates
+                    metadata={},  # Empty metadata for content-only updates
                 )
                 processed_chunks = self.document_preprocessor.process_documents([doc])
                 if processed_chunks:
@@ -418,7 +416,7 @@ class KnowledgeBaseTable:
         query: str,
         keywords: List[str] = None,
         metadata: Dict[str, str] = None,
-        distance_function=DistanceFunction.COSINE_DISTANCE
+        distance_function=DistanceFunction.COSINE_DISTANCE,
     ) -> pd.DataFrame:
         query_df = pd.DataFrame.from_records([{TableField.CONTENT.value: query}])
         embeddings_df = self._df_to_embeddings(query_df)
@@ -427,14 +425,14 @@ class KnowledgeBaseTable:
         embeddings = embeddings_df.iloc[0][TableField.EMBEDDINGS.value]
         keywords_query = None
         if keywords is not None:
-            keywords_query = ' '.join(keywords)
+            keywords_query = " ".join(keywords)
         db_handler = self.get_vector_db()
         return db_handler.hybrid_search(
             self._kb.vector_database_table,
             embeddings,
             query=keywords_query,
             metadata=metadata,
-            distance_function=distance_function
+            distance_function=distance_function,
         )
 
     def clear(self):
@@ -467,7 +465,7 @@ class KnowledgeBaseTable:
 
         # 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])
+        content_columns = self._kb.params.get("content_columns", [TableField.CONTENT.value])
 
         # Convert DataFrame rows to documents, creating separate documents for each content column
         raw_documents = []
@@ -485,15 +483,11 @@ class KnowledgeBaseTable:
 
                     metadata = {
                         **base_metadata,
-                        'original_row_index': str(idx),  # provide link to original row index
-                        'content_column': col,
+                        "original_row_index": str(idx),  # provide link to original row index
+                        "content_column": col,
                     }
 
-                    raw_documents.append(Document(
-                        content=content_str,
-                        id=doc_id,
-                        metadata=metadata
-                    ))
+                    raw_documents.append(Document(content=content_str, id=doc_id, metadata=metadata))
 
         # Apply preprocessing to all documents if preprocessor exists
         if self.document_preprocessor:
@@ -502,11 +496,16 @@ class KnowledgeBaseTable:
             processed_chunks = raw_documents  # Use raw documents if no preprocessing
 
         # Convert processed chunks back to DataFrame with standard structure
-        df = pd.DataFrame([{
-            TableField.CONTENT.value: chunk.content,
-            TableField.ID.value: chunk.id,
-            TableField.METADATA.value: chunk.metadata
-        } for chunk in processed_chunks])
+        df = pd.DataFrame(
+            [
+                {
+                    TableField.CONTENT.value: chunk.content,
+                    TableField.ID.value: chunk.id,
+                    TableField.METADATA.value: chunk.metadata,
+                }
+                for chunk in processed_chunks
+            ]
+        )
 
         if df.empty:
             logger.warning("No valid content found in any content columns")
@@ -517,17 +516,17 @@ class KnowledgeBaseTable:
         df = pd.concat([df, df_emb], axis=1)
         db_handler = self.get_vector_db()
 
-        if params is not None and params.get('kb_no_upsert', False):
+        if params is not None and params.get("kb_no_upsert", False):
             # speed up inserting by disable checking existing records
             db_handler.insert(self._kb.vector_database_table, df)
         else:
             db_handler.do_upsert(self._kb.vector_database_table, df)
 
     def _adapt_column_names(self, df: pd.DataFrame) -> pd.DataFrame:
-        '''
+        """
         Convert input columns for vector db input
         - id, content and metadata
-        '''
+        """
         # Debug incoming data
         logger.debug(f"Input DataFrame columns: {df.columns}")
         logger.debug(f"Input DataFrame first row: {df.iloc[0].to_dict()}")
@@ -536,7 +535,7 @@ class KnowledgeBaseTable:
         columns = list(df.columns)
 
         # -- prepare id --
-        id_column = params.get('id_column')
+        id_column = params.get("id_column")
         if id_column is not None and id_column not in columns:
             id_column = None
 
@@ -546,8 +545,8 @@ class KnowledgeBaseTable:
         # Also check for case-insensitive 'id' column
         if id_column is None:
             column_map = {col.lower(): col for col in columns}
-            if 'id' in column_map:
-                id_column = column_map['id']
+            if "id" in column_map:
+                id_column = column_map["id"]
 
         if id_column is not None:
             columns.remove(id_column)
@@ -562,8 +561,8 @@ class KnowledgeBaseTable:
             logger.debug(f"Added IDs: {df_out[TableField.ID.value].tolist()}")
 
         # -- prepare content and metadata --
-        content_columns = params.get('content_columns', [TableField.CONTENT.value])
-        metadata_columns = params.get('metadata_columns')
+        content_columns = params.get("content_columns", [TableField.CONTENT.value])
+        metadata_columns = params.get("metadata_columns")
 
         logger.debug(f"Processing with: content_columns={content_columns}, metadata_columns={metadata_columns}")
 
@@ -571,25 +570,19 @@ class KnowledgeBaseTable:
         if content_columns:
             # Ensure content columns are case-insensitive
             column_map = {col.lower(): col for col in columns}
-            content_columns = [
-                column_map.get(col.lower(), col)
-                for col in content_columns
-            ]
+            content_columns = [column_map.get(col.lower(), col) for col in content_columns]
             logger.debug(f"Mapped content columns: {content_columns}")
 
         if metadata_columns:
             # Ensure metadata columns are case-insensitive
             column_map = {col.lower(): col for col in columns}
-            metadata_columns = [
-                column_map.get(col.lower(), col)
-                for col in metadata_columns
-            ]
+            metadata_columns = [column_map.get(col.lower(), col) for col in metadata_columns]
             logger.debug(f"Mapped metadata columns: {metadata_columns}")
 
         if content_columns is not None:
             content_columns = list(set(content_columns).intersection(columns))
             if len(content_columns) == 0:
-                raise ValueError(f'Content columns {params.get("content_columns")} not found in dataset: {columns}')
+                raise ValueError(f"Content columns {params.get('content_columns')} not found in dataset: {columns}")
 
             if metadata_columns is not None:
                 metadata_columns = list(set(metadata_columns).intersection(columns))
@@ -603,12 +596,13 @@ class KnowledgeBaseTable:
 
         # Add metadata
         if metadata_columns and len(metadata_columns) > 0:
+
             def convert_row_to_metadata(row):
                 metadata = {}
                 for col in metadata_columns:
                     value = row[col]
                     # Convert numpy/pandas types to Python native types
-                    if pd.api.types.is_datetime64_any_dtype(value) or isinstance(value, pd.Timestamp):
+                    if safe_pandas_is_datetime(value) or isinstance(value, pd.Timestamp):
                         value = str(value)
                     elif pd.api.types.is_integer_dtype(value):
                         value = int(value)
@@ -648,7 +642,7 @@ class KnowledgeBaseTable:
         if self._vector_db is None:
             database = db.Integration.query.get(self._kb.vector_database_id)
             if database is None:
-                raise ValueError('Vector database not found. Is it deleted?')
+                raise ValueError("Vector database not found. Is it deleted?")
             database_name = database.name
             self._vector_db = self.session.integration_controller.get_data_handler(database_name)
         return self._vector_db
@@ -673,6 +667,15 @@ 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 = copy.deepcopy(config.get("default_embedding_model", {}))
+            embedding_params.update(self._kb.params["embedding_model"])
+            results = self.call_litellm_embedding(self.session, embedding_params, messages)
+            results = [[val] for val in results]
+            return pd.DataFrame(results, columns=[TableField.EMBEDDINGS.value])
+
         # get the input columns
         model_rec = db.session.query(db.Predictor).filter_by(id=model_id).first()
 
@@ -681,19 +684,15 @@ class KnowledgeBaseTable:
 
         project_datanode = self.session.datahub.get(model_project.name)
 
-        model_using = model_rec.learn_args.get('using', {})
-        input_col = model_using.get('question_column')
+        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')
+            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})
 
-        df_out = project_datanode.predict(
-            model_name=model_rec.name,
-            df=df,
-            params=self.model_params
-        )
+        df_out = project_datanode.predict(model_name=model_rec.name, df=df, params=self.model_params)
 
         target = model_rec.to_predict[0]
         if target != TableField.EMBEDDINGS.value:
@@ -714,6 +713,23 @@ class KnowledgeBaseTable:
         res = self._df_to_embeddings(df)
         return res[TableField.EMBEDDINGS.value][0]
 
+    @staticmethod
+    def call_litellm_embedding(session, model_params, messages):
+        args = copy.deepcopy(model_params)
+
+        llm_model = args.pop("model_name")
+        engine = args.pop("provider")
+
+        llm_model = f"{engine}/{llm_model}"
+
+        if "base_url" in args:
+            args["api_base"] = args.pop("base_url")
+
+        module = session.integration_controller.get_handler_module("litellm")
+        if module is None or module.Handler is None:
+            raise ValueError(f'Unable to use "{engine}" provider. Litellm handler is not installed')
+        return module.Handler.embeddings(llm_model, messages, args)
+
     def build_rag_pipeline(self, retrieval_config: dict):
         """
         Builds a RAG pipeline with returned sources
@@ -729,10 +745,10 @@ class KnowledgeBaseTable:
         """
         # Get embedding model from knowledge base
         embeddings_model = None
-        embedding_model_params = get_model_params(self._kb.params.get('embedding_model', {}), 'default_embedding_model')
+        embedding_model_params = get_model_params(self._kb.params.get("embedding_model", {}), "default_embedding_model")
         if self._kb.embedding_model:
             # Extract embedding model args from knowledge base table
-            embedding_args = self._kb.embedding_model.learn_args.get('using', {})
+            embedding_args = self._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}")
@@ -744,21 +760,17 @@ class KnowledgeBaseTable:
             logger.debug("Using default embedding model as knowledge base has no embedding model")
 
         # Update retrieval config with knowledge base parameters
-        kb_params = {
-            'vector_store_config': {
-                'kb_table': self
-            }
-        }
+        kb_params = {"vector_store_config": {"kb_table": self}}
 
         # Load and validate config
         try:
             rag_config = load_rag_config(retrieval_config, kb_params, embeddings_model)
 
             # Build LLM if specified
-            if 'llm_model_name' in rag_config:
+            if "llm_model_name" in rag_config:
                 llm_args = {"model_name": rag_config.llm_model_name}
                 if not rag_config.llm_provider:
-                    llm_args['provider'] = get_llm_provider(llm_args)
+                    llm_args["provider"] = get_llm_provider(llm_args)
                 else:
                     llm_args["provider"] = rag_config.llm_provider
                 rag_config.llm = create_chat_model(llm_args)
@@ -779,6 +791,7 @@ class KnowledgeBaseTable:
         if isinstance(base_metadata, str):
             try:
                 import ast
+
                 return ast.literal_eval(base_metadata)
             except (SyntaxError, ValueError):
                 logger.warning(f"Could not parse metadata: {base_metadata}. Using empty dict.")
@@ -788,6 +801,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, provided_id=provided_id)
 
     def _convert_metadata_value(self, value):
@@ -820,6 +834,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:
     """
@@ -831,14 +854,14 @@ class KnowledgeBaseController:
         self.session = session
 
     def add(
-            self,
-            name: str,
-            project_name: str,
-            embedding_model: Identifier,
-            storage: Identifier,
-            params: dict,
-            preprocessing_config: Optional[dict] = None,
-            if_not_exists: bool = False
+        self,
+        name: str,
+        project_name: str,
+        storage: Identifier,
+        params: dict,
+        preprocessing_config: Optional[dict] = None,
+        if_not_exists: bool = False,
+        # embedding_model: Identifier = None, # Legacy: Allow MindsDB models to be passed as embedding_model.
     ) -> db.KnowledgeBase:
         """
         Add a new knowledge base to the database
@@ -846,15 +869,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
         """
+        # 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
+            params["preprocessing"] = preprocessing_config
 
         # Check if vector_size is provided when using sparse vectors
-        is_sparse = params.get('is_sparse')
-        vector_size = params.get('vector_size')
+        is_sparse = params.get("is_sparse")
+        vector_size = params.get("vector_size")
         if is_sparse and vector_size is None:
             raise ValueError("vector_size is required when is_sparse=True")
 
@@ -871,41 +897,45 @@ class KnowledgeBaseController:
                 return kb
             raise EntityExistsError("Knowledge base already exists", name)
 
-        embedding_params = copy.deepcopy(config.get('default_embedding_model', {}))
-
-        model_name = None
-        model_project = project
-        if embedding_model:
-            model_name = embedding_model.parts[-1]
-            if len(embedding_model.parts) > 1:
-                model_project = self.session.database_controller.get_project(embedding_model.parts[-2])
-
-        elif 'embedding_model' in params:
-            if isinstance(params['embedding_model'], str):
-                # it is model name
-                model_name = params['embedding_model']
-            else:
-                # it is params for model
-                embedding_params.update(params['embedding_model'])
-
-        if model_name is None:
-            model_name = self._create_embedding_model(
-                project.name,
-                params=embedding_params,
-                kb_name=name,
-            )
-            params['created_embedding_model'] = model_name
+        embedding_params = copy.deepcopy(config.get("default_embedding_model", {}))
+
+        # Legacy
+        # model_name = None
+        # model_project = project
+        # if embedding_model:
+        #     model_name = embedding_model.parts[-1]
+        #     if len(embedding_model.parts) > 1:
+        #         model_project = self.session.database_controller.get_project(embedding_model.parts[-2])
+
+        # elif "embedding_model" in params:
+        #     if isinstance(params["embedding_model"], str):
+        #         # it is model name
+        #         model_name = params["embedding_model"]
+        #     else:
+        #         # it is params for model
+        #         embedding_params.update(params["embedding_model"])
+
+        if "embedding_model" in params:
+            if not isinstance(params["embedding_model"], dict):
+                raise ValueError("embedding_model should be JSON object with model parameters.")
+            embedding_params.update(params["embedding_model"])
+
+        # if model_name is None:  # Legacy
+        model_name = self._create_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=model_project.name
-            )
-            model_record = db.Predictor.query.get(model['id'])
+            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
 
-        reranking_model_params = get_model_params(params.get('reranking_model', {}), 'default_llm')
+        reranking_model_params = get_model_params(params.get("reranking_model", {}), "default_reranking_model")
         if reranking_model_params:
             # Get reranking model from params.
             # This is called here to check validaity of the parameters.
@@ -913,17 +943,17 @@ class KnowledgeBaseController:
 
         # search for the vector database table
         if storage is None:
-            cloud_pg_vector = os.environ.get('KB_PGVECTOR_URL')
+            cloud_pg_vector = os.environ.get("KB_PGVECTOR_URL")
             if cloud_pg_vector:
                 vector_table_name = name
                 # Add sparse vector support for pgvector
                 vector_db_params = {}
                 # Check both explicit parameter and model configuration
-                is_sparse = is_sparse or model_record.learn_args.get('using', {}).get('sparse')
+                is_sparse = is_sparse or model_record.learn_args.get("using", {}).get("sparse")
                 if is_sparse:
-                    vector_db_params['is_sparse'] = True
+                    vector_db_params["is_sparse"] = True
                     if vector_size is not None:
-                        vector_db_params['vector_size'] = vector_size
+                        vector_db_params["vector_size"] = vector_size
                 vector_db_name = self._create_persistent_pgvector(vector_db_params)
 
             else:
@@ -931,26 +961,22 @@ class KnowledgeBaseController:
                 vector_table_name = "default_collection"
                 vector_db_name = self._create_persistent_chroma(name)
                 # memorize to remove it later
-                params['default_vector_storage'] = vector_db_name
+                params["default_vector_storage"] = vector_db_name
         elif len(storage.parts) != 2:
-            raise ValueError('Storage param has to be vector db with table')
+            raise ValueError("Storage param has to be vector db with table")
         else:
             vector_db_name, vector_table_name = storage.parts
 
         # create table in vectordb before creating KB
-        self.session.datahub.get(vector_db_name).integration_handler.create_table(
-            vector_table_name
-        )
-        vector_database_id = self.session.integration_controller.get(vector_db_name)['id']
+        self.session.datahub.get(vector_db_name).integration_handler.create_table(vector_table_name)
+        vector_database_id = self.session.integration_controller.get(vector_db_name)["id"]
 
         # Store sparse vector settings in params if specified
         if is_sparse:
             params = params or {}
-            params['vector_config'] = {
-                'is_sparse': is_sparse
-            }
+            params["vector_config"] = {"is_sparse": is_sparse}
             if vector_size is not None:
-                params['vector_config']['vector_size'] = vector_size
+                params["vector_config"]["vector_size"] = vector_size
 
         kb = db.KnowledgeBase(
             name=name,
@@ -972,7 +998,7 @@ class KnowledgeBaseController:
         if self.session.integration_controller.get(vector_store_name):
             return vector_store_name
 
-        self.session.integration_controller.add(vector_store_name, 'pgvector', params or {})
+        self.session.integration_controller.add(vector_store_name, "pgvector", params or {})
         return vector_store_name
 
     def _create_persistent_chroma(self, kb_name, engine="chromadb"):
@@ -990,7 +1016,7 @@ 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=''):
+    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}"
 
@@ -1002,42 +1028,47 @@ class KnowledgeBaseController:
         except PredictorRecordNotFound:
             pass
 
-        if 'provider' in params:
-            engine = params.pop('provider').lower()
+        if params.get("provider", None) not in ("openai", "azure"):
+            # try use litellm
+            KnowledgeBaseTable.call_litellm_embedding(self.session, params, ["test"])
+            return
 
-        api_key = get_api_key(engine, params, strict=False) or params.pop('api_key')
+        if "provider" in params:
+            engine = params.pop("provider").lower()
 
-        if engine == 'azure_openai':
-            engine = 'openai'
-            params['provider'] = 'azure'
+        api_key = get_api_key(engine, params, strict=False) or params.pop("api_key")
 
-        if engine == 'openai':
-            if 'question_column' not in params:
-                params['question_column'] = 'content'
+        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 'base_url' in params:
-                params['api_base'] = params.pop('base_url')
+                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'
+        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])
-            ]
+            targets=[Identifier(parts=[TableField.EMBEDDINGS.value])],
         )
 
         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'])
+        if record["STATUS"] == "error":
+            raise ValueError("Embedding model error:" + record["ERROR"])
         return model_name
 
     def delete(self, name: str, project_name: int, if_exists: bool = False) -> None:
@@ -1064,16 +1095,16 @@ class KnowledgeBaseController:
         db.session.commit()
 
         # drop objects if they were created automatically
-        if 'default_vector_storage' in kb.params:
+        if "default_vector_storage" in kb.params:
             try:
-                handler = self.session.datahub.get(kb.params['default_vector_storage']).integration_handler
+                handler = self.session.datahub.get(kb.params["default_vector_storage"]).integration_handler
                 handler.drop_table(kb.vector_database_table)
-                self.session.integration_controller.delete(kb.params['default_vector_storage'])
+                self.session.integration_controller.delete(kb.params["default_vector_storage"])
             except EntityNotExistsError:
                 pass
-        if 'created_embedding_model' in kb.params:
+        if "created_embedding_model" in kb.params:
             try:
-                self.session.model_controller.delete_model(kb.params['created_embedding_model'], project_name)
+                self.session.model_controller.delete_model(kb.params["created_embedding_model"], project_name)
             except EntityNotExistsError:
                 pass
 
@@ -1104,11 +1135,11 @@ class KnowledgeBaseController:
         if kb is not None:
             table = KnowledgeBaseTable(kb, self.session)
             if params:
-                table.model_params = params.get('model')
+                table.model_params = params.get("model")
 
             # Always configure preprocessing - either from params or default
-            if kb.params and 'preprocessing' in kb.params:
-                table.configure_preprocessing(kb.params['preprocessing'])
+            if kb.params and "preprocessing" in kb.params:
+                table.configure_preprocessing(kb.params["preprocessing"])
             else:
                 table.configure_preprocessing(None)  # This ensures default preprocessor is created
 
@@ -1124,35 +1155,38 @@ class KnowledgeBaseController:
         if project_name is not None:
             projects = [p for p in projects if p.name == project_name]
 
-        query = (
-            db.session.query(db.KnowledgeBase)
-            .filter(db.KnowledgeBase.project_id.in_(list([p.id for p in projects])))
+        query = db.session.query(db.KnowledgeBase).filter(
+            db.KnowledgeBase.project_id.in_(list([p.id for p in projects]))
         )
 
         data = []
-        project_names = {
-            i.id: i.name
-            for i in project_controller.get_list()
-        }
+        project_names = {i.id: i.name for i in project_controller.get_list()}
 
         for record in query:
             vector_database = record.vector_database
             embedding_model = record.embedding_model
 
-            data.append({
-                'id': record.id,
-                'name': record.name,
-                'project_id': record.project_id,
-                'project_name': project_names[record.project_id],
-                '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
-            })
+            data.append(
+                {
+                    "id": record.id,
+                    "name": record.name,
+                    "project_id": record.project_id,
+                    "project_name": project_names[record.project_id],
+                    "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,
+                }
+            )
 
         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