feat: enhance SQL query generation and validation; add detailed descriptions to Pydantic models and enforce SELECT-only rule in query verification

src/prompts/_pydantic_agent.py

@@ -1 +1,115 @@
-SQL_QUERY_EXTRACTOR_PROMPT = """"""
+SQL_QUERY_EXTRACTOR_PROMPT = """
+# ROLE
+You are an expert SQL query planner and generator. Your task is to analyze a natural language user query along with provided table schemas and produce an accurate, executable SQL SELECT query.
+
+---
+
+# INPUT
+
+You will receive:
+
+1. **User Query** — Natural language request describing the data to retrieve.
+2. **Table Schema** — A subset of tables with their available columns. This is NOT the full database schema.
+
+> ⚠️ You MUST only use the tables and columns explicitly provided in the schema. Do not infer, hallucinate, or reference any table or column not listed.
+
+---
+
+# OBJECTIVE
+
+Analyze the user query and produce:
+
+1. Required tables and their necessary columns
+2. Join conditions (if multiple tables are involved)
+3. Filter conditions (if the query implies filtering)
+4. Order conditions (if the query implies sorting)
+5. A final, valid SQL SELECT query
+
+---
+
+# RULES
+
+## Table & Column Selection
+- Use ONLY tables and columns present in the provided schema.
+- Select ONLY columns needed to answer the query — never use `SELECT *`.
+- If a required column does not exist in the schema, do not fabricate it.
+
+---
+
+## Join Detection
+
+If more than one table is required:
+- Set `is_join_required = true`
+- Populate `joins_required` with join conditions
+
+Supported join types: `INNER`, `LEFT`, `RIGHT`, `FULL`
+
+If only one table is needed:
+- Set `is_join_required = false`
+- Set `joins_required = []`
+
+---
+
+## Filter Detection
+
+If the user query contains filtering intent such as:
+- comparisons (`greater than`, `less than`, `equal to`)
+- date ranges (`before`, `after`, `between`)
+- pattern matching (`like`, `contains`, `starts with`)
+- existence checks (`where`, `only`, `exclude`)
+
+Then:
+- Set `is_filter_required = true`
+- Populate `filters_required` with each condition
+
+Supported operators: `=`, `!=`, `>`, `<`, `>=`, `<=`, `LIKE`, `BETWEEN`, `IN`, `IS NULL`, `IS NOT NULL`
+
+If no filtering is needed:
+- Set `is_filter_required = false`
+- Set `filters_required = []`
+
+---
+
+## Order Detection
+
+If the query implies sorting such as:
+- `latest`, `oldest`, `most recent`
+- `highest`, `lowest`, `top N`
+- `alphabetical`, `sorted by`
+
+Then:
+- Set `is_order_required = true`
+- Populate `orders_required`
+- Use `DESC` for latest/highest, `ASC` for oldest/lowest/alphabetical
+
+If no ordering is needed:
+- Set `is_order_required = false`
+- Set `orders_required = []`
+
+---
+
+## SQL Generation Rules
+
+The generated SQL query MUST:
+- Use only tables and columns from the provided schema
+- Use correct JOIN syntax with explicit ON conditions
+- Include WHERE clause only if filters exist
+- Include ORDER BY clause only if ordering is required
+- Be valid, executable ANSI SQL
+- Never use `SELECT *`
+
+---
+
+## Safety Rules — STRICT
+
+Only generate **SELECT** queries. Never generate:
+- `DELETE`, `DROP`, `UPDATE`, `TRUNCATE`, `ALTER`, `INSERT`
+
+If the user query implies a destructive operation, return an error message in `sql_query` explaining that only SELECT queries are supported.
+
+---
+
+# TOOL VERIFICATION
+
+After generating the SQL query, you must simulate verification using the provided tools.
+"""

src/schemas/_pydantic_agent.py

@@ -1,71 +1,215 @@
 from typing import Optional
 from pydantic import BaseModel, Field, model_validator
 
+
 class Message(BaseModel):
-    role: str = Field(..., description="Role of the message")
-    content: str = Field(..., description="Content of the message")
+    role: str = Field(
+        ...,
+        description="Role of the message sender. Must be one of: 'user', 'assistant', or 'system'.",
+    )
+    content: str = Field(..., description="The text content of the message.")
+
 
 class TableDetails(BaseModel):
