---
title: "Agent"
description: "Abstract base class for stateful AI agents"
---

The `Agent` class is the foundation for building stateful AI agents with persistent storage, WebSocket support, and lifecycle hooks.

## Import

```python
from ai_query.agents import Agent
```

## Class Definition

```python
class Agent(ABC, Generic[State]):
    initial_state: State = {}
    
    def __init__(self, agent_id: str, *, env: Any = None)
```

### Parameters

<ParamField path="agent_id" type="str" required>
  Unique identifier for this agent instance. Used as the key for state persistence.
</ParamField>

<ParamField path="env" type="Any">
  Optional environment bindings. Used for Cloudflare Durable Objects integration.
</ParamField>

## Properties

### state

```python
@property
def state(self) -> State
```

Current agent state. Raises `RuntimeError` if accessed before `start()` is called.

### messages

```python
@property
def messages(self) -> list[Message]
```

Conversation history for this agent.

### id

```python
@property
def id(self) -> str
```

The agent's unique identifier.

## Methods

### set_state

```python
async def set_state(self, state: State) -> None
```

Update the agent's state. This will:
1. Update in-memory state
2. Persist to storage
3. Call `on_state_update` hook
4. Broadcast new state to all WebSocket clients

### save_messages

```python
async def save_messages(self, messages: list[Message]) -> None
```

Persist messages to storage.

### clear_messages

```python
async def clear_messages(self) -> None
```

Clear the conversation history.

### broadcast

```python
async def broadcast(self, message: str | bytes) -> None
```

Send a message to all connected WebSocket clients.

### stream_to_sse

```python
async def stream_to_sse(self, event: str, data: str) -> None
```

Send an SSE event to all connected SSE clients.

**Parameters:**
- `event` - Event type (e.g., `"ai_start"`, `"ai_chunk"`, `"ai_end"`)
- `data` - Event data to send

## Actor Communication

Methods for agent-to-agent communication and event-driven architecture.

### invoke

```python
async def invoke(
    self, 
    agent_id: str, 
    payload: dict[str, Any],
    timeout: float = 30.0
) -> dict[str, Any]
```

Call another agent and wait for response. Requires a transport to be configured (automatic when using `AgentServer`).

**Parameters:**
- `agent_id` - Target agent's identifier
- `payload` - Request data to send
- `timeout` - Max wait time in seconds (default: 30)

**Returns:** Response dict from the target agent.

**Example:**
```python
result = await self.invoke("summarizer", {
    "task": "summarize",
    "text": long_document
})
print(result["summary"])
```

### emit

```python
async def emit(self, event: str, data: dict[str, Any]) -> None
```

Fire an event to all subscribers. Events are namespaced: `"{agent_id}:{event}"`.

**Example:**
```python
await self.emit("task.complete", {"result": analysis_result})
```

### handle_invoke

```python
async def handle_invoke(self, payload: dict[str, Any]) -> dict[str, Any]
```

Override to handle incoming `invoke()` calls from other agents.

**Parameters:**

<ParamField path="payload" type="dict" required>
  The data sent by the calling agent. **You define the structure.** A common pattern is to use a `task` field to route to different handlers.
</ParamField>

**Returns:** A dict that gets sent back to the calling agent.

**Example:**
```python
async def handle_invoke(self, payload: dict) -> dict:
    # payload structure is whatever the caller sends
    # Common pattern: use "task" to route
    task = payload.get("task")
    
    if task == "summarize":
        text = payload.get("text", "")
        summary = await self.chat(f"Summarize: {text}")
        return {"status": "ok", "summary": summary}
    
    if task == "translate":
        text = payload.get("text", "")
        lang = payload.get("language", "Spanish")
        result = await self.chat(f"Translate to {lang}: {text}")
        return {"status": "ok", "translation": result}
    
    return {"status": "error", "message": f"Unknown task: {task}"}
```

### handle_request

```python
async def handle_request(self, request: dict[str, Any]) -> dict[str, Any]
```

Serverless lifecycle handler. Use for stateless environments (Lambda, Cloud Run).

**Request Format:**

| Field | Type | Description |
|-------|------|-------------|
| `action` | `str` | What to do: `"chat"`, `"invoke"`, or `"state"` |
| `message` | `str` | Required for `action="chat"` |
| `payload` | `dict` | Required for `action="invoke"` (passed to `handle_invoke`) |

**Response Format:**

| Action | Response |
|--------|----------|
| `chat` | `{"agent_id": "...", "response": "AI response..."}` |
| `invoke` | `{"agent_id": "...", "result": {...}}` |
| `state` | `{"agent_id": "...", "state": {...}}` |

**Example (AWS Lambda):**
```python
def handler(event, context):
    agent_id = event["pathParameters"]["agent_id"]
    body = json.loads(event.get("body", "{}"))
    
    agent = MyAgent(agent_id)
    result = asyncio.run(agent.handle_request({
        "action": body.get("action", "chat"),
        "message": body.get("message", ""),
        "payload": body.get("payload", {}),
    }))
    
    return {"statusCode": 200, "body": json.dumps(result)}
```

### on_receive

```python
async def on_receive(self, message: IncomingMessage) -> Any
```

Unified handler for all incoming messages (WebSocket, HTTP, or agent invoke). Override for a single entry point.

**Example:**
```python
async def on_receive(self, message: IncomingMessage) -> Any:
    if message.source == "client":
        response = await self.chat(message.content)
        await message.reply(response)
    elif message.source == "agent":
        return await self.process_request(message.content)
```

### serve

```python
def serve(
    self,
    host: str = "localhost",
    port: int = 8080,
    path: str = "/ws",
) -> None
```

Start a WebSocket + SSE server for this agent. This is a blocking call.

