Initial commit: Boltzmann LLM proxy with compression + classifier middleware

- 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
2026-06-15 15:12:25 +02:00
commit c121428c58
8 changed files with 712 additions and 0 deletions
@@ -0,0 +1,379 @@
+#!/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()