Create-knowledge-find-decode/app/routes/query_generation.py

"""Formal query-generation preview API routes."""
from __future__ import annotations

from typing import Any
from uuid import UUID

from fastapi import APIRouter, Depends, HTTPException, Query

from acquisition.queries.builder import QueryBuildOptions, TREES, build_creation_query_batch
from app.dependencies import _env_file, get_acquisition_repository
from app.routes.acquisition import _model_dump
from app.schemas import QuerySchema, TopicTableQueryPreviewRequest
from core.config import Settings
from core.llm import LLMError, chat_json
from query_planning.sources import (
    PatternTopicBuildConfig,
    PatternTopicBuildSource,
    TopicBuildSourceError,
)
from query_planning.topic_table import (
    PROMPT_NAME,
    PROMPT_VERSION,
    ROUTES,
    TopicTableGenerationError,
    generate_topic_table_preview,
    preview_to_dict,
    topic_table_source_preview_to_dict,
    topic_table_prompt,
)

router = APIRouter(prefix="/api/query-generation", tags=["query-generation"])


def _summary(generated: dict[str, Any]) -> dict[str, Any]:
    families = generated.get("families") or []
    return {
        "family_count": len(families),
        "family_keys": [family.get("key") for family in families],
        "query_count": sum(len(family.get("items") or []) for family in families),
        "kept_count": sum(
            1
            for family in families
            for item in family.get("items") or []
            if item.get("keep", True)
        ),
    }


LIGHT_ITEM_METADATA_KEYS = {
    "page_index",
    "page_rank",
    "source_cursor",
    "content_mode",
    "search_provider",
    "detail_provider",
    "acquisition_match_status",
    "matched_unique_key",
    "skip_reason",
    "unsupported_raw_type",
    "video_url_missing",
}


def _light_metadata(metadata: dict[str, Any] | None) -> dict[str, Any]:
    if not isinstance(metadata, dict):
        return {}
    return {key: metadata[key] for key in LIGHT_ITEM_METADATA_KEYS if key in metadata}


def _normalize_light_query_detail(detail: dict[str, Any]) -> dict[str, Any]:
    media_by_item = {media["item_id"]: media for media in detail.get("media_assets") or []}
    classification_by_item = {row["item_id"]: row for row in detail.get("classifications") or []}
    decode_by_item = {row["item_id"]: row for row in detail.get("decode_summaries") or []}
    platforms: dict[str, dict[str, Any]] = {}
    for raw_job in detail.get("jobs") or []:
        job = _model_dump(raw_job)
        platform = job.get("platform")
        if not platform:
            continue
        platforms.setdefault(
            platform,
            {
                "platform": platform,
                "status": job.get("status") or "pending",
                "attempt_count": job.get("attempt_count"),
                "display_limit": job.get("display_limit"),
                "search_limit": job.get("search_limit"),
                "error_message": job.get("error_message"),
                "item_ids": [],
                "item_count": 0,
            },
        )

    items: list[dict[str, Any]] = []
    for row in detail.get("items") or []:
        item = _model_dump(row)
        item_id = item["id"]
        classification = classification_by_item.get(item_id)
        payload = {
            "id": item_id,
            "platform": item.get("platform"),
            "title": item.get("title"),
            "raw_summary": item.get("raw_summary"),
            "status": item.get("status"),
            "content_mode": item.get("content_mode"),
            "metadata": _light_metadata(item.get("metadata")),
            "classification": _model_dump(classification) if classification else None,
            "decode_summary": _model_dump(decode_by_item[item_id]) if item_id in decode_by_item else None,
            "media_assets": [],
        }
        if item_id in media_by_item:
            payload["media_assets"] = [_model_dump(media_by_item[item_id])]
        items.append(payload)
        platform = payload["platform"]
        group = platforms.setdefault(
            platform,
            {"platform": platform, "status": "done", "item_ids": [], "item_count": 0},
        )
        if group.get("status") in {None, "pending"}:
            group["status"] = "done"
        group.setdefault("item_ids", []).append(item_id)
        group["item_count"] = len(group["item_ids"])

    return {
        "query": QuerySchema.model_validate(detail["query"]).model_dump(mode="json"),
        "run": _model_dump(detail["run"]) if detail.get("run") else None,
        "jobs": [_model_dump(job) for job in detail.get("jobs") or []],
        "items": items,
        "platforms": platforms,
    }


