SQL agent

The SQL Agent is the core of our system. It takes natural language questions and converts them to valid SQLite queries using the schema context.

Text-to-SQL pipeline

How the SQL agent converts natural language to SQL queries

Without the schema, the LLM would hallucinate table and column names. By including SCHEMA_INFO, we give it the exact tables, columns, and types available, so it can only reference real database objects. This dramatically improves query accuracy.

text2sql_agent.py
python
def sql_agent(state: AgentState) -> AgentState:
    """Generate SQL query from natural language question"""
    question = state["question"]
    iteration = state.get("iteration", 0)

    prompt = f"""You are a SQL expert. Convert this question into a valid SQLite query.

{SCHEMA_INFO}

Question: {question}

Important Guidelines:
1. Use only the tables and columns in the schema
2. Use proper JOIN clauses for multiple tables
3. Return JSON with "sql_query" and "reasoning" fields
4. Use aggregate functions (COUNT, SUM, AVG) appropriately
5. Add LIMIT 10 unless user specifies otherwise
6. For date comparisons, dates are TEXT in ISO format

Respond in JSON format:
{{"reasoning": "...", "sql_query": "SELECT ..."}}
"""

The prompt injects the full schema (SCHEMA_INFO) and guidelines. The guidelines prevent common SQL mistakes like missing JOINs or unbounded queries.

With the prompt ready, we call the LLM using the same pattern as the Guardrails Agent: temperature=0 for deterministic output and JSON mode for structured responses.

text2sql_agent.py
python
    # Call the LLM with JSON mode
    response = completion(
        model=DEFAULT_MODEL,
        messages=[
            {"role": "system", "content": AGENT_CONFIGS["sql_agent"]["system_prompt"]},
            {"role": "user", "content": prompt}
        ],
        temperature=0,
        response_format={"type": "json_object"}
    )

Same pattern as the Guardrails Agent: temperature=0 for consistent SQL generation, JSON mode for structured output.

Some LLMs wrap SQL in markdown code fences (``sql ... ``) even when using JSON mode. We strip those out to avoid syntax errors when executing the query.

text2sql_agent.py
python
    # Parse and clean the generated SQL
    result = json.loads(response.choices[0].message.content)
    sql_query = result.get("sql_query", "").strip()

    # Clean up any markdown formatting the LLM might add
    sql_query = sql_query.replace("```sql", "").replace("```", "").strip()

    state["sql_query"] = sql_query
    state["sql_reason"] = result.get("reasoning", "")
    state["iteration"] = iteration + 1

    return state

After cleaning, the SQL query and reasoning are stored in state. The iteration counter tracks retry attempts for error recovery.

Even with JSON mode, the LLM sometimes wraps the SQL query in markdown code blocks like ``sql ... ``. If we pass that directly to SQLite, it will fail because the backticks are not valid SQL. The strip step is a simple but essential cleanup.

Flashcards: Flashcards

Loading practice…

AI prompt: Try it with AI

Loading practice…

Checkpoint: SQL agent knowledge check

Loading practice…

The SQL Agent can now convert natural language to valid SQLite queries. Next, we will build the Executor that actually runs these queries against the database.