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

mindsdb/interfaces/agents/langchain_agent.py

@@ -11,10 +11,8 @@ import pandas as pd
 from langchain.agents import AgentExecutor
 from langchain.agents.initialize import initialize_agent
 from langchain.chains.conversation.memory import ConversationSummaryBufferMemory
-from langchain_community.chat_models import (
-    ChatAnyscale,
-    ChatLiteLLM,
-    ChatOllama)
+from langchain_community.chat_models import ChatAnyscale, ChatLiteLLM, ChatOllama
+from langchain_writer import ChatWriter
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_core.agents import AgentAction, AgentStep
 from langchain_core.callbacks.base import BaseCallbackHandler
@@ -27,7 +25,9 @@ from langchain_core.tools import Tool
 from mindsdb.integrations.libs.llm.utils import get_llm_config
 from mindsdb.integrations.utilities.handler_utils import get_api_key
 from mindsdb.integrations.utilities.rag.settings import DEFAULT_RAG_PROMPT_TEMPLATE
-from mindsdb.interfaces.agents.event_dispatch_callback_handler import EventDispatchCallbackHandler
+from mindsdb.interfaces.agents.event_dispatch_callback_handler import (
+    EventDispatchCallbackHandler,
+)
 from mindsdb.interfaces.agents.constants import AGENT_CHUNK_POLLING_INTERVAL_SECONDS
 from mindsdb.utilities import log
 from mindsdb.utilities.context_executor import ContextThreadPoolExecutor
@@ -54,7 +54,10 @@ from mindsdb.interfaces.agents.constants import (
     NVIDIA_NIM_CHAT_MODELS,
     USER_COLUMN,
     ASSISTANT_COLUMN,
-    CONTEXT_COLUMN, TRACE_ID_COLUMN
+    CONTEXT_COLUMN,
+    TRACE_ID_COLUMN,
+    DEFAULT_AGENT_SYSTEM_PROMPT,
+    WRITER_CHAT_MODELS,
 )
 from mindsdb.interfaces.skills.skill_tool import skill_tool, SkillData
 from langchain_anthropic import ChatAnthropic
@@ -87,6 +90,9 @@ def get_llm_provider(args: Dict) -> str:
         return "nvidia_nim"
     if args["model_name"] in GOOGLE_GEMINI_CHAT_MODELS:
         return "google"
+    # Check for writer models
+    if args["model_name"] in WRITER_CHAT_MODELS:
+        return "writer"
 
     # For vLLM, require explicit provider specification
     raise ValueError("Invalid model name. Please define a supported llm provider")
@@ -100,21 +106,21 @@ def get_embedding_model_provider(args: Dict) -> str:
     # Check for explicit embedding model provider
     if "embedding_model_provider" in args:
         provider = args["embedding_model_provider"]
-        if provider == 'vllm':
-            if not (args.get('openai_api_base') and args.get('model')):
+        if provider == "vllm":
+            if not (args.get("openai_api_base") and args.get("model")):
                 raise ValueError(
                     "VLLM embeddings configuration error:\n"
                     "- Missing required parameters: 'openai_api_base' and/or 'model'\n"
                     "- Example: openai_api_base='http://localhost:8003/v1', model='your-model-name'"
                 )
             logger.info("Using custom VLLMEmbeddings class")
-            return 'vllm'
+            return "vllm"
         return provider
 
     # Check if LLM provider is vLLM
-    llm_provider = args.get('provider', DEFAULT_EMBEDDINGS_MODEL_PROVIDER)
-    if llm_provider == 'vllm':
-        if not (args.get('openai_api_base') and args.get('model')):
+    llm_provider = args.get("provider", DEFAULT_EMBEDDINGS_MODEL_PROVIDER)
+    if llm_provider == "vllm":
+        if not (args.get("openai_api_base") and args.get("model")):
             raise ValueError(
                 "VLLM embeddings configuration error:\n"
                 "- Missing required parameters: 'openai_api_base' and/or 'model'\n"
@@ -122,7 +128,7 @@ def get_embedding_model_provider(args: Dict) -> str:
                 "- Example: openai_api_base='http://localhost:8003/v1', model='your-model-name'"
             )
         logger.info("Using custom VLLMEmbeddings class")
-        return 'vllm'
+        return "vllm"
 
     # Default to LLM provider
     return llm_provider
@@ -131,14 +137,15 @@ def get_embedding_model_provider(args: Dict) -> str:
 def get_chat_model_params(args: Dict) -> Dict:
     model_config = args.copy()
     # Include API keys.
