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

mindsdb/interfaces/skills/skill_tool.py

@@ -1,5 +1,6 @@
 import enum
 import inspect
+
 from dataclasses import dataclass
 from collections import defaultdict
 from typing import List, Dict, Optional
@@ -14,7 +15,7 @@ from mindsdb.utilities.config import config
 from mindsdb.interfaces.storage import db
 from mindsdb.interfaces.skills.sql_agent import SQLAgent
 from mindsdb.integrations.libs.vectordatabase_handler import TableField
-
+from mindsdb.interfaces.agents.constants import DEFAULT_TEXT2SQL_DATABASE
 
 _DEFAULT_TOP_K_SIMILARITY_SEARCH = 5
 _MAX_CACHE_SIZE = 1000
@@ -120,7 +121,6 @@ class SkillToolController:
         try:
             from mindsdb.interfaces.agents.mindsdb_database_agent import MindsDBSQL
             from mindsdb.interfaces.skills.custom.text2sql.mindsdb_sql_toolkit import MindsDBSQLToolkit
-            from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
         except ImportError:
             raise ImportError(
                 'To use the text-to-SQL skill, please install langchain with `pip install mindsdb[langchain]`')
@@ -132,75 +132,220 @@ class SkillToolController:
             return f'`{name}`'
 
         tables_list = []
+        knowledge_bases_list = []
+        ignore_knowledge_bases_list = []
+
+        # Track databases extracted from dot notation
+        extracted_databases = set()
+
+        # Initialize knowledge_base_database with default value
+        knowledge_base_database = DEFAULT_TEXT2SQL_DATABASE  # Default to mindsdb project
+
+        # First pass: collect all database and knowledge base parameters
         for skill in skills:
-            database = skill.params['database']
-            restriction_on_tables = skill.restriction_on_tables
-            if restriction_on_tables is None:
-                handler = command_executor.session.integration_controller.get_data_handler(database)
-                if 'all' in inspect.signature(handler.get_tables).parameters:
-                    response = handler.get_tables(all=True)
-                else:
-                    response = handler.get_tables()
-                # no restrictions
-                columns = [c.lower() for c in response.data_frame.columns]
-                name_idx = columns.index('table_name') if 'table_name' in columns else 0
-
-                if 'table_schema' in response.data_frame.columns:
-                    for _, row in response.data_frame.iterrows():
-                        tables_list.append(f"{database}.{row['table_schema']}.{escape_table_name(row[name_idx])}")
-                else:
-                    for table_name in response.data_frame.iloc[:, name_idx]:
-                        tables_list.append(f"{database}.{escape_table_name(table_name)}")
+            # Update knowledge_base_database if specified in any skill
+            if skill.params.get('knowledge_base_database'):
+                knowledge_base_database = skill.params.get('knowledge_base_database')
+
+            # Extract databases from include_tables with dot notation
+            if skill.params.get('include_tables'):
+                include_tables = skill.params.get('include_tables')
+                if isinstance(include_tables, str):
+                    include_tables = [t.strip() for t in include_tables.split(',')]
+
+                # Extract database names from dot notation
+                for table in include_tables:
+                    if '.' in table:
+                        db_name = table.split('.')[0]
+                        extracted_databases.add(db_name)
+
+            # Extract databases from include_knowledge_bases with dot notation
+            if skill.params.get('include_knowledge_bases'):
+                include_kbs = skill.params.get('include_knowledge_bases')
+                if isinstance(include_kbs, str):
+                    include_kbs = [kb.strip() for kb in include_kbs.split(',')]
+
+                # Extract database names from dot notation
+                for kb in include_kbs:
+                    if '.' in kb:
+                        db_name = kb.split('.')[0]
+                        if db_name != knowledge_base_database:
+                            # Only update if it's different from the default
+                            knowledge_base_database = db_name
+
+        # Second pass: collect all tables and knowledge base restrictions
+        for skill in skills:
+            # Get database for tables (this is an actual database connection)
+            database = skill.params.get('database', DEFAULT_TEXT2SQL_DATABASE)
+
+            # Add databases extracted from dot notation if no explicit database is provided
+            if not database and extracted_databases:
+                # Use the first extracted database if no explicit database is provided
+                database = next(iter(extracted_databases))
+                # Update the skill params with the extracted database
+                skill.params['database'] = database
+
+            # Extract knowledge base restrictions if they exist in the skill params
+            if skill.params.get('include_knowledge_bases'):
+                # Convert to list if it's a string
+                include_kbs = skill.params.get('include_knowledge_bases')
+                if isinstance(include_kbs, str):
+                    include_kbs = [kb.strip() for kb in include_kbs.split(',')]
+
+                # Process each knowledge base name
+                for kb in include_kbs:
+                    # If it doesn't have a dot, prefix it with the knowledge_base_database
+                    if '.' not in kb:
+                        knowledge_bases_list.append(f"{knowledge_base_database}.{kb}")
+                    else:
+                        knowledge_bases_list.append(kb)
+
+            # Collect ignore_knowledge_bases
+            if skill.params.get('ignore_knowledge_bases'):
+                # Convert to list if it's a string
+                ignore_kbs = skill.params.get('ignore_knowledge_bases')
+                if isinstance(ignore_kbs, str):
+                    ignore_kbs = [kb.strip() for kb in ignore_kbs.split(',')]
+
+                # Process each knowledge base name to ignore
+                for kb in ignore_kbs:
+                    # If it doesn't have a dot, prefix it with the knowledge_base_database
+                    if '.' not in kb:
+                        ignore_knowledge_bases_list.append(f"{knowledge_base_database}.{kb}")
+                    else:
+                        ignore_knowledge_bases_list.append(kb)
+
+            # Skip if no database specified
+            if not database:
                 continue