@router.get("/preview")
def query_generation_preview(
    per: int = Query(default=0, ge=0, le=10000),
    batch_n: int = Query(default=0, ge=0, le=1000),
) -> dict[str, Any]:
    """Preview the currently active formal query families without writing DB rows."""

    settings = Settings.from_env(_env_file())
    generated = build_creation_query_batch(
        settings,
        tree_path=TREES,
        options=QueryBuildOptions(
            per=per,
            batch_n=batch_n,
            active_family_keys=("f1", "f2"),
        ),
    )
    generated["summary"] = _summary(generated)
    return generated


def _topic_source(env_file: str) -> PatternTopicBuildSource:
    return PatternTopicBuildSource(PatternTopicBuildConfig.from_env(env_file))


def _topic_query_chat(system: str, user: str, *, settings: Settings) -> dict[str, Any]:
    return chat_json(
        system,
        user,
        settings=settings,
        timeout=settings.bailian_timeout_seconds,
        trace_stage="query_planning",
        trace_substage="topic_table_query_generation",
        prompt_name=PROMPT_NAME,
    )


@router.get("/topic-table/prompt")
def topic_table_generation_prompt() -> dict[str, Any]:
    """Return the exact versioned Prompt shown by the topic-table preview UI."""

    return {
        "name": PROMPT_NAME,
        "version": PROMPT_VERSION,
        "routes": [
            {"key": key, "label": value["label"], "default_priority": value["priority"]}
            for key, value in ROUTES.items()
        ],
        "system_prompt": topic_table_prompt(),
        "example_request": {
            "topic_build_id": 1229,
            "topic_id": 1392,
            "max_queries": 18,
        },
    }


@router.post("/topic-table/preview")
def topic_table_query_preview(request: TopicTableQueryPreviewRequest) -> dict[str, Any]:
    """Generate Query preview from one topic table without persisting or searching."""

    env_file = _env_file()
    try:
        source_payload = _topic_source(env_file).fetch(request.topic_build_id)
        settings = Settings.from_env(env_file)
        preview = generate_topic_table_preview(
            source_payload,
            topic_id=request.topic_id,
            max_queries=request.max_queries,
            chat_fn=lambda system, user: _topic_query_chat(
                system,
                user,
                settings=settings,
            ),
        )
        return preview_to_dict(preview)
    except TopicBuildSourceError as exc:
        raise HTTPException(status_code=502, detail=str(exc)) from exc
    except LLMError as exc:
        raise HTTPException(status_code=502, detail=f"Query 生成模型调用失败：{exc}") from exc
    except TopicTableGenerationError as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc


@router.post("/topic-table/source-preview")
def topic_table_source_preview(request: TopicTableQueryPreviewRequest) -> dict[str, Any]:
    """Read all topic-table source fields without calling an LLM."""

    try:
        source_payload = _topic_source(_env_file()).fetch(request.topic_build_id)
        return topic_table_source_preview_to_dict(
            source_payload,
            topic_id=request.topic_id,
        )
    except TopicBuildSourceError as exc:
        raise HTTPException(status_code=502, detail=str(exc)) from exc
    except TopicTableGenerationError as exc:
        raise HTTPException(status_code=422, detail=str(exc)) from exc


@router.get("/latest-singleton")
def latest_singleton_overview(repo: Any = Depends(get_acquisition_repository)) -> dict[str, Any]:
    """Return links from the query preview board to the latest real singleton run."""

    getter = getattr(repo, "get_latest_singleton_overview", None)
    if getter is None:
        return {"batch": None, "run": None, "queries": [], "decoded_items": []}
    return getter()


@router.get("/latest/queries/{query_id}")
def latest_query_detail(
    query_id: UUID,
    repo: Any = Depends(get_acquisition_repository),
) -> dict[str, Any]:
    """Return search material for one query in the latest real query board batch."""

    detail_getter = getattr(repo, "get_latest_query_result_list", None)
    if detail_getter is None:
        return {"query": None, "jobs": [], "items": [], "platforms": {}}
    return _normalize_light_query_detail(detail_getter(query_id))