136 lines
4.8 KiB
Python
136 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Force a real Pi HTTP provider response to prove response-hook timing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import threading
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from pathlib import Path
|
|
|
|
from pi_gate0_run import PiRpc, jsonl
|
|
|
|
HERE = Path(__file__).resolve().parent
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
protocol_version = "HTTP/1.1"
|
|
|
|
def log_message(self, _format: str, *_args: object) -> None:
|
|
return
|
|
|
|
def do_POST(self) -> None: # noqa: N802
|
|
length = int(self.headers.get("content-length", "0"))
|
|
self.rfile.read(length)
|
|
chunks = [
|
|
{
|
|
"id": "gate0-response",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gate0-model",
|
|
"choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
|
|
},
|
|
{
|
|
"id": "gate0-response",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gate0-model",
|
|
"choices": [
|
|
{"index": 0, "delta": {"content": "TIMING_OK"}, "finish_reason": None}
|
|
],
|
|
},
|
|
{
|
|
"id": "gate0-response",
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": "gate0-model",
|
|
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12},
|
|
},
|
|
]
|
|
body = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + "data: [DONE]\n\n"
|
|
encoded = body.encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "text/event-stream")
|
|
self.send_header("Content-Length", str(len(encoded)))
|
|
self.send_header("X-Gate0-Response", "headers-before-stream")
|
|
self.end_headers()
|
|
self.wfile.write(encoded)
|
|
self.wfile.flush()
|
|
|
|
|
|
def main() -> None:
|
|
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
|
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
thread.start()
|
|
port = server.server_address[1]
|
|
|
|
with tempfile.TemporaryDirectory(prefix="gate0-p2-timing-") as temp:
|
|
root = Path(temp)
|
|
workspace = root / "workspace"
|
|
workspace.mkdir()
|
|
log = root / "hooks.jsonl"
|
|
env = os.environ.copy()
|
|
env.update(
|
|
{
|
|
"GATE0_PI_LOG": str(log),
|
|
"GATE0_LOCAL_PROVIDER_URL": f"http://127.0.0.1:{port}/v1",
|
|
"MOSAIC_PI_FORCE_SKILLS": "",
|
|
"PI_SKIP_VERSION_CHECK": "1",
|
|
}
|
|
)
|
|
command = [
|
|
"mosaic",
|
|
"yolo",
|
|
"pi",
|
|
"--mode",
|
|
"rpc",
|
|
"--no-session",
|
|
"--no-extensions",
|
|
"--no-context-files",
|
|
"--no-prompt-templates",
|
|
"--provider",
|
|
"gate0-local",
|
|
"--model",
|
|
"gate0-model",
|
|
"--extension",
|
|
str(HERE / "pi_gate0_extension.ts"),
|
|
]
|
|
pi = PiRpc(command, workspace, env)
|
|
try:
|
|
pi.prompt_and_settle("timing", "Reply with TIMING_OK")
|
|
records = jsonl(log)
|
|
selected = [
|
|
record
|
|
for record in records
|
|
if record["event"] in {"before_provider_request", "after_provider_response", "message_end"}
|
|
and (record["event"] != "message_end" or record.get("role") == "assistant")
|
|
]
|
|
print("$ python3 docs/compaction-refresh/probes/p2_provider_timing_run.py")
|
|
print(f"local_http_endpoint=http://127.0.0.1:{port}/v1/chat/completions")
|
|
for record in selected:
|
|
print(json.dumps(record, sort_keys=True))
|
|
after = next(record for record in selected if record["event"] == "after_provider_response")
|
|
message = next(record for record in selected if record["event"] == "message_end")
|
|
if not (
|
|
after["seq"] < message["seq"]
|
|
and after["assistantContentAvailableAtThisHook"] is False
|
|
and message["assistantContentObserved"] is True
|
|
):
|
|
raise AssertionError("provider response/content observation ordering failed")
|
|
print("machine_assertions=PASS")
|
|
print(f"after_provider_response_seq={after['seq']}")
|
|
print(f"message_end_seq={message['seq']}")
|
|
print(f"headers_hook_precedes_completed_message={after['seq'] < message['seq']}")
|
|
finally:
|
|
pi.close()
|
|
|
|
server.shutdown()
|
|
server.server_close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|