-    model_config["api_keys"] = {
-        p: get_api_key(p, model_config, None, strict=False) for p in SUPPORTED_PROVIDERS
-    }
-    llm_config = get_llm_config(
-        args.get("provider", get_llm_provider(args)), model_config
-    )
-    config_dict = llm_config.model_dump()
+    model_config["api_keys"] = {p: get_api_key(p, model_config, None, strict=False) for p in SUPPORTED_PROVIDERS}
+    llm_config = get_llm_config(args.get("provider", get_llm_provider(args)), model_config)
+    config_dict = llm_config.model_dump(by_alias=True)
     config_dict = {k: v for k, v in config_dict.items() if v is not None}
+
+    # If provider is writer, ensure the API key is passed as 'api_key'
+    if args.get("provider") == "writer" and "writer_api_key" in config_dict:
+        config_dict["api_key"] = config_dict.pop("writer_api_key")
+
     return config_dict
 
 
@@ -166,20 +173,26 @@ def create_chat_model(args: Dict):
         return ChatNVIDIA(**model_kwargs)
     if args["provider"] == "google":
         return ChatGoogleGenerativeAI(**model_kwargs)
+    if args["provider"] == "writer":
+        return ChatWriter(**model_kwargs)
     if args["provider"] == "mindsdb":
         return ChatMindsdb(**model_kwargs)
-    raise ValueError(f'Unknown provider: {args["provider"]}')
+    raise ValueError(f"Unknown provider: {args['provider']}")
 
 
 def prepare_prompts(df, base_template, input_variables, user_column=USER_COLUMN):
     empty_prompt_ids = np.where(df[input_variables].isna().all(axis=1).values)[0]
-    base_template = base_template.replace('{{', '{').replace('}}', '}')
+
+    # Combine system prompt with user-provided template
+    base_template = f"{DEFAULT_AGENT_SYSTEM_PROMPT}\n\n{base_template}"
+
+    base_template = base_template.replace("{{", "{").replace("}}", "}")
     prompts = []
 
     for i, row in df.iterrows():
         if i not in empty_prompt_ids:
             prompt = PromptTemplate(input_variables=input_variables, template=base_template)
-            kwargs = {col: row[col] if row[col] is not None else '' for col in input_variables}
+            kwargs = {col: row[col] if row[col] is not None else "" for col in input_variables}
             prompts.append(prompt.format(**kwargs))
         elif row.get(user_column):
             prompts.append(row[user_column])
@@ -194,8 +207,9 @@ def prepare_callbacks(self, args):
     return callbacks, context_callback
 
 
-def handle_agent_error(e):
-    error_message = f"An error occurred during agent execution: {str(e)}"
+def handle_agent_error(e, error_message=None):
+    if error_message is None:
+        error_message = f"An error occurred during agent execution: {str(e)}"
     logger.error(error_message, exc_info=True)
     return error_message
 
@@ -212,9 +226,7 @@ def process_chunk(chunk):
 
 
 class LangchainAgent:
-
     def __init__(self, agent: db.Agents, model: dict = None):
-
         self.agent = agent
         self.model = model
 
@@ -237,18 +249,14 @@ class LangchainAgent:
         args = self.agent.params.copy()
         args["model_name"] = self.agent.model_name
         args["provider"] = self.agent.provider
-        args["embedding_model_provider"] = args.get(
-            "embedding_model", get_embedding_model_provider(args)
-        )
+        args["embedding_model_provider"] = args.get("embedding_model", get_embedding_model_provider(args))
 
         # agent is using current langchain model
         if self.agent.provider == "mindsdb":
             args["model_name"] = self.agent.model_name
 
             # get prompt
-            prompt_template = (
-                self.model["problem_definition"].get("using", {}).get("prompt_template")
-            )
+            prompt_template = self.model["problem_definition"].get("using", {}).get("prompt_template")
             if prompt_template is not None:
                 # only update prompt_template if it is set on the model
                 args["prompt_template"] = prompt_template
@@ -257,24 +265,23 @@ class LangchainAgent:
             if args.get("mode") == "retrieval":
                 args["prompt_template"] = DEFAULT_RAG_PROMPT_TEMPLATE
             else:
