diff --git a/lambda_agent_api/server.py b/lambda_agent_api/server.py index e5882f7..c0f01df 100644 --- a/lambda_agent_api/server.py +++ b/lambda_agent_api/server.py @@ -4,10 +4,18 @@ from typing import Literal, Annotated, Union, Any, Dict, Optional __all__ = [ - 'EServerMessage', 'MsgStatus', 'MsgError', 'MsgGracefulDisconnect', - 'MsgEventTextChunk', 'MsgEventToolCallChunk', 'MsgEventToolResult', - 'MsgEventCustomUpdate', 'MsgEventEnd', - 'AgentEventUnion', 'ServerMessage' + "EServerMessage", + "MsgStatus", + "MsgError", + "MsgGracefulDisconnect", + "MsgEventTextChunk", + "MsgEventToolCallChunk", + "MsgEventToolResult", + "MsgEventCustomUpdate", + "MsgEventSendFile", + "MsgEventEnd", + "AgentEventUnion", + "ServerMessage", ] @@ -19,18 +27,21 @@ class EServerMessage(str, Enum): # Ивенты агента 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_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 @@ -38,16 +49,23 @@ class MsgError(BaseModel): class MsgGracefulDisconnect(BaseModel): """Отправляется перед завершением работы контейнера с агентом.""" - type: Literal[EServerMessage.GRACEFUL_DISCONNECT] = EServerMessage.GRACEFUL_DISCONNECT + + 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 + + type: Literal[EServerMessage.AGENT_EVENT_TEXT_CHUNK] = ( + EServerMessage.AGENT_EVENT_TEXT_CHUNK + ) text: str # Новое: "main" (главный агент) или "tools:..." (субагент, если будем использовать) source: str = "main" # пока везде будет main @@ -55,17 +73,23 @@ class MsgEventTextChunk(BaseModel): class MsgEventToolCallChunk(BaseModel): """Агент решил использовать инструмент и генерирует аргументы.""" - type: Literal[EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK] = EServerMessage.AGENT_EVENT_TOOL_CALL_CHUNK + + 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-аргументов") + 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 + + type: Literal[EServerMessage.AGENT_EVENT_TOOL_RESULT] = ( + EServerMessage.AGENT_EVENT_TOOL_RESULT + ) tool_name: str result: Any # Может быть строкой, словарем или списком source: str = "main" @@ -73,14 +97,28 @@ class MsgEventToolResult(BaseModel): class MsgEventCustomUpdate(BaseModel): """Кастомный прогресс (например, скачивание файла) изнутри инструмента.""" - type: Literal[EServerMessage.AGENT_EVENT_CUSTOM_UPDATE] = EServerMessage.AGENT_EVENT_CUSTOM_UPDATE + + type: Literal[EServerMessage.AGENT_EVENT_CUSTOM_UPDATE] = ( + EServerMessage.AGENT_EVENT_CUSTOM_UPDATE + ) payload: Dict[str, Any] = Field( - ..., description="Любые данные о прогрессе (status, progress и т.д.)") + ..., 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 @@ -95,25 +133,24 @@ AgentEventUnion = Union[ MsgEventToolCallChunk, MsgEventToolResult, MsgEventCustomUpdate, - MsgEventEnd + MsgEventSendFile, + MsgEventEnd, ] -# Обновлено: добавили новые модели в Union адаптера -ServerMessage = TypeAdapter(Annotated[ - Union[ - MsgStatus, - MsgError, - MsgGracefulDisconnect, - MsgEventTextChunk, - MsgEventToolCallChunk, - MsgEventToolResult, - MsgEventCustomUpdate, - MsgEventEnd - ], - Field(discriminator="type") -]) +# ServerMessage использует AgentEventUnion + остальные типы +ServerMessage = TypeAdapter( + Annotated[ + Union[ + MsgStatus, + MsgError, + MsgGracefulDisconnect, + AgentEventUnion, + ], + Field(discriminator="type"), + ] +) """ -Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации. +TypeAdapter для десериализации всех входящих сообщений (от сервера). Pydantic сам определит нужный тип в зависимости от поля `type`. Использование: msg = ServerMessage.validate_json(json_str)