адаптация для установки пакетом

lambda_agent_api/models.py

@@ -0,0 +1,121 @@
+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)
+"""