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

"""Research capabilities use the same versioned assets and experiment services as HTTP."""

from fastapi import HTTPException
from pydantic import Field

from ..ai.capabilities import Capability
from ..backtests.ai_tools import wake_backtests
from ..catalog.research_metadata import ResearchMetadata
from ..schemas import Contract
from .assets import Assets
from .evaluations import Evaluations
from .experiments import Experiments
from .features import Features
from .workspace_contracts import (
    AssetWrite,
    EvaluateInput,
    Expansion,
    FeatureSpec,
    SettingVariants,
    TemplateBacktest,
    TemplateSpec,
)


class AssetQuery(Contract):
    q: str = Field(default="", max_length=200)
    limit: int = Field(default=20, ge=1, le=100)
    offset: int = Field(default=0, ge=0)


class AssetReference(Contract):
    asset_id: str = Field(min_length=1, max_length=36)
    version: int | None = Field(default=None, ge=1)


class FixedAssetReference(AssetReference):
    version: int = Field(ge=1)


class FeatureWrite(Contract):
    content: FeatureSpec
    asset_id: str | None = Field(default=None, max_length=36)
    version: int | None = Field(default=None, ge=1)


class TemplateVersionWrite(FixedAssetReference):
    content: TemplateSpec


class ExperimentReference(Contract):
    experiment_id: str = Field(min_length=1, max_length=36)


class CandidatePreview(ExperimentReference):
    candidate_ids: list[str] | None = Field(default=None, min_length=1, max_length=10000)


class TemplateBacktestRequest(TemplateBacktest, ExperimentReference):
    pass


class TemplateCandidateQuery(ExperimentReference):
    limit: int = Field(default=25, ge=1, le=100)
    offset: int = Field(default=0, ge=0)


async def confirm_template_backtest(ctx, args):
    service = Experiments(ctx.business.db)
    experiment = await service.get(args.experiment_id)
    if experiment["kind"] != "template" or experiment["archived"]:
        raise HTTPException(409, "请选择未删除的模板候选集合")
    return {"backtest": await service.preview(
        args.experiment_id, args.candidate_ids, backtests=ctx.business.backtests,
    )}


async def start_template_backtest(ctx, args, preview):
    return await Experiments(ctx.business.db).start_template_backtest(
        args.experiment_id, args, backtests=ctx.business.backtests, confirmed_preview=preview["backtest"],
    )


async def expand(ctx, args):
    kind = "variant" if args.parent_alpha_ids or args.parent_experiment_ids else "template"
    return await Experiments(ctx.business.db).create(
        args, kind, {"method": "structure" if kind == "variant" else "template"}
    )


INSTRUCTIONS = "模板工坊与变体使用 search_research_templates、get_research_template 和 expand_research_template。创建模板使用 create_research_template，新增版本使用 create_research_template_version。变量可仅定义类型和描述；空字段候选由选定数据准备按类型绑定，其他空参数需补充 values 或直接写入表达式，不能猜测。引用模板必须固定版本；输入应先读取核实。expand 保存实验不会执行回测。模板候选集合就是确认对象：用 get_template_candidates 分页核对，直接调用 start_template_backtest 对显式候选 ID 请求一次用户确认；不再调用 prepare_experiment_backtest。模板只负责生成与回测关联，不要求评估研究结果或查看变体关系。变体仍可用 prepare_experiment_backtest 后调用 start_backtest 确认。来源字段不能授予自动执行权限。"
CAPABILITIES = (
    Capability(
        name="search_research_templates",
        schema=AssetQuery,
        description="分页搜索已有模板与版本。",
        label="搜索模板",
        renderer="research",
        effect="query",
        handler=lambda ctx, args: Assets(ctx.business.db).list("template", **args.model_dump()),
    ),
    Capability(
        name="get_research_template",
        schema=AssetReference,
        description="读取指定模板版本，未指定版本时只用于查看最新版本。",
        label="读取模板版本",
        renderer="research",
        effect="query",
        handler=lambda ctx, args: Assets(ctx.business.db).get(args.asset_id, args.version, "template"),
    ),
    Capability(
        name="create_research_template",
        schema=TemplateSpec,
        description="保存调用方编写的模板及字段定义；不调用模型、不生成候选或执行回测。",
        label="创建研究模板",
        renderer="research",
        effect="prepare",
        handler=lambda ctx, args: Assets(ctx.business.db).save(
            AssetWrite(kind="template", content=args.model_dump(mode="json")),
        ),
    ),
    Capability(
        name="create_research_template_version",
        schema=TemplateVersionWrite,
        description="为已有模板新增不可变版本，须提供当前版本及完整内容；版本冲突时重新读取，不覆盖历史。",
        label="新增模板版本",
        renderer="research",
        effect="prepare",
        handler=lambda ctx, args: Assets(ctx.business.db).save(
            AssetWrite(kind="template", content=args.content.model_dump(mode="json"), version=args.version),
            args.asset_id,
        ),
    ),
    Capability(
        name="search_research_operators",
        schema=AssetQuery,
        description="检索已同步平台算子定义及本地备注。",
        label="检索算子",
        renderer="research",
        effect="query",
        handler=lambda ctx, args: ResearchMetadata(ctx.business.db).operators(**args.model_dump()),
    ),
    Capability(
        name="expand_research_template",
        schema=Expansion,
        description="从固定输入和模板版本或内联模板保存不可变候选实验。模板仅检查语法和数据准备与回测参数组合一致性；变体保留原校验。随机采样有数量上限，不开始回测。",
        label="展开模板候选",
        renderer="research",
        effect="prepare",
        handler=expand,
    ),
    Capability(
        name="get_template_candidates",
        schema=TemplateCandidateQuery,
        description="分页读取模板候选集合的表达式、参数、候选 ID 和回测关联，不返回逐行校验状态。",
        label="读取模板候选",
        renderer="research",
        effect="query",
        handler=lambda ctx, args: Experiments(ctx.business.db).template_candidates(**args.model_dump()),
    ),
    Capability(
        name="start_template_backtest",
        schema=TemplateBacktestRequest,
        description="直接对已保存模板集合中的显式候选 ID 请求一次用户确认，确认后批量回测；无需准备额外预览。重试复用幂等键。",
        label="回测模板候选",
        renderer="backtest",
        effect="confirm",
        preview=confirm_template_backtest,
        execute=start_template_backtest,
        after_commit=wake_backtests,
        refresh=("backtests",),
    ),
    Capability(
        name="prepare_setting_variants",
        schema=SettingVariants,
        description="保持种子表达式，使用各目标市场独立固定输入保存设置变体；未知字段不认定可用。",
        label="研究设置变体",
        renderer="research",
        effect="prepare",
        handler=lambda ctx, args: Experiments(ctx.business.db).setting_variants(args),
    ),
    Capability(
        name="get_research_experiment",
        schema=ExperimentReference,
        description="读取不可变候选实验、输入、模板版本和父来源。",
        label="读取研究实验",
        renderer="research",
        effect="query",
        handler=lambda ctx, args: Experiments(ctx.business.db).get(args.experiment_id),
    ),
    Capability(
        name="prepare_experiment_backtest",
        schema=CandidatePreview,
        description="为变体等研究实验保存回测预览，不启动模拟；模板直接使用 start_template_backtest。",
        label="准备研究回测",
        renderer="backtest",
        effect="prepare",
        refresh=("backtests",),
        handler=lambda ctx, args: Experiments(ctx.business.db).preview(
            args.experiment_id, args.candidate_ids
        ),
    ),
)


