routes/unified_query.py

from __future__ import annotations

import asyncio
from datetime import datetime
from fastapi import HTTPException, Request
from services.unified_query import plan_unified_query

def _inject(context):
    globals().update(context)

async def unified_query_preview(request: Request):
    try:
        body=await request.json()
    except Exception as exc:
        raise HTTPException(400,"invalid json") from exc

    prompt=str((body or {}).get("prompt") or "").strip()
    if not prompt:
        raise HTTPException(400,"prompt is empty")

    await resolve_request_user(request, str((body or {}).get("user_id") or "").strip())
    plan=plan_unified_query(prompt)

    ocean_result={"requested":False,"sources":[],"errors":[]}
    fisheries_result={"requested":False,"matches":[],"errors":[]}

    if "ocean" in plan["domains"]:
        ocean_result["requested"]=True
        paths=("/catalog","/status/ocean","/domains")
        responses=await asyncio.gather(
            *(_marine_api_get(path) for path in paths),
            return_exceptions=True,
        )
        payloads={}
        for path,value in zip(paths,responses):
            if isinstance(value,Exception):
                ocean_result["errors"].append(f"{path}: {str(value)[:240]}")
            else:
                payloads[path]=value
        # Return the configured source registry plus live endpoint evidence;
        # this is a preview, not an export/download action.
        wanted=set(plan.get("variables") or [])
        for key,name,name_zh,variables in _OCEAN_CATALOG:
            score=len(wanted.intersection(set(variables or [])))
            if wanted and score==0:
                continue
            ocean_result["sources"].append({
                "key":key,
                "name":name,
                "name_zh":name_zh,
                "variables":list(variables or []),
                "match_score":score,
            })
        ocean_result["sources"].sort(key=lambda x:(-x["match_score"],x["name"]))
        ocean_result["live_endpoints"]={
            "catalog":"/catalog" in payloads,
            "status":"/status/ocean" in payloads,
            "domains":"/domains" in payloads,
        }

    if any(x in plan["domains"] for x in ("tuna","squid","fisheries")):
        fisheries_result["requested"]=True
        try:
            files,repo_errors=await hf_all_live_files(force=False)
            fisheries_result["errors"].extend(
                f"{repo}: {msg}" for repo,msg in (repo_errors or {}).items()
            )
            wanted_domains=set(x for x in plan["domains"] if x in ("tuna","squid"))
            candidates=[]
            for item in files:
                repo_domain=str(item.get("repository_domain") or "")
                if wanted_domains and repo_domain not in wanted_domains:
                    continue
                candidates.append(item)

            # Query terms improve ranking but do not hide every file if the
            # repository has usable data and the user query is broad.
            low=prompt.lower()
            tokens=[
                x for x in ("cpue","catch","effort","species","gear","tuna","squid","柔鱼","金枪鱼")
                if x in low
            ]
            def score(item):
                p=str(item.get("path_lower") or item.get("path") or "").lower()
                return sum(1 for t in tokens if t in p)

            candidates.sort(key=lambda x:(-score(x), int(x.get("size_bytes") or 0), str(x.get("path") or "")))
            fisheries_result["matches"]=[
                {
                    "repository":x.get("repository"),
                    "repository_domain":x.get("repository_domain"),
                    "path":x.get("path"),
                    "size_bytes":int(x.get("size_bytes") or 0),
                    "match_score":score(x),
                }
                for x in candidates[:40]
            ]
            fisheries_result["total_candidate_files"]=len(candidates)
        except Exception as exc:
            fisheries_result["errors"].append(str(exc)[:500])

    return {
        "checked_at":datetime.now().astimezone().isoformat(timespec="seconds"),
        "plan":plan,
        "ocean":ocean_result,
        "fisheries":fisheries_result,
        "note":"这是统一查询预览：只读取目录/目录元数据，不执行大文件下载或 Ocean 导出。",
    }

def register_unified_query_routes(app, context):
    _inject(context)
    app.add_api_route("/api/query/preview", unified_query_preview, methods=["POST"])