Create AI Agent tool for query DB

Step 1: create ducks.db file

Step 2: Seeding data

import duckdb

# Connect to or create a DuckDB database file
con = duckdb.connect("ducks.duckdb")

# Create the ducks table
con.execute("""
CREATE TABLE ducks (
    id VARCHAR,
    color VARCHAR,
    firstName VARCHAR,
    lastName VARCHAR,
    gender VARCHAR
);
""")

# Insert data
con.execute("""
INSERT INTO ducks VALUES
('kA0KgL', 'red', 'Marty', 'McFly', 'male'),
('dx3ngL', 'teal', 'Duckota', 'Fanning', 'female'),
('FQ4dU1', 'yellow', 'Duck', 'Norris', 'male'),
('JqS7ZZ', 'red', 'James', 'Pond', 'male'),
('ZM5uJL', 'black', 'Darth', 'Wader', 'male'),
('05FuKa', 'yellow', 'Clint', 'Beakwood', 'male'),
('wKq9zD', 'yellow', 'Mary', 'Quackens', 'female'),
('QCab5l', 'orange', 'Ducky', 'Balboa', 'male'),
('eKiyA5', 'orange', 'Captain', 'Quack', 'male'),
('YiSGQl', 'teal', 'Wonder', 'Duck', 'female');
""")

# Confirm the data
result = con.execute("SELECT * FROM ducks LIMIT 10").fetchall()
for row in result:
    print(row)

# Close connection
con.close()

Step 3: Build AI Agent

import duckdb
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI
con = duckdb.connect("ducks.duckdb")
schema = con.execute(f"DESCRIBE ducks").fetchdf()
schema_str = schema.to_string(index=False)

@tool
def query(query:str):
    """
    queries the database for information and returns the results
    
    Args:
        query: str: the query to be executed
    """
    return str(con.execute(query).fetchone()[0])

google_llm = ChatGoogleGenerativeAI(
    model="gemini-2.0-flash",
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key="xxx",
).bind_tools([query])


    
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
def query_ai_agent(input):
    
    # Initialize the prompt template
    prompt = ChatPromptTemplate.from_messages([
        ("human", f"You are provied You are given a DuckDB schema for table 'ducks'\n\n: {schema_str}\n\n.\n\n Answer the user query: {input} in the single sentence"),
        MessagesPlaceholder(variable_name="agent_scratchpad"),
    ])

    # Create the agent and executor with out llm, tools and prompt
    agent = create_tool_calling_agent(google_llm, [query],prompt)
    agent_executor = AgentExecutor(agent=agent, tools=[query], verbose=True)

    # Run our query 
    res = agent_executor.invoke({"input": input})
    print(res["output"])    
    
query_ai_agent("What are the color of duck with first name 'Marty'?")    

Here is output:

> Entering new AgentExecutor chain...

Invoking: `query` with `{'query': "SELECT color FROM ducks WHERE firstName = 'Marty'"}`


redThe color of the duck with the first name 'Marty' is red.