worldquant-alpha-system/backend/app/ai/tools.py

"""Explicit business tool catalog. This module has no database or provider credentials."""

from datetime import datetime
from typing import Literal

from pydantic import Field

from ..backtests.contracts import ControlInput, PreviewInput, RerunInput, StartInput
from ..catalog.contracts import UNIVERSES, CatalogFilters, Scope
from ..research.contracts import (
    ChatboxResearchInput,
    InputPageArgs,
    ResearchInputSelection,
    ResearchPreviewInput,
)
from ..schemas import AlphaFilters, BulkInput, BulkUpdate, Contract, JobInput, ResearchInput, ResearchUpdate


class EmptyArgs(Contract):
    pass


class SearchArgs(Contract):
    filters: AlphaFilters = Field(default_factory=AlphaFilters)


class AlphaArgs(Contract):
    alpha_id: str = Field(min_length=1, max_length=100, pattern=r"^[A-Za-z0-9_-]+$")


class JobArgs(Contract):
    job_id: str = Field(min_length=1, max_length=100)


class ResearchArgs(AlphaArgs):
    changes: ResearchInput


class ResultMetadata(Contract):
    source: Literal["local_database"] = "local_database"
    observed_at: datetime
    nulls: str = "null 表示来源未提供，不等于零"
    units: dict[str, str] = Field(
        default_factory=lambda: {
            "turnover": "比例，0.15 = 15%",
            "returns": "比例",
            "drawdown": "比例",
            "margin": "比例",
            "pnl": "供应商原始累计值，未提供货币/规模单位",
        }
    )


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


class BacktestListArgs(Contract):
    limit: int = Field(default=20, ge=1, le=100)
    offset: int = Field(default=0, ge=0)
    source: str | None = Field(default=None, max_length=100)
    reference: str | None = Field(default=None, max_length=200)
    research_id: str | None = Field(default=None, max_length=200)


class CatalogSearchArgs(Contract):
    filters: CatalogFilters
    dataset_id: str | None = Field(default=None, min_length=1, max_length=200)


class CatalogDetailArgs(Contract):
    scope: Scope
    dataset_id: str = Field(min_length=1, max_length=200)
    field_id: str = Field(default="", max_length=200)


class BacktestDraftArgs(Contract):
    draft_id: str = Field(min_length=1, max_length=36)
    limit: int = Field(default=25, ge=1, le=100)
    offset: int = Field(default=0, ge=0)


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


class BacktestResultsArgs(BacktestRunArgs):
    limit: int = Field(default=20, ge=1, le=100)
    offset: int = Field(default=0, ge=0)


class BacktestPreviewArgs(Contract):
    preview_id: str = Field(min_length=1, max_length=36)
    limit: int = Field(default=20, ge=1, le=100)
    offset: int = Field(default=0, ge=0)


class BacktestControlArgs(BacktestRunArgs):
    action: Literal["pause", "resume", "stop", "recover"]


class BacktestRerunArgs(BacktestRunArgs):
    item_ids: list[str] = Field(min_length=1, max_length=100)


