корректные pydantic модели для автоматического определения класса по полю type

README.md

@@ -1,7 +1,238 @@
-## Contributing
+# Lambda Agent API
 
-- В качестве таск-трекера - раздел "Задачи" (Issues).
+WebSocket API SDK для взаимодействия с AI-агентом.
-- В названии коммита обязательно указывать ID задачи. Пример: "#1 Описание коммита"
+
-- При закрытии задачи в комментарии делать ссылку на PR.
+## Установка
-- Одна задача, одна ветка, один PR.
+
-- Формат названия ветки: #<ID задачи>-<короткое описание>. Пример: "#1-task-description".
+```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.MsgError):
+        print(f"\n[Ошибка: {message.code}] {message.details}")
+    elif isinstance(message, OM.MsgStatus):
+        print("✓ Agent status update")
+    elif isinstance(message, OM.MsgGracefulDisconnect):
+        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.MsgEventTextChunk):
+                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.MsgEventTextChunk):
+            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.MsgStatus):
+        print("Agent status update")
+    elif isinstance(message, OM.MsgError):
+        print(f"Agent error: {message.code} - {message.details}")
+    elif isinstance(message, OM.MsgGracefulDisconnect):
+        print("Agent disconnecting gracefully")
+```
+
+## Классический подход (низкоуровневый)
+
+```python
+import asyncio
+import websockets
+from lambda_agent_api.server 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.MsgEventEnd():
+                    print(f"\n[Завершено, использовано токенов: {data.tokens_used}]")
+                    break
+                case OM.MsgError():
+                    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

api/README.md

@@ -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

api/models.py

@@ -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)
-"""

lambda_agent_api/__init__.py

@@ -7,14 +7,13 @@
     - IM, OM, IncomingMessage, OutgoingMessage: Pydantic модели контракта
 """
 
-from .agent_api import AgentApi, AgentException
+from lambda_agent_api.agent_api import AgentApi, AgentException
-from .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",
 ]

lambda_agent_api/agent_api.py

@@ -2,7 +2,9 @@ import logging
 from typing import Callable, Optional, AsyncIterator
 import aiohttp
 import asyncio
-from models import CM, SM, ClientMessage, ServerMessage
+
+from lambda_agent_api.server import *
+from lambda_agent_api.client import *
 
 logger = logging.getLogger(__name__)
 
@@ -51,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)
+                status_msg = ServerMessage.validate_json(msg.data)
-                if isinstance(status_msg, SM.Status):
+                if isinstance(status_msg, MsgStatus):
                     logger.info(f"Agent {self.id} is ready")
                 else:
                     raise AgentException(
@@ -129,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.
@@ -150,8 +152,8 @@ class AgentApi:
         try:
             self._current_queue = asyncio.Queue()
 
-            message = CM.UserMessage(
+            message = MsgUserMessage(
-                type=CM.Type.USER_MESSAGE,
+                type=EClientMessage.USER_MESSAGE,
                 text=text
             )
 
@@ -167,7 +169,7 @@ class AgentApi:
                     raise chunk
 
                 # Если конец ответа
-                if isinstance(chunk, SM.EventEnd):
+                if isinstance(chunk, MsgEventEnd):
                     # Если вам нужны tokens_used, можно сохранить их в атрибут self
                     break
 
@@ -221,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)
                             # Если очереди нет (клиент отменил запрос), но токены идут — шлем их в коллбек
@@ -234,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(
@@ -247,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(

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)
+"""

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)
+"""

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"
@@ -7,4 +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"]

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())

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

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"
@@ -107,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"
@@ -157,6 +151,32 @@ 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"
+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]]
 name = "multidict"
 version = "6.7.1"
@@ -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"