-            for schema_name, tables in restriction_on_tables.items():
-                for table in tables:
-                    if schema_name is None:
-                        tables_list.append(f'{database}.{escape_table_name(table)}')
+
+            # Process include_tables with dot notation
+            if skill.params.get('include_tables'):
+                include_tables = skill.params.get('include_tables')
+                if isinstance(include_tables, str):
+                    include_tables = [t.strip() for t in include_tables.split(',')]
+
+                for table in include_tables:
+                    # If table already has a database prefix, use it as is
+                    if '.' in table:
+                        # Check if the table already has backticks
+                        if '`' in table:
+                            tables_list.append(table)
+                        else:
+                            # Apply escape_table_name only to the table part
+                            parts = table.split('.')
+                            if len(parts) == 2:
+                                # Format: database.table
+                                tables_list.append(f"{parts[0]}.{escape_table_name(parts[1])}")
+                            elif len(parts) == 3:
+                                # Format: database.schema.table
+                                tables_list.append(f"{parts[0]}.{parts[1]}.{escape_table_name(parts[2])}")
+                            else:
+                                # Unusual format, escape the whole thing
+                                tables_list.append(escape_table_name(table))
                     else:
-                        tables_list.append(f'{database}.{schema_name}.{escape_table_name(table)}')
+                        # Otherwise, prefix with the database
+                        tables_list.append(f"{database}.{escape_table_name(table)}")
+
+                # Skip further table processing if include_tables is specified
+                continue
+
+            restriction_on_tables = skill.restriction_on_tables
+
+            if restriction_on_tables is None and database:
+                try:
+                    handler = command_executor.session.integration_controller.get_data_handler(database)
+                    if 'all' in inspect.signature(handler.get_tables).parameters:
+                        response = handler.get_tables(all=True)
+                    else:
+                        response = handler.get_tables()
+                    # no restrictions
+                    columns = [c.lower() for c in response.data_frame.columns]
+                    name_idx = columns.index('table_name') if 'table_name' in columns else 0
+
+                    if 'table_schema' in response.data_frame.columns:
+                        for _, row in response.data_frame.iterrows():
+                            tables_list.append(f"{database}.{row['table_schema']}.{escape_table_name(row[name_idx])}")
+                    else:
+                        for table_name in response.data_frame.iloc[:, name_idx]:
+                            tables_list.append(f"{database}.{escape_table_name(table_name)}")
+                except Exception as e:
+                    logger.warning(f"Could not get tables from database {database}: {str(e)}")
+                continue
+
+            # Handle table restrictions
+            if restriction_on_tables and database:
+                for schema_name, tables in restriction_on_tables.items():
+                    for table in tables:
+                        # Check if the table already has dot notation (e.g., 'postgresql_conn.home_rentals')
+                        if '.' in table:
+                            # Table already has database prefix, add it directly
+                            tables_list.append(escape_table_name(table))
+                        else:
+                            # No dot notation, apply schema and database as needed
+                            if schema_name is None:
+                                tables_list.append(f'{database}.{escape_table_name(table)}')
+                            else:
+                                tables_list.append(f'{database}.{schema_name}.{escape_table_name(table)}')
+                continue
+
+        # Remove duplicates from lists
+        tables_list = list(set(tables_list))
+        knowledge_bases_list = list(set(knowledge_bases_list))
+        ignore_knowledge_bases_list = list(set(ignore_knowledge_bases_list))
+
+        # Determine knowledge base parameters to pass to SQLAgent
+        include_knowledge_bases = knowledge_bases_list if knowledge_bases_list else None
+        ignore_knowledge_bases = ignore_knowledge_bases_list if ignore_knowledge_bases_list else None
+
+        # If both include and ignore lists exist, include takes precedence
+        if include_knowledge_bases:
+            ignore_knowledge_bases = None
+
+        # # Get all databases from skills and extracted databases
+        # all_databases = list(set([s.params.get('database', DEFAULT_TEXT2SQL_DATABASE) for s in skills if s.params.get('database')] + list(extracted_databases)))
+        #
+        #
+        # # If no databases were specified or extracted, use 'mindsdb' as a default
+        # if not all_databases:
+        #     all_databases = [DEFAULT_TEXT2SQL_DATABASE]
+        #
+
+        all_databases = []
+        # Filter out None values
+        all_databases = [db for db in all_databases if db is not None]
+
+        # Create a databases_struct dictionary that includes all extracted databases
+        databases_struct = {}
+
+        # First, add databases from skills with explicit database parameters
+        for skill in skills:
+            if skill.params.get('database'):
+                databases_struct[skill.params['database']] = skill.restriction_on_tables
+
+        # Then, add all extracted databases with no restrictions
+        for db_name in extracted_databases:
+            if db_name not in databases_struct:
+                databases_struct[db_name] = None
 
         sql_agent = SQLAgent(
             command_executor=command_executor,
-            databases=list(set(s.params['database'] for s in skills)),
-            databases_struct={
-                skill.params['database']: skill.restriction_on_tables
-                for skill in skills
-            },
+            databases=all_databases,
+            databases_struct=databases_struct,
             include_tables=tables_list,
             ignore_tables=None,
+            include_knowledge_bases=include_knowledge_bases,
+            ignore_knowledge_bases=ignore_knowledge_bases,
+            knowledge_base_database=knowledge_base_database,
             sample_rows_in_table_info=3,
+
             cache=get_cache('agent', max_size=_MAX_CACHE_SIZE)
         )
         db = MindsDBSQL.custom_init(
             sql_agent=sql_agent
         )
