DatabaseMixin

Source Code: src/gaia/database/mixin.py

Detailed Spec: spec/database-mixin The DatabaseMixin provides SQLite database access for GAIA agents. It uses Python’s built-in sqlite3 module with zero external dependencies, making it lightweight and fast.

Two Approaches

1. DatabaseAgent (Recommended for prototyping)

Use DatabaseAgent when you want the LLM to have direct database access:

from gaia import DatabaseAgent

class MyAgent(DatabaseAgent):
    def __init__(self, **kwargs):
        super().__init__(db_path="data/app.db", **kwargs)

        if not self.table_exists("items"):
            self.execute("CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)")

    def _get_system_prompt(self) -> str:
        return "You manage items. Use the database tools to query and modify data."

    def _register_tools(self):
        pass  # db_query, db_insert, db_update, db_delete auto-registered

Auto-registered tools: db_query, db_insert, db_update, db_delete, db_tables, db_schema

Composing with Other Mixins

DatabaseAgent can be extended with additional mixins for more capabilities:

from gaia import DatabaseAgent
from gaia.agents.tools import RAGToolsMixin, FileToolsMixin

class MyAgent(DatabaseAgent, RAGToolsMixin, FileToolsMixin):
    """Agent with database, RAG, and file tools."""

    def __init__(self, **kwargs):
        super().__init__(db_path="data/app.db", **kwargs)

    def _get_system_prompt(self) -> str:
        return "You can query databases, search documents, and manage files."

    def _register_tools(self):
        pass  # All tools auto-registered from parent classes

2. DatabaseMixin (Recommended for production)

Use DatabaseMixin when you want to expose domain-specific tools:

from gaia import Agent, tool, DatabaseMixin

class MyAgent(Agent, DatabaseMixin):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.init_db("data/myagent.db")

        if not self.table_exists("items"):
            self.execute('''
                CREATE TABLE items (
                    id INTEGER PRIMARY KEY,
                    name TEXT NOT NULL
                )
            ''')

    def _get_system_prompt(self) -> str:
        return "You manage items."

    def _register_tools(self):
        @tool
        def add_item(name: str) -> dict:
            """Add a new item."""
            item_id = self.insert("items", {"name": name})
            return {"id": item_id}

        @tool
        def list_items() -> dict:
            """List all items."""
            items = self.query("SELECT * FROM items")
            return {"items": items}

API Reference

Initialization

Method	Description
`init_db(path)`	Initialize database. Use `":memory:"` for in-memory DB.
`close_db()`	Close connection. Safe to call multiple times.
`db_ready`	Property. `True` if database is initialized.

CRUD Operations

Method	Returns	Description
`query(sql, params, one)`	`list[dict]` or `dict`	Execute SELECT query
`insert(table, data)`	`int` (row ID)	Insert a row
`update(table, data, where, params)`	`int` (affected count)	Update rows
`delete(table, where, params)`	`int` (deleted count)	Delete rows

Schema & Transactions

Method	Description
`execute(sql)`	Execute raw SQL (CREATE TABLE, etc.)
`table_exists(name)`	Check if table exists
`transaction()`	Context manager for atomic operations

Examples

Basic CRUD

# Insert
user_id = self.insert("users", {
    "name": "Alice",
    "email": "[email protected]"
})

# Query all
users = self.query("SELECT * FROM users")

# Query one
user = self.query(
    "SELECT * FROM users WHERE id = :id",
    {"id": user_id},
    one=True
)

# Update
self.update(
    "users",
    {"email": "[email protected]"},
    "id = :id",
    {"id": user_id}
)

# Delete
self.delete("users", "id = :id", {"id": user_id})

Transactions

Use transactions when multiple operations must succeed or fail together:

with self.transaction():
    user_id = self.insert("users", {"name": "Alice"})
    self.insert("profiles", {"user_id": user_id, "bio": "Hello"})
    self.insert("settings", {"user_id": user_id, "theme": "dark"})
    # If any insert fails, all are rolled back

Schema Initialization

def __init__(self, **kwargs):
    super().__init__(**kwargs)
    self.init_db("data/app.db")

    if not self.table_exists("users"):
        self.execute('''
            CREATE TABLE users (
                id INTEGER PRIMARY KEY,
                name TEXT NOT NULL,
                email TEXT UNIQUE,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE posts (
                id INTEGER PRIMARY KEY,
                user_id INTEGER REFERENCES users(id),
                content TEXT NOT NULL
            );
        ''')

Best Practices

Use Parameterized Queries

Always use :param placeholders to prevent SQL injection and syntax errors:

# Good - parameterized
self.query("SELECT * FROM users WHERE name = :name", {"name": user_input})

# Bad - string formatting (SQL injection risk)
self.query(f"SELECT * FROM users WHERE name = '{user_input}'")

