aboutsummaryrefslogtreecommitdiff
path: root/tests/fake-model.py
diff options
context:
space:
mode:
Diffstat (limited to 'tests/fake-model.py')
-rwxr-xr-xtests/fake-model.py59
1 files changed, 59 insertions, 0 deletions
diff --git a/tests/fake-model.py b/tests/fake-model.py
new file mode 100755
index 0000000..a02b98c
--- /dev/null
+++ b/tests/fake-model.py
@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""A fake OpenAI-compatible model for tests/pi-live.sh: streams a short reply
+to /v1/chat/completions, after a delay taken from the last user message
+("slow 3" waits 3 seconds). Prints its port, then serves until killed."""
+import json
+import re
+import sys
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+
+def last_user_text(body):
+ for message in reversed(body.get("messages", [])):
+ if message.get("role") == "user":
+ content = message.get("content")
+ if isinstance(content, list):
+ return " ".join(part.get("text", "") for part in content if isinstance(part, dict))
+ return content or ""
+ return ""
+
+
+class Handler(BaseHTTPRequestHandler):
+ def log_message(self, *args):
+ pass
+
+ def do_GET(self):
+ self.send_response(200)
+ self.send_header("Content-Type", "application/json")
+ self.end_headers()
+ self.wfile.write(json.dumps({"object": "list", "data": [{"id": "fake-model", "object": "model"}]}).encode())
+
+ def do_POST(self):
+ body = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}")
+ text = last_user_text(body)
+ match = re.search(r"slow (\d+(?:\.\d+)?)", text)
+ time.sleep(float(match.group(1)) if match else 0.2)
+ self.send_response(200)
+ self.send_header("Content-Type", "text/event-stream")
+ self.send_header("Cache-Control", "no-cache")
+ self.end_headers()
+
+ def chunk(delta, finish=None, usage=None):
+ data = {"id": "fake", "object": "chat.completion.chunk", "created": int(time.time()), "model": "fake-model",
+ "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}
+ if usage:
+ data["usage"] = usage
+ self.wfile.write(f"data: {json.dumps(data)}\n\n".encode())
+ self.wfile.flush()
+
+ chunk({"role": "assistant", "content": "Done: "})
+ chunk({"content": "fake reply."})
+ chunk({}, "stop", {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13})
+ self.wfile.write(b"data: [DONE]\n\n")
+ self.wfile.flush()
+
+
+server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+print(server.server_address[1], flush=True)
+server.serve_forever()