diff --git a/README.md b/README.md index 496c876..50769bd 100644 --- a/README.md +++ b/README.md @@ -31,14 +31,6 @@ 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() @@ -47,8 +39,6 @@ async def main(): asyncio.run(main()) ``` -> `AgentApi.send_message()` возвращает стриминг-итерируемый объект, который может выдавать не только текстовые чанки, но и события инструментов (`MsgEventToolCallChunk`, `MsgEventToolResult`, `MsgEventCustomUpdate`) и финальный `MsgEventEnd`. - ## Предполагаемое использование ```python @@ -115,97 +105,37 @@ async def on_telegram_message(from_user: int, text: str): } ``` -#### AGENT_EVENT_TEXT_CHUNK +#### AGENT_EVENT + +Базовый класс для ивентов, которые стримит агент во время генерации ответа. Конкретный класс для ивента определяется по `subtype`. + +##### TEXT_CHUNK Чанк текста ответа агента. ```json { - "type": "AGENT_EVENT_TEXT_CHUNK", - "text": "Фрагмент текста", - "source": "main" + "type": "AGENT_EVENT", + "subtype": "TEXT_CHUNK", + "text": "Фрагмент текста" } ``` -| Поле | Тип | Описание | -|--------|--------|-----------------------------------------------| -| 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 +##### END Агент закончил генерацию ответа. ```json { - "type": "AGENT_EVENT_END", + "type": "AGENT_EVENT", + "subtype": "END", "tokens_used": 42 } ``` -| Поле | Тип | Описание | -|-------------|--------|-----------------------------------------------| -| type | string | Всегда `AGENT_EVENT_END` | -| tokens_used | int | Количество использованных токенов | +| Поле | Тип | Описание | +|-------------|--------|-----------------------| +| tokens_used | int | Количество использованных токенов | #### ERROR @@ -234,29 +164,4 @@ 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 d130bcf..544c090 100644 --- a/lambda_agent_api/agent_api.py +++ b/lambda_agent_api/agent_api.py @@ -226,19 +226,20 @@ class AgentApi: outgoing_msg = ServerMessage.validate_json( msg.data) - if isinstance(outgoing_msg, (MsgEventTextChunk, - MsgEventToolCallChunk, - MsgEventToolResult, - MsgEventCustomUpdate, - MsgEventEnd)): + if isinstance(outgoing_msg, MsgEventTextChunk): 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/tests/test_models.py b/tests/test_models.py index bbb91c9..005e56a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -33,9 +33,6 @@ 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), @@ -50,8 +47,6 @@ 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"},