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

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

from pydantic import Field

from ..ai.capabilities import Capability
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


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 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)


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。引用模板必须固定版本；输入应先读取核实。expand 保存实验不会执行回测。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="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="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="从实验内已校验的固定候选保存回测确认预览，不启动模拟。",
        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="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 不授予自动研究执行权限。"