-                raise ValueError(
-                    "Please provide a `prompt_template` or set `mode=retrieval`"
-                )
+                raise ValueError("Please provide a `prompt_template` or set `mode=retrieval`")
 
         return args
 
     def get_metadata(self) -> Dict:
         return {
-            'provider': self.provider,
-            'model_name': self.args["model_name"],
-            'embedding_model_provider': self.args.get('embedding_model_provider',
-                                                      get_embedding_model_provider(self.args)),
-            'skills': get_skills(self.agent),
-            'user_id': ctx.user_id,
-            'session_id': ctx.session_id,
-            'company_id': ctx.company_id,
-            'user_class': ctx.user_class,
-            'email_confirmed': ctx.email_confirmed
+            "provider": self.provider,
+            "model_name": self.args["model_name"],
+            "embedding_model_provider": self.args.get(
+                "embedding_model_provider", get_embedding_model_provider(self.args)
+            ),
+            "skills": get_skills(self.agent),
+            "user_id": ctx.user_id,
+            "session_id": ctx.session_id,
+            "company_id": ctx.company_id,
+            "user_class": ctx.user_class,
+            "email_confirmed": ctx.email_confirmed,
         }
 
     def get_tags(self) -> List:
@@ -283,14 +290,13 @@ class LangchainAgent:
         ]
 
     def get_completion(self, messages, stream: bool = False):
-
         # Get metadata and tags to be used in the trace
         metadata = self.get_metadata()
         tags = self.get_tags()
 
         # Set up trace for the API completion in Langfuse
         self.langfuse_client_wrapper.setup_trace(
-            name='api-completion',
+            name="api-completion",
             input=messages,
             tags=tags,
             metadata=metadata,
@@ -299,9 +305,7 @@ class LangchainAgent:
         )
 
         # Set up trace for the run completion in Langfuse
-        self.run_completion_span = self.langfuse_client_wrapper.start_span(
-            name='run-completion',
-            input=messages)
+        self.run_completion_span = self.langfuse_client_wrapper.start_span(name="run-completion", input=messages)
 
         if stream:
             return self._get_completion_stream(messages)
@@ -339,7 +343,7 @@ class LangchainAgent:
 
         df = pd.DataFrame(messages)
 
-        self.embedding_model_provider = args.get('embedding_model_provider', get_embedding_model_provider(args))
+        self.embedding_model_provider = args.get("embedding_model_provider", get_embedding_model_provider(args))
         # Back compatibility for old models
         self.provider = args.get("provider", get_llm_provider(args))
 
@@ -392,7 +396,7 @@ class LangchainAgent:
             agent=agent_type,
             # Use custom output parser to handle flaky LLMs that don't ALWAYS conform to output format.
             agent_kwargs={"output_parser": SafeOutputParser()},
-            # Calls the agent’s LLM Chain one final time to generate a final answer based on the previous steps
+            # Calls the agent's LLM Chain one final time to generate a final answer based on the previous steps
             early_stopping_method="generate",
             handle_parsing_errors=self._handle_parsing_errors,
             # Timeout per agent invocation.
@@ -400,11 +404,9 @@ class LangchainAgent:
                 "timeout_seconds",
                 args.get("timeout_seconds", DEFAULT_AGENT_TIMEOUT_SECONDS),
             ),
-            max_iterations=args.get(
-                "max_iterations", args.get("max_iterations", DEFAULT_MAX_ITERATIONS)
-            ),
+            max_iterations=args.get("max_iterations", args.get("max_iterations", DEFAULT_MAX_ITERATIONS)),
             memory=memory,
-            verbose=args.get("verbose", args.get("verbose", False))
+            verbose=args.get("verbose", args.get("verbose", False)),
         )
         return agent_executor
 
@@ -416,7 +418,7 @@ class LangchainAgent:
                 type=rel.skill.type,
                 params=rel.skill.params,
                 project_id=rel.skill.project_id,
-                agent_tables_list=(rel.parameters or {}).get('tables')
+                agent_tables_list=(rel.parameters or {}).get("tables"),
             )
             for rel in self.agent.skills_relationships
         ]