-
-        # Users probably don't need to configure this for now.
-        sql_database_tools = MindsDBSQLToolkit(db=db, llm=llm).get_tools()
-        descriptions = []
-        for skill in skills:
-            description = skill.params.get('description', '')
-            if description:
-                descriptions.append(description)
-
-        for i, tool in enumerate(sql_database_tools):
-            if isinstance(tool, QuerySQLDataBaseTool):
-                # Add our own custom description so our agent knows when to query this table.
-                original_description = tool.description
-                tool.description = ''
-                if len(descriptions) > 0:
-                    tool.description += f'Use this tool if you need data about {" OR ".join(descriptions)}.\n'
-                tool.description += 'Use the conversation context to decide which table to query.\n'
-                if len(tables_list) > 0:
-                    f'These are the available tables: {",".join(tables_list)}.\n'
-                tool.description += (
-                    'ALWAYS consider these special cases:\n'
-                    ' - For TIMESTAMP type columns, make sure you include the time portion in your query (e.g. WHERE date_column = "2020-01-01 12:00:00")\n'
-                    'Here are the rest of the instructions:\n'
-                    f'{original_description}'
-                )
-                sql_database_tools[i] = tool
-        return sql_database_tools
+        toolkit = MindsDBSQLToolkit(db=db, llm=llm)
+        return toolkit.get_tools()
 
     def _make_retrieval_tools(self, skill: db.Skills, llm, embedding_model):
         """

mindsdb/interfaces/skills/sql_agent.py

@@ -1,7 +1,7 @@
-
 import re
 import csv
 import inspect
+import traceback
 from io import StringIO
 from typing import Iterable, List, Optional, Any
 
@@ -13,6 +13,7 @@ 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
+from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
 
 logger = log.getLogger(__name__)
 
@@ -34,7 +35,7 @@ def list_to_csv_str(array: List[List[Any]]) -> str:
 
 
 def split_table_name(table_name: str) -> List[str]:
-    """Split table name from llm to parst
+    """Split table name from llm to parts
 
     Args:
         table_name (str): input table name
@@ -72,34 +73,69 @@ def split_table_name(table_name: str) -> List[str]:
     if current:
         result.append(current.strip('`'))
 
+    # ensure we split the table name
+    result = [r.split(".") for r in result][0]
+
     return result
 
 
 class SQLAgent:
+    """
+    SQLAgent is a class that handles SQL queries for agents.
+    """
+
     def __init__(
-            self,
-            command_executor,
-            databases: List[str],
-            databases_struct: dict,
-            include_tables: Optional[List[str]] = None,
-            ignore_tables: Optional[List[str]] = None,
-            sample_rows_in_table_info: int = 3,
-            cache: Optional[dict] = None
+        self,
+        command_executor,
+        databases: List[str],
+        databases_struct: dict,
+        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
     ):
+        """
+        Initialize SQLAgent.
+
+        Args:
+            command_executor: Executor for SQL commands
+            databases (List[str]): List of databases to use
+            databases_struct (dict): Dictionary of database structures
+            knowledge_base_database (str): Project name where knowledge bases are stored (defaults to 'mindsdb')
+            include_tables (List[str]): Tables to include
+            ignore_tables (List[str]): Tables to ignore
+            include_knowledge_bases (List[str]): Knowledge bases to include
+            ignore_knowledge_bases (List[str]): Knowledge bases to ignore
+            sample_rows_in_table_info (int): Number of sample rows to include in table info
+            cache (Optional[dict]): Cache for query results
+        """
         self._command_executor = command_executor
         self._mindsdb_db_struct = databases_struct
-
+        self.knowledge_base_database = knowledge_base_database  # This is a project name, not a database connection
         self._sample_rows_in_table_info = int(sample_rows_in_table_info)
 
         self._tables_to_include = include_tables
         self._tables_to_ignore = []
+        self._knowledge_bases_to_include = include_knowledge_bases
+        self._knowledge_bases_to_ignore = []
         self._databases = databases
         if not self._tables_to_include:
             # ignore_tables and include_tables should not be used together.
             # include_tables takes priority if it's set.
             self._tables_to_ignore = ignore_tables or []
+        if not self._knowledge_bases_to_include:
+            # ignore_knowledge_bases and include_knowledge_bases should not be used together.
+            # include_knowledge_bases takes priority if it's set.
+            self._knowledge_bases_to_ignore = ignore_knowledge_bases or []
         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('`'))
@@ -107,7 +143,10 @@ class SQLAgent:
 
         if database is None:
             # if we use tables with prefixes it should work for any database
-            database = self._databases[0]
+            if self._databases is not None:
+                # if we have multiple databases, we need to check which one to use
+                # 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,
@@ -131,9 +170,28 @@ class SQLAgent:
 
             def _check_f(node, is_table=None, **kwargs):
                 if is_table and isinstance(node, Identifier):
-                    if node.parts not in tables_parts:
-                        raise ValueError(f"Table {'.'.join(node.parts)} not found. Available tables: {', '.join(self._tables_to_include)}")
-
+                    table_name = '.'.join(node.parts)
+
+                    # Get the list of available knowledge bases
+                    kb_names = self.get_usable_knowledge_base_names()
+
+                    # Check if this table is a knowledge base
+                    is_kb = table_name in kb_names
+
+                    # If it's a knowledge base and we have knowledge base restrictions
+                    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)}")
+                    # 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)}")
+                    # 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,6 +258,78 @@ class SQLAgent:
             self._cache.set(cache_key, set(result_tables))
         return result_tables
 
+    def get_usable_knowledge_base_names(self) -> Iterable[str]:
+        """Get a list of knowledge bases that the agent has access to.
+
+        Returns:
+            Iterable[str]: list with knowledge base names
+        """
+        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)
+        #     if cached_kbs:
+        #         return cached_kbs
+
+        if self._knowledge_bases_to_include:
+            return self._knowledge_bases_to_include
+
+        try:
+            # Query to get all knowledge bases
+            query = f"SHOW KNOWLEDGE_BASES FROM {self.knowledge_base_database};"
+            try:
+                result = self._call_engine(query, database=self.knowledge_base_database)
+            except Exception as e:
+                # If the direct query fails, try a different approach
+                # This handles the case where knowledge_base_database is not a valid integration
+                logger.warning(f"Error querying knowledge bases from {self.knowledge_base_database}: {str(e)}")
+                # Try to get knowledge bases directly from the project database
+                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()]
+
+                    # 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}")
+                            return []
+
+                        return kb_names
+
+                    # Filter knowledge bases based on ignore list
+                    kb_names = [kb_name for kb_name in kb_names if kb_name not in self._knowledge_bases_to_ignore]
+
+                    if self._cache:
+                        self._cache.set(cache_key, set(kb_names))
+
+                    return kb_names
+                except Exception as inner_e:
+                    logger.error(f"Error getting knowledge bases from kb_controller: {str(inner_e)}")
+                    return []
+
+            if not result:
+                return []
+
+            # Filter knowledge bases based on ignore list
+            kb_names = []
+            for row in result:
+                kb_name = row['name']
+                if kb_name not in self._knowledge_bases_to_ignore:
+                    kb_names.append(kb_name)
+
+            if self._cache:
+                self._cache.set(cache_key, set(kb_names))
+
+            return kb_names
+        except Exception as e:
+            # If there's an error, log it and return an empty list
+            logger.error(f"Error in get_usable_knowledge_base_names: {str(e)}")
+            return []
+
     def _resolve_table_names(self, table_names: List[str], all_tables: List[Identifier]) -> List[Identifier]:
         """
         Tries to find table (which comes directly from an LLM) by its name
