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

mindsdb/interfaces/knowledge_base/controller.py

@@ -27,7 +27,7 @@ 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.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
@@ -37,11 +37,12 @@ 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.config import config
 from mindsdb.utilities.context import context as ctx
 
 from mindsdb.api.executor.command_executor import ExecuteCommands
 from mindsdb.utilities import log
-from mindsdb.integrations.utilities.rag.rerankers.reranker_compressor import LLMReranker
+from mindsdb.integrations.utilities.rag.rerankers.base_reranker import BaseLLMReranker
 
 logger = log.getLogger(__name__)
 
@@ -52,6 +53,18 @@ KB_TO_VECTORDB_COLUMNS = {
 }
 
 
+def get_model_params(model_params: dict, default_config_key: str):
+    """
+    Get model parameters by combining default config with user provided parameters.
+    """
+    combined_model_params = copy.deepcopy(config.get(default_config_key, {}))
+
+    if model_params:
+        combined_model_params.update(model_params)
+
+    return combined_model_params
+
+
 def get_embedding_model_from_params(embedding_model_params: dict):
     """
     Create embedding model from parameters.
@@ -89,7 +102,7 @@ def get_reranking_model_from_params(reranking_model_params: dict):
         params_copy["api_key"] = get_api_key(provider, params_copy, strict=False)
     params_copy['model'] = params_copy.pop('model_name', None)
 
-    return LLMReranker(**params_copy)
+    return BaseLLMReranker(**params_copy)
 
 
 class KnowledgeBaseTable:
@@ -215,7 +228,7 @@ class KnowledgeBaseTable:
     def add_relevance(self, df, query_text, relevance_threshold=None):
         relevance_column = TableField.RELEVANCE.value
 
-        reranking_model_params = self._kb.params.get("reranking_model")
+        reranking_model_params = get_model_params(self._kb.params.get("reranking_model"), "default_llm")
         if reranking_model_params and query_text and len(df) > 0:
             # Use reranker for relevance score
             try:
@@ -657,47 +670,34 @@ 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
-        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()
-
-            project_datanode = self.session.datahub.get(model_project.name)
 
-            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')
+        # get the input columns
+        model_rec = db.session.query(db.Predictor).filter_by(id=model_id).first()
 
-            if input_col is not None and input_col != TableField.CONTENT.value:
-                df = df.rename(columns={TableField.CONTENT.value: input_col})
+        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()
 
-            df_out = project_datanode.predict(
-                model_name=model_rec.name,
-                df=df,
-                params=self.model_params
-            )
+        project_datanode = self.session.datahub.get(model_project.name)
 
-            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})
+        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')
 
-        elif self._kb.params.get('embedding_model'):
-            embedding_model = get_embedding_model_from_params(self._kb.params.get('embedding_model'))
+        if input_col is not None and input_col != TableField.CONTENT.value:
+            df = df.rename(columns={TableField.CONTENT.value: input_col})
 
-            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})
+        df_out = project_datanode.predict(
+            model_name=model_rec.name,
+            df=df,
+            params=self.model_params
+        )
 
-        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]]
 
@@ -728,14 +728,15 @@ 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')
         if self._kb.embedding_model:
             # Extract embedding model args from knowledge base table
             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}")
-        elif self._kb.params.get('embedding_model'):
-            embeddings_model = get_embedding_model_from_params(self._kb.params['embedding_model'])
+        elif embedding_model_params:
+            embeddings_model = get_embedding_model_from_params(embedding_model_params)
             logger.debug(f"Using knowledge base embedding model from params: {self._kb.params['embedding_model']}")
         else:
             embeddings_model = DEFAULT_EMBEDDINGS_MODEL_CLASS()
@@ -869,35 +870,33 @@ class KnowledgeBaseController:
                 return kb
             raise EntityExistsError("Knowledge base already exists", name)
 
-        embedding_model_params = params.get('embedding_model', None)
-        reranking_model_params = params.get('reranking_model', None)
+        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_params:
-            # Get embedding model from params.
-            # This is called here to check validaity of the parameters.
-            get_embedding_model_from_params(
-                embedding_model_params
-            )
+        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'])
 
-        else:
-            model_name = self._get_default_embedding_model(
+        if model_name is None:
+            model_name = self._create_embedding_model(
                 project.name,
-                params=params
+                params=embedding_params,
+                kb_name=name,
             )
-            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])
-        elif not embedding_model_params:
-            model_project = project
+            params['created_embedding_model'] = model_name
 
         embedding_model_id = None
-        if model_project:
+        if model_name is not None:
             model = self.session.model_controller.get_model(
                 name=model_name,
                 project_name=model_project.name
@@ -905,6 +904,7 @@ class KnowledgeBaseController:
             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')
         if reranking_model_params:
             # Get reranking model from params.
             # This is called here to check validaity of the parameters.
@@ -989,38 +989,52 @@ class KnowledgeBaseController:
         self.session.integration_controller.add(vector_store_name, engine, connection_args)
         return vector_store_name
 
-    def _get_default_embedding_model(self, project_name, engine="langchain_embedding", params: dict = None):
+    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 = "kb_default_embedding_model"
+        model_name = f"kb_embedding_{kb_name}"
 
-        # check exists
+        # drop if exists - parameters can be different
         try:
             model = self.session.model_controller.get_model(model_name, project_name=project_name)
             if model is not None:
-                return model_name
+                self.session.model_controller.delete_model(model_name, project_name)
         except PredictorRecordNotFound:
             pass
 
-        using_args = {
-            'engine': engine
-        }
-        if engine == 'langchain_embedding':
-            # Use default embeddings.
-            using_args['class'] = 'openai'
+        if 'provider' in params:
+            engine = params.pop('provider').lower()
+
+        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' in params:
+                params[f"{engine}_api_key"] = 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.
-        using_args.update({k: v for k, v in params.items() if 'api_key' in k})
         statement = CreatePredictor(
             name=Identifier(parts=[project_name, model_name]),
-            using=using_args,
+            using=params,
             targets=[
                 Identifier(parts=[TableField.EMBEDDINGS.value])
             ]
         )
 
         command_executor = ExecuteCommands(self.session)
-        command_executor.answer_create_predictor(statement, project_name)
-
+        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
 
     def delete(self, name: str, project_name: int, if_exists: bool = False) -> None:
@@ -1054,9 +1068,9 @@ class KnowledgeBaseController:
                 self.session.integration_controller.delete(kb.params['default_vector_storage'])
             except EntityNotExistsError:
                 pass
-        if 'default_embedding_model' in kb.params:
+        if 'created_embedding_model' in kb.params:
             try:
-                self.session.model_controller.delete_model(kb.params['default_embedding_model'], project_name)
+                self.session.model_controller.delete_model(kb.params['created_embedding_model'], project_name)
             except EntityNotExistsError:
                 pass
 

mindsdb/interfaces/query_context/context_controller.py

@@ -15,6 +15,7 @@ from mindsdb.utilities.cache import get_cache
 
 from mindsdb.interfaces.storage import db
 from mindsdb.utilities.context import context as ctx
+from mindsdb.utilities.config import config
 
 from .last_query import LastQuery
 
@@ -24,9 +25,12 @@ class RunningQuery:
       Query in progres
     """
 
