platform/agent_api
5
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch '#13-file-transfer'

Browse source 
# Conflicts:
#	README.md
#	lambda_agent_api/agent_api.py
#	tests/manual.py
This commit is contained in:
Егор Кандрушин 2026-04-19 22:08:48 +03:00
commit 483a63b999
4 changed files with 144 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

79
lambda_agent_api/agent_api.py
View file

@ -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: