157 lines
5.1 KiB
Python
157 lines
5.1 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",
|
||
"MsgEventSendFile",
|
||
"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_SEND_FILE = "AGENT_EVENT_SEND_FILE" # Новое
|
||
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 MsgEventSendFile(BaseModel):
|
||
"""Агент отправляет файл пользователю."""
|
||
|
||
type: Literal[EServerMessage.AGENT_EVENT_SEND_FILE] = (
|
||
EServerMessage.AGENT_EVENT_SEND_FILE
|
||
)
|
||
path: str = Field(..., description="Путь к файлу относительно /workspace")
|
||
|
||
|
||
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,
|
||
MsgEventSendFile,
|
||
MsgEventEnd,
|
||
]
|
||
|
||
# ServerMessage использует AgentEventUnion + остальные типы
|
||
ServerMessage = TypeAdapter(
|
||
Annotated[
|
||
Union[
|
||
MsgStatus,
|
||
MsgError,
|
||
MsgGracefulDisconnect,
|
||
AgentEventUnion,
|
||
],
|
||
Field(discriminator="type"),
|
||
]
|
||
)
|
||
"""
|
||
TypeAdapter для десериализации всех входящих сообщений (от сервера).
|
||
Pydantic сам определит нужный тип в зависимости от поля `type`.
|
||
Использование:
|
||
msg = ServerMessage.validate_json(json_str)
|
||
"""
|