72 lines
No EOL
2.8 KiB
Python
72 lines
No EOL
2.8 KiB
Python
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)
|
||
""" |