@@ -507,24 +509,34 @@ AI: {response}"""
         return f"Agent failed with error:\n{str(error)}..."
 
     def run_agent(self, df: pd.DataFrame, agent: AgentExecutor, args: Dict) -> pd.DataFrame:
-        base_template = args.get('prompt_template', args['prompt_template'])
-        return_context = args.get('return_context', True)
+        base_template = args.get("prompt_template", args["prompt_template"])
+        return_context = args.get("return_context", True)
         input_variables = re.findall(r"{{(.*?)}}", base_template)
 
-        prompts, empty_prompt_ids = prepare_prompts(df, base_template, input_variables,
-                                                    args.get('user_column', USER_COLUMN))
+        prompts, empty_prompt_ids = prepare_prompts(
+            df, base_template, input_variables, args.get("user_column", USER_COLUMN)
+        )
 
         def _invoke_agent_executor_with_prompt(agent_executor, prompt):
             if not prompt:
                 return {CONTEXT_COLUMN: [], ASSISTANT_COLUMN: ""}
             try:
                 callbacks, context_callback = prepare_callbacks(self, args)
-                result = agent_executor.invoke(prompt, config={'callbacks': callbacks})
+                result = agent_executor.invoke(prompt, config={"callbacks": callbacks})
                 captured_context = context_callback.get_contexts()
-                output = result['output'] if isinstance(result, dict) and 'output' in result else str(result)
+                output = result["output"] if isinstance(result, dict) and "output" in result else str(result)
                 return {CONTEXT_COLUMN: captured_context, ASSISTANT_COLUMN: output}
             except Exception as e:
-                return {CONTEXT_COLUMN: [], ASSISTANT_COLUMN: handle_agent_error(e)}
+                error_message = str(e)
+                # Special handling for API key errors
+                if "API key" in error_message and ("not found" in error_message or "missing" in error_message):
+                    # Format API key error more clearly
+                    logger.error(f"API Key Error: {error_message}")
+                    error_message = f"API Key Error: {error_message}"
+                return {
+                    CONTEXT_COLUMN: [],
+                    ASSISTANT_COLUMN: handle_agent_error(e, error_message),
+                }
 
         completions = []
         contexts = []
@@ -533,10 +545,7 @@ AI: {response}"""
         agent_timeout_seconds = args.get("timeout", DEFAULT_AGENT_TIMEOUT_SECONDS)
 
         with ContextThreadPoolExecutor(max_workers=max_workers) as executor:
-            futures = [
-                executor.submit(_invoke_agent_executor_with_prompt, agent, prompt)
-                for prompt in prompts
-            ]
+            futures = [executor.submit(_invoke_agent_executor_with_prompt, agent, prompt) for prompt in prompts]
             try:
                 for future in as_completed(futures, timeout=agent_timeout_seconds):
                     result = future.result()
@@ -550,9 +559,7 @@ AI: {response}"""
                     contexts.append(result[CONTEXT_COLUMN])
             except TimeoutError:
                 timeout_message = "I'm sorry! I couldn't come up with a response in time. Please try again."
-                logger.warning(
-                    f"Agent execution timed out after {agent_timeout_seconds} seconds"
-                )
+                logger.warning(f"Agent execution timed out after {agent_timeout_seconds} seconds")
                 for _ in range(len(futures) - len(completions)):
                     completions.append(timeout_message)
                     contexts.append([])
@@ -566,10 +573,8 @@ AI: {response}"""
         pred_df = pd.DataFrame(
             {
                 ASSISTANT_COLUMN: completions,
-                CONTEXT_COLUMN: [
-                    json.dumps(ctx) for ctx in contexts
-                ],  # Serialize context to JSON string
-                TRACE_ID_COLUMN: self.langfuse_client_wrapper.get_trace_id()
+                CONTEXT_COLUMN: [json.dumps(ctx) for ctx in contexts],  # Serialize context to JSON string
+                TRACE_ID_COLUMN: self.langfuse_client_wrapper.get_trace_id(),
             }
         )
 
