Create-knowledge-find-decode/query_planning/topic_table.py

"""Application service for deriving Query candidates from one topic-table snapshot."""
from __future__ import annotations

import json
import re
from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Callable

from core.prompts import load_prompt
from query_planning.domain import GenerationRequest, GenerationResult, GeneratorKind
from query_planning.service import UnifiedQueryGenerationService


PROMPT_NAME = "topic_table_query_generation"
PROMPT_VERSION = "topic_table_query_v4"

ROUTES: dict[str, dict[str, Any]] = {
    "unique_hook": {
        "label": "灵感点",
        "source_description": "来自选题的灵感点字段",
        "priority": 0,
    },
    "goal_execution": {
        "label": "创作目标",
        "source_description": "来自目的点字段",
        "priority": 0,
    },
    "key_execution": {
        "label": "关键镜头和动作",
        "source_description": "来自关键点字段",
        "priority": 0,
    },
    "cross_dimension": {
        "label": "实质、形式、意图及元素关系",
        "source_description": "来自组成元素、最终选题描述和元素关系字段",
        "priority": 0,
    },
    "account_fit": {
        "label": "账号内容偏好",
        "source_description": "来自账号的实质偏好、形式偏好和意图偏好字段",
        "priority": 0,
    },
    "evidence_gap": {
        "label": "参考案例和选取理由",
        "source_description": "来自每个元素引用的参考来源说明字段",
        "priority": 0,
    },
}


class TopicTableGenerationError(ValueError):
    """Raised when a topic snapshot or LLM result cannot become a Query plan."""


@dataclass(frozen=True)
class TopicTablePreview:
    result: GenerationResult
    source: dict[str, Any]
    topic: dict[str, Any]
    field_catalog: list[dict[str, Any]]
    llm_output: dict[str, Any]


def topic_table_prompt() -> str:
    return load_prompt(PROMPT_NAME)


def _select_topic(payload: dict[str, Any], topic_id: int | None) -> dict[str, Any]:
    topics = [topic for topic in payload.get("topics") or [] if isinstance(topic, dict)]
    if topic_id is not None:
        for topic in topics:
            if int(topic.get("id") or 0) == topic_id:
                return topic
        raise TopicTableGenerationError(f"选题构建中找不到 topic_id={topic_id}")
    mature = next((topic for topic in topics if topic.get("status") == "mature"), None)
    if mature is not None:
        return mature
    if topics:
        return topics[0]
    raise TopicTableGenerationError("选题构建中没有可用 topic")


def _source_summary(source: dict[str, Any]) -> dict[str, Any]:
    return {
        "derivation_type": source.get("derivation_type"),
        "source_type": source.get("source_type"),
        "source_reference_id": source.get("source_reference_id"),
        "reason": source.get("reason"),
        "dataset_from": source.get("dataset_from"),
    }


def _compact_snapshot(
    payload: dict[str, Any],
    topic: dict[str, Any],
) -> dict[str, Any]:
    build = payload.get("build") or {}
    items = []
    for item in topic.get("composition_items") or []:
        if not isinstance(item, dict) or item.get("is_active") is False:
            continue
        items.append(
            {
                "id": item.get("id"),
                "item_level": item.get("item_level"),
                "dimension": item.get("dimension"),
                "point_type": item.get("point_type"),
                "element_name": item.get("element_name"),
                "category_path": item.get("category_path"),
                "reason": item.get("reason"),
                "is_active": item.get("is_active"),
                "sources": [
                    _source_summary(source)
                    for source in item.get("sources") or []
                    if isinstance(source, dict)
                ],
            }
        )
    return {
        "topic_build_id": build.get("id"),
        "topic_id": topic.get("id"),
        "build": {
            "demand": build.get("demand"),
            "demand_constraints": build.get("demand_constraints"),
            "strategies_config": build.get("strategies_config"),
            "status": build.get("status"),
        },
        "topic": {
            "status": topic.get("status"),
            "result": topic.get("result"),
            "topic_direction": topic.get("topic_direction"),
            "points": [
                deepcopy(point)
                for point in topic.get("points") or []
                if isinstance(point, dict) and point.get("is_active") is not False
            ],
            "composition_items": items,
            "item_relations": deepcopy(topic.get("item_relations") or []),
        },
    }


