Files
z-bar-qt/scripts/mock-llama-server.py
T
zach 7bc1db7f19
C++ / fmt (pull_request) Successful in 4s
C++ / build (pull_request) Failing after 9s
JS/TS / fmt (pull_request) Successful in 8s
C++ / clang-tidy (pull_request) Failing after 26s
JS/TS / lint (pull_request) Successful in 8s
Python / static (pull_request) Successful in 26s
Rust / build (pull_request) Successful in 48s
Rust / fmt (pull_request) Successful in 28s
Python / verify (pull_request) Successful in 1m40s
Rust / clippy (pull_request) Successful in 45s
fix: format and apply lint fixes in python script
2026-09-19 10:13:46 +02:00

640 lines
15 KiB
Python

import argparse
import contextlib
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
MODEL_ID = "mock-model"
CTX_SIZE = 8192
_host = "127.0.0.1"
_port = 8080
_state_lock = threading.Lock()
_state = {
"loading": False,
"no_models": False,
"chat_mode": None,
}
def _now():
return int(time.time())
def _chunk(cid, delta, finish_reason=None):
return {
"id": cid,
"object": "chat.completion.chunk",
"created": _now(),
"model": MODEL_ID,
"system_fingerprint": "mock",
"choices": [
{
"index": 0,
"delta": delta,
"finish_reason": finish_reason,
}
],
}
def _usage_chunk(cid, prompt_tokens, completion_tokens):
return {
"id": cid,
"object": "chat.completion.chunk",
"created": _now(),
"model": MODEL_ID,
"choices": [],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
}
def _error(code, message, etype):
return {"error": {"code": code, "message": message, "type": etype}}
def _sse(obj):
return ("data: " + json.dumps(obj, separators=(",", ":")) + "\n\n").encode()
def _sse_done():
return b"data: [DONE]\n\n"
def _behaviour_for(text):
t = text.lower()
if "instant error" in t:
return "instant_error"
if "context overflow" in t:
return "context_overflow"
if "error during" in t:
return "error_during"
if "no done" in t or "truncate" in t:
return "no_done"
if "200 error" in t:
return "200_error"
if "two tool calls" in t or "two tool call" in t:
return "tool_call_duo"
if "some tool calls" in t:
return "tool_call_some"
if "many tool calls" in t:
return "tool_call_many"
if "unknown tool" in t:
return "tool_call_unknown"
if "tool call" in t:
return "tool_call"
return "normal"
_TOOL_MODES = (
"tool_call",
"tool_call_duo",
"tool_call_some",
"tool_call_many",
"tool_call_unknown",
)
def _tool_call_delta(index, call_id=None, name=None, arguments=None):
call = {"index": index}
if call_id is not None:
call["id"] = call_id
if name is not None or arguments is not None:
function = {}
if name is not None:
function["name"] = name
if arguments is not None:
function["arguments"] = arguments
call["function"] = function
return {"tool_calls": [call]}
def _stream_tool_calls(w, cid, reasoning, calls):
w(_sse(_chunk(cid, {"role": "assistant"})))
if reasoning:
for piece in reasoning:
w(_sse(_chunk(cid, {"reasoning_content": piece})))
time.sleep(0.05)
for i, (call_id, name, arguments) in enumerate(calls):
w(_sse(_chunk(cid, _tool_call_delta(i, call_id=call_id, name=name))))
time.sleep(0.05)
for fragment in arguments:
w(_sse(_chunk(cid, _tool_call_delta(i, arguments=fragment))))
time.sleep(0.05)
w(_sse(_chunk(cid, {}, finish_reason="tool_calls")))
w(_sse(_usage_chunk(cid, 24, len(reasoning or []) + len(calls) * 2)))
w(_sse_done())
def _tool_result_snippets(messages, limit=160):
snippets = []
for msg in messages or []:
if msg.get("role") != "tool":
continue
content = msg.get("content")
if isinstance(content, list):
content = " ".join(
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
)
flat = " ".join(str(content or "").split())
snippets.append(flat[:limit] + ("..." if len(flat) > limit else ""))
return snippets
def _stream_tool_round_done(w, cid, messages):
w(_sse(_chunk(cid, {"role": "assistant"})))
parts = ["mock:", "tool", "round", "complete."]
for i, snippet in enumerate(_tool_result_snippets(messages), 1):
parts.append(f"result {i}:")
parts.append(snippet or "(empty)")
for word in parts:
w(_sse(_chunk(cid, {"content": word + " "})))
time.sleep(0.02)
w(_sse(_chunk(cid, {}, finish_reason="stop")))
w(_sse(_usage_chunk(cid, 40, len(parts) + 1)))
w(_sse_done())
def _stream_normal(w, cid, words):
w(_sse(_chunk(cid, {"role": "assistant"})))
for word in words:
w(_sse(_chunk(cid, {"content": word + " "})))
time.sleep(0.05)
w(_sse(_chunk(cid, {}, finish_reason="stop")))
w(_sse(_usage_chunk(cid, 12, len(words) + 1)))
w(_sse_done())
def _stream_error_during(w, cid, words):
w(_sse(_chunk(cid, {"role": "assistant"})))
for word in words:
w(_sse(_chunk(cid, {"content": word + " "})))
time.sleep(0.05)
w(_sse(_error(500, "Mock: generation failed mid-stream", "server_error")))
def _stream_no_done(w, cid, words):
w(_sse(_chunk(cid, {"role": "assistant"})))
for word in words:
w(_sse(_chunk(cid, {"content": word + " "})))
time.sleep(0.05)
def _last_user_text(messages):
for msg in reversed(messages or []):
if msg.get("role") == "user":
content = msg.get("content")
if isinstance(content, list):
parts = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
parts.append(part.get("text", ""))
content = " ".join(parts)
return content or ""
return ""
class Handler(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
server_version = "MockLlama/1.0"
def _send_json(self, status, obj):
body = json.dumps(obj, separators=(",", ":")).encode()
self.send_response(status)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.send_header("Connection", "close")
self.end_headers()
self.wfile.write(body)
def _send_loading_error(self):
self._send_json(503, _error(503, "Loading model", "unavailable_error"))
def _not_ready(self):
with _state_lock:
return _state["loading"]
def _stream(self, status, writer):
self.send_response(status)
self.send_header("Content-Type", "text/event-stream")
self.send_header("Cache-Control", "no-cache")
self.send_header("Connection", "close")
self.end_headers()
try:
writer(self.wfile.write)
except (BrokenPipeError, ConnectionResetError):
pass
finally:
with contextlib.suppress(BrokenPipeError, ConnectionResetError):
self.wfile.flush()
def do_GET(self):
path = self.path.split("?", 1)[0].rstrip("/") or "/"
if path in ("/health", "/v1/health"):
if self._not_ready():
return self._send_loading_error()
return self._send_json(200, {"status": "ok"})
if path in ("/props", "/v1/props"):
if self._not_ready():
return self._send_loading_error()
return self._send_json(
200,
{
"total_slots": 1,
"model_alias": MODEL_ID,
"model_path": "/mock/model.gguf",
"default_generation_settings": {
"params": {"n_predict": 2048},
"n_ctx": CTX_SIZE,
},
"is_sleeping": False,
},
)
if path in ("/models", "/v1/models"):
if self._not_ready():
return self._send_loading_error()
with _state_lock:
no_models = _state["no_models"]
data = (
[]
if no_models
else [
{
"id": MODEL_ID,
"object": "model",
"created": _now(),
"owned_by": "mock",
}
]
)
return self._send_json(200, {"object": "list", "data": data})
if path == "/":
return self._send_json(200, {"status": "ok", "mock": True})
return self._send_json(404, _error(404, "not found", "not_found_error"))
def do_POST(self):
path = self.path.split("?", 1)[0].rstrip("/") or "/"
length = int(self.headers.get("Content-Length") or 0)
raw = self.rfile.read(length) if length else b""
try:
body = json.loads(raw) if raw else {}
except json.JSONDecodeError:
body = {}
if path == "/_mock":
return self._mock_control(body)
if path in ("/chat/completions", "/v1/chat/completions"):
if self._not_ready():
return self._send_loading_error()
return self._chat(body)
return self._send_json(404, _error(404, "not found", "not_found_error"))
def _mock_control(self, body):
with _state_lock:
if body.get("reset"):
_state["loading"] = False
_state["no_models"] = False
_state["chat_mode"] = None
if "loading" in body:
_state["loading"] = bool(body["loading"])
if "no_models" in body:
_state["no_models"] = bool(body["no_models"])
if "chat_mode" in body:
_state["chat_mode"] = body["chat_mode"]
snapshot = dict(_state)
self.log_message("control -> %s", json.dumps(snapshot))
return self._send_json(200, {"ok": True, "state": snapshot})
def _chat(self, body):
text = _last_user_text(body.get("messages"))
with _state_lock:
forced = _state["chat_mode"]
mode = forced or _behaviour_for(text)
stream = bool(body.get("stream", True))
cid = "chatcmpl-mock"
words = ["mock", "says", "hello", "from", "the", "mock", "server"]
self.log_message("chat mode=%s stream=%s text=%r", mode, stream, text)
if mode == "instant_error":
return self._send_json(
500, _error(500, "Mock: instant server error", "server_error")
)
if mode == "context_overflow":
return self._send_json(
400,
_error(
400,
"request (100000 tokens) exceeds the available context "
"size (2048 tokens), try increasing it",
"exceed_context_size_error",
),
)
if mode == "200_error":
return self._send_json(
200, _error(400, "Mock: 200-with-json-error", "server_error")
)
if mode in _TOOL_MODES:
if any(m.get("role") == "tool" for m in body.get("messages") or []):
return self._stream(
200,
lambda w: _stream_tool_round_done(
w, cid, body.get("messages")
),
)
if mode == "tool_call":
calls = [
(
"call_mock_1",
"readfile",
['{"path": "/etc/', 'os-release"}'],
)
]
reasoning = [
"The user wants a tool call. ",
"I will read /etc/os-release.",
]
elif mode == "tool_call_duo":
calls = [
(
"call_mock_1",
"readfile",
['{"path": "/etc/', 'hostname"}'],
),
(
"call_mock_2",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
]
reasoning = [
"Two independent tool calls, ",
"batched together.",
]
elif mode == "tool_call_some":
calls = [
(
"call_mock_1",
"readfile",
['{"path": "/etc/', 'hostname"}'],
),
(
"call_mock_2",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_3",
"readfile",
['{"path": "/etc/', 'os-release"}'],
),
(
"call_mock_4",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_5",
"readfile",
['{"path": "/proc/', 'uptime"}'],
),
(
"call_mock_6",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
]
reasoning = [
"Two independent tool calls, ",
"batched together.",
]
elif mode == "tool_call_many":
calls = [
(
"call_mock_1",
"readfile",
['{"path": "/etc/', 'hostname"}'],
),
(
"call_mock_2",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_3",
"readfile",
['{"path": "/etc/', 'os-release"}'],
),
(
"call_mock_4",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_5",
"readfile",
['{"path": "/proc/', 'uptime"}'],
),
(
"call_mock_6",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_1",
"readfile",
['{"path": "/etc/', 'hostname"}'],
),
(
"call_mock_2",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_3",
"readfile",
['{"path": "/etc/', 'os-release"}'],
),
(
"call_mock_4",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_5",
"readfile",
['{"path": "/proc/', 'uptime"}'],
),
(
"call_mock_6",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_1",
"readfile",
['{"path": "/etc/', 'hostname"}'],
),
(
"call_mock_2",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_3",
"readfile",
['{"path": "/etc/', 'os-release"}'],
),
(
"call_mock_4",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
(
"call_mock_5",
"readfile",
['{"path": "/proc/', 'uptime"}'],
),
(
"call_mock_6",
"webfetch",
[
f'{{"url": "http://{_host}:{_port}/", ',
'"format": "text"}',
],
),
]
reasoning = [
"Two independent tool calls, ",
"batched together.",
]
else:
calls = [
(
"call_mock_3",
"nonexistent_tool",
['{"target": "test"}'],
)
]
reasoning = ["Calling a tool the client ", "does not have."]
self.log_message(
"tool mode=%s calls=%s", mode, [c[1] for c in calls]
)
return self._stream(
200, lambda w: _stream_tool_calls(w, cid, reasoning, calls)
)
if not stream:
content = " ".join(words)
return self._send_json(
200,
{
"id": cid,
"object": "chat.completion",
"created": _now(),
"model": MODEL_ID,
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": content,
},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 12,
"completion_tokens": len(words),
"total_tokens": 12 + len(words),
},
},
)
if mode == "error_during":
return self._stream(
200, lambda w: _stream_error_during(w, cid, words)
)
if mode == "no_done":
return self._stream(200, lambda w: _stream_no_done(w, cid, words))
return self._stream(200, lambda w: _stream_normal(w, cid, words))
def log_message(self, format, *args):
print(f"{self.log_date_time_string()} - {format % args}", flush=True)
def main():
ap = argparse.ArgumentParser(
description="Mock llama-server for exercising zshell LLM error handling"
)
ap.add_argument("--port", type=int, default=8080)
ap.add_argument("--host", default="127.0.0.1")
args = ap.parse_args()
global _host, _port
_host = args.host
_port = args.port
server = ThreadingHTTPServer((args.host, args.port), Handler)
server.daemon_threads = True
print(
f"mock llama-server on http://{args.host}:{args.port} "
f"(model={MODEL_ID}, ctx={CTX_SIZE})\n"
"type trigger words in the chat UI, or curl /_mock to change state",
flush=True,
)
with contextlib.suppress(KeyboardInterrupt):
server.serve_forever()
if __name__ == "__main__":
main()