From b34cbaf677ad43d6a59203d05ffa8cdffed9a3e7 Mon Sep 17 00:00:00 2001 From: MrKan Date: Wed, 1 Apr 2026 23:30:00 +0300 Subject: [PATCH 1/2] =?UTF-8?q?=D0=B0=D0=B4=D0=B0=D0=BF=D1=82=D0=B0=D1=86?= =?UTF-8?q?=D0=B8=D1=8F=20=D0=B4=D0=BB=D1=8F=20=D1=83=D1=81=D1=82=D0=B0?= =?UTF-8?q?=D0=BD=D0=BE=D0=B2=D0=BA=D0=B8=20=D0=BF=D0=B0=D0=BA=D0=B5=D1=82?= =?UTF-8?q?=D0=BE=D0=BC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- api/.python-version => .python-version | 0 README.md | 243 ++++++++++++++++++++++++- api/README.md | 233 ------------------------ {api/docs => docs}/schema.png | Bin {api => lambda_agent_api}/__init__.py | 4 +- {api => lambda_agent_api}/agent_api.py | 3 +- {api => lambda_agent_api}/models.py | 0 api/pyproject.toml => pyproject.toml | 5 +- api/uv.lock => uv.lock | 30 +-- 9 files changed, 260 insertions(+), 258 deletions(-) rename api/.python-version => .python-version (100%) delete mode 100644 api/README.md rename {api/docs => docs}/schema.png (100%) rename {api => lambda_agent_api}/__init__.py (77%) rename {api => lambda_agent_api}/agent_api.py (99%) rename {api => lambda_agent_api}/models.py (100%) rename api/pyproject.toml => pyproject.toml (75%) rename api/uv.lock => uv.lock (99%) diff --git a/api/.python-version b/.python-version similarity index 100% rename from api/.python-version rename to .python-version diff --git a/README.md b/README.md index 8032115..72a44cf 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,238 @@ -## Contributing +# Lambda Agent API -- В качестве таск-трекера - раздел "Задачи" (Issues). -- В названии коммита обязательно указывать ID задачи. Пример: "#1 Описание коммита" -- При закрытии задачи в комментарии делать ссылку на PR. -- Одна задача, одна ветка, один PR. -- Формат названия ветки: #-<короткое описание>. Пример: "#1-task-description". +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 diff --git a/api/README.md b/api/README.md deleted file mode 100644 index 4796b1d..0000000 --- a/api/README.md +++ /dev/null @@ -1,233 +0,0 @@ -# Lambda Agent API - -WebSocket API SDK для взаимодействия с AI-агентом. - -## Установка - -```bash -pip install . -``` - -Требуется Python 3.14+. - -## Быстрый старт (с использованием AgentApi) - -```python -import asyncio -from 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 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 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 diff --git a/api/docs/schema.png b/docs/schema.png similarity index 100% rename from api/docs/schema.png rename to docs/schema.png diff --git a/api/__init__.py b/lambda_agent_api/__init__.py similarity index 77% rename from api/__init__.py rename to lambda_agent_api/__init__.py index a449215..fadda1f 100644 --- a/api/__init__.py +++ b/lambda_agent_api/__init__.py @@ -7,8 +7,8 @@ - IM, OM, IncomingMessage, OutgoingMessage: Pydantic модели контракта """ -from .agent_api import AgentApi, AgentException -from .models import CM, SM, ClientMessage, ServerMessage +from lambda_agent_api.agent_api import AgentApi, AgentException +from lambda_agent_api.models import CM, SM, ClientMessage, ServerMessage __all__ = [ "AgentApi", diff --git a/api/agent_api.py b/lambda_agent_api/agent_api.py similarity index 99% rename from api/agent_api.py rename to lambda_agent_api/agent_api.py index 21a30d2..c38e01e 100644 --- a/api/agent_api.py +++ b/lambda_agent_api/agent_api.py @@ -2,7 +2,8 @@ import logging from typing import Callable, Optional, AsyncIterator import aiohttp import asyncio -from models import CM, SM, ClientMessage, ServerMessage + +from lambda_agent_api.models import CM, SM, ClientMessage, ServerMessage logger = logging.getLogger(__name__) diff --git a/api/models.py b/lambda_agent_api/models.py similarity index 100% rename from api/models.py rename to lambda_agent_api/models.py diff --git a/api/pyproject.toml b/pyproject.toml similarity index 75% rename from api/pyproject.toml rename to pyproject.toml index 3b284b5..e1f135e 100644 --- a/api/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [project] -name = "api" +name = "lambda_agent_api" version = "0.1.0" description = "WebSocket API SDK для взаимодействия с AI-агентом" readme = "README.md" @@ -8,3 +8,6 @@ dependencies = [ "aiohttp>=3.13.4", "pydantic>=2.12.5", ] + +[tool.setuptools] +packages = ["lambda_agent_api"] \ No newline at end of file diff --git a/api/uv.lock b/uv.lock similarity index 99% rename from api/uv.lock rename to uv.lock index 4e51458..baed3b5 100644 --- a/api/uv.lock +++ b/uv.lock @@ -83,21 +83,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] -[[package]] -name = "api" -version = "0.1.0" -source = { virtual = "." } -dependencies = [ - { name = "aiohttp" }, - { name = "pydantic" }, -] - -[package.metadata] -requires-dist = [ - { name = "aiohttp", specifier = ">=3.13.4" }, - { name = "pydantic", specifier = ">=2.12.5" }, -] - [[package]] name = "attrs" version = "26.1.0" @@ -157,6 +142,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, ] +[[package]] +name = "lambda-agent-api" +version = "0.1.0" +source = { virtual = "." } +dependencies = [ + { name = "aiohttp" }, + { name = "pydantic" }, +] + +[package.metadata] +requires-dist = [ + { name = "aiohttp", specifier = ">=3.13.4" }, + { name = "pydantic", specifier = ">=2.12.5" }, +] + [[package]] name = "multidict" version = "6.7.1" From 1e256a545be0e3eb56fe5e0cac36c0a184eb3a31 Mon Sep 17 00:00:00 2001 From: MrKan Date: Thu, 2 Apr 2026 00:40:30 +0300 Subject: [PATCH 2/2] =?UTF-8?q?=D0=BA=D0=BE=D1=80=D1=80=D0=B5=D0=BA=D1=82?= =?UTF-8?q?=D0=BD=D1=8B=D0=B5=20pydantic=20=D0=BC=D0=BE=D0=B4=D0=B5=D0=BB?= =?UTF-8?q?=D0=B8=20=D0=B4=D0=BB=D1=8F=20=D0=B0=D0=B2=D1=82=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D1=82=D0=B8=D1=87=D0=B5=D1=81=D0=BA=D0=BE=D0=B3=D0=BE=20?= =?UTF-8?q?=D0=BE=D0=BF=D1=80=D0=B5=D0=B4=D0=B5=D0=BB=D0=B5=D0=BD=D0=B8?= =?UTF-8?q?=D1=8F=20=D0=BA=D0=BB=D0=B0=D1=81=D1=81=D0=B0=20=D0=BF=D0=BE=20?= =?UTF-8?q?=D0=BF=D0=BE=D0=BB=D1=8E=20type?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 22 +++---- lambda_agent_api/__init__.py | 5 +- lambda_agent_api/agent_api.py | 25 +++---- lambda_agent_api/client.py | 33 ++++++++++ lambda_agent_api/models.py | 121 ---------------------------------- lambda_agent_api/server.py | 72 ++++++++++++++++++++ pyproject.toml | 3 +- tests/__init__.py | 0 tests/manual.py | 24 +++++++ tests/test_models.py | 74 +++++++++++++++++++++ uv.lock | 63 ++++++++++++++++++ 11 files changed, 294 insertions(+), 148 deletions(-) create mode 100644 lambda_agent_api/client.py delete mode 100644 lambda_agent_api/models.py create mode 100644 lambda_agent_api/server.py create mode 100644 tests/__init__.py create mode 100644 tests/manual.py create mode 100644 tests/test_models.py diff --git a/README.md b/README.md index 72a44cf..eb11610 100644 --- a/README.md +++ b/README.md @@ -18,11 +18,11 @@ from lambda_agent_api.agent_api import AgentApi, OM def my_callback(message): - if isinstance(message, OM.Error): + if isinstance(message, OM.MsgError): print(f"\n[Ошибка: {message.code}] {message.details}") - elif isinstance(message, OM.Status): + elif isinstance(message, OM.MsgStatus): print("✓ Agent status update") - elif isinstance(message, OM.GracefulDisconnect): + elif isinstance(message, OM.MsgGracefulDisconnect): print("✓ Agent gracefully requested disconnect") @@ -34,7 +34,7 @@ async def main(): response = await api.send_message("Привет, агент!") async for chunk in response: - if isinstance(chunk, OM.EventTextChunk): + if isinstance(chunk, OM.MsgEventTextChunk): print(chunk.text, end="", flush=True) # После окончания Generation возможно получить EventEnd в очереди и сохранить tokens @@ -75,7 +75,7 @@ try: response = await api.send_message("Your question here") async for chunk in response: - if isinstance(chunk, OM.EventTextChunk): + if isinstance(chunk, OM.MsgEventTextChunk): print(chunk.text, end="", flush=True) print("\nDone!") @@ -93,11 +93,11 @@ Callback функция для обработки событий вне гене ```python def my_callback(message): - if isinstance(message, OM.Status): + if isinstance(message, OM.MsgStatus): print("Agent status update") - elif isinstance(message, OM.Error): + elif isinstance(message, OM.MsgError): print(f"Agent error: {message.code} - {message.details}") - elif isinstance(message, OM.GracefulDisconnect): + elif isinstance(message, OM.MsgGracefulDisconnect): print("Agent disconnecting gracefully") ``` @@ -106,7 +106,7 @@ def my_callback(message): ```python import asyncio import websockets -from lambda_agent_api.models import ServerMessage, OM +from lambda_agent_api.server import ServerMessage, OM async def main(): @@ -128,10 +128,10 @@ async def main(): match data: case OM.AgentEvent(subtype=OM.AgentEventType.TEXT_CHUNK): print(data.text, end="", flush=True) - case OM.EventEnd(): + case OM.MsgEventEnd(): print(f"\n[Завершено, использовано токенов: {data.tokens_used}]") break - case OM.Error(): + case OM.MsgError(): print(f"\n[Ошибка: {data.code}] {data.details}") break diff --git a/lambda_agent_api/__init__.py b/lambda_agent_api/__init__.py index fadda1f..d00fbe9 100644 --- a/lambda_agent_api/__init__.py +++ b/lambda_agent_api/__init__.py @@ -8,13 +8,12 @@ """ from lambda_agent_api.agent_api import AgentApi, AgentException -from lambda_agent_api.models import CM, SM, ClientMessage, ServerMessage +from lambda_agent_api.server import ServerMessage +from lambda_agent_api.client import ClientMessage __all__ = [ "AgentApi", "AgentException", - "CM", - "SM", "ClientMessage", "ServerMessage", ] diff --git a/lambda_agent_api/agent_api.py b/lambda_agent_api/agent_api.py index c38e01e..544c090 100644 --- a/lambda_agent_api/agent_api.py +++ b/lambda_agent_api/agent_api.py @@ -3,7 +3,8 @@ from typing import Callable, Optional, AsyncIterator import aiohttp import asyncio -from lambda_agent_api.models import CM, SM, ClientMessage, ServerMessage +from lambda_agent_api.server import * +from lambda_agent_api.client import * logger = logging.getLogger(__name__) @@ -52,8 +53,8 @@ class AgentApi: msg = await asyncio.wait_for(self._ws.receive(), timeout=5.0) if msg.type == aiohttp.WSMsgType.TEXT: - status_msg = ServerMessage.model_validate_json(msg.data) - if isinstance(status_msg, SM.Status): + status_msg = ServerMessage.validate_json(msg.data) + if isinstance(status_msg, MsgStatus): logger.info(f"Agent {self.id} is ready") else: raise AgentException( @@ -130,7 +131,7 @@ class AgentApi: except Exception as e: logger.error(f"Error in on_disconnect: {e}") - async def send_message(self, text: str) -> AsyncIterator[SM.AgentEvent]: + async def send_message(self, text: str) -> AsyncIterator[AgentEventUnion]: """ Нативный асинхронный генератор. Не требует отдельного класса ResponseIterator. @@ -151,8 +152,8 @@ class AgentApi: try: self._current_queue = asyncio.Queue() - message = CM.UserMessage( - type=CM.Type.USER_MESSAGE, + message = MsgUserMessage( + type=EClientMessage.USER_MESSAGE, text=text ) @@ -168,7 +169,7 @@ class AgentApi: raise chunk # Если конец ответа - if isinstance(chunk, SM.EventEnd): + if isinstance(chunk, MsgEventEnd): # Если вам нужны tokens_used, можно сохранить их в атрибут self break @@ -222,10 +223,10 @@ class AgentApi: async for msg in self._ws: if msg.type == aiohttp.WSMsgType.TEXT: try: - outgoing_msg = ServerMessage.model_validate_json( + outgoing_msg = ServerMessage.validate_json( msg.data) - if isinstance(outgoing_msg, SM.AgentEvent): + if isinstance(outgoing_msg, MsgEventTextChunk): if self._current_queue: await self._current_queue.put(outgoing_msg) # Если очереди нет (клиент отменил запрос), но токены идут — шлем их в коллбек @@ -235,11 +236,11 @@ class AgentApi: logger.warning( f"[{self.id}] AgentEvent without active request") - elif isinstance(outgoing_msg, SM.EventEnd): + elif isinstance(outgoing_msg, MsgEventEnd): if self._current_queue: await self._current_queue.put(outgoing_msg) - elif isinstance(outgoing_msg, SM.Error): + elif isinstance(outgoing_msg, MsgError): if self.callback: self.callback(outgoing_msg) error = AgentException( @@ -248,7 +249,7 @@ class AgentApi: if self._current_queue: await self._current_queue.put(error) - elif isinstance(outgoing_msg, SM.GracefulDisconnect): + elif isinstance(outgoing_msg, MsgGracefulDisconnect): if self.callback: self.callback(outgoing_msg) logger.info( diff --git a/lambda_agent_api/client.py b/lambda_agent_api/client.py new file mode 100644 index 0000000..df672c0 --- /dev/null +++ b/lambda_agent_api/client.py @@ -0,0 +1,33 @@ +from enum import Enum +from pydantic import BaseModel, Field, TypeAdapter +from typing import Annotated, Union, Literal + + +__all__ = ['EClientMessage', 'MsgUserMessage', 'ClientMessage'] + + +class EClientMessage(str, Enum): + USER_MESSAGE = "USER_MESSAGE" + + +class MsgUserMessage(BaseModel): + """ + Полное сообщение от пользователя. + """ + type: Literal[EClientMessage.USER_MESSAGE] = EClientMessage.USER_MESSAGE + text: str + """ + Текст сообщения. + """ + + +ClientMessage = TypeAdapter(Annotated[ + Union[MsgUserMessage,], + Field(discriminator="type") +]) +""" +Объединяет все типы входящих сообщений в одно для удобной автоматической десериализации.\n +Pydantic сам определит нужный тип в зависимости от поля ``type``.\n +Использование:\n +msg = ClientMessage.model_validate_json(json) +""" \ No newline at end of file diff --git a/lambda_agent_api/models.py b/lambda_agent_api/models.py deleted file mode 100644 index f9d88a3..0000000 --- a/lambda_agent_api/models.py +++ /dev/null @@ -1,121 +0,0 @@ -from pydantic import BaseModel, Field -from enum import Enum -from typing import Literal, Annotated, Union - - -class CM: - """ - Namespace для моделей входящих сообщений (от клиента к серверу).\n - CM = Client Message - """ - - class Type(str, Enum): - USER_MESSAGE = "USER_MESSAGE" - - # noinspection PyPep8Naming - - class UserMessage(BaseModel): - """ - Полное сообщение от пользователя. - """ - type: Literal[CM.Type.USER_MESSAGE] - text: str - """ - Текст сообщения. - """ - - -ClientMessage = Annotated[ - Union[CM.UserMessage,], - Field(discriminator="type") -] -""" -Объединяет все типы входящих сообщений в одно для удобной автоматической десериализации.\n -Pydantic сам определит нужный тип в зависимости от поля ``type``.\n -Использование:\n -msg = ClientMessage.model_validate_json(json) -""" - - -class SM: - """ - Namespace для моделей исходящих сообщений (от сервера к клиенту).\n - SM = Server Message - """ - - class Type(str, Enum): - STATUS = "STATUS" - AGENT_EVENT = "AGENT_EVENT" - ERROR = "ERROR" - GRACEFUL_DISCONNECT = "GRACEFUL_DISCONNECT" - - class Status(BaseModel): - """ - Отправляется сервером при открытии соединения с клиентом. - Будет дополнен информацией о готовности агента принимать сообщения. - """ - type: Literal[SM.Type.STATUS] - - class AgentEventType(str, Enum): - TEXT_CHUNK = "TEXT_CHUNK" - END = "END" - - class AgentEvent(BaseModel): - """ - Базовый класс для ивентов, которые стримит агент во время генерации ответа. - Конкретный класс для ивента определяется по ``subtype``. - """ - type: Literal[SM.Type.AGENT_EVENT] - subtype: SM.AgentEventType - - class EventTextChunk(AgentEvent): - """ - Чанк текста ответа агента. - """ - subtype: Literal[SM.AgentEventType.TEXT_CHUNK] - text: str - - class EventEnd(AgentEvent): - """ - Агент закончил генерацию ответа. - """ - subtype: Literal[SM.AgentEventType.END] - tokens_used: int - - class Error(BaseModel): - """ - Неопределенная ошибка в работе агента. - """ - type: Literal[SM.Type.ERROR] - code: str - details: str - - class GracefulDisconnect(BaseModel): - """ - Отправляется перед завершением работы контейнера с агентом. Например, при долгом бездействии. - Нужно, чтобы отделять обрыв соединения из-за ошибки с необходимостью повторного подключения. - Приход этого сообщения означает, что агент осознанно завершает работу с клиентом по какой-то причине. - Для дальнейшего взаимодействия нужно снова обратиться к мастеру. - """ - type: Literal[SM.Type.GRACEFUL_DISCONNECT] - - -AgentEventUnion = Annotated[ - Union[ - SM.EventTextChunk, - SM.EventEnd, - ], - Field(discriminator="subtype") -] - - -ServerMessage = Annotated[ - Union[SM.Status, AgentEventUnion, SM.Error, SM.GracefulDisconnect], - Field(discriminator="type") -] -""" -Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.\n -Pydantic сам определит нужный тип в зависимости от поля ``type``.\n -Использование:\n -msg = ServerMessage.model_validate_json(json) -""" diff --git a/lambda_agent_api/server.py b/lambda_agent_api/server.py new file mode 100644 index 0000000..dbeaaa5 --- /dev/null +++ b/lambda_agent_api/server.py @@ -0,0 +1,72 @@ +from pydantic import BaseModel, Field, TypeAdapter +from enum import Enum +from typing import Literal, Annotated, Union + + +__all__ = ['EServerMessage', 'MsgStatus', 'MsgError', 'MsgEventTextChunk', 'MsgEventEnd', 'AgentEventUnion', 'ServerMessage'] + + +class EServerMessage(str, Enum): + STATUS = "STATUS" + ERROR = "ERROR" + GRACEFUL_DISCONNECT = "GRACEFUL_DISCONNECT" + AGENT_EVENT_TEXT_CHUNK = "AGENT_EVENT_TEXT_CHUNK" + AGENT_EVENT_END = "AGENT_EVENT_END" + + +class MsgStatus(BaseModel): + """ + Отправляется сервером при открытии соединения с клиентом. + Будет дополнен информацией о готовности агента принимать сообщения. + """ + type: Literal[EServerMessage.STATUS] = EServerMessage.STATUS + + +class MsgEventTextChunk(BaseModel): + """ + Чанк текста ответа агента. + """ + type: Literal[EServerMessage.AGENT_EVENT_TEXT_CHUNK] = EServerMessage.AGENT_EVENT_TEXT_CHUNK + text: str + + +class MsgEventEnd(BaseModel): + """ + Агент закончил генерацию ответа. + """ + type: Literal[EServerMessage.AGENT_EVENT_END] = EServerMessage.AGENT_EVENT_END + tokens_used: int + + +class MsgError(BaseModel): + """ + Неопределенная ошибка в работе агента. + """ + type: Literal[EServerMessage.ERROR] = EServerMessage.ERROR + code: str + details: str + + +class MsgGracefulDisconnect(BaseModel): + """ + Отправляется перед завершением работы контейнера с агентом. Например, при долгом бездействии. + Нужно, чтобы отделять обрыв соединения из-за ошибки с необходимостью повторного подключения. + Приход этого сообщения означает, что агент осознанно завершает работу с клиентом по какой-то причине. + Для дальнейшего взаимодействия нужно снова обратиться к мастеру. + """ + type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = EServerMessage.GRACEFUL_DISCONNECT + + +AgentEventUnion = Union[MsgEventTextChunk, MsgEventEnd] + + +ServerMessage = TypeAdapter(Annotated[ + Union[MsgStatus, MsgEventTextChunk, MsgEventEnd, MsgError, MsgGracefulDisconnect], + Field(discriminator="type") +]) +""" +Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.\n +Pydantic сам определит нужный тип в зависимости от поля ``type``.\n +Использование:\n +msg = ServerMessage.model_validate_json(json) +""" \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index e1f135e..3ee4887 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,8 @@ requires-python = ">=3.14" dependencies = [ "aiohttp>=3.13.4", "pydantic>=2.12.5", + "pytest>=9.0.2", ] [tool.setuptools] -packages = ["lambda_agent_api"] \ No newline at end of file +packages = ["lambda_agent_api"] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/manual.py b/tests/manual.py new file mode 100644 index 0000000..3b920e3 --- /dev/null +++ b/tests/manual.py @@ -0,0 +1,24 @@ +import asyncio + +from lambda_agent_api.agent_api import AgentApi +from lambda_agent_api.server import MsgEventTextChunk + + +def my_callback(message): + print(f"Callback: {message}") + + +async def main(): + api = AgentApi("agent-1", "ws://localhost:8000/agent_ws/", callback=my_callback) + + await api.connect() + try: + async for chunk in api.send_message("Привет, агент!"): + if isinstance(chunk, MsgEventTextChunk): + print(chunk.text, end="", flush=True) + + finally: + await api.close() + + +asyncio.run(main()) \ No newline at end of file diff --git a/tests/test_models.py b/tests/test_models.py new file mode 100644 index 0000000..005e56a --- /dev/null +++ b/tests/test_models.py @@ -0,0 +1,74 @@ +import pytest +from pydantic import ValidationError + +from lambda_agent_api.server import * +from lambda_agent_api.client import * + + +@pytest.mark.parametrize( + "data, expected_type", + [ + ({"type": "USER_MESSAGE", "text": "hello"}, MsgUserMessage), + ], +) +def test_client_message_valid(data, expected_type): + msg = ClientMessage.validate_python(data) + assert isinstance(msg, expected_type) + + +@pytest.mark.parametrize( + "data", + [ + {"type": "UNKNOWN", "text": "hello"}, + {"type": "USER_MESSAGE"}, # нет text + ], +) +def test_client_message_invalid(data): + with pytest.raises(ValidationError): + ClientMessage.validate_python(data) + + +@pytest.mark.parametrize( + "data, expected_type", + [ + ({"type": "STATUS"}, MsgStatus), + ({"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hi"}, MsgEventTextChunk), + ({"type": "AGENT_EVENT_END", "tokens_used": 10}, MsgEventEnd), + ({"type": "ERROR", "code": "E1", "details": "fail"}, MsgError), + ({"type": "GRACEFUL_DISCONNECT"}, MsgGracefulDisconnect), + ], +) +def test_server_message_valid(data, expected_type): + msg = ServerMessage.validate_python(data) + assert isinstance(msg, expected_type) + + +@pytest.mark.parametrize( + "data", + [ + {"type": "AGENT_EVENT_TEXT_CHUNK"}, # нет text + {"type": "AGENT_EVENT_END"}, # нет tokens_used + {"type": "ERROR", "code": "E1"}, # нет details + {"type": "UNKNOWN"}, + ], +) +def test_server_message_invalid(data): + with pytest.raises(ValidationError): + ServerMessage.validate_python(data) + + +def test_validate_json(): + json_data = '{"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hello"}' + msg = ServerMessage.validate_json(json_data) + + assert isinstance(msg, MsgEventTextChunk) + assert msg.text == "hello" + + +def test_model_dump_roundtrip(): + data = {"type": "ERROR", "code": "E1", "details": "fail"} + + msg = ServerMessage.validate_python(data) + dumped = msg.model_dump() + + assert dumped == data \ No newline at end of file diff --git a/uv.lock b/uv.lock index baed3b5..01160bd 100644 --- a/uv.lock +++ b/uv.lock @@ -92,6 +92,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, ] +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + [[package]] name = "frozenlist" version = "1.8.0" @@ -142,6 +151,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "lambda-agent-api" version = "0.1.0" @@ -149,12 +167,14 @@ source = { virtual = "." } dependencies = [ { name = "aiohttp" }, { name = "pydantic" }, + { name = "pytest" }, ] [package.metadata] requires-dist = [ { name = "aiohttp", specifier = ">=3.13.4" }, { name = "pydantic", specifier = ">=2.12.5" }, + { name = "pytest", specifier = ">=9.0.2" }, ] [[package]] @@ -202,6 +222,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/08/7036c080d7117f28a4af526d794aab6a84463126db031b007717c1a6676e/multidict-6.7.1-py3-none-any.whl", hash = "sha256:55d97cc6dae627efa6a6e548885712d4864b81110ac76fa4e534c03819fa4a56", size = 12319, upload-time = "2026-01-26T02:46:44.004Z" }, ] +[[package]] +name = "packaging" +version = "26.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/65/ee/299d360cdc32edc7d2cf530f3accf79c4fca01e96ffc950d8a52213bd8e4/packaging-26.0.tar.gz", hash = "sha256:00243ae351a257117b6a241061796684b084ed1c516a08c48a3f7e147a9d80b4", size = 143416, upload-time = "2026-01-21T20:50:39.064Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "propcache" version = "0.4.1" @@ -295,6 +333,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9f/ed/068e41660b832bb0b1aa5b58011dea2a3fe0ba7861ff38c4d4904c1c1a99/pydantic_core-2.41.5-cp314-cp314t-win_arm64.whl", hash = "sha256:35b44f37a3199f771c3eaa53051bc8a70cd7b54f333531c59e29fd4db5d15008", size = 1974769, upload-time = "2025-11-04T13:42:01.186Z" }, ] +[[package]] +name = "pygments" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, +] + +[[package]] +name = "pytest" +version = "9.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0"