MindsDB 25.9.2.0a1__py3-none-any.whl → 25.9.3rc1__py3-none-any.whl

mindsdb/interfaces/skills/retrieval_tool.py

@@ -1,32 +1,32 @@
-import traceback
+from langchain_core.documents import Document
+from langchain_core.tools import Tool
 
 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.rag.settings import RAGPipelineModel
 from mindsdb.integrations.utilities.sql_utils import FilterCondition, FilterOperator
-
+from mindsdb.integrations.libs.response import RESPONSE_TYPE
+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.skills.skill_tool import skill_tool
 from mindsdb.interfaces.storage import db
 from mindsdb.interfaces.storage.db import KnowledgeBase
 from mindsdb.utilities import log
-from langchain_core.documents import Document
-from langchain_core.tools import Tool
-from mindsdb.integrations.libs.response import RESPONSE_TYPE
-from mindsdb.integrations.handlers.langchain_embedding_handler.langchain_embedding_handler import construct_model_from_args
 
 logger = log.getLogger(__name__)
 
 
 def _load_rag_config(tool: dict, pred_args: dict, skill: db.Skills) -> RAGPipelineModel:
-    tools_config = tool['config']
+    tools_config = tool["config"]
     tools_config.update(pred_args)
 
     kb_params = {}
     embeddings_model = None
 
-    if 'source' in tool:
-        kb_name = tool['source']
+    if "source" in tool:
+        kb_name = tool["source"]
         executor = skill_tool.get_command_executor()
         kb = _get_knowledge_base(kb_name, skill.project_id, executor)
 
@@ -34,30 +34,26 @@ def _load_rag_config(tool: dict, pred_args: dict, skill: db.Skills) -> RAGPipeli
             raise ValueError(f"Knowledge base not found: {kb_name}")
 
         kb_table = executor.session.kb_controller.get_table(kb.name, kb.project_id)
-        vector_store_config = {
-            'kb_table': kb_table
-        }
-        is_sparse = tools_config.pop('is_sparse', None)
-        vector_size = tools_config.pop('vector_size', None)
+        vector_store_config = {"kb_table": kb_table}
+        is_sparse = tools_config.pop("is_sparse", None)
+        vector_size = tools_config.pop("vector_size", None)
         if is_sparse is not None:
-            vector_store_config['is_sparse'] = is_sparse
+            vector_store_config["is_sparse"] = is_sparse
         if vector_size is not None:
-            vector_store_config['vector_size'] = vector_size
-        kb_params = {
-            'vector_store_config': vector_store_config
-        }
+            vector_store_config["vector_size"] = vector_size
+        kb_params = {"vector_store_config": vector_store_config}
 
         # Get embedding model from knowledge base table
         if kb_table._kb.embedding_model:
             # Extract embedding model args from knowledge base table
-            embedding_args = kb_table._kb.embedding_model.learn_args.get('using', {})
+            embedding_args = kb_table._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}")
         else:
             embeddings_model = DEFAULT_EMBEDDINGS_MODEL_CLASS()
             logger.debug("Using default embedding model as knowledge base has no embedding model")
-    elif 'embedding_model' not in tools_config:
+    elif "embedding_model" not in tools_config:
         embeddings_model = DEFAULT_EMBEDDINGS_MODEL_CLASS()
         logger.debug("Using default embedding model as no knowledge base provided")
 
@@ -75,29 +71,28 @@ def _build_rag_pipeline_tool(tool: dict, pred_args: dict, skill: db.Skills):
         try:
             result = rag_pipeline(query)
             logger.debug(f"RAG pipeline result: {result}")
-            return result['answer']
+            return result["answer"]
         except Exception as e:
-            logger.error(f"Error in RAG pipeline: {str(e)}")
-            logger.error(traceback.format_exc())
+            logger.exception("Error in RAG pipeline:")
             return f"Error in retrieval: {str(e)}"
 
     # Create RAG tool
-    tools_config = tool['config']
+    tools_config = tool["config"]
     tools_config.update(pred_args)
     return Tool(
         func=rag_wrapper,
-        name=tool['name'],
-        description=tool['description'],
-        response_format='content',
+        name=tool["name"],
+        description=tool["description"],
+        response_format="content",
         # Return directly by default since we already use an LLM against retrieved context to generate a response.
-        return_direct=tools_config.get('return_direct', True)
+        return_direct=tools_config.get("return_direct", True),
     )
 
 
 def _build_name_lookup_tool(tool: dict, pred_args: dict, skill: db.Skills):
-    if 'source' not in tool:
+    if "source" not in tool:
         raise ValueError("Knowledge base for tool not found")
