Skip to content

SQLDatabase

pip install datapizza-ai-tools-sqldatabase

datapizza.tools.SQLDatabase.SQLDatabase

A collection of tools to interact with a SQL database using SQLAlchemy. This class is a container for methods that are exposed as tools.

__init__ 

__init__(db_uri)

Initializes the SQLDatabase tool container.

Parameters:

Name Type Description Default
db_uri str

The database URI for connection (e.g., "sqlite:///my_database.db").

 required

get_table_schema 

get_table_schema(table_name)

Returns the schema of a specific table in a human-readable format.

Parameters:

Name Type Description Default
table_name str

The name of the table to inspect.

 required

list_tables 

list_tables()

Returns a newline-separated string of available table names in the database.

run_sql_query 

run_sql_query(query)

Executes a SQL query and returns the result. For SELECT statements, it returns a JSON string of the rows. For other statements (INSERT, UPDATE, DELETE), it returns a success message.

Parameters:

Name Type Description Default
query str

The SQL query to execute.

 required

Overview

The SQLDatabase tool provides a powerful interface for AI agents to interact with any SQL database supported by SQLAlchemy. This allows models to query structured, relational data to answer questions, providing more accurate and fact-based responses.

Features

  • Broad Database Support: Connect to any database with a SQLAlchemy driver (SQLite, PostgreSQL, MySQL, etc.).
  • Schema Inspection: Allows the agent to view table schemas to understand data structure before querying.
  • Table Listing: Lets the agent list all available tables to get context of the database.

Integration with Agents

from datapizza.agents import Agent
from datapizza.clients.openai import OpenAIClient
from datapizza.tools.SQLDatabase import SQLDatabase

db_uri = "sqlite:///company.db"

# 1. Initialize the SQLDatabase tool
db_tool = SQLDatabase(db_uri=db_uri)

# 2. Create an agent and provide it with the database tool's methods
agent = Agent(
    name="database_expert",
    client=OpenAIClient(api_key="YOUR_API_KEY"),
    system_prompt="You are a database expert. Use the available tools to answer questions about the database.",
    tools=[
        db_tool.list_tables,
        db_tool.get_table_schema,
        db_tool.run_sql_query
    ]
)

# 3. Run the agent
response = agent.run("How many people work in the Engineering department?")
print(response)