feat: add SQL table schema extractor and describer agent with auto-generation script

Author: Shrijeeth-Suresh

Makefile

@@ -13,6 +13,9 @@ format:
 upsert-data:
 	cd backend && python -m scripts.clear_vector_db && python -m scripts.add_sample_queries_qdrant && python -m scripts.add_table_schemas
 
+generate-data:
+	cd backend && python -m scripts.generate_tables_json
+
 run-backend:
 	cd backend && uvicorn main:app --port 8000
 

README.md

@@ -21,6 +21,11 @@ QueryGPT-ADK is an open-source, multi-agent system for natural language to SQL q
 
 - Natural Language to SQL: Converts user questions into SQL queries using LLM-based agents.
 - Multi-Agent Architecture: Modular agents for query formation, explanation, table/column selection, and validation.
+- **Automatic Schema Extraction & Table Description:**
+  - New agent (`sql_table_describer_agent`) for extracting and describing SQL table schemas, including sub-agents for table-level tasks.
+  - Includes a script to auto-generate JSON documentation of all tables, views, and enums from your database (`generate_tables_json.py`).
+- Robust SQL Query Validation: Ensures only valid SELECT statements are executed, supporting both MySQL and PostgreSQL.
+
 - Vector Search Integration: Uses Qdrant for semantic search over sample queries and table schemas.
 - Validation: Ensures only valid SELECT queries are generated and executed.
 - **Production-Ready Security:**
@@ -44,6 +49,12 @@ QueryGPT-ADK is an open-source, multi-agent system for natural language to SQL q
 
 ## Changelog
 
+### 2025-05-23
+- Added `sql_table_describer_agent` for automated SQL table schema extraction and description.
+- Introduced `backend/scripts/generate_tables_json.py` for generating up-to-date JSON documentation of all database tables, views, and enums.
+- Enhanced SQL query validation tools to support both MySQL and PostgreSQL, with improved error handling.
+- Updated configuration and helper utilities for multi-database and LLM support.
+
 ### 2025-05-20
 
 - **PostgreSQL Support**: Added support for PostgreSQL alongside MySQL for query validation.
@@ -92,7 +103,16 @@ POSTGRES_DB_URL=postgresql+asyncpg://user:password@localhost:5432/querygpt
 REDIS_URL=redis://localhost:6379/1
 ```
 
-4. Set up environment variables (Frontend):
+4. Generate up-to-date table and enum documentation (Backend):
+
+```bash
+cd backend
+python -m scripts.generate_tables_json
+```
+
+This command will connect to your configured database and auto-generate the latest tables, views, and enums documentation in `backend/scripts/data/tables.json`. Before running this command, make sure your database is up and running and the connection string in `.env` is correct for validation database. You also need LLM provider and API key for this command to work. Set `SYNTHETIC_DATA_LLM_MODEL` and `SYNTHETIC_DATA_LLM_API_KEY` in `.env` for this to work.
+
+5. Set up environment variables (Frontend):
 
 ```bash
 cd frontend
@@ -105,7 +125,7 @@ Edit `.env` and set your API base URL as needed. Example:
 NEXT_PUBLIC_API_BASE_URL=http://localhost:8000
 ```
 
-5. Set up the Frontend (Next.js UI):
+6. Set up the Frontend (Next.js UI):
 
 ```bash
 cd frontend