+    OBJECT_TYPE = 'query'
+
     def __init__(self, record: db.Queries):
         self.record = record
         self.sql = record.sql
+        self.database = record.database or config.get('default_project')
 
     def get_partition_query(self, step_num: int, query: Select) -> Select:
         """
@@ -67,6 +71,44 @@ class RunningQuery:
 
         return query
 
+    def get_info(self):
+        record = self.record
+        return {
+            'id': record.id,
+            'sql': record.sql,
+            'database': record.database,
+            'started_at': record.started_at,
+            'finished_at': record.finished_at,
+            'parameters': record.parameters,
+            'context': record.context,
+            'processed_rows': record.processed_rows,
+            'error': record.error,
+            'updated_at': record.updated_at,
+        }
+
+    def add_to_task(self):
+
+        task_record = db.Tasks(
+            company_id=ctx.company_id,
+            user_class=ctx.user_class,
+
+            object_type=self.OBJECT_TYPE,
+            object_id=self.record.id,
+        )
+        db.session.add(task_record)
+        db.session.commit()
+
+    def remove_from_task(self):
+        task = db.Tasks.query.filter(
+            db.Tasks.object_type == self.OBJECT_TYPE,
+            db.Tasks.object_id == self.record.id,
+            db.Tasks.company_id == ctx.company_id,
+        ).first()
+
+        if task is not None:
+            db.session.delete(task)
+            db.session.commit()
+
     def set_params(self, params: dict):
         """
             Store parameters of the step which is about to be split into partitions
