231 lines
8.4 KiB
Python
231 lines
8.4 KiB
Python
# platform/mock.py
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import random
|
||
import uuid
|
||
from datetime import UTC, datetime
|
||
from typing import Any, AsyncIterator, Literal
|
||
|
||
import structlog
|
||
|
||
from sdk.interface import (
|
||
AgentEvent,
|
||
Attachment,
|
||
MessageChunk,
|
||
MessageResponse,
|
||
User,
|
||
UserSettings,
|
||
WebhookReceiver,
|
||
)
|
||
|
||
logger = structlog.get_logger(__name__)
|
||
|
||
|
||
class MockPlatformClient:
|
||
"""
|
||
Заглушка SDK платформы Lambda.
|
||
|
||
Реализует PlatformClient Protocol. При подключении реального SDK
|
||
заменяется только этот файл — core/ и адаптеры не трогаются.
|
||
|
||
attachment_mode — симулирует разные варианты передачи файлов:
|
||
"url" — платформа получает URL, скачивает сама (текущий план)
|
||
"binary" — бинарные данные в теле (резерв)
|
||
"s3" — pre-signed S3 URL (резерв)
|
||
|
||
Webhook: зарегистрируй WebhookReceiver через register_webhook_receiver(),
|
||
вызови simulate_agent_event() чтобы имитировать входящее уведомление.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
attachment_mode: Literal["url", "binary", "s3"] = "url",
|
||
) -> None:
|
||
self.attachment_mode = attachment_mode
|
||
self._users: dict[str, dict] = {}
|
||
self._messages: dict[str, list] = {} # "{user_id}:{chat_id}" → messages
|
||
self._settings: dict[str, dict] = {}
|
||
self._webhook_receiver: WebhookReceiver | None = None
|
||
logger.info("MockPlatformClient initialized", attachment_mode=attachment_mode)
|
||
|
||
# ------------------------------------------------------------------ users
|
||
|
||
async def get_or_create_user(
|
||
self,
|
||
external_id: str,
|
||
platform: str,
|
||
display_name: str | None = None,
|
||
) -> User:
|
||
await self._latency()
|
||
key = f"{platform}:{external_id}"
|
||
is_new = key not in self._users
|
||
if is_new:
|
||
self._users[key] = {
|
||
"user_id": f"usr-{platform}-{external_id}",
|
||
"external_id": external_id,
|
||
"platform": platform,
|
||
"display_name": display_name,
|
||
"created_at": "2025-01-01T00:00:00Z",
|
||
"is_new": True,
|
||
}
|
||
data = {**self._users[key], "is_new": is_new}
|
||
return User(**data)
|
||
|
||
# --------------------------------------------------------------- messages
|
||
|
||
async def send_message(
|
||
self,
|
||
user_id: str,
|
||
chat_id: str,
|
||
text: str,
|
||
attachments: list[Attachment] | None = None,
|
||
) -> MessageResponse:
|
||
await self._latency(200, 600)
|
||
message_id, response, tokens = self._build_response(user_id, chat_id, text, attachments)
|
||
logger.info("send_message", user_id=user_id, chat_id=chat_id, message_id=message_id)
|
||
return MessageResponse(
|
||
message_id=message_id,
|
||
response=response,
|
||
tokens_used=tokens,
|
||
finished=True,
|
||
)
|
||
|
||
async def stream_message(
|
||
self,
|
||
user_id: str,
|
||
chat_id: str,
|
||
text: str,
|
||
attachments: list[Attachment] | None = None,
|
||
) -> AsyncIterator[MessageChunk]:
|
||
"""
|
||
Сейчас: один чанк с полным ответом (sync под капотом).
|
||
При реальном SDK: заменить на SSE/WebSocket итератор в platform/mock.py.
|
||
Адаптеры переписывать не нужно.
|
||
"""
|
||
await self._latency(200, 600)
|
||
message_id, response, tokens = self._build_response(user_id, chat_id, text, attachments)
|
||
logger.info("stream_message", user_id=user_id, chat_id=chat_id, message_id=message_id)
|
||
|
||
async def _gen() -> AsyncIterator[MessageChunk]:
|
||
yield MessageChunk(
|
||
message_id=message_id,
|
||
delta=response,
|
||
finished=True,
|
||
tokens_used=tokens,
|
||
)
|
||
|
||
return _gen()
|
||
|
||
# --------------------------------------------------------------- settings
|
||
|
||
async def get_settings(self, user_id: str) -> UserSettings:
|
||
await self._latency()
|
||
stored = self._settings.get(user_id, {})
|
||
return UserSettings(
|
||
skills=stored.get("skills", {
|
||
"web-search": True,
|
||
"fetch-url": True,
|
||
"email": False,
|
||
"browser": False,
|
||
"image-gen": False,
|
||
"files": True,
|
||
}),
|
||
connectors=stored.get("connectors", {}),
|
||
soul=stored.get("soul", {"name": "Лямбда", "instructions": ""}),
|
||
safety=stored.get("safety", {
|
||
"email-send": True,
|
||
"file-delete": True,
|
||
"social-post": True,
|
||
}),
|
||
plan=stored.get("plan", {
|
||
"name": "Beta",
|
||
"tokens_used": 0,
|
||
"tokens_limit": 1000,
|
||
}),
|
||
)
|
||
|
||
async def update_settings(self, user_id: str, action: Any) -> None:
|
||
await self._latency()
|
||
settings = self._settings.setdefault(user_id, {})
|
||
|
||
if action.action == "toggle_skill":
|
||
skills = settings.setdefault("skills", {})
|
||
skills[action.payload["skill"]] = action.payload.get("enabled", True)
|
||
elif action.action == "set_soul":
|
||
soul = settings.setdefault("soul", {})
|
||
soul[action.payload["field"]] = action.payload["value"]
|
||
elif action.action == "set_safety":
|
||
safety = settings.setdefault("safety", {})
|
||
safety[action.payload["trigger"]] = action.payload.get("enabled", True)
|
||
|
||
logger.info("Settings updated", user_id=user_id, action=action.action)
|
||
|
||
# --------------------------------------------------------------- webhooks
|
||
|
||
def register_webhook_receiver(self, receiver: WebhookReceiver) -> None:
|
||
"""Бот регистрирует свой обработчик входящих событий от платформы."""
|
||
self._webhook_receiver = receiver
|
||
logger.info("WebhookReceiver registered")
|
||
|
||
async def simulate_agent_event(
|
||
self,
|
||
user_id: str,
|
||
chat_id: str,
|
||
event_type: Literal["task_done", "task_error", "task_progress"] = "task_done",
|
||
payload: dict | None = None,
|
||
) -> None:
|
||
"""Имитирует входящий webhook от платформы. Используется в тестах и ручном QA."""
|
||
if self._webhook_receiver is None:
|
||
logger.warning("simulate_agent_event: no WebhookReceiver registered")
|
||
return
|
||
event = AgentEvent(
|
||
event_id=str(uuid.uuid4()),
|
||
user_id=user_id,
|
||
chat_id=chat_id,
|
||
event_type=event_type,
|
||
payload=payload or {"message": "[MOCK] Долгая задача выполнена"},
|
||
)
|
||
await self._webhook_receiver.on_agent_event(event)
|
||
|
||
# ------------------------------------------------------------------ utils
|
||
|
||
def get_stats(self) -> dict:
|
||
return {
|
||
"total_users": len(self._users),
|
||
"total_messages": sum(len(msgs) for msgs in self._messages.values()),
|
||
"attachment_mode": self.attachment_mode,
|
||
}
|
||
|
||
def _build_response(
|
||
self,
|
||
user_id: str,
|
||
chat_id: str,
|
||
text: str,
|
||
attachments: list[Attachment] | None,
|
||
) -> tuple[str, str, int]:
|
||
key = f"{user_id}:{chat_id}"
|
||
if key not in self._messages:
|
||
self._messages[key] = []
|
||
|
||
message_id = str(uuid.uuid4())
|
||
preview = text[:50] + ("..." if len(text) > 50 else "")
|
||
attachment_note = ""
|
||
if attachments:
|
||
names = [a.filename or a.mime_type for a in attachments]
|
||
attachment_note = f" [вложения: {', '.join(names)}]"
|
||
response = f"[MOCK] Ответ на: «{preview}»{attachment_note}"
|
||
tokens = len(text.split()) * 2
|
||
|
||
self._messages[key].append({
|
||
"message_id": message_id,
|
||
"user_text": text,
|
||
"response": response,
|
||
"tokens_used": tokens,
|
||
"finished": True,
|
||
"created_at": datetime.now(UTC).isoformat(),
|
||
})
|
||
return message_id, response, tokens
|
||
|
||
async def _latency(self, min_ms: int = 10, max_ms: int = 80) -> None:
|
||
await asyncio.sleep(random.randint(min_ms, max_ms) / 1000)
|