SupplyAgent/api/app.py

"""FastAPI application — category tree API on port 8080."""
from __future__ import annotations

from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

from api.services.category_tree import build_category_tree
from api.services.demand_belong_category import list_demand_belong_categories
from api.services.demand_grade import list_demand_grades
from api.services.demand_grade_videos import list_videos_for_demand_grade
from api.services.demand_videos import list_videos_for_demand_belong
from api.services.oss_logs import list_demand_belong_oss_logs
from api.services.scheduler import (
    get_scheduler_job_run,
    list_triggerable_jobs,
    run_scheduler_job,
    run_supply_pipeline,
)
from api.services.video_discovery import (
    get_video_discovery_demand,
    list_video_discovery_demands,
)
from supply_infra.db import init_db
from supply_infra.scheduler.app import get_scheduler_status, start_scheduler, stop_scheduler


@asynccontextmanager
async def lifespan(_app: FastAPI):
    init_db()
    start_scheduler()
    yield
    stop_scheduler()


app = FastAPI(title="SupplyAgent API", version="0.1.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:5173",
        "http://127.0.0.1:5173",
        "http://localhost:4173",
        "http://127.0.0.1:4173",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/api/scheduler/status")
def scheduler_status() -> dict:
    """Return scheduler enabled/running state and next run times."""
    return get_scheduler_status()


class RunPipelineBody(BaseModel):
    biz_dt: str | None = Field(
        default=None,
        pattern=r"^\d{8}$",
        description="业务日 YYYYMMDD；省略则取当天",
    )


class TriggerSchedulerJobBody(BaseModel):
    biz_dt: str | None = Field(
        default=None,
        pattern=r"^\d{8}$",
        description="业务日 YYYYMMDD；省略则取当天",
    )
    partition_date: str | None = Field(
        default=None,
        pattern=r"^\d{8}$",
        description="ODPS 分区日 YYYYMMDD；全局树默认取 biz_dt 前一日",
    )
    wait: bool = Field(
        default=False,
        description="是否同步等待执行完成；默认 false 表示后台执行",
    )


@app.post("/api/pipeline/run", status_code=202)
def run_pipeline(
    body: RunPipelineBody | None = None,
    biz_dt: str | None = Query(
        default=None,
        pattern=r"^\d{8}$",
        description="业务日 YYYYMMDD（与 body 二选一，body 优先）",
    ),
) -> dict:
    """
    异步一键执行供给数据全流程，立即返回 run_id，后台串行执行：

    全局树同步 → 需求池同步 → 需求分级 → 视频点位拓展 → find_agent 找片 → AIGC 发布。
    """
    resolved_biz_dt = body.biz_dt if body and body.biz_dt is not None else biz_dt
    return run_supply_pipeline(biz_dt=resolved_biz_dt)


@app.get("/api/scheduler/jobs")
def scheduler_jobs() -> dict:
    """Return jobs that can be triggered manually."""
    return {"items": list_triggerable_jobs()}


@app.post("/api/scheduler/jobs/{job_id}/run")
def trigger_scheduler_job(
    job_id: str,
    body: TriggerSchedulerJobBody | None = None,
    biz_dt: str | None = Query(
        default=None,
        pattern=r"^\d{8}$",
        description="业务日 YYYYMMDD（与 body 二选一，body 优先）",
    ),
    partition_date: str | None = Query(
        default=None,
        pattern=r"^\d{8}$",
        description="ODPS 分区日 YYYYMMDD（与 body 二选一，body 优先）",
    ),
    wait: bool = Query(
        default=False,
        description="是否同步等待执行完成",
    ),
) -> dict:
    """Manually trigger a scheduler job."""
    resolved_biz_dt = body.biz_dt if body and body.biz_dt is not None else biz_dt
    resolved_partition_date = (
        body.partition_date if body and body.partition_date is not None else partition_date
    )
    resolved_wait = body.wait if body is not None else wait

    try:
        return run_scheduler_job(
            job_id,
            biz_dt=resolved_biz_dt,
            partition_date=resolved_partition_date,
            wait=resolved_wait,
        )
    except KeyError:
        raise HTTPException(status_code=404, detail=f"scheduler job not found: {job_id}") from None


@app.get("/api/scheduler/runs/{run_id}")
def scheduler_job_run(run_id: str) -> dict:
    """Return manual job run status (in-memory; cleared on process restart)."""
    result = get_scheduler_job_run(run_id)
    if result is None:
        raise HTTPException(status_code=404, detail="scheduler run not found")
    return result


@app.get("/api/category-tree")
def category_tree(
    biz_dt: str | None = Query(
        default=None,
        description="业务日 YYYYMMDD；省略则取 category_tree_weight 最新一日",
    ),
) -> dict:
    """Return nested global_tree_category with per-dim avg score."""
    return build_category_tree(biz_dt=biz_dt)


@app.get("/api/demand-belong-category")
def demand_belong_category() -> dict:
    """Return all active demand_belong_category rows in one response."""
    items = list_demand_belong_categories()
    return {"items": items}


@app.get("/api/demand-belong-category/{belong_id}/videos")
def demand_belong_videos(belong_id: int) -> dict:
    """Return videos linked to a demand_belong_category row (vid + title + points JSON)."""
    result = list_videos_for_demand_belong(belong_id)
    if result is None:
        raise HTTPException(status_code=404, detail="demand_belong_category not found")
    return result


@app.get("/api/demand-grade")
def demand_grade(
    biz_dt: str | None = Query(
        default=None,
        description="业务日 YYYYMMDD；省略则取 demand_grade 最新一日",
    ),
) -> dict:
    """Return demand_grade rows (one per category_id) for the given/latest biz_dt."""
    items = list_demand_grades(biz_dt=biz_dt)
    return {"items": items}


@app.get("/api/demand-grade/{demand_grade_id}/videos")
def demand_grade_videos(demand_grade_id: int) -> dict:
    """Return expansion videos/points for a demand_grade row (by its biz_dt)."""
    result = list_videos_for_demand_grade(demand_grade_id)
    if result is None:
        raise HTTPException(status_code=404, detail="demand_grade not found")
    return result


@app.get("/api/video-discovery/demands")
def video_discovery_demands(
    biz_dt: str | None = Query(
        default=None,
        description="业务日 YYYYMMDD；省略则取 demand_grade 最新一日",
    ),
) -> dict:
    """Return demand-first video discovery cards for the selected/latest day."""
    return list_video_discovery_demands(biz_dt=biz_dt)


@app.get("/api/video-discovery/demands/{demand_grade_id}")
def video_discovery_demand(demand_grade_id: int) -> dict:
    """Return one demand and its videos/points resolved from the source tables."""
    result = get_video_discovery_demand(demand_grade_id)
    if result is None:
        raise HTTPException(status_code=404, detail="demand_grade not found")
    return result


@app.get("/api/demand-belong-oss-logs")
def demand_belong_oss_logs() -> dict:
    """Return demand_belong_category_agent oss_logs ordered by create_time desc."""
    items = list_demand_belong_oss_logs()
    return {"items": items}


_web_dist = Path(__file__).resolve().parent.parent / "web" / "dist"
if _web_dist.is_dir():
    app.mount("/", StaticFiles(directory=_web_dist, html=True), name="web")