@@ -126,14 +168,21 @@ class RunningQuery:
 
         db.session.commit()
 
-    def clear_error(self):
+    def mark_as_run(self):
         """
-            Reset error of the query in database
+            Mark query as running and reset error of the query
         """
+        if self.record.finished_at is not None:
+            raise RuntimeError('The query already finished')
 
-        if self.record.error is not None:
+        if self.record.started_at is None:
+            self.record.started_at = dt.datetime.now()
+            db.session.commit()
+        elif self.record.error is not None:
             self.record.error = None
             db.session.commit()
+        else:
+            raise RuntimeError('The query might be running already')
 
     def get_state(self) -> dict:
         """
@@ -448,7 +497,7 @@ class QueryContextController:
             raise RuntimeError(f'Query not found: {query_id}')
         return RunningQuery(rec)
 
-    def create_query(self, query: ASTNode) -> RunningQuery:
+    def create_query(self, query: ASTNode, database: str = None) -> RunningQuery:
         """
            Create a new running query from AST query
         """
@@ -463,6 +512,7 @@ class QueryContextController:
 
         rec = db.Queries(
             sql=str(query),
+            database=database,
             company_id=ctx.company_id,
         )
 
@@ -479,17 +529,7 @@ class QueryContextController:
             db.Queries.company_id == ctx.company_id
         )
         return [
-            {
-                'id': record.id,
-                'sql': record.sql,
-                'started_at': record.started_at,
-                'finished_at': record.finished_at,
-                'parameters': record.parameters,
-                'context': record.context,
-                'processed_rows': record.processed_rows,
-                'error': record.error,
-                'updated_at': record.updated_at,
-            }
+            RunningQuery(record).get_info()
             for record in query
         ]
 

mindsdb/interfaces/query_context/query_task.py

@@ -0,0 +1,19 @@
+from mindsdb.api.executor.sql_query import SQLQuery
+from mindsdb.interfaces.query_context.context_controller import query_context_controller
+from mindsdb.api.executor.controllers.session_controller import SessionController
+from mindsdb.interfaces.tasks.task import BaseTask
+
+
+class QueryTask(BaseTask):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.query_id = self.object_id
+
+    def run(self, stop_event):
+
+        try:
+            session = SessionController()
+            SQLQuery(None, query_id=self.query_id, session=session, stop_event=stop_event)
+        finally:
+            # clear task
+            query_context_controller.get_query(self.query_id).remove_from_task()

mindsdb/interfaces/storage/db.py

@@ -571,9 +571,9 @@ class Queries(Base):
     company_id: int = Column(Integer, nullable=True)
 
     sql: str = Column(String, nullable=False)
-    # step_data: JSON = Column(JSON, nullable=True)
+    database: str = Column(String, nullable=True)
 
-    started_at: datetime.datetime = Column(DateTime, default=datetime.datetime.now)
+    started_at: datetime.datetime = Column(DateTime)
     finished_at: datetime.datetime = Column(DateTime)
 
     parameters = Column(JSON, default={})

mindsdb/interfaces/tasks/task_monitor.py

@@ -2,6 +2,7 @@ import datetime as dt
 import os
 import socket
 import time
+from threading import Event
 
 import sqlalchemy as sa
 
@@ -22,7 +23,7 @@ class TaskMonitor:
     def __init__(self):
         self._active_tasks = {}
 
-    def start(self):
+    def start(self, stop_event: Event = None):
         config = Config()
         db.init()
         self.config = config
@@ -42,6 +43,9 @@ class TaskMonitor:
                 logger.error(e)
                 db.session.rollback()
 
+            if stop_event is not None and stop_event.is_set():
+                return
+
     def stop_all_tasks(self):
 
         active_tasks = list(self._active_tasks.keys())

mindsdb/interfaces/tasks/task_thread.py

@@ -6,6 +6,7 @@ from mindsdb.utilities import log
 
 from mindsdb.interfaces.triggers.trigger_task import TriggerTask
 from mindsdb.interfaces.chatbot.chatbot_task import ChatBotTask
+from mindsdb.interfaces.query_context.query_task import QueryTask
 
 logger = log.getLogger(__name__)
 
@@ -28,6 +29,7 @@ class TaskThread(threading.Thread):
         ctx.company_id = task_record.company_id
         if task_record.user_class is not None:
             ctx.user_class = task_record.user_class
+        ctx.task_id = task_record.id
 
         self.object_type = task_record.object_type
         self.object_id = task_record.object_id
@@ -43,6 +45,10 @@ class TaskThread(threading.Thread):
                 bot = ChatBotTask(self.task_id, self.object_id)
                 bot.run(self._stop_event)
 
+            elif self.object_type == 'query':
+                query = QueryTask(self.task_id, self.object_id)
+                query.run(self._stop_event)
+
         except Exception:
             logger.error(traceback.format_exc())
             task_record.last_error = str(traceback.format_exc())

mindsdb/migrations/versions/2025-04-22_53502b6d63bf_query_database.py

@@ -0,0 +1,27 @@
+"""query_database
+
+Revision ID: 53502b6d63bf
+Revises: fda503400e43
+Create Date: 2025-04-22 16:30:15.139978
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import mindsdb.interfaces.storage.db  # noqa
+
+
+# revision identifiers, used by Alembic.
+revision = '53502b6d63bf'
+down_revision = 'fda503400e43'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    with op.batch_alter_table('queries', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('database', sa.String(), nullable=True))
+
+
+def downgrade():
+    with op.batch_alter_table('queries', schema=None) as batch_op:
+        batch_op.drop_column('database')

