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

mindsdb/interfaces/skills/sql_agent.py

@@ -14,6 +14,8 @@ 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
 from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
+from mindsdb.utilities.config import config
+from mindsdb.interfaces.data_catalog.data_catalog_reader import DataCatalogReader
 
 logger = log.getLogger(__name__)
 
@@ -28,7 +30,7 @@ def list_to_csv_str(array: List[List[Any]]) -> str:
         str: The array formatted as a CSV string using Excel dialect
     """
     output = StringIO()
-    writer = csv.writer(output, dialect='excel')
+    writer = csv.writer(output, dialect="excel")
     str_array = [[str(item) for item in row] for row in array]
     writer.writerows(str_array)
     return output.getvalue()
@@ -55,23 +57,23 @@ def split_table_name(table_name: str) -> List[str]:
         'input': '`aaa`.`bbb.ccc`', 'output': ['aaa', 'bbb.ccc']
     """
     result = []
-    current = ''
+    current = ""
     in_backticks = False
 
     i = 0
     while i < len(table_name):
-        if table_name[i] == '`':
+        if table_name[i] == "`":
             in_backticks = not in_backticks
-        elif table_name[i] == '.' and not in_backticks:
+        elif table_name[i] == "." and not in_backticks:
             if current:
-                result.append(current.strip('`'))
-                current = ''
+                result.append(current.strip("`"))
+                current = ""
         else:
             current += table_name[i]
         i += 1
 
     if current:
-        result.append(current.strip('`'))
+        result.append(current.strip("`"))
 
     # ensure we split the table name
     result = [r.split(".") for r in result][0]
@@ -89,13 +91,13 @@ class SQLAgent:
         command_executor,
         databases: List[str],
         databases_struct: dict,
-        knowledge_base_database: str = 'mindsdb',
+        knowledge_base_database: str = "mindsdb",
         include_tables: Optional[List[str]] = None,
         ignore_tables: Optional[List[str]] = None,
         include_knowledge_bases: Optional[List[str]] = None,
         ignore_knowledge_bases: Optional[List[str]] = None,
         sample_rows_in_table_info: int = 3,
-        cache: Optional[dict] = None
+        cache: Optional[dict] = None,
     ):
         """
         Initialize SQLAgent.
@@ -133,12 +135,13 @@ class SQLAgent:
         self._cache = cache
 
         from mindsdb.interfaces.skills.skill_tool import SkillToolController
+
         # Initialize the skill tool controller from MindsDB
         self.skill_tool = SkillToolController()
 
     def _call_engine(self, query: str, database=None):
         # switch database
-        ast_query = parse_sql(query.strip('`'))
+        ast_query = parse_sql(query.strip("`"))
         self._check_permissions(ast_query)
 
         if database is None:
@@ -148,10 +151,7 @@ class SQLAgent:
                 # for now, we will just use the first one
                 database = self._databases[0] if self._databases else "mindsdb"
 
-        ret = self._command_executor.execute_command(
-            ast_query,
-            database_name=database
-        )
+        ret = self._command_executor.execute_command(ast_query, database_name=database)
         return ret
 
     def _check_permissions(self, ast_query):
@@ -170,7 +170,7 @@ class SQLAgent:
 
             def _check_f(node, is_table=None, **kwargs):
                 if is_table and isinstance(node, Identifier):
-                    table_name = '.'.join(node.parts)
+                    table_name = ".".join(node.parts)
 
                     # Get the list of available knowledge bases
                     kb_names = self.get_usable_knowledge_base_names()
@@ -182,16 +182,21 @@ class SQLAgent:
                     if is_kb and self._knowledge_bases_to_include:
                         kb_parts = [split_table_name(x) for x in self._knowledge_bases_to_include]
                         if node.parts not in kb_parts:
-                            raise ValueError(f"Knowledge base {table_name} not found. Available knowledge bases: {', '.join(self._knowledge_bases_to_include)}")
+                            raise ValueError(
+                                f"Knowledge base {table_name} not found. Available knowledge bases: {', '.join(self._knowledge_bases_to_include)}"
+                            )
                     # Regular table check
                     elif not is_kb and self._tables_to_include and node.parts not in tables_parts:
-                        raise ValueError(f"Table {table_name} not found. Available tables: {', '.join(self._tables_to_include)}")
+                        raise ValueError(
+                            f"Table {table_name} not found. Available tables: {', '.join(self._tables_to_include)}"
+                        )
                     # Check if it's a restricted knowledge base
                     elif is_kb and table_name in self._knowledge_bases_to_ignore:
                         raise ValueError(f"Knowledge base {table_name} is not allowed.")
                     # Check if it's a restricted table
                     elif not is_kb and table_name in self._tables_to_ignore:
                         raise ValueError(f"Table {table_name} is not allowed.")
+
             query_traversal(ast_query, _check_f)
 
     def get_usable_table_names(self) -> Iterable[str]:
@@ -200,7 +205,7 @@ class SQLAgent:
         Returns:
             Iterable[str]: list with table names
         """