CATALOG = {
    "get_catalog_scopes": (EmptyArgs, "读取本版支持的研究范围组合，不表示账户已获平台权限。"),
    "search_catalog": (
        CatalogSearchArgs,
        "分页查询本地目录。省略 dataset_id 查询数据集；提供 dataset_id 查询其字段、类型和完整集合版本。无缓存时说明需在数据集页同步，不编造字段。",
    ),
    "get_catalog_detail": (CatalogDetailArgs, "读取指定范围的数据集或字段详情；field_id 为空表示数据集。"),
    "prepare_research_input": (
        ResearchInputSelection,
        "把明确选出的 1–100 个字段固定为研究输入快照；必须提供已读取的集合版本。只保存本地输入，不启动同步或回测。已有输入引用时直接读取，不重建。",
    ),
    "get_research_input": (
        InputPageArgs,
        "分页读取不可变研究输入及原版本字段说明。q 仅搜索字段 ID，类型筛选不改变输入。field_count 是输入总数，total 是筛选匹配数。",
    ),
    "prepare_research_backtest": (
        ChatboxResearchInput,
        "从研究输入快照构建固定回测预览：提供假设、候选模板（如 rank({price})）、具名字段绑定及真实类型、明确模拟参数。服务端校验归属/类型/范围并替换占位符；不验证算子语义、不自动聚合 VECTOR。来源由服务端标记为当前 chatbox 研究。",
    ),
    "get_backtest_draft": (
        BacktestDraftArgs,
        "分页读取已有候选草稿、版本和来源，随后按 draft_id/draft_version 准备预览。",
    ),
    "get_alpha_sources": (
        AlphaSourcesArgs,
        "分页读取 Alpha 的全部已保存研究来源、关联回测、聊天会话和输入快照。没有记录不推断来源。",
    ),
    "get_backtest_capabilities": (
        EmptyArgs,
        "读取回测输入 schema、支持类型和本地调度配置，不代表平台剩余额度。",
    ),
    "prepare_backtest": (
        PreviewInput,
        "准备服务端固定回测预览，可用 inline 候选或草稿引用；只保存预览，不提交平台，不需要执行确认。",
    ),
    "get_backtest_preview": (BacktestPreviewArgs, "分页读取完整固定预览，确认前核对表达式和最终参数。"),
    "start_backtest": (
        StartInput,
        "对已保存预览请求一次用户确认，确认后后台运行全部固定候选，立即返回运行 ID；禁止循环等待。",
    ),
    "list_backtests": (BacktestListArgs, "分页查询回测运行与统计，可按来源筛选。"),
    "get_backtest": (BacktestRunArgs, "查询指定运行的真实进度，不循环等待完成。"),
    "get_backtest_results": (
        BacktestResultsArgs,
        "分页读取逐项状态、历史指标和错误；未知结果不能推测为成功。",
    ),
    "control_backtest": (
        BacktestControlArgs,
        "预览并确认暂停/继续/停止剩余项/找回原任务；不远端取消，不重新提交。",
    ),
    "prepare_backtest_rerun": (
        BacktestRerunArgs,
        "从明确指定的已结束回测项准备新预览，保留来源；不会自动启动。",
    ),
    "search_alphas": (
        SearchArgs,
        "按明确筛选条件查询本地 Alpha。Turnover 0.15 表示 15%；支持分页，禁止把当前页当作全部结果。",
    ),
    "get_alpha_facets": (EmptyArgs, "获取可用地区、类型、状态、标签与本地 Alpha 总数。"),
    "get_alpha": (AlphaArgs, "读取指定 Alpha 的表达式、指标、已有检查和本地研究记录。"),
    "get_alpha_pnl": (AlphaArgs, "读取指定 Alpha 的 PnL 缓存摘要。缺失缓存时说明情况，不自动刷新。"),
    "list_jobs": (EmptyArgs, "查询最近的同步任务，不要循环轮询等待。"),
    "get_job_status": (JobArgs, "查询指定任务的状态、目标和错误，不要循环等待任务完成。"),
    "update_research": (
        ResearchArgs,
        "提出指定 Alpha 的本地研究记录修改。仅传要修改的字段，等待用户在界面确认。",
    ),
    "bulk_update_research": (BulkInput, "提出固定 1–100 个 Alpha 的批量标签或研究状态修改，等待用户确认。"),
    "create_sync_job": (
        JobInput,
        "提出同步或本地自相关任务，等待确认。full_sync 仅同步已提交；待提交必须用 daily_sync 并指定 submission、date_from/date_to（UTC），待提交按创建日、已提交按提交日逐天同步。alpha_refresh/pnl_refresh/self_correlation 使用固定 alpha_ids；自相关缺失 PnL 时自动补取，不触发平台检查。创建后立即返回任务 ID。",
    ),
    "cancel_job": (JobArgs, "提出取消指定同步任务，等待用户确认。"),
    "retry_job": (JobArgs, "提出重试指定失败或暂停的同步任务，等待用户确认。"),
}
WRITES = {
    "update_research",
    "bulk_update_research",
    "create_sync_job",
    "cancel_job",
    "retry_job",
    "start_backtest",
    "control_backtest",
}


def bounded(value):
    if isinstance(value, str):
        return value if len(value) <= 2000 else value[:2000] + "…（已截断）"
    if isinstance(value, list):
        return [bounded(v) for v in value[:100]]
    if isinstance(value, dict):
        return {k: bounded(v) for k, v in list(value.items())[:100]}
    return value


