"""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), }