-        cache_key = f'{ctx.company_id}_{",".join(self._databases)}_tables'
+        cache_key = f"{ctx.company_id}_{','.join(self._databases)}_tables"
 
         # first check cache and return if found
         if self._cache:
@@ -218,13 +223,13 @@ class SQLAgent:
 
             schemas_names = list(self._mindsdb_db_struct[db_name].keys())
             if len(schemas_names) > 1 and None in schemas_names:
-                raise Exception('default schema and named schemas can not be used in same filter')
+                raise Exception("default schema and named schemas can not be used in same filter")
 
             if None in schemas_names:
                 # get tables only from default schema
                 response = handler.get_tables()
                 tables_in_default_schema = list(response.data_frame.table_name)
-                schema_tables_restrictions = self._mindsdb_db_struct[db_name][None]     # None - is default schema
+                schema_tables_restrictions = self._mindsdb_db_struct[db_name][None]  # None - is default schema
                 if schema_tables_restrictions is None:
                     for table_name in tables_in_default_schema:
                         result_tables.append([db_name, table_name])
@@ -233,27 +238,27 @@ class SQLAgent:
                         if table_name in tables_in_default_schema:
                             result_tables.append([db_name, table_name])
             else:
-                if 'all' in inspect.signature(handler.get_tables).parameters:
+                if "all" in inspect.signature(handler.get_tables).parameters:
                     response = handler.get_tables(all=True)
                 else:
                     response = handler.get_tables()
                 response_schema_names = list(response.data_frame.table_schema.unique())
                 schemas_intersection = set(schemas_names) & set(response_schema_names)
                 if len(schemas_intersection) == 0:
-                    raise Exception('There are no allowed schemas in ds')
+                    raise Exception("There are no allowed schemas in ds")
 
                 for schema_name in schemas_intersection:
-                    schema_sub_df = response.data_frame[response.data_frame['table_schema'] == schema_name]
+                    schema_sub_df = response.data_frame[response.data_frame["table_schema"] == schema_name]
                     if self._mindsdb_db_struct[db_name][schema_name] is None:
                         # all tables from schema allowed
                         for row in schema_sub_df:
-                            result_tables.append([db_name, schema_name, row['table_name']])
+                            result_tables.append([db_name, schema_name, row["table_name"]])
                     else:
                         for table_name in self._mindsdb_db_struct[db_name][schema_name]:
-                            if table_name in schema_sub_df['table_name'].values:
+                            if table_name in schema_sub_df["table_name"].values:
                                 result_tables.append([db_name, schema_name, table_name])
 
-        result_tables = ['.'.join(x) for x in result_tables]
+        result_tables = [".".join(x) for x in result_tables]
         if self._cache:
             self._cache.set(cache_key, set(result_tables))
         return result_tables
