Fix agent work with sql skill

mindsdb/api/executor/datahub/datanodes/project_datanode.py

@@ -132,7 +132,10 @@ def query(self, query=None, native_query=None, session=None) -> DataHubResponse:
                 case [query_table, str(version)], [is_quoted, _] if version.isdigit():
                     ...
                 case _:
-                    raise ValueError("Table name should contain only one part")
+                    raise EntityNotExistsError(
+                        f"Table '{query.from_table}' not found in the database. The project database support only single-part names",
+                        self.project.name,
+                    )
 
             if not is_quoted:
                 query_table = query_table.lower()

mindsdb/interfaces/agents/constants.py

@@ -260,8 +260,6 @@
 
 For factual questions, ALWAYS use the available tools to look up information rather than relying on your internal knowledge.
 
-Here is the user's question: {{question}}   
-
 TOOLS:
 ------
 

mindsdb/interfaces/skills/custom/text2sql/mindsdb_sql_toolkit.py

@@ -55,7 +55,7 @@ def get_tools(self, prefix="") -> List[BaseTool]:
             f"""\
             Input: A detailed and well-structured SQL query. The query must be enclosed between the symbols $START$ and $STOP$.
             Output: Database result or error message. For errors, rewrite and retry the query. For 'Unknown column' errors, use '{info_sql_database_tool.name}' to check table fields.
-            This system is a highly intelligent and reliable PostgreSQL SQL skill designed to work with databases.
+            This system is a highly intelligent and reliable SQL skill designed to work with databases.
             Follow these instructions with utmost precision:
             1. Final Response Format:
                - Assume the frontend fully supports Markdown unless the user specifies otherwise.
@@ -73,7 +73,7 @@ def get_tools(self, prefix="") -> List[BaseTool]:
                - Let the user know they can request additional results and/or specify how they would like the results ordered or grouped.
             5. Date Handling:
                - **System current date and time: {current_date_time} (UTC or local timezone based on server settings).**
-               - **Always** use PostgreSQL-compatible `CURRENT_DATE` or `NOW()` functions when working with dates—never assume or guess the current date.
+               - **Always** use `CURRENT_DATE` or `NOW()` functions when working with dates—never assume or guess the current date.
                - For any date-related comparisons in the query, *always* ensure that your query casts the column being compared using `column_name::DATE [operator] ..`
                - Do not compare date values without casting columns to date.
                - For date interval operations, use Interval units as keywords. You can use keywords to specify units like days, hours, months, years, etc., directly without quotes. Examples:
@@ -95,6 +95,8 @@ def get_tools(self, prefix="") -> List[BaseTool]:
             8. Identity and Purpose:
                - When asked about yourself or your maker, state that you are a Data-Mind, created by MindsDB to help answer data questions.
                - When asked about your purpose or how you can help, explore the available data sources and then explain that you can answer questions based on the connected data. Provide a few relevant example questions that you could answer for the user about their data.
+            9. Important: you can use only mysql quoting rules to compose queries: backticks (`) for identifiers, and single quotes (') for constants
+
             Adhere to these guidelines for all queries and responses. Ask for clarification if needed.
         """
         )
@@ -110,15 +112,15 @@ def get_tools(self, prefix="") -> List[BaseTool]:
             "If the query is correct, it will be parsed and returned. "
             f"ALWAYS run this tool before executing a query with {query_sql_database_tool.name}. "
         )
-        mindsdb_sql_parser_tool = MindsDBSQLParserTool(
+        mindsdb_sql_parser_tool = MindsDBSQLParserTool(  # noqa: F841
             name=f"mindsdb_sql_parser_tool{prefix}", description=mindsdb_sql_parser_tool_description
         )
 
         sql_tools = [
             query_sql_database_tool,
             info_sql_database_tool,
             list_sql_database_tool,
-            mindsdb_sql_parser_tool,
+            # mindsdb_sql_parser_tool,
         ]
         if not self.include_knowledge_base_tools:
             return sql_tools

mindsdb/interfaces/skills/sql_agent.py

@@ -254,8 +254,16 @@ def _check_f(node, is_table=None, **kwargs):
                             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
@@ -483,9 +491,9 @@ def get_table_info(self, table_names: Optional[List[str]] = None) -> str:
                 # 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))
 
@@ -585,7 +593,7 @@ def _get_single_table_info(self, table: Identifier) -> str:
 
     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()

tests/unit/executor/test_agent.py

@@ -917,15 +917,15 @@ def test_agent_accept_wrong_quoting(self, mock_litellm_embedding, mock_openai):
               openai_api_key='--',
               data = {
                  "knowledge_bases": ["kb1"],
-                 "tables": ["files.file1"]
+                 "tables": ["files.file1", "files.file2.*"]
               }
          """)
         self.run_sql("""
             insert into kb1
             select id, planet_name content from files.file1
         """)
 
-        # exposed
+        # # exposed
         set_openai_completion(
             mock_openai,
             [
@@ -939,6 +939,44 @@ def test_agent_accept_wrong_quoting(self, mock_litellm_embedding, mock_openai):
         assert "Jupiter" in mock_openai.agent_calls[1]
         assert "Jupiter" in mock_openai.agent_calls[2]
 
+    @patch("openai.OpenAI")
+    @patch("mindsdb.integrations.handlers.postgres_handler.Handler")
+    def test_3_part_table(self, mock_pg, mock_openai):
+        df = get_dataset_planets()
+        self.set_handler(mock_pg, name="pg", tables={"planets": df}, schema="public")
+
+        self.run_sql("""
+            CREATE AGENT my_agent
+            USING
+              model = "gpt-3.5-turbo",
+              openai_api_key='--',
+              data = {
+                 "tables": ["pg.public.*"] 
+              }
+        """)
+
+        set_openai_completion(
+            mock_openai,
+            [
+                # test getting table info
+                self._action("sql_db_schema", "pg.public.planets"),
+                # test wrong quoting
+                self._action("sql_db_query", "SELECT * FROM pg.public.planets WHERE id = '1000'"),
+                self._action("sql_db_query", "SELECT * FROM `pg.public`.planets WHERE id = '1000'"),
+                self._action("sql_db_query", "SELECT * FROM `pg.public`.`planets` WHERE id = '1000'"),
+                "Hi!",
+            ],
+        )
+        self.run_sql("select * from my_agent where question = 'test'")
+
+        # result of sql_db_schema
+        assert "planets" in mock_openai.agent_calls[1]  # table name
+        assert "Jupiter" in mock_openai.agent_calls[1]  # column
+        # results of sql_db_query
+        assert "Moon" in mock_openai.agent_calls[2]
+        assert "Moon" in mock_openai.agent_calls[3]
+        assert "Moon" in mock_openai.agent_calls[4]
+
     @patch("openai.OpenAI")
     @patch(
         "mindsdb.interfaces.agents.langchain_agent.LangchainAgent.run_agent",