@@ -579,17 +584,22 @@ AI: {response}"""
         return pred_df
 
     def add_chunk_metadata(self, chunk: Dict) -> Dict:
-        logger.debug(f'Adding metadata to chunk: {chunk}')
-        logger.debug(f'Trace ID: {self.langfuse_client_wrapper.get_trace_id()}')
+        logger.debug(f"Adding metadata to chunk: {chunk}")
+        logger.debug(f"Trace ID: {self.langfuse_client_wrapper.get_trace_id()}")
         chunk["trace_id"] = self.langfuse_client_wrapper.get_trace_id()
         return chunk
 
-    def _stream_agent_executor(self, agent_executor: AgentExecutor, prompt: str, callbacks: List[BaseCallbackHandler]):
+    def _stream_agent_executor(
+        self,
+        agent_executor: AgentExecutor,
+        prompt: str,
+        callbacks: List[BaseCallbackHandler],
+    ):
         chunk_queue = queue.Queue()
         # Add event dispatch callback handler only to streaming completions.
         event_dispatch_callback_handler = EventDispatchCallbackHandler(chunk_queue)
         callbacks.append(event_dispatch_callback_handler)
-        stream_iterator = agent_executor.stream(prompt, config={'callbacks': callbacks})
+        stream_iterator = agent_executor.stream(prompt, config={"callbacks": callbacks})
 
         agent_executor_finished_event = threading.Event()
 
@@ -604,7 +614,10 @@ AI: {response}"""
 
         # Enqueue Langchain agent streaming chunks in a separate thread to not block event chunks.
         executor_stream_thread = threading.Thread(
-            target=stream_worker, daemon=True, args=(ctx.dump(),), name='LangchainAgent.stream_worker'
+            target=stream_worker,
+            daemon=True,
+            args=(ctx.dump(),),
+            name="LangchainAgent.stream_worker",
         )
         executor_stream_thread.start()
 
@@ -613,24 +626,24 @@ AI: {response}"""
                 chunk = chunk_queue.get(block=True, timeout=AGENT_CHUNK_POLLING_INTERVAL_SECONDS)
             except queue.Empty:
                 continue
-            logger.debug(f'Processing streaming chunk {chunk}')
+            logger.debug(f"Processing streaming chunk {chunk}")
             processed_chunk = self.process_chunk(chunk)
-            logger.info(f'Processed chunk: {processed_chunk}')
+            logger.info(f"Processed chunk: {processed_chunk}")
             yield self.add_chunk_metadata(processed_chunk)
             chunk_queue.task_done()
 
     def stream_agent(self, df: pd.DataFrame, agent_executor: AgentExecutor, args: Dict) -> Iterable[Dict]:
-        base_template = args.get('prompt_template', args['prompt_template'])
+        base_template = args.get("prompt_template", args["prompt_template"])
         input_variables = re.findall(r"{{(.*?)}}", base_template)
-        return_context = args.get('return_context', True)
+        return_context = args.get("return_context", True)
 
-        prompts, _ = prepare_prompts(df, base_template, input_variables, args.get('user_column', USER_COLUMN))
+        prompts, _ = prepare_prompts(df, base_template, input_variables, args.get("user_column", USER_COLUMN))
 
         callbacks, context_callback = prepare_callbacks(self, args)
 
         yield self.add_chunk_metadata({"type": "start", "prompt": prompts[0]})
 
-        if not hasattr(agent_executor, 'stream') or not callable(agent_executor.stream):
+        if not hasattr(agent_executor, "stream") or not callable(agent_executor.stream):
             raise AttributeError("The agent_executor does not have a 'stream' method")
 
         stream_iterator = self._stream_agent_executor(agent_executor, prompts[0], callbacks)
@@ -659,21 +672,19 @@ AI: {response}"""
         if isinstance(chunk, AgentAction):
             # Format agent actions properly for streaming.
             return {
-                'tool': LangchainAgent.process_chunk(chunk.tool),
-                'tool_input': LangchainAgent.process_chunk(chunk.tool_input),
-                'log': LangchainAgent.process_chunk(chunk.log)
+                "tool": LangchainAgent.process_chunk(chunk.tool),
+                "tool_input": LangchainAgent.process_chunk(chunk.tool_input),
+                "log": LangchainAgent.process_chunk(chunk.log),
             }
         if isinstance(chunk, AgentStep):
             # Format agent steps properly for streaming.
             return {
-                'action': LangchainAgent.process_chunk(chunk.action),
-                'observation': LangchainAgent.process_chunk(chunk.observation) if chunk.observation else ''
+                "action": LangchainAgent.process_chunk(chunk.action),
+                "observation": LangchainAgent.process_chunk(chunk.observation) if chunk.observation else "",
             }
         if issubclass(chunk.__class__, BaseMessage):
             # Extract content from message subclasses properly for streaming.
-            return {
-                'content': chunk.content
-            }
+            return {"content": chunk.content}
         if isinstance(chunk, (str, int, float, bool, type(None))):
             return chunk
         return str(chunk)

mindsdb/interfaces/agents/mindsdb_database_agent.py

