refactor: use thin upstream transport adapter
This commit is contained in:
parent
569824ead1
commit
0c2884c2b1
8 changed files with 420 additions and 255 deletions
|
|
@ -1,6 +1,5 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
|
@ -27,11 +26,6 @@ class AgentApiWrapper(AgentApi):
|
|||
self._base_url = self._normalize_base_url(base_url)
|
||||
self._init_kwargs = dict(kwargs)
|
||||
self.chat_id = chat_id
|
||||
if not self._supports_modern_constructor():
|
||||
raise RuntimeError(
|
||||
"Pinned platform-agent_api is expected to support base_url + chat_id"
|
||||
)
|
||||
|
||||
super().__init__(
|
||||
agent_id=agent_id,
|
||||
base_url=self._base_url,
|
||||
|
|
@ -39,21 +33,13 @@ class AgentApiWrapper(AgentApi):
|
|||
**kwargs,
|
||||
)
|
||||
|
||||
@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, "", ""))
|
||||
|
||||
def for_chat(self, chat_id: int | str) -> "AgentApiWrapper":
|
||||
def for_chat(self, chat_id: int | str) -> AgentApiWrapper:
|
||||
return type(self)(
|
||||
agent_id=self.id,
|
||||
base_url=self._base_url,
|
||||
|
|
|
|||
152
sdk/real.py
152
sdk/real.py
|
|
@ -1,10 +1,11 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import AsyncIterator
|
||||
from pathlib import Path
|
||||
|
||||
from lambda_agent_api.server import MsgEventSendFile, MsgEventTextChunk
|
||||
|
||||
from sdk.agent_api_wrapper import AgentApiWrapper
|
||||
from sdk.interface import (
|
||||
Attachment,
|
||||
|
|
@ -40,14 +41,10 @@ class RealPlatformClient(PlatformClient):
|
|||
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)
|
||||
chat_api = self._agent_api.for_chat(chat_key)
|
||||
await chat_api.connect()
|
||||
self._chat_apis[chat_key] = chat_api
|
||||
return chat_api
|
||||
|
|
@ -80,48 +77,36 @@ class RealPlatformClient(PlatformClient):
|
|||
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):
|
||||
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:
|
||||
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)
|
||||
await self._prototype_state.set_last_tokens_used(str(chat_id), 0)
|
||||
|
||||
response_kwargs = {
|
||||
"message_id": message_id,
|
||||
"response": "".join(response_parts),
|
||||
"tokens_used": tokens_used,
|
||||
"tokens_used": 0,
|
||||
"finished": True,
|
||||
"attachments": sent_attachments,
|
||||
}
|
||||
if self._message_response_accepts_attachments():
|
||||
response_kwargs["attachments"] = sent_attachments
|
||||
return MessageResponse(**response_kwargs)
|
||||
|
||||
async def stream_message(
|
||||
|
|
@ -134,44 +119,27 @@ class RealPlatformClient(PlatformClient):
|
|||
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):
|
||||
if isinstance(event, MsgEventTextChunk):
|
||||
yield MessageChunk(
|
||||
message_id=user_id,
|
||||
delta=getattr(event, "text", ""),
|
||||
delta=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:
|
||||
elif isinstance(event, MsgEventSendFile):
|
||||
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,
|
||||
)
|
||||
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)
|
||||
|
|
@ -195,10 +163,6 @@ class RealPlatformClient(PlatformClient):
|
|||
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,
|
||||
|
|
@ -206,12 +170,8 @@ class RealPlatformClient(PlatformClient):
|
|||
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)
|
||||
event_stream = chat_api.send_message(text, attachments=attachment_paths or None)
|
||||
async for event in event_stream:
|
||||
yield event
|
||||
|
||||
|
|
@ -231,61 +191,9 @@ class RealPlatformClient(PlatformClient):
|
|||
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)
|
||||
def _attachment_from_send_file_event(event: MsgEventSendFile) -> Attachment:
|
||||
location = str(event.path)
|
||||
filename = Path(location).name or None
|
||||
workspace_path = location
|
||||
if workspace_path.startswith("/workspace/"):
|
||||
workspace_path = workspace_path[len("/workspace/") :]
|
||||
|
|
@ -293,18 +201,8 @@ class RealPlatformClient(PlatformClient):
|
|||
workspace_path = ""
|
||||
return Attachment(
|
||||
url=location,
|
||||
mime_type=mime_type,
|
||||
size=size,
|
||||
mime_type="application/octet-stream",
|
||||
size=None,
|
||||
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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue