356 lines
12 KiB
Python
356 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from backend.database.connection import Database
|
|
from backend.database.migrations import MIGRATIONS, MigrationRunner
|
|
from backend.features.accounts.models import MembershipRecord, Principal, UserRecord
|
|
from backend.features.mentor.context import (
|
|
MentorContextBuilder,
|
|
_matched_identifiers,
|
|
_requires_dragon_context,
|
|
)
|
|
from backend.features.mentor.repository import MentorRepository
|
|
from backend.features.mentor.skills import MentorSkillRegistry
|
|
from backend.llm.gateway import LLMGateway, LLMGatewayError
|
|
from backend.llm.provider import ProviderFailure
|
|
from backend.llm.repository import LLMRepository
|
|
from backend.llm.streaming import TextAccumulator
|
|
|
|
|
|
def _skill(root: Path, identifier: str, title: str, grade: str = "A") -> None:
|
|
directory = root / identifier
|
|
directory.mkdir(parents=True)
|
|
content = (
|
|
f'---\nname: {identifier}\ndescription: "用途：{title}复盘"\n---\n'
|
|
f'# {title}\n\n> "守住边界"\n'
|
|
)
|
|
(directory / "SKILL.md").write_text(
|
|
content,
|
|
encoding="utf-8",
|
|
)
|
|
(root / "mentor_catalog.json").write_text(
|
|
'{"mentors":{"'
|
|
+ identifier
|
|
+ '":{"evidence":{"grade":"'
|
|
+ grade
|
|
+ '","label":"公开资料","note":"已核验"}}}}',
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def _database(tmp_path: Path) -> Database:
|
|
database = Database(tmp_path / "mentor.db")
|
|
MigrationRunner(database).upgrade(MIGRATIONS)
|
|
with database.transaction() as connection:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO users (
|
|
id, username, username_key, password_hash, is_admin,
|
|
status, created_at, updated_at
|
|
) VALUES (1, 'member', 'member', 'hash', 0, 'active', ?, ?)
|
|
""",
|
|
("2026-07-30T00:00:00+00:00", "2026-07-30T00:00:00+00:00"),
|
|
)
|
|
for model_id in (1, 2):
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO llm_models (
|
|
id, display_name, display_name_key, base_url, model_identifier,
|
|
encrypted_api_key, created_at, updated_at, updated_by
|
|
) VALUES (?, ?, ?, 'https://example.invalid/v1', ?, 'secret', ?, ?, 1)
|
|
""",
|
|
(
|
|
model_id,
|
|
f"model-{model_id}",
|
|
f"model-{model_id}",
|
|
f"model-{model_id}",
|
|
"2026-07-30T00:00:00+00:00",
|
|
"2026-07-30T00:00:00+00:00",
|
|
),
|
|
)
|
|
return database
|
|
|
|
|
|
def _principal() -> Principal:
|
|
now = datetime.now(UTC)
|
|
user = UserRecord(1, "member", "member", "hash", False, "active", now, now)
|
|
membership = MembershipRecord(1, "active", None, True, 50, now, 1)
|
|
return Principal("token", "csrf", user, membership)
|
|
|
|
|
|
class _Memberships:
|
|
def view_for(self, _principal: Principal):
|
|
return SimpleNamespace(quota_exempt=False, daily_limit=50)
|
|
|
|
def can_use_smart_features(self, _principal: Principal) -> bool:
|
|
return True
|
|
|
|
|
|
class _ModelPool:
|
|
def runtime_config(self):
|
|
return SimpleNamespace(
|
|
primary=SimpleNamespace(
|
|
id=1,
|
|
base_url="https://example.invalid/v1",
|
|
model_identifier="primary",
|
|
),
|
|
fallback=SimpleNamespace(
|
|
id=2,
|
|
base_url="https://example.invalid/v1",
|
|
model_identifier="fallback",
|
|
),
|
|
)
|
|
|
|
def decrypt_api_key(self, _record) -> str:
|
|
return "secret"
|
|
|
|
|
|
class _ScriptedProvider:
|
|
def __init__(self, scripts: list[object]) -> None:
|
|
self.scripts = scripts
|
|
self.calls: list[str] = []
|
|
|
|
def stream(self, profile, _messages):
|
|
self.calls.append(profile.role)
|
|
script = self.scripts.pop(0)
|
|
if isinstance(script, Exception):
|
|
raise script
|
|
yield from script
|
|
|
|
|
|
def _gateway(database: Database, provider: _ScriptedProvider) -> LLMGateway:
|
|
return LLMGateway(
|
|
database,
|
|
LLMRepository(),
|
|
_Memberships(),
|
|
_ModelPool(),
|
|
provider,
|
|
)
|
|
|
|
|
|
def test_skill_registry_discovers_public_and_protects_private_overrides(tmp_path) -> None:
|
|
public = tmp_path / "public"
|
|
private = tmp_path / "private"
|
|
_skill(public, "same", "公开模型", "B")
|
|
_skill(private, "same", "私有模型", "A")
|
|
|
|
registry = MentorSkillRegistry(public, private)
|
|
public_only = registry.list(False)
|
|
admin = registry.list(True)
|
|
|
|
assert [(item.id, item.name, item.grade, item.private) for item in public_only] == [
|
|
("same", "公开模型", "B", False)
|
|
]
|
|
assert [(item.id, item.name, item.grade, item.private) for item in admin] == [
|
|
("same", "私有模型", "A", True)
|
|
]
|
|
assert "score" not in admin[0].public()
|
|
|
|
|
|
def test_preferences_messages_and_clear_are_strictly_scoped(tmp_path) -> None:
|
|
database = _database(tmp_path)
|
|
repository = MentorRepository()
|
|
with database.transaction() as connection:
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO users (
|
|
id, username, username_key, password_hash, is_admin,
|
|
status, created_at, updated_at
|
|
) VALUES (2, 'other', 'other', 'hash', 0, 'active', 'now', 'now')
|
|
"""
|
|
)
|
|
repository.save_preferences(connection, 1, ["a", "b"], {"b"}, "now")
|
|
repository.save_preferences(connection, 2, ["b", "a"], {"a"}, "now")
|
|
for user_id, mentor_id, trade_date in (
|
|
(1, "a", "2026-07-30"),
|
|
(1, "b", "2026-07-30"),
|
|
(1, "a", "2026-07-29"),
|
|
(2, "a", "2026-07-30"),
|
|
):
|
|
repository.add_message(
|
|
connection,
|
|
user_id=user_id,
|
|
mentor_id=mentor_id,
|
|
trade_date=trade_date,
|
|
role="user",
|
|
content="question",
|
|
request_id=None,
|
|
status="complete",
|
|
created_at="now",
|
|
)
|
|
deleted = repository.clear_messages(connection, 1, "a", "2026-07-30")
|
|
|
|
with database.read() as connection:
|
|
assert repository.preferences(connection, 1)["b"]["pinned"] == 1
|
|
assert repository.preferences(connection, 2)["a"]["pinned"] == 1
|
|
assert deleted == 1
|
|
assert repository.messages(connection, 1, "a", "2026-07-30") == ()
|
|
assert len(repository.messages(connection, 1, "b", "2026-07-30")) == 1
|
|
assert len(repository.messages(connection, 1, "a", "2026-07-29")) == 1
|
|
assert len(repository.messages(connection, 2, "a", "2026-07-30")) == 1
|
|
|
|
|
|
def test_main_failure_before_first_delta_falls_back_and_counts_once(tmp_path) -> None:
|
|
database = _database(tmp_path)
|
|
provider = _ScriptedProvider([ProviderFailure("capacity"), ["回", "答"]])
|
|
gateway = _gateway(database, provider)
|
|
call = gateway.prepare(
|
|
_principal(),
|
|
feature="mentor",
|
|
prompt_version="test",
|
|
business_id="mentor:date",
|
|
input_chars=10,
|
|
)
|
|
|
|
events = list(gateway.stream(call, [{"role": "user", "content": "问题"}]))
|
|
|
|
assert provider.calls == ["primary", "fallback"]
|
|
assert "".join(item.content for item in events if item.type == "delta") == "回答"
|
|
with database.read() as connection:
|
|
request = connection.execute("SELECT * FROM llm_requests").fetchone()
|
|
attempts = connection.execute("SELECT * FROM llm_attempts ORDER BY id").fetchall()
|
|
usage = connection.execute("SELECT successful_calls FROM llm_usage_daily").fetchone()
|
|
assert request["status"] == "success"
|
|
assert [row["status"] for row in attempts] == ["failed", "success"]
|
|
assert usage["successful_calls"] == 1
|
|
|
|
|
|
def test_failure_after_first_delta_keeps_partial_and_never_falls_back(tmp_path) -> None:
|
|
database = _database(tmp_path)
|
|
|
|
def interrupted():
|
|
yield "部分"
|
|
raise ProviderFailure("network")
|
|
|
|
provider = _ScriptedProvider([interrupted(), ["不应调用"]])
|
|
gateway = _gateway(database, provider)
|
|
call = gateway.prepare(
|
|
_principal(),
|
|
feature="mentor",
|
|
prompt_version="test",
|
|
business_id="mentor:date",
|
|
input_chars=10,
|
|
)
|
|
stream = gateway.stream(call, [{"role": "user", "content": "问题"}])
|
|
|
|
first = next(stream)
|
|
assert first.content == "部分"
|
|
with pytest.raises(LLMGatewayError, match="连接中断") as captured:
|
|
list(stream)
|
|
assert captured.value.partial is True
|
|
assert provider.calls == ["primary"]
|
|
with database.read() as connection:
|
|
request = connection.execute("SELECT * FROM llm_requests").fetchone()
|
|
usage = connection.execute("SELECT * FROM llm_usage_daily").fetchone()
|
|
assert request["status"] == "failed"
|
|
assert request["output_chars"] == 2
|
|
assert usage is None
|
|
|
|
|
|
def test_closing_after_first_delta_marks_request_stopped(tmp_path) -> None:
|
|
database = _database(tmp_path)
|
|
provider = _ScriptedProvider([["第一段", "第二段"]])
|
|
gateway = _gateway(database, provider)
|
|
call = gateway.prepare(
|
|
_principal(),
|
|
feature="mentor",
|
|
prompt_version="test",
|
|
business_id="mentor:date",
|
|
input_chars=10,
|
|
)
|
|
stream = gateway.stream(call, [])
|
|
|
|
assert next(stream).content == "第一段"
|
|
stream.close()
|
|
|
|
with database.read() as connection:
|
|
request = connection.execute("SELECT * FROM llm_requests").fetchone()
|
|
attempt = connection.execute("SELECT * FROM llm_attempts").fetchone()
|
|
assert request["status"] == "stopped"
|
|
assert attempt["status"] == "stopped"
|
|
assert request["output_chars"] == 3
|
|
|
|
|
|
def test_unexpected_provider_error_closes_reservation(tmp_path) -> None:
|
|
database = _database(tmp_path)
|
|
gateway = _gateway(database, _ScriptedProvider([RuntimeError("private detail")]))
|
|
call = gateway.prepare(
|
|
_principal(),
|
|
feature="mentor",
|
|
prompt_version="test",
|
|
business_id="mentor:date",
|
|
input_chars=10,
|
|
)
|
|
|
|
with pytest.raises(LLMGatewayError, match="暂不可用"):
|
|
list(gateway.stream(call, []))
|
|
with database.read() as connection:
|
|
request = connection.execute("SELECT * FROM llm_requests").fetchone()
|
|
assert request["status"] == "failed"
|
|
assert "private detail" not in request["error_type"]
|
|
|
|
|
|
def test_stream_accumulator_does_not_repeat_final_snapshot() -> None:
|
|
accumulator = TextAccumulator()
|
|
|
|
chunks = [
|
|
accumulator.feed({"delta": {"content": "第一"}}),
|
|
accumulator.feed({"delta": {"content": "段"}}),
|
|
accumulator.feed({"message": {"content": "第一段"}}),
|
|
]
|
|
|
|
assert "".join(chunks) == "第一段"
|
|
assert accumulator.text == "第一段"
|
|
|
|
|
|
def test_context_profiles_are_distinct_and_dragon_data_is_question_driven() -> None:
|
|
summary = {
|
|
"overview": {"limit_up_count": 20},
|
|
"sentiment": {"temperature": 42},
|
|
"limits": [{"code": "000001", "streak": 3, "amount": 10}],
|
|
"broken": [{"code": "000002"}],
|
|
"sector_rotation": [{"name": "银行"}],
|
|
"ladders": [{"level": 3}],
|
|
}
|
|
popularity = {"combined": [{"code": "000001"}]}
|
|
indexes = [{"code": "000300.SH", "available": False}]
|
|
etfs = [{"code": "510300.SH", "available": False}]
|
|
leader: dict = {}
|
|
macro: dict = {}
|
|
|
|
MentorContextBuilder._apply_profile(
|
|
leader, "leader", summary, popularity, indexes, etfs
|
|
)
|
|
MentorContextBuilder._apply_profile(
|
|
macro, "macro", summary, popularity, indexes, etfs
|
|
)
|
|
|
|
assert leader["multi_board_leaders"][0]["code"] == "000001"
|
|
assert leader["popularity_core"][0]["code"] == "000001"
|
|
assert "broad_indexes" not in leader
|
|
assert macro["broad_indexes"] == indexes
|
|
assert macro["core_etfs"] == etfs
|
|
assert "multi_board_leaders" not in macro
|
|
assert _requires_dragon_context("看看龙虎榜席位") is True
|
|
assert _requires_dragon_context("看看市场情绪") is False
|
|
|
|
|
|
def test_question_stock_matching_is_stable_and_limited_to_two() -> None:
|
|
directory = {
|
|
"000001.SZ": {"symbol": "000001", "name": "平安银行"},
|
|
"000002.SZ": {"symbol": "000002", "name": "万科A"},
|
|
"000003.SZ": {"symbol": "000003", "name": "国华网安"},
|
|
}
|
|
|
|
matched = _matched_identifiers(
|
|
directory,
|
|
"比较平安银行、000002和国华网安的强弱",
|
|
)
|
|
|
|
assert matched == ["000001.SZ", "000002.SZ"]
|