# Lambda Agent API

WebSocket API SDK для взаимодействия с AI-агентом.

## Установка

```bash
pip install .
```

Требуется Python 3.14+.

## Быстрый старт (с использованием AgentApi)

```python
import asyncio
from lambda_agent_api.agent_api import AgentApi, OM


def my_callback(message):
    if isinstance(message, OM.Error):
        print(f"\n[Ошибка: {message.code}] {message.details}")
    elif isinstance(message, OM.Status):
        print("✓ Agent status update")
    elif isinstance(message, OM.GracefulDisconnect):
        print("✓ Agent gracefully requested disconnect")


async def main():
    api = AgentApi("agent-1", "ws://localhost:8000/ws", callback=my_callback)

    await api.connect()
    try:
        response = await api.send_message("Привет, агент!")

        async for chunk in response:
            if isinstance(chunk, OM.EventTextChunk):
                print(chunk.text, end="", flush=True)

        # После окончания Generation возможно получить EventEnd в очереди и сохранить tokens
        # (в current implementation: `response` - генератор, для токенов смотрите `EventEnd` в callback)

    finally:
        await api.close()


asyncio.run(main())
```

## AgentApi - Асинхронный Python клиент

Новая библиотека `AgentApi` предоставляет типизированный асинхронный клиент для WebSocket взаимодействия с агентом.

### Характеристики

- ✅ **Асинхронный клиент** на основе `aiohttp`
- ✅ **Явное подключение/закрытие** через `connect()`/`close()`
- ✅ **Защита от параллельных запросов** через `AgentBusyException`
- ✅ **ResponseIterator** для асинхронной итерации по чанкам ответа
- ✅ **Callback** для обработки событий вне генерации ответа (`Status`, `Error`, `GracefulDisconnect`)
- ✅ **Типизированные сообщения** через Pydantic с дискриминированными объединениями
- ✅ **Обработка ошибок** с кастомным исключением `AgentException`
- ✅ **Логирование** на всех уровнях операций
- ✅ **Полная документация** всех методов

### Использование

```python
from lambda_agent_api.agent_api import AgentApi, OM

api = AgentApi("agent-1", "ws://localhost:8000/ws", callback=my_callback)

await api.connect()
try:
    response = await api.send_message("Your question here")

    async for chunk in response:
        if isinstance(chunk, OM.EventTextChunk):
            print(chunk.text, end="", flush=True)

    print("\nDone!")
finally:
    await api.close()
```

# Обработка ошибок

- `AgentBusyException` возникает, если отправить `send_message` пока предыдущий запрос ещё в процессе.
- `AgentException` возникает, если агент возвращает `ERROR` или есть проблемы с подключением.
- `on_disconnect` callback вызывается один раз при закрытии/разрыве соединения.

Callback функция для обработки событий вне генерации:

```python
def my_callback(message):
    if isinstance(message, OM.Status):
        print("Agent status update")
    elif isinstance(message, OM.Error):
        print(f"Agent error: {message.code} - {message.details}")
    elif isinstance(message, OM.GracefulDisconnect):
        print("Agent disconnecting gracefully")
```

## Классический подход (низкоуровневый)

```python
import asyncio
import websockets
from lambda_agent_api.models import ServerMessage, OM


async def main():
    uri = "ws://localhost:8000/ws"

    async with websockets.connect(uri) as ws:
        # 1. Ждём STATUS - подтверждение готовности
        status = await ws.recv()
        print(f"Connected: {status}")

        # 2. Отправляем сообщение
        await ws.send('{"type": "USER_MESSAGE", "text": "Привет!"}')

        # 3. Читаем ответ в виде потока событий
        while True:
            msg = await ws.recv()
            data = ServerMessage.model_validate_json(msg)

            match data:
                case OM.AgentEvent(subtype=OM.AgentEventType.TEXT_CHUNK):
                    print(data.text, end="", flush=True)
                case OM.EventEnd():
                    print(f"\n[Завершено, использовано токенов: {data.tokens_used}]")
                    break
                case OM.Error():
                    print(f"\n[Ошибка: {data.code}] {data.details}")
                    break


asyncio.run(main())
```

## Протокол

### Клиент → Сервер

#### USER_MESSAGE

Полное сообщение от пользователя.

```json
{
  "type": "USER_MESSAGE",
  "text": "Текст сообщения"
}
```

| Поле | Тип   | Описание          |
|------|-------|-------------------|
| type | string | Всегда `USER_MESSAGE` |
| text | string | Текст сообщения  |

### Сервер → Клиент

#### STATUS

Отправляется сервером при открытии соединения с клиентом. Будет дополнен информацией о готовности агента принимать сообщения.

```json
{
  "type": "STATUS"
}
```

#### AGENT_EVENT

Базовый класс для ивентов, которые стримит агент во время генерации ответа. Конкретный класс для ивента определяется по `subtype`.

##### TEXT_CHUNK

Чанк текста ответа агента.

```json
{
  "type": "AGENT_EVENT",
  "subtype": "TEXT_CHUNK",
  "text": "Фрагмент текста"
}
```

##### END

Агент закончил генерацию ответа.

```json
{
  "type": "AGENT_EVENT",
  "subtype": "END",
  "tokens_used": 42
}
```

| Поле        | Тип    | Описание              |
|-------------|--------|-----------------------|
| tokens_used | int    | Количество использованных токенов |

#### ERROR

Неопределенная ошибка в работе агента.

```json
{
  "type": "ERROR",
  "code": "error_code",
  "details": "Описание ошибки"
}
```

| Поле    | Тип   | Описание       |
|---------|-------|----------------|
| code    | string | Код ошибки    |
| details | string | Подробности   |

#### GRACEFUL_DISCONNECT

Отправляется перед завершением работы контейнера с агентом. Например, при долгом бездействии. Нужно, чтобы отделять обрыв соединения из-за ошибки с необходимостью повторного подключения. Приход этого сообщения означает, что агент осознанно завершает работу с клиентом по какой-то причине. Для дальнейшего взаимодействия нужно снова обратиться к мастеру.

```json
{
  "type": "GRACEFUL_DISCONNECT"
}
```

![Схема взаимодействия](docs/schema.png)

## Зависимости

- Python 3.14+
- pydantic >= 2.12.5