diff --git a/README.md b/README.md index 23841e6..ccc0115 100644 --- a/README.md +++ b/README.md @@ -3,11 +3,12 @@ WebSocket API SDK для взаимодействия с AI-агентом. ## Release Notes -# v1.1 +# v1.1 - **CRITICAL**: `AgentAPI` вместо `url` принимает `base_url` и сам дописывает нужный эндпоинт. -Раньше: `AgentAPI(url="ws://localhost:8000/agent_ws/")`. Сейчас: `AgentAPI(base_url="ws://localhost:8000/")` + Раньше: `AgentAPI(url="ws://localhost:8000/agent_ws/")`. Сейчас: `AgentAPI(base_url="ws://localhost:8000/")` - Добавлен параметр `chat_id` в конструктор `AgentAPI`. Нужен для разделения истории сообщений по чатам/веткам. - `AgentAPI.connect()` вызывает `AgentBusyException`, если выбранный чат уже занят другим API клиентом. +- Добавлен новый тип события `MsgEventSendFile` для отправки файлов пользователю. Поле `path` — путь к файлу относительно `/workspace`. ## Установка В `master` всегда будет актуальная рабочая версия. @@ -164,6 +165,22 @@ async def on_telegram_message(from_user: int, text: str): | payload | object | Любые данные о прогрессе | | source | string | Источник события (по умолчанию "main") | +#### AGENT_EVENT_SEND_FILE + +Агент отправляет файл пользователю. Путь к файлу относительно `/workspace`. + +```json +{ + "type": "AGENT_EVENT_SEND_FILE", + "path": "reports/2024/report.pdf" +} +``` + +| Поле | Тип | Описание | +|--------|--------|-----------------------------------------------| +| type | string | Всегда `AGENT_EVENT_SEND_FILE` | +| path | string | Путь к файлу относительно `/workspace` | + #### AGENT_EVENT_END Агент закончил генерацию ответа. diff --git a/lambda_agent_api/agent_api.py b/lambda_agent_api/agent_api.py index e22a0b8..9dd4608 100644 --- a/lambda_agent_api/agent_api.py +++ b/lambda_agent_api/agent_api.py @@ -30,8 +30,8 @@ class AgentApi: agent_id: str, base_url: str, callback: Optional[Callable[[ServerMessage], None]] = None, - on_disconnect: Optional[Callable[['AgentApi'], None]] = None, - chat_id: int = 0 # значение по умолчанию для обратной совместимости + on_disconnect: Optional[Callable[["AgentApi"], None]] = None, + chat_id: int = 0, # значение по умолчанию для обратной совместимости # значение по умолчанию для обратной совместимости ): self.id = agent_id # ID агента для словаря self.chat_id = chat_id @@ -67,10 +67,12 @@ class AgentApi: logger.info(f"Agent {self.id} is ready") else: raise AgentException( - "INVALID_STATUS", f"Expected SM.Status, got {status_msg.type}") + "INVALID_STATUS", f"Expected SM.Status, got {status_msg.type}" + ) else: - raise AgentException("UNEXPECTED_MSG_TYPE", - f"Unexpected message type: {msg.type}") + raise AgentException( + "UNEXPECTED_MSG_TYPE", f"Unexpected message type: {msg.type}" + ) self._connected = True self._listen_task = asyncio.create_task(self._listen()) @@ -83,7 +85,8 @@ class AgentApi: await self._session.close() raise AgentException( - "TIMEOUT", "Agent did not send initial Status message within 5 seconds") from e + "TIMEOUT", "Agent did not send initial Status message within 5 seconds" + ) from e except AgentException: # Перехватываем наши собственные ошибки (INVALID_STATUS, UNEXPECTED_MSG_TYPE), @@ -121,8 +124,10 @@ class AgentApi: await self._session.close() # Можно оставить RuntimeError, а можно тоже завернуть в AgentException - raise AgentException(code="CONNECTION_ERROR", - details=f"Failed to connect agent {self.id}: {e}") from e + raise AgentException( + code="CONNECTION_ERROR", + details=f"Failed to connect agent {self.id}: {e}", + ) from e async def close(self): """Явное ручное закрытие соединения.""" @@ -161,17 +166,17 @@ class AgentApi: async def send_message(self, text: str) -> AsyncIterator[AgentEventUnion]: """ - Нативный асинхронный генератор. + Нативный асинхронный генератор. Не требует отдельного класса ResponseIterator. Гарантированно освобождает блокировку. """ if not self._connected or not self._ws: - raise AgentException(code="NOT_CONNECTED", - details="Not connected. Call connect() first.") + raise AgentException( + code="NOT_CONNECTED", details="Not connected. Call connect() first." + ) if self._request_lock.locked(): - raise AgentBusyException( - "Agent is currently processing another request") + raise AgentBusyException("Agent is currently processing another request") # Блокируем параллельные запросы # если идет стриминг ответа, то при попытки отправить новое сообщение будет ошибка - ее рейзим(делаем AgentBusyError) @@ -180,10 +185,7 @@ class AgentApi: try: self._current_queue = asyncio.Queue() - message = MsgUserMessage( - type=EClientMessage.USER_MESSAGE, - text=text - ) + message = MsgUserMessage(type=EClientMessage.USER_MESSAGE, text=text) await self._ws.send_str(message.model_dump_json()) logger.debug(f"[{self.id}] Sent message: {text[:50]}...") @@ -226,7 +228,8 @@ class AgentApi: # Если это исключение, просто логируем, в коллбек его кидать не стоит if isinstance(orphan_msg, Exception): logger.debug( - f"[{self.id}] Dropped exception from queue during cleanup: {orphan_msg}") + f"[{self.id}] Dropped exception from queue during cleanup: {orphan_msg}" + ) continue # 3. Отправляем "мусорные/осиротевшие" куски в callback @@ -234,7 +237,8 @@ class AgentApi: self.callback(orphan_msg) else: logger.debug( - f"[{self.id}] Dropped orphaned message during cleanup") + f"[{self.id}] Dropped orphaned message during cleanup" + ) except asyncio.QueueEmpty: break @@ -251,27 +255,34 @@ class AgentApi: async for msg in self._ws: if msg.type == aiohttp.WSMsgType.TEXT: try: - outgoing_msg = ServerMessage.validate_json( - msg.data) + outgoing_msg = ServerMessage.validate_json(msg.data) - if isinstance(outgoing_msg, (MsgEventTextChunk, - MsgEventToolCallChunk, - MsgEventToolResult, - MsgEventCustomUpdate, - MsgEventEnd)): + if isinstance( + outgoing_msg, + ( + MsgEventTextChunk, + MsgEventToolCallChunk, + MsgEventToolResult, + MsgEventCustomUpdate, + MsgEventSendFile, + 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") + f"[{self.id}] AgentEvent without active request" + ) elif isinstance(outgoing_msg, MsgError): if self.callback: self.callback(outgoing_msg) error = AgentException( - outgoing_msg.code, outgoing_msg.details) + outgoing_msg.code, outgoing_msg.details + ) logger.error(f"[{self.id}] Agent error: {error}") if self._current_queue: await self._current_queue.put(error) @@ -279,8 +290,7 @@ class AgentApi: elif isinstance(outgoing_msg, MsgGracefulDisconnect): if self.callback: self.callback(outgoing_msg) - logger.info( - f"[{self.id}] Gracefully disconnecting") + logger.info(f"[{self.id}] Gracefully disconnecting") break # Выход из цикла приведет к finally -> _cleanup else: @@ -292,17 +302,14 @@ class AgentApi: self.callback(outgoing_msg) except Exception as e: - logger.error( - f"[{self.id}] Failed to deserialize message: {e}") + logger.error(f"[{self.id}] Failed to deserialize message: {e}") if self._current_queue: await self._current_queue.put( - AgentException( - "PARSE_ERROR", f"Validation failed: {e}") + AgentException("PARSE_ERROR", f"Validation failed: {e}") ) elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED): - logger.error( - f"[{self.id}] WebSocket closed/error: {msg.type}") + logger.error(f"[{self.id}] WebSocket closed/error: {msg.type}") break except asyncio.CancelledError: 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) diff --git a/tests/manual.py b/tests/manual.py index 10f3005..44adea2 100644 --- a/tests/manual.py +++ b/tests/manual.py @@ -2,8 +2,12 @@ import asyncio import traceback from lambda_agent_api.agent_api import AgentApi, AgentBusyException -from lambda_agent_api.server import MsgEventTextChunk, MsgEventToolCallChunk, MsgEventToolResult - +from lambda_agent_api.server import ( + MsgEventTextChunk, + MsgEventToolCallChunk, + MsgEventToolResult, + MsgEventSendFile, +) async def main(): @@ -28,12 +32,18 @@ async def main(): print(chunk.text, end="", flush=True) case MsgEventToolCallChunk(): if not is_tool: - print(f"\n\n### TOOL CALL: ({chunk.tool_name}) ", end="", flush=True) + print( + f"\n\n### TOOL CALL: ({chunk.tool_name}) ", + end="", + flush=True, + ) is_tool = True print(chunk.args_chunk, end="", flush=True) case MsgEventToolResult(): is_tool = False print(f"\nResult: {chunk.result}\n\n", end="", flush=True) + case MsgEventSendFile(): + print(f"\n[SEND FILE] {chunk.path}\n", end="", flush=True) print("\n") except KeyboardInterrupt: @@ -44,4 +54,4 @@ async def main(): await api.close() -asyncio.run(main()) \ No newline at end of file +asyncio.run(main())