MindsDB 25.5.4.2__py3-none-any.whl → 25.6.2.0__py3-none-any.whl

mindsdb/interfaces/database/projects.py

@@ -11,9 +11,10 @@ from mindsdb_sql_parser.ast import Select, Star, Constant, Identifier
 from mindsdb_sql_parser import parse_sql
 
 from mindsdb.interfaces.storage import db
-from mindsdb.utilities.config import Config
 from mindsdb.interfaces.model.model_controller import ModelController
 from mindsdb.interfaces.database.views import ViewController
+from mindsdb.utilities import log
+from mindsdb.utilities.config import Config
 from mindsdb.utilities.context import context as ctx
 from mindsdb.utilities.exception import EntityExistsError, EntityNotExistsError
 import mindsdb.utilities.profiler as profiler
@@ -21,6 +22,8 @@ from mindsdb.api.executor.sql_query import SQLQuery
 from mindsdb.api.executor.utilities.sql import query_df
 from mindsdb.interfaces.query_context.context_controller import query_context_controller
 
+logger = log.getLogger(__name__)
+
 
 class Project:
     @staticmethod
@@ -39,11 +42,10 @@ class Project:
         company_id = ctx.company_id if ctx.company_id is not None else 0
 
         existing_record = db.Integration.query.filter(
-            sa.func.lower(db.Integration.name) == name,
-            db.Integration.company_id == ctx.company_id
+            sa.func.lower(db.Integration.name) == name, db.Integration.company_id == ctx.company_id
         ).first()
         if existing_record is not None:
-            raise EntityExistsError('Database exists with this name ', name)
+            raise EntityExistsError("Database exists with this name ", name)
 
         existing_record = db.Project.query.filter(
             (sa.func.lower(db.Project.name) == name)
@@ -51,12 +53,9 @@ class Project:
             & (db.Project.deleted_at == sa.null())
         ).first()
         if existing_record is not None:
-            raise EntityExistsError('Project already exists', name)
+            raise EntityExistsError("Project already exists", name)
 
-        record = db.Project(
-            name=name,
-            company_id=company_id
-        )
+        record = db.Project(name=name, company_id=company_id)
 
         self.record = record
         self.name = name
@@ -68,18 +67,20 @@ class Project:
         self.id = record.id
 
     def delete(self):
-        if self.record.metadata_ and self.record.metadata_.get('is_default', False):
+        if self.record.metadata_ and self.record.metadata_.get("is_default", False):
             raise Exception(
                 f"Project '{self.name}' can not be deleted, because it is default project."
                 "The default project can be changed in the config file or by setting the environment variable MINDSDB_DEFAULT_PROJECT."
             )
 
         tables = self.get_tables()
-        tables = [key for key, val in tables.items() if val['type'] != 'table']
+        tables = [key for key, val in tables.items() if val["type"] != "table"]
         if len(tables) > 0:
-            raise Exception(f"Project '{self.name}' can not be deleted, because it contains tables: {', '.join(tables)}")
+            raise Exception(
+                f"Project '{self.name}' can not be deleted, because it contains tables: {', '.join(tables)}"
+            )
 
-        is_cloud = Config().get('cloud', False)
+        is_cloud = Config().get("cloud", False)
         if is_cloud is True:
             self.record.deleted_at = datetime.datetime.now()
         else:
@@ -91,60 +92,36 @@ class Project:
         db.session.commit()
 
     def drop_model(self, name: str):
-        ModelController().delete_model(
-            name,
-            project_name=self.name
-        )
+        ModelController().delete_model(name, project_name=self.name)
 
     def drop_view(self, name: str):
-        ViewController().delete(
-            name,
-            project_name=self.name
-        )
+        ViewController().delete(name, project_name=self.name)
 
     def create_view(self, name: str, query: str):
-        ViewController().add(
-            name,
-            query=query,
-            project_name=self.name
-        )
+        ViewController().add(name, query=query, project_name=self.name)
 
     def update_view(self, name: str, query: str):
-        ViewController().update(
-            name,
-            query=query,
-            project_name=self.name
-        )
+        ViewController().update(name, query=query, project_name=self.name)
 
     def delete_view(self, name: str):
