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

import copy
import json
from contextlib import asynccontextmanager

import httpx
import pytest
from pydantic_ai.messages import ModelResponse, TextPart, UserPromptPart
from pydantic_ai.models.function import FunctionModel

from app.alphas import upsert_alpha
from app.models import Account, AISettings, Alpha, Job, Research, SelfCorrelation
from app.security import cipher
from app.submission import HEADINGS, Description, comparable_description, fingerprint, source
from app.worldquant import WqClient
from tests.conftest import alpha

FIELDS = {
    "idea": "Short term price reversal is a hypothesis for this signal.",
    "data_rationale": "Close prices represent the observed price history of each instrument.",
    "operator_rationale": "The delta measures five day change, rank compares stocks and negation reverses the ordering.",
}
# Captured from the failing gateway response: 649 characters including template labels.
ACCRUAL_FIELDS = {
    "idea": "Combines ranked accrual-related signals and subtracts a ranked score signal, hypothesizing that higher accrual measures and lower fscore values identify relatively attractive stocks.",
    "data_rationale": "The selected fields are named cashflow accruals, WC accruals, and fscore; their exact definitions are not provided, so interpretation is qualified. A 120-period backfill supplies missing observations.",
    "operator_rationale": "ts_backfill extends the latest available value, rank makes fields cross-sectionally comparable, addition combines the accrual signals, and subtraction gives the fscore component a negative contribution.",
}
CHECKS = [{"name": "PROD_CORRELATION", "result": "FAIL", "value": 0.8, "limit": 0.7}]


@pytest.mark.parametrize(
    "text",
    [
        "plain text",
        "Idea: actual idea\nRationale for data used:\nRationale for operators used:",
        "Idea:\nRationale for data used: actual data\nRationale for operators used:",
        "Idea:\nRationale for data used:\nRationale for operators used: actual operator",
        "Idea:\nRationale for data used:",
        "Rationale for data used:\nIdea:\nRationale for operators used:",
    ],
)
def test_substantive_or_unknown_descriptions_remain_exact(text):
    assert comparable_description(text) == text
    original = source(alpha(regular={"code": "rank(close)", "description": text}))
    changed = copy.deepcopy(original)
    changed["sections"]["regular"]["description"] = text + " "
    assert fingerprint(original) != fingerprint(changed)
    assert original["sections"]["regular"]["description"] == text


class Platform:
    def __init__(self, raw):
        self.raw = copy.deepcopy(raw)
        self.calls = []
        self.patches = []
        self.pending = 0
        self.fail_patch = False
        self.unknown_patch = False
        self.fail_check = False

    def __call__(self, request):
        self.calls.append((request.method, request.url.path))
        if request.url.path == "/authentication":
            return httpx.Response(201, json={})
        if request.method == "PATCH" and request.url.path == "/alphas/alpha1":
            self.patches.append(json.loads(request.content))
            if self.fail_patch:
                return httpx.Response(400, json={"secret": "not-for-client"})
            for section, value in self.patches[-1].items():
                self.raw[section].update(value)
            if self.unknown_patch:
                self.unknown_patch = False
                raise httpx.ReadTimeout("secret", request=request)
            return httpx.Response(204)
        if request.method == "GET" and request.url.path == "/alphas/alpha1":
            return httpx.Response(200, json=self.raw)
        if request.method == "GET" and request.url.path == "/alphas/alpha1/check":
            if self.fail_check:
                return httpx.Response(403, json={})
            if self.pending:
                self.pending -= 1
                return httpx.Response(202, headers={"Retry-After": "0"}, json={})
            return httpx.Response(200, json={"is": {"checks": CHECKS}})
        raise AssertionError(f"Unexpected platform operation {request.method} {request.url.path}")


