diff --git a/lambda_agent_api/server.py b/lambda_agent_api/server.py index dbeaaa5..e5882f7 100644 --- a/lambda_agent_api/server.py +++ b/lambda_agent_api/server.py @@ -1,72 +1,120 @@ from pydantic import BaseModel, Field, TypeAdapter from enum import Enum -from typing import Literal, Annotated, Union +from typing import Literal, Annotated, Union, Any, Dict, Optional -__all__ = ['EServerMessage', 'MsgStatus', 'MsgError', 'MsgEventTextChunk', 'MsgEventEnd', 'AgentEventUnion', 'ServerMessage'] +__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 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] +# ------------------------------------------------------------------ +# 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, MsgEventTextChunk, MsgEventEnd, MsgError, MsgGracefulDisconnect], + Union[ + MsgStatus, + MsgError, + MsgGracefulDisconnect, + MsgEventTextChunk, + MsgEventToolCallChunk, + MsgEventToolResult, + MsgEventCustomUpdate, + MsgEventEnd + ], Field(discriminator="type") ]) """ -Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.\n -Pydantic сам определит нужный тип в зависимости от поля ``type``.\n -Использование:\n -msg = ServerMessage.model_validate_json(json) -""" \ No newline at end of file +Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации. +Pydantic сам определит нужный тип в зависимости от поля `type`. +Использование: +msg = ServerMessage.validate_json(json_str) +"""