273 lines
9.1 KiB
Python
273 lines
9.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import re
|
|
from collections.abc import AsyncIterator
|
|
from pathlib import Path
|
|
from urllib.parse import urljoin, urlsplit, urlunsplit
|
|
|
|
import structlog
|
|
|
|
from sdk.interface import (
|
|
Attachment,
|
|
MessageChunk,
|
|
MessageResponse,
|
|
PlatformClient,
|
|
PlatformError,
|
|
User,
|
|
UserSettings,
|
|
)
|
|
from sdk.prototype_state import PrototypeStateStore
|
|
from sdk.upstream_agent_api import AgentApi, MsgEventSendFile, MsgEventTextChunk
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
|
|
|
|
def _ws_debug_enabled() -> bool:
|
|
value = os.environ.get("SURFACES_DEBUG_WS", "")
|
|
return value.strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
|
|
class RealPlatformClient(PlatformClient):
|
|
def __init__(
|
|
self,
|
|
agent_id: str,
|
|
agent_base_url: str,
|
|
prototype_state: PrototypeStateStore,
|
|
platform: str = "matrix",
|
|
agent_api_cls=AgentApi,
|
|
) -> None:
|
|
self._agent_id = agent_id
|
|
self._raw_agent_base_url = agent_base_url
|
|
self._agent_base_url = self._normalize_agent_base_url(agent_base_url)
|
|
self._agent_api_cls = agent_api_cls
|
|
self._prototype_state = prototype_state
|
|
self._platform = platform
|
|
self._chat_send_locks: dict[str, asyncio.Lock] = {}
|
|
if _ws_debug_enabled():
|
|
logger.warning(
|
|
"agent_client_initialized",
|
|
agent_id=self._agent_id,
|
|
platform=self._platform,
|
|
raw_base_url=self._raw_agent_base_url,
|
|
normalized_base_url=self._agent_base_url,
|
|
)
|
|
|
|
@property
|
|
def agent_id(self) -> str:
|
|
return self._agent_id
|
|
|
|
@property
|
|
def agent_base_url(self) -> str:
|
|
return self._agent_base_url
|
|
|
|
def _get_chat_send_lock(self, chat_id: str) -> asyncio.Lock:
|
|
chat_key = str(chat_id)
|
|
lock = self._chat_send_locks.get(chat_key)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
self._chat_send_locks[chat_key] = lock
|
|
return lock
|
|
|
|
async def get_or_create_user(
|
|
self,
|
|
external_id: str,
|
|
platform: str,
|
|
display_name: str | None = None,
|
|
) -> User:
|
|
return await self._prototype_state.get_or_create_user(
|
|
external_id=external_id,
|
|
platform=platform,
|
|
display_name=display_name,
|
|
)
|
|
|
|
async def send_message(
|
|
self,
|
|
user_id: str,
|
|
chat_id: str,
|
|
text: str,
|
|
attachments: list[Attachment] | None = None,
|
|
) -> MessageResponse:
|
|
response_parts: list[str] = []
|
|
sent_attachments: list[Attachment] = []
|
|
message_id = user_id
|
|
|
|
lock = self._get_chat_send_lock(chat_id)
|
|
async with lock:
|
|
chat_api = self._build_chat_api(chat_id)
|
|
try:
|
|
await chat_api.connect()
|
|
async for event in self._stream_agent_events(
|
|
chat_api, text, attachments=attachments
|
|
):
|
|
message_id = user_id
|
|
if isinstance(event, MsgEventTextChunk) and event.text:
|
|
response_parts.append(event.text)
|
|
elif isinstance(event, MsgEventSendFile):
|
|
attachment = self._attachment_from_send_file_event(event)
|
|
if attachment is not None:
|
|
sent_attachments.append(attachment)
|
|
except Exception as exc:
|
|
raise self._to_platform_error(exc) from exc
|
|
finally:
|
|
await self._close_chat_api(chat_api)
|
|
await self._prototype_state.set_last_tokens_used(str(chat_id), 0)
|
|
|
|
response_kwargs = {
|
|
"message_id": message_id,
|
|
"response": "".join(response_parts),
|
|
"tokens_used": 0,
|
|
"finished": True,
|
|
"attachments": sent_attachments,
|
|
}
|
|
return MessageResponse(**response_kwargs)
|
|
|
|
async def stream_message(
|
|
self,
|
|
user_id: str,
|
|
chat_id: str,
|
|
text: str,
|
|
attachments: list[Attachment] | None = None,
|
|
) -> AsyncIterator[MessageChunk]:
|
|
lock = self._get_chat_send_lock(chat_id)
|
|
async with lock:
|
|
chat_api = self._build_chat_api(chat_id)
|
|
try:
|
|
await chat_api.connect()
|
|
async for event in self._stream_agent_events(
|
|
chat_api, text, attachments=attachments
|
|
):
|
|
if isinstance(event, MsgEventTextChunk):
|
|
yield MessageChunk(
|
|
message_id=user_id,
|
|
delta=event.text,
|
|
finished=False,
|
|
)
|
|
elif isinstance(event, MsgEventSendFile):
|
|
continue
|
|
except Exception as exc:
|
|
raise self._to_platform_error(exc) from exc
|
|
finally:
|
|
await self._close_chat_api(chat_api)
|
|
await self._prototype_state.set_last_tokens_used(str(chat_id), 0)
|
|
yield MessageChunk(
|
|
message_id=user_id,
|
|
delta="",
|
|
finished=True,
|
|
tokens_used=0,
|
|
)
|
|
|
|
async def get_settings(self, user_id: str) -> UserSettings:
|
|
return await self._prototype_state.get_settings(user_id)
|
|
|
|
async def update_settings(self, user_id: str, action) -> None:
|
|
await self._prototype_state.update_settings(user_id, action)
|
|
|
|
async def disconnect_chat(self, chat_id: str) -> None:
|
|
self._chat_send_locks.pop(str(chat_id), None)
|
|
|
|
async def close(self) -> None:
|
|
self._chat_send_locks.clear()
|
|
|
|
async def _stream_agent_events(
|
|
self,
|
|
chat_api,
|
|
text: str,
|
|
attachments: list[Attachment] | None = None,
|
|
) -> AsyncIterator[object]:
|
|
attachment_paths = self._attachment_paths(attachments)
|
|
event_stream = chat_api.send_message(text, attachments=attachment_paths or None)
|
|
chunk_index = 0
|
|
async for event in event_stream:
|
|
if isinstance(event, MsgEventTextChunk):
|
|
logger.debug("agent_chunk", index=chunk_index, text=repr(event.text[:40]))
|
|
chunk_index += 1
|
|
else:
|
|
logger.debug("agent_event", index=chunk_index, type=type(event).__name__)
|
|
yield event
|
|
|
|
def _build_chat_api(self, chat_id: str):
|
|
if _ws_debug_enabled():
|
|
logger.warning(
|
|
"agent_chat_api_build",
|
|
agent_id=self._agent_id,
|
|
chat_id=str(chat_id),
|
|
normalized_base_url=self._agent_base_url,
|
|
ws_url=urljoin(self._agent_base_url, f"v1/agent_ws/{chat_id}/"),
|
|
)
|
|
return self._agent_api_cls(
|
|
agent_id=self._agent_id,
|
|
base_url=self._agent_base_url,
|
|
chat_id=str(chat_id),
|
|
)
|
|
|
|
@staticmethod
|
|
def _normalize_agent_base_url(base_url: str) -> str:
|
|
parsed = urlsplit(base_url)
|
|
path = re.sub(r"(?:/v1)?/agent_ws(?:/[^/]+)?/?$", "", parsed.path.rstrip("/"))
|
|
if path:
|
|
path = f"{path}/"
|
|
return urlunsplit((parsed.scheme, parsed.netloc, path, "", ""))
|
|
|
|
@staticmethod
|
|
async def _close_chat_api(chat_api) -> None:
|
|
close = getattr(chat_api, "close", None)
|
|
if callable(close):
|
|
try:
|
|
await close()
|
|
except Exception:
|
|
pass
|
|
|
|
@staticmethod
|
|
def _to_platform_error(exc: Exception) -> PlatformError:
|
|
code = getattr(exc, "code", None) or "PLATFORM_CONNECTION_ERROR"
|
|
return PlatformError(str(exc), code=code)
|
|
|
|
@staticmethod
|
|
def _normalize_workspace_path(location: str) -> str | None:
|
|
if not location:
|
|
return None
|
|
|
|
path = Path(location)
|
|
if not path.is_absolute():
|
|
normalized = path.as_posix()
|
|
return normalized or None
|
|
|
|
parts = path.parts
|
|
if len(parts) >= 2 and parts[1] == "workspace":
|
|
relative = Path(*parts[2:]).as_posix()
|
|
return relative or None
|
|
if len(parts) >= 3 and parts[1] == "agents":
|
|
relative = Path(*parts[3:]).as_posix()
|
|
return relative or None
|
|
|
|
relative = path.as_posix().lstrip("/")
|
|
return relative or None
|
|
|
|
@staticmethod
|
|
def _attachment_paths(attachments: list[Attachment] | None) -> list[str]:
|
|
if not attachments:
|
|
return []
|
|
paths = []
|
|
for attachment in attachments:
|
|
if attachment.workspace_path:
|
|
normalized = RealPlatformClient._normalize_workspace_path(
|
|
attachment.workspace_path
|
|
)
|
|
if normalized:
|
|
paths.append(normalized)
|
|
return paths
|
|
|
|
@staticmethod
|
|
def _attachment_from_send_file_event(event: MsgEventSendFile) -> Attachment:
|
|
location = str(event.path)
|
|
filename = Path(location).name or None
|
|
workspace_path = RealPlatformClient._normalize_workspace_path(location)
|
|
return Attachment(
|
|
url=location,
|
|
mime_type="application/octet-stream",
|
|
size=None,
|
|
filename=filename,
|
|
workspace_path=workspace_path or None,
|
|
)
|