SHA256
68 lines
1.9 KiB
Python
68 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Fake ollama for end-to-end tests of the sjcode daemon + client.
|
|
|
|
Serves a scripted conversation: every /api/generate call returns the next
|
|
reply from SCRIPT. Each prompt it receives is appended to prompts.log so
|
|
the test can assert what the client actually fed back to the model.
|
|
|
|
Usage: python3 fake_ollama.py PORT PROMPT_LOG
|
|
"""
|
|
import json
|
|
import sys
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
|
|
PORT = int(sys.argv[1])
|
|
LOG = sys.argv[2]
|
|
|
|
RUN_REPLY = (
|
|
"I will create the file.\n"
|
|
"```run\n"
|
|
"printf 'hi' > hello.txt\n"
|
|
"```\n"
|
|
)
|
|
READ_REPLY = (
|
|
"Checking it.\n"
|
|
"```sj:read\n"
|
|
'{"path": "hello.txt"}\n'
|
|
"```\n"
|
|
)
|
|
DONE_REPLY = "The file is in place.\nDone."
|
|
|
|
SCRIPT = [RUN_REPLY, READ_REPLY, DONE_REPLY]
|
|
calls = {"n": 0}
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def log_message(self, *args): # quiet
|
|
pass
|
|
|
|
def _json(self, obj):
|
|
body = json.dumps(obj).encode()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def do_GET(self):
|
|
if self.path == "/api/tags":
|
|
self._json({"models": [{"name": "fake-model"}]})
|
|
else:
|
|
self._json({"error": "not found"})
|
|
|
|
def do_POST(self):
|
|
if self.path != "/api/generate":
|
|
self._json({"error": "not found"})
|
|
return
|
|
length = int(self.headers.get("Content-Length", "0"))
|
|
payload = json.loads(self.rfile.read(length).decode())
|
|
with open(LOG, "a") as fh:
|
|
fh.write(json.dumps(payload.get("prompt", "")) + "\n")
|
|
i = min(calls["n"], len(SCRIPT) - 1)
|
|
calls["n"] += 1
|
|
self._json({"model": "fake-model", "response": SCRIPT[i], "done": True})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ThreadingHTTPServer(("127.0.0.1", PORT), Handler).serve_forever()
|