def build_topic_table_user_prompt(
    snapshot: dict[str, Any],
    field_catalog: list[dict[str, Any]],
    *,
    target_query_count: int,
) -> str:
    return (
        "请从 field_catalog 中选择真正能产生创作方法搜索词的具体字段。"
        "field_catalog 已经按照六种数据来源标明 source_group_key 和 source_group_label。"
        "每一项 field_mapping 只能引用一个真实 field_key；不能合并多个字段，也不能自己编字段。"
        "六种数据来源必须全部覆盖，每种来源至少选择一个字段。"
        f"所有字段合计生成约 {target_query_count} 条具体搜索词，每个字段生成 1～3 条。"
        "judgement 必须是一句不超过 45 个汉字的大白话，只说明大模型针对该字段要判断什么。"
        "不要提前编造答案、参数、比例或效果数据，不要启动搜索。\n\n"
        + json.dumps(
            {
                "topic_build_id": snapshot["topic_build_id"],
                "topic_id": snapshot["topic_id"],
                "field_catalog": field_catalog,
            },
            ensure_ascii=False,
            indent=2,
            default=str,
        )
    )


def _field_catalog(snapshot: dict[str, Any]) -> list[dict[str, Any]]:
    topic = snapshot["topic"]
    topic_id = snapshot["topic_id"]
    fields: list[dict[str, Any]] = []

    def add(
        *,
        field_key: str,
        field_path: str,
        field_label: str,
        field_value: Any,
        route: str,
        point_id: Any = None,
        item_ids: list[Any] | None = None,
        is_current: bool = True,
    ) -> None:
        value = str(field_value or "").strip()
        if not value:
            return
        fields.append(
            {
                "field_key": field_key,
                "field_path": field_path,
                "field_label": field_label,
                "field_value": value,
                "source_group_key": route,
                "source_group_label": ROUTES[route]["label"],
                "point_id": point_id,
                "item_ids": list(item_ids or []),
                "is_current": is_current,
            }
        )

    add(
        field_key="topic.result",
        field_path=f"topics[id={topic_id}].result",
        field_label="最终选题描述",
        field_value=topic.get("result"),
        route="cross_dimension",
    )
    direction = topic.get("topic_direction") or {}
    for preference, label in (
        ("实质偏好", "账号实质偏好"),
        ("形式偏好", "账号形式偏好"),
        ("意图偏好", "账号意图偏好"),
    ):
        add(
            field_key=f"topic_direction.persona_dimensions.{preference}",
            field_path=f"topics[id={topic_id}].topic_direction.persona_dimensions.{preference}",
            field_label=label,
            field_value=(direction.get("persona_dimensions") or {}).get(preference),
            route="account_fit",
        )

    point_routes = {
        "灵感点": "unique_hook",
        "目的点": "goal_execution",
        "关键点": "key_execution",
    }
    for point in topic.get("points") or []:
        point_id = point.get("id")
        point_type = str(point.get("point_type") or "选题点")
        add(
            field_key=f"point:{point_id}:point_result",
            field_path=f"topics[id={topic_id}].points[id={point_id}].point_result",
            field_label=f"{point_type}字段",
            field_value=point.get("point_result"),
            route=point_routes.get(point_type, "cross_dimension"),
            point_id=point_id,
            item_ids=list(point.get("item_ids") or []),
        )

    for item in topic.get("composition_items") or []:
        item_id = item.get("id")
        dimension = str(item.get("dimension") or "元素")
        add(
            field_key=f"item:{item_id}:element_name",
            field_path=(
                f"topics[id={topic_id}].composition_items[id={item_id}].element_name"
            ),
            field_label=f"{dimension}字段",
            field_value=item.get("element_name"),
            route="cross_dimension",
            item_ids=[item_id],
        )
        for source_index, source in enumerate(item.get("sources") or []):
            add(
                field_key=f"item:{item_id}:source:{source_index}:reason",
                field_path=(
                    f"topics[id={topic_id}].composition_items[id={item_id}]"
                    f".sources[{source_index}].reason"
                ),
                field_label="参考来源说明",
                field_value=source.get("reason"),
                route="evidence_gap",
                item_ids=[item_id],
            )
    active_item_ids = {
        item.get("id") for item in topic.get("composition_items") or []
    }
    for relation_index, relation in enumerate(topic.get("item_relations") or []):
        relation_item_ids = [
            value
            for value in (
                relation.get("source_item_id"),
                relation.get("target_item_id"),
            )
            if value is not None
        ]
        add(
            field_key=f"relation:{relation_index}:reason",
            field_path=f"topics[id={topic_id}].item_relations[{relation_index}].reason",
            field_label="元素关系说明",
            field_value=relation.get("reason"),
            route="cross_dimension",
            item_ids=relation_item_ids,
            is_current=all(item_id in active_item_ids for item_id in relation_item_ids),
        )
    return fields


def _safe_need_key(value: Any, route: str, index: int) -> str:
    raw = re.sub(r"[^a-zA-Z0-9_-]+", "_", str(value or "").strip()).strip("_")
    return raw[:80] or f"{route}_{index + 1}"


def _planning_payload(
    llm_output: dict[str, Any],
    *,
    snapshot: dict[str, Any],
    field_catalog: list[dict[str, Any]],
) -> dict[str, Any]:
    raw_mappings = llm_output.get("field_mappings")
    if not isinstance(raw_mappings, list) or not raw_mappings:
        raise TopicTableGenerationError("大模型没有返回字段与搜索词的对应关系")
    topic_build_id = snapshot["topic_build_id"]
    topic_id = snapshot["topic_id"]
    catalog_by_key = {field["field_key"]: field for field in field_catalog}
    needs: list[dict[str, Any]] = []
    route_mapping_candidates: dict[str, list[list[dict[str, Any]]]] = {
        route: [] for route in ROUTES
    }
    used_field_keys: set[str] = set()
    used_need_keys: set[str] = set()
    route_mapping_counts: dict[str, int] = {route: 0 for route in ROUTES}
    selected_labels_by_route: dict[str, set[str]] = {
        route: set() for route in ROUTES
    }
    for index, raw in enumerate(raw_mappings):
        if not isinstance(raw, dict):
            continue
        field_key = str(raw.get("field_key") or "").strip()
        field = catalog_by_key.get(field_key)
        if field is None:
            raise TopicTableGenerationError(f"大模型引用了不存在的选题表字段：{field_key}")
        if field_key in used_field_keys:
            raise TopicTableGenerationError(f"大模型重复引用了同一个选题表字段：{field_key}")
        used_field_keys.add(field_key)
        judgement = str(raw.get("judgement") or "").strip()
        if not judgement:
            continue
        raw_queries = raw.get("queries") or []
        if isinstance(raw_queries, str):
            raw_queries = [raw_queries]
        query_texts = [
            str(query or "").strip()
            for query in raw_queries
            if str(query or "").strip()
        ]
        if not query_texts:
            continue
        route = field["source_group_key"]
        route_mapping_counts[route] += 1
        selected_labels_by_route[route].add(field["field_label"])
        priority = ROUTES[route]["priority"]
        need_key = _safe_need_key(field_key, route, index)
        if need_key in used_need_keys:
            need_key = f"{need_key}_{index + 1}"
        used_need_keys.add(need_key)
        source_ref = {
            "generator_kind": GeneratorKind.TOPIC_TABLE.value,
            "topic_build_id": topic_build_id,
            "topic_id": topic_id,
            "route": route,
            "field_key": field["field_key"],
            "field_path": field["field_path"],
            "field_label": field["field_label"],
            "field_value": field["field_value"],
            "point_id": field.get("point_id"),
            "item_ids": list(field.get("item_ids") or []),
        }
        needs.append(
            {
                "need_key": need_key,
                "decision_context": judgement,
                "unknown_information": judgement,
                "source_ref": source_ref,
                "priority": priority,
                "metadata": {
                    "route": route,
                    "route_label": ROUTES[route]["label"],
                    "source_field": field,
                },
            }
        )
        mapping_candidates: list[dict[str, Any]] = []
        for query_text in query_texts:
            mapping_candidates.append(
                {
                    "query_text": query_text,
                    "axes": {
                        "选题表字段": field["field_path"],
                        "字段值": field["field_value"],
                        "大模型判断": judgement,
                    },
                    "priority": priority,
                    "original_position": 0,
                    "source_refs": [source_ref],
                    "knowledge_need_keys": [need_key],
                    "metadata": {
                        "family_key": GeneratorKind.TOPIC_TABLE.value,
                        "route": route,
                        "route_label": ROUTES[route]["label"],
                        "need_key": need_key,
                    },
                }
            )
        route_mapping_candidates[route].append(mapping_candidates)
    missing_routes = [
        ROUTES[route]["label"]
        for route, count in route_mapping_counts.items()
        if count == 0
    ]
    if missing_routes:
        raise TopicTableGenerationError(
            "大模型没有覆盖全部六种数据来源：" + "、".join(missing_routes)
        )
    required_dimension_labels = {
        field["field_label"]
        for field in field_catalog
        if field["source_group_key"] == "cross_dimension"
        and field["field_label"] in {"实质字段", "形式字段", "意图字段"}
    }
    missing_dimension_labels = sorted(
        required_dimension_labels - selected_labels_by_route["cross_dimension"]
    )
    if missing_dimension_labels:
        raise TopicTableGenerationError(
            "组成元素没有同时覆盖实质、形式、意图字段："
            + "、".join(missing_dimension_labels)
        )
    required_account_labels = {
        field["field_label"]
        for field in field_catalog
        if field["source_group_key"] == "account_fit"
    }
    missing_account_labels = sorted(
        required_account_labels - selected_labels_by_route["account_fit"]
    )
    if missing_account_labels:
        raise TopicTableGenerationError(
            "账号内容偏好没有覆盖完整：" + "、".join(missing_account_labels)
        )
    route_candidates: dict[str, list[dict[str, Any]]] = {
        route: [] for route in ROUTES
    }
    for route, mapping_lists in route_mapping_candidates.items():
        max_length = max((len(items) for items in mapping_lists), default=0)
        for query_index in range(max_length):
            for mapping_candidates in mapping_lists:
                if query_index < len(mapping_candidates):
                    route_candidates[route].append(mapping_candidates[query_index])
    candidates: list[dict[str, Any]] = []
    position = 0
    while any(route_candidates.values()):
        for route in ROUTES:
            if not route_candidates[route]:
                continue
            candidate = route_candidates[route].pop(0)
            candidate["original_position"] = position
            candidates.append(candidate)
            position += 1
    if not needs or not candidates:
        raise TopicTableGenerationError("大模型返回内容无法形成字段与搜索词的对应关系")
    return {
        "input_snapshot": snapshot,
        "knowledge_needs": needs,
        "candidates": candidates,
        "generator_config": {
            "prompt_name": PROMPT_NAME,
            "prompt_version": PROMPT_VERSION,
            "field_count": len(field_catalog),
        },
    }


