from lambda_agent_api import AgentApi

# Lambda Agent API

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

## Установка
В `master` всегда будет актуальная рабочая версия.
```bash
pip install git+https://git.lambda.coredump.ru/platform/agent_api.git
```

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

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

```python
import asyncio

from lambda_agent_api.agent_api import AgentApi
from lambda_agent_api.server import MsgEventTextChunk


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

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

        async for chunk in response:
            if isinstance(chunk, MsgEventTextChunk):
                print(chunk.text, end="", flush=True)
            elif isinstance(chunk, MsgEventToolCallChunk):
                print(f"Tool call started: {chunk.tool_name}")
            elif isinstance(chunk, MsgEventToolResult):
                print(f"Tool result: {chunk.result}")
            elif isinstance(chunk, MsgEventCustomUpdate):
                print(f"Progress update: {chunk.payload}")
            elif isinstance(chunk, MsgEventEnd):
                print(f"Generation ended, tokens used: {chunk.tokens_used}")

    finally:
        await api.close()


asyncio.run(main())
```

> `AgentApi.send_message()` возвращает стриминг-итерируемый объект, который может выдавать не только текстовые чанки, но и события инструментов (`MsgEventToolCallChunk`, `MsgEventToolResult`, `MsgEventCustomUpdate`) и финальный `MsgEventEnd`.

## Предполагаемое использование

```python
from lambda_agent_api.agent_api import AgentApi

connected_agents: dict[str, AgentApi] = {}


# агент автоматически выключится после нескольких минут бездействия
# этот колбек вызовется при разрыве соединения
def on_agent_disconnect(agent: AgentApi):
    del connected_agents[agent.id]


async def on_telegram_message(from_user: int, text: str):
    agent_id = get_agent_id_by_user(user_id)

    agent = connected_agents.get(agent_id, None)
    if not agent:
        agent = AgentApi(agent_id, get_agent_url(agent_id), on_disconnect=on_agent_disconnect)
        await agent.connect()
        connected_agents[agent_id] = agent

    async for event in agent.send_message(text):
        ...
```

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

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


## Протокол

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

#### USER_MESSAGE

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

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

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

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

#### STATUS

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

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

#### AGENT_EVENT_TEXT_CHUNK

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

```json
{
  "type": "AGENT_EVENT_TEXT_CHUNK",
  "text": "Фрагмент текста",
  "source": "main"
}
```

| Поле   | Тип    | Описание                                      |
|--------|--------|-----------------------------------------------|
| type   | string | Всегда `AGENT_EVENT_TEXT_CHUNK`              |
| text   | string | Фрагмент текста ответа агента                |
| source | string | Источник события (по умолчанию "main")       |

#### AGENT_EVENT_TOOL_CALL_CHUNK

Агент решил использовать инструмент и генерирует аргументы.

```json
{
  "type": "AGENT_EVENT_TOOL_CALL_CHUNK",
  "tool_name": "имя_инструмента",
  "args_chunk": "{\"key\": \"value\"}",
  "source": "main"
}
```

| Поле        | Тип     | Описание                                      |
|-------------|---------|-----------------------------------------------|
| type        | string  | Всегда `AGENT_EVENT_TOOL_CALL_CHUNK`         |
| tool_name   | string  | Имя инструмента (может быть null в первом чанке) |
| args_chunk  | string  | Кусок JSON-аргументов (может быть null)      |
| source      | string  | Источник события (по умолчанию "main")       |

#### AGENT_EVENT_TOOL_RESULT

Инструмент отработал и вернул результат.

```json
{
  "type": "AGENT_EVENT_TOOL_RESULT",
  "tool_name": "имя_инструмента",
  "result": "результат выполнения",
  "source": "main"
}
```

| Поле       | Тип    | Описание                                      |
|------------|--------|-----------------------------------------------|
| type       | string | Всегда `AGENT_EVENT_TOOL_RESULT`             |
| tool_name  | string | Имя инструмента                              |
| result     | any    | Результат выполнения (строка, объект или массив) |
| source     | string | Источник события (по умолчанию "main")       |

#### AGENT_EVENT_CUSTOM_UPDATE

Кастомный прогресс (например, скачивание файла) изнутри инструмента.

```json
{
  "type": "AGENT_EVENT_CUSTOM_UPDATE",
  "payload": {"status": "in_progress", "progress": 50},
  "source": "main"
}
```

| Поле     | Тип             | Описание                                      |
|----------|-----------------|-----------------------------------------------|
| type     | string          | Всегда `AGENT_EVENT_CUSTOM_UPDATE`           |
| payload  | object          | Любые данные о прогрессе                      |
| source   | string          | Источник события (по умолчанию "main")       |

#### AGENT_EVENT_END

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

```json
{
  "type": "AGENT_EVENT_END",
  "tokens_used": 42
}
```

| Поле        | Тип    | Описание                                      |
|-------------|--------|-----------------------------------------------|
| type        | string | Всегда `AGENT_EVENT_END`                     |
| tokens_used | int    | Количество использованных токенов            |

#### ERROR

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

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

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

#### GRACEFUL_DISCONNECT

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

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

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

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

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

#### GRACEFUL_DISCONNECT

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

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

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