mindsdb/utilities/config.py

@@ -230,7 +230,9 @@ class Config:
             "tasks": {
                 "disable": False
             },
-            "default_project": "mindsdb"
+            "default_project": "mindsdb",
+            "default_llm": {},
+            "default_embedding_model": {}
         }
         # endregion
 
@@ -369,6 +371,15 @@ class Config:
         if os.environ.get('MINDSDB_DEFAULT_PROJECT', '') != '':
             self._env_config['default_project'] = os.environ['MINDSDB_DEFAULT_PROJECT'].lower()
 
+        if os.environ.get('MINDSDB_DEFAULT_LLM_API_KEY', '') != '':
+            self._env_config['default_llm'] = {
+                'api_key': os.environ['MINDSDB_DEFAULT_LLM_API_KEY']
+            }
+        if os.environ.get('MINDSDB_DEFAULT_EMBEDDING_MODEL_API_KEY', '') != '':
+            self._env_config['default_embedding_model'] = {
+                'api_key': os.environ['MINDSDB_DEFAULT_EMBEDDING_MODEL_API_KEY']
+            }
+
     def parse_cmd_args(self) -> None:
         """Collect cmd args to self._cmd_args (accessable as self.cmd_args)
         """

mindsdb/utilities/context.py

@@ -18,6 +18,7 @@ class Context:
             'user_id': None,
             'company_id': None,
             'session_id': "",
+            'task_id': None,
             'user_class': 0,
             'profiling': {
                 'level': 0,