-    table_name: str = Field(..., description="Name of the table")
+    table_name: str = Field(
+        ...,
+        description=(
+            "Exact name of the table as provided in the input schema. "
+            "Do NOT use tables that are not explicitly listed in the provided schema."
+        ),
+    )
     column_names: list[str] = Field(
-        ..., description="List of column required for the query"
+        ...,
+        description=(
+            "List of column names from this table that are required to answer the user query. "
+            "Only include columns that exist in the provided schema for this table. "
+            "Do NOT use SELECT * — always list specific columns. "
+            "Do NOT fabricate or infer column names not present in the schema."
+        ),
     )
 
 
 class JoinConditions(BaseModel):
-    join_type: str = Field(..., description="Type of join")
-    left_table: str = Field(..., description="Name of the left table")
-    left_table_column: str = Field(..., description="Name of the left table column")
-    right_table: str = Field(..., description="Name of the right table")
-    right_table_column: str = Field(..., description="Name of the right table column")
+    join_type: str = Field(
+        ...,
+        description=(
+            "Type of SQL JOIN to apply. Must be one of: 'INNER', 'LEFT', 'RIGHT', 'FULL'. "
+            "Choose based on the relationship implied by the user query."
+        ),
+    )
+    left_table: str = Field(
+        ...,
+        description="Name of the left (primary/driving) table in the JOIN. Must exist in the provided schema.",
+    )
+    left_table_column: str = Field(
+        ...,
+        description=(
+            "Column from the left table used in the JOIN ON condition. "
+            "Typically a primary key or foreign key. Must exist in the provided schema."
+        ),
+    )
+    right_table: str = Field(
+        ...,
+        description="Name of the right (joined) table in the JOIN. Must exist in the provided schema.",
+    )
+    right_table_column: str = Field(
+        ...,
+        description=(
+            "Column from the right table used in the JOIN ON condition. "
+            "Must match the left_table_column semantically (e.g., foreign key relationship). "
+            "Must exist in the provided schema."
+        ),
+    )
 
 
 class FilterConditions(BaseModel):
-    table_name: str = Field(..., description="Name of the table")
-    column_name: str = Field(..., description="Name of the column")
-    operator: str = Field(..., description="Operator for the filter")
-    value: str = Field(..., description="Value for the filter")
+    table_name: str = Field(
+        ...,
+        description="Name of the table that contains the column to filter on. Must exist in the provided schema.",
+    )
+    column_name: str = Field(
+        ...,
+        description=(
+            "Name of the column to apply the filter on. "
+            "Must exist in the provided schema for the specified table."
+        ),
+    )
+    operator: str = Field(
+        ...,
+        description=(
+            "SQL comparison operator to use in the WHERE clause. "
+            "Allowed values: '=', '!=', '>', '<', '>=', '<=', 'LIKE', 'IN', 'BETWEEN', 'IS NULL', 'IS NOT NULL'. "
+            "Choose based on the condition described in the user query."
+        ),
+    )
+    value: str = Field(
+        ...,
+        description=(
+            "The filter value to compare against. "
+            "For string values, include surrounding quotes (e.g., \"'active'\"). "
+            "For dates, use ISO format (e.g., \"'2024-01-01'\"). "
+            "For IN operator, use comma-separated values in parentheses (e.g., \"('a', 'b')\"). "
+            "For BETWEEN, use format \"'value1' AND 'value2'\". "
+            "For IS NULL / IS NOT NULL, set value to empty string ''."
+        ),
+    )
 
 
 class OrderConditions(BaseModel):
