feat: compose native research workflows in QuantFlow
This commit is contained in:
yuxuanhui
@@ -0,0 +1,121 @@
|
||||
"""Stage-four PostgreSQL acceptance in dedicated databases only."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
NAME = "wq_research_stage4_test"
|
||||
RESTORE = "wq_research_restore_stage4"
|
||||
os.environ.update(
|
||||
DATABASE_URL=f"postgresql+asyncpg://postgres:research-test-only@127.0.0.1:18436/{NAME}",
|
||||
ADMIN_PASSWORD="research-acceptance-only",
|
||||
ENCRYPTION_KEY=Fernet.generate_key().decode(),
|
||||
)
|
||||
|
||||
|
||||
def docker(*args, **kwargs):
|
||||
return subprocess.run(["docker", "exec", "-i", "wq-research-acceptance-pg", *args], check=True, **kwargs)
|
||||
|
||||
|
||||
async def acceptance():
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.config import Settings
|
||||
from app.main import create_app
|
||||
from app.models import TemplateInput
|
||||
from app.research.workspace_contracts import TemplateSpec
|
||||
from tests.test_ai import configure
|
||||
from tests.test_backtests import setup
|
||||
from tests.test_quantflow import graph, launch, save
|
||||
from tests.test_research_flows import drive
|
||||
from tests.test_research_workspace import template
|
||||
|
||||
app = create_app(Settings(_env_file=None, enable_runner=False, public_origin="http://testserver"))
|
||||
async with app.router.lifespan_context(app):
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app), base_url="http://testserver", headers={"X-WQ-Request": "1"}
|
||||
) as client:
|
||||
assert (
|
||||
await client.post(
|
||||
"/api/v1/auth/login", json={"username": "admin", "password": "research-acceptance-only"}
|
||||
)
|
||||
).status_code == 200
|
||||
# Configure a deterministic model; no provider or real platform network.
|
||||
from tests.ai_fake import fake_model
|
||||
|
||||
app.state.ai.model_factory = fake_model
|
||||
await configure(app, client)
|
||||
platform, lane = await setup(app)
|
||||
calls = []
|
||||
|
||||
async def model(ai, context, output_type, revision):
|
||||
calls.append(context)
|
||||
value = template()
|
||||
value["expression"] = f"rank({{field}}) + {len(calls)}"
|
||||
return TemplateSpec.model_validate(value), {
|
||||
"model": "fixture",
|
||||
"revision": revision,
|
||||
"usage": {"requests": 1},
|
||||
}
|
||||
|
||||
async with app.state.sessions() as db:
|
||||
fixed = await db.scalar(select(TemplateInput))
|
||||
body = {
|
||||
"request_id": "finite-run",
|
||||
"name": "PG 有限研究",
|
||||
"input_ids": [fixed.id],
|
||||
"hypothesis": "排名稳定性",
|
||||
"settings": {"region": "USA", "universe": "TOP3000", "delay": 1},
|
||||
"budget": {"max_rounds": 2, "max_simulations": 4, "max_model_calls": 3},
|
||||
"batch_candidates": 2,
|
||||
}
|
||||
body["request_id"] = "pg-quantflow"
|
||||
fixed_template = await save(client, "template", template())
|
||||
definition = graph("input", "expand", "backtest", "evaluate", "filter", "iterate")
|
||||
definition["nodes"][1]["config"] = {"asset_id": fixed_template["id"], "version": 1}
|
||||
definition["nodes"][4]["config"] = {"verdicts": ["pass", "review", "block"]}
|
||||
with patch("app.research.runtime.request_model", model):
|
||||
run = await launch(client, body, definition)
|
||||
result = await drive(app, client, run["id"], lane, ticks=60)
|
||||
assert result["status"] == "completed", result
|
||||
assert (
|
||||
result["round"] == 2
|
||||
and result["simulations_used"] == 4
|
||||
and result["model_calls_used"] == 0
|
||||
)
|
||||
assert not calls and len(platform.posts) == 2
|
||||
await app.state.research.recover()
|
||||
await app.state.research.advance(run["id"])
|
||||
assert len(platform.posts) == 2
|
||||
print("PASS PostgreSQL: versioned native graph, filtering, bounded iteration and replay")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
docker("createdb", "-U", "postgres", NAME)
|
||||
with Path("/tmp/wq-research-stage3.dump").open("rb") as source:
|
||||
docker("pg_restore", "-U", "postgres", "-d", NAME, stdin=source)
|
||||
config = Config("alembic.ini")
|
||||
command.upgrade(config, "0008")
|
||||
command.check(config)
|
||||
asyncio.run(acceptance())
|
||||
dump = Path("/tmp/wq-research-stage4.dump")
|
||||
with dump.open("wb") as output:
|
||||
docker("pg_dump", "-U", "postgres", "-Fc", NAME, stdout=output)
|
||||
docker("createdb", "-U", "postgres", RESTORE)
|
||||
with dump.open("rb") as source:
|
||||
docker("pg_restore", "-U", "postgres", "-d", RESTORE, stdin=source)
|
||||
query = "SELECT (SELECT md5(string_agg(row_to_json(t)::text, '' ORDER BY id)) FROM research_flow_runs t),(SELECT md5(string_agg(row_to_json(t)::text, '' ORDER BY id)) FROM research_step_runs t),(SELECT count(*) FROM research_revisions),(SELECT note FROM research WHERE alpha_id='OLD_RESEARCH')"
|
||||
a = docker("psql", "-U", "postgres", "-d", NAME, "-Atc", query, capture_output=True).stdout
|
||||
b = docker("psql", "-U", "postgres", "-d", RESTORE, "-Atc", query, capture_output=True).stdout
|
||||
assert a == b
|
||||
print(
|
||||
"PASS PostgreSQL 17: head 0008 unchanged, pg_dump/pg_restore preserves complete run/step snapshots and old research notes"
|
||||
)
|
||||
@@ -0,0 +1,286 @@
|
||||
"""Native graph execution shares immutable inputs, artifacts and budgeted backtests."""
|
||||
|
||||
import copy
|
||||
|
||||
import pytest
|
||||
|
||||
from app.alphas import upsert_alpha
|
||||
from app.models import Alpha, CatalogResource
|
||||
from app.research.workspace_contracts import FeatureSpec
|
||||
from tests.conftest import alpha
|
||||
from tests.test_catalog import SCOPE, prepare, sync
|
||||
from tests.test_research_flows import begin, drive, flow_setup, get
|
||||
from tests.test_research_workspace import catalog, research_input, template
|
||||
|
||||
__all__ = ["flow_setup", "catalog", "research_input"]
|
||||
|
||||
|
||||
def graph(*kinds):
|
||||
return {
|
||||
"name": "原生节点研究",
|
||||
"nodes": [
|
||||
{
|
||||
"id": f"n{i}",
|
||||
"type": kind,
|
||||
"label": kind,
|
||||
"config": {"max_rounds": 2} if kind == "iterate" else {},
|
||||
}
|
||||
for i, kind in enumerate(kinds)
|
||||
],
|
||||
"edges": [{"source": f"n{i}", "target": f"n{i + 1}"} for i in range(len(kinds) - 1)],
|
||||
}
|
||||
|
||||
|
||||
async def save(client, kind, content):
|
||||
r = await client.post("/api/v1/research/assets", json={"kind": kind, "content": content})
|
||||
assert r.status_code == 201, r.text
|
||||
return r.json()
|
||||
|
||||
|
||||
async def launch(client, body, definition):
|
||||
asset = await save(client, "workflow", definition)
|
||||
return await begin(client, {**body, "workflow_id": asset["id"], "workflow_version": asset["version"]})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"case", ["cycle", "type", "duplicate", "disconnected", "loop", "config", "filter", "ref"]
|
||||
)
|
||||
async def test_invalid_graphs_rejected_before_save(logged_in, case):
|
||||
value = graph("input", "generate", "expand", "backtest", "evaluate")
|
||||
if case == "cycle":
|
||||
value["nodes"] += [{"id": "loop_a", "type": "generate"}, {"id": "loop_b", "type": "expand"}]
|
||||
value["edges"] += [{"source": "loop_a", "target": "loop_b"}, {"source": "loop_b", "target": "loop_a"}]
|
||||
if case == "type":
|
||||
value["edges"][0]["target"] = "n3"
|
||||
if case == "duplicate":
|
||||
value["edges"].append(value["edges"][0])
|
||||
if case == "disconnected":
|
||||
value["edges"].pop()
|
||||
if case == "loop":
|
||||
value["nodes"].append({"id": "iterate", "type": "iterate", "config": {"max_rounds": True}})
|
||||
value["edges"].append({"source": "n4", "target": "iterate"})
|
||||
if case == "config":
|
||||
value["nodes"][1]["config"] = {"command": "bash"}
|
||||
if case == "filter":
|
||||
value["nodes"].append({"id": "filter", "type": "filter", "config": {"verdicts": [{}]}})
|
||||
value["edges"].append({"source": "n4", "target": "filter"})
|
||||
if case == "ref":
|
||||
value["nodes"][2]["config"] = {"asset_id": "x", "version": True}
|
||||
for url in ("/api/v1/research/flows/validate", "/api/v1/research/assets"):
|
||||
r = await logged_in.post(
|
||||
url, json=value if url.endswith("validate") else {"kind": "workflow", "content": value}
|
||||
)
|
||||
assert r.status_code == 422, r.text
|
||||
|
||||
|
||||
async def test_versioned_template_branch_join_and_replayed_events(app, logged_in, flow_setup):
|
||||
body, platform, lane, calls = flow_setup
|
||||
fixed = await save(logged_in, "template", template())
|
||||
definition = graph("input", "expand", "backtest", "evaluate", "condition")
|
||||
definition["nodes"][1]["config"] = {"asset_id": fixed["id"], "version": 1}
|
||||
for verdict in ("pass", "review", "block"):
|
||||
definition["nodes"].append({"id": verdict, "type": "summarize"})
|
||||
definition["edges"].append({"source": "n4", "target": verdict, "branch": verdict})
|
||||
# A shared join receives the condition's evidence as well as direct evaluation.
|
||||
definition["nodes"].append({"id": "summary", "type": "summarize"})
|
||||
definition["edges"] += [{"source": "n4", "target": "summary"}, {"source": "n3", "target": "summary"}]
|
||||
run = await launch(logged_in, body, definition)
|
||||
workflow = run["authorization"]["workflow"]
|
||||
modified = copy.deepcopy(definition)
|
||||
modified["nodes"][1]["config"]["version"] = 2
|
||||
assert (
|
||||
await logged_in.put(
|
||||
f"/api/v1/research/assets/{workflow['id']}",
|
||||
json={"kind": "workflow", "version": 1, "content": modified},
|
||||
)
|
||||
).status_code == 200
|
||||
value = template()
|
||||
value["expression"] = "unknown({field})"
|
||||
assert (
|
||||
await logged_in.put(
|
||||
f"/api/v1/research/assets/{fixed['id']}",
|
||||
json={"kind": "template", "version": 1, "content": value},
|
||||
)
|
||||
).status_code == 200
|
||||
final = await drive(app, logged_in, run["id"], lane)
|
||||
assert final["status"] == "completed", final
|
||||
assert final["model_calls_used"] == 0 and not calls and final["simulations_used"] == 2
|
||||
branches = [s for s in final["steps"] if s["node_id"] in ("pass", "review", "block")]
|
||||
verdict = next(s for s in final["steps"] if s["node_id"] == "n3")["output"]["verdict"]
|
||||
assert [s["node_id"] for s in branches if s["status"] == "completed"] == [verdict]
|
||||
assert sum(s["status"] == "skipped" for s in branches) == 2
|
||||
assert len(next(s for s in final["steps"] if s["node_id"] == "summary")["output"]["artifacts"]) == 2
|
||||
assert final["authorization"]["node_assets"]["n1"]["version"] == 1
|
||||
for _ in range(3):
|
||||
await app.state.research.advance(run["id"])
|
||||
assert len(platform.posts) == 1
|
||||
assert (await logged_in.get("/api/v1/research/flows/runs?kind=quantflow")).json()["total"] == 1
|
||||
assert (await logged_in.get("/api/v1/research/flows/runs?kind=pipeline")).json()["total"] == 0
|
||||
|
||||
|
||||
async def test_feature_reference_conversion_and_bounded_loop(app, logged_in, flow_setup):
|
||||
body, _, lane, calls = flow_setup
|
||||
feature = await save(
|
||||
logged_in,
|
||||
"feature",
|
||||
{
|
||||
"name": "字段方案",
|
||||
"hypothesis": "排名稳定性",
|
||||
"input_ids": body["input_ids"],
|
||||
"steps": [{"name": "排名", "rationale": "截面比较", "expression": "rank(TEST_FIN_001)"}],
|
||||
"template": template(),
|
||||
},
|
||||
)
|
||||
definition = graph("input", "feature", "expand", "backtest", "evaluate", "filter", "iterate")
|
||||
definition["nodes"][1]["config"] = {"asset_id": feature["id"], "version": 1}
|
||||
definition["nodes"][5]["config"] = {"verdicts": ["pass", "review", "block"]}
|
||||
run = await launch(logged_in, body, definition)
|
||||
final = await drive(app, logged_in, run["id"], lane, ticks=60)
|
||||
assert final["status"] == "completed", final
|
||||
assert final["round"] == 2 and final["simulations_used"] == 4 and not calls
|
||||
outputs = [s["output"] for s in final["steps"] if s["node_id"] == "n1"]
|
||||
assert all(o["feature"]["id"] == feature["id"] for o in outputs)
|
||||
template_asset = (await logged_in.get(f"/api/v1/research/assets/{outputs[0]['template']['id']}")).json()
|
||||
assert template_asset["provenance"]["feature"]["version"] == 1
|
||||
|
||||
|
||||
async def test_native_feature_model_and_structure_variant(app, logged_in, flow_setup, monkeypatch):
|
||||
body, _, lane, calls = flow_setup
|
||||
|
||||
async def feature_model(ai, context, output_type, revision):
|
||||
assert output_type is FeatureSpec
|
||||
calls.append(context)
|
||||
return FeatureSpec.model_validate(
|
||||
{
|
||||
"name": "模型特征",
|
||||
"hypothesis": "排名",
|
||||
"input_ids": body["input_ids"],
|
||||
"steps": [{"name": "排名", "rationale": "比较", "expression": "rank(TEST_FIN_001)"}],
|
||||
"template": template(),
|
||||
}
|
||||
), {"model": "fixture", "revision": revision}
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
patch.setattr("app.research.runtime.request_model", feature_model)
|
||||
run = await launch(logged_in, body, graph("input", "feature", "expand", "backtest", "evaluate"))
|
||||
final = await drive(app, logged_in, run["id"], lane)
|
||||
assert final["status"] == "completed" and final["model_calls_used"] == 1, final
|
||||
assert final["steps"][1]["output"]["feature"]
|
||||
async with app.state.sessions.begin() as db:
|
||||
await upsert_alpha(db, alpha("seed", regular={"code": "rank(TEST_FIN_001)"}))
|
||||
run = await launch(
|
||||
logged_in,
|
||||
{**body, "request_id": "variant", "parent_alpha_ids": ["seed"]},
|
||||
graph("input", "variant", "backtest", "evaluate"),
|
||||
)
|
||||
final = await drive(app, logged_in, run["id"], lane)
|
||||
assert final["status"] == "completed" and final["model_calls_used"] == 1, final
|
||||
experiment = (
|
||||
await logged_in.get(f"/api/v1/research/experiments/{final['steps'][1]['output']['experiment_id']}")
|
||||
).json()
|
||||
assert experiment["parents"][0]["id"] == "seed" and experiment["kind"] == "quantflow"
|
||||
|
||||
|
||||
async def test_settings_variant_uses_frozen_seed_and_target_scope(app, logged_in, research_input, catalog):
|
||||
from tests.test_backtests import setup
|
||||
|
||||
body = {
|
||||
"request_id": "scope-run",
|
||||
"name": "范围变体",
|
||||
"input_ids": [research_input["id"]],
|
||||
"hypothesis": "跨股票池比较",
|
||||
"settings": {"region": "USA", "universe": "TOP3000", "delay": 1},
|
||||
"budget": {"max_rounds": 1, "max_simulations": 2, "max_model_calls": 1},
|
||||
}
|
||||
target = {**SCOPE, "universe": "TOP1000"}
|
||||
await sync(catalog, scope=target)
|
||||
version = (await sync(catalog, "TEST_FIN", scope=target))["id"]
|
||||
fixed = (await prepare(logged_in, version, scope=target)).json()
|
||||
async with app.state.sessions.begin() as db:
|
||||
await upsert_alpha(db, alpha("seed", regular={"code": "x = TEST_FIN_001; rank(x)"}))
|
||||
metadata = await db.get(CatalogResource, "settings")
|
||||
metadata.content = {
|
||||
"items": metadata.content["items"] + [{**target, "neutralizations": ["INDUSTRY"]}]
|
||||
}
|
||||
_, lane = await setup(app)
|
||||
definition = graph("input", "variant", "backtest", "evaluate")
|
||||
definition["nodes"][1]["config"] = {"method": "settings"}
|
||||
run = await launch(
|
||||
logged_in,
|
||||
{**body, "input_ids": body["input_ids"] + [fixed["id"]], "parent_alpha_ids": ["seed"]},
|
||||
definition,
|
||||
)
|
||||
async with app.state.sessions.begin() as db:
|
||||
seed = await db.get(Alpha, "seed")
|
||||
seed.expression = "rank(unknown)"
|
||||
final = await drive(app, logged_in, run["id"], lane)
|
||||
assert (
|
||||
final["status"] == "completed" and final["simulations_used"] == 1 and final["model_calls_used"] == 0
|
||||
), final
|
||||
experiment = (
|
||||
await logged_in.get(f"/api/v1/research/experiments/{final['steps'][1]['output']['experiment_id']}")
|
||||
).json()
|
||||
c = experiment["candidates"][0]
|
||||
assert c["expression"] == "x = TEST_FIN_001; rank(x)" and c["input_ids"] == [fixed["id"]]
|
||||
assert c["settings"]["universe"] == "TOP1000" and len(experiment["inputs"]) == 2
|
||||
|
||||
|
||||
async def test_empty_filter_skips_downstream_simulation(app, logged_in, flow_setup):
|
||||
body, platform, lane, _ = flow_setup
|
||||
definition = graph(
|
||||
"input", "generate", "expand", "backtest", "evaluate", "filter", "backtest", "evaluate"
|
||||
)
|
||||
definition["nodes"][5]["config"] = {"verdicts": []}
|
||||
run = await launch(logged_in, body, definition)
|
||||
final = await drive(app, logged_in, run["id"], lane)
|
||||
assert final["status"] == "completed" and len(platform.posts) == 1, final
|
||||
assert [s["status"] for s in final["steps"][-3:]] == ["skipped"] * 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("action", ["pause", "stop", "account"])
|
||||
async def test_collected_variant_output_waits_for_authorized_next_tick(app, logged_in, flow_setup, action):
|
||||
from sqlalchemy import func, select
|
||||
|
||||
from app.models import Account, ResearchExperiment
|
||||
from app.research.workspace_contracts import TemplateSpec
|
||||
|
||||
body, platform, lane, _ = flow_setup
|
||||
async with app.state.sessions.begin() as db:
|
||||
await upsert_alpha(db, alpha("seed", regular={"code": "rank(TEST_FIN_001)"}))
|
||||
run = await launch(
|
||||
logged_in, {**body, "parent_alpha_ids": ["seed"]}, graph("input", "variant", "backtest", "evaluate")
|
||||
)
|
||||
await app.state.research.advance(run["id"])
|
||||
work = await app.state.research.prepare(run["id"])
|
||||
current = await get(logged_in, run["id"])
|
||||
if action == "account":
|
||||
async with app.state.sessions.begin() as db:
|
||||
account = await db.get(Account, 1)
|
||||
account.wq_user_id = "OTHER_ACCOUNT"
|
||||
else:
|
||||
assert (
|
||||
await logged_in.post(
|
||||
f"/api/v1/research/flows/runs/{run['id']}/control",
|
||||
json={"action": action, "version": current["version"]},
|
||||
)
|
||||
).status_code == 200
|
||||
await app.state.research.finish_model(
|
||||
run["id"], work[0], TemplateSpec.model_validate(template()), {"model": "fixture"}
|
||||
)
|
||||
await app.state.research.advance(run["id"])
|
||||
current = await get(logged_in, run["id"])
|
||||
assert current["steps"][1]["status"] == "generated" and current["model_calls_used"] == 1
|
||||
assert current["status"] == {"pause": "paused", "stop": "stopped", "account": "interrupted"}[action]
|
||||
async with app.state.sessions() as db:
|
||||
assert await db.scalar(select(func.count()).select_from(ResearchExperiment)) == 0
|
||||
assert not platform.posts
|
||||
if action == "pause":
|
||||
assert (
|
||||
await logged_in.post(
|
||||
f"/api/v1/research/flows/runs/{run['id']}/control",
|
||||
json={"action": "resume", "version": current["version"]},
|
||||
)
|
||||
).status_code == 200
|
||||
await app.state.research.recover()
|
||||
final = await drive(app, logged_in, run["id"], lane)
|
||||
assert final["status"] == "completed" and final["model_calls_used"] == 1
|
||||
Reference in New Issue
Block a user