from __future__ import annotations from datetime import datetime, timezone from decimal import Decimal from uuid import uuid4 from schemas.enums import AgentChatMessageRole, AgentChatSessionStatus, SessionType from models.agent_chat_message import AgentChatMessage from models.agent_chat_session import AgentChatSession from v1.agent.anonymizer import ( _aggregate_latency, _extract_derived_fields, _extract_keywords, _extract_model_code, _extract_question_type, _extract_sign_level, _extract_tool_name, _truncate_to_day, anonymize, ) def _make_session(**overrides: object) -> AgentChatSession: defaults: dict[str, object] = { "id": uuid4(), "user_id": uuid4(), "session_type": SessionType.CHAT, "status": AgentChatSessionStatus.COMPLETED, "message_count": 3, "total_tokens": 1500, "total_cost": Decimal("0.05"), "created_at": datetime(2026, 4, 15, 14, 32, 0, tzinfo=timezone.utc), "last_activity_at": datetime(2026, 4, 15, 14, 45, 0, tzinfo=timezone.utc), "job_id": None, "title": "Will I get the job?", "state_snapshot": None, "updated_at": datetime(2026, 4, 15, 14, 45, 0, tzinfo=timezone.utc), "deleted_at": None, } defaults.update(overrides) return AgentChatSession(**defaults) def _make_message( *, session_id: object | None = None, role: AgentChatMessageRole = AgentChatMessageRole.ASSISTANT, metadata_json: dict[str, object] | None = None, model_code: str | None = None, tool_name: str | None = None, latency_ms: int | None = None, ) -> AgentChatMessage: return AgentChatMessage( id=uuid4(), session_id=session_id or uuid4(), seq=1, role=role, content="some content", model_code=model_code, tool_name=tool_name, input_tokens=100, output_tokens=200, cost=Decimal("0.02"), latency_ms=latency_ms, visibility_mask=0, metadata_json=metadata_json, created_at=datetime(2026, 4, 15, 14, 33, 0, tzinfo=timezone.utc), updated_at=datetime(2026, 4, 15, 14, 33, 0, tzinfo=timezone.utc), deleted_at=None, ) def test_truncate_to_day() -> None: dt = datetime(2026, 4, 15, 14, 32, 45, 123456, tzinfo=timezone.utc) result = _truncate_to_day(dt) assert result == datetime(2026, 4, 15, 0, 0, 0, 0, tzinfo=timezone.utc) def test_extract_derived_fields_found() -> None: msg = _make_message( metadata_json={ "agent_output": { "divination_derived": { "guaName": "乾", "questionType": "career", "hasChangingYao": True, } } } ) derived = _extract_derived_fields([msg]) assert derived.get("guaName") == "乾" assert derived.get("questionType") == "career" def test_extract_derived_fields_missing() -> None: msg = _make_message(metadata_json={"run_id": "abc"}) derived = _extract_derived_fields([msg]) assert derived == {} def test_extract_sign_level() -> None: msg = _make_message(metadata_json={"agent_output": {"sign_level": "中上签"}}) assert _extract_sign_level([msg]) == "中上签" def test_extract_sign_level_none() -> None: msg = _make_message(metadata_json={"agent_output": {}}) assert _extract_sign_level([msg]) is None def test_extract_keywords() -> None: msg = _make_message(metadata_json={"agent_output": {"keywords": ["事业", "贵人"]}}) assert _extract_keywords([msg]) == ["事业", "贵人"] def test_extract_question_type_from_derived() -> None: msg = _make_message( metadata_json={ "agent_output": { "divination_derived": {"questionType": "career"}, } } ) assert _extract_question_type([msg]) == "career" def test_extract_question_type_from_agent_output() -> None: msg = _make_message(metadata_json={"agent_output": {"questionType": "love"}}) assert _extract_question_type([msg]) == "love" def test_extract_model_code() -> None: msg = _make_message(model_code="qwen3.5-flash") assert _extract_model_code([msg]) == "qwen3.5-flash" def test_extract_tool_name() -> None: msg = _make_message(tool_name="liuyao") assert _extract_tool_name([msg]) == "liuyao" def test_aggregate_latency() -> None: msg1 = _make_message(latency_ms=500) msg2 = _make_message(latency_ms=300) assert _aggregate_latency([msg1, msg2]) == 800 def test_aggregate_latency_none() -> None: msg = _make_message(latency_ms=None) assert _aggregate_latency([msg]) is None def test_anonymize_full_snapshot() -> None: session = _make_session() msg = _make_message( session_id=session.id, role=AgentChatMessageRole.ASSISTANT, model_code="qwen3.5-flash", tool_name="liuyao", latency_ms=1200, metadata_json={ "agent_output": { "sign_level": "上上签", "keywords": ["事业", "贵人"], "divination_derived": { "questionType": "career", "guaName": "乾", "guaNameHant": "乾", "targetGuaName": "姤", "hasChangingYao": True, }, } }, ) user_msg = _make_message( session_id=session.id, role=AgentChatMessageRole.USER, latency_ms=None, ) snapshot = anonymize(session=session, messages=[msg, user_msg]) assert snapshot.session_type == "chat" assert snapshot.message_count == 3 assert snapshot.status == "completed" assert snapshot.question_type == "career" assert snapshot.tool_name == "liuyao" assert snapshot.model_code == "qwen3.5-flash" assert snapshot.gua_name == "乾" assert snapshot.gua_name_hant == "乾" assert snapshot.target_gua_name == "姤" assert snapshot.has_changing_yao is True assert snapshot.sign_level == "上上签" assert snapshot.keywords == ["事业", "贵人"] assert snapshot.total_tokens == 1500 assert snapshot.total_cost == Decimal("0.05") assert snapshot.total_latency_ms == 1200 assert snapshot.created_at == datetime(2026, 4, 15, 0, 0, 0, tzinfo=timezone.utc) assert snapshot.last_activity_at == datetime( 2026, 4, 15, 0, 0, 0, tzinfo=timezone.utc ) assert snapshot.anonymous_id is not None assert snapshot.id is not None def test_anonymize_no_metadata() -> None: session = _make_session() msg = _make_message(session_id=session.id, metadata_json=None) snapshot = anonymize(session=session, messages=[msg]) assert snapshot.question_type is None assert snapshot.gua_name is None assert snapshot.sign_level is None assert snapshot.keywords is None assert snapshot.has_changing_yao is None