Enhance database agent with schema and query logging features

- Introduced QuestionMeta model to encapsulate question, SQL, and prompt.
- Updated _parse_schema_to_knowledge to save DDL and DML records in JSON format.
- Modified question_to_sql to return QuestionMeta instead of raw SQL.
- Improved error handling in database manager to raise SQLExecutionError.
- Refactored prompt generation for question inference to use predefined templates.

Author: coolbeevip

README.md

@@ -29,7 +29,7 @@ cd camel-database-agent
 pip install uv ruff mypy
 uv venv .venv --python=3.10
 source .venv/bin/activate
-uv pip install -e ".[dev,test]"
+uv sync --all-extras
 ````
 
 #### Music Database
@@ -141,34 +141,8 @@ Run the Spider 2.0-Lite evaluation.
 
 ```shell
 cd spider2_lite
-export API_KEY=sk-xx
+export OPENAI_API_KEY=sk-xxx
+export OPENAI_API_BASE_URL=https://api.openai.com/v1/
+export MODEL_NAME=gpt-4o-mini
 python spider2_run.py
-```
-
-## Development
-
-Install the development dependencies.
-```shell
-pip install uv ruff mypy
-uv pip install -e ".[dev]"
-```
-
-Run code formatters
-```shell
-make format
-```
-
-Run code linters
-```shell
-make lint
-```
-
-Run unit tests
-```shell
-make test
-```
-
-Create a uv.lock file from pyproject.toml
-```shell
-uv pip compile pyproject.toml -o uv.lock --resolution=highest
 ```

camel_database_agent/database/database_manager.py

@@ -61,7 +61,7 @@ def __init__(self, db_url: str, read_only_model: bool = True):
     @with_session
     def select(
         self, session: Session, sql: str, bind_pd: bool = False
-    ) -> Union[List[dict], pd.DataFrame, SQLExecutionError]:
+    ) -> Union[List[dict], pd.DataFrame]:
         """Execute Query SQL"""
         self._check_sql(sql)
         try:
