120 lines
4.7 KiB
Python
120 lines
4.7 KiB
Python
from pydantic import BaseModel, Field, TypeAdapter
|
||
from enum import Enum
|
||
from typing import Literal, Annotated, Union, Any, Dict, Optional
|
||
|
||
|
||
__all__ = [
|
||
'EServerMessage', 'MsgStatus', 'MsgError', 'MsgGracefulDisconnect',
|
||
'MsgEventTextChunk', 'MsgEventToolCallChunk', 'MsgEventToolResult',
|
||
'MsgEventCustomUpdate', '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_TOOL_CALL_CHUNK = "AGENT_EVENT_TOOL_CALL_CHUNK" # Новое
|
||
AGENT_EVENT_TOOL_RESULT = "AGENT_EVENT_TOOL_RESULT" # Новоеы
|
||
AGENT_EVENT_CUSTOM_UPDATE = "AGENT_EVENT_CUSTOM_UPDATE" # Новое
|
||
AGENT_EVENT_END = "AGENT_EVENT_END"
|
||
|
||
|
||
class MsgStatus(BaseModel):
|
||
"""Отправляется сервером при открытии соединения с клиентом."""
|
||
type: Literal[EServerMessage.STATUS] = EServerMessage.STATUS
|
||
|
||
|
||
class MsgError(BaseModel):
|
||
"""Неопределенная ошибка в работе агента."""
|
||
type: Literal[EServerMessage.ERROR] = EServerMessage.ERROR
|
||
code: str
|
||
details: str
|
||
|
||
|
||
class MsgGracefulDisconnect(BaseModel):
|
||
"""Отправляется перед завершением работы контейнера с агентом."""
|
||
type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = EServerMessage.GRACEFUL_DISCONNECT
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# AGENT EVENTS (События генерации)
|
||
# ------------------------------------------------------------------
|
||
|
||
class MsgEventTextChunk(BaseModel):
|
||
"""Чанк текста ответа агента."""
|
||
type: Literal[EServerMessage.AGENT_EVENT_TEXT_CHUNK] = EServerMessage.AGENT_EVENT_TEXT_CHUNK
|
||
text: str
|
||
# Новое: "main" (главный агент) или "tools:..." (субагент, если будем использовать)
|
||
source: str = "main" # пока везде будет main
|
||
|
||
|
||
class MsgEventToolCallChunk(BaseModel):
|
||
"""Агент решил использовать инструмент и генерирует аргументы."""
|
||
type: Literal[EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK] = EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK
|
||
tool_name: Optional[str] = Field(
|
||
None, description="Имя инструмента (приходит обычно в первом чанке)")
|
||
args_chunk: Optional[str] = Field(
|
||
None, description="Кусок JSON-аргументов")
|
||
source: str = "main"
|
||
|
||
|
||
class MsgEventToolResult(BaseModel):
|
||
"""Инструмент отработал и вернул результат."""
|
||
type: Literal[EServerMessage.AGENT_EVENT_TOOL_RESULT] = EServerMessage.AGENT_EVENT_TOOL_RESULT
|
||
tool_name: str
|
||
result: Any # Может быть строкой, словарем или списком
|
||
source: str = "main"
|
||
|
||
|
||
class MsgEventCustomUpdate(BaseModel):
|
||
"""Кастомный прогресс (например, скачивание файла) изнутри инструмента."""
|
||
type: Literal[EServerMessage.AGENT_EVENT_CUSTOM_UPDATE] = EServerMessage.AGENT_EVENT_CUSTOM_UPDATE
|
||
payload: Dict[str, Any] = Field(
|
||
..., description="Любые данные о прогрессе (status, progress и т.д.)")
|
||
source: str = "main"
|
||
|
||
|
||
class MsgEventEnd(BaseModel):
|
||
"""Агент закончил генерацию ответа."""
|
||
type: Literal[EServerMessage.AGENT_EVENT_END] = EServerMessage.AGENT_EVENT_END
|
||
tokens_used: int
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# UNIONS & ADAPTERS
|
||
# ------------------------------------------------------------------
|
||
|
||
# Обновлено: добавили новые модели в Union
|
||
AgentEventUnion = Union[
|
||
MsgEventTextChunk,
|
||
MsgEventToolCallChunk,
|
||
MsgEventToolResult,
|
||
MsgEventCustomUpdate,
|
||
MsgEventEnd
|
||
]
|
||
|
||
# Обновлено: добавили новые модели в Union адаптера
|
||
ServerMessage = TypeAdapter(Annotated[
|
||
Union[
|
||
MsgStatus,
|
||
MsgError,
|
||
MsgGracefulDisconnect,
|
||
MsgEventTextChunk,
|
||
MsgEventToolCallChunk,
|
||
MsgEventToolResult,
|
||
MsgEventCustomUpdate,
|
||
MsgEventEnd
|
||
],
|
||
Field(discriminator="type")
|
||
])
|
||
"""
|
||
Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.
|
||
Pydantic сам определит нужный тип в зависимости от поля `type`.
|
||
Использование:
|
||
msg = ServerMessage.validate_json(json_str)
|
||
"""
|