-        ViewController().delete(
-            name,
-            project_name=self.name
-        )
+        ViewController().delete(name, project_name=self.name)
 
     def get_view_meta(self, query: ASTNode) -> ASTNode:
         view_name = query.from_table.parts[-1]
-        view_meta = ViewController().get(
-            name=view_name,
-            project_name=self.name
-        )
-        view_meta['query_ast'] = parse_sql(view_meta['query'])
+        view_meta = ViewController().get(name=view_name, project_name=self.name)
+        view_meta["query_ast"] = parse_sql(view_meta["query"])
         return view_meta
 
     def query_view(self, query, session):
-
         view_meta = self.get_view_meta(query)
 
-        query_context_controller.set_context('view', view_meta['id'])
+        query_context_controller.set_context("view", view_meta["id"])
 
         try:
-            sqlquery = SQLQuery(
-                view_meta['query_ast'],
-                session=session
-            )
+            sqlquery = SQLQuery(view_meta["query_ast"], session=session)
             df = sqlquery.fetched_data.to_df()
         finally:
-            query_context_controller.release_context('view', view_meta['id'])
+            query_context_controller.release_context("view", view_meta["id"])
 
         # remove duplicated columns
         df = df.loc[:, ~df.columns.duplicated()]
@@ -160,73 +137,62 @@ class Project:
         if (
             predictor_record.training_start_at is not None
             and predictor_record.training_stop_at is None
-            and predictor_record.status != 'error'
+            and predictor_record.status != "error"
         ):
             training_time = round((datetime.datetime.now() - predictor_record.training_start_at).total_seconds(), 3)
-        elif (
-            predictor_record.training_start_at is not None
-            and predictor_record.training_stop_at is not None
-        ):
-            training_time = round((predictor_record.training_stop_at - predictor_record.training_start_at).total_seconds(), 3)
+        elif predictor_record.training_start_at is not None and predictor_record.training_stop_at is not None:
+            training_time = round(
+                (predictor_record.training_stop_at - predictor_record.training_start_at).total_seconds(), 3
+            )
 
         # regon Hide sensitive info
         training_options = predictor_record.learn_args
         handler_module = integration_controller.get_handler_module(integraion_record.engine)
 
         if with_secrets is False and handler_module:
-
-            model_using_args = getattr(handler_module, 'model_using_args', None)
+            model_using_args = getattr(handler_module, "model_using_args", None)
             if (
                 isinstance(model_using_args, dict)
                 and isinstance(training_options, dict)
-                and isinstance(training_options.get('using'), dict)
+                and isinstance(training_options.get("using"), dict)
             ):
-                training_options['using'] = deepcopy(training_options['using'])
+                training_options["using"] = deepcopy(training_options["using"])
                 for key, value in model_using_args.items():
-                    if key in training_options['using'] and value.get('secret', False):
-                        training_options['using'][key] = '******'
+                    if key in training_options["using"] and value.get("secret", False):
+                        training_options["using"][key] = "******"
         # endregion
 
         predictor_meta = {
-            'type': 'model',
-            'id': predictor_record.id,
-            'engine': integraion_record.engine,
-            'engine_name': integraion_record.name,
-            'active': predictor_record.active,
-            'version': predictor_record.version,
-            'status': predictor_record.status,
-            'accuracy': None,
-            'predict': predictor_record.to_predict[0],
-            'update_status': predictor_record.update_status,
-            'mindsdb_version': predictor_record.mindsdb_version,
-            'error': predictor_data.get('error'),
-            'select_data_query': predictor_record.fetch_data_query,
-            'training_options': training_options,
-            'deletable': True,
-            'label': predictor_record.label,
-            'current_training_phase': predictor_record.training_phase_current,
-            'total_training_phases': predictor_record.training_phase_total,
-            'training_phase_name': predictor_record.training_phase_name,
-            'training_time': training_time
-        }
-        if predictor_data.get('accuracies', None) is not None:
-            if len(predictor_data['accuracies']) > 0:
-                predictor_meta['accuracy'] = float(np.mean(list(predictor_data['accuracies'].values())))
-        return {
-            'name': predictor_record.name,
-            'metadata': predictor_meta,
-            'created_at': predictor_record.created_at
+            "type": "model",
+            "id": predictor_record.id,
+            "engine": integraion_record.engine,
+            "engine_name": integraion_record.name,
+            "active": predictor_record.active,
+            "version": predictor_record.version,
+            "status": predictor_record.status,
+            "accuracy": None,
+            "predict": predictor_record.to_predict[0],
+            "update_status": predictor_record.update_status,
+            "mindsdb_version": predictor_record.mindsdb_version,
+            "error": predictor_data.get("error"),
+            "select_data_query": predictor_record.fetch_data_query,
+            "training_options": training_options,
+            "deletable": True,
+            "label": predictor_record.label,
+            "current_training_phase": predictor_record.training_phase_current,
+            "total_training_phases": predictor_record.training_phase_total,
+            "training_phase_name": predictor_record.training_phase_name,
+            "training_time": training_time,
         }
