Mikhail Putilovskij
b52fdc4670
- add save/load/reset/context handlers and matrix interception flows - persist current session and last token usage in prototype state
88 lines
2.7 KiB
Python
88 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import AsyncIterator
|
|
|
|
from sdk.agent_api_wrapper import AgentApiWrapper
|
|
from sdk.interface import Attachment, MessageChunk, MessageResponse, PlatformClient, 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
|
|
|
|
@property
|
|
def agent_api(self) -> AgentApiWrapper:
|
|
return self._agent_api
|
|
|
|
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
|
|
message_id = user_id
|
|
|
|
async for chunk in self.stream_message(user_id, chat_id, text, attachments=attachments):
|
|
message_id = chunk.message_id
|
|
if chunk.delta:
|
|
response_parts.append(chunk.delta)
|
|
if chunk.finished:
|
|
tokens_used = chunk.tokens_used
|
|
|
|
return MessageResponse(
|
|
message_id=message_id,
|
|
response="".join(response_parts),
|
|
tokens_used=tokens_used,
|
|
finished=True,
|
|
)
|
|
|
|
async def stream_message(
|
|
self,
|
|
user_id: str,
|
|
chat_id: str,
|
|
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):
|
|
yield MessageChunk(
|
|
message_id=user_id,
|
|
delta=event.text,
|
|
finished=False,
|
|
)
|
|
await self._prototype_state.set_last_tokens_used(user_id, self._agent_api.last_tokens_used)
|
|
yield MessageChunk(
|
|
message_id=user_id,
|
|
delta="",
|
|
finished=True,
|
|
tokens_used=self._agent_api.last_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)
|