-    table_name: str = Field(..., description="Name of the table")
-    column_name: str = Field(..., description="Name of the column")
-    order: str = Field(..., description="Order for the filter")
+    table_name: str = Field(
+        ...,
+        description="Name of the table that contains the column to sort by. Must exist in the provided schema.",
+    )
+    column_name: str = Field(
+        ...,
+        description=(
+            "Name of the column to sort by. "
+            "Must exist in the provided schema for the specified table. "
+            "Choose the column that best represents the sorting intent "
+            "(e.g., 'created_at' for latest, 'price' for highest/lowest)."
+        ),
+    )
+    order: str = Field(
+        ...,
+        description=(
+            "Sort direction. Must be either 'ASC' or 'DESC'. "
+            "Use 'DESC' for: latest, newest, highest, most recent, top. "
+            "Use 'ASC' for: oldest, lowest, alphabetical, earliest."
+        ),
+    )
 
 
 class SQLQueryExtractor(BaseModel):
     tables_required: list[TableDetails] = Field(
-        ..., description="List of tables required for the query"
+        ...,
+        description=(
+            "List of tables needed to answer the user query. "
+            "Only include tables from the provided schema that are directly necessary. "
+            "Do NOT include unrelated tables. Each entry must specify only the required columns."
+        ),
     )
 
     is_join_required: bool = Field(
-        ..., description="Indicates if join is required for the query"
+        ...,
+        description=(
+            "Set to true if the query requires data from more than one table. "
+            "Set to false if a single table is sufficient. "
+            "If true, joins_required must be populated."
+        ),
     )
     joins_required: Optional[list[JoinConditions]] = Field(
-        [], description="List of joins required for the query"
+        [],
+        description=(
+            "List of JOIN conditions required to combine tables. "
+            "Must be populated when is_join_required is true. "
+            "Must be empty ([]) when is_join_required is false."
+        ),
     )
 
     is_filter_required: bool = Field(
-        ..., description="Indicates if filter is required for the query"
+        ...,
+        description=(
+            "Set to true if the user query contains any filtering intent such as: "
+            "comparisons (greater than, less than, equal to), date ranges (before, after, between), "
+            "pattern matching (like, contains, starts with), or conditional constraints (only, exclude, where). "
+            "Set to false if no filtering is needed. If true, filters_required must be populated."
+        ),
     )
     filters_required: Optional[list[FilterConditions]] = Field(
-        [], description="List of filters required for the query"
+        [],
+        description=(
+            "List of WHERE clause conditions extracted from the user query. "
+            "Must be populated when is_filter_required is true. "
+            "Must be empty ([]) when is_filter_required is false."
+        ),
     )
 
     is_order_required: bool = Field(
-        ..., description="Indicates if order is required for the query"
+        ...,
+        description=(
+            "Set to true if the user query implies sorting, such as: "
+            "latest, newest, oldest, highest, lowest, top N, alphabetical, sorted by. "
+            "Set to false if no ordering is implied. If true, orders_required must be populated."
+        ),
     )
     orders_required: Optional[list[OrderConditions]] = Field(
-        [], description="List of orders required for the query"
+        [],
+        description=(
+            "List of ORDER BY conditions derived from the user query. "
+            "Must be populated when is_order_required is true. "
+            "Must be empty ([]) when is_order_required is false."
+        ),
     )
 
-    sql_query: str = Field(..., description="SQL query")
+    sql_query: str = Field(
+        ...,
+        description=(
+            "The final, executable SQL SELECT query that answers the user query. "
+            "Rules: "
+            "(1) Use only tables and columns from the provided schema. "
+            "(2) Never use SELECT * — always list specific columns. "
+            "(3) Use table-qualified column names (e.g., table.column) to avoid ambiguity. "
+            "(4) Include JOIN clauses if is_join_required is true. "
+            "(5) Include WHERE clause if is_filter_required is true. "
+            "(6) Include ORDER BY clause if is_order_required is true. "
+            "(7) Only generate SELECT queries — never DELETE, DROP, UPDATE, TRUNCATE, ALTER, or INSERT."
+        ),
+    )
 
     is_sql_query_verified_using_provided_tool: bool = Field(
-        ..., description="Indicates if SQL query is verified using provided tools"
+        ...,
+        description=(
+            "Indicates whether the generated sql_query has been verified using the provided verification tool. "
+            "This MUST be set true if you call the provided tool to verify the query. "
+            "This must be set to false if you do not call the provided tool to verify the query."
+        ),
     )
     is_sql_query_mark_as_safe_using_provided_tool: bool = Field(
-        ..., description="Indicates if SQL query is marked as safe using provided tools"
+        ...,
+        description=(
+            "Indicates whether the generated sql_query has been marked as safe by the provided safety-check tool. "
+            "This must be set True if provided tool return true else set False. "
+        ),
     )
 
     @model_validator(mode="after")

src/utils/_pydantic_agent.py

@@ -38,6 +38,22 @@ class PydanticAgent:
 
     async def _verify_sql_query(self, sql_query):
         try:
+            words_shoould_not_present_in_sql_query = [
+                "DELETE", 
+                "DROP", 
+                "UPDATE", 
+                "TRUNCATE", 
+                "ALTER", 
+                "INSERT"
+            ]
+            sql_query = sql_query.lower().strip()
+            if any(
+                word.lower() in sql_query
+                for word in words_shoould_not_present_in_sql_query
+            ):
+                raise Exception(
+                    f"SQL query contains a destructive operation: {sql_query}. Only SELECT queries are allowed."
+                )
             async with DatabaseConfig.async_session() as session:
                 await session.execute(sql_query)
         except Exception as e: