Post cover image

February 22, 2025

Building a LLM Agent to Directly Interact with a Database

Large Language Models (LLMs) have revolutionized the way we interact with data and build intelligent applications. In this guide, I will…

By Ayush Gupta

10 min read

Large Language Models (LLMs) have revolutionized the way we interact with data and build intelligent applications. In this guide, I will walk you through the process of creating an LLM-powered Database agent using Google's Gemini model and LangGraph that can directly interact with a database to query and retrieve data efficiently.

Prerequisites

  • Basic knowledge of Python, LangGraph and Langchain
  • Access to Google Gemini API key [You can use any LLM you want]
  • Database setup (e.g., SQLite, PostgreSQL, MySQL, etc.)

So before we dive in, I'll assume you're familiar with the basics I mentioned above— but even if you don't, I'll still keep things as simple and clear as possible. So, grab a cup of coffee, settle in, and get ready — because we're about to explore some seriously cool AI stuff. Let's go! 🚀☕

None

Understanding the Workflow

The application flow begins when the user enters a natural language query (e.g., "What were the total sales last month?"). This input is passed to the query_gen node, where the LLM, with the help of the provided database schema (tables, columns, and their relationships), converts the user's question into an SQL query.

The generated SQL query is then sent to the query_check node, which acts as a verification layer. This step is like a double-check process to ensure that the query is syntactically and logically correct, reducing the chance of runtime errors during execution.

Once validated, the query moves to the query_execute node, where it is executed on the database using the db_exec_tool. This tool runs the SQL query and retrieves the result from the database.

Finally, the LLM interprets the raw query output into a human-friendly response, making the data easier for the user to understand (e.g., "The total sales last month were $50,000"). The response is then returned to the user as the final output.

None

Should we start with the coding part now ? Yeah we should !

Let's start by setting up the LLM

import os
from langchain_google_genai import ChatGoogleGenerativeAI
# Set your API key
os.environ["GOOGLE_API_KEY"] = "YOUR_KEY"
# Initialize the model
llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash")
llm.invoke("hi")

# cool now our llm is ready to handle tasks

In this application, we are going to use three tools. Tools in an LLM agent are essentially functions that allow the agent to interact with the outside world, such as databases, APIs, or other systems. The three tools used in this application are:

  1. list_tables_tool — This tool is responsible for listing all the tables in the connected database. It helps the agent understand the structure of the database.
  2. get_schema_tool — This tool retrieves the schema of a specific table, providing details about the columns and their data types.
  3. db_exec_tool — This tool executes SQL queries on the database and returns the results.

These tools collectively enable the LLM agent to understand the database structure and interact with it dynamically to retrieve information or perform operations.

db = SQLDatabase.from_uri("Your_database_url") # connecting database
from langchain_community.agent_toolkits import SQLDatabaseToolkit
toolkit=SQLDatabaseToolkit(db=db,llm=llm)
tools=toolkit.get_tools()

# mapping through all the tools to get the tools we need

# tool1
list_tables_tool = next((tool for tool in tools if tool.name == "sql_db_list_tables"), None)
# list_tables_tool.invoke("") fetches names of all table

# tool2
get_schema_tool = next((tool for tool in tools if tool.name == "sql_db_schema"), None)
# get_schema_tool.invoke("table_name") fetches schema for that table

@tool # this decorator specifies that this function is a tool
def db_exec_tool(query : str)-> str:
    """
    Execute a SQL query against the database and return the result.
    If the query is invalid or returns no result, an error message will be returned.
    In case of an error, the user is advised to rewrite the query and try again.
    """
    # this description is very important for llm to understand what this tool does

    # Remove ```sql and ``` if present
    query = query.replace("```sql", "").replace("```", "").strip()

    print("Executing query:")
    print(query)

    result = db.run_no_throw(query)

    print("Query result:")
    # print(result)
    return {"result": result}

That's it, all 3 tools are ready now we will create our 3 nodes query_gen, query_checker and query_executor, these nodes will use the tools to give final output.

Let's start with query_gen node

This node will take MessagesState as input. The MessagesState is an object that contains a key called messages, which holds an array. This array acts as a conversation log, storing the user's inputs as well as every response generated by the LLM and outputs from various nodes in the workflow. Each time a new message is produced—whether it's from user input, an LLM response, or a node's output—it will be appended to this messages array, ensuring that the complete interaction history is maintained and accessible throughout the agent's execution. This allows the agent to maintain context and make informed decisions based on the conversation's flow.

since it's an array….so state["messages"][-1].content will naturally return the value at last index which will have the last response generated by llm or previous node, the output of any last activity will always be at last index.