@@ -264,14 +269,14 @@ class SQLAgent:
         Returns:
             Iterable[str]: list with knowledge base names
         """
-        cache_key = f'{ctx.company_id}_{self.knowledge_base_database}_knowledge_bases'
+        cache_key = f"{ctx.company_id}_{self.knowledge_base_database}_knowledge_bases"
 
         # todo we need to fix the cache, file cache can potentially store out of data information
         # # first check cache and return if found
         # if self._cache:
-        #     cached_kbs = self._cache.get(cache_key)
+        #    cached_kbs = self._cache.get(cache_key)
         #     if cached_kbs:
-        #         return cached_kbs
+        #        return cached_kbs
 
         if self._knowledge_bases_to_include:
             return self._knowledge_bases_to_include
@@ -289,13 +294,15 @@ class SQLAgent:
                 try:
                     # Get knowledge bases from the project database
                     kb_controller = self._command_executor.session.kb_controller
-                    kb_names = [kb['name'] for kb in kb_controller.list()]
+                    kb_names = [kb["name"] for kb in kb_controller.list()]
 
                     # Filter knowledge bases based on include list
                     if self._knowledge_bases_to_include:
                         kb_names = [kb_name for kb_name in kb_names if kb_name in self._knowledge_bases_to_include]
                         if not kb_names:
-                            logger.warning(f"No knowledge bases found in the include list: {self._knowledge_bases_to_include}")
+                            logger.warning(
+                                f"No knowledge bases found in the include list: {self._knowledge_bases_to_include}"
+                            )
                             return []
 
                         return kb_names
@@ -317,7 +324,7 @@ class SQLAgent:
             # Filter knowledge bases based on ignore list
             kb_names = []
             for row in result:
-                kb_name = row['name']
+                kb_name = row["name"]
                 if kb_name not in self._knowledge_bases_to_ignore:
                     kb_names.append(kb_name)
 
@@ -368,7 +375,7 @@ class SQLAgent:
         return tables
 
     def get_knowledge_base_info(self, kb_names: Optional[List[str]] = None) -> str:
-        """ Get information about specified knowledge bases.
+        """Get information about specified knowledge bases.
         Follows best practices as specified in: Rajkumar et al, 2022 (https://arxiv.org/abs/2204.00498)
         If `sample_rows_in_table_info`, the specified number of sample rows will be
         appended to each table description. This can increase performance as demonstrated in the paper.
@@ -388,38 +395,56 @@ class SQLAgent:
         return "\n\n".join(kbs_info)
 
     def get_table_info(self, table_names: Optional[List[str]] = None) -> str:
-        """ Get information about specified tables.
+        """Get information about specified tables.
         Follows best practices as specified in: Rajkumar et al, 2022 (https://arxiv.org/abs/2204.00498)
         If `sample_rows_in_table_info`, the specified number of sample rows will be
         appended to each table description. This can increase performance as demonstrated in the paper.
         """
+        if config.get("data_catalog", {}).get("enabled", False):
+            database_table_map = {}
+            for name in self.get_usable_table_names():
+                name = name.replace("`", "")
 
-        all_tables = []
-        for name in self.get_usable_table_names():
-            # remove backticks
-            name = name.replace("`", "")
+                # TODO: Can there be situations where the database name is returned from the above method?
+                parts = name.split(".", 1)
+                database_table_map[parts[0]] = database_table_map.get(parts[0], []) + [parts[1]]
 
-            split = name.split(".")
-            if len(split) > 1:
-                all_tables.append(Identifier(parts=[split[0], split[1]]))
-            else:
-                all_tables.append(Identifier(name))
+            data_catalog_str = ""
+            for database_name, table_names in database_table_map.items():
+                data_catalog_reader = DataCatalogReader(database_name=database_name, table_names=table_names)
 
-        # if table_names is not None:
-        #     all_tables = self._resolve_table_names(table_names, all_tables)
+                data_catalog_str += data_catalog_reader.read_metadata_as_string()
 
-        tables_info = []
-        for table in all_tables:
-            key = f"{ctx.company_id}_{table}_info"
-            table_info = self._cache.get(key) if self._cache else None
-            if True or table_info is None:
-                table_info = self._get_single_table_info(table)
-                if self._cache:
-                    self._cache.set(key, table_info)
+            return data_catalog_str
 
