Surfaces_team/surfaces
4
0
Fork 0
Code Issues 5 Pull requests Projects Releases Packages Wiki Activity Actions

Fix real client chat cache compatibility

Browse source
This commit is contained in:
Mikhail Putilovskij 2026-04-19 17:07:52 +03:00
parent 414a8645bd
commit 730ea70f78
2 changed files with 77 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

30
sdk/real.py
View file

@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
from typing import AsyncIterator
from sdk.agent_api_wrapper import AgentApiWrapper
@ -18,18 +19,26 @@ class RealPlatformClient(PlatformClient):
self._prototype_state = prototype_state
self._platform = platform
self._chat_apis: dict[str, AgentApiWrapper] = {}
self._chat_api_lock = asyncio.Lock()
@property
def agent_api(self) -> AgentApiWrapper:
return self._agent_api
async def _get_chat_api(self, chat_id: str) -> AgentApiWrapper:
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 = self._agent_api.for_chat(chat_key)
await chat_api.connect()
self._chat_apis[chat_key] = chat_api
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
async def get_or_create_user(
@ -77,7 +86,8 @@ class RealPlatformClient(PlatformClient):
attachments: list[Attachment] | None = None,
) -> AsyncIterator[MessageChunk]:
chat_api = await self._get_chat_api(chat_id)
chat_api.last_tokens_used = 0
if hasattr(chat_api, "last_tokens_used"):
chat_api.last_tokens_used = 0
async for event in chat_api.send_message(text):
yield MessageChunk(
message_id=user_id,
@ -88,7 +98,7 @@ class RealPlatformClient(PlatformClient):
message_id=user_id,
delta="",
finished=True,
tokens_used=chat_api.last_tokens_used,
tokens_used=getattr(chat_api, "last_tokens_used", 0),
)
async def get_settings(self, user_id: str) -> UserSettings:
@ -99,5 +109,11 @@ class RealPlatformClient(PlatformClient):
async def close(self) -> None:
for chat_api in list(self._chat_apis.values()):
await chat_api.close()
close = getattr(chat_api, "close", None)
if callable(close):
await close()
self._chat_apis.clear()
if not callable(getattr(self._agent_api, "for_chat", None)):
close = getattr(self._agent_api, "close", None)
if callable(close):
await close()