def query_gen(state: MessagesState)-> Command[Literal[ "query_check"]]:
    """
    Query Generation Node to convert natural language database queries into PostgreSQL queries.

    Args:
        state (MessagesState): The current state containing the conversation history with a natural language database query.

    Returns:
        Command: A command to update the state with the generated PostgreSQL query.
    """
    # print("query_generator")
    # print(state["messages"][-1].content) # user given input will be at last index

    query_agent = create_react_agent(
            llm,  # The language model instance used by the agent
            tools=[list_tables_tool, get_schema_tool],  # List of database tools the agent can utilize
            state_modifier=(
                "You are an expert database query generator specialized in PostgreSQL. "
                "You are provided with tools to list the tables in the database and get the schema of specific tables. "
                "Always use the list_tables_tool first to get an overview of available tables. "
                "Then, for any relevant table(s), use the get_schema_tool to retrieve their schema before constructing the query. "
                "Ensure that the SQL query you generate is syntactically correct and follows PostgreSQL standards. "
                "Your final output should only be the SQL query, without any additional explanation or commentary."
                "If empty response is returned by database ,deal accordingly by telling that to user"
            )
        )

    # Example of how the agent might be invoked (depends on the exact agent framework you use)
    result = query_agent.invoke(state)
    # print(result)
    return Command(
     update={
        "messages": [
            # Append the reason (supervisor's response) to the state, tagged with "supervisor"
            HumanMessage(content=result["messages"][-1].content, name="supervisor")
        ]
    },
    goto="query_check",  # Specify the next node in the workflow
  )

In the query_gen function, we are creating what is called a REACT Agent. Think of this agent as an intelligent assistant that can both think and act. It can decide when to think using the LLM (language model) and when to act using the tools provided to it.

Inputs Needed to Create the Agent:

  • state_modifier (prompt): Instructions guiding the agent to list tables, fetch schemas, and generate a valid query.
  • LLM: The language model that helps with reasoning and query generation.
  • Tools : list_tables_tool (Gets the names of all tables ) , get_schema_tool: (Gets the structure of a table).

The agent first checks the table names, then fetches the table's schema if needed, and finally generates the SQL query using this information. The result is added to the state's messages array, and the flow moves to the next step (query_check).

How does the agent know which to call ?

Remember the description we provide in an agent function, that description is used by LLM to understand what a specific tool does, what input it takes and what it returns, in our case list_tables_tool and get_schema_tool are inbuilt tools hence no description is needed for them but for db_exec_tool we have added this description.

In the end of query_gen function we are returning The Command object which is used to instruct the system on how to update the conversation state and where to move next in the workflow.

It updates the messages array in the MessagesState by appending the agent's final output (the generated SQL query) as a HumanMessage with the name "supervisor". This labels the message as coming from the "supervisor" or the system, indicating it is a response generated by the agent. After updating the state, the goto="query_check" tells the system to move to the next node in the workflow named query_check, ensuring the conversation progresses smoothly.

return Command(
     update={
        "messages": [
            # Append the reason (supervisor's response) to the state, tagged with "supervisor"
            HumanMessage(content=result["messages"][-1].content, name="supervisor")
        ]
    },
    goto="query_check",  # Specify the next node in the workflow
  )

okay so now we have 2 items in our MessagesState object array , one is for user input and other item is the output given by query_gennode. This MessagesState object , containing this array is passed to query_check node which also expect MessagesState as input.

def query_check(state: MessagesState)-> Command[Literal["query_execute"]]:
    """
    This tool checks if the provided SQL query is correct.
    If incorrect, it returns the corrected query; otherwise, it returns the original query.
    """
    query_check_system = """You are a SQL expert with a strong attention to detail.
    You work with PostgreSQL, SQLite, and other relational databases.
    Your task is to carefully review the provided SQL query for any mistakes, including:
    - Quoting identifiers correctly (e.g., "Snippet" vs Snippet in PostgreSQL)
    - Data type mismatches
    - Using the correct number of arguments in functions
    - Ensuring joins use valid columns
    - Checking for NULL handling issues
    - Ensuring correct usage of UNION vs UNION ALL
    - Proper casting and type usage
    Make sure that the final query is in a postgres acceptable format
    If there is an issue, respond with the **corrected query only**.
    If the query is already correct, simply return the **original query**.
    """
    query = state["messages"][-1].content

    full_prompt = f"{query_check_system}\n\nQuery:\n{query}"

    # LLM invocation
    response = llm.with_structured_output(QueryChecker).invoke(full_prompt)

    print("query_Check")
    print(response)
    return Command(
     update={
        "messages": [
            # Append the reason (supervisor's response) to the state, tagged with "supervisor"
            HumanMessage(content=response.query, name="supervisor")
        ]
    },
    goto="query_execute",  # Specify the next node in the workflow
  )