-            tables_info.append(table_info)
+        else:
+            # TODO: Improve old logic without data catalog
+            all_tables = []
+            for name in self.get_usable_table_names():
+                # remove backticks
+                name = name.replace("`", "")
+
+                split = name.split(".")
+                if len(split) > 1:
+                    all_tables.append(Identifier(parts=[split[0], split[1]]))
+                else:
+                    all_tables.append(Identifier(name))
+
+            # if table_names is not None:
+            #     all_tables = self._resolve_table_names(table_names, all_tables)
+
+            tables_info = []
+            for table in all_tables:
+                key = f"{ctx.company_id}_{table}_info"
+                table_info = self._cache.get(key) if self._cache else None
+                if True or table_info is None:
+                    table_info = self._get_single_table_info(table)
+                    if self._cache:
+                        self._cache.set(key, table_info)
+
+                tables_info.append(table_info)
 
-        return "\n\n".join(tables_info)
+            return "\n\n".join(tables_info)
 
     def get_kb_sample_rows(self, kb_name: str) -> str:
         """Get sample rows from a knowledge base.
@@ -430,7 +455,7 @@ class SQLAgent:
         Returns:
             str: A string containing the sample rows from the knowledge base.
         """
-        logger.info(f'_get_sample_rows: knowledge base={kb_name}')
+        logger.info(f"_get_sample_rows: knowledge base={kb_name}")
         command = f"select * from {kb_name} limit 10;"
         try:
             ret = self._call_engine(command)
@@ -438,13 +463,12 @@ class SQLAgent:
 
             def truncate_value(val):
                 str_val = str(val)
-                return str_val if len(str_val) < 100 else (str_val[:100] + '...')
+                return str_val if len(str_val) < 100 else (str_val[:100] + "...")
 
-            sample_rows = list(
-                map(lambda row: [truncate_value(value) for value in row], sample_rows))
+            sample_rows = list(map(lambda row: [truncate_value(value) for value in row], sample_rows))
             sample_rows_str = "\n" + f"{kb_name}:" + list_to_csv_str(sample_rows)
         except Exception as e:
-            logger.info(f'_get_sample_rows error: {e}')
+            logger.info(f"_get_sample_rows error: {e}")
             sample_rows_str = "\n" + "\t [error] Couldn't retrieve sample rows!"
 
         return sample_rows_str