def generate_topic_table_preview(
    source_payload: dict[str, Any],
    *,
    topic_id: int | None,
    max_queries: int | None,
    chat_fn: Callable[[str, str], dict[str, Any]],
) -> TopicTablePreview:
    topic = _select_topic(source_payload, topic_id)
    snapshot = _compact_snapshot(source_payload, topic)
    field_catalog = _field_catalog(snapshot)
    system_prompt = topic_table_prompt()
    user_prompt = build_topic_table_user_prompt(
        snapshot,
        field_catalog,
        target_query_count=max_queries or 18,
    )
    llm_output: dict[str, Any] = {}
    payload: dict[str, Any] | None = None
    validation_error: TopicTableGenerationError | None = None
    for attempt in range(2):
        current_user_prompt = user_prompt
        if attempt and validation_error is not None:
            current_user_prompt += (
                "\n\n上一次输出没有满足字段覆盖要求。"
                f"具体问题：{validation_error}。"
                "请重新检查六种数据来源、实质/形式/意图字段和三种账号偏好后，"
                "重新输出完整 JSON。"
            )
        candidate_output = chat_fn(system_prompt, current_user_prompt)
        if not isinstance(candidate_output, dict):
            validation_error = TopicTableGenerationError("大模型返回格式不正确")
            continue
        llm_output = candidate_output
        try:
            payload = _planning_payload(
                llm_output,
                snapshot=snapshot,
                field_catalog=field_catalog,
            )
            break
        except TopicTableGenerationError as exc:
            validation_error = exc
    if payload is None:
        raise validation_error or TopicTableGenerationError("大模型返回格式不正确")
    request = GenerationRequest(
        generator_kind=GeneratorKind.TOPIC_TABLE,
        payload=payload,
        name=f"topic-table-{snapshot['topic_build_id']}-{snapshot['topic_id']}",
        max_queries=max_queries,
        metadata={
            "source": "pattern_topic_build_api",
            "prompt_version": PROMPT_VERSION,
        },
    )
    result = UnifiedQueryGenerationService().generate(request)
    return TopicTablePreview(
        result=result,
        source={
            "topic_build_id": snapshot["topic_build_id"],
            "topic_id": snapshot["topic_id"],
            "source_url": source_payload.get("source_url"),
        },
        topic=snapshot["topic"],
        field_catalog=field_catalog,
        llm_output=llm_output,
    )


