company-ai-platform/tests/test_personalization.py

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()