@@ -471,11 +495,9 @@ class SQLAgent:
 
             fields = df[INF_SCHEMA_COLUMNS_NAMES.COLUMN_NAME].to_list()
             dtypes = [
-                mysql_data_type.value if isinstance(mysql_data_type, MYSQL_DATA_TYPE) else (data_type or 'UNKNOWN')
-                for mysql_data_type, data_type
-                in zip(
-                    df[INF_SCHEMA_COLUMNS_NAMES.MYSQL_DATA_TYPE],
-                    df[INF_SCHEMA_COLUMNS_NAMES.DATA_TYPE]
+                mysql_data_type.value if isinstance(mysql_data_type, MYSQL_DATA_TYPE) else (data_type or "UNKNOWN")
+                for mysql_data_type, data_type in zip(
+                    df[INF_SCHEMA_COLUMNS_NAMES.MYSQL_DATA_TYPE], df[INF_SCHEMA_COLUMNS_NAMES.DATA_TYPE]
                 )
             ]
         except Exception as e:
@@ -492,16 +514,18 @@ class SQLAgent:
             logger.warning(f"Could not get sample rows for {table_str}: {e}")
             sample_rows_info = "\n\t [error] Couldn't retrieve sample rows!"
 
-        info = f'Table named `{table_str}`:\n'
+        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"
         info += sample_rows_info + "\n"
-        info += '\nColumn data types: ' + ",\t".join(
-            [f'\n`{field}` : `{dtype}`' for field, dtype in zip(fields, dtypes)]
-        ) + '\n'
+        info += (
+            "\nColumn data types: "
+            + ",\t".join([f"\n`{field}` : `{dtype}`" for field, dtype in zip(fields, dtypes)])
+            + "\n"
+        )
         return info
 
     def _get_sample_rows(self, table: str, fields: List[str]) -> str:
-        logger.info(f'_get_sample_rows: table={table} fields={fields}')
+        logger.info(f"_get_sample_rows: table={table} fields={fields}")
         command = f"select {', '.join(fields)} from {table} limit {self._sample_rows_in_table_info};"
         try:
             ret = self._call_engine(command)
@@ -509,20 +533,19 @@ class SQLAgent:
 
             def truncate_value(val):
                 str_val = str(val)
-                return str_val if len(str_val) < 100 else (str_val[:100] + '...')
+                return str_val if len(str_val) < 100 else (str_val[:100] + "...")
 
-            sample_rows = list(
-                map(lambda row: [truncate_value(value) for value in row], sample_rows))
+            sample_rows = list(map(lambda row: [truncate_value(value) for value in row], sample_rows))
             sample_rows_str = "\n" + list_to_csv_str([fields] + sample_rows)
         except Exception as e:
-            logger.info(f'_get_sample_rows error: {e}')
+            logger.info(f"_get_sample_rows error: {e}")
             sample_rows_str = "\n" + "\t [error] Couldn't retrieve sample rows!"
 
         return sample_rows_str
 
     def _clean_query(self, query: str) -> str:
         # Sometimes LLM can input markdown into query tools.
-        cmd = re.sub(r'```(sql)?', '', query)
+        cmd = re.sub(r"```(sql)?", "", query)
         return cmd
 
     def query(self, command: str, fetch: str = "all") -> str:
@@ -534,16 +557,16 @@ class SQLAgent:
         def _repr_result(ret):
             limit_rows = 30
 
-            columns_str = ', '.join([repr(col.name) for col in ret.columns])
-            res = f'Output columns: {columns_str}\n'
+            columns_str = ", ".join([repr(col.name) for col in ret.columns])
+            res = f"Output columns: {columns_str}\n"
 
             data = ret.to_lists()
             if len(data) > limit_rows:
                 df = pd.DataFrame(data, columns=[col.name for col in ret.columns])
 
-                res += f'Result has {len(data)} rows. Description of data:\n'
-                res += str(df.describe(include='all')) + '\n\n'
-                res += f'First {limit_rows} rows:\n'
+                res += f"Result has {len(data)} rows. Description of data:\n"
+                res += str(df.describe(include="all")) + "\n\n"
+                res += f"First {limit_rows} rows:\n"
 
             else:
                 res += "Result in CSV format (dialect is 'excel'):\n"
@@ -562,20 +585,20 @@ class SQLAgent:
 
     def get_table_info_safe(self, table_names: Optional[List[str]] = None) -> str:
         try:
-            logger.info(f'get_table_info_safe: {table_names}')
+            logger.info(f"get_table_info_safe: {table_names}")
             return self.get_table_info(table_names)
         except Exception as e:
-            logger.info(f'get_table_info_safe error: {e}')
+            logger.info(f"get_table_info_safe error: {e}")
             return f"Error: {e}"
 
     def query_safe(self, command: str, fetch: str = "all") -> str:
         try:
-            logger.info(f'query_safe (fetch={fetch}): {command}')
+            logger.info(f"query_safe (fetch={fetch}): {command}")
             return self.query(command, fetch)
         except Exception as e:
             logger.error(f"Error in query_safe: {str(e)}\n{traceback.format_exc()}")
-            logger.info(f'query_safe error: {e}')
+            logger.info(f"query_safe error: {e}")
             msg = f"Error: {e}"
-            if 'does not exist' in msg and ' relation ' in msg:
-                msg += '\nAvailable tables: ' + ', '.join(self.get_usable_table_names())
+            if "does not exist" in msg and " relation " in msg:
+                msg += "\nAvailable tables: " + ", ".join(self.get_usable_table_names())
             return msg