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

"""Manual query batch submission API."""
from __future__ import annotations

import os
import subprocess
import sys
from datetime import datetime
from pathlib import Path
from typing import Any
from uuid import UUID

from fastapi import APIRouter, Body, Depends, HTTPException
from pydantic import BaseModel, Field

from acquisition.repositories.postgres import PostgresAcquisitionRepository
from app.dependencies import _env_file, get_creation_db_config
from app.routes.acquisition import _model_dump
from core.config import CreationDbConfig
from core.db_session import transaction
from pipeline.postgres import PostgresPipelineRepository
from pipeline.tracing import TraceContext, new_trace_writer
from query_planning import (
    GenerationRequest,
    GeneratorKind,
    QueryBatchWriteSpec,
    QueryBatchWriter,
    UnifiedQueryGenerationService,
    planning_store_for_repository,
)

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

ROOT = Path(__file__).resolve().parents[2]
RUNTIME_MANUAL_DIR = ROOT / "runtime" / "manual_query_runs"
DEFAULT_PLATFORMS = ("xiaohongshu", "weixin", "douyin")
QUERY_MAX_COUNT = 500
QUERY_MAX_CHARS = 500


class ManualQueryEntry(BaseModel):
    query_text: str
    axes: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)
    filter_reason: str | None = None


class ManualQueryBatchRequest(BaseModel):
    name: str | None = None
    queries: list[Any] | None = None
    families: list[dict[str, Any]] | None = None
    target_platforms: list[str] | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    search_limit: int = Field(default=10, ge=1, le=50)
    display_limit: int = Field(default=5, ge=1, le=50)
    decode_limit: int = Field(default=100, ge=0, le=500)


def _as_request(payload: Any) -> ManualQueryBatchRequest:
    if isinstance(payload, list):
        payload = {"queries": payload}
    if not isinstance(payload, dict):
        raise HTTPException(status_code=422, detail="请求体必须是对象或 query 数组")
    return ManualQueryBatchRequest.model_validate(payload)


def _query_text(value: Any) -> str:
    if value is None:
        return ""
    return str(value).strip()


def _entry_from_value(value: Any, *, family: dict[str, Any] | None = None) -> ManualQueryEntry | None:
    if isinstance(value, str):
        query_text = _query_text(value)
        axes: dict[str, Any] = {}
        metadata: dict[str, Any] = {}
        filter_reason = None
    elif isinstance(value, dict):
        query_text = _query_text(value.get("query_text") or value.get("query"))
        axes = value.get("axes") or value.get("parts") or {}
        if not isinstance(axes, dict):
            axes = {}
        metadata = value.get("metadata") or {}
        if not isinstance(metadata, dict):
            metadata = {}
        filter_reason = value.get("filter_reason") or value.get("reason")
    else:
        return None

    if not query_text:
        return None
    if len(query_text) > QUERY_MAX_CHARS:
        raise HTTPException(status_code=422, detail=f"query 过长，最多 {QUERY_MAX_CHARS} 字符")

    merged_metadata = dict(metadata)
    merged_metadata["family_key"] = "manual"
    if family:
        merged_metadata["source_family_key"] = family.get("key") or family.get("family_key")
        merged_metadata["source_family_name"] = family.get("name") or family.get("title")
    return ManualQueryEntry(
        query_text=query_text,
        axes=axes,
        metadata=merged_metadata,
        filter_reason=_query_text(filter_reason) or None,
    )


def _normalize_entries(request: ManualQueryBatchRequest) -> list[ManualQueryEntry]:
    entries: list[ManualQueryEntry] = []
    for value in request.queries or []:
        entry = _entry_from_value(value)
        if entry:
            entries.append(entry)

    for family in request.families or []:
        if not isinstance(family, dict):
            continue
        for value in family.get("items") or []:
            entry = _entry_from_value(value, family=family)
            if entry:
                entries.append(entry)

    if not entries:
        raise HTTPException(status_code=422, detail="没有可提交的 query")
    return entries


def _normalize_platforms(platforms: list[str] | None) -> list[str]:
    values = platforms or list(DEFAULT_PLATFORMS)
    if not values:
        raise HTTPException(status_code=422, detail="至少选择一个平台")
    seen: set[str] = set()
    out: list[str] = []
    for platform in values:
        key = str(platform).strip()
        if key not in DEFAULT_PLATFORMS:
            raise HTTPException(status_code=422, detail=f"不支持的平台：{key}")
        if key not in seen:
            seen.add(key)
            out.append(key)
    return out


def _pipeline_command(
    *,
    batch_id: str,
    run_key: str,
    platforms: list[str],
    search_limit: int,
    display_limit: int,
    decode_limit: int,
) -> list[str]:
    cmd = [
        sys.executable,
        str(ROOT / "scripts" / "run_creation_pipeline.py"),
        "--batch-id",
        batch_id,
        "--search-limit",
        str(search_limit),
        "--display-limit",
        str(display_limit),
        "--decode-limit",
        str(decode_limit),
        "--run-key",
        run_key,
        "--env-file",
        _env_file(),
    ]
    for platform in platforms:
        cmd.extend(["--platform", platform])
    return cmd


def _create_pipeline_trace(
    db_config: CreationDbConfig,
    *,
    batch_id: Any,
    run_key: str,
    platforms: list[str],
    request: ManualQueryBatchRequest,
    log_path: Path,
) -> str | None:
    try:
        with transaction(db_config) as conn:
            run = PostgresPipelineRepository(conn).create_pipeline_run(
                run_key=run_key,
                batch_id=batch_id,
                status="pending",
                current_stage="query",
                config={
                    "platforms": platforms,
                    "search_limit": request.search_limit,
                    "display_limit": request.display_limit,
                    "decode_limit": request.decode_limit,
                },
                metadata={
                    "source": "manual_api",
                    "log_path": str(log_path),
                },
            )
            return str(run.id) if run.id else None
    except Exception:
        return None