def topic_table_source_preview_to_dict(
    source_payload: dict[str, Any],
    *,
    topic_id: int | None,
) -> dict[str, Any]:
    """Return all six source groups without calling an LLM or generating Query."""

    topic = _select_topic(source_payload, topic_id)
    snapshot = _compact_snapshot(source_payload, topic)
    field_catalog = _field_catalog(snapshot)
    source_groups = []
    for route, config in ROUTES.items():
        fields = [
            field
            for field in field_catalog
            if field["source_group_key"] == route
        ]
        source_groups.append(
            {
                "key": route,
                "label": config["label"],
                "source_description": config["source_description"],
                "field_count": len(fields),
                "selected_field_count": 0,
                "historical_field_count": sum(
                    1 for field in fields if not field.get("is_current", True)
                ),
                "fields": [
                    {
                        "field_key": field["field_key"],
                        "field_path": field["field_path"],
                        "field_label": field["field_label"],
                        "field_value": field["field_value"],
                        "is_current": field.get("is_current", True),
                        "used_for_query": False,
                        "judgement": None,
                        "queries": [],
                    }
                    for field in fields
                ],
                "field_mappings": [],
            }
        )
    return {
        "generator_kind": GeneratorKind.TOPIC_TABLE.value,
        "source": {
            "topic_build_id": snapshot["topic_build_id"],
            "topic_id": snapshot["topic_id"],
            "source_url": source_payload.get("source_url"),
        },
        "topic": {
            "status": snapshot["topic"].get("status"),
            "result": snapshot["topic"].get("result"),
        },
        "source_groups": source_groups,
        "field_mappings": [],
        "queries": [],
        "summary": {
            "field_count": len(field_catalog),
            "selected_count": 0,
            "query_generation_completed": False,
            "persisted": False,
            "search_started": False,
        },
        "prompt": {"name": PROMPT_NAME, "version": PROMPT_VERSION},
    }