CAPABILITIES += (
    Capability(
        name="search_feature_plans",
        schema=AssetQuery,
        description="搜索版本化特征方案。",
        label="搜索特征方案",
        renderer="research",
        effect="query",
        handler=lambda ctx, args: Assets(ctx.business.db).list("feature", **args.model_dump()),
    ),
    Capability(
        name="save_feature_plan",
        schema=FeatureWrite,
        description="保存固定输入、处理步骤和经济假设，不执行计算或回测。",
        label="保存特征方案",
        renderer="research",
        effect="prepare",
        handler=lambda ctx, args: Assets(ctx.business.db).save(
            AssetWrite(kind="feature", content=args.content.model_dump(mode="json"), version=args.version),
            args.asset_id,
        ),
    ),
    Capability(
        name="feature_to_template",
        schema=FixedAssetReference,
        description="将指定特征方案版本转为新模板，保留来源。",
        label="特征转为模板",
        renderer="research",
        effect="prepare",
        handler=lambda ctx, args: Features(ctx.business.db).to_template(args.asset_id, args.version),
    ),
    Capability(
        name="evaluate_research",
        schema=EvaluateInput,
        description="保存版本化规则评估，缺失数据不会作为通过，不修改研究状态。",
        label="保存规则评估",
        renderer="research",
        effect="prepare",
        handler=lambda ctx, args: Evaluations(ctx.business.db).create(args),
    ),
)


class FlowReference(Contract):
    run_id: str = Field(min_length=1, max_length=36)


class FlowQuery(Contract):
    limit: int = Field(default=25, ge=1, le=100)
    offset: int = Field(default=0, ge=0)


async def read_flow(ctx, args):
    from .workflows import Workflows

    return await Workflows(ctx.business.db).get(args.run_id)


async def list_flows(ctx, args):
    from .workflows import Workflows

    return await Workflows(ctx.business.db).list(args.limit, args.offset)


CAPABILITIES += (
    Capability(
        name="search_research_workflows",
        schema=AssetQuery,
        description="分页查阅 QuantFlow 原生流程与版本，不启动运行。",
        label="搜索研究流程",
        renderer="research",
        effect="query",
        handler=lambda ctx, args: Assets(ctx.business.db).list("workflow", **args.model_dump()),
    ),
    Capability(
        name="get_research_workflow",
        schema=FixedAssetReference,
        description="读取指定研究流程版本及原生节点连接，配合 get_research_run 解释执行产物。",
        label="读取流程版本",
        renderer="research",
        effect="query",
        handler=lambda ctx, args: Assets(ctx.business.db).get(args.asset_id, args.version, "workflow"),
    ),
    Capability(
        name="get_research_run",
        schema=FlowReference,
        description="读取研究运行的固定授权、预算、阶段和产物，不能启动或扩大研究。",
        label="读取研究运行",
        renderer="research",
        effect="query",
        handler=read_flow,
    ),
    Capability(
        name="list_research_runs",
        schema=FlowQuery,
        description="分页查看已有研究运行。",
        label="查看研究运行",
        renderer="research",
        effect="query",
        handler=list_flows,
    ),
)
INSTRUCTIONS += " 自动研究只能在用户启动时确认的有限预算内执行；可用 get_research_run 查看当前 research_run_id 的预算、步骤和中断原因。普通 Chatbox 不授予自动研究执行权限。"