worldquant-alpha-system/backend/app/research/routes.py

"""Authenticated research workspace; previewing never starts a platform simulation."""

from fastapi import APIRouter, Depends, HTTPException, Query, Request

from ..backtests.contracts import RunOutput
from ..security import require_auth
from .assets import Assets
from .comparisons import compare
from .evaluations import Evaluations
from .experiments import Experiments
from .features import Features
from .model import OUTPUTS, Advice, request_model
from .workspace_contracts import (
    AssetWrite,
    CompareInput,
    EvaluateInput,
    Expansion,
    ExperimentPreview,
    FeatureConversion,
    FlowControl,
    FlowStart,
    Generation,
    ImportCommit,
    ImportPreview,
    SettingVariants,
    TemplateBacktest,
    WorkflowSpec,
)

router = APIRouter(prefix="/api/v1/research", tags=["research"], dependencies=[Depends(require_auth)])


@router.get("/assets")
async def assets(
    request: Request,
    kind: str = "template",
    q: str = "",
    limit: int = Query(25, ge=1, le=100),
    offset: int = Query(0, ge=0),
):
    if kind not in ("template", "feature", "view", "workflow"):
        raise HTTPException(422, "当前素材类型尚未开放")
    async with request.app.state.sessions() as db:
        return await Assets(db).list(kind, q, limit, offset)


@router.post("/assets", status_code=201)
async def save_asset(body: AssetWrite, request: Request):
    if body.kind not in ("template", "feature", "view", "workflow"):
        raise HTTPException(422, "当前素材类型尚未开放")
    async with request.app.state.sessions.begin() as db:
        return await Assets(db).save(body)


@router.get("/assets/{asset_id}")
async def asset(asset_id: str, request: Request, version: int | None = Query(None, ge=1)):
    async with request.app.state.sessions() as db:
        return await Assets(db).get(asset_id, version)


@router.put("/assets/{asset_id}")
async def update_asset(asset_id: str, body: AssetWrite, request: Request):
    if body.kind not in ("template", "feature", "view", "workflow"):
        raise HTTPException(422, "当前素材类型尚未开放")
    async with request.app.state.sessions.begin() as db:
        return await Assets(db).save(body, asset_id)


@router.delete("/assets/{asset_id}")
async def archive_asset(asset_id: str, request: Request, version: int = Query(ge=1)):
    async with request.app.state.sessions.begin() as db:
        return await Assets(db).archive(asset_id, version)


@router.get("/assets/{asset_id}/versions")
async def versions(asset_id: str, request: Request):
    async with request.app.state.sessions() as db:
        return await Assets(db).versions(asset_id)


@router.post("/templates/import-preview")
async def import_preview(body: ImportPreview, request: Request):
    async with request.app.state.sessions() as db:
        return await Assets(db).import_preview(body.templates)


@router.post("/templates/import", status_code=201)
async def import_commit(body: ImportCommit, request: Request):
    async with request.app.state.sessions.begin() as db:
        return await Assets(db).import_commit(body)


@router.post("/generate", status_code=201)
async def generate(body: Generation, request: Request):
    async with request.app.state.sessions.begin() as db:
        context = await Experiments(db).generation_context(body)
    result, evidence = await request_model(request.app.state.ai, context, OUTPUTS[body.method])
    if body.method == "feature" and (result.preparation_refs or set(result.input_ids) != set(body.input_ids)):
        raise HTTPException(422, "模型不能改变已固定的输入范围")
    async with request.app.state.sessions.begin() as db:
        asset = await Assets(db).save(
            AssetWrite(
                kind="feature" if body.method == "feature" else "template",
                content=result.model_dump(mode="json"),
            ),
            provenance={"generation": evidence, "context": context},
        )
    return {
        **asset,
        "generation": evidence,
        "parent_alpha_ids": body.parent_alpha_ids,
        "parent_experiment_ids": body.parent_experiment_ids,
    }


@router.post("/experiments", status_code=201)
async def expand(body: Expansion, request: Request):
    async with request.app.state.sessions.begin() as db:
        kind = "variant" if body.parent_alpha_ids or body.parent_experiment_ids else "template"
        return await Experiments(db).create(
            body, kind, {"method": "structure" if kind == "variant" else "template"}
        )


@router.get("/experiments")
async def experiments(
    request: Request,
    kind: str | None = None,
    limit: int = Query(25, ge=1, le=100),
    offset: int = Query(0, ge=0),
):
    async with request.app.state.sessions() as db:
        return await Experiments(db).list(kind, limit, offset)


@router.get("/experiments/{experiment_id}")
async def experiment(experiment_id: str, request: Request):
    async with request.app.state.sessions() as db:
        return await Experiments(db).get(experiment_id)


@router.delete("/experiments/{experiment_id}")
async def archive_experiment(experiment_id: str, request: Request):
    async with request.app.state.sessions.begin() as db:
        return await Experiments(db).archive(experiment_id)


@router.post("/experiments/{experiment_id}/preview", status_code=201)
async def preview(experiment_id: str, body: ExperimentPreview, request: Request):
    async with request.app.state.sessions.begin() as db:
        return await Experiments(db).preview(experiment_id, body.candidate_ids)


