Create-knowledge-find-decode/tests/test_topic_table_query_generation.py

from __future__ import annotations

from typing import Any

import pytest
from fastapi.testclient import TestClient

from app.api import app
from app.routes import query_generation
from query_planning.sources.pattern_topic_build import (
    PatternTopicBuildConfig,
    PatternTopicBuildSource,
)
from query_planning.topic_table import (
    ROUTES,
    TopicTableGenerationError,
    generate_topic_table_preview,
    preview_to_dict,
)


def _topic_payload() -> dict[str, Any]:
    return {
        "success": True,
        "source_url": "https://pattern.aiddit.com/api/pattern/topic_builds/1229",
        "build": {
            "id": 1229,
            "demand": "为小羊糊涂啊创作新选题",
            "demand_constraints": None,
            "strategies_config": {"always_on": []},
            "status": "success",
        },
        "topics": [
            {
                "id": 1392,
                "status": "mature",
                "result": "樱花树下低角度拍少女与小羊，传递治愈感并分享穿搭。",
                "topic_direction": {
                    "account_name": "小羊糊涂啊",
                    "persona_dimensions": {
                        "实质偏好": "樱花与穿搭",
                        "形式偏好": "低角度和动态摆拍",
                        "意图偏好": "治愈和种草",
                    },
                },
                "points": [
                    {
                        "id": 8371,
                        "point_type": "灵感点",
                        "point_result": "低角度仰拍小羊和彩虹光晕",
                        "is_active": True,
                        "item_ids": [25684, 25688],
                    },
                    {
                        "id": 8372,
                        "point_type": "目的点",
                        "point_result": "传递治愈感",
                        "is_active": True,
                        "item_ids": [25700],
                    },
                    {
                        "id": 8374,
                        "point_type": "关键点",
                        "point_result": "逆光氛围感特写",
                        "is_active": True,
                        "item_ids": [25684],
                    },
                ],
                "composition_items": [
                    {
                        "id": 25684,
                        "item_level": "element",
                        "dimension": "形式",
                        "point_type": "灵感点",
                        "element_name": "低角度仰拍",
                        "category_path": "/呈现/视觉/低角度",
                        "reason": "增强视觉张力",
                        "is_active": True,
                        "sources": [
                            {
                                "derivation_type": "external_search",
                                "source_type": "search_case",
                                "source_reference_id": "case-1",
                                "reason": "草坪案例证明低角度有动感",
                                "dataset_from": "",
                                "source_reference_data": {"large": "must not enter prompt"},
                            }
                        ],
                    },
                    {
                        "id": 25688,
                        "item_level": "element",
                        "dimension": "实质",
                        "point_type": "灵感点",
                        "element_name": "羊",
                        "reason": "与账号名形成双关",
                        "is_active": True,
                        "sources": [],
                    },
                    {
                        "id": 25700,
                        "item_level": "element",
                        "dimension": "意图",
                        "point_type": "目的点",
                        "element_name": "治愈感",
                        "reason": "满足粉丝对美好生活的向往",
                        "is_active": True,
                        "sources": [],
                    },
                ],
                "item_relations": [
                    {
                        "source_item_id": 25684,
                        "target_item_id": 25688,
                        "reason": "低角度用于突出人与小羊互动",
                    }
                ],
            }
        ],
    }


def _llm_output() -> dict[str, Any]:
    fields = [
        ("point:8371:point_result", "灵感点怎么拍出来"),
        ("point:8372:point_result", "治愈感怎么落实到画面"),
        ("point:8374:point_result", "逆光特写怎么执行"),
        ("item:25688:element_name", "小羊怎么自然入镜"),
        ("item:25684:element_name", "低角度仰拍怎么执行"),
        ("item:25700:element_name", "治愈感怎么通过画面表现"),
        ("topic_direction.persona_dimensions.实质偏好", "怎么保持账号原来的选材"),
        ("topic_direction.persona_dimensions.形式偏好", "怎么保持账号原来的拍法"),
        ("topic_direction.persona_dimensions.意图偏好", "怎么保持账号原来的表达效果"),
        ("item:25684:source:0:reason", "旧案例换场景后怎么调整"),
    ]
    return {
        "field_mappings": [
            {
                "field_key": field_key,
                "judgement": judgement,
                "queries": [
                    f"{judgement}有哪些方法",
                    "逆光人像怎么拍出彩虹光晕" if index == 0 else f"{judgement}常见失败原因",
                ],
            }
            for index, (field_key, judgement) in enumerate(fields)
        ]
    }


