Persistent Chatbot

import asyncio
from ai_query import tool, Field
from ai_query.agents import Agent, SQLiteStorage
from ai_query.providers import openai

class SupportBot(Agent):
    def __init__(self, user_id: str):
        super().__init__(
            user_id,
            model=openai("gpt-4o"),
            system="You are a helpful customer support agent.",
            storage=SQLiteStorage("./support_bot.db"),
            initial_state={"ticket_count": 0}
        )

    @tool(description="Create a support ticket")
    async def create_ticket(
        self,
        subject: str = Field(description="Ticket subject")
    ) -> str:
        count = self.state.get("ticket_count", 0) + 1
        await self.update_state(ticket_count=count)
        return f"Created ticket #{count}: {subject}"

async def main():
    user_id = "user-123"
    
    # Session 1
    bot = SupportBot(user_id)
    async with bot:
        print(await bot.chat("Hi, I'm Alice. Create a ticket for my keyboard."))

    # Session 2 (Persistent!)
    bot2 = SupportBot(user_id)
    async with bot2:
        print(await bot2.chat("What was that ticket number again?"))

if __name__ == "__main__":
    asyncio.run(main())

agent = Agent("id", storage=SQLiteStorage("db.sqlite"))