72 lines
2.3 KiB
Python
Executable file
72 lines
2.3 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
"""Serial HTTP/SSE fixture for the authentic prefill-cancellation runner."""
|
|
|
|
import argparse
|
|
import http.server
|
|
import json
|
|
import time
|
|
|
|
|
|
class Handler(http.server.BaseHTTPRequestHandler):
|
|
cancellation_started = False
|
|
|
|
def do_GET(self):
|
|
if self.path != "/health":
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
return
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.end_headers()
|
|
self.wfile.write(b'{"status":"ok"}')
|
|
|
|
def do_POST(self):
|
|
if self.path != "/v1/chat/completions":
|
|
self.send_response(404)
|
|
self.end_headers()
|
|
return
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
try:
|
|
request = json.loads(self.rfile.read(length))
|
|
assert request["stream"] is True
|
|
except (AssertionError, KeyError, json.JSONDecodeError):
|
|
self.send_response(400)
|
|
self.end_headers()
|
|
return
|
|
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "text/event-stream")
|
|
self.send_header("Cache-Control", "no-cache")
|
|
self.end_headers()
|
|
if not Handler.cancellation_started:
|
|
Handler.cancellation_started = True
|
|
self.wfile.write(b": keepalive\n\n")
|
|
self.wfile.flush()
|
|
time.sleep(2.0)
|
|
try:
|
|
self.wfile.write(b": keepalive\n\n")
|
|
self.wfile.flush()
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
pass
|
|
return
|
|
|
|
time.sleep(0.05)
|
|
self.wfile.write(
|
|
b'data: {"choices":[{"delta":{"content":"OK"},'
|
|
b'"finish_reason":null}]}\n\n'
|
|
)
|
|
self.wfile.write(
|
|
b'data: {"choices":[{"delta":{},"finish_reason":"stop"}],'
|
|
b'"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}\n\n'
|
|
)
|
|
self.wfile.write(b"data: [DONE]\n\n")
|
|
self.wfile.flush()
|
|
|
|
def log_message(self, _format, *_args):
|
|
pass
|
|
|
|
|
|
parser = argparse.ArgumentParser(add_help=False)
|
|
parser.add_argument("--port", type=int, required=True)
|
|
arguments, _ = parser.parse_known_args()
|
|
http.server.HTTPServer(("127.0.0.1", arguments.port), Handler).serve_forever()
|