from datetime import timedelta import pytest from fastapi import HTTPException from sqlalchemy import create_engine, select, update from sqlalchemy.orm import Session, sessionmaker from app.core.config import get_settings from app.core.database import Base from app.core.utils.time import utc_now from app.modules.ai_memory.constants import AIMemoryScope from app.modules.ai_memory.models import AIMemoryEntry from app.modules.ai_memory.service import AIMemoryService from app.modules.business.models import MarketWatchlist from app.modules.feishu_users.models import FeishuUser from app.modules.market.service import MarketService from app.modules.personalization.constants import ( CONVERSATION_MAX_MESSAGES, PersonalizationContextKey, ) from app.modules.personalization.models import ( AIConversation, AIConversationMessage, PersonalDataErasureRequest, UserPreference, ) from app.modules.personalization.services import ( ConversationService, PersonalDataErasureService, PersonalizationContextService, PreferenceService, ) def _session_factory(): engine = create_engine("sqlite://") Base.metadata.create_all(engine) return engine, sessionmaker(bind=engine, expire_on_commit=False) def _user(db: Session, suffix: str, *, open_id: str | None = None) -> FeishuUser: record = FeishuUser( code=f"USR-{suffix}", tenant_key=f"tenant-{suffix}", open_id=open_id or f"open-{suffix}", ) db.add(record) db.commit() db.refresh(record) return record def test_rules_memory_and_watchlists_are_owner_scoped(monkeypatch) -> None: monkeypatch.setenv("AI_MEMORY_ENABLED", "true") monkeypatch.setenv("AI_MEMORY_AUTO_WRITE_ENABLED", "true") get_settings.cache_clear() engine, factory = _session_factory() try: with factory() as db: first = _user(db, "first", open_id="same-open-id") second = _user(db, "second", open_id="same-open-id") memory = AIMemoryService(db) company_rule = memory.create_rule( content="公司规则", scope=AIMemoryScope.GLOBAL, subject="company", priority=90, tags=[], actor="service", ) first_rule = memory.create_rule( content="相同个人规则", scope=AIMemoryScope.USER, subject="profile", priority=50, tags=[], actor="first", owner_id=first.id, ) second_rule = memory.create_rule( content="相同个人规则", scope=AIMemoryScope.USER, subject="profile", priority=50, tags=[], actor="second", owner_id=second.id, ) assert [item["code"] for item in memory.list_rules()] == [company_rule["code"]] assert [item["code"] for item in memory.list_rules(owner_id=first.id)] == [ first_rule["code"] ] assert [item["code"] for item in memory.list_rules(owner_id=second.id)] == [ second_rule["code"] ] with pytest.raises(HTTPException) as exc_info: memory.update_rule( code=second_rule["code"], content="越权修改", priority=None, tags=None, enabled=None, actor="first", owner_id=first.id, ) assert exc_info.value.status_code == 404 first_memory = memory.auto_write( prompt="Remember concise project risk summaries", context={"scope": "user", "subject": "profile"}, answer="Use concise project risk summaries.", actor="first", owner_id=first.id, ) second_memory = memory.auto_write( prompt="Remember concise project risk summaries", context={"scope": "user", "subject": "profile"}, answer="Use concise project risk summaries.", actor="second", owner_id=second.id, ) assert first_memory is not None assert second_memory is not None assert first_memory.code != second_memory.code assert first_memory.fingerprint != second_memory.fingerprint market = MarketService(db) market.add_watchlist("same-open-id", "600000", owner_id=first.id) market.add_watchlist("same-open-id", "600000", owner_id=second.id) assert market.watchlist("ignored", owner_id=first.id) == [{"symbol": "600000.SH"}] assert market.watchlist("ignored", owner_id=second.id) == [{"symbol": "600000.SH"}] market.add_watchlist("legacy-open", "000001") assert market.watchlist("legacy-open") == [{"symbol": "000001.SZ"}] assert market.claim_legacy_watchlist(first.id, "legacy-open") == 1 assert market.watchlist("legacy-open") == [] assert market.watchlist("ignored", owner_id=first.id) == [ {"symbol": "600000.SH"}, {"symbol": "000001.SZ"}, ] finally: get_settings.cache_clear() engine.dispose() def test_preferences_are_allowlisted_sensitive_safe_and_owner_scoped() -> None: engine, factory = _session_factory() try: with factory() as db: first = _user(db, "pref-first") second = _user(db, "pref-second") service = PreferenceService(db) first_pref = service.upsert(first.id, "tone", "简洁直接") second_pref = service.upsert(second.id, "tone", "简洁直接") assert first_pref["code"] != second_pref["code"] assert service.list_preferences(first.id) == [first_pref] with pytest.raises(HTTPException) as exc_info: service.delete(first.id, second_pref["code"]) assert exc_info.value.status_code == 404 with pytest.raises(HTTPException) as exc_info: service.upsert(first.id, "topic", "记住我的银行账号 123456") assert exc_info.value.status_code == 422 assert ( service.save_auto_extraction( first.id, provider_name="noop", user_text="以后请用中文", structured_payload={ "preferences": [{"category": "language", "value": "中文"}] }, ) == [] ) assert ( service.save_auto_extraction( first.id, provider_name="direct_llm", user_text="今天怎么样", structured_payload={ "preferences": [{"category": "language", "value": "中文"}] }, ) == [] ) saved = service.save_auto_extraction( first.id, provider_name="direct_llm", user_text="以后请用中文，并记住我的健康诊断", structured_payload={ "preferences": [ {"category": "language", "value": "中文"}, {"category": "topic", "value": "我的健康诊断"}, ] }, ) assert [(item["category"], item["value"]) for item in saved] == [ ("language", "中文") ] finally: engine.dispose() def test_conversations_keep_twenty_turns_and_isolate_sessions() -> None: engine, factory = _session_factory() try: with factory() as db: first = _user(db, "conversation-first") second = _user(db, "conversation-second") service = ConversationService(db) assert ( service.record_turn( second.id, "private", "chat-1", user_content="placeholder", assistant_content="placeholder", provider_name="noop", ) is False ) for index in range(21): assert service.record_turn( first.id, "private", "chat-1", user_content=f"user-{index}", assistant_content=f"assistant-{index}", provider_name="direct_llm", ) history = service.history(first.id, "private", "chat-1") assert len(history) == CONVERSATION_MAX_MESSAGES assert history[0]["content"] == "user-1" assert history[-1]["content"] == "assistant-20" assert service.history(second.id, "private", "chat-1") == [] assert service.provider_session_id(first.id, "private", "chat-1") == ( service.provider_session_id(first.id, "p2p", "chat-1") ) assert service.provider_session_id(first.id, "group", "chat-1") != ( service.provider_session_id(second.id, "group", "chat-1") ) conversation = db.execute( select(AIConversation).where(AIConversation.owner_id == first.id) ).scalar_one() db.execute( update(AIConversationMessage) .where(AIConversationMessage.conversation_id == conversation.id) .values(created_at=utc_now() - timedelta(days=31)) ) db.commit() assert service.history(first.id, "private", "chat-1") == [] finally: engine.dispose() def test_context_order_and_erasure_core_hooks() -> None: engine, factory = _session_factory() try: with factory() as db: owner = _user(db, "context-owner") other = _user(db, "context-other") memory = AIMemoryService(db) memory.create_rule( content="公司规则优先", scope="global", subject="company", priority=100, tags=[], actor="service", ) memory.create_rule( content="个人规则", scope="global", subject="profile", priority=80, tags=[], actor="owner", owner_id=owner.id, ) PreferenceService(db).upsert(owner.id, "tone", "简洁") PreferenceService(db).upsert(owner.id, "interest", "人工智能") MarketService(db).add_watchlist("owner-open", "600000", owner_id=owner.id) ConversationService(db).record_turn( owner.id, "private", "chat-context", user_content="上一问", assistant_content="上一答", provider_name="direct_llm", ) context = PersonalizationContextService(db).build( owner_id=owner.id, request="当前请求", system_constraints="只读安全规则", chat_type="private", chat_key="chat-context", subject="profile", actor="owner", ) assert [str(key) for key in context.as_ordered_dict()] == [ PersonalizationContextKey.SYSTEM_CONSTRAINTS, PersonalizationContextKey.COMPANY_RULES, PersonalizationContextKey.PERSONAL_RULES, PersonalizationContextKey.CURRENT_REQUEST, PersonalizationContextKey.PREFERENCES, PersonalizationContextKey.INTERESTS, PersonalizationContextKey.PERSONAL_MEMORY, PersonalizationContextKey.CONVERSATION_HISTORY, ] assert context.company_rules[0]["rule"] == "公司规则优先" assert context.personal_rules[0]["rule"] == "个人规则" assert {item["value"] for item in context.interests} == { "人工智能", "600000.SH", } assert context.conversation_history[-1]["content"] == "上一答" confirmation = PersonalDataErasureService(db).request_confirmation(owner.id) with pytest.raises(HTTPException): PersonalDataErasureService(db).confirm_and_erase(owner.id, "BAD-CODE") seen: dict[str, str | int] = {} def integration_hook( _db: Session, owner_id: int, anonymous_id: str, ) -> dict[str, int]: seen.update(owner_id=owner_id, anonymous_id=anonymous_id) return {"subscriptions": 0} erased = PersonalDataErasureService(db).confirm_and_erase( owner.id, confirmation.confirmation_code, extra_hooks=(integration_hook,), ) assert erased.deleted["preferences"] == 2 assert erased.deleted["conversations"] == 1 assert erased.deleted["ai_memory"] == 1 assert erased.deleted["watchlist"] == 1 assert seen["owner_id"] == owner.id assert seen["anonymous_id"] == erased.anonymous_id assert db.scalar( select(UserPreference).where(UserPreference.owner_id == owner.id) ) is None assert db.scalar( select(AIMemoryEntry).where(AIMemoryEntry.owner_id == owner.id) ) is None assert db.scalar( select(MarketWatchlist).where(MarketWatchlist.owner_id == owner.id) ) is None assert db.scalar( select(PersonalDataErasureRequest).where( PersonalDataErasureRequest.owner_id == owner.id ) ) is None assert db.get(FeishuUser, other.id) is not None assert memory.active_rules()[0]["rule"] == "公司规则优先" finally: engine.dispose()