модели передаваемых сообщений

api/models.py

@@ -0,0 +1,119 @@
+from pydantic import BaseModel, Field
+from enum import Enum
+from typing import Literal, Annotated, Union
+
+
+class IM:
+    """
+    Namespace для моделей входящих сообщений (от клиента к серверу).\n
+    IM = Incoming Message
+    """
+
+    class Type(str, Enum):
+        USER_MESSAGE = "USER_MESSAGE"
+
+
+    # noinspection PyPep8Naming
+    class UserMessage(BaseModel):
+        """
+        Полное сообщение от пользователя.
+        """
+        type: Literal[IM.Type.USER_MESSAGE]
+        text: str
+        """
+        Текст сообщения.
+        """
+
+
+IncomingMessage = Annotated[
+    Union[IM.UserMessage,],
+    Field(discriminator="type")
+]
+"""
+Объединяет все типы входящих сообщений в одно для удобной автоматической десериализации.\n
+Pydantic сам определит нужный тип в зависимости от поля ``type``.\n
+Использование:\n
+msg = IncomingMessage.model_validate_json(json)
+"""
+
+
+class OM:
+    """
+    Namespace для моделей исходящих сообщений (от сервера к клиенту).\n
+    OM = Outgoing Message
+    """
+
+    class Type(str, Enum):
+        STATUS = "STATUS"
+        AGENT_EVENT = "AGENT_EVENT"
+        ERROR = "ERROR"
+        GRACEFUL_DISCONNECT = "GRACEFUL_DISCONNECT"
+
+
+    class Status(BaseModel):
+        """
+        Отправляется сервером при открытии соединения с клиентом.
+        Будет дополнен информацией о готовности агента принимать сообщения.
+        """
+        type: Literal[OM.Type.STATUS]
+
+
+    class AgentEventType(str, Enum):
+        TEXT_CHUNK = "TEXT_CHUNK"
+        END = "END"
+
+
+    class AgentEvent(BaseModel):
+        """
+        Базовый класс для ивентов, которые стримит агент во время генерации ответа.
+        Конкретный класс для ивента определяется по ``subtype``.
+        """
+        type: Literal[OM.Type.AGENT_EVENT]
+        subtype: OM.AgentEventType
+
+
+    class EventTextChunk(AgentEvent):
+        """
+        Чанк текста ответа агента.
+        """
+        subtype: Literal[OM.AgentEventType.TEXT_CHUNK]
+        text: str
+
+
+    class EventEnd(AgentEvent):
+        """
+        Агент закончил генерацию ответа.
+        """
+        subtype: Literal[OM.AgentEventType.END]
+        tokens_used: int
+
+
+    class Error(BaseModel):
+        """
+        Неопределенная ошибка в работе агента.
+        """
+        type: Literal[OM.Type.ERROR]
+        code: str
+        details: str
+
+
+    class GracefulDisconnect(BaseModel):
+        """
+        Отправляется перед завершением работы контейнера с агентом. Например, при долгом бездействии.
+        Нужно, чтобы отделять обрыв соединения из-за ошибки с необходимостью повторного подключения.
+        Приход этого сообщения означает, что агент осознанно завершает работу с клиентом по какой-то причине.
+        Для дальнейшего взаимодействия нужно снова обратиться к мастеру.
+        """
+        type: Literal[OM.Type.GRACEFUL_DISCONNECT]
+
+
+OutgoingMessage = Annotated[
+    Union[OM.Status, OM.AgentEvent, OM.Error, OM.GracefulDisconnect],
+    Field(discriminator="type")
+]
+"""
+Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.\n
+Pydantic сам определит нужный тип в зависимости от поля ``type``.\n
+Использование:\n
+msg = OutgoingMessage.model_validate_json(json)
+"""