245 lines
9.1 KiB
Python
245 lines
9.1 KiB
Python
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")
|