LLM-Powered SQL Agent Example

---

This script demonstrates how to build a stepwise SQL agent using LangChain, LangGraph, and Ollama.
Given a user’s natural language question, the agent will:

1. Generate a syntactically correct SQL query using an LLM, based on the database schema and question.
2. Execute the generated SQL query against a live database.
3. Summarize the SQL result into a natural language answer.

Key Features:

• Uses a state graph to manage the workflow: question → query → result → answer.
• Enforces SQL best practices and schema awareness via a detailed system prompt.
• Ensures only allowed tables/columns are used and never selects all columns.
• Modular and extensible for more complex agentic workflows.

Dependencies:

• langchain
• langchain-ollama
• langchain-community
• langgraph
• A running Ollama server and SQL database
• A get_db() function in db.py that returns a LangChain-compatible database object

---

from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama.chat_models import ChatOllama
from db import get_db
from typing_extensions import TypedDict, Annotated
from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool
from langgraph.graph import START, StateGraph

# Define the agent's state, which tracks the question, generated query, SQL result, and final answer.
class State(TypedDict):
    question: str
    query: str
    result: str
    answer: str

# Initialize the LLM (Ollama) for SQL generation and answer summarization.
llm = ChatOllama(
    model="llama3:latest",
    base_url="http://localhost:11434",
    stream=False,
    temperature=0
)

# System prompt to instruct the LLM on how to generate safe, correct SQL queries.
system_message = """
Given an input question, create a syntactically correct {dialect} query to
run to help find the answer. Unless the user specifies in his question a
specific number of examples they wish to obtain, always limit your query to
at most {top_k} results. You can order the results by a relevant column to
return the most interesting examples in the database.

Never query for all the columns from a specific table, only ask for the
few relevant columns given the question.

Pay attention to use only the column names that you can see in the schema
description. Be careful to not query for columns that do not exist. Also,
pay attention to which column is in which table.

Strictly follow these rules:
- Return a syntactically correct query using MySQL dialect.
- Use JOINs properly when a question spans multiple tables.
- Use `GROUP BY`, `ORDER BY`, and `LIMIT` appropriately when asking for "top", "most", "average", etc.
- When a question involves filtering by category (e.g., actors, short films), use specific column filters.
- Never use `SELECT *`. Always specify only the relevant columns.
- Use table aliases (`n`, `t`, `r`, etc.) to reduce verbosity.
- Qualify columns if ambiguity exists.
- Provide correlated aggregates when the question implies "by person", "by film", or "by group".

Only use the following tables:
{table_info}
"""

# User prompt template for the question.
user_prompt = "Question: {input}"

# Combine system and user prompts into a chat prompt template.
query_prompt_template = ChatPromptTemplate(
    [("system", system_message), ("user", user_prompt)]
)

# Define the expected output structure from the LLM for SQL generation.
class QueryOutput(TypedDict):
    query: Annotated[str, ..., "Syntactically valid SQL query."]

def write_query(state: State):
    """
    Given the user's question, generate a syntactically correct SQL query using the LLM.
    """
    db = get_db()
    prompt = query_prompt_template.invoke(
        {
            "dialect": db.dialect,
            "top_k": 10,
            "table_info": db.get_table_info(),
            "input": state["question"],
        }
    )
    structured_llm = llm.with_structured_output(QueryOutput)
    result = structured_llm.invoke(prompt)
    return {"query": result["query"]}

def execute_query(state: State):
    """
    Execute the generated SQL query against the database and return the result.
    """
    db = get_db()
    execute_query_tool = QuerySQLDatabaseTool(db=db)
    return {"result": execute_query_tool.invoke(state["query"])}

def generate_answer(state: State):
    """
    Given the original question, SQL query, and SQL result, generate a natural language answer.
    """
    prompt = (
        "Given the following user question, corresponding SQL query, "
        "and SQL result, answer the user question.\n\n"
        f'Question: {state["question"]}\n'
        f'SQL Query: {state["query"]}\n'
        f'SQL Result: {state["result"]}'
    )
    response = llm.invoke(prompt)
    return {"answer": response.content}

# Build the state graph to orchestrate the workflow.
graph_builder = StateGraph(State).add_sequence(
    [write_query, execute_query, generate_answer]
)
graph_builder.add_edge(START, "write_query")
graph = graph_builder.compile()