async def setup(app, raw=None):
    raw = raw or alpha()
    async with app.state.sessions.begin() as db:
        await upsert_alpha(db, raw)
        account = await db.get(Account, 1)
        account.email = "test@example.com"
        account.password_encrypted = cipher(app.state.settings).encrypt(b"fake-password").decode()
        account.connection_status = "connected"
        db.add(SelfCorrelation(alpha_id="alpha1", region="USA", result={"status": "low"}, stale=False))
        (await db.get(Research, "alpha1")).note = "preserve local research"
    platform = Platform(raw)
    await app.state.runner.client.close()
    app.state.runner.client = WqClient(app.state.settings, transport=httpx.MockTransport(platform))
    return platform


async def enqueue(client, descriptions=None):
    state = (await client.get("/api/v1/alphas/alpha1/submission")).json()
    return await client.post(
        "/api/v1/alphas/alpha1/submission-check",
        json={
            "snapshot": state["snapshot"],
            "descriptions": descriptions or {"regular": FIELDS},
        },
    )


@pytest.mark.parametrize("kind", ["REGULAR", "SUPER"])
async def test_write_check_poll_and_preserve_snapshot(app, logged_in, kind):
    raw = (
        alpha()
        if kind == "REGULAR"
        else alpha(type="SUPER", selection={"code": "rank(close)"}, combo={"code": "alpha"})
    )
    platform = await setup(app, raw)
    platform.pending = 2
    descriptions = {key: FIELDS for key in (["regular"] if kind == "REGULAR" else ["selection", "combo"])}
    response = await enqueue(logged_in, descriptions)
    assert response.status_code == 202, response.text
    job_id = response.json()["id"]
    duplicate = await enqueue(logged_in, descriptions)
    assert duplicate.json()["id"] == job_id
    await app.state.runner.execute(job_id)
    async with app.state.sessions() as db:
        job = await db.get(Job, job_id)
        assert job.status == "completed", job.error
        assert job.processed == 1 and job.checkpoint["phase"] == "checked"
        item = await db.get(Alpha, "alpha1")
        assert item.check_type == "FAIL_1" and item.prod_correlation == 0.8
        assert item.checks == CHECKS and item.is_metrics["checks"] == CHECKS
        assert item.expression == raw["regular"]["code"] and item.sharpe == 1.5
        assert item.raw["settings"] == raw["settings"]
        assert (await db.get(Research, "alpha1")).note == "preserve local research"
        assert not (await db.get(SelfCorrelation, "alpha1")).stale
    assert platform.patches == [{key: {"description": Description(**FIELDS).text()} for key in descriptions}]
    assert platform.calls.count(("GET", "/alphas/alpha1/check")) == 3
    assert not any(path.endswith("/submit") for _, path in platform.calls)
    # Durable completion does not repeat either operation.
    await app.state.runner.execute(job_id)
    assert len(platform.patches) == 1
    state = (await logged_in.get("/api/v1/alphas/alpha1/submission")).json()
    assert state["job"]["status"] == "completed"


@pytest.mark.parametrize(
    "status,stale", [("high", False), ("partial", False), ("insufficient_data", False), ("low", True)]
)
async def test_local_correlation_blocks_upstream(app, logged_in, status, stale):
    platform = await setup(app)
    async with app.state.sessions.begin() as db:
        row = await db.get(SelfCorrelation, "alpha1")
        row.result, row.stale = {"status": status}, stale
    response = await enqueue(logged_in)
    assert response.status_code == 409
    assert platform.calls == []


@pytest.mark.parametrize("kind", ["REGULAR", "SUPER"])
@pytest.mark.parametrize("refresh", [False, True])
@pytest.mark.parametrize(
    "before,after",
    [
        ("", "Idea: \nRationale for data used: \nRationale for operators used:"),
        ("Idea:\r\n Rationale for data used:\t Rationale for operators used: ", ""),
        (None, " \r\n\t\u00a0"),
    ],
)
async def test_empty_description_representations_allow_check(app, logged_in, kind, refresh, before, after):
    sections = ["regular"] if kind == "REGULAR" else ["selection", "combo"]
    raw = alpha(type=kind, **{key: {"code": "rank(close)", "description": before} for key in sections})
    platform = await setup(app, raw)
    state = (await logged_in.get("/api/v1/alphas/alpha1/submission")).json()
    for key in sections:
        platform.raw[key]["description"] = after
    if refresh:
        async with app.state.sessions.begin() as db:
            await upsert_alpha(db, platform.raw)
    target = Description(**FIELDS).text()
    response = await logged_in.post(
        "/api/v1/alphas/alpha1/submission-check",
        json={"snapshot": state["snapshot"], "descriptions": dict.fromkeys(sections, target)},
    )
    assert response.status_code == 202, response.text
    await app.state.runner.execute(response.json()["id"])
    async with app.state.sessions() as db:
        job = await db.get(Job, response.json()["id"])
        assert job.status == "completed", job.error
    assert platform.patches == [{key: {"description": target} for key in sections}]
    assert ("GET", "/alphas/alpha1/check") in platform.calls


@pytest.mark.parametrize("change", ["description", "code", "settings", "status"])
async def test_remote_conflict_never_overwrites(app, logged_in, change):
    platform = await setup(app)
    response = await enqueue(logged_in)
    if change in ("description", "code"):
        platform.raw["regular"][change] = "changed by another user"
    elif change == "settings":
        platform.raw["settings"]["delay"] = 0
    else:
        platform.raw["status"] = "ACTIVE"
    await app.state.runner.execute(response.json()["id"])
    async with app.state.sessions() as db:
        assert (await db.get(Job, response.json()["id"])).status == "failed"
    assert platform.patches == []
    assert ("GET", "/alphas/alpha1/check") not in platform.calls


@pytest.mark.parametrize("failure", ["fail_patch", "unknown_patch", "fail_check"])
async def test_retry_reconciles_partial_writes(app, logged_in, failure):
    platform = await setup(app)
    setattr(platform, failure, True)
    response = await enqueue(logged_in)
    job_id = response.json()["id"]
    await app.state.runner.execute(job_id)
    async with app.state.sessions() as db:
        assert (await db.get(Job, job_id)).status == "failed"
        item = await db.get(Alpha, "alpha1")
        assert item.checks != CHECKS
        if failure == "fail_check":
            assert item.raw["regular"]["description"] == Description(**FIELDS).text()
    if failure == "fail_patch":
        assert ("GET", "/alphas/alpha1/check") not in platform.calls
    setattr(platform, failure, False)
    await app.state.runner.execute(job_id)
    async with app.state.sessions() as db:
        assert (await db.get(Job, job_id)).status == "completed"
    assert len(platform.patches) == (2 if failure == "fail_patch" else 1)


async def test_complete_description_reused_verbatim(app, logged_in):
    original = Description(**FIELDS).text().replace("\n", "\n\n")
    platform = await setup(app, alpha(regular={"code": "rank(close)", "description": original}))
    response = await enqueue(logged_in)
    await app.state.runner.execute(response.json()["id"])
    assert not platform.patches
    assert ("GET", "/alphas/alpha1/check") in platform.calls


@pytest.mark.parametrize("kind", ["REGULAR", "SUPER"])
@pytest.mark.parametrize(
    "text",
    [
        "Brief draft",
        "\n".join(HEADINGS),
        Description(**ACCRUAL_FIELDS).text(),
        "  An idea.\n\nMore detail.  ",
    ],
    ids=["short_plain_text", "empty_template", "recorded_over_500", "whitespace"],
)
async def test_generated_text_is_returned_verbatim_without_content_validation(app, logged_in, kind, text):
    raw = (
        alpha()
        if kind == "REGULAR"
        else alpha(type="SUPER", selection={"code": "rank(close)"}, combo={"code": "alpha"})
    )
    platform = await setup(app, raw)
    sections = ["regular"] if kind == "REGULAR" else ["selection", "combo"]
    seen = []
    prompts = []

    def answer(messages, info):
        assert not info.output_tools and not info.function_tools
        prompts.append(
            json.loads(next(p.content for m in messages for p in m.parts if isinstance(p, UserPromptPart)))
        )
        return ModelResponse(parts=[TextPart(text + prompts[-1]["section"])])

    @asynccontextmanager
    async def model_factory(config, settings):
        seen.append(
            (
                config.base_url,
                config.model,
                config.protocol,
                cipher(settings).decrypt(config.api_key_encrypted.encode()).decode(),
            )
        )
        yield FunctionModel(function=answer)

    app.state.ai.model_factory = model_factory
    response = await logged_in.put(
        "/api/v1/ai/settings",
        json={
            "base_url": "https://model.test/v1",
            "model": "bot-model",
            "description_model": "description-model",
            "protocol": "responses",
            "api_key": "shared-secret",
        },
    )
    assert response.status_code == 200 and "shared-secret" not in response.text
    state = (await logged_in.get("/api/v1/alphas/alpha1/submission")).json()
    generated = await logged_in.post(
        "/api/v1/alphas/alpha1/description/generate", json={"snapshot": state["snapshot"]}
    )
    assert generated.status_code == 200, generated.text
    assert generated.json()["descriptions"] == {key: text + key for key in sections}
    assert seen == [("https://model.test/v1", "description-model", "responses", "shared-secret")]
    assert [prompt["section"] for prompt in prompts] == sections
    assert [prompt["code"] for prompt in prompts] == [raw[key]["code"] for key in sections]
    assert not platform.calls
    async with app.state.sessions() as db:
        assert (await db.get(AISettings, 1)).model == "bot-model"
        assert "description" not in (await db.get(Alpha, "alpha1")).raw["regular"]


async def test_generation_requires_config_and_safely_reports_provider_failure(app, logged_in, caplog):
    platform = await setup(app)
    state = (await logged_in.get("/api/v1/alphas/alpha1/submission")).json()
    path = "/api/v1/alphas/alpha1/description/generate"
    assert (await logged_in.post(path, json={"snapshot": state["snapshot"]})).status_code == 409
    await logged_in.put(
        "/api/v1/ai/settings",
        json={
            "base_url": "https://model.test/v1",
            "model": "bot",
            "description_model": "description",
            "api_key": "shared-secret",
        },
    )

    def fail(messages, info):
        raise httpx.ReadTimeout("shared-secret private provider response")

    @asynccontextmanager
    async def broken(config, settings):
        yield FunctionModel(function=fail)

    app.state.ai.model_factory = broken
    response = await logged_in.post(path, json={"snapshot": state["snapshot"]})
    assert response.status_code == 502
    assert "超时" in response.json()["detail"]
    assert "ReadTimeout" in caplog.text
    assert "shared-secret" not in response.text + caplog.text
    assert "private provider response" not in response.text + caplog.text
    assert not platform.calls


@pytest.mark.parametrize("protocol", ["chat_completions", "responses"])
@pytest.mark.parametrize("returned_fields", [FIELDS, ACCRUAL_FIELDS], ids=["short", "recorded_over_500"])
async def test_description_text_output_through_real_provider(app, logged_in, protocol, returned_fields):
    from app.ai.provider import model_connection

    platform = await setup(app)
    await logged_in.put(
        "/api/v1/ai/settings",
        json={
            "base_url": "https://model.test/v1",
            "model": "bot",
            "description_model": "description-wire-model",
            "protocol": protocol,
            "api_key": "synthetic-key",
        },
    )
    paths = []

    def gateway(request):
        paths.append(request.url.path)
        body = json.loads(request.content)
        assert request.headers["authorization"] == "Bearer synthetic-key"
        assert body["model"] == "description-wire-model" and body["stream"] is False
        assert not body.get("tools") and not body.get("response_format")
        text = Description(**returned_fields).text()
        if protocol == "chat_completions":
            return httpx.Response(
                200,
                json={
                    "id": "chat-description",
                    "object": "chat.completion",
                    "created": 1789000000,
                    "model": body["model"],
                    "choices": [
                        {
                            "index": 0,
                            "finish_reason": "stop",
                            "message": {
                                "role": "assistant",
                                "content": text,
                            },
                        }
                    ],
                },
            )
        return httpx.Response(
            200,
            json={
                "id": "resp-description",
                "object": "response",
                "created_at": 1789000000,
                "model": body["model"],
                "status": "completed",
                "output": [
                    {
                        "id": "msg-description",
                        "type": "message",
                        "role": "assistant",
                        "content": [{"type": "output_text", "text": text, "annotations": []}],
                        "status": "completed",
                    }
                ],
            },
        )

    @asynccontextmanager
    async def model_factory(config, settings):
        async with model_connection(config, settings, httpx.MockTransport(gateway)) as model:
            yield model

    app.state.ai.model_factory = model_factory
    state = (await logged_in.get("/api/v1/alphas/alpha1/submission")).json()
    response = await logged_in.post(
        "/api/v1/alphas/alpha1/description/generate", json={"snapshot": state["snapshot"]}
    )
    assert response.status_code == 200, response.text
    assert response.json()["descriptions"] == {"regular": Description(**returned_fields).text()}
    assert paths == ["/v1/" + ("chat/completions" if protocol == "chat_completions" else "responses")]
    assert not platform.calls


@pytest.mark.parametrize("format", ["fields", "complete_text"])
async def test_manual_description_over_500_is_written_completely(app, logged_in, format):
    platform = await setup(app)
    fields = {**FIELDS, "idea": "x" * 600}
    text = "\n\n".join(heading + value for heading, value in zip(HEADINGS, fields.values()))
    response = await enqueue(logged_in, {"regular": fields if format == "fields" else text})
    assert response.status_code == 202, response.text
    await app.state.runner.execute(response.json()["id"])
    assert platform.patches == [{"regular": {"description": text}}]


async def test_invalid_snapshot_and_section_rejected(app, logged_in):
    platform = await setup(app)
    response = await logged_in.post(
        "/api/v1/alphas/alpha1/submission-check",
        json={
            "snapshot": "0" * 64,
            "descriptions": {"regular": FIELDS},
        },
    )
    assert response.status_code == 409
    assert (await enqueue(logged_in, {"combo": FIELDS})).status_code == 422
    assert platform.calls == []


async def test_description_model_does_not_invalidate_bot_test(app, logged_in):
    from tests.test_ai import CONFIG, configure

    await configure(app, logged_in)
    response = await logged_in.put(
        "/api/v1/ai/settings",
        json={
            **{k: v for k, v in CONFIG.items() if k != "api_key"},
            "enabled": True,
            "description_model": "  separate-description-model  ",
        },
    )
    assert response.json()["ready"] and response.json()["enabled"]
    assert response.json()["description_model"] == "separate-description-model"
    # Older clients saving bot settings do not clear the independently configured model.
    response = await logged_in.put(
        "/api/v1/ai/settings",
        json={
            **{k: v for k, v in CONFIG.items() if k != "api_key"},
            "enabled": True,
        },
    )
    assert response.json()["description_model"] == "separate-description-model"


def test_description_model_migration_preserves_existing_config(tmp_path):
    import importlib.util
    from pathlib import Path

    import sqlalchemy as sa
    from alembic.migration import MigrationContext
    from alembic.operations import Operations

    path = Path(__file__).parents[1] / "migrations/versions/0012_description_model.py"
    spec = importlib.util.spec_from_file_location("description_migration", path)
    migration = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(migration)
    engine = sa.create_engine(f"sqlite:///{tmp_path}/migration.db")
    with engine.begin() as connection:
        connection.exec_driver_sql("CREATE TABLE ai_settings (id INTEGER PRIMARY KEY, model VARCHAR(200))")
        connection.exec_driver_sql("INSERT INTO ai_settings VALUES (1, 'keep-bot-model')")
        with Operations.context(MigrationContext.configure(connection)):
            migration.upgrade()
            assert connection.exec_driver_sql("SELECT model, description_model FROM ai_settings").one() == (
                "keep-bot-model",
                "",
            )
            migration.downgrade()
            assert connection.exec_driver_sql("SELECT model FROM ai_settings").scalar() == "keep-bot-model"
            assert "description_model" not in {
                c["name"] for c in sa.inspect(connection).get_columns("ai_settings")
            }
    engine.dispose()


