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_v2"

ROUTES: dict[str, dict[str, Any]] = {
    "unique_hook": {"label": "灵感点和特殊组合", "priority": 100},
    "goal_execution": {"label": "目标和想达到的效果", "priority": 90},
    "key_execution": {"label": "关键镜头和动作要求", "priority": 80},
    "cross_dimension": {"label": "三类元素怎么放在一起", "priority": 70},
    "account_fit": {"label": "账号原来的内容风格", "priority": 60},
    "evidence_gap": {"label": "参考案例哪些地方不能照搬", "priority": 50},
}


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]
    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],
    *,
    target_query_count: int,
) -> str:
    return (
        "请根据下面的真实选题表快照，从六类信息中各判断一个需要搜索的问题。"
        "knowledge_needs 数组必须正好有 6 项，每类信息正好 1 项，不能把搜索词数量误当成问题数量。"
        f"六个问题合计生成约 {target_query_count} 条具体搜索词；数量允许时，每个问题给 2～3 条。"
        "decision_context 必须是一句不超过 45 个汉字的大白话，只说大模型要判断什么，不要先解释背景。"
        "只输出待搜索的问题，不要在 decision_context 或 unknown_information 中提前编造答案、参数、比例或效果数据。"
        "不要启动搜索。\n\n"
        + json.dumps(snapshot, ensure_ascii=False, indent=2, default=str)
    )


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],
) -> dict[str, Any]:
    raw_needs = llm_output.get("knowledge_needs")
    if not isinstance(raw_needs, list) or not raw_needs:
        raise TopicTableGenerationError("大模型没有返回六类搜索问题")
    topic_build_id = snapshot["topic_build_id"]
    topic_id = snapshot["topic_id"]
    needs: list[dict[str, Any]] = []
    candidates: list[dict[str, Any]] = []
    used_keys: set[str] = set()
    route_counts: dict[str, int] = {route: 0 for route in ROUTES}
    position = 0
    for index, raw in enumerate(raw_needs):
        if not isinstance(raw, dict):
            continue
        route = str(raw.get("route") or "").strip()
        if route not in ROUTES:
            continue
        route_counts[route] += 1
        need_key = _safe_need_key(raw.get("need_key"), route, index)
        if need_key in used_keys:
            need_key = f"{need_key}_{index + 1}"
        used_keys.add(need_key)
        decision_context = str(raw.get("decision_context") or "").strip()
        unknown_information = str(raw.get("unknown_information") or "").strip()
        if not decision_context or not unknown_information:
            continue
        priority = max(0, min(100, int(raw.get("priority") or ROUTES[route]["priority"])))
        raw_ref = raw.get("source_ref") if isinstance(raw.get("source_ref"), dict) else {}
        source_ref = {
            "generator_kind": GeneratorKind.TOPIC_TABLE.value,
            "topic_build_id": topic_build_id,
            "topic_id": topic_id,
            "route": route,
            "point_id": raw_ref.get("point_id"),
            "item_ids": list(raw_ref.get("item_ids") or []),
        }
        needs.append(
            {
                "need_key": need_key,
                "decision_context": decision_context,
                "unknown_information": unknown_information,
                "source_ref": source_ref,
                "priority": priority,
                "metadata": {"route": route, "route_label": ROUTES[route]["label"]},
            }
        )
        raw_queries = raw.get("queries") or []
        if isinstance(raw_queries, str):
            raw_queries = [raw_queries]
        for query in raw_queries:
            query_text = str(query or "").strip()
            if not query_text:
                continue
            candidates.append(
                {
                    "query_text": query_text,
                    "axes": {
                        "需求路径": ROUTES[route]["label"],
                        "创作决定": decision_context,
                        "未知信息": unknown_information,
                    },
                    "priority": priority,
                    "original_position": position,
                    "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,
                    },
                }
            )
            position += 1
    invalid_routes = [
        ROUTES[route]["label"]
        for route, count in route_counts.items()
        if count != 1
    ]
    if invalid_routes:
        raise TopicTableGenerationError(
            "大模型返回的六类问题不完整或有重复：" + "、".join(invalid_routes)
        )
    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,
            "routes": list(ROUTES),
        },
    }


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)
    llm_output = chat_fn(
        topic_table_prompt(),
        build_topic_table_user_prompt(
            snapshot,
            target_query_count=max_queries or 18,
        ),
    )
    if not isinstance(llm_output, dict):
        raise TopicTableGenerationError("大模型返回格式不正确")
    payload = _planning_payload(llm_output, snapshot=snapshot)
    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"],
        llm_output=llm_output,
    )


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)
    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
        ],
        "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,
            "persisted": False,
            "search_started": False,
        },
        "prompt": {"name": PROMPT_NAME, "version": PROMPT_VERSION},
    }