Back to Examples
Intermediate

Database Query Agent

Build an agent that can query PostgreSQL databases autonomously using the database plugin

PluginsPostgreSQLDatabase

#Overview

Learn how to build an agent that can autonomously query and analyze data from PostgreSQL databases. This example shows how plugins extend agent capabilities and enable complex database operations.

#What You'll Learn

  • Using the PostgreSQL plugin with agents
  • How agents autonomously query databases
  • Plugin tool integration
  • Database security best practices

#Prerequisites

  • PostgreSQL database running locally or remotely
  • Database credentials
  • asyncpg installed: pip install asyncpg

#Step 1: Set Up the PostgreSQL Plugin

The PostgreSQL plugin provides database tools that agents can use:

python
from daita import Agent
from daita.plugins import postgresql
import asyncio
 
# Create the database plugin
db = postgresql(
    host="localhost",
    port=5432,
    database="analytics",
    username="postgres",
    password="your_password"
)
 
# The plugin automatically exposes these tools to the agent:
# - query_database: Execute SELECT queries
# - list_tables: Get all table names
# - get_table_schema: Inspect table structure
# - execute_sql: Run INSERT/UPDATE/DELETE

Security note:

  • Never hardcode credentials
  • Use environment variables: os.getenv("DB_PASSWORD")
  • Use read-only accounts when possible

#Step 2: Create the Agent with Database Access

Pass the plugin to the agent to enable database operations:

python
async def main():
    # Create database plugin
    db = postgresql(
        host="localhost",
        database="analytics",
        username="postgres",
        password="your_password"
    )
 
    # Create agent with database tools
    agent = Agent(
        name="Data Analyst",
        prompt="""You are a data analyst. Help users query and analyze
        data from the database. Use the available database tools to
        explore tables, understand schemas, and query data.""",
        llm_provider="openai",
        model="gpt-4",
        tools=[db]  # Pass plugin as tool
    )
 
    await agent.start()

What happens:

  • Plugin registers its tools with the agent
  • Agent can now see query_database, list_tables, etc.
  • LLM receives tool descriptions and signatures

#Step 3: Let the Agent Explore the Database

The agent can autonomously discover and query database structure:

python
async def main():
    db = postgresql(host="localhost", database="analytics")
 
    agent = Agent(
        name="Data Analyst",
        prompt="You are a data analyst. Help users understand database data.",
        llm_provider="openai",
        model="gpt-4",
        tools=[db]
    )
 
    await agent.start()
 
    try:
        # Agent will autonomously:
        # 1. Call list_tables() to see available tables
        # 2. Call get_table_schema() to understand structure
        # 3. Formulate and provide a summary
        result = await agent.run("What tables are in this database?")
        print(result)
    finally:
        await agent.stop()

Agent's autonomous process:

  1. Receives task: "What tables are in this database?"
  2. Recognizes need for database tool
  3. Calls list_tables()
  4. Receives table list
  5. Formats and presents results to user

#Step 4: Complex Data Analysis

Agents can handle multi-step analysis tasks:

python
async def main():
    db = postgresql(host="localhost", database="analytics")
 
    agent = Agent(
        name="Data Analyst",
        prompt="""You are a data analyst. Query databases to answer
        questions. Always verify table structure before querying.""",
        llm_provider="openai",
        model="gpt-4",
        tools=[db]
    )
 
    await agent.start()
 
    try:
        # Agent performs multi-step analysis:
        # 1. Checks if 'users' table exists (list_tables)
        # 2. Inspects schema (get_table_schema)
        # 3. Queries for active users (query_database)
        # 4. Analyzes and presents results
        result = await agent.run("""
        How many active users do we have?
        Show me the top 5 users by activity count.
        """)
 
        print(result)
    finally:
        await agent.stop()
 
if __name__ == "__main__":
    asyncio.run(main())

#Step 5: Monitor Tool Execution

See exactly how the agent interacts with the database:

python
async def main():
    db = postgresql(host="localhost", database="analytics")
 
    agent = Agent(
        name="Data Analyst",
        prompt="You are a data analyst.",
        llm_provider="openai",
        model="gpt-4",
        tools=[db]
    )
 
    await agent.start()
 
    try:
        result = await agent.run_detailed("""
        Find all users who signed up in the last 30 days.
        """)
 
        print(f"Answer: {result['result']}\n")
        print(f"Tools used: {len(result['tool_calls'])}\n")
 
        # See each database operation
        for i, call in enumerate(result['tool_calls'], 1):
            print(f"{i}. {call['tool']}")
            print(f"   Args: {call['args']}")
            if 'sql' in call['args']:
                print(f"   SQL: {call['args']['sql']}")
    finally:
        await agent.stop()
 
if __name__ == "__main__":
    asyncio.run(main())

#Complete Example with Error Handling

python
from daita import Agent
from daita.plugins import postgresql
import asyncio
import os
 
async def main():
    # Create database plugin with environment variables
    db = postgresql(
        host=os.getenv("DB_HOST", "localhost"),
        database=os.getenv("DB_NAME", "analytics"),
        username=os.getenv("DB_USER", "postgres"),
        password=os.getenv("DB_PASSWORD")
    )
 
    # Create agent
    agent = Agent(
        name="Data Analyst",
        prompt="""You are a data analyst. Help users query and analyze
        database data. Always verify table schemas before querying.""",
        llm_provider="openai",
        model="gpt-4",
        tools=[db]
    )
 
    await agent.start()
 
    try:
        # Complex analysis task
        result = await agent.run("""
        Analyze our user base:
        1. Total number of users
        2. Users who signed up in the last 30 days
        3. Most active users this month
        """)
 
        print("Analysis Results:")
        print("=" * 50)
        print(result)
 
    except Exception as e:
        print(f"Error: {e}")
    finally:
        await agent.stop()
 
if __name__ == "__main__":
    asyncio.run(main())

#Framework Internals

How plugins work with agents:

  1. Plugin Initialization: Plugin connects to database, registers tools
  2. Tool Registration: Each plugin method becomes an available tool
  3. LLM Integration: Tool definitions sent to LLM with parameters
  4. Autonomous Execution: Agent decides which tools to call based on task
  5. Connection Management: Plugin handles database connections automatically

Available PostgreSQL tools:

  • query_database(sql, params): Execute SELECT queries
  • execute_sql(sql, params): Run INSERT/UPDATE/DELETE
  • list_tables(): Get all table names
  • get_table_schema(table_name): Inspect columns and types

#Security Best Practices

  1. Use read-only accounts for query-only agents
  2. Environment variables for credentials
  3. Limit agent permissions - don't give DROP/ALTER access
  4. Validate queries - monitor what the agent executes
  5. Rate limiting - prevent excessive queries

#Key Takeaways

  1. Plugins = Pre-built tools: PostgreSQL plugin adds database operations
  2. Autonomous exploration: Agents can discover and understand database structure
  3. Multi-step workflows: Agents handle complex analysis automatically
  4. Connection management: Plugin handles pooling and cleanup
  5. Security matters: Always use appropriate permissions and credentials

#Next Steps