@@ -237,6 +367,26 @@ class SQLAgent:
 
         return tables
 
+    def get_knowledge_base_info(self, kb_names: Optional[List[str]] = None) -> str:
+        """ 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.
+        """
+
+        kbs_info = []
+        for kb in kb_names:
+            key = f"{ctx.company_id}_{kb}_info"
+            kb_info = self._cache.get(key) if self._cache else None
+            if True or kb_info is None:
+                kb_info = self.get_kb_sample_rows(kb)
+                if self._cache:
+                    self._cache.set(key, kb_info)
+
+            kbs_info.append(kb_info)
+
+        return "\n\n".join(kbs_info)
+
     def get_table_info(self, table_names: Optional[List[str]] = None) -> str:
         """ Get information about specified tables.
         Follows best practices as specified in: Rajkumar et al, 2022 (https://arxiv.org/abs/2204.00498)
@@ -244,16 +394,25 @@ class SQLAgent:
         appended to each table description. This can increase performance as demonstrated in the paper.
         """
 
-        all_tables = [Identifier(name) for name in self.get_usable_table_names()]
+        all_tables = []
+        for name in self.get_usable_table_names():
+            # remove backticks
+            name = name.replace("`", "")
 
-        if table_names is not None:
-            all_tables = self._resolve_table_names(table_names, all_tables)
+            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 table_info is 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)
@@ -262,6 +421,34 @@ class SQLAgent:
 
         return "\n\n".join(tables_info)
 