**Endpoints:**
- `ws://{host}:{port}{path}` - WebSocket for messages
- `http://{host}:{port}/events` - SSE for AI streaming

**Example:**
```python
MyAgent("agent-1").serve(port=8080)
```

### serve_async

```python
async def serve_async(
    self,
    host: str = "localhost",
    port: int = 8080,
    path: str = "/ws",
) -> None
```

Async version of `serve()`.

### serve_many (class method)

```python
@classmethod
def serve_many(
    cls,
    host: str = "localhost",
    port: int = 8080,
    config: AgentServerConfig | None = None,
) -> None
```

Start a multi-agent server for this agent class. Each client connects to a unique agent instance via URL path:

- `ws://{host}:{port}/agent/{agent_id}/ws` (WebSocket)
- `http://{host}:{port}/agent/{agent_id}/events` (SSE)
- `http://{host}:{port}/agent/{agent_id}/state` (REST API)

See [Multi-Agent Server](/core/agent-server) for details.

### start

```python
async def start(self) -> None
```

Initialize the agent. Loads state/messages from storage, starts the message processor, and calls `on_start` hook.

### stop

```python
async def stop(self) -> None
```

Stop the agent's message processing loop. Drains remaining messages and cancels the processor task. Called automatically when using the context manager.

## Actor Model Methods

Agents implement the Actor model - all messages are processed sequentially through a mailbox, eliminating race conditions.

### enqueue

```python
def enqueue(
    self,
    kind: str,
    payload: Any,
    connection: Connection | None = None,
    ctx: ConnectionContext | None = None,
) -> None
```

Enqueue a message for sequential processing (fire-and-forget).

**Parameters:**
- `kind` - Message type: `"message"`, `"connect"`, `"close"`, `"error"`
- `payload` - The message payload
- `connection` - Associated WebSocket connection
- `ctx` - Connection context (for `"connect"` kind)

**Example:**
```python
# Typically called by the server, not directly
agent.enqueue("message", "Hello", connection=conn)
```

### enqueue_invoke

```python
async def enqueue_invoke(
    self,
    payload: dict[str, Any],
    timeout: float = 30.0,
) -> dict[str, Any]
```

Enqueue an invoke and wait for the result with timeout enforcement.

**Parameters:**
- `payload` - The invoke payload
- `timeout` - Maximum time to wait in seconds (default: 30)

**Returns:** The response from `handle_invoke`.

**Raises:** `asyncio.TimeoutError` if the invoke times out.

**Example:**
```python
# Used internally by transports - typically use invoke() instead
result = await agent.enqueue_invoke({"task": "analyze"}, timeout=60.0)
```

## Lifecycle Hooks

Override these methods in your subclass:

### on_start

```python
async def on_start(self) -> None
```

Called when the agent starts (after state is loaded).

### on_state_update

```python
def on_state_update(self, state: State, source: str | Connection) -> None
```

Called when state changes. `source` is `"server"` for server-initiated changes.

## WebSocket Hooks

### on_connect

```python
async def on_connect(self, connection: Connection, ctx: ConnectionContext) -> None
```

Called when a WebSocket client connects.

### on_message

```python
async def on_message(self, connection: Connection, message: str | bytes) -> None
```

Called when a message is received from a client.

### on_close

```python
async def on_close(self, connection: Connection, code: int, reason: str) -> None
```

Called when a WebSocket client disconnects.

### on_error

```python
async def on_error(self, connection: Connection, error: Exception) -> None
```

Called on WebSocket errors.

## Abstract Methods

These must be implemented by storage backends:

```python
@abstractmethod
async def _load_state(self) -> State | None

@abstractmethod
async def _save_state(self, state: State) -> None

@abstractmethod
async def _load_messages(self) -> list[Message]

@abstractmethod
async def _save_messages(self, messages: list[Message]) -> None
```

## Context Manager

Agents can be used as async context managers:

```python
async with MyAgent("agent-1") as agent:
    await agent.set_state({"key": "value"})
```

---

## Connection Protocol

```python
from ai_query.agents import Connection
```

Protocol that any WebSocket must implement to work with agents.

```python
@runtime_checkable
class Connection(Protocol):
    async def send(self, message: str | bytes) -> None:
        """Send a message to the client."""
        ...
    
    async def close(self, code: int = 1000, reason: str = "") -> None:
        """Close the connection."""
        ...
```

**Implement this to use agents with any WebSocket library.**

---

## ConnectionContext

```python
from ai_query.agents import ConnectionContext
```

Context passed to `on_connect` with connection metadata.

```python
@dataclass
class ConnectionContext:
    request: Any = None      # Original HTTP request
    metadata: dict[str, Any] = {}  # Query params, headers, etc.
```

---

## AioHttpConnection

```python
from ai_query.agents import AioHttpConnection
```

Built-in `Connection` implementation for aiohttp WebSockets. Used internally by `serve()`.

```python
class AioHttpConnection(Connection):
    def __init__(self, ws: WebSocketResponse, request: Request)
    
    async def send(self, message: str | bytes) -> None
    async def close(self, code: int = 1000, reason: str = "") -> None
    
    username: str | None  # Convenience attribute
```

---

## Example

```python
from ai_query.agents import InMemoryAgent

class CounterAgent(InMemoryAgent):
    initial_state = {"count": 0}
    
    async def on_start(self):
        print(f"Started with count: {self.state['count']}")
    
    async def increment(self):
        await self.set_state({"count": self.state["count"] + 1})

# As context manager
async with CounterAgent("counter-1") as agent:
    await agent.increment()
    print(agent.state["count"])  # 1

# Or as server
CounterAgent("counter-1").serve(port=8080)
```
