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"]