@router.post("/experiments/{experiment_id}/backtest", status_code=202, response_model=RunOutput)
async def template_backtest(experiment_id: str, body: TemplateBacktest, request: Request):
    async with request.app.state.sessions.begin() as db:
        result = await Experiments(db).start_template_backtest(experiment_id, body)
    request.app.state.runner.backtests.wake.set()
    return result


@router.post("/variants/settings", status_code=201)
async def settings_variants(body: SettingVariants, request: Request):
    async with request.app.state.sessions.begin() as db:
        return await Experiments(db).setting_variants(body)


@router.post("/compare")
async def comparison(body: CompareInput, request: Request):
    async with request.app.state.sessions() as db:
        return await compare(db, body.alpha_ids)


@router.post("/features/{asset_id}/template", status_code=201)
async def feature_template(asset_id: str, body: FeatureConversion, request: Request):
    async with request.app.state.sessions.begin() as db:
        return await Features(db).to_template(asset_id, body.version)


@router.post("/evaluations", status_code=201)
async def evaluate(body: EvaluateInput, request: Request):
    async with request.app.state.sessions.begin() as db:
        return await Evaluations(db).create(body)


@router.get("/evaluations")
async def evaluations(
    request: Request,
    alpha_id: str | None = None,
    experiment_id: str | None = None,
    limit: int = Query(25, ge=1, le=100),
    offset: int = Query(0, ge=0),
):
    async with request.app.state.sessions() as db:
        return await Evaluations(db).list(alpha_id, experiment_id, limit, offset)


@router.get("/evaluations/{evaluation_id}")
async def evaluation(evaluation_id: str, request: Request):
    async with request.app.state.sessions() as db:
        return await Evaluations(db).get(evaluation_id)


@router.post("/evaluations/{evaluation_id}/advice", status_code=201)
async def evaluation_advice(evaluation_id: str, request: Request):
    async with request.app.state.sessions() as db:
        original = await Evaluations(db).get(evaluation_id)
    report = original["report"]
    context = {
        "evaluation_id": evaluation_id,
        "rules": report["rules"],
        "verdict": report["verdict"],
        "total_records": len(report["records"]),
        "records": [
            {
                key: row.get(key)
                for key in ("alpha_id", "client_item_id", "evidence", "missing", "failed", "verdict")
            }
            for row in report["records"][:20]
        ],
        "record_limit": 20,
    }
    result, evidence = await request_model(request.app.state.ai, context, Advice)
    evidence["context"] = context
    async with request.app.state.sessions.begin() as db:
        return await Evaluations(db).add_advice(evaluation_id, result.model_dump(mode="json"), evidence)


@router.get("/lineage")
async def research_lineage(
    request: Request,
    alpha_id: str | None = None,
    experiment_id: str | None = None,
    limit: int = Query(25, ge=1, le=100),
    offset: int = Query(0, ge=0),
):
    from .lineage import lineage

    async with request.app.state.sessions() as db:
        return await lineage(db, alpha_id, experiment_id, limit, offset)


@router.get("/flows/recipe")
async def fixed_recipe():
    from .workflows import fixed_workflow

    return fixed_workflow().model_dump(mode="json")


@router.post("/flows/runs", status_code=201)
async def start_flow(body: FlowStart, request: Request):
    from .workflows import Workflows

    async with request.app.state.sessions.begin() as db:
        from .workflows import fixed_workflow
        from .workspace_contracts import WorkflowSpec

        graph = (
            WorkflowSpec.model_validate(
                (await Assets(db).get(body.workflow_id, body.workflow_version, "workflow"))["content"]
            )
            if body.workflow_id
            else fixed_workflow(body.budget.max_rounds)
        )
        needs_model = any(
            n.type == "generate"
            or (n.type == "feature" and not n.config.get("asset_id"))
            or (n.type == "variant" and n.config.get("method", "structure") == "structure")
            for n in graph.nodes
        )
        config = await request.app.state.ai.config(db) if needs_model else None
        result = await Workflows(db).start(body, config.revision if config else None)
    request.app.state.research.wake.set()
    return result


@router.get("/flows/runs")
async def flow_runs(
    request: Request,
    limit: int = Query(25, ge=1, le=100),
    offset: int = Query(0, ge=0),
    kind: str | None = None,
):
    from .workflows import Workflows

    async with request.app.state.sessions() as db:
        if kind not in (None, "pipeline", "quantflow"):
            raise HTTPException(422, "未知研究运行类型")
        return await Workflows(db).list(limit, offset, kind)


@router.get("/flows/runs/{run_id}")
async def flow_run(run_id: str, request: Request):
    from .workflows import Workflows

    async with request.app.state.sessions() as db:
        return await Workflows(db).get(run_id)


@router.post("/flows/runs/{run_id}/control")
async def control_flow(run_id: str, body: FlowControl, request: Request):
    from .runtime import control
    from .workflows import Workflows

    async with request.app.state.sessions.begin() as db:
        await control(db, run_id, body)
        result = await Workflows(db).get(run_id)
    request.app.state.research.wake.set()
    request.app.state.runner.backtests.wake.set()
    return result


@router.get("/flows/nodes")
async def flow_nodes():
    from .workflows import NODE_TYPES

    return {"items": [{"type": key, **value} for key, value in NODE_TYPES.items()]}


@router.post("/flows/validate")
async def validate_flow(body: WorkflowSpec):
    from .workflows import validate_graph

    return {"valid": True, "order": validate_graph(body)}