@pytest.mark.parametrize("cached", [False, True])
async def test_no_local_references_does_not_block_platform_check(app, logged_in, cached):
    platform = await setup(app)
    async with app.state.sessions.begin() as db:
        result = await db.get(SelfCorrelation, "alpha1")
        if cached:
            result.result = {"status": "insufficient_data", "candidate_count": 0}
        else:
            await db.delete(result)
    state = (await logged_in.get("/api/v1/alphas/alpha1/submission")).json()
    assert state["can_check"] is True
    response = await enqueue(logged_in)
    assert response.status_code == 202, response.text
    await app.state.runner.execute(response.json()["id"])
    async with app.state.sessions() as db:
        job = await db.get(Job, response.json()["id"])
        assert job.status == "completed", job.error
    assert ("GET", "/alphas/alpha1/check") in platform.calls
    assert not any(path.endswith("/submit") for _, path in platform.calls)


async def test_new_reference_blocks_empty_cached_result_at_enqueue_and_execution(app, logged_in):
    platform = await setup(app)
    async with app.state.sessions.begin() as db:
        (await db.get(SelfCorrelation, "alpha1")).result = {
            "status": "insufficient_data",
            "candidate_count": 0,
        }
    response = await enqueue(logged_in)
    assert response.status_code == 202
    async with app.state.sessions.begin() as db:
        await upsert_alpha(db, alpha(id="peer", status="ACTIVE"))
    assert not (await logged_in.get("/api/v1/alphas/alpha1/submission")).json()["can_check"]
    assert (await enqueue(logged_in)).status_code == 409
    await app.state.runner.execute(response.json()["id"])
    async with app.state.sessions() as db:
        assert (await db.get(Job, response.json()["id"])).status == "failed"
    assert not platform.patches
    assert ("GET", "/alphas/alpha1/check") not in platform.calls


async def test_complete_text_input_is_written_as_three_paragraphs(app, logged_in):
    platform = await setup(app)
    text = Description(**FIELDS).text()
    response = await enqueue(logged_in, {"regular": text})
    assert response.status_code == 202
    await app.state.runner.execute(response.json()["id"])
    assert platform.patches == [{"regular": {"description": text}}]
    assert len(text.split("\n\n")) == 3
    state = (await logged_in.get("/api/v1/alphas/alpha1/submission")).json()
    assert state["descriptions"] == {"regular": text}


@pytest.mark.parametrize("heading,field", list(zip(HEADINGS, FIELDS)))
@pytest.mark.parametrize("failure", ["missing_label", "empty_content"])
async def test_incomplete_template_can_be_reviewed_and_written(app, logged_in, heading, field, failure):
    platform = await setup(app)
    text = Description(**FIELDS).text()
    bad = text.replace(heading, "") if failure == "missing_label" else text.replace(FIELDS[field], " \n ")
    response = await enqueue(logged_in, {"regular": bad})
    assert response.status_code == 202, response.text
    await app.state.runner.execute(response.json()["id"])
    assert platform.patches == [{"regular": {"description": bad}}]


@pytest.mark.parametrize("text", ["", "short", "x" * 600, "  Idea: draft\nMore details.  "])
async def test_reviewed_plain_text_is_written_verbatim(app, logged_in, text):
    platform = await setup(app)
    response = await enqueue(logged_in, {"regular": text})
    assert response.status_code == 202, response.text
    await app.state.runner.execute(response.json()["id"])
    assert platform.patches == ([{"regular": {"description": text}}] if text else [])