119 lines
4.1 KiB
Python
119 lines
4.1 KiB
Python
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)
|
||
"""
|