DatabaseMixin

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

from gaia import DatabaseAgent
from gaia.agents.chat.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

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}

# 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})

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

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
            );
        ''')

# 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}'")

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

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})

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

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

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"