310 lines
11 KiB
Python
310 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
from collections.abc import AsyncIterator
|
|
from pathlib import Path
|
|
|
|
from sdk.agent_api_wrapper import AgentApiWrapper
|
|
from sdk.interface import (
|
|
Attachment,
|
|
MessageChunk,
|
|
MessageResponse,
|
|
PlatformClient,
|
|
PlatformError,
|
|
User,
|
|
UserSettings,
|
|
)
|
|
from sdk.prototype_state import PrototypeStateStore
|
|
|
|
|
|
class RealPlatformClient(PlatformClient):
|
|
def __init__(
|
|
self,
|
|
agent_api: AgentApiWrapper,
|
|
prototype_state: PrototypeStateStore,
|
|
platform: str = "matrix",
|
|
) -> None:
|
|
self._agent_api = agent_api
|
|
self._prototype_state = prototype_state
|
|
self._platform = platform
|
|
self._chat_apis: dict[str, AgentApiWrapper] = {}
|
|
self._chat_api_lock = asyncio.Lock()
|
|
self._chat_send_locks: dict[str, asyncio.Lock] = {}
|
|
|
|
@property
|
|
def agent_api(self) -> AgentApiWrapper:
|
|
return self._agent_api
|
|
|
|
async def _get_chat_api(self, chat_id: str):
|
|
chat_key = str(chat_id)
|
|
chat_api = self._chat_apis.get(chat_key)
|
|
if chat_api is None:
|
|
chat_api_factory = getattr(self._agent_api, "for_chat", None)
|
|
if not callable(chat_api_factory):
|
|
return self._agent_api
|
|
|
|
async with self._chat_api_lock:
|
|
chat_api = self._chat_apis.get(chat_key)
|
|
if chat_api is None:
|
|
chat_api = chat_api_factory(chat_key)
|
|
await chat_api.connect()
|
|
self._chat_apis[chat_key] = chat_api
|
|
return chat_api
|
|
|
|
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] = []
|
|
tokens_used = 0
|
|
sent_attachments: list[Attachment] = []
|
|
message_id = user_id
|
|
saw_end_event = False
|
|
|
|
lock = self._get_chat_send_lock(chat_id)
|
|
async with lock:
|
|
chat_api = await self._get_chat_api(chat_id)
|
|
if hasattr(chat_api, "last_tokens_used"):
|
|
chat_api.last_tokens_used = 0
|
|
|
|
try:
|
|
async for event in self._stream_agent_events(
|
|
chat_api, text, attachments=attachments
|
|
):
|
|
message_id = user_id
|
|
if self._is_text_event(event):
|
|
chunk_text = getattr(event, "text", "")
|
|
if chunk_text:
|
|
response_parts.append(chunk_text)
|
|
elif self._is_end_event(event):
|
|
tokens_used = getattr(event, "tokens_used", tokens_used)
|
|
saw_end_event = True
|
|
elif self._is_send_file_event(event):
|
|
attachment = self._attachment_from_send_file_event(event)
|
|
if attachment is not None:
|
|
sent_attachments.append(attachment)
|
|
except Exception as exc:
|
|
await self._handle_chat_api_failure(chat_id, exc)
|
|
|
|
if not saw_end_event:
|
|
tokens_used = getattr(chat_api, "last_tokens_used", tokens_used)
|
|
await self._prototype_state.set_last_tokens_used(str(chat_id), tokens_used)
|
|
|
|
response_kwargs = {
|
|
"message_id": message_id,
|
|
"response": "".join(response_parts),
|
|
"tokens_used": tokens_used,
|
|
"finished": True,
|
|
}
|
|
if self._message_response_accepts_attachments():
|
|
response_kwargs["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 = await self._get_chat_api(chat_id)
|
|
if hasattr(chat_api, "last_tokens_used"):
|
|
chat_api.last_tokens_used = 0
|
|
saw_end_event = False
|
|
try:
|
|
async for event in self._stream_agent_events(
|
|
chat_api, text, attachments=attachments
|
|
):
|
|
if self._is_text_event(event):
|
|
yield MessageChunk(
|
|
message_id=user_id,
|
|
delta=getattr(event, "text", ""),
|
|
finished=False,
|
|
)
|
|
elif self._is_end_event(event):
|
|
tokens_used = getattr(event, "tokens_used", 0)
|
|
saw_end_event = True
|
|
await self._prototype_state.set_last_tokens_used(str(chat_id), tokens_used)
|
|
yield MessageChunk(
|
|
message_id=user_id,
|
|
delta="",
|
|
finished=True,
|
|
tokens_used=tokens_used,
|
|
)
|
|
elif self._is_send_file_event(event):
|
|
continue
|
|
else:
|
|
continue
|
|
except Exception as exc:
|
|
await self._handle_chat_api_failure(chat_id, exc)
|
|
if not saw_end_event:
|
|
tokens_used = getattr(chat_api, "last_tokens_used", 0)
|
|
await self._prototype_state.set_last_tokens_used(str(chat_id), tokens_used)
|
|
yield MessageChunk(
|
|
message_id=user_id,
|
|
delta="",
|
|
finished=True,
|
|
tokens_used=tokens_used,
|
|
)
|
|
|
|
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:
|
|
chat_key = str(chat_id)
|
|
chat_api = self._chat_apis.pop(chat_key, None)
|
|
self._chat_send_locks.pop(chat_key, None)
|
|
if chat_api is not None:
|
|
close = getattr(chat_api, "close", None)
|
|
if callable(close):
|
|
await close()
|
|
|
|
async def close(self) -> None:
|
|
for chat_api in list(self._chat_apis.values()):
|
|
close = getattr(chat_api, "close", None)
|
|
if callable(close):
|
|
await close()
|
|
self._chat_apis.clear()
|
|
self._chat_send_locks.clear()
|
|
if not callable(getattr(self._agent_api, "for_chat", None)):
|
|
close = getattr(self._agent_api, "close", None)
|
|
if callable(close):
|
|
await close()
|
|
|
|
async def _stream_agent_events(
|
|
self,
|
|
chat_api,
|
|
text: str,
|
|
attachments: list[Attachment] | None = None,
|
|
) -> AsyncIterator[object]:
|
|
send_message = chat_api.send_message
|
|
attachment_paths = self._attachment_paths(attachments)
|
|
if attachment_paths and self._send_message_accepts_attachments(send_message):
|
|
event_stream = send_message(text, attachments=attachment_paths)
|
|
else:
|
|
event_stream = send_message(text)
|
|
async for event in event_stream:
|
|
yield event
|
|
|
|
async def _handle_chat_api_failure(self, chat_id: str, exc: Exception) -> None:
|
|
await self.disconnect_chat(chat_id)
|
|
code = getattr(exc, "code", None) or "PLATFORM_CONNECTION_ERROR"
|
|
raise PlatformError(str(exc), code=code) from exc
|
|
|
|
@staticmethod
|
|
def _attachment_paths(attachments: list[Attachment] | None) -> list[str]:
|
|
if not attachments:
|
|
return []
|
|
paths = []
|
|
for attachment in attachments:
|
|
if attachment.workspace_path:
|
|
paths.append(attachment.workspace_path)
|
|
return paths
|
|
|
|
@staticmethod
|
|
def _send_message_accepts_attachments(send_message) -> bool:
|
|
try:
|
|
parameters = inspect.signature(send_message).parameters
|
|
except (TypeError, ValueError):
|
|
return False
|
|
return "attachments" in parameters or any(
|
|
parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()
|
|
)
|
|
|
|
@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_text_event(cls, event: object) -> bool:
|
|
return hasattr(event, "text") or "TEXT_CHUNK" in cls._event_kind(event)
|
|
|
|
@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:
|
|
kind = cls._event_kind(event)
|
|
return "SEND_FILE" in kind
|
|
|
|
@staticmethod
|
|
def _attachment_from_send_file_event(event: object) -> Attachment | None:
|
|
location = None
|
|
for attr in ("url", "workspace_path", "path", "file_path", "uri"):
|
|
value = getattr(event, attr, None)
|
|
if value:
|
|
location = str(value)
|
|
break
|
|
if location is None:
|
|
return None
|
|
|
|
mime_type = getattr(event, "mime_type", None) or "application/octet-stream"
|
|
filename = getattr(event, "filename", None) or Path(location).name or None
|
|
size = getattr(event, "size", None)
|
|
workspace_path = location
|
|
if workspace_path.startswith("/workspace/"):
|
|
workspace_path = workspace_path[len("/workspace/") :]
|
|
elif workspace_path == "/workspace":
|
|
workspace_path = ""
|
|
return Attachment(
|
|
url=location,
|
|
mime_type=mime_type,
|
|
size=size,
|
|
filename=filename,
|
|
workspace_path=workspace_path or None,
|
|
)
|
|
|
|
@staticmethod
|
|
def _message_response_accepts_attachments() -> bool:
|
|
fields = getattr(MessageResponse, "model_fields", None)
|
|
if isinstance(fields, dict):
|
|
return "attachments" in fields
|
|
try:
|
|
return "attachments" in inspect.signature(MessageResponse).parameters
|
|
except (TypeError, ValueError):
|
|
return False
|