@@ -2,6 +2,7 @@
     Wrapper around MindsDB's executor and integration controller following the implementation of the original
     langchain.sql_database.SQLDatabase class to partly replicate its behavior.
 """
+import traceback
 from typing import Any, Iterable, List, Optional
 
 from mindsdb.utilities import log
@@ -57,8 +58,20 @@ class MindsDBSQL(SQLDatabase):
         """Information about all tables in the database."""
         return self._sql_agent.get_table_info()
 
+    @property
+    def kb_info(self) -> str:
+        """Information about all knowledge bases in the database."""
+        return self._sql_agent.get_knowledge_base_info()
+
     def get_usable_table_names(self) -> Iterable[str]:
-        return self._sql_agent.get_usable_table_names()
+        """Information about all tables in the database."""
+        try:
+            return self._sql_agent.get_usable_table_names()
+        except Exception as e:
+            # If there's an error accessing tables through the integration, handle it gracefully
+            logger.warning(f"Error getting usable table names: {str(e)}")
+            # Return an empty list instead of raising an exception
+            return []
 
     def get_table_info_no_throw(self, table_names: Optional[List[str]] = None) -> str:
         for i in range(len(table_names)):
@@ -66,5 +79,91 @@ class MindsDBSQL(SQLDatabase):
         return self._sql_agent.get_table_info_safe(table_names)
 
     def run_no_throw(self, command: str, fetch: str = "all") -> str:
+        """Execute a SQL command and return the result as a string.
+
+        This method catches any exceptions and returns an error message instead of raising an exception.
+
+        Args:
+            command: The SQL command to execute
+            fetch: Whether to fetch 'all' results or just 'one'
+
+        Returns:
+            A string representation of the result or an error message
+        """
         command = extract_essential(command)
-        return self._sql_agent.query_safe(command)
+
+        try:
+
+            # Log the query for debugging
+            logger.info(f"Executing SQL query: {command}")
+
+            # remove backticks
+            command = command.replace('`', '')
+
+            # Parse the SQL string to an AST object first
+            from mindsdb_sql_parser import parse_sql
+            ast_query = parse_sql(command)
+
+            # Now execute the parsed query
+            result = self._sql_agent.skill_tool.get_command_executor().execute_command(ast_query, database_name="mindsdb")
+
+            # Convert ExecuteAnswer to a DataFrame for easier manipulation
+            df = None
+            if hasattr(result, 'data') and hasattr(result.data, 'data_frame'):
+                df = result.data.data_frame
+            else:
+                # Fallback to to_df when data_frame attr not available
+                try:
+                    df = result.data.to_df()
+                except Exception:
+                    df = None
+
+            # Default behaviour (string)
+            if df is not None:
+                if not df.empty:
+                    return df.to_string(index=False)
+                else:
+                    return "Query executed successfully, but returned no data."
+
+            return str(result)
+
+        except Exception as e:
+            logger.error(f"Error executing SQL command: {str(e)}\n{traceback.format_exc()}")
+            # If this is a knowledge base query, provide a more helpful error message
+            if "knowledge_base" in command.lower() or any(kb in command for kb in self._sql_agent.get_usable_knowledge_base_names()):
+                return f"Error executing knowledge base query: {str(e)}. Please check that the knowledge base exists and your query syntax is correct."
+            return f"Error: {str(e)}"
+
+    # def run_no_throw(self, command: str, fetch: str = "all") -> str:
+    #     """Execute a SQL command and return the result as a string.
+    #
+    #     This method catches any exceptions and returns an error message instead of raising an exception.
+    #
+    #     Args:
+    #         command: The SQL command to execute
+    #         fetch: Whether to fetch 'all' results or just 'one'
+    #
+    #     Returns:
+    #         A string representation of the result or an error message
+    #     """
+    #     command = extract_essential(command)
+    #     try:
+    #         return self._sql_agent.query_safe(command)
+    #     except Exception as e:
+    #         logger.error(f"Error executing SQL command: {str(e)}")
+    #         # If this is a knowledge base query, provide a more helpful error message
+    #         if "knowledge_base" in command.lower() or any(kb in command for kb in self._sql_agent.get_usable_knowledge_base_names()):
+    #             return f"Error executing knowledge base query: {str(e)}. Please check that the knowledge base exists and your query syntax is correct."
+    #         return f"Error: {str(e)}"
+
+    def get_usable_knowledge_base_names(self) -> List[str]:
+        """Get a list of usable knowledge base names.
+
+        Returns:
+            A list of knowledge base names that can be used in queries
+        """
+        try:
+            return self._sql_agent.get_usable_knowledge_base_names()
+        except Exception as e:
+            logger.error(f"Error getting usable knowledge base names: {str(e)}")
+            return []