+        if predictor_data.get("accuracies", None) is not None:
+            if len(predictor_data["accuracies"]) > 0:
+                predictor_meta["accuracy"] = float(np.mean(list(predictor_data["accuracies"].values())))
+        return {"name": predictor_record.name, "metadata": predictor_meta, "created_at": predictor_record.created_at}
 
     def get_model(self, name: str):
         record = (
-            db.session.query(db.Predictor, db.Integration).filter_by(
-                project_id=self.id,
-                active=True,
-                name=name,
-                deleted_at=sa.null(),
-                company_id=ctx.company_id
-            )
+            db.session.query(db.Predictor, db.Integration)
+            .filter_by(project_id=self.id, active=True, name=name, deleted_at=sa.null(), company_id=ctx.company_id)
             .join(db.Integration, db.Integration.id == db.Predictor.integration_id)
             .order_by(db.Predictor.name, db.Predictor.id)
             .first()
@@ -237,12 +203,8 @@ class Project:
 
     def get_model_by_id(self, model_id: int):
         record = (
-            db.session.query(db.Predictor, db.Integration).filter_by(
-                project_id=self.id,
-                id=model_id,
-                deleted_at=sa.null(),
-                company_id=ctx.company_id
-            )
+            db.session.query(db.Predictor, db.Integration)
+            .filter_by(project_id=self.id, id=model_id, deleted_at=sa.null(), company_id=ctx.company_id)
             .join(db.Integration, db.Integration.id == db.Predictor.integration_id)
             .order_by(db.Predictor.name, db.Predictor.id)
             .first()
@@ -253,45 +215,38 @@ class Project:
 
     def get_models(self, active: bool = True, with_secrets: bool = True):
         query = db.session.query(db.Predictor, db.Integration).filter_by(
-            project_id=self.id,
-            deleted_at=sa.null(),
-            company_id=ctx.company_id
+            project_id=self.id, deleted_at=sa.null(), company_id=ctx.company_id
         )
         if isinstance(active, bool):
             query = query.filter_by(active=active)
 
-        query = query.join(
-            db.Integration, db.Integration.id == db.Predictor.integration_id
-        ).order_by(db.Predictor.name, db.Predictor.id)
+        query = query.join(db.Integration, db.Integration.id == db.Predictor.integration_id).order_by(
+            db.Predictor.name, db.Predictor.id
+        )
 
         data = []
 
         for predictor_record, integraion_record in query.all():
-            data.append(
-                self._get_model_data(predictor_record, integraion_record, with_secrets)
-            )
+            data.append(self._get_model_data(predictor_record, integraion_record, with_secrets))
 
         return data
 
     def get_agents(self):
         records = (
-            db.session.query(db.Agents).filter(
+            db.session.query(db.Agents)
+            .filter(
                 db.Agents.project_id == self.id,
                 db.Agents.company_id == ctx.company_id,
-                db.Agents.deleted_at == sa.null()
+                db.Agents.deleted_at == sa.null(),
             )
             .order_by(db.Agents.name)
             .all()
         )
         data = [
             {
-                'name': record.name,
-                'query': record.query,
-                'metadata': {
-                    'type': 'agent',
-                    'id': record.id,
-                    'deletable': True
-                }
+                "name": record.name,
+                "query": record.query,
+                "metadata": {"type": "agent", "id": record.id, "deletable": True},
             }
             for record in records
         ]
@@ -299,113 +254,122 @@ class Project:
 
     def get_knowledge_bases(self):
         from mindsdb.api.executor.controllers.session_controller import SessionController
+
         session = SessionController()
 
         return {
-            kb['name']: {
-                'type': 'knowledge_base',
-                'id': kb['id'],
-                'deletable': True
-            }
+            kb["name"]: {"type": "knowledge_base", "id": kb["id"], "deletable": True}
             for kb in session.kb_controller.list(self.name)
         }
 
     def get_views(self):
         records = (
-            db.session.query(db.View).filter_by(
-                project_id=self.id,
-                company_id=ctx.company_id
-            )
+            db.session.query(db.View)
+            .filter_by(project_id=self.id, company_id=ctx.company_id)
             .order_by(db.View.name, db.View.id)
             .all()
         )
-        data = [{
-            'name': view_record.name,
-            'query': view_record.query,
-            'metadata': {
-                'type': 'view',
-                'id': view_record.id,
-                'deletable': True
-            }}
+        data = [
+            {
+                "name": view_record.name,
+                "query": view_record.query,
+                "metadata": {"type": "view", "id": view_record.id, "deletable": True},
+            }
             for view_record in records
         ]
         return data
 
     def get_view(self, name):
-        view_record = db.session.query(db.View).filter(
-            db.View.project_id == self.id,
-            db.View.company_id == ctx.company_id,
-            sa.func.lower(db.View.name) == name.lower(),
-        ).one_or_none()
+        view_record = (
+            db.session.query(db.View)
+            .filter(
+                db.View.project_id == self.id,
+                db.View.company_id == ctx.company_id,
+                sa.func.lower(db.View.name) == name.lower(),
+            )
+            .one_or_none()
+        )
         if view_record is None:
             return view_record
         return {
-            'name': view_record.name,
-            'query': view_record.query,
-            'metadata': {
-                'type': 'view',
-                'id': view_record.id,
-                'deletable': True
-            }
+            "name": view_record.name,
+            "query": view_record.query,
+            "metadata": {"type": "view", "id": view_record.id, "deletable": True},
         }
 
     @profiler.profile()
     def get_tables(self):
         data = OrderedDict()
-        data['models'] = {'type': 'table', 'deletable': False}
+        data["models"] = {"type": "table", "deletable": False}
 
         models = self.get_models()
         for model in models:
-            if model['metadata']['active'] is True:
-                data[model['name']] = model['metadata']
+            if model["metadata"]["active"] is True:
+                data[model["name"]] = model["metadata"]
 
         views = self.get_views()
         for view in views:
-            data[view['name']] = view['metadata']
+            data[view["name"]] = view["metadata"]
 
         agents = self.get_agents()
         for agent in agents:
-            data[agent['name']] = agent['metadata']
+            data[agent["name"]] = agent["metadata"]
 
         data.update(self.get_knowledge_bases())
 
         return data
 
-    def get_columns(self, table_name: str):
-        # at the moment it works only for models
-        predictor_record = db.Predictor.query.filter_by(
-            company_id=ctx.company_id,
-            project_id=self.id,
-            name=table_name
-        ).first()
+    def get_columns(self, table_name: str) -> list[str] | None:
         columns = []
-        if predictor_record is not None:
-            if isinstance(predictor_record.dtype_dict, dict):
-                columns = list(predictor_record.dtype_dict.keys())
-            elif predictor_record.to_predict is not None:
-                # no dtype_dict, use target
-                columns = predictor_record.to_predict
-                if not isinstance(columns, list):
-                    columns = [columns]
+        tables = self.get_tables()
+        table = None
+        for key, value in tables.items():
+            if key.lower() == table_name.lower():
+                table_name = key
+                table = value
+        if table is None:
             return columns
-        if self.get_view(table_name):
-            query = Select(targets=[Star()], from_table=Identifier(table_name), limit=Constant(1))
-
-            from mindsdb.api.executor.controllers.session_controller import SessionController
-            session = SessionController()
-            session.database = self.name
-            df = self.query_view(query, session)
-            return df.columns
-        else:
-            # is it agent?
-            agent = db.Agents.query.filter_by(
-                company_id=ctx.company_id,
-                project_id=self.id,
-                name=table_name
-            ).first()
-            if agent is not None:
-                from mindsdb.interfaces.agents.constants import ASSISTANT_COLUMN, USER_COLUMN
-                columns = [ASSISTANT_COLUMN, USER_COLUMN]
+
+        match str(table["type"]).upper():
+            case "MODEL":
+                predictor_record = db.Predictor.query.filter_by(
+                    company_id=ctx.company_id, project_id=self.id, name=table_name
+                ).first()
+                columns = []
+                if predictor_record is not None:
+                    if isinstance(predictor_record.dtype_dict, dict):
+                        columns = list(predictor_record.dtype_dict.keys())
+                    elif predictor_record.to_predict is not None:
+                        # no dtype_dict, use target
+                        columns = predictor_record.to_predict
+                        if not isinstance(columns, list):
+                            columns = [columns]
+            case "VIEW":
+                query = Select(targets=[Star()], from_table=Identifier(table_name), limit=Constant(1))
+
+                from mindsdb.api.executor.controllers.session_controller import SessionController
+
+                session = SessionController()
+                session.database = self.name
+                df = self.query_view(query, session)
+                columns = df.columns
+            case "AGENT":
+                agent = db.Agents.query.filter_by(
+                    company_id=ctx.company_id, project_id=self.id, name=table_name
+                ).first()
+                if agent is not None:
+                    from mindsdb.interfaces.agents.constants import ASSISTANT_COLUMN, USER_COLUMN
+
+                    columns = [ASSISTANT_COLUMN, USER_COLUMN]
+            case "KNOWLEDGE_BASE":
+                from mindsdb.interfaces.knowledge_base.controller import KB_TO_VECTORDB_COLUMNS
+
+                columns = list(KB_TO_VECTORDB_COLUMNS.keys()) + ["metadata", "relevance", "distance"]
+            case "TABLE":
+                # like 'mindsdb.models'
+                pass
+            case _:
+                logger.warning(f"Unknown table type: {table['type']}")
 
         return columns
 
@@ -417,13 +381,14 @@ class ProjectController:
     def get_list(self) -> List[Project]:
         company_id = ctx.company_id if ctx.company_id is not None else 0
         records = db.Project.query.filter(
-            (db.Project.company_id == company_id)
-            & (db.Project.deleted_at == sa.null())
+            (db.Project.company_id == company_id) & (db.Project.deleted_at == sa.null())
         ).order_by(db.Project.name)
 
         return [Project.from_record(x) for x in records]
 
-    def get(self, id: Optional[int] = None, name: Optional[str] = None, deleted: bool = False, is_default: bool = False) -> Project:
+    def get(
+        self, id: Optional[int] = None, name: Optional[str] = None, deleted: bool = False, is_default: bool = False
+    ) -> Project:
         if id is not None and name is not None:
             raise ValueError("Both 'id' and 'name' can't be provided at the same time")
 
@@ -433,9 +398,7 @@ class ProjectController:
         if id is not None:
             q = q.filter_by(id=id)
         elif name is not None:
-            q = q.filter(
-                (sa.func.lower(db.Project.name) == sa.func.lower(name))
-            )
+            q = q.filter((sa.func.lower(db.Project.name) == sa.func.lower(name)))
 
         if deleted is True:
             q = q.filter((db.Project.deleted_at != sa.null()))
@@ -443,12 +406,12 @@ class ProjectController:
             q = q.filter_by(deleted_at=sa.null())
 
         if is_default:
-            q = q.filter(db.Project.metadata_['is_default'].as_boolean() == is_default)
+            q = q.filter(db.Project.metadata_["is_default"].as_boolean() == is_default)
 
         record = q.first()
 
         if record is None:
-            raise EntityNotExistsError(f'Project not found: {name}')
+            raise EntityNotExistsError(f"Project not found: {name}")
         return Project.from_record(record)
 
     def add(self, name: str) -> Project: