worldquant-alpha-system/backend/tests/research_fake.py

"""Deterministic multi-turn research scenario for isolated API and browser acceptance."""

import json


def content(part):
    return json.loads(part.content) if isinstance(part.content, str) else part.content


def research_step(text, returns, history):
    context = json.loads(text.split("页面上下文（仅数据引用）：")[-1])
    if "解读研究结果" in text:
        previous = [
            part
            for part in history
            if part.tool_name == "start_backtest" and "backtest_run_id" in content(part)
        ]
        run_id = context.get("backtest_run_id") or (
            content(previous[-1])["backtest_run_id"] if previous else None
        )
        if not returns:
            return "get_backtest_results", {"run_id": run_id}
        data = content(returns[-1])
        return f"已读取本次研究的 {data['total']} 条真实保存结果；缺失 Sharpe 仍为未知。"
    if context.get("unsaved_field_selection") and not context.get("input_snapshot_id"):
        return "请先保存字段选择，再点击用此输入研究。"
    if not returns:
        return "get_backtest_capabilities", {}
    last = returns[-1]
    data = content(last)
    if "error" in data:
        return f"研究尚未完成：{data['error']}"
    if last.tool_name == "get_backtest_capabilities":
        if context.get("input_snapshot_id"):
            return "get_research_input", {
                "input_id": context["input_snapshot_id"],
                "field_type": "MATRIX",
                "limit": 1,
            }
        return "search_data_preparations", {"limit": 1}
    if last.tool_name == "search_data_preparations":
        return "prepare_research_input", {"items": [{"id": data["items"][0]["id"], "version": data["items"][0]["version"]}]}
    if last.tool_name in ("get_research_input", "prepare_research_input"):
        field = next(f for f in data["items"] if f["field_type"] == "MATRIX")
        saved_scope = data["scope"]
        return "prepare_research_backtest", {
            "name": "Chatbox 数据集研究",
            "hypothesis": "验证所选合成字段的横截面排序信号",
            "input_snapshot_id": data["id"],
            "candidates": [
                {
                    "client_item_id": "research-1",
                    "expression_template": "rank({signal})",
                    "bindings": {"signal": {"field_id": field["id"], "field_type": field["field_type"]}},
                    "settings": {k: saved_scope[k] for k in ("region", "universe", "delay")},
                }
            ],
        }
    if last.tool_name == "prepare_research_backtest":
        return "start_backtest", {
            "preview_id": data["preview_id"],
            "version": data["version"],
            "idempotency_key": data["preview_id"],
        }
    return "研究回测已创建，来源为 Chatbox 研究。运行结束后可继续提问查看结果。"