isokit/tools/mock_llm.py

"""A stand-in for an OpenAI-compatible model server, for the agent window tests.

    python tools/mock_llm.py [port]        (default 7932; base URL http://127.0.0.1:<port>/v1)

Wants the key `test-key`. Answers like a model that uses tools:
- a user message "plant a tree at X, Y" -> reasoning, then two tool calls:
  list_objects {type: tree} and add_object {type: tree, x: X, y: Y, name: "Agent tree"}
- tool results -> a short answer naming the new object's id in backticks
- a user message containing "slow" -> a long, slow answer (to test Stop)
- a user message "show me X" -> an answer and a select_object call for X in one step
  (that ends the turn: no request follows)
- anything else -> a plain answer
Streams (SSE) when asked to. Every page origin may call it (CORS, also the
preflight). GET /_requests lists the requests seen (headers and bodies).
"""
import json
import re
import sys
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

KEY = "test-key"
SEEN = []


class H(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def log_message(self, *a):
        pass

    def cors(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Headers", "authorization, content-type")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")

    def reply(self, code, obj):
        b = json.dumps(obj).encode()
        self.send_response(code)
        self.cors()
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(b)))
        self.end_headers()
        self.wfile.write(b)

    def authorized(self):
        if self.headers.get("Authorization") == "Bearer " + KEY:
            return True
        self.reply(401, {"error": {"message": "invalid api key"}})
        return False

    def do_OPTIONS(self):
        self.send_response(204)
        self.cors()
        self.send_header("Content-Length", "0")
        self.end_headers()

    def do_GET(self):
        if self.path == "/_requests":
            return self.reply(200, SEEN)
        if self.path == "/v1/models":
            if self.authorized():
                self.reply(200, {"object": "list", "data": [{"id": "mock-1", "object": "model"}, {"id": "mock-2", "object": "model"}]})
            return
        self.reply(404, {"error": {"message": "not found"}})

    def do_POST(self):
        n = int(self.headers.get("Content-Length") or 0)
        body = json.loads(self.rfile.read(n) or b"{}")
        SEEN.append({"path": self.path, "auth": self.headers.get("Authorization"), "body": body})
        if self.path != "/v1/chat/completions":
            return self.reply(404, {"error": {"message": "not found"}})
        if not self.authorized():
            return
        msgs = body.get("messages", [])
        last = msgs[-1] if msgs else {}
        deltas, finish = [], "stop"
        if last.get("role") == "user" and "slow" in last.get("content", ""):
            deltas = [{"content": "word%d " % i} for i in range(80)]
            return self.stream(body, deltas, finish, pause=0.1)
        show = re.search(r"show me (\S+)", last.get("content", "")) if last.get("role") == "user" else None
        if show:
            deltas = [{"content": "Here is `%s`." % show.group(1)},
                      {"tool_calls": [{"index": 0, "id": "call_s", "type": "function",
                                       "function": {"name": "select_object", "arguments": json.dumps({"id": show.group(1)})}}]}]
            return self.stream(body, deltas, "tool_calls")
        m = re.search(r"tree at (-?[\d.]+),\s*(-?[\d.]+)", last.get("content", "")) if last.get("role") == "user" else None
        if m:
            x, y = float(m.group(1)), float(m.group(2))
            deltas = [{"reasoning": "The user wants a tree. "}, {"reasoning": "I look at the trees first, then add one."},
                      {"tool_calls": [{"index": 0, "id": "call_a", "type": "function", "function": {"name": "list_objects", "arguments": ""}}]},
                      {"tool_calls": [{"index": 0, "function": {"arguments": '{"type": "tree"}'}}]},
                      {"tool_calls": [{"index": 1, "id": "call_b", "type": "function", "function": {"name": "add_object", "arguments": '{"type": "tree", '}}]},
                      {"tool_calls": [{"index": 1, "function": {"arguments": '"x": %s, "y": %s, "name": "Agent tree"}' % (x, y)}}]}]
            finish = "tool_calls"
        elif last.get("role") == "tool":
            new_id = "?"
            for mm in reversed(msgs):
                if mm.get("role") == "tool":
                    try:
                        new_id = json.loads(mm["content"]).get("id") or new_id
                    except (ValueError, AttributeError):
                        pass
                    if new_id != "?":
                        break
            deltas = [{"content": "Planted "}, {"content": "`%s` for you." % new_id}, {"content": "\n\n- one tree\n- **no** overlaps"}]
        else:
            deltas = [{"content": "Hello from the mock model."}]
        self.stream(body, deltas, finish)

    def stream(self, body, deltas, finish, pause=0.02):
        if not body.get("stream"):
            msg = {"role": "assistant", "content": "".join(d.get("content", "") for d in deltas)}
            return self.reply(200, {"choices": [{"index": 0, "message": msg, "finish_reason": finish}]})
        self.send_response(200)
        self.cors()
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("Connection", "close")
        self.end_headers()
        try:
            for d in [{"role": "assistant", "content": ""}] + deltas:
                ch = {"object": "chat.completion.chunk", "model": body.get("model") or "mock-1",
                      "choices": [{"index": 0, "delta": d, "finish_reason": None}]}
                self.wfile.write(("data: %s\n\n" % json.dumps(ch)).encode())
                self.wfile.flush()
                time.sleep(pause)
            end = {"object": "chat.completion.chunk", "choices": [{"index": 0, "delta": {}, "finish_reason": finish}]}
            self.wfile.write(("data: %s\n\ndata: [DONE]\n\n" % json.dumps(end)).encode())
            self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError):
            SEEN.append({"aborted": True})
        self.close_connection = True


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 7932
    print(f"mock model on http://127.0.0.1:{port}/v1", flush=True)
    ThreadingHTTPServer(("127.0.0.1", port), H).serve_forever()