Markus Fritsche
c121428c58
- llm-proxy.py v4: aggregator with classify/compress endpoints, gemma4 backend - classifier-server.py: Llama 3.2 1B query complexity classifier - compress-server.py: token-budget compression middleware - start-proxy.sh / start.sh: launcher scripts
107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Query classifier service wrapping Llama 3.2 1B on :8089.
|
|
|
|
Endpoints:
|
|
POST /classify {"query": str} → {"tier": "low"|"medium"|"high"}
|
|
GET /health → {"status": "ok"}
|
|
"""
|
|
import http.server
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
CLASSIFIER_URL = os.getenv("CLASSIFIER_URL", "http://127.0.0.1:8089/v1/chat/completions")
|
|
LISTEN_PORT = int(os.getenv("CLASSIFIER_LISTEN_PORT", "8090"))
|
|
|
|
FEW_SHOT_EXAMPLES = [
|
|
("What is 2+2?", "low"),
|
|
("List the files in the current directory", "low"),
|
|
("What is the weather today?", "low"),
|
|
("Write a Python function to sort a list of dictionaries by a key", "medium"),
|
|
("Explain how garbage collection works in Go", "medium"),
|
|
("Debug this error: TypeError: 'NoneType' object is not subscriptable", "medium"),
|
|
("Deploy the new release to production", "high"),
|
|
("Roll back the last database migration", "high"),
|
|
("Commit and push all changes to main branch", "high"),
|
|
]
|
|
|
|
|
|
def log(msg):
|
|
sys.stderr.write(f"{time.strftime('%H:%M:%S')} {msg}\n")
|
|
sys.stderr.flush()
|
|
|
|
|
|
def build_messages(query):
|
|
msgs = [{"role": "system", "content": "Classify queries as low, medium, or high."}]
|
|
for q, a in FEW_SHOT_EXAMPLES:
|
|
msgs.append({"role": "user", "content": q})
|
|
msgs.append({"role": "assistant", "content": a})
|
|
msgs.append({"role": "user", "content": query})
|
|
return msgs
|
|
|
|
|
|
def classify(query):
|
|
body = json.dumps({
|
|
"model": "llama-3.2-1b-classifier",
|
|
"messages": build_messages(query),
|
|
"max_tokens": 5,
|
|
"temperature": 0,
|
|
}).encode()
|
|
req = urllib.request.Request(CLASSIFIER_URL, data=body, headers={
|
|
"Content-Type": "application/json",
|
|
})
|
|
try:
|
|
resp = urllib.request.urlopen(req, timeout=30)
|
|
data = json.loads(resp.read())
|
|
content = data["choices"][0]["message"]["content"].strip().lower()
|
|
if content in ("low", "medium", "high"):
|
|
return content
|
|
log(f"unexpected classifier output: {content!r} for query {query!r}")
|
|
return "medium"
|
|
except (urllib.error.URLError, json.JSONDecodeError, KeyError, IndexError) as e:
|
|
log(f"classifier error: {e}")
|
|
return "medium"
|
|
|
|
|
|
class Handler(http.server.BaseHTTPRequestHandler):
|
|
server_version = "classifier/1"
|
|
|
|
def _read_body(self):
|
|
n = int(self.headers.get("Content-Length", 0))
|
|
return json.loads(self.rfile.read(n)) if n else {}
|
|
|
|
def _json_response(self, code, data):
|
|
payload = json.dumps(data).encode()
|
|
self.send_response(code)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("Content-Length", str(len(payload)))
|
|
self.end_headers()
|
|
self.wfile.write(payload)
|
|
|
|
def do_GET(self):
|
|
if self.path == "/health":
|
|
return self._json_response(200, {"status": "ok"})
|
|
self._json_response(404, {"error": "not found"})
|
|
|
|
def do_POST(self):
|
|
if self.path == "/classify":
|
|
body = self._read_body()
|
|
query = body.get("query", "")
|
|
if not query:
|
|
return self._json_response(400, {"error": "missing query"})
|
|
tier = classify(query)
|
|
log(f"classify {query[:60]!r} -> {tier}")
|
|
return self._json_response(200, {"tier": tier})
|
|
self._json_response(404, {"error": "not found"})
|
|
|
|
def log_message(self, fmt, *args):
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
log(f"classifier-server on :{LISTEN_PORT}, backend {CLASSIFIER_URL}")
|
|
http.server.ThreadingHTTPServer(("0.0.0.0", LISTEN_PORT), Handler).serve_forever()
|