+    def get_kb_sample_rows(self, kb_name: str) -> str:
+        """Get sample rows from a knowledge base.
+
+        Args:
+            kb_name (str): The name of the knowledge base.
+
+        Returns:
+            str: A string containing the sample rows from the knowledge base.
+        """
+        logger.info(f'_get_sample_rows: knowledge base={kb_name}')
+        command = f"select * from {kb_name} limit 10;"
+        try:
+            ret = self._call_engine(command)
+            sample_rows = ret.data.to_lists()
+
+            def truncate_value(val):
+                str_val = str(val)
+                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_str = "\n" + f"{kb_name}:" + list_to_csv_str(sample_rows)
+        except Exception as 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 _get_single_table_info(self, table: Identifier) -> str:
         if len(table.parts) < 2:
             raise ValueError(f"Database is required for table: {table}")
@@ -276,19 +463,41 @@ class SQLAgent:
         dn = self._command_executor.session.datahub.get(integration)
 
         fields, dtypes = [], []
-        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])
+        try:
+            df = dn.get_table_columns_df(table_name, schema_name)
+            if not isinstance(df, pd.DataFrame) or df.empty:
+                logger.warning(f"Received empty or invalid DataFrame for table columns of {table_str}")
+                return f"Table named `{table_str}`:\n [No column information available]"
+
+            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]
+                )
+            ]
+        except Exception as e:
+            logger.error(f"Failed processing column info for {table_str}: {e}", exc_info=True)
+            raise ValueError(f"Failed to process column info for {table_str}") from e
+
+        if not fields:
+            logger.error(f"Could not extract column fields for {table_str}.")
+            return f"Table named `{table_str}`:\n [Could not extract column information]"
+
+        try:
+            sample_rows_info = self._get_sample_rows(table_str, fields)
+        except Exception as e:
+            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"\nSample with first {self._sample_rows_in_table_info} rows from table {table_str} in CSV format (dialect is 'excel'):\n"
-        info += self._get_sample_rows(table_str, fields) + "\n"
+        info += sample_rows_info + "\n"
         info += '\nColumn data types: ' + ",\t".join(
-            [f'\n`{field}` : `{dtype}`' for field, dtype in zip(fields, dtypes)]) + '\n'  # noqa
+            [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:
@@ -364,6 +573,7 @@ class SQLAgent:
             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}')
             msg = f"Error: {e}"
             if 'does not exist' in msg and ' relation ' in msg: