feat: compose native research workflows in QuantFlow

2026-09-08 22:42:29 +08:00
parent 7860434b04
commit a6e36e50ec
22 changed files with 2274 additions and 85 deletions
@@ -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"
+    )