def test_topic_source_calls_read_only_detail_endpoint():
    calls = []

    class Response:
        def raise_for_status(self):
            return None

        def json(self):
            return _topic_payload()

    source = PatternTopicBuildSource(
        PatternTopicBuildConfig(base_url="https://pattern.example", timeout_seconds=8),
        http_get=lambda url, **kwargs: calls.append((url, kwargs)) or Response(),
    )

    payload = source.fetch(1229)

    assert calls == [
        ("https://pattern.example/api/pattern/topic_builds/1229", {"timeout": 8})
    ]
    assert payload["build"]["id"] == 1229
    assert payload["source_url"].endswith("/1229")


def test_topic_table_generation_maps_real_fields_and_uses_unified_budget():
    seen = {}

    def chat(system, user):
        seen["system"] = system
        seen["user"] = user
        return _llm_output()

    preview = generate_topic_table_preview(
        _topic_payload(),
        topic_id=1392,
        max_queries=8,
        chat_fn=chat,
    )
    data = preview_to_dict(preview)

    assert data["generator_kind"] == "topic_table"
    assert len(data["knowledge_needs"]) == 10
    assert set(data["summary"]["route_counts"]) == set(ROUTES)
    assert data["summary"]["selected_count"] == 8
    assert data["summary"]["persisted"] is False
    assert data["summary"]["search_started"] is False
    assert len(data["source_groups"]) == 6
    assert all(group["field_mappings"] for group in data["source_groups"])
    assert [group["field_count"] for group in data["source_groups"]] == [1, 1, 1, 5, 3, 1]
    assert sum(len(group["fields"]) for group in data["source_groups"]) == 12
    topic_result_field = next(
        field
        for field in data["source_groups"][3]["fields"]
        if field["field_key"] == "topic.result"
    )
    assert topic_result_field["used_for_query"] is False
    assert topic_result_field["judgement"] is None
    assert topic_result_field["queries"] == []
    assert data["topic"]["dimensions"] == {
        "实质": ["羊"],
        "形式": ["低角度仰拍"],
        "意图": ["治愈感"],
    }
    assert "source_reference_data" not in seen["user"]
    assert "field_catalog" in seen["user"]
    assert "topics[id=1392].composition_items[id=25688].element_name" in seen["user"]
    assert "废弃" not in seen["user"]
    assert "大模型要判断什么" in seen["system"]
    assert data["topic"]["input_groups"]["points"]["灵感点"] == [
        "低角度仰拍小羊和彩虹光晕"
    ]
    assert data["topic"]["input_groups"]["reference_notes"] == [
        {"element": "低角度仰拍", "note": "草坪案例证明低角度有动感"}
    ]
    assert data["field_mappings"][0] == {
        "field_key": "point:8371:point_result",
        "field_path": "topics[id=1392].points[id=8371].point_result",
        "field_label": "灵感点字段",
        "field_value": "低角度仰拍小羊和彩虹光晕",
        "judgement": "灵感点怎么拍出来",
        "source_group_key": "unique_hook",
        "source_group_label": "灵感点",
        "queries": [
            "灵感点怎么拍出来有哪些方法",
            "逆光人像怎么拍出彩虹光晕",
        ],
    }


def test_topic_table_generation_rejects_unknown_field_key():
    malformed = _llm_output()
    malformed["field_mappings"][-1]["field_key"] = "invented-field"

    with pytest.raises(TopicTableGenerationError, match="不存在的选题表字段"):
        generate_topic_table_preview(
            _topic_payload(),
            topic_id=1392,
            max_queries=18,
            chat_fn=lambda _system, _user: malformed,
        )


