AgentLab/backend/app/main.py

from __future__ import annotations

import logging
from contextlib import asynccontextmanager
from typing import AsyncIterator

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from app.api import (
    artifacts, conversations, copilot, datasources, evidence, governance, knowledge, runs,
    sandbox, settings as settings_api, table_imports, tools, workflows,
)
from app.api.coded import CodedHTTPException, coded_handler
from app.core.config import settings
from app.db.base import init_db
from app.engine.runner import run_manager
from app.seed import seed_defaults

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)-7s %(name)s | %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("agentlab")


@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
    from app.data.engine import sqlite_dqs_self_check

    sqlite_dqs_self_check()  # 写错的双引号列名必须报错；做不到就记日志并让启动失败
    settings.ensure_dirs()
    await init_db()
    await seed_defaults()

    # 把存的 embedding 配置读进来。不读的话重启后会悄悄退回本地哈希向量，
    # 而库里的向量是上次那个模型建的——检索会整体退回关键词，
    # 用户只看到"重启之后搜得不准了"
    from app.db.base import SessionLocal
    from app.memory.embeddings import load_setting

    async with SessionLocal() as _session:
        try:
            await load_setting(_session)
        except Exception as e:  # noqa: BLE001 - key 没了不该拦住整个服务起不来
            logger.warning("embedding 配置没能生效，暂用本地向量：%s", e)

    # 上次被强杀时留下的"处理中"文档。后台任务不会自己回来，留着它只会在
    # 界面上永远转圈——和 runner 对 running 的 run 是同一个道理
    from sqlalchemy import update as _update

    from app.db.models import Document as _Doc

    async with SessionLocal() as _s:
        await _s.execute(
            _update(_Doc).where(_Doc.status == "processing")
            .values(status="failed", error="服务重启，这份文档没有处理完，请重新上传")
        )
        await _s.commit()

    # 上传表格的版本底座：清理上次没发布完的文件、给老上传补初始版本、回收没人用的版本。
    # 出错只记日志：最坏是留下几个多余的文件，不该让服务起不来
    from app.data import table_versions

    async with SessionLocal() as _s:
        try:
            await table_versions.startup(_s)
        except Exception:  # noqa: BLE001
            logger.exception("上传表格的版本整理失败，服务照常启动")

    await run_manager.setup()
    from app.sandbox.manager import sandbox_manager as _sbm

    await _sbm.start_reaper()
    logger.info("AgentLab 就绪 · 数据目录 %s", settings.data_dir)
    try:
        yield
    finally:
        await run_manager.shutdown()
        from app.sandbox.manager import sandbox_manager

        await sandbox_manager.close()
        logger.info("已停止")


app = FastAPI(
    title="AgentLab",
    description="可视化 Agent 编排实验台",
    version="0.1.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=settings.cors_origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

for router in (
    workflows.router,
    runs.router,
    runs.approvals_router,
    evidence.router,
    settings_api.router,
    tools.router,
    tools.custom_router,
    tools.mcp_router,
    knowledge.memory_router,
    knowledge.kb_router,
    knowledge.skills_router,
    sandbox.router,
    copilot.router,
    conversations.router,
    artifacts.router,
    governance.router,
    # 按配方导入的 /api/datasources/imports/… 要在 datasources 的通配 /{source_id} 之前登记
    table_imports.router,
    datasources.router,
):
    app.include_router(router)


app.add_exception_handler(CodedHTTPException, coded_handler)


@app.exception_handler(ValueError)
async def value_error_handler(_: Request, exc: ValueError) -> JSONResponse:
    return JSONResponse(status_code=400, content={"detail": str(exc)})


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


# 生产部署：前端构建产物由服务端一并托管（见 app/web.py）。页面回退是通配路由，
# 必须在所有 /api 路由之后注册。没设 AGENTLAB_WEB_DIST 时什么都不挂，开发模式照旧
if settings.web_dist is not None:
    from app.web import mount_web

    mount_web(app, settings.web_dist)