Add per-chat real client routing

This commit is contained in:
Mikhail Putilovskij 2026-04-19 17:03:48 +03:00
parent 5782001d3d
commit 414a8645bd
3 changed files with 138 additions and 18 deletions

View file

@ -17,11 +17,21 @@ class RealPlatformClient(PlatformClient):
self._agent_api = agent_api
self._prototype_state = prototype_state
self._platform = platform
self._chat_apis: dict[str, AgentApiWrapper] = {}
@property
def agent_api(self) -> AgentApiWrapper:
return self._agent_api
async def _get_chat_api(self, chat_id: str) -> AgentApiWrapper:
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
return chat_api
async def get_or_create_user(
self,
external_id: str,
@ -66,8 +76,9 @@ class RealPlatformClient(PlatformClient):
text: str,
attachments: list[Attachment] | None = None,
) -> AsyncIterator[MessageChunk]:
self._agent_api.last_tokens_used = 0
async for event in self._agent_api.send_message(text):
chat_api = await self._get_chat_api(chat_id)
chat_api.last_tokens_used = 0
async for event in chat_api.send_message(text):
yield MessageChunk(
message_id=user_id,
delta=event.text,
@ -77,7 +88,7 @@ class RealPlatformClient(PlatformClient):
message_id=user_id,
delta="",
finished=True,
tokens_used=self._agent_api.last_tokens_used,
tokens_used=chat_api.last_tokens_used,
)
async def get_settings(self, user_id: str) -> UserSettings:
@ -85,3 +96,8 @@ class RealPlatformClient(PlatformClient):
async def update_settings(self, user_id: str, action) -> None:
await self._prototype_state.update_settings(user_id, action)
async def close(self) -> None:
for chat_api in list(self._chat_apis.values()):
await chat_api.close()
self._chat_apis.clear()