Unlike the query_gen function, this node does not require a REACT agent or tools because it only involves a single LLM call to check the correctness of a query.

The function starts by fetching the last message from the state, which contains the SQL query generated in the previous query_gen node.

query = state["messages"][-1].content

A system prompt is prepared, instructing the LLM to review the SQL query for common issues like quoting, data types, joins, NULL handling, and other PostgreSQL-specific mistakes. The query is then combined with this prompt and sent to the LLM for validation.

The LLM returns either the same query if it is correct or a corrected version if any issues are found. Finally, the validated query is appended to the messages array in the state, and the workflow moves to the next node called query_execute.

The query_execute function is responsible for executing the final SQL query and returning the result in a human-readable format. this function is the last node in our workflow and also takes the MessagesState as input , this state now contains an array where value at last index is our SQL query to be executed.

def query_execute(state : MessagesState):
  """
    This node executes the provided SQL query and return the response.
    It returns the response in simple human understandable format.
  """
  # print("state from query execute")
  # print(state)
  executing_agent = create_react_agent(
          llm,  # The language model instance used by the agent
          tools=[db_exec_tool],  # List of database tools the agent can utilize
          state_modifier=(
              "You are an expert PostgreSQL query executor. "
              "you can use db_exec_tool for execution of sql query"
              "Your primary task is to execute the provided SQL query accurately and return the result. "
              "Ensure that the query is executed against the database without modification, and the response is returned in a clear, human-understandable format. "
              "Do not generate new queries, retrieve schemas, or list tables unless explicitly asked. "
              "Your output should only contain the query execution result in a human readable manner with some explanation or commentary."
              "Make sure that you show user the response and not just explanation of response"
          ),

      )

  final_result = executing_agent.invoke(state)
  print(final_result["messages"][-1].content)

This node creates a REACT agent that uses the db_exec_tool—a tool specifically for executing SQL queries on the database. The agent is configured with a state modifier (prompt), instructing it to execute the provided query as it is, without modifying or generating new queries. It is also instructed to return the execution result in a clear, easy-to-understand format, including both the response data and some explanation if needed.

The agent processes the state containing the validated query from the previous node (query_check), executes it against the database using the db_exec_tool, and gives us a final result. This result is printed to the Terminal , you can return it using same Command object used in previous functions.

Now that our nodes and tools are set up, it's time to connect everything together and define the flow of our process and test the agent. We start by creating a StateGraph instance and specifying MessagesState as the state type that will flow through the nodes.

builder = StateGraph(MessagesState)
# adding all our nodes to graph
builder.add_node("query_gen", query_gen)  # Add the supervisor node to the graph
builder.add_node("query_check", query_check)  # Add the supervisor node to the graph
builder.add_node("query_execute", query_execute)  # Add the supervisor node to the graph

# Addin edges to define the workflow of the graph
builder.add_edge(START, "query_gen")  # Connect the start node to the supervisor node
# rest of edges between nodes are already added by Command object that we were returning in each node
# Compile the graph to finalize its structure
graph = builder.compile()

To test the application, we create an input structured as a dictionary with a key called messages. This key holds an array that represents the conversation state. Each element in the array is a tuple where the first part is the role (such as "user"), and the second part is the message content (like the user's query). This format aligns with the MessagesState, which keeps track of the entire conversation, appending every input and output as the process moves through each node.

We then pass this input into the graph using a loop, which streams outputs from each node. The responses from the nodes are processed, and if any output is not None, it is printed.

import pprint

inputs = {
    "messages": [
        ("user", "how many users have connected their github account?"),
    ]
}
for output in graph.stream(inputs):
    for key, value in output.items():
        if value is None:
            continue
        # pprint.pprint(f"Output from node '{key}':")
        # pprint.pprint(value, indent=2, width=80, depth=None)

And the output for the query will be something like this :

None

Kaboom!! you are done , now you can make a terminal application out of it or use FastApi to convert it into an API or make it as a chatBOT , just go and tinker around to build whatever you want. If you have any doubts or queries you can dm me on twitter/X

The link to complete code file is here :

github — https://github.com/ayushgupta4002/database-talks

Thank you for staying till the end, Do give it a star and leave a comment below if you liked this project and learned something new ; )

None