-    kb_name = tool['source']
+    kb_name = tool["source"]
     executor = skill_tool.get_command_executor()
     kb = _get_knowledge_base(kb_name, skill.project_id, executor)
     if not kb:
@@ -110,16 +105,16 @@ def _build_name_lookup_tool(tool: dict, pred_args: dict, skill: db.Skills):
 
     def _get_document_by_name(name: str):
         if metadata_config.name_column_index is not None:
-            tsquery_str = ' & '.join(name.split(' '))
+            tsquery_str = " & ".join(name.split(" "))
             documents_response = vector_db_handler.native_query(
-                f'SELECT * FROM {metadata_config.table} WHERE {metadata_config.name_column_index} @@ to_tsquery(\'{tsquery_str}\') LIMIT 1;'
+                f"SELECT * FROM {metadata_config.table} WHERE {metadata_config.name_column_index} @@ to_tsquery('{tsquery_str}') LIMIT 1;"
             )
         else:
             documents_response = vector_db_handler.native_query(
-                f'SELECT * FROM {metadata_config.table} WHERE "{metadata_config.name_column}" ILIKE \'%{name}%\' LIMIT 1;'
+                f"SELECT * FROM {metadata_config.table} WHERE \"{metadata_config.name_column}\" ILIKE '%{name}%' LIMIT 1;"
             )
         if documents_response.resp_type == RESPONSE_TYPE.ERROR:
-            raise RuntimeError(f'There was an error looking up documents: {documents_response.error_message}')
+            raise RuntimeError(f"There was an error looking up documents: {documents_response.error_message}")
         if documents_response.data_frame.empty:
             return None
         document_row = documents_response.data_frame.head(1)
@@ -127,38 +122,36 @@ def _build_name_lookup_tool(tool: dict, pred_args: dict, skill: db.Skills):
         id_filter_condition = FilterCondition(
             f"{metadata_config.embeddings_metadata_column}->>'{metadata_config.doc_id_key}'",
             FilterOperator.EQUAL,
-            str(document_row.get(metadata_config.id_column).item())
+            str(document_row.get(metadata_config.id_column).item()),
         )
         document_chunks_df = vector_db_handler.select(
-            metadata_config.embeddings_table,
-            conditions=[id_filter_condition]
+            metadata_config.embeddings_table, conditions=[id_filter_condition]
         )
         if document_chunks_df.empty:
             return None
-        sort_col = 'chunk_id' if 'chunk_id' in document_chunks_df.columns else 'id'
+        sort_col = "chunk_id" if "chunk_id" in document_chunks_df.columns else "id"
         document_chunks_df.sort_values(by=sort_col)
-        content = ''
+        content = ""
         for _, chunk in document_chunks_df.iterrows():
             if len(content) > metadata_config.max_document_context:
                 break
-            content += chunk.get(metadata_config.content_column, '')
+            content += chunk.get(metadata_config.content_column, "")
 
-        return Document(
-            page_content=content,
-            metadata=document_row.to_dict(orient='records')[0]
-        )
+        return Document(page_content=content, metadata=document_row.to_dict(orient="records")[0])
 
     def _lookup_document_by_name(name: str):
         found_document = _get_document_by_name(name)
         if found_document is None:
-            return f'I could not find any document with name {name}. Please make sure the document name matches exactly.'
+            return (
+                f"I could not find any document with name {name}. Please make sure the document name matches exactly."
+            )
         return f"I found document {found_document.metadata.get(metadata_config.id_column)} with name {found_document.metadata.get(metadata_config.name_column)}. Here is the full document to use as context:\n\n{found_document.page_content}"
 
     return Tool(
         func=_lookup_document_by_name,
-        name=tool.get('name', '') + '_name_lookup',
-        description='You must use this tool ONLY when the user is asking about a specific document by name or title. The input should be the exact name of the document the user is looking for.',
-        return_direct=False
+        name=tool.get("name", "") + "_name_lookup",
+        description="You must use this tool ONLY when the user is asking about a specific document by name or title. The input should be the exact name of the document the user is looking for.",
+        return_direct=False,
     )
 
 
@@ -181,7 +174,7 @@ def build_retrieval_tools(tool: dict, pred_args: dict, skill: db.Skills):
     try:
         rag_config = _load_rag_config(tool, pred_args, skill)
     except Exception as e:
-        logger.error(f"Error building RAG pipeline: {str(e)}")
+        logger.exception("Error building RAG pipeline:")
         raise ValueError(f"Failed to build RAG pipeline: {str(e)}")
     tools = [_build_rag_pipeline_tool(tool, pred_args, skill)]
     if rag_config.metadata_config is None:

