315 lines
12 KiB
Python
315 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import logging
|
|
import os
|
|
import re
|
|
import sys
|
|
from collections.abc import AsyncIterator
|
|
from pathlib import Path
|
|
from urllib.parse import urlsplit, urlunsplit
|
|
|
|
import aiohttp
|
|
|
|
_api_root = Path(__file__).resolve().parents[1] / "external" / "platform-agent_api"
|
|
if str(_api_root) not in sys.path:
|
|
sys.path.insert(0, str(_api_root))
|
|
|
|
from lambda_agent_api.agent_api import AgentApi, AgentBusyException, AgentException # noqa: E402
|
|
from lambda_agent_api.client import EClientMessage, MsgUserMessage # noqa: E402
|
|
from lambda_agent_api.server import AgentEventUnion, MsgEventEnd, ServerMessage # noqa: E402
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_DEBUG_STREAM = os.environ.get("SURFACES_AGENT_DEBUG_STREAM", "").strip().lower() in {
|
|
"1",
|
|
"true",
|
|
"yes",
|
|
}
|
|
_POST_END_DRAIN_MS = int(os.environ.get("SURFACES_AGENT_POST_END_DRAIN_MS", "120"))
|
|
_STREAM_IDLE_TIMEOUT_MS = int(os.environ.get("SURFACES_AGENT_IDLE_TIMEOUT_MS", "60000"))
|
|
|
|
|
|
class AgentApiWrapper(AgentApi):
|
|
"""Capture tokens_used from MsgEventEnd without patching upstream code."""
|
|
|
|
def __init__(
|
|
self,
|
|
agent_id: str,
|
|
base_url: str | None = None,
|
|
*,
|
|
chat_id: int | str = 0,
|
|
url: str | None = None,
|
|
**kwargs,
|
|
) -> None:
|
|
if base_url is None and url is None:
|
|
raise TypeError("AgentApiWrapper requires base_url or url")
|
|
|
|
self._base_url = self._normalize_base_url(base_url or url or "")
|
|
self._init_kwargs = dict(kwargs)
|
|
self.chat_id = chat_id
|
|
if self._supports_modern_constructor():
|
|
super().__init__(
|
|
agent_id=agent_id,
|
|
base_url=self._base_url,
|
|
chat_id=chat_id,
|
|
**kwargs,
|
|
)
|
|
else:
|
|
super().__init__(
|
|
agent_id=agent_id,
|
|
url=self._build_ws_url(self._base_url, chat_id),
|
|
**kwargs,
|
|
)
|
|
self.last_tokens_used = 0
|
|
|
|
@staticmethod
|
|
def _supports_modern_constructor() -> bool:
|
|
try:
|
|
parameters = inspect.signature(AgentApi.__init__).parameters
|
|
except (TypeError, ValueError):
|
|
return False
|
|
|
|
return "base_url" in parameters and "chat_id" in parameters
|
|
|
|
@staticmethod
|
|
def _normalize_base_url(base_url: str) -> str:
|
|
parsed = urlsplit(base_url)
|
|
path = re.sub(r"(?:/v1)?/agent_ws(?:/[^/]+)?$", "", parsed.path.rstrip("/"))
|
|
return urlunsplit((parsed.scheme, parsed.netloc, path, "", ""))
|
|
|
|
@staticmethod
|
|
def _build_ws_url(base_url: str, chat_id: int | str) -> str:
|
|
return base_url.rstrip("/") + f"/agent_ws/?thread_id={chat_id}"
|
|
|
|
def for_chat(self, chat_id: int | str) -> AgentApiWrapper:
|
|
return type(self)(
|
|
agent_id=self.id,
|
|
base_url=self._base_url,
|
|
chat_id=chat_id,
|
|
**self._init_kwargs,
|
|
)
|
|
|
|
@staticmethod
|
|
def _event_kind(event: object) -> str:
|
|
raw_kind = getattr(event, "type", None)
|
|
if hasattr(raw_kind, "value"):
|
|
raw_kind = raw_kind.value
|
|
if raw_kind is None:
|
|
raw_kind = event.__class__.__name__
|
|
|
|
kind = str(raw_kind).replace("-", "_")
|
|
if "_" in kind:
|
|
return kind.upper()
|
|
|
|
normalized = []
|
|
for index, char in enumerate(kind):
|
|
if index and char.isupper() and not kind[index - 1].isupper():
|
|
normalized.append("_")
|
|
normalized.append(char)
|
|
return "".join(normalized).upper()
|
|
|
|
@classmethod
|
|
def _is_kind(cls, event: object, *needles: str) -> bool:
|
|
kind = cls._event_kind(event)
|
|
return any(needle in kind for needle in needles)
|
|
|
|
@classmethod
|
|
def _is_text_event(cls, event: object) -> bool:
|
|
return hasattr(event, "text") or cls._is_kind(event, "TEXT_CHUNK")
|
|
|
|
@classmethod
|
|
def _is_end_event(cls, event: object) -> bool:
|
|
kind = cls._event_kind(event)
|
|
return kind == "END" or kind.endswith("_END")
|
|
|
|
@classmethod
|
|
def _is_send_file_event(cls, event: object) -> bool:
|
|
return "SEND_FILE" in cls._event_kind(event)
|
|
|
|
async def _publish_event(self, event: object, *, queue_event: object | None = None) -> None:
|
|
if self.callback:
|
|
self.callback(event)
|
|
if self._current_queue:
|
|
await self._current_queue.put(queue_event if queue_event is not None else event)
|
|
|
|
async def _publish_error(self, event: object) -> None:
|
|
if self.callback:
|
|
self.callback(event)
|
|
if self._current_queue and hasattr(event, "code") and hasattr(event, "details"):
|
|
await self._current_queue.put(AgentException(event.code, event.details))
|
|
|
|
async def _listen(self):
|
|
try:
|
|
async for msg in self._ws:
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
try:
|
|
outgoing_msg = ServerMessage.validate_json(msg.data)
|
|
|
|
if self._is_text_event(outgoing_msg):
|
|
if _DEBUG_STREAM:
|
|
logger.warning(
|
|
"[%s] text chunk queue=%s text=%r",
|
|
self.id,
|
|
self._current_queue is not None,
|
|
getattr(outgoing_msg, "text", "")[:80],
|
|
)
|
|
if self._current_queue:
|
|
await self._current_queue.put(outgoing_msg)
|
|
elif self.callback:
|
|
self.callback(outgoing_msg)
|
|
else:
|
|
logger.warning("[%s] AgentEvent without active request", self.id)
|
|
|
|
elif self._is_end_event(outgoing_msg):
|
|
self.last_tokens_used = outgoing_msg.tokens_used
|
|
if _DEBUG_STREAM:
|
|
logger.warning(
|
|
"[%s] end event queue=%s tokens=%s",
|
|
self.id,
|
|
self._current_queue is not None,
|
|
getattr(outgoing_msg, "tokens_used", None),
|
|
)
|
|
await self._publish_event(outgoing_msg)
|
|
|
|
elif self._is_kind(outgoing_msg, "ERROR"):
|
|
error = AgentException(outgoing_msg.code, outgoing_msg.details)
|
|
logger.error("[%s] Agent error: %s", self.id, error)
|
|
await self._publish_error(outgoing_msg)
|
|
|
|
elif self._is_kind(outgoing_msg, "GRACEFUL_DISCONNECT"):
|
|
await self._publish_event(outgoing_msg)
|
|
logger.info("[%s] Gracefully disconnecting", self.id)
|
|
break
|
|
|
|
else:
|
|
await self._publish_event(outgoing_msg)
|
|
|
|
except Exception as exc:
|
|
logger.error("[%s] Failed to deserialize message: %s", self.id, exc)
|
|
if self._current_queue:
|
|
await self._current_queue.put(
|
|
AgentException("PARSE_ERROR", f"Validation failed: {exc}")
|
|
)
|
|
|
|
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
|
|
logger.error("[%s] WebSocket closed/error: %s", self.id, msg.type)
|
|
break
|
|
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as exc:
|
|
logger.error("[%s] Error in listen loop: %s", self.id, exc)
|
|
finally:
|
|
await self._cleanup()
|
|
|
|
async def send_message(
|
|
self, text: str, attachments: list[str] | None = None
|
|
) -> AsyncIterator[AgentEventUnion]:
|
|
if not self._connected or not self._ws:
|
|
raise AgentException(
|
|
code="NOT_CONNECTED", details="Not connected. Call connect() first."
|
|
)
|
|
|
|
if self._request_lock.locked():
|
|
raise AgentBusyException("Agent is currently processing another request")
|
|
|
|
await self._request_lock.acquire()
|
|
try:
|
|
self._current_queue = asyncio.Queue()
|
|
|
|
message = MsgUserMessage(
|
|
type=EClientMessage.USER_MESSAGE,
|
|
text=text,
|
|
attachments=attachments or [],
|
|
)
|
|
|
|
await self._ws.send_str(message.model_dump_json())
|
|
logger.debug("[%s] Sent message: %s...", self.id, text[:50])
|
|
|
|
while True:
|
|
try:
|
|
chunk = await asyncio.wait_for(
|
|
self._current_queue.get(),
|
|
timeout=max(_STREAM_IDLE_TIMEOUT_MS, 0) / 1000,
|
|
)
|
|
except TimeoutError as exc:
|
|
raise AgentException(
|
|
"TIMEOUT",
|
|
(
|
|
"Timed out waiting for the next agent stream event "
|
|
f"after {max(_STREAM_IDLE_TIMEOUT_MS, 0)}ms"
|
|
),
|
|
) from exc
|
|
|
|
if isinstance(chunk, Exception):
|
|
raise chunk
|
|
|
|
if isinstance(chunk, MsgEventEnd):
|
|
self.last_tokens_used = chunk.tokens_used
|
|
async for late_chunk in self._drain_post_end_events():
|
|
yield late_chunk
|
|
break
|
|
|
|
yield chunk
|
|
|
|
finally:
|
|
if self._current_queue:
|
|
orphan_queue = self._current_queue
|
|
self._current_queue = None
|
|
|
|
while not orphan_queue.empty():
|
|
try:
|
|
orphan_msg = orphan_queue.get_nowait()
|
|
if isinstance(orphan_msg, Exception):
|
|
logger.debug(
|
|
"[%s] Dropped exception from queue during cleanup: %s",
|
|
self.id,
|
|
orphan_msg,
|
|
)
|
|
continue
|
|
|
|
if self.callback:
|
|
self.callback(orphan_msg)
|
|
else:
|
|
logger.debug("[%s] Dropped orphaned message during cleanup", self.id)
|
|
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
if self._request_lock.locked():
|
|
self._request_lock.release()
|
|
|
|
async def _drain_post_end_events(self) -> AsyncIterator[AgentEventUnion]:
|
|
if self._current_queue is None:
|
|
return
|
|
|
|
timeout_s = max(_POST_END_DRAIN_MS, 0) / 1000
|
|
while True:
|
|
try:
|
|
chunk = await asyncio.wait_for(self._current_queue.get(), timeout=timeout_s)
|
|
except TimeoutError:
|
|
break
|
|
|
|
if isinstance(chunk, Exception):
|
|
logger.warning("[%s] dropping post-END exception: %s", self.id, chunk)
|
|
continue
|
|
|
|
if isinstance(chunk, MsgEventEnd):
|
|
self.last_tokens_used = chunk.tokens_used
|
|
if _DEBUG_STREAM:
|
|
logger.warning(
|
|
"[%s] dropped duplicate END tokens=%s",
|
|
self.id,
|
|
chunk.tokens_used,
|
|
)
|
|
continue
|
|
|
|
if _DEBUG_STREAM and self._is_text_event(chunk):
|
|
logger.warning(
|
|
"[%s] recovered post-END text chunk=%r",
|
|
self.id,
|
|
getattr(chunk, "text", "")[:80],
|
|
)
|
|
|
|
yield chunk
|