@@ -116,22 +136,22 @@ npm run dev
 
 This will start the frontend at [http://localhost:3000](http://localhost:3000).
 
-6. Clear Qdrant database (Optional):
+7. Clear Qdrant database (Optional):
 
 ```bash
 cd backend
 python scripts/clear_vector_db.py
 ```
 
-7. Load db data into Qdrant:
+8. Load db data into Qdrant:
 
 ```bash
 cd backend
 python scripts/add_table_schemas.py
 python scripts/add_sample_queries_qdrant.py
 ```
 
-8. (Optional) Use Your Own Data:
+9. (Optional) Use Your Own Data:
 
 To use QueryGPT-ADK with your own data, replace the provided `sample_queries.json` and `tables.json` files in the `backend/scripts/data/` directory with your own files. Make sure your files follow the same structure as the samples. Then, rerun the data loading scripts:
 
@@ -388,7 +408,7 @@ curl -X POST "http://localhost:8000/query" \
 - Enable schema and collection import for semantic search and retrieval.
 - Ensure integration of custom vector DBs with agentic workflow and query tools.
 
-### Generate Synthetic Data for a Database (To Do)
+### Generate Synthetic Data for a Database (In Progress)
 
 - Add support for generating synthetic data for a given database.
 - Provide a UI and/or CLI for users to generate synthetic data for their own database.

backend/agents/query_agent/sub_agents/query_validation_agent/tools.py

@@ -1,58 +1,39 @@
-import mysql.connector
-import psycopg2
 from config import get_settings
 from google.adk.tools import FunctionTool
 from mysql.connector import Error as MySQLError
 from psycopg2 import Error as PostgresError
+from utils.helpers import (
+    get_validate_db_mysql_connection,
+    get_validate_db_postgres_connection,
+)
 
 
-def validate_sql_query_mysql(query: str) -> dict:
-    conn = mysql.connector.connect(
-        host=get_settings().VALIDATE_DB_HOST,
-        port=get_settings().VALIDATE_DB_PORT,
-        user=get_settings().VALIDATE_DB_USER,
-        password=get_settings().VALIDATE_DB_PASSWORD,
-        database=get_settings().VALIDATE_DB_DATABASE,
-        autocommit=False,
-    )
-    cursor = conn.cursor()
+def validate_sql_query_mysql(query: str) -> tuple[dict, object]:
+    conn, cursor = get_validate_db_mysql_connection()
     query_type = query.strip().split()[0].upper()
 
     if query_type == "SELECT":
         try:
             cursor.execute(f"EXPLAIN {query}")
-            conn.close()
-            return {"valid": "true", "error": ""}
+            return {"valid": "true", "error": ""}, conn
         except MySQLError as e:
-            conn.close()
-            return {"valid": "false", "error": str(e)}
+            return {"valid": "false", "error": str(e)}, conn
     else:
-        conn.close()
-        return {"valid": "false", "error": "Only SELECT statements are allowed"}
+        return {"valid": "false", "error": "Only SELECT statements are allowed"}, conn
 
 
-def validate_sql_query_postgres(query: str) -> dict:
-    conn = psycopg2.connect(
-        host=get_settings().VALIDATE_DB_HOST,
-        port=get_settings().VALIDATE_DB_PORT,
-        user=get_settings().VALIDATE_DB_USER,
-        password=get_settings().VALIDATE_DB_PASSWORD,
-        database=get_settings().VALIDATE_DB_DATABASE,
-    )
-    cursor = conn.cursor()
+def validate_sql_query_postgres(query: str) -> tuple[dict, object]:
+    conn, cursor = get_validate_db_postgres_connection()
     query_type = query.strip().split()[0].upper()
 
     if query_type == "SELECT":
         try:
             cursor.execute(f"EXPLAIN {query}")
-            conn.close()
-            return {"valid": "true", "error": ""}
+            return {"valid": "true", "error": ""}, conn
         except PostgresError as e:
-            conn.close()
-            return {"valid": "false", "error": str(e)}
+            return {"valid": "false", "error": str(e)}, conn
     else:
-        conn.close()
-        return {"valid": "false", "error": "Only SELECT statements are allowed"}
+        return {"valid": "false", "error": "Only SELECT statements are allowed"}, conn
 
 
 def validate_sql_query(query: str) -> dict:
@@ -67,15 +48,20 @@ def validate_sql_query(query: str) -> dict:
     Returns:
         dict: A dictionary containing the validation result (valid: bool) and error message (error: str)
     """
+    conn = None
     try:
         if get_settings().VALIDATE_DB_TYPE == "mysql":
-            return validate_sql_query_mysql(query)
+            result, conn = validate_sql_query_mysql(query)
         elif get_settings().VALIDATE_DB_TYPE == "postgresql":
-            return validate_sql_query_postgres(query)
+            result, conn = validate_sql_query_postgres(query)
         else:
             return {"valid": "false", "error": "Unsupported database type"}
+        return result
     except Exception as e:
         return {"valid": "false", "error": str(e)}
+    finally:
+        if conn:
+            conn.close()
 
 
 validate_sql_query_tool = FunctionTool(

backend/agents/sql_table_describer_agent/__init__.py

@@ -0,0 +1,3 @@
+from . import agent
+
+__all__ = ["agent"]

backend/agents/sql_table_describer_agent/agent.py

@@ -0,0 +1,12 @@
+from google.adk.agents import LoopAgent
+
+from .sub_agents import table_count_check_agent, table_description_agent
+
+root_agent = LoopAgent(
+    name="sql_table_describer_agent",
+    sub_agents=[
+        table_description_agent,
+        table_count_check_agent,
+    ],
+    max_iterations=3,
+)

backend/agents/sql_table_describer_agent/sub_agents/__init__.py

@@ -0,0 +1,4 @@
+from .table_count_check_agent.agent import root_agent as table_count_check_agent
+from .table_description_agent.agent import root_agent as table_description_agent
+
+__all__ = ["table_count_check_agent", "table_description_agent"]

backend/agents/sql_table_describer_agent/sub_agents/table_count_check_agent/__init__.py

@@ -0,0 +1,34 @@
+from typing import AsyncGenerator
+
+from google.adk.agents import BaseAgent
+from google.adk.agents.invocation_context import InvocationContext
+from google.adk.events import Event, EventActions
+
+
+class ValidateTableCount(BaseAgent):
+    async def _run_async_impl(
+        self, ctx: InvocationContext
+    ) -> AsyncGenerator[Event, None]:
+        try:
+            table_descriptions = ctx.session.state.get("sql_table_description", "")
+            input_query = ctx.session.events[-1].content.parts[0].text
+            input_tables_count = len(input_query.strip().split(";"))
+            table_descriptions_count = len(table_descriptions["table_descriptions"])
+            should_stop = input_tables_count == table_descriptions_count
+            yield Event(
+                author=self.name,
+                actions=EventActions(
+                    escalate=should_stop,
+                ),
+            )
+        except Exception as e:
+            print("Exception: ", e)
+            yield Event(
+                author=self.name,
+                actions=EventActions(
+                    escalate=False,
+                ),
+            )
+
+
+root_agent = ValidateTableCount(name="validate_table_count")

backend/agents/sql_table_describer_agent/sub_agents/table_count_check_agent/agent.py

@@ -0,0 +1,26 @@
+from config import get_settings
+from google.adk.agents import LlmAgent
+from google.adk.models.lite_llm import LiteLlm
+from google.genai import types
+
+from .prompts import agent_description, agent_instruction
+from .types import TableDescriptionSchema
+
+AGENT_MODEL = LiteLlm(
+    model=get_settings().SYNTHETIC_DATA_LLM_MODEL,
+    api_key=get_settings().SYNTHETIC_DATA_LLM_API_KEY,
+    response_format=TableDescriptionSchema,
+)
+
+
+root_agent = LlmAgent(
+    name="sql_table_describer_agent",
+    model=AGENT_MODEL,
+    description=agent_description(version=1),
+    instruction=agent_instruction(version=1),
+    output_key="sql_table_description",
+    generate_content_config=types.GenerateContentConfig(
+        temperature=get_settings().SYNTHETIC_DATA_LLM_TEMPERATURE,
+    ),
+    output_schema=TableDescriptionSchema,
+)