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 sqlalchemy import text
from starlette.exceptions import HTTPException as StarletteHTTPException

from api.routers.pipeline import router as pipeline_router
from api.services.agent_catalog import (
    get_agent_detail,
    update_agent_document_injection,
)
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_agent_oss_logs, list_demand_belong_oss_logs
from api.services.scheduler import (
    get_scheduler_job_run,
    list_triggerable_jobs,
    run_scheduler_job,
    run_supply_pipeline,
    scheduler_status as get_persistent_scheduler_status,
)
from api.services.video_discovery import (
    get_video_discovery_demand,
    list_video_discovery_demands,
)
from supply_infra.config import get_infra_settings
from supply_infra.db import dispose_engine, get_session, init_db


class SPAStaticFiles(StaticFiles):
    """Serve index.html for browser routes handled by Vue Router."""

    async def get_response(self, path: str, scope):
        try:
            return await super().get_response(path, scope)
        except StarletteHTTPException as exc:
            is_backend_path = (
                path == "api"
                or path.startswith("api/")
                or path == "health"
                or path.startswith("health/")
            )
            if exc.status_code != 404 or is_backend_path or Path(path).suffix:
                raise
            return await super().get_response("index.html", scope)


@asynccontextmanager
async def lifespan(_app: FastAPI):
    if get_infra_settings().database_auto_create:
        init_db()
    yield
    dispose_engine()


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

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("/health/live")
def health_live() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/health/ready")
def health_ready() -> dict[str, str]:
    try:
        with get_session() as session:
            session.execute(text("SELECT 1"))
    except Exception as exc:
        raise HTTPException(status_code=503, detail="database unavailable") from exc
    return {"status": "ready"}


@app.get("/api/scheduler/status")
def scheduler_status() -> dict:
    """Deprecated compatibility view backed by MySQL pipeline state."""
    return get_persistent_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 表示后台执行",
    )


class AgentDocumentInjectionBody(BaseModel):
    content: str = Field(
        default="",
        max_length=100_000,
        description="追加到 Agent System Prompt 后的业务文档",
    )
    enabled: bool = Field(
        default=True,
        description="是否在新建 Agent 时启用该文档注入",
    )


@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
    except ValueError as exc:
        raise HTTPException(status_code=409, detail=str(exc)) from None


@app.get("/api/scheduler/runs/{run_id}")
def scheduler_job_run(run_id: str) -> dict:
    """Return durable pipeline run status."""
    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}


@app.get("/api/agent-oss-logs")
def agent_oss_logs(
    agent_name: str | None = Query(
        default=None,
        description="Agent 名称；省略则返回全部 Agent 日志",
    ),
) -> dict:
    """Return Agent oss_logs and the available Agent names."""
    return list_agent_oss_logs(agent_name=agent_name)


@app.get("/api/agents/{agent_name}")
def agent_detail(agent_name: str) -> dict:
    """Return one business Agent's current system prompt and tools."""
    result = get_agent_detail(agent_name)
    if result is None:
        raise HTTPException(status_code=404, detail=f"agent not found: {agent_name}")
    return result


@app.put("/api/agents/{agent_name}/document-injection")
def agent_document_injection(
    agent_name: str,
    body: AgentDocumentInjectionBody,
) -> dict:
    """Create or update the web-managed document injected into an Agent prompt."""
    result = update_agent_document_injection(
        agent_name,
        content=body.content,
        enabled=body.enabled,
    )
    if result is None:
        raise HTTPException(status_code=404, detail=f"agent not found: {agent_name}")
    return result


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