Storage Backends

Built-in storage backends for agents

ai-query provides built-in storage backends for different use cases. Storage backends implement the Storage protocol and are injected into agents via the storage parameter.

Import

from ai_query import Storage
from ai_query.agents import MemoryStorage
from ai_query.agents import SQLiteStorage

Storage Protocol

All storage backends implement this protocol:

from typing import Protocol, Any

class Storage(Protocol):
    async def get(self, key: str) -> Any | None: ...
    async def set(self, key: str, value: Any) -> None: ...
    async def delete(self, key: str) -> None: ...
    async def keys(self, prefix: str = "") -> list[str]: ...

---

MemoryStorage

In-memory storage for development and testing. Data is lost when the process exits.

from ai_query.agents import MemoryStorage

Constructor

MemoryStorage()

No parameters required. Creates an empty in-memory store.

Usage

from ai_query.agents import Agent, MemoryStorage
from ai_query.providers import openai

agent = Agent(
    "my-agent",
    model=openai("gpt-4o"),
    storage=MemoryStorage(),
    initial_state={"count": 0}
)

async with agent:
    await agent.update_state(count=1)
    print(agent.state)  # {"count": 1}

Notes

• Data is shared across all agents using the same MemoryStorage instance
• Thread-safe for concurrent access
• Fast for development and testing
• Not suitable for production persistence

---

SQLiteStorage

Async SQLite persistence using aiosqlite. Non-blocking database operations that won’t freeze your event loop.

from ai_query.agents import SQLiteStorage

Constructor

SQLiteStorage(path: str)

path · str · required

Path to the SQLite database file. Use :memory: for in-memory database.

Usage

from ai_query.agents import Agent, SQLiteStorage
from ai_query.providers import openai

agent = Agent(
    "user-123",
    model=openai("gpt-4o"),
    storage=SQLiteStorage("agents.db"),
    initial_state={"preferences": {}}
)

async with agent:
    await agent.update_state(theme="dark")
    print(await agent.chat("Hello!"))

# Later, in a new process...
async with agent:
    print(agent.state)  # {"preferences": {}, "theme": "dark"}

Methods

sql

async def sql(self, query: str, *params: Any) -> list[dict[str, Any]]

Execute a SQL query against the storage database asynchronously.

Parameters:

• query: SQL query string with ? placeholders
• *params: Query parameters

Returns: List of rows as dictionaries.

Example:

storage = SQLiteStorage("./data/agents.db")

# Custom queries on the database
results = await storage.sql(
    "SELECT * FROM agent_state WHERE key LIKE ?",
    "user-%"
)

Notes

• Uses lazy connection initialization
• Connection is automatically closed when the agent stops
• Tables are created automatically on first use
• State and messages are stored as JSON

---

Creating Custom Storage Backends

Implement the Storage protocol to use any backend:

from ai_query import Storage
import redis.asyncio as redis
import json

class RedisStorage:
    def __init__(self, url: str):
        self.client = redis.from_url(url)

    async def get(self, key: str) -> Any | None:
        value = await self.client.get(key)
        return json.loads(value) if value else None

    async def set(self, key: str, value: Any) -> None:
        await self.client.set(key, json.dumps(value))

    async def delete(self, key: str) -> None:
        await self.client.delete(key)

    async def keys(self, prefix: str = "") -> list[str]:
        return await self.client.keys(f"{prefix}*")

# Use it
from ai_query.agents import Agent
from ai_query.providers import openai

agent = Agent(
    "bot",
    model=openai("gpt-4o"),
    storage=RedisStorage("redis://localhost")
)

DynamoDB Example

import boto3
import json
from typing import Any

class DynamoDBStorage:
    def __init__(self, table_name: str):
        self.table = boto3.resource("dynamodb").Table(table_name)

    async def get(self, key: str) -> Any | None:
        response = self.table.get_item(Key={"pk": key})
        item = response.get("Item")
        return json.loads(item["value"]) if item else None

    async def set(self, key: str, value: Any) -> None:
        self.table.put_item(Item={"pk": key, "value": json.dumps(value)})

    async def delete(self, key: str) -> None:
        self.table.delete_item(Key={"pk": key})

    async def keys(self, prefix: str = "") -> list[str]:
        response = self.table.scan(
            FilterExpression="begins_with(pk, :prefix)",
            ExpressionAttributeValues={":prefix": prefix}
        )
        return [item["pk"] for item in response.get("Items", [])]

PostgreSQL Example

import asyncpg
import json
from typing import Any

class PostgreSQLStorage:
    def __init__(self, dsn: str):
        self.dsn = dsn
        self._pool = None

    async def _get_pool(self):
        if self._pool is None:
            self._pool = await asyncpg.create_pool(self.dsn)
        return self._pool

    async def get(self, key: str) -> Any | None:
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT value FROM agent_storage WHERE key = $1",
                key
            )
            return json.loads(row["value"]) if row else None

    async def set(self, key: str, value: Any) -> None:
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO agent_storage (key, value) VALUES ($1, $2)
                ON CONFLICT (key) DO UPDATE SET value = $2
                """,
                key, json.dumps(value)
            )

    async def delete(self, key: str) -> None:
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            await conn.execute("DELETE FROM agent_storage WHERE key = $1", key)

    async def keys(self, prefix: str = "") -> list[str]:
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            rows = await conn.fetch(
                "SELECT key FROM agent_storage WHERE key LIKE $1",
                f"{prefix}%"
            )
            return [row["key"] for row in rows]

---

Comparison

Backend	Persistence	Async	Best For
MemoryStorage	No	Yes	Development, testing, ephemeral state
SQLiteStorage	Yes	Yes	Local apps, prototypes, single-instance servers
Custom (Redis)	Yes	Yes	Production, distributed, multi-instance
Custom (DynamoDB)	Yes	Yes	Serverless, AWS deployments
Custom (PostgreSQL)	Yes	Yes	Production, complex queries

Choosing a Backend

• Development/Testing: Use MemoryStorage - fast, no setup
• Local Deployment: Use SQLiteStorage - simple persistence, works out of box
• Multi-Instance/Scaling: Create a custom backend with Redis, DynamoDB, or PostgreSQL
• Serverless: Create a custom backend with DynamoDB or other serverless databases