zk-data-agent/e2e/stub_provider.py

from __future__ import annotations

import json
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any


def _last_user_text(messages: list[dict[str, Any]]) -> str:
    for message in reversed(messages):
        if message.get("role") != "user":
            continue
        content = message.get("content", "")
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            return " ".join(str(part.get("text", "")) for part in content if isinstance(part, dict))
    return ""


def _completion(model: str, message: dict[str, Any], finish_reason: str) -> dict[str, Any]:
    return {
        "id": f"chatcmpl-stub-{uuid.uuid4().hex}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
        "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
    }


class StubHandler(BaseHTTPRequestHandler):
    server_version = "K1412E2EStub/1"

    def log_message(self, format: str, *args: Any) -> None:
        return

    def _json(self, status: int, payload: dict[str, Any]) -> None:
        encoded = json.dumps(payload, ensure_ascii=False).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(encoded)))
        self.end_headers()
        self.wfile.write(encoded)

    def do_GET(self) -> None:
        if self.path == "/health":
            self._json(200, {"status": "ok"})
            return
        self._json(404, {"error": "not found"})

    def do_POST(self) -> None:
        if self.path != "/v1/chat/completions":
            self._json(404, {"error": "not found"})
            return
        length = int(self.headers.get("Content-Length", "0"))
        payload = json.loads(self.rfile.read(length) or b"{}")
        model = str(payload.get("model", "stub"))
        messages = payload.get("messages") or []
        tools = payload.get("tools") or []
        has_tool_result = any(message.get("role") == "tool" for message in messages)
        called_tools = {
            str((call.get("function") or {}).get("name", ""))
            for message in messages
            if message.get("role") == "assistant"
            for call in (message.get("tool_calls") or [])
        }
        available_names = {
            str((tool.get("function") or {}).get("name", "")) for tool in tools if isinstance(tool, dict)
        }
        is_agent = "update_plan" in available_names

        if tools and not has_tool_result:
            arguments = json.dumps(
                {"path": "work-proof.txt", "content": _last_user_text(messages) or "e2e"},
                ensure_ascii=False,
            )
            message = {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": "call_e2e_write",
                        "type": "function",
                        "function": {"name": "write_file", "arguments": arguments},
                    }
                ],
            }
            self._respond(payload, _completion(model, message, "tool_calls"))
            return

        if is_agent and "write_file" in called_tools and "read_file" not in called_tools:
            message = {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": "call_e2e_verify",
                        "type": "function",
                        "function": {
                            "name": "read_file",
                            "arguments": json.dumps({"path": "work-proof.txt"}),
                        },
                    }
                ],
            }
            self._respond(payload, _completion(model, message, "tool_calls"))
            return

        text = "E2E provider completed after tool execution." if has_tool_result else "E2E provider response."
        self._respond(payload, _completion(model, {"role": "assistant", "content": text}, "stop"))

    def _respond(self, request: dict[str, Any], completion: dict[str, Any]) -> None:
        if not request.get("stream"):
            self._json(200, completion)
            return

        message = completion["choices"][0]["message"]
        finish_reason = completion["choices"][0]["finish_reason"]
        chunks: list[dict[str, Any]] = []
        if message.get("tool_calls"):
            call = message["tool_calls"][0]
            chunks.append(
                {
                    "id": completion["id"],
                    "object": "chat.completion.chunk",
                    "created": completion["created"],
                    "model": completion["model"],
                    "choices": [
                        {
                            "index": 0,
                            "delta": {
                                "role": "assistant",
                                "tool_calls": [{"index": 0, **call}],
                            },
                            "finish_reason": None,
                        }
                    ],
                }
            )
        else:
            chunks.append(
                {
                    "id": completion["id"],
                    "object": "chat.completion.chunk",
                    "created": completion["created"],
                    "model": completion["model"],
                    "choices": [
                        {
                            "index": 0,
                            "delta": {"role": "assistant", "content": message.get("content", "")},
                            "finish_reason": None,
                        }
                    ],
                }
            )
        chunks.append(
            {
                "id": completion["id"],
                "object": "chat.completion.chunk",
                "created": completion["created"],
                "model": completion["model"],
                "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}],
            }
        )

        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.end_headers()
        for chunk in chunks:
            self.wfile.write(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode())
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()


if __name__ == "__main__":
    ThreadingHTTPServer(("0.0.0.0", 9000), StubHandler).serve_forever()  # noqa: S104