Build a Custom SQL Agent

from langchain.chat_models import init_chat_model
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit

model = init_chat_model("gpt-5.4")
db = SQLDatabase.from_uri("sqlite:///Chinook.db")
toolkit = SQLDatabaseToolkit(db=db, llm=model)
tools = toolkit.get_tools()

def list_tables(state: MessagesState):
    list_tables_tool = next(t for t in tools if t.name == "sql_db_list_tables")
    tool_call = AIMessage(content="", tool_calls=[{
        "name": "sql_db_list_tables", "args": {}, "id": "abc123", "type": "tool_call"
    }])
    tool_message = list_tables_tool.invoke(tool_call.tool_calls[0])
    response = AIMessage(f"Available tables: {tool_message.content}")
    return {"messages": [tool_call, tool_message, response]}

get_schema_tool = next(t for t in tools if t.name == "sql_db_schema")
get_schema_node = ToolNode([get_schema_tool], name="get_schema")

def call_get_schema(state: MessagesState):
    llm_with_tools = model.bind_tools([get_schema_tool], tool_choice="any")
    response = llm_with_tools.invoke(state["messages"])
    return {"messages": [response]}

run_query_tool = next(t for t in tools if t.name == "sql_db_query")

def generate_query(state: MessagesState):
    system_msg = {"role": "system", "content": "Create a syntactically correct SQL query..."}
    llm_with_tools = model.bind_tools([run_query_tool])
    response = llm_with_tools.invoke([system_msg] + state["messages"])
    return {"messages": [response]}

def should_continue(state: MessagesState) -> Literal[END, "check_query"]:
    if not state["messages"][-1].tool_calls:
        return END
    return "check_query"

builder = StateGraph(MessagesState)
builder.add_node(list_tables)
builder.add_node(call_get_schema)
builder.add_node("get_schema", get_schema_node)
builder.add_node(generate_query)
builder.add_node(check_query)
builder.add_node("run_query", run_query_node)

builder.add_edge(START, "list_tables")
builder.add_edge("list_tables", "call_get_schema")
builder.add_edge("call_get_schema", "get_schema")
builder.add_edge("get_schema", "generate_query")
builder.add_conditional_edges("generate_query", should_continue)
builder.add_edge("check_query", "run_query")
builder.add_edge("run_query", "generate_query")

agent = builder.compile()

from langgraph.types import interrupt, Command

def run_query_tool_with_interrupt(config: RunnableConfig, **tool_input):
    response = interrupt([{
        "action": run_query_tool.name,
        "args": tool_input,
        "description": "Please review the SQL query"
    }])
    if response["type"] == "accept":
        return run_query_tool.invoke(tool_input, config)
    elif response["type"] == "edit":
        return run_query_tool.invoke(response["args"]["args"], config)
    elif response["type"] == "response":
        return response["args"]

# Compile with checkpointer
from langgraph.checkpoint.memory import InMemorySaver
agent = builder.compile(checkpointer=InMemorySaver())

# Resume after review
agent.stream(Command(resume={"type": "accept"}), config)