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

"""Synthetic OpenAI wire protocol server; no outbound network or business access."""

import json


def model_events(body, protocol):
    """Return concrete SSE frames used by the real SDK adapters in acceptance tests."""
    inputs = body.get("messages", body.get("input", []))
    outputs = [p for p in inputs if p.get("role") == "tool" or p.get("type") == "function_call_output"]
    prompt = json.dumps(inputs, ensure_ascii=False)
    tool = None
    text = "READY"
    if outputs:
        text = outputs[-1].get("content", outputs[-1].get("output", ""))
    elif "capability_probe" in json.dumps(body.get("tools", [])):
        tool = ("capability_probe", {})
    elif "修改" in prompt:
        tool = ("update_research", {"alpha_id": "DOCKER_ACCEPTANCE", "changes": {"note": "AI verified note"}})
    elif "查询" in prompt:
        tool = ("search_alphas", {"filters": {"limit": 5, "turnover_max": 0.15}})
    elif "SLOW" in prompt:
        text = "STREAM READY"
    if protocol == "chat_completions":
        base = {
            "id": "chat-mock",
            "object": "chat.completion.chunk",
            "created": 1788739200,
            "model": body["model"],
        }
        delta = (
            {
                "tool_calls": [
                    {
                        "index": 0,
                        "id": "call-mock",
                        "type": "function",
                        "function": {"name": tool[0], "arguments": json.dumps(tool[1])},
                    }
                ]
            }
            if tool
            else {"content": text}
        )
        frames = [
            base | {"choices": [{"index": 0, "delta": delta, "finish_reason": None}]},
            base
            | {
                "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls" if tool else "stop"}],
                "usage": {"prompt_tokens": 12, "completion_tokens": 6, "total_tokens": 18},
            },
        ]
        return ["data: " + json.dumps(frame) + "\n\n" for frame in frames] + ["data: [DONE]\n\n"]
    response = {
        "id": "resp-mock",
        "object": "response",
        "created_at": 1788739200,
        "model": body["model"],
        "status": "in_progress",
        "output": [],
        "error": None,
        "incomplete_details": None,
        "usage": None,
    }
    frames = [{"type": "response.created", "response": response}]
    if tool:
        item = {
            "id": "fc-mock",
            "type": "function_call",
            "call_id": "call-mock",
            "name": tool[0],
            "arguments": "",
            "status": "in_progress",
        }
        frames += [
            {"type": "response.output_item.added", "output_index": 0, "item": item},
            {
                "type": "response.function_call_arguments.delta",
                "item_id": "fc-mock",
                "output_index": 0,
                "delta": json.dumps(tool[1]),
            },
            {
                "type": "response.output_item.done",
                "output_index": 0,
                "item": item | {"arguments": json.dumps(tool[1]), "status": "completed"},
            },
        ]
    else:
        item = {
            "id": "msg-mock",
            "type": "message",
            "role": "assistant",
            "status": "in_progress",
            "content": [],
        }
        frames += [
            {"type": "response.output_item.added", "output_index": 0, "item": item},
            {
                "type": "response.output_text.delta",
                "item_id": "msg-mock",
                "output_index": 0,
                "content_index": 0,
                "delta": text,
                "logprobs": [],
            },
        ]
    frames.append(
        {
            "type": "response.completed",
            "response": response
            | {
                "status": "completed",
                "usage": {
                    "input_tokens": 12,
                    "output_tokens": 6,
                    "total_tokens": 18,
                    "input_tokens_details": {"cached_tokens": 0},
                    "output_tokens_details": {"reasoning_tokens": 0},
                },
            },
        }
    )
    return [
        "event: " + frame["type"] + "\ndata: " + json.dumps(frame | {"sequence_number": i}) + "\n\n"
        for i, frame in enumerate(frames)
    ]


if __name__ == "__main__":
    import time
    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

    class Handler(BaseHTTPRequestHandler):
        def do_POST(self):
            body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
            self.send_response(200)
            self.send_header("Content-Type", "text/event-stream")
            self.end_headers()
            for frame in model_events(
                body, "responses" if self.path.endswith("/responses") else "chat_completions"
            ):
                self.wfile.write(frame.encode())
                self.wfile.flush()
                if "SLOW" in json.dumps(body):
                    time.sleep(1)

        def log_message(self, *args):
            pass

    ThreadingHTTPServer(("127.0.0.1", 19010), Handler).serve_forever()