640 lines
15 KiB
Python
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()
|