Initialize Schema Once

Check table_exists() before creating tables to avoid errors on restart:

if not self.table_exists("users"):
    self.execute("CREATE TABLE users (...)")

Ensure data consistency by wrapping related operations in a transaction:

with self.transaction():
    order_id = self.insert("orders", {"customer_id": cust_id})
    for item in items:
        self.insert("order_items", {"order_id": order_id, **item})

Close on Shutdown

Call close_db() when the agent is done (optional but clean):

try:
    agent.run()
finally:
    agent.close_db()

Complete Example: Notes Agent

A simple but complete agent that manages notes:

from gaia import Agent, DatabaseMixin, tool

class NotesAgent(Agent, DatabaseMixin):
    """Agent that manages personal notes."""

    def __init__(self, db_path: str = "data/notes.db", **kwargs):
        super().__init__(**kwargs)
        self.init_db(db_path)

        # Create schema on first run
        if not self.table_exists("notes"):
            self.execute("""
                CREATE TABLE notes (
                    id INTEGER PRIMARY KEY,
                    title TEXT NOT NULL,
                    content TEXT,
                    created_at TEXT DEFAULT CURRENT_TIMESTAMP
                )
            """)

    def _get_system_prompt(self) -> str:
        return """You are a notes assistant. You can:
        - Create new notes with add_note
        - List all notes with list_notes
        - Search notes with search_notes
        - Delete notes with delete_note
        """

    def _register_tools(self):
        agent = self

        @tool
        def add_note(title: str, content: str = "") -> dict:
            """Create a new note."""
            note_id = agent.insert("notes", {"title": title, "content": content})
            return {"id": note_id, "message": f"Created note: {title}"}

        @tool
        def list_notes() -> dict:
            """List all notes."""
            notes = agent.query("SELECT id, title, created_at FROM notes ORDER BY created_at DESC")
            return {"notes": notes, "count": len(notes)}

        @tool
        def search_notes(query: str) -> dict:
            """Search notes by title or content."""
            notes = agent.query(
                "SELECT * FROM notes WHERE title LIKE :q OR content LIKE :q",
                {"q": f"%{query}%"}
            )
            return {"results": notes, "count": len(notes)}

        @tool
        def delete_note(note_id: int) -> dict:
            """Delete a note by ID."""
            count = agent.delete("notes", "id = :id", {"id": note_id})
            return {"deleted": count > 0}

# Usage
if __name__ == "__main__":
    agent = NotesAgent()
    # Agent is ready to use with chat interface

Testing Your Agent

Use the temp_db fixture for isolated tests:

from gaia.database import temp_db

def test_my_agent():
    schema = "CREATE TABLE items (id INTEGER PRIMARY KEY, name TEXT)"

    with temp_db(schema) as db_path:
        agent = MyAgent(skip_lemonade=True)
        agent.init_db(db_path)  # Use temp database
        agent.insert("items", {"name": "test"})

        items = agent.query("SELECT * FROM items")
        assert len(items) == 1
        assert items[0]["name"] == "test"

Design Notes

Choice	Reason
`sqlite3`	Built into Python, zero dependencies
No ORM	Simpler, more control, less memory
Dict results	Pythonic, easy to serialize to JSON
8 methods	Complete but minimal API surface
Context manager transactions	Clean, idiomatic Python

Getting Started

C++ Framework

Python Framework

Two Approaches

1. DatabaseAgent (Recommended for prototyping)

Composing with Other Mixins

2. DatabaseMixin (Recommended for production)

API Reference

Initialization

CRUD Operations

Schema & Transactions

Examples

Basic CRUD

Transactions

Schema Initialization

Best Practices

Use Parameterized Queries

Initialize Schema Once

Close on Shutdown

Complete Example: Notes Agent

Testing Your Agent

Design Notes

​Two Approaches

​1. DatabaseAgent (Recommended for prototyping)

​Composing with Other Mixins

​2. DatabaseMixin (Recommended for production)

​API Reference

​Initialization

​CRUD Operations

​Schema & Transactions

​Examples

​Basic CRUD

​Transactions

​Schema Initialization

​Best Practices

​Use Parameterized Queries

​Initialize Schema Once

​Use Transactions for Related Operations

​Close on Shutdown

​Complete Example: Notes Agent

​Testing Your Agent

​Design Notes

Two Approaches

1. DatabaseAgent (Recommended for prototyping)

Composing with Other Mixins

2. DatabaseMixin (Recommended for production)

API Reference

Initialization

CRUD Operations

Schema & Transactions

Examples

Basic CRUD

Transactions

Schema Initialization

Best Practices

Use Parameterized Queries

Initialize Schema Once

Use Transactions for Related Operations

Close on Shutdown

Complete Example: Notes Agent

Testing Your Agent

Design Notes