def preview_to_dict(preview: TopicTablePreview) -> dict[str, Any]:
    result = preview.result
    route_counts: dict[str, int] = {route: 0 for route in ROUTES}
    for need in result.knowledge_needs:
        route = str(need.metadata.get("route") or "")
        if route in route_counts:
            route_counts[route] += 1
    dimensions: dict[str, list[str]] = {"实质": [], "形式": [], "意图": []}
    reference_notes: list[dict[str, str]] = []
    for item in preview.topic.get("composition_items") or []:
        dimension = item.get("dimension")
        name = item.get("element_name")
        if dimension in dimensions and name and item.get("is_active"):
            dimensions[dimension].append(name)
        for source in item.get("sources") or []:
            note = str(source.get("reason") or item.get("reason") or "").strip()
            if name and note:
                reference_notes.append({"element": str(name), "note": note})
    point_groups: dict[str, list[str]] = {"灵感点": [], "目的点": [], "关键点": []}
    for point in preview.topic.get("points") or []:
        point_type = str(point.get("point_type") or "")
        point_result = str(point.get("point_result") or "").strip()
        if point_type in point_groups and point_result:
            point_groups[point_type].append(point_result)
    queries_by_need: dict[str, list[str]] = {}
    for query in result.selected_queries:
        need_keys = query.knowledge_need_keys or [str(query.metadata.get("need_key") or "")]
        for need_key in need_keys:
            if need_key:
                queries_by_need.setdefault(need_key, []).append(query.query_text)
    selected_by_field_key = {
        str(need.source_ref.get("field_key")): {
            "judgement": need.decision_context,
            "queries": queries_by_need.get(need.need_key, []),
        }
        for need in result.knowledge_needs
        if need.source_ref.get("field_key")
    }
    return {
        "generator_kind": GeneratorKind.TOPIC_TABLE.value,
        "source": preview.source,
        "topic": {
            "status": preview.topic.get("status"),
            "result": preview.topic.get("result"),
            "topic_direction": preview.topic.get("topic_direction"),
            "dimensions": dimensions,
            "points": preview.topic.get("points") or [],
            "input_groups": {
                "points": point_groups,
                "reference_notes": reference_notes[:8],
            },
        },
        "knowledge_needs": [
            {
                "need_key": need.need_key,
                "route": need.metadata.get("route"),
                "route_label": need.metadata.get("route_label"),
                "decision_context": need.decision_context,
                "unknown_information": need.unknown_information,
                "priority": need.priority,
                "source_ref": need.source_ref,
            }
            for need in result.knowledge_needs
        ],
        "field_mappings": [
            {
                "field_key": need.source_ref.get("field_key"),
                "field_path": need.source_ref.get("field_path"),
                "field_label": need.source_ref.get("field_label"),
                "field_value": need.source_ref.get("field_value"),
                "judgement": need.decision_context,
                "source_group_key": need.metadata.get("route"),
                "source_group_label": need.metadata.get("route_label"),
                "queries": queries_by_need.get(need.need_key, []),
            }
            for need in result.knowledge_needs
            if queries_by_need.get(need.need_key)
        ],
        "source_groups": [
            {
                "key": route,
                "label": config["label"],
                "source_description": config["source_description"],
                "field_count": sum(
                    1
                    for field in preview.field_catalog
                    if field["source_group_key"] == route
                ),
                "selected_field_count": sum(
                    1
                    for field in preview.field_catalog
                    if field["source_group_key"] == route
                    and field["field_key"] in selected_by_field_key
                    and selected_by_field_key[field["field_key"]]["queries"]
                ),
                "historical_field_count": sum(
                    1
                    for field in preview.field_catalog
                    if field["source_group_key"] == route
                    and not field.get("is_current", True)
                ),
                "fields": [
                    {
                        "field_key": field["field_key"],
                        "field_path": field["field_path"],
                        "field_label": field["field_label"],
                        "field_value": field["field_value"],
                        "is_current": field.get("is_current", True),
                        "used_for_query": bool(
                            selected_by_field_key.get(field["field_key"], {}).get(
                                "queries"
                            )
                        ),
                        "judgement": selected_by_field_key.get(
                            field["field_key"], {}
                        ).get("judgement"),
                        "queries": selected_by_field_key.get(
                            field["field_key"], {}
                        ).get("queries", []),
                    }
                    for field in preview.field_catalog
                    if field["source_group_key"] == route
                ],
                "field_mappings": [
                    {
                        "field_key": need.source_ref.get("field_key"),
                        "field_path": need.source_ref.get("field_path"),
                        "field_label": need.source_ref.get("field_label"),
                        "field_value": need.source_ref.get("field_value"),
                        "judgement": need.decision_context,
                        "queries": queries_by_need.get(need.need_key, []),
                    }
                    for need in result.knowledge_needs
                    if need.metadata.get("route") == route
                    and queries_by_need.get(need.need_key)
                ],
            }
            for route, config in ROUTES.items()
        ],
        "queries": [
            {
                "query": query.query_text,
                "priority": query.priority,
                "route": query.metadata.get("route"),
                "route_label": query.metadata.get("route_label"),
                "need_key": query.metadata.get("need_key"),
                "axes": query.axes,
                "source_refs": query.source_refs,
            }
            for query in result.selected_queries
        ],
        "summary": {
            "generated_count": result.stats.generated_count,
            "unique_count": result.stats.unique_count,
            "selected_count": result.stats.selected_count,
            "dropped_count": result.stats.dropped_count,
            "route_counts": route_counts,
            "query_generation_completed": True,
            "persisted": False,
            "search_started": False,
        },
        "prompt": {"name": PROMPT_NAME, "version": PROMPT_VERSION},
    }