def test_topic_table_generation_rejects_duplicate_field_key():
    malformed = _llm_output()
    malformed["field_mappings"][-1]["field_key"] = malformed["field_mappings"][0][
        "field_key"
    ]

    with pytest.raises(TopicTableGenerationError, match="重复引用"):
        generate_topic_table_preview(
            _topic_payload(),
            topic_id=1392,
            max_queries=18,
            chat_fn=lambda _system, _user: malformed,
        )


def test_topic_table_generation_rejects_missing_source_group():
    malformed = _llm_output()
    malformed["field_mappings"] = malformed["field_mappings"][:-1]

    with pytest.raises(TopicTableGenerationError, match="全部六种数据来源"):
        generate_topic_table_preview(
            _topic_payload(),
            topic_id=1392,
            max_queries=18,
            chat_fn=lambda _system, _user: malformed,
        )


def test_topic_table_generation_retries_once_when_a_source_group_is_missing():
    malformed = _llm_output()
    malformed["field_mappings"] = malformed["field_mappings"][:-1]
    outputs = iter([malformed, _llm_output()])
    prompts = []

    preview = generate_topic_table_preview(
        _topic_payload(),
        topic_id=1392,
        max_queries=18,
        chat_fn=lambda _system, user: prompts.append(user) or next(outputs),
    )

    assert len(preview.result.selected_queries) == 18
    assert len(prompts) == 2
    assert "上一次输出没有满足字段覆盖要求" in prompts[1]


def test_topic_table_prompt_and_preview_api_do_not_persist_or_search(monkeypatch):
    class Source:
        def fetch(self, topic_build_id):
            assert topic_build_id == 1229
            return _topic_payload()

    monkeypatch.setattr(query_generation, "_topic_source", lambda _env: Source())
    monkeypatch.setattr(
        query_generation,
        "_topic_query_chat",
        lambda system, user, settings: _llm_output(),
    )
    monkeypatch.setattr(
        query_generation.Settings,
        "from_env",
        classmethod(lambda cls, _env: object()),
    )
    client = TestClient(app)

    prompt = client.get("/api/query-generation/topic-table/prompt")
    assert prompt.status_code == 200
    assert prompt.json()["example_request"]["topic_build_id"] == 1229
    assert len(prompt.json()["routes"]) == 6

    response = client.post(
        "/api/query-generation/topic-table/preview",
        json={"topic_build_id": 1229, "topic_id": 1392, "max_queries": 10},
    )
    assert response.status_code == 200
    body = response.json()
    assert body["source"]["topic_id"] == 1392
    assert body["summary"]["selected_count"] == 10
    assert body["summary"]["persisted"] is False
    assert body["summary"]["search_started"] is False


def test_topic_table_source_preview_returns_every_field_without_calling_llm(monkeypatch):
    class Source:
        def fetch(self, topic_build_id):
            assert topic_build_id == 1229
            return _topic_payload()

    monkeypatch.setattr(query_generation, "_topic_source", lambda _env: Source())
    monkeypatch.setattr(
        query_generation,
        "_topic_query_chat",
        lambda *_args, **_kwargs: (_ for _ in ()).throw(
            AssertionError("source preview must not call the LLM")
        ),
    )
    client = TestClient(app)

    response = client.post(
        "/api/query-generation/topic-table/source-preview",
        json={"topic_build_id": 1229, "topic_id": 1392, "max_queries": 18},
    )

    assert response.status_code == 200
    body = response.json()
    assert body["summary"] == {
        "field_count": 12,
        "selected_count": 0,
        "query_generation_completed": False,
        "persisted": False,
        "search_started": False,
    }
    assert [group["field_count"] for group in body["source_groups"]] == [1, 1, 1, 5, 3, 1]
    assert all(
        not field["used_for_query"]
        for group in body["source_groups"]
        for field in group["fields"]
    )