diff --git a/README.md b/README.md index 50769bd..496c876 100644 --- a/README.md +++ b/README.md @@ -31,6 +31,14 @@ async def main(): async for chunk in response: if isinstance(chunk, MsgEventTextChunk): print(chunk.text, end="", flush=True) + elif isinstance(chunk, MsgEventToolCallChunk): + print(f"Tool call started: {chunk.tool_name}") + elif isinstance(chunk, MsgEventToolResult): + print(f"Tool result: {chunk.result}") + elif isinstance(chunk, MsgEventCustomUpdate): + print(f"Progress update: {chunk.payload}") + elif isinstance(chunk, MsgEventEnd): + print(f"Generation ended, tokens used: {chunk.tokens_used}") finally: await api.close() @@ -39,6 +47,8 @@ async def main(): asyncio.run(main()) ``` +> `AgentApi.send_message()` возвращает стриминг-итерируемый объект, который может выдавать не только текстовые чанки, но и события инструментов (`MsgEventToolCallChunk`, `MsgEventToolResult`, `MsgEventCustomUpdate`) и финальный `MsgEventEnd`. + ## Предполагаемое использование ```python @@ -105,37 +115,97 @@ async def on_telegram_message(from_user: int, text: str): } ``` -#### AGENT_EVENT - -Базовый класс для ивентов, которые стримит агент во время генерации ответа. Конкретный класс для ивента определяется по `subtype`. - -##### TEXT_CHUNK +#### AGENT_EVENT_TEXT_CHUNK Чанк текста ответа агента. ```json { - "type": "AGENT_EVENT", - "subtype": "TEXT_CHUNK", - "text": "Фрагмент текста" + "type": "AGENT_EVENT_TEXT_CHUNK", + "text": "Фрагмент текста", + "source": "main" } ``` -##### END +| Поле | Тип | Описание | +|--------|--------|-----------------------------------------------| +| type | string | Всегда `AGENT_EVENT_TEXT_CHUNK` | +| text | string | Фрагмент текста ответа агента | +| source | string | Источник события (по умолчанию "main") | + +#### AGENT_EVENT_TOOL_CALL_CHUNK + +Агент решил использовать инструмент и генерирует аргументы. + +```json +{ + "type": "AGENT_EVENT_TOOL_CALL_CHUNK", + "tool_name": "имя_инструмента", + "args_chunk": "{\"key\": \"value\"}", + "source": "main" +} +``` + +| Поле | Тип | Описание | +|-------------|---------|-----------------------------------------------| +| type | string | Всегда `AGENT_EVENT_TOOL_CALL_CHUNK` | +| tool_name | string | Имя инструмента (может быть null в первом чанке) | +| args_chunk | string | Кусок JSON-аргументов (может быть null) | +| source | string | Источник события (по умолчанию "main") | + +#### AGENT_EVENT_TOOL_RESULT + +Инструмент отработал и вернул результат. + +```json +{ + "type": "AGENT_EVENT_TOOL_RESULT", + "tool_name": "имя_инструмента", + "result": "результат выполнения", + "source": "main" +} +``` + +| Поле | Тип | Описание | +|------------|--------|-----------------------------------------------| +| type | string | Всегда `AGENT_EVENT_TOOL_RESULT` | +| tool_name | string | Имя инструмента | +| result | any | Результат выполнения (строка, объект или массив) | +| source | string | Источник события (по умолчанию "main") | + +#### AGENT_EVENT_CUSTOM_UPDATE + +Кастомный прогресс (например, скачивание файла) изнутри инструмента. + +```json +{ + "type": "AGENT_EVENT_CUSTOM_UPDATE", + "payload": {"status": "in_progress", "progress": 50}, + "source": "main" +} +``` + +| Поле | Тип | Описание | +|----------|-----------------|-----------------------------------------------| +| type | string | Всегда `AGENT_EVENT_CUSTOM_UPDATE` | +| payload | object | Любые данные о прогрессе | +| source | string | Источник события (по умолчанию "main") | + +#### AGENT_EVENT_END Агент закончил генерацию ответа. ```json { - "type": "AGENT_EVENT", - "subtype": "END", + "type": "AGENT_EVENT_END", "tokens_used": 42 } ``` -| Поле | Тип | Описание | -|-------------|--------|-----------------------| -| tokens_used | int | Количество использованных токенов | +| Поле | Тип | Описание | +|-------------|--------|-----------------------------------------------| +| type | string | Всегда `AGENT_EVENT_END` | +| tokens_used | int | Количество использованных токенов | #### ERROR @@ -164,4 +234,29 @@ async def on_telegram_message(from_user: int, text: str): } ``` +Неопределенная ошибка в работе агента. + +```json +{ + "type": "ERROR", + "code": "error_code", + "details": "Описание ошибки" +} +``` + +| Поле | Тип | Описание | +|---------|-------|----------------| +| code | string | Код ошибки | +| details | string | Подробности | + +#### GRACEFUL_DISCONNECT + +Отправляется перед завершением работы контейнера с агентом. Например, при долгом бездействии. Нужно, чтобы отделять обрыв соединения из-за ошибки с необходимостью повторного подключения. Приход этого сообщения означает, что агент осознанно завершает работу с клиентом по какой-то причине. Для дальнейшего взаимодействия нужно снова обратиться к мастеру. + +```json +{ + "type": "GRACEFUL_DISCONNECT" +} +``` + ![Схема взаимодействия](docs/schema.png) diff --git a/lambda_agent_api/agent_api.py b/lambda_agent_api/agent_api.py index 544c090..d130bcf 100644 --- a/lambda_agent_api/agent_api.py +++ b/lambda_agent_api/agent_api.py @@ -226,20 +226,19 @@ class AgentApi: outgoing_msg = ServerMessage.validate_json( msg.data) - if isinstance(outgoing_msg, MsgEventTextChunk): + if isinstance(outgoing_msg, (MsgEventTextChunk, + MsgEventToolCallChunk, + MsgEventToolResult, + MsgEventCustomUpdate, + MsgEventEnd)): if self._current_queue: await self._current_queue.put(outgoing_msg) - # Если очереди нет (клиент отменил запрос), но токены идут — шлем их в коллбек elif self.callback: self.callback(outgoing_msg) else: logger.warning( f"[{self.id}] AgentEvent without active request") - elif isinstance(outgoing_msg, MsgEventEnd): - if self._current_queue: - await self._current_queue.put(outgoing_msg) - elif isinstance(outgoing_msg, MsgError): if self.callback: self.callback(outgoing_msg) 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) +""" diff --git a/tests/test_models.py b/tests/test_models.py index 005e56a..bbb91c9 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -33,6 +33,9 @@ def test_client_message_invalid(data): [ ({"type": "STATUS"}, MsgStatus), ({"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hi"}, MsgEventTextChunk), + ({"type": "AGENT_EVENT_TOOL_CALL_CHUNK", "tool_name": "search", "args_chunk": "{\"q\": \"hello\"}"}, MsgEventToolCallChunk), + ({"type": "AGENT_EVENT_TOOL_RESULT", "tool_name": "search", "result": {"items": [1, 2, 3]}}, MsgEventToolResult), + ({"type": "AGENT_EVENT_CUSTOM_UPDATE", "payload": {"status": "in_progress", "progress": 50}}, MsgEventCustomUpdate), ({"type": "AGENT_EVENT_END", "tokens_used": 10}, MsgEventEnd), ({"type": "ERROR", "code": "E1", "details": "fail"}, MsgError), ({"type": "GRACEFUL_DISCONNECT"}, MsgGracefulDisconnect), @@ -47,6 +50,8 @@ def test_server_message_valid(data, expected_type): "data", [ {"type": "AGENT_EVENT_TEXT_CHUNK"}, # нет text + {"type": "AGENT_EVENT_TOOL_RESULT", "tool_name": "search"}, # нет result + {"type": "AGENT_EVENT_CUSTOM_UPDATE"}, # нет payload {"type": "AGENT_EVENT_END"}, # нет tokens_used {"type": "ERROR", "code": "E1"}, # нет details {"type": "UNKNOWN"},