from __future__ import annotations from collections.abc import Iterator from dataclasses import dataclass from datetime import date, datetime from typing import Any from zoneinfo import ZoneInfo from backend.data.gateway import DataGateway from backend.database.connection import Database from backend.features.accounts.models import Principal from backend.features.mentor.context import MentorContextBuilder from backend.features.mentor.prompt import PROMPT_VERSION from backend.features.mentor.prompt import messages as prompt_messages from backend.features.mentor.repository import MentorRepository from backend.features.mentor.skills import MentorSkill, MentorSkillError, MentorSkillRegistry from backend.llm.gateway import LLMCall, LLMGateway, LLMGatewayError SHANGHAI = ZoneInfo("Asia/Shanghai") class MentorError(RuntimeError): pass @dataclass(frozen=True, slots=True) class PreparedChat: call: LLMCall skill: MentorSkill trade_date: str prompt: list[dict[str, str]] class MentorService: def __init__( self, database: Database, repository: MentorRepository, registry: MentorSkillRegistry, context: MentorContextBuilder, gateway: DataGateway, llm: LLMGateway, ) -> None: self._database = database self._repository = repository self._registry = registry self._context = context self._gateway = gateway self._llm = llm def setup(self, principal: Principal, requested_date: str) -> dict[str, Any]: trade_date = self._trade_date(requested_date) skills = self._registry.list(principal.user.is_admin) with self._database.read() as connection: preferences = self._repository.preferences(connection, principal.user.id) rows = [] for fallback_order, skill in enumerate(skills): preference = preferences.get(skill.id, {}) rows.append( { **skill.public(), "pinned": bool(preference.get("pinned")), "sort_order": int(preference.get("sort_order", 10_000 + fallback_order)), } ) rows.sort(key=lambda item: (not item["pinned"], item["sort_order"], item["name"])) return {"trade_date": trade_date, "mentors": rows} def save_preferences( self, principal: Principal, order: list[str], pinned: list[str] ) -> dict[str, Any]: available = [item.id for item in self._registry.list(principal.user.is_admin)] available_set = set(available) normalized = [] for mentor_id in order: if mentor_id not in available_set: raise MentorError("排序中包含不可用的思维模型。") if mentor_id not in normalized: normalized.append(mentor_id) normalized.extend(item for item in available if item not in normalized) pinned_set = set(pinned) if not pinned_set <= available_set: raise MentorError("置顶列表中包含不可用的思维模型。") with self._database.transaction() as connection: self._repository.save_preferences( connection, principal.user.id, normalized, pinned_set, _now(), ) return {"order": normalized, "pinned": [item for item in normalized if item in pinned_set]} def messages( self, principal: Principal, mentor_id: str, requested_date: str ) -> list[dict[str, Any]]: self._skill(principal, mentor_id) trade_date = self._valid_date(requested_date) with self._database.read() as connection: rows = self._repository.messages( connection, principal.user.id, mentor_id, trade_date ) return [ { "id": int(row["id"]), "role": str(row["role"]), "content": str(row["content"]), "status": str(row["status"]), "created_at": str(row["created_at"]), } for row in rows ] def clear(self, principal: Principal, mentor_id: str, requested_date: str) -> int: self._skill(principal, mentor_id) trade_date = self._valid_date(requested_date) with self._database.transaction() as connection: return self._repository.clear_messages( connection, principal.user.id, mentor_id, trade_date ) def prepare_chat( self, principal: Principal, mentor_id: str, requested_date: str, question: str, ) -> PreparedChat: skill = self._skill(principal, mentor_id) trade_date = self._trade_date(requested_date) normalized = " ".join(question.split()) if not normalized or len(normalized) > 2000: raise MentorError("问题应为1至2000个字符。") history = self._history(principal.user.id, mentor_id, trade_date) context = self._context.build(trade_date, normalized, skill) prompt = prompt_messages(skill, context, history, normalized) input_chars = sum(len(item["content"]) for item in prompt) call = self._llm.prepare( principal, feature="mentor", prompt_version=PROMPT_VERSION, business_id=f"{mentor_id}:{trade_date}", input_chars=input_chars, ) with self._database.transaction() as connection: self._repository.add_message( connection, user_id=principal.user.id, mentor_id=mentor_id, trade_date=trade_date, role="user", content=normalized, request_id=call.request_id, status="complete", created_at=_now(), ) return PreparedChat(call, skill, trade_date, prompt) def stream_chat(self, prepared: PreparedChat) -> Iterator[dict[str, Any]]: answer = "" saved = False stream = self._llm.stream(prepared.call, prepared.prompt) try: for event in stream: if event.type == "delta": answer += event.content yield { "type": "delta", "content": event.content, "request_id": event.request_id, } elif event.type == "done": self._save_answer(prepared, answer, "complete") saved = True yield {"type": "done", "request_id": event.request_id} except GeneratorExit: stream.close() if answer and not saved: self._save_answer(prepared, answer, "stopped") raise except LLMGatewayError as exc: if answer: self._save_answer(prepared, answer, "error") yield { "type": "error", "code": exc.code, "message": str(exc), "partial": exc.partial, "request_id": prepared.call.request_id, } def _history(self, user_id: int, mentor_id: str, trade_date: str) -> list[dict[str, str]]: with self._database.read() as connection: rows = self._repository.messages(connection, user_id, mentor_id, trade_date, 40) result = [] total = 0 for row in reversed(rows): if str(row["status"]) != "complete": continue content = str(row["content"]) if total + len(content) > 24_000: break result.append({"role": str(row["role"]), "content": content}) total += len(content) if len(result) == 10: break return list(reversed(result)) def _save_answer(self, prepared: PreparedChat, answer: str, status: str) -> None: if not answer: return with self._database.transaction() as connection: self._repository.add_message( connection, user_id=prepared.call.user_id, mentor_id=prepared.skill.id, trade_date=prepared.trade_date, role="assistant", content=answer, request_id=prepared.call.request_id, status=status, created_at=_now(), ) def _skill(self, principal: Principal, mentor_id: str) -> MentorSkill: try: return self._registry.get(mentor_id, principal.user.is_admin) except MentorSkillError as exc: raise MentorError(str(exc)) from exc def _trade_date(self, requested_date: str) -> str: requested = self._valid_date(requested_date) return self._gateway.trade_context(requested).actual_date or requested @staticmethod def _valid_date(value: str) -> str: try: return date.fromisoformat(value).isoformat() except ValueError as exc: raise MentorError("日期格式无效。") from exc def _now() -> str: return datetime.now(SHANGHAI).isoformat(timespec="seconds")