async def read_tool(business, name, args):
    from datetime import timezone

    if name == "get_catalog_scopes":
        data = {"universes": UNIVERSES, "instrument_type": "EQUITY", "delays": [0, 1]}
    elif name == "search_catalog":
        data = await business.catalog.search(args.filters, args.dataset_id)
        data.update(
            scope=args.filters.model_dump(include=set(Scope.model_fields)), dataset_id=args.dataset_id
        )
    elif name == "get_catalog_detail":
        data = await business.catalog.detail(args.scope, args.dataset_id, args.field_id)
        # Saved notes are not required for selection; unsaved drafts never cross this interface.
        data.pop("research", None)
    elif name == "prepare_research_input":
        data = await business.research_builder.select_input(args)
    elif name == "get_research_input":
        data = await business.research_builder.input_page(**args.model_dump())
    elif name == "prepare_research_backtest":
        data = await business.research_builder.prepare(ResearchPreviewInput(**args.model_dump()))
    elif name == "get_backtest_draft":
        data = await business.backtests.draft(args.draft_id)
        candidates = data.pop("candidates")
        data.update(
            items=candidates[args.offset : args.offset + args.limit],
            total=len(candidates),
            limit=args.limit,
            offset=args.offset,
            has_more=args.offset + args.limit < len(candidates),
        )
    elif name == "get_alpha_sources":
        data = await business.get_alpha_sources(**args.model_dump())
    elif name == "get_backtest_capabilities":
        data = await business.backtests.capabilities()
    elif name == "prepare_backtest":
        data = await business.backtests.preview(args)
    elif name == "get_backtest_preview":
        data = await business.backtests.get_preview(**args.model_dump())
    elif name == "list_backtests":
        data = await business.backtests.runs(**args.model_dump())
    elif name == "get_backtest":
        data = await business.backtests.run(args.run_id)
    elif name == "get_backtest_results":
        data = await business.backtests.results(**args.model_dump())
        # The complete historical response remains available through the business endpoint.
        for item in data["items"]:
            if item["result"]:
                snapshot = item["result"].pop("snapshot")
                item["result"].update({k: snapshot.get(k) for k in ("is", "os", "checks", "dateCreated")})
    elif name == "prepare_backtest_rerun":
        data = await business.backtests.rerun(args.run_id, RerunInput(item_ids=args.item_ids))
    elif name == "search_alphas":
        data = await business.search_alphas(args.filters)
        data["filters"] = args.filters.model_dump(mode="json")
    elif name == "get_alpha_pnl":
        data = await business.get_alpha_pnl(args.alpha_id)
        points = data.pop("points")
        data.update(
            alpha_id=args.alpha_id,
            count=len(points),
            first=points[0] if points else None,
            last=points[-1] if points else None,
            null_count=sum(p["value"] is None for p in points),
        )
    elif name in ("get_alpha", "get_job_status"):
        data = await getattr(business, name)(*args.model_dump().values())
    else:
        data = await getattr(business, name)()
        if isinstance(data, list):
            data = {"items": data[:20]}
    data["_meta"] = ResultMetadata(observed_at=datetime.now(timezone.utc)).model_dump(mode="json")
    return bounded(data)


async def preview_tool(business, name, args):
    if name == "start_backtest":
        return {"backtest": await business.backtests.get_preview(args.preview_id)}
    if name == "control_backtest":
        return {"backtest_run": await business.backtests.run(args.run_id), "action": args.action}
    if name in ("update_research", "bulk_update_research"):
        ids = [args.alpha_id] if name == "update_research" else args.alpha_ids
        targets, versions = [], {}
        for alpha_id in ids:
            detail = await business.get_alpha(alpha_id)
            before = detail["research"]
            versions[alpha_id] = before["version"]
            if name == "update_research":
                after = {**before, **args.changes.model_dump(exclude_unset=True)}
            else:
                after = {
                    **before,
                    "tags": sorted((set(before["tags"]) | set(args.add_tags)) - set(args.remove_tags)),
                }
                if args.state:
                    after["state"] = args.state
            # Preview and execution use the same validation rules.
            ResearchInput.model_validate({k: after[k] for k in ("note", "tags", "favorite", "state")})
            targets.append({"alpha_id": alpha_id, "before": before, "after": after})
        return {"targets": targets, "versions": versions}
    if name in ("cancel_job", "retry_job"):
        return {"job": await business.get_job_status(args.job_id)}
    return {"operation": args.model_dump(mode="json")}


async def execute_tool(business, name, args, preview):
    if name == "start_backtest":
        current = await business.backtests.get_preview(args.preview_id)
        if current["digest"] != preview["backtest"]["digest"] or current["version"] != args.version:
            from fastapi import HTTPException

            raise HTTPException(409, "回测预览不匹配，请重新确认")
        return await business.backtests.start(args)
    if name == "control_backtest":
        return await business.backtests.control(
            args.run_id, ControlInput(action=args.action, version=preview["backtest_run"]["version"])
        )
    if name == "update_research":
        body = ResearchUpdate(
            **args.changes.model_dump(exclude_unset=True), version=preview["versions"][args.alpha_id]
        )
        return await business.update_research(args.alpha_id, body)
    if name == "bulk_update_research":
        return await business.bulk_update_research(
            BulkUpdate(**args.model_dump(), versions=preview["versions"])
        )
    if name == "create_sync_job":
        return await business.create_sync_job(args)
    current = await business.get_job_status(args.job_id)
    if current["updated_at"] != preview["job"]["updated_at"] or current["status"] != preview["job"]["status"]:
        from fastapi import HTTPException

        raise HTTPException(409, "任务状态已变化，请重新确认操作")
    return await getattr(business, name)(args.job_id)