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

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