EvoScientist-Multi/tests/test_gateway_proxy.py

import json

import httpx
import pytest
from langchain_core.messages import HumanMessage

from EvoScientist.llm.contracts import EvoRuntimeError
from EvoScientist.llm.gateway_proxy import GatewayProxyChatModel


class _FakeStream:
    def __init__(self, lines):
        self._lines = lines

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    def raise_for_status(self):
        pass

    async def aiter_lines(self):
        for line in self._lines:
            yield line


class _FakeClient:
    def __init__(self, lines):
        self._lines = lines
        self.sent = None

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc):
        return False

    def stream(self, method, url, json=None, headers=None):
        assert method == "POST"
        assert url.endswith("/api/internal/recoverable-runs/model/stream")
        self.sent = json
        self.headers = headers
        return _FakeStream(self._lines)


def test_runtime_error_repr_preserves_only_stable_code():
    error = EvoRuntimeError(
        "UPSTREAM_RATE_LIMITED",
        "safe display message",
        details=({"provider_request": "must-not-persist"},),
    )

    assert repr(error) == "EvoRuntimeError(code='UPSTREAM_RATE_LIMITED')"
    assert "safe display message" not in repr(error)
    assert "must-not-persist" not in repr(error)


@pytest.mark.anyio
async def test_astream_yields_chunks_from_sse(monkeypatch):
    monkeypatch.setenv("AI4SCI_EVO_RUNTIME_GRANT_SECRET", "runtime-service-secret")
    model = GatewayProxyChatModel(
        gateway_url="http://gw",
        run_id="run-1",
        envelope_signature="sig",
    )
    msg = {"type": "AIMessageChunk", "data": {"content": "hello"}}
    lines = [
        f"data: {json.dumps({'delta': {'message': msg}})}\n",
        'data: {"delta": {"message": {"type": "AIMessageChunk", "data": {"content": " world"}}}}\n',
        "data: [DONE]\n",
    ]
    fake = _FakeClient(lines)
    monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: fake)

    chunks = [c async for c in model._astream([HumanMessage(content="hi")])]

    assert fake.sent["stream"] is True
    assert fake.headers == {"X-Ai4Sci-Service-Token": "runtime-service-secret"}
    assert len(chunks) == 2
    assert chunks[0].message.content == "hello"


@pytest.mark.anyio
async def test_astream_roundtrips_streaming_tool_call_chunks(monkeypatch):
    model = GatewayProxyChatModel(
        gateway_url="http://gw",
        run_id="run-1",
        envelope_signature="sig",
    )
    msg = {
        "type": "AIMessageChunk",
        "data": {
            "content": "",
            "tool_call_chunks": [
                {
                    "name": "read_file",
                    "args": '{"path":',
                    "id": "call-1",
                    "index": 0,
                    "type": "tool_call_chunk",
                }
            ],
        },
    }
    lines = [f"data: {json.dumps({'delta': {'message': msg}})}\n", "data: [DONE]\n"]
    fake = _FakeClient(lines)
    monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: fake)

    chunks = [c async for c in model._astream([HumanMessage(content="hi")])]

    assert len(chunks) == 1
    assert chunks[0].message.tool_call_chunks[0]["name"] == "read_file"
    assert chunks[0].message.tool_call_chunks[0]["id"] == "call-1"


@pytest.mark.anyio
async def test_astream_raises_on_missing_done(monkeypatch):
    model = GatewayProxyChatModel(
        gateway_url="http://gw",
        run_id="run-1",
        envelope_signature="sig",
    )
    lines = [
        'data: {"delta": {"message": {"type": "AIMessageChunk", "data": {"content": "hi"}}}}\n'
    ]
    fake = _FakeClient(lines)
    monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: fake)

    with pytest.raises(RuntimeError, match="AI4SCI_MODEL_STREAM_INCOMPLETE"):
        _ = [c async for c in model._astream([HumanMessage(content="hi")])]


@pytest.mark.anyio
async def test_astream_projects_gateway_error_frame(monkeypatch):
    model = GatewayProxyChatModel(
        gateway_url="http://gw",
        run_id="run-1",
        envelope_signature="sig",
    )
    lines = [
        'data: {"type":"error","code":"UPSTREAM_RATE_LIMITED",'
        '"status":429,"message":"模型服务请求频率超限，请稍后重试或切换模型。"}\n'
    ]
    fake = _FakeClient(lines)
    monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: fake)

    with pytest.raises(EvoRuntimeError) as exc_info:
        _ = [c async for c in model._astream([HumanMessage(content="hi")])]

    assert exc_info.value.code == "UPSTREAM_RATE_LIMITED"
    assert exc_info.value.details == ({"http_status": 429},)