---
title: "Storage Backends"
description: "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

```python
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:

```python
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.

```python
from ai_query.agents import MemoryStorage
```

### Constructor

```python
MemoryStorage()
```

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

### Usage

```python
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.

```python
from ai_query.agents import SQLiteStorage
```

### Constructor

```python
SQLiteStorage(path: str)
```

<ParamField path="path" type="str" required>
  Path to the SQLite database file. Use `:memory:` for in-memory database.
</ParamField>

### Usage

```python
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

```python
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:**

```python
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:

```python
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

```python
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

```python
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
