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

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/skills/skill_tool.py

@@ -243,7 +243,13 @@ class SkillToolController:
             kb_table = session_controller.kb_controller.get_table(knowledge_base_name, skill.project_id)
 
             res = kb_table.select_query(query)
-            return '\n'.join(res.content)
+            # Handle both chunk_content and content column names
+            if hasattr(res, 'chunk_content'):
+                return '\n'.join(res.chunk_content)
+            elif hasattr(res, 'content'):
+                return '\n'.join(res.content)
+            else:
+                return "No content or chunk_content found in knowledge base response"
 
         return _answer_question
 

mindsdb/interfaces/skills/sql_agent.py

@@ -12,6 +12,7 @@ from mindsdb_sql_parser.ast import Select, Show, Describe, Explain, Identifier
 from mindsdb.utilities import log
 from mindsdb.utilities.context import context as ctx
 from mindsdb.integrations.utilities.query_traversal import query_traversal
+from mindsdb.integrations.libs.response import INF_SCHEMA_COLUMNS_NAMES
 
 logger = log.getLogger(__name__)
 
@@ -275,9 +276,13 @@ class SQLAgent:
         dn = self._command_executor.session.datahub.get(integration)
 
         fields, dtypes = [], []
-        for column in dn.get_table_columns(table_name, schema_name):
-            fields.append(column['name'])
-            dtypes.append(column.get('type', ''))
+        for df in dn.get_table_columns_df(table_name, schema_name):
+            df_records = df.to_dict(orient='records')
+            fields.append(df_records[INF_SCHEMA_COLUMNS_NAMES.COLUMN_NAME])
+            if df_records[INF_SCHEMA_COLUMNS_NAMES.MYSQL_DATA_TYPE] is not None:
+                dtypes.append(df_records[INF_SCHEMA_COLUMNS_NAMES.MYSQL_DATA_TYPE].value)
+            else:
+                dtypes.append(df_records[INF_SCHEMA_COLUMNS_NAMES.DATA_TYPE])
 
         info = f'Table named `{table_str}`:\n'
         info += f"\nSample with first {self._sample_rows_in_table_info} rows from table {table_str} in CSV format (dialect is 'excel'):\n"

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

@@ -209,6 +209,10 @@ class Config:
                     "restart_on_failure": True,
                     "max_restart_count": 1,
                     "max_restart_interval_seconds": 60
+                },
+                "litellm": {
+                    "host": "0.0.0.0",  # API server binds to all interfaces by default
+                    "port": "8000"
                 }
             },
             "cache": {
@@ -226,7 +230,9 @@ class Config:
             "tasks": {
                 "disable": False
             },
-            "default_project": "mindsdb"
+            "default_project": "mindsdb",
+            "default_llm": {},
+            "default_embedding_model": {}
         }
         # endregion
 
@@ -365,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)
         """
@@ -378,7 +393,8 @@ class Config:
         ):
             self._cmd_args = argparse.Namespace(
                 api=None, config=None, install_handlers=None, verbose=False,
-                no_studio=False, version=False, ml_task_queue_consumer=None
+                no_studio=False, version=False, ml_task_queue_consumer=None,
+                agent=None, project=None
             )
             return
 
@@ -390,6 +406,8 @@ class Config:
         parser.add_argument('--no_studio', action='store_true')
         parser.add_argument('-v', '--version', action='store_true')
         parser.add_argument('--ml_task_queue_consumer', action='store_true', default=None)
+        parser.add_argument('--agent', type=str, default=None, help='Name of the agent to use with litellm APIs')
+        parser.add_argument('--project', type=str, default=None, help='Project containing the agent (default: mindsdb)')
         self._cmd_args = parser.parse_args()
 
     def fetch_auto_config(self) -> bool:

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,

mindsdb/utilities/starters.py

@@ -38,3 +38,10 @@ def start_mcp(*args, **kwargs):
     from mindsdb.api.mcp.start import start
 
     start(*args, **kwargs)
+
+
+def start_litellm(*args, **kwargs):
+    """Start the LiteLLM server"""
+    from mindsdb.api.litellm.start import start
+
+    start(*args, **kwargs)