@router.post("/manual")
def create_manual_query_batch(
    payload: Any = Body(...),
    db_config: CreationDbConfig = Depends(get_creation_db_config),
) -> dict[str, Any]:
    request = _as_request(payload)
    entries = _normalize_entries(request)
    platforms = _normalize_platforms(request.target_platforms)
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    name = request.name or f"manual-query-{timestamp}"

    generation_result = UnifiedQueryGenerationService().generate(
        GenerationRequest(
            generator_kind=GeneratorKind.MANUAL,
            name=name,
            target_platforms=tuple(platforms),
            payload={
                "candidates": [
                    {
                        "query_text": entry.query_text,
                        "axes": entry.axes,
                        "filter_reason": entry.filter_reason,
                        "metadata": entry.metadata,
                        "source_refs": [
                            {
                                "generator_kind": GeneratorKind.MANUAL.value,
                                "original_position": index,
                                "axes": entry.axes,
                                "source_family_key": entry.metadata.get("source_family_key"),
                                "source_family_name": entry.metadata.get("source_family_name"),
                            }
                        ],
                    }
                    for index, entry in enumerate(entries)
                ],
                "generator_config": {
                    "request_shape": "manual_query_api_v1",
                },
                "input_snapshot": {
                    "submitted_count": len(entries),
                    "target_platforms": platforms,
                },
            },
            metadata={"source": "manual_api"},
        )
    )
    if generation_result.stats.selected_count > QUERY_MAX_COUNT:
        raise HTTPException(status_code=422, detail=f"一次最多提交 {QUERY_MAX_COUNT} 条 query")

    batch_metadata = {
        "family_key": "manual",
        "source": "manual_api",
        "query_count": generation_result.stats.selected_count,
        "platforms": platforms,
        **request.metadata,
    }

    with transaction(db_config) as conn:
        repo = PostgresAcquisitionRepository(conn)
        write_result = QueryBatchWriter(
            legacy_sink=repo,
            planning_store=planning_store_for_repository(repo),
        ).write(
            generation_result,
            QueryBatchWriteSpec(
                name=name,
                source_type="manual",
                generation_method="manual_query_api_v1",
                target_platforms=tuple(platforms),
                metadata=batch_metadata,
            ),
        )
        batch = write_result.batch
        queries = write_result.queries
        run_key = f"manual-api:{batch.id}:{timestamp}"
        log_path = RUNTIME_MANUAL_DIR / f"{run_key.replace(':', '-')}.log"
        cmd = _pipeline_command(
            batch_id=str(batch.id),
            run_key=run_key,
            platforms=platforms,
            search_limit=request.search_limit,
            display_limit=request.display_limit,
            decode_limit=request.decode_limit,
        )
        run_metadata = {
            "source": "manual_api",
            "batch_id": str(batch.id),
            "platforms": platforms,
            "search_limit": request.search_limit,
            "display_limit": request.display_limit,
            "decode_limit": request.decode_limit,
            "dry_ingest_record": True,
            "log_path": str(log_path),
            "command": cmd,
        }
        run = repo.create_acquisition_run(
            batch_id=batch.id,
            run_key=run_key,
            status="pending",
            note="manual query API queued",
            metadata=run_metadata,
        )

    pipeline_run_id = _create_pipeline_trace(
        db_config,
        batch_id=batch.id,
        run_key=run_key,
        platforms=platforms,
        request=request,
        log_path=log_path,
    )
    if pipeline_run_id:
        trace_writer = new_trace_writer(db_config, env_file=_env_file())
        trace_writer.event(
            context=TraceContext(
                pipeline_run_id=UUID(pipeline_run_id),
                acquisition_run_id=run.id,
                stage="query",
            ),
            stage="query",
            event_type="manual_query_queued",
            status="pending",
            payload={"query_count": len(queries), "platforms": platforms},
        )

    RUNTIME_MANUAL_DIR.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    env["PYTHONPATH"] = f"{ROOT}:{env.get('PYTHONPATH', '')}".rstrip(":")
    try:
        with log_path.open("ab") as stream:
            process = subprocess.Popen(
                cmd,
                cwd=str(ROOT),
                env=env,
                stdout=stream,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
        if pipeline_run_id:
            trace_writer = new_trace_writer(db_config, env_file=_env_file())
            trace_writer.event(
                context=TraceContext(
                    pipeline_run_id=UUID(pipeline_run_id),
                    acquisition_run_id=run.id,
                    stage="query",
                ),
                stage="query",
                event_type="process_started",
                status="running",
                payload={"pid": process.pid, "log_path": str(log_path)},
            )
    except OSError as exc:
        with transaction(db_config) as conn:
            PostgresAcquisitionRepository(conn).update_acquisition_run(
                run.id,
                status="failed",
                error_message=f"启动手动 query pipeline 失败：{exc}",
                metadata={"log_path": str(log_path), "command": cmd},
        )
        raise HTTPException(status_code=500, detail="启动后台 pipeline 失败") from exc

    return {
        "status": "queued",
        "batch_id": str(batch.id),
        "pipeline_run_id": pipeline_run_id,
        "run_id": str(run.id),
        "run_key": run_key,
        "pid": process.pid,
        "query_count": len(queries),
        "queries": [_model_dump(query) for query in queries],
        "summary_url": f"/api/acquisition/runs/{run.id}/summary",
        "log_path": str(log_path),
    }