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
380 lines
14 KiB
Python
Executable File
380 lines
14 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Boltzmann-local LLM aggregator with compression + classifier middleware.
|
|
|
|
Behaviour (unchanged from v2):
|
|
- GET /v1/models : merged model list from every reachable backend
|
|
- POST /v1/chat/completions : routes by model name; compresses gemma4 messages
|
|
- POST /v1/completions : same routing rule
|
|
|
|
New:
|
|
- POST /v1/classify : calls local Llama 3.2 1B classifier
|
|
- POST /v1/compress : token-budget compression of chat messages
|
|
"""
|
|
import http.server
|
|
import http.client
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
BACKENDS = [
|
|
("gemma4-12b-q4km", "127.0.0.1", 8087),
|
|
("classifier-1b", "127.0.0.1", 8089),
|
|
("coder-1.5b", "127.0.0.1", 8081),
|
|
("llama-3.1", "127.0.0.1", 8083),
|
|
("nemo-12b", "127.0.0.1", 8084),
|
|
("qwen3-30b-a3b", "127.0.0.1", 8085),
|
|
("qwen3.5-9b-npu", "127.0.0.1", 8086),
|
|
]
|
|
CONNECT_TIMEOUT = 30
|
|
READ_TIMEOUT = 1800
|
|
LISTEN_PORT = int(os.getenv("LLM_PROXY_PORT", "8082"))
|
|
CLASSIFIER_URL = os.getenv("CLASSIFIER_URL", "http://127.0.0.1:8090/classify")
|
|
COMPRESS_URL = os.getenv("COMPRESS_URL", "http://127.0.0.1:8091/compress")
|
|
|
|
# Gemma 4 12B context limit - reserve 4096 for completion
|
|
GEMMA_MAX_INPUT_TOKENS = 65536 - 4096
|
|
|
|
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.stdout.write(f"{time.strftime('%H:%M:%S')} {msg}\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
# ── helpers ────────────────────────────────────────────────────────
|
|
|
|
def _simple_token_estimate(text):
|
|
"""Rough token estimate: 4 chars per token (standard heuristic)."""
|
|
return len(text) // 4
|
|
|
|
|
|
def compress_messages(messages, max_tokens=GEMMA_MAX_INPUT_TOKENS):
|
|
"""Drop oldest non-system messages to stay within token budget."""
|
|
if not messages:
|
|
return messages
|
|
|
|
system_msgs = [m for m in messages if m.get("role") == "system"]
|
|
history = [m for m in messages if m.get("role") != "system"]
|
|
|
|
# Count current tokens
|
|
total = sum(_simple_token_estimate(json.dumps(m)) for m in messages)
|
|
|
|
if total <= max_tokens:
|
|
return messages
|
|
|
|
# Drop oldest history messages, keep most recent
|
|
while history and total > max_tokens:
|
|
dropped = history.pop(0)
|
|
total -= _simple_token_estimate(json.dumps(dropped))
|
|
log(f"compress: dropped {dropped.get('role','?')} msg ({total} est. tokens remain)")
|
|
|
|
compressed = system_msgs + history
|
|
return compressed if compressed else messages
|
|
|
|
|
|
def classify_query(query):
|
|
"""Send query to the local classifier service. Returns tier string."""
|
|
body = json.dumps({"query": query}).encode()
|
|
req = urllib.request.Request(
|
|
CLASSIFIER_URL, data=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
try:
|
|
resp = urllib.request.urlopen(req, timeout=15)
|
|
data = json.loads(resp.read())
|
|
return data.get("tier", "medium")
|
|
except (urllib.error.URLError, json.JSONDecodeError, OSError) as e:
|
|
log(f"classifier call failed: {e}")
|
|
return "medium"
|
|
|
|
|
|
def compress_via_service(messages, rate=0.4):
|
|
"""Send messages to the ML compression service. Returns compressed messages or None."""
|
|
body = json.dumps({"messages": messages, "rate": rate}).encode()
|
|
req = urllib.request.Request(
|
|
COMPRESS_URL, data=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
try:
|
|
resp = urllib.request.urlopen(req, timeout=5) # fast timeout, fall back to truncation
|
|
data = json.loads(resp.read())
|
|
return data.get("compressed")
|
|
except (urllib.error.URLError, json.JSONDecodeError, OSError) as e:
|
|
log(f"compress service call failed: {e}")
|
|
return None
|
|
|
|
|
|
def classify_local(query):
|
|
"""Fallback: classify directly against Llama 3.2 1B if classifier service is unavailable."""
|
|
url = "http://127.0.0.1:8089/v1/chat/completions"
|
|
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})
|
|
body = json.dumps({
|
|
"model": "llama-3.2-1b-classifier",
|
|
"messages": msgs,
|
|
"max_tokens": 5,
|
|
"temperature": 0,
|
|
}).encode()
|
|
req = urllib.request.Request(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()
|
|
return content if content in ("low", "medium", "high") else "medium"
|
|
except Exception as e:
|
|
log(f"direct classify failed: {e}")
|
|
return "medium"
|
|
|
|
|
|
# ── model discovery ────────────────────────────────────────────────
|
|
|
|
def fetch_models(backend):
|
|
name, host, port = backend
|
|
try:
|
|
c = http.client.HTTPConnection(host, port, timeout=CONNECT_TIMEOUT)
|
|
c.request("GET", "/v1/models")
|
|
r = c.getresponse()
|
|
if r.status != 200:
|
|
return None
|
|
data = json.loads(r.read())
|
|
c.close()
|
|
ids = []
|
|
for m in data.get("data", []) or data.get("models", []):
|
|
mid = m.get("id") or m.get("name") or m.get("model")
|
|
if mid:
|
|
ids.append(mid)
|
|
return ids
|
|
except (OSError, ValueError, http.client.HTTPException):
|
|
return None
|
|
|
|
|
|
def discover():
|
|
out = []
|
|
for b in BACKENDS:
|
|
ids = fetch_models(b)
|
|
if ids is not None:
|
|
out.append((b, ids))
|
|
return out
|
|
|
|
|
|
def find_backend_for_model(model_id):
|
|
for b in BACKENDS:
|
|
ids = fetch_models(b)
|
|
if ids and model_id in ids:
|
|
return b
|
|
return None
|
|
|
|
|
|
# ── HTTP handler ──────────────────────────────────────────────────
|
|
|
|
class Handler(http.server.BaseHTTPRequestHandler):
|
|
server_version = "boltzmann-llm-proxy/3"
|
|
|
|
def _read_body(self):
|
|
n = int(self.headers.get("Content-Length", 0))
|
|
return self.rfile.read(n) if n else b""
|
|
|
|
def _json_body(self):
|
|
try:
|
|
return json.loads(self._read_body())
|
|
except (ValueError, TypeError):
|
|
return {}
|
|
|
|
def _fwd_headers(self, body_len):
|
|
h = {}
|
|
for k in ("Content-Type", "Authorization", "Accept", "Accept-Encoding"):
|
|
v = self.headers.get(k)
|
|
if v:
|
|
h[k] = v
|
|
if body_len:
|
|
h["Content-Length"] = str(body_len)
|
|
return h
|
|
|
|
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 _proxy_to(self, backend, method, body, extra_headers=None):
|
|
name, host, port = backend
|
|
try:
|
|
headers = self._fwd_headers(len(body))
|
|
if extra_headers:
|
|
headers.update(extra_headers)
|
|
c = http.client.HTTPConnection(host, port, timeout=READ_TIMEOUT)
|
|
c.request(method, self.path, body=body, headers=headers)
|
|
r = c.getresponse()
|
|
data = r.read()
|
|
self.send_response(r.status)
|
|
for hname in ("Content-Type", "Content-Encoding", "Cache-Control", "ETag", "Last-Modified"):
|
|
v = r.getheader(hname)
|
|
if v:
|
|
self.send_header(hname, v)
|
|
self.send_header("Content-Length", str(len(data)))
|
|
self.send_header("X-Backend", name)
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
c.close()
|
|
log(f"{method} {self.path} -> {name} ({r.status})")
|
|
except (OSError, http.client.HTTPException) as e:
|
|
self.send_response(502)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.send_header("X-Backend", name)
|
|
self.end_headers()
|
|
self.wfile.write(json.dumps({"error": f"{name} unreachable: {e}"}).encode())
|
|
log(f"{method} {self.path} -> {name} FAILED ({e})")
|
|
|
|
# ── GET ──
|
|
|
|
def do_GET(self):
|
|
path = self.path.rstrip("/")
|
|
if path == "/v1/models":
|
|
merged = []
|
|
seen = set()
|
|
for backend, ids in discover():
|
|
name = backend[0]
|
|
for mid in ids:
|
|
if mid in seen:
|
|
continue
|
|
seen.add(mid)
|
|
merged.append({
|
|
"id": mid,
|
|
"object": "model",
|
|
"owned_by": f"boltzmann/{name}",
|
|
})
|
|
self._json_response(200, {"object": "list", "data": merged})
|
|
log(f"GET /v1/models -> aggregator ({len(merged)} models)")
|
|
return
|
|
|
|
if path == "/health":
|
|
self._json_response(200, {"status": "ok", "version": 3})
|
|
return
|
|
|
|
for b in BACKENDS:
|
|
if fetch_models(b) is not None:
|
|
return self._proxy_to(b, "GET", b"")
|
|
self._json_response(502, {"error": "no backend up"})
|
|
|
|
# ── POST ──
|
|
|
|
def do_POST(self):
|
|
path = self.path.rstrip("/")
|
|
|
|
# ── Classification endpoint ──
|
|
if path == "/v1/classify":
|
|
body = self._json_body()
|
|
query = body.get("query", "")
|
|
if not query:
|
|
return self._json_response(400, {"error": "missing query"})
|
|
tier = classify_local(query)
|
|
log(f"classify {query[:60]!r} -> {tier}")
|
|
return self._json_response(200, {"tier": tier})
|
|
|
|
# ── Compression endpoint ──
|
|
if path == "/v1/compress":
|
|
body = self._json_body()
|
|
messages = body.get("messages", [])
|
|
max_tokens = body.get("max_tokens", GEMMA_MAX_INPUT_TOKENS)
|
|
compressed = compress_messages(messages, max_tokens)
|
|
ratio = 0
|
|
if messages:
|
|
orig = sum(_simple_token_estimate(json.dumps(m)) for m in messages)
|
|
new = sum(_simple_token_estimate(json.dumps(m)) for m in compressed)
|
|
ratio = round((1 - new / orig) * 100, 1) if orig else 0
|
|
return self._json_response(200, {
|
|
"messages": compressed,
|
|
"original_tokens": orig if messages else 0,
|
|
"compressed_tokens": sum(_simple_token_estimate(json.dumps(m)) for m in compressed),
|
|
"ratio_pct": ratio,
|
|
})
|
|
|
|
# ── Chat / Completions ──
|
|
if path in ("/v1/chat/completions", "/v1/completions"):
|
|
body = self._read_body()
|
|
model = None
|
|
try:
|
|
model = json.loads(body).get("model")
|
|
except (ValueError, TypeError):
|
|
pass
|
|
|
|
target = None
|
|
if model:
|
|
target = find_backend_for_model(model)
|
|
if target is None:
|
|
available = sorted({m for b in BACKENDS for m in (fetch_models(b) or [])})
|
|
return self._json_response(404, {"error": {
|
|
"message": f"model '{model}' not available on this aggregator",
|
|
"type": "model_not_found",
|
|
"available": available,
|
|
}})
|
|
|
|
if target is None:
|
|
for b in BACKENDS:
|
|
if fetch_models(b) is not None:
|
|
target = b
|
|
break
|
|
if target is None:
|
|
return self._json_response(502, {"error": "no backend up"})
|
|
|
|
# Compression: try ML service (fast timeout), fall back to token-budget truncation
|
|
compressed_body = body
|
|
if path == "/v1/chat/completions" and "gemma4" in (model or ""):
|
|
try:
|
|
req = json.loads(body)
|
|
if "messages" in req and len(req["messages"]) > 2:
|
|
orig_n = len(req["messages"])
|
|
orig_t = sum(_simple_token_estimate(json.dumps(m)) for m in req["messages"])
|
|
compressed = compress_via_service(req["messages"], rate=0.4)
|
|
if compressed:
|
|
req["messages"] = compressed
|
|
compressed_body = json.dumps(req).encode()
|
|
new_t = sum(_simple_token_estimate(json.dumps(m)) for m in compressed)
|
|
log(f"ML compress {model}: ~{orig_t}->{new_t} tokens")
|
|
else:
|
|
truncated = compress_messages(req["messages"], GEMMA_MAX_INPUT_TOKENS)
|
|
if len(truncated) < len(req["messages"]):
|
|
req["messages"] = truncated
|
|
compressed_body = json.dumps(req).encode()
|
|
log(f"truncate compress {model}: {orig_n}->{len(truncated)} msgs")
|
|
except (ValueError, TypeError, json.JSONDecodeError) as e:
|
|
log(f"compression skipped (parse error): {e}")
|
|
|
|
return self._proxy_to(target, "POST", compressed_body)
|
|
|
|
# ── Anything else: forward to default backend ──
|
|
body = self._read_body()
|
|
for b in BACKENDS:
|
|
if fetch_models(b) is not None:
|
|
return self._proxy_to(b, "POST" if body else "GET", body)
|
|
self._json_response(502, {"error": "no backend up"})
|
|
|
|
def log_message(self, fmt, *args):
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
log(f"boltzmann-llm-proxy v3 on :{LISTEN_PORT}")
|
|
log(f" classifier: {CLASSIFIER_URL}")
|
|
log(f" backends: {', '.join(n for n, h, p in BACKENDS)}")
|
|
http.server.ThreadingHTTPServer(("0.0.0.0", LISTEN_PORT), Handler).serve_forever()
|