@@ -74,7 +74,7 @@ def select(
                 rows = [dict(zip(column_names, row)) for row in result]
                 return rows
         except OperationalError as e:
-            return SQLExecutionError(sql, str(e))
+            raise SQLExecutionError(sql, str(e))
 
     @with_session
     def execute(

camel_database_agent/database/database_schema_parse.py

@@ -55,12 +55,12 @@ def __init__(
     def parse_ddl_record(self, text: str) -> List[DDLRecord]:
         """Parsing DDL SQL statements"""
         prompt = (
-            "Here are some DDL statements from which you need to "
-            "refer to table names, field names, data types, default "
-            "values, etc., to generate summary information and extract "
-            "the SQL statements for each table.\n\n"
+            "The following are some DDL script. Please read the script in its "
+            "entirety and provide descriptions for the tables and fields to "
+            "generate summary information and extract the SQL script for each "
+            "table.\n\n"
         )
-        prompt += f"```sql\n{text}```\n"
+        prompt += f"```sql\n{text}```\n\n"
         prompt += "Please output the summary information and SQL script in JSON format."
         response = self.parsing_agent.step(prompt, response_format=DDLRecordResponseFormat)
         ddl_record_response = DDLRecordResponseFormat.model_validate_json(response.msgs[0].content)

camel_database_agent/database/dialect/database_schema_dialect.py

@@ -6,6 +6,7 @@
 from camel.models import BaseModelBackend
 
 from camel_database_agent.database.database_manager import DatabaseManager
+from camel_database_agent.database_prompt import POLISH_SCHEMA_OUTPUT_EXAMPLE
 
 logger = logging.getLogger(__name__)
 
@@ -50,16 +51,9 @@ def get_dialect(
 
     def get_polished_schema(self, language: str = "English") -> str:
         if self.schema_polish_agent:
-            prompt = (
-                f"Please optimize the SQL schema of the database in {language}, "
-                f"ensuring it includes table name comments, field comments, "
-                f"foreign key explanations, etc., to make it more readable.\n\n"
-            )
-            prompt += f"```sql\n{self.get_schema()}```\n\n"
-            prompt += (
-                "Now, please directly output the optimized SQL Schema. "
-                "Do not explain the process and optimization ideas."
-            )
+            prompt = POLISH_SCHEMA_OUTPUT_EXAMPLE.replace(
+                "{{ddl_sql}}", self.get_schema()
+            ).replace("{{language}}", language)
             response = self.schema_polish_agent.step(prompt)
             return response.msgs[0].content
         else:

camel_database_agent/database_agent.py

@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import random
@@ -26,13 +27,15 @@
     HumanMessage,
     MessageLog,
     MessageLogToEmpty,
+    SQLExecutionError,
     TrainLevel,
     messages_log,
     strip_sql_code_block,
     timing,
 )
 from camel_database_agent.database_prompt import (
     DATABASE_SUMMARY_OUTPUT_EXAMPLE,
+    QUESTION_CONVERT_SQL,
 )
 from camel_database_agent.datagen.sql_query_inference_pipeline import (
     DataQueryInferencePipeline,
@@ -45,6 +48,12 @@
 logger = logging.getLogger(__name__)
 
 
+class QuestionMeta(BaseModel):
+    question: str
+    sql: str
+    prompt: str
+
+
 class DatabaseAgentResponse(BaseModel):
     ask: str
     dataset: Optional[Any] = None
@@ -167,20 +176,41 @@ def __init__(
                 self.data_sql = f.read()
 
     @timing
-    def _parse_schema_to_knowledge(self, polish: bool = True) -> None:
+    def _parse_schema_to_knowledge(self, polish: bool = False) -> None:
         """Generate schema data to knowledge"""
         self.ddl_sql = (
             self.dialect.get_polished_schema(self.language)
             if polish
             else self.dialect.get_schema()
         )
+        # Save the schema to a file
+        with open(
+            os.path.join(self.knowledge_path, "ddl_origin.sql"),
+            "w",
+            encoding="utf-8",
+        ) as f:
+            f.write(self.dialect.get_schema())
+
+        # Save the polished schema to a file
         with open(
             os.path.join(self.knowledge_path, "ddl_sql.sql"),
             "w",
             encoding="utf-8",
         ) as f:
             f.write(self.ddl_sql)
+
         ddl_records: List[DDLRecord] = self.schema_parse.parse_ddl_record(self.ddl_sql)
+        with open(
+            os.path.join(self.knowledge_path, "ddl_records.json"),
+            "w",
+            encoding="utf-8",
+        ) as f:
+            f.write(
+                json.dumps(
+                    [record.model_dump() for record in ddl_records], ensure_ascii=False, indent=4
+                )
+            )
+
         self.database_knowledge_backend.add(ddl_records)
 
     @timing
@@ -194,6 +224,18 @@ def _parse_sampled_data_to_knowledge(self, data_samples_size: int = 5) -> None:
         ) as f:
             f.write(self.data_sql)
         dml_records: List[DMLRecord] = self.schema_parse.parse_dml_record(self.data_sql)
+
+        with open(
+            os.path.join(self.knowledge_path, "data_records.json"),
+            "w",
+            encoding="utf-8",
+        ) as f:
+            f.write(
+                json.dumps(
+                    [record.model_dump() for record in dml_records], ensure_ascii=False, indent=4
+                )
+            )
+
         self.database_knowledge_backend.add(dml_records)
 
     @timing
@@ -210,6 +252,15 @@ def _parse_query_to_knowledge(self, query_samples_size: int = 20) -> None:
             query_records: List[QueryRecord] = []
             while len(query_records) < query_samples_size:
                 query_records.extend(pipeline.generate(query_samples_size=query_samples_size))
+
+            with open(
+                os.path.join(self.knowledge_path, "question_sql.txt"),
+                "w",
+                encoding="utf-8",
+            ) as f:
+                for query_record in query_records:
+                    f.write(f"QUESTION: {query_record.question}\nSQL: {query_record.sql}\n\n")
+
             self.database_knowledge_backend.add(query_records)
         else:
             raise ValueError("ddl_sql and data_sql must be provided")
@@ -292,6 +343,10 @@ def train_knowledge(
 
         if reset_train and os.path.exists(self.knowledge_path):
             self.database_knowledge_backend.clear()
+            self.ddl_sql = None
+            self.data_sql = None
+            self.database_summary = ""
+            self.recommendation_question = ""
             logger.info("Reset knowledge...")
 
         if (
@@ -319,41 +374,36 @@ def train_knowledge(
             self.generate_database_summary(query_samples_size=query_samples_size)
 
     @timing
-    def question_to_sql(self, question: str, dialect_name: str) -> str:
+    def question_to_sql(self, question: str, dialect_name: str) -> QuestionMeta:
         """Question to SQL"""
-        prompt = (
-            f"The following is the table structure in the database and "
-            f"some common query SQL statements. Please convert the user's "
-            f"question into an SQL query statement. Note to comply "
-            f"with {dialect_name} syntax. Do not explain, "
-            f"just provide the SQL directly.\n\n"
-        )
-        prompt += "## Table Schema\n"
+        prompt = QUESTION_CONVERT_SQL.replace("{{dialect_name}}", dialect_name)
+
         ddl_records: List[DDLRecord] = self.database_knowledge_backend.query_ddl(question)
-        prompt += "```sql\n"
-        for ddl_record in ddl_records:
-            prompt += f"{ddl_record.sql}\n"
-        prompt += "```\n\n"
+        prompt = prompt.replace(
+            "{{table_schema}}", "\n".join([record.sql for record in ddl_records])
+        )
 
-        prompt += "## Data Example\n"
-        prompt += "```sql\n"
         data_records: List[DMLRecord] = self.database_knowledge_backend.query_data(question)
-        for data_record in data_records:
-            prompt += f"```{data_record.sql}\n"
-        prompt += "```\n\n"
+        prompt = prompt.replace(
+            "{{sample_data}}", "\n".join([record.sql for record in data_records])
+        )
 
-        # some few shot
         query_records: List[QueryRecord] = self.database_knowledge_backend.query_query(question)
-        for query_record in query_records:
-            prompt += f"Question: {query_record.question}\n"
-            prompt += f"SQL: {query_record.sql}\n\n"
+        prompt = prompt.replace(
+            "{{qa_pairs}}",
+            "\n".join(
+                [f"QUESTION: {record.question}\nSQL: {record.sql}\n\n" for record in query_records]
+            ),
+        )
 
-        prompt += f"Question: {question}\n"
-        prompt += "SQL: "
-        logger.debug(Fore.GREEN + "PROMPT:", prompt)
+        prompt = prompt.replace("{{question}}", question)
+        logger.debug(Fore.GREEN + "PROMPT:" + prompt)
         self.agent.reset()
         response = self.agent.step(prompt)
-        return strip_sql_code_block(response.msgs[0].content)
+
+        return QuestionMeta(
+            question=question, sql=strip_sql_code_block(response.msgs[0].content), prompt=prompt
+        )
 
     @messages_log
     def ask(
@@ -366,27 +416,37 @@ def ask(
         if not message_log:
             message_log = MessageLogToEmpty()
         message_log.messages_writer(HumanMessage(session_id=session_id, content=question))
-        sql = self.question_to_sql(
+        question_meta = self.question_to_sql(
             question=question,
             dialect_name=self.database_manager.dialect_name(),
         )
-        message_log.messages_writer(AssistantMessage(session_id=session_id, content=sql))
         try:
-            dataset = self.database_manager.select(sql=sql, bind_pd=bind_pd)
+            message_log.messages_writer(
+                AssistantMessage(session_id=session_id, content=question_meta.sql)
+            )
+            dataset = self.database_manager.select(sql=question_meta.sql, bind_pd=bind_pd)
             message_log.messages_writer(
                 AssistantMessage(
                     session_id=session_id,
                     content=tabulate(dataset, headers="keys", tablefmt="psql"),
                 )
             )
-            return DatabaseAgentResponse(ask=question, dataset=dataset, sql=sql)
+            return DatabaseAgentResponse(ask=question, dataset=dataset, sql=question_meta.sql)
+        except SQLExecutionError as e:
+            message_log.messages_writer(AssistantMessage(session_id=session_id, content=str(e)))
+            return DatabaseAgentResponse(
+                ask=question,
+                dataset=None,
+                sql=e.sql,
+                success=False,
+                error=e.error_message,
+            )
         except Exception as e:
-            logger.error(e)
             message_log.messages_writer(AssistantMessage(session_id=session_id, content=str(e)))
             return DatabaseAgentResponse(
                 ask=question,
                 dataset=None,
-                sql=sql,
+                sql=question_meta.sql,
                 success=False,
                 error=str(e),
             )

camel_database_agent/database_base.py

@@ -102,12 +102,12 @@ def timing_wrapper(*args: Any, **kwargs: Any) -> Any:
         try:
             spinner_thread.start()
             result = func(*args, **kwargs)
-            end_time = time.perf_counter()
-            total_time = end_time - start_time
         finally:
             # sys.stdout.write('\r' + ' ' * 100 + '\r')
             stop_spinner.set()
             spinner_thread.join()
+            end_time = time.perf_counter()
+            total_time = end_time - start_time
             logger.info(f"\r{info} Took {Fore.GREEN}{total_time:.4f} seconds{Fore.RESET}")
         return result