#7 Добавить основные модели ивентов агента. Добавлены основные инветы

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)
-"""
+Объединяет все типы исходящих сообщений в одно для удобной автоматической десериализации.
+Pydantic сам определит нужный тип в зависимости от поля `type`.
+Использование:
+msg = ServerMessage.validate_json(json_str)
+"""