mindsdb/interfaces/skills/skill_tool.py

@@ -274,8 +274,8 @@ class SkillToolController:
                     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)}")
+                except Exception:
+                    logger.warning(f"Could not get tables from database {database}:", exc_info=True)
                 continue
 
             # Handle table restrictions

mindsdb/interfaces/skills/sql_agent.py

@@ -1,7 +1,6 @@
 import re
 import csv
 import inspect
-import traceback
 from io import StringIO
 from typing import Iterable, List, Optional, Any, Tuple
 from collections import defaultdict
@@ -254,8 +253,16 @@ class SQLAgent:
                             self.check_table_permission(node)
                         except ValueError as origin_exc:
                             # was it badly quoted by llm?
-                            if len(node.parts) == 1 and node.is_quoted[0] and "." in node.parts[0]:
-                                node2 = Identifier(node.parts[0])
+                            #
+                            if "." in node.parts[0]:
+                                # extract quoted parts (with dots) to sub-parts
+                                parts = []
+                                for i, item in enumerate(node.parts):
+                                    if node.is_quoted[i] and "." in item:
+                                        parts.extend(Identifier(item).parts)
+                                    else:
+                                        parts.append(item)
+                                node2 = Identifier(parts=parts)
                                 try:
                                     _check_f(node2, is_table=True)
                                     return node2
@@ -382,9 +389,9 @@ class SQLAgent:
             #     self._cache.set(cache_key, set(kb_names))
 
             return kb_names
-        except Exception as e:
+        except Exception:
             # If there's an error, log it and return an empty list
-            logger.error(f"Error in get_usable_knowledge_base_names: {str(e)}")
+            logger.exception("Error in get_usable_knowledge_base_names")
             return []
 
     def _resolve_table_names(self, table_names: List[str], all_tables: List[Identifier]) -> List[Identifier]:
@@ -483,9 +490,9 @@ class SQLAgent:
                 # remove backticks
                 name = name.replace("`", "")
 
-                split = name.split(".")
-                if len(split) > 1:
-                    all_tables.append(Identifier(parts=[split[0], split[-1]]))
+                parts = name.split(".")
+                if len(parts) > 1:
+                    all_tables.append(Identifier(parts=parts))
                 else:
                     all_tables.append(Identifier(name))
 
@@ -526,8 +533,8 @@ class SQLAgent:
 
             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}")
+        except Exception:
+            logger.info("_get_sample_rows error:", exc_info=True)
             sample_rows_str = "\n" + "\t [error] Couldn't retrieve sample rows!"
 
         return sample_rows_str
@@ -560,7 +567,7 @@ class SQLAgent:
                 )
             ]
         except Exception as e:
-            logger.error(f"Failed processing column info for {table_str}: {e}", exc_info=True)
+            logger.exception(f"Failed processing column info for {table_str}:")
             raise ValueError(f"Failed to process column info for {table_str}") from e
 
         if not fields:
@@ -569,8 +576,8 @@ class SQLAgent:
 
         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}")
+        except Exception:
+            logger.warning(f"Could not get sample rows for {table_str}:", exc_info=True)
             sample_rows_info = "\n\t [error] Couldn't retrieve sample rows!"
 
         info = f"Table named `{table_str}`:\n"
@@ -585,7 +592,7 @@ class SQLAgent:
 
     def _get_sample_rows(self, table: str, fields: List[str]) -> str:
         logger.info(f"_get_sample_rows: table={table} fields={fields}")
-        command = f"select {', '.join(fields)} from {table} limit {self._sample_rows_in_table_info};"
+        command = f"select * from {table} limit {self._sample_rows_in_table_info};"
         try:
             ret = self._call_engine(command)
             sample_rows = ret.data.to_lists()
@@ -596,8 +603,8 @@ class SQLAgent:
 
             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}")
+        except Exception:
+            logger.info("_get_sample_rows error:", exc_info=True)
             sample_rows_str = "\n" + "\t [error] Couldn't retrieve sample rows!"
 
         return sample_rows_str
@@ -647,7 +654,7 @@ class SQLAgent:
             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("get_table_info_safe error:", exc_info=True)
             return f"Error: {e}"
 
     def query_safe(self, command: str, fetch: str = "all") -> str:
@@ -655,8 +662,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}")
+            logger.exception("Error in query_safe:")
             msg = f"Error: {e}"
             if "does not exist" in msg and " relation " in msg:
                 msg += "\nAvailable tables: " + ", ".join(self.get_usable_table_names())