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
This commit is contained in:
@@ -0,0 +1,5 @@
|
||||
*.gguf
|
||||
*.log
|
||||
llamafile
|
||||
*.bak-*
|
||||
__pycache__/
|
||||
@@ -0,0 +1,106 @@
|
||||
#!/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()
|
||||
@@ -0,0 +1,176 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Prompt compression service using XLM-RoBERTa-large (LLMLingua-2).
|
||||
|
||||
Scoring approach: runs the classification head on each token, drops tokens
|
||||
with low keep-probability. The compressed text is reconstructed from the
|
||||
kept token spans.
|
||||
|
||||
Endpoints:
|
||||
POST /compress {"messages": [...], "rate": 0.5} → {"compressed": [...], "stats": {...}}
|
||||
GET /health → {"status": "ok"}
|
||||
"""
|
||||
import http.server
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import re
|
||||
import math
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
from transformers import AutoTokenizer, AutoModelForTokenClassification
|
||||
|
||||
MODEL_NAME = os.getenv("COMPRESS_MODEL", "microsoft/llmlingua-2-xlm-roberta-large-meetingbank")
|
||||
LISTEN_PORT = int(os.getenv("COMPRESS_LISTEN_PORT", "8091"))
|
||||
DEVICE = "cpu"
|
||||
DTYPE = torch.float16
|
||||
|
||||
|
||||
def log(msg):
|
||||
sys.stderr.write(f"{time.strftime('%H:%M:%S')} {msg}\n")
|
||||
sys.stderr.flush()
|
||||
|
||||
|
||||
# ── model loading ─────────────────────────────────────────────────
|
||||
|
||||
log(f"loading model {MODEL_NAME} ...")
|
||||
t0 = time.time()
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
|
||||
model = AutoModelForTokenClassification.from_pretrained(
|
||||
MODEL_NAME, torch_dtype=DTYPE
|
||||
).to(DEVICE).eval()
|
||||
log(f"model loaded in {time.time()-t0:.1f}s ({sum(p.numel() for p in model.parameters())/1e6:.0f}M params)")
|
||||
|
||||
|
||||
# ── compression core ──────────────────────────────────────────────
|
||||
|
||||
def _merge_token_spans(text, keep_mask, tokens):
|
||||
"""Reconstruct text from kept tokens, merging subwords cleanly."""
|
||||
kept = []
|
||||
for i, (tok, keep) in enumerate(zip(tokens, keep_mask)):
|
||||
if not keep:
|
||||
continue
|
||||
piece = tok
|
||||
# Remove prefix space indicator for merged tokens
|
||||
if i > 0 and piece.startswith("##"):
|
||||
piece = piece[2:]
|
||||
elif i > 0 and not piece.startswith(" ") and not re.match(r'^[^\w]', piece):
|
||||
piece = " " + piece
|
||||
kept.append(piece)
|
||||
return "".join(kept).strip()
|
||||
|
||||
|
||||
def compress_text(text, rate=0.5):
|
||||
"""Compress a single text string, targeting `rate` compression ratio (0-1)."""
|
||||
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=2048).to(DEVICE)
|
||||
input_ids = inputs["input_ids"][0]
|
||||
n_tokens = len(input_ids)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(**inputs)
|
||||
logits = outputs.logits[0] # (seq_len, num_labels)
|
||||
|
||||
# For binary classification: label 0 = keep, 1 = drop
|
||||
keep_probs = torch.softmax(logits, dim=-1)[:, 0] # probability of keep
|
||||
probs = keep_probs.cpu().numpy()
|
||||
|
||||
# Don't drop special tokens ([CLS], [SEP], [PAD])
|
||||
special_ids = {tokenizer.cls_token_id, tokenizer.sep_token_id,
|
||||
tokenizer.pad_token_id, tokenizer.bos_token_id,
|
||||
tokenizer.eos_token_id, tokenizer.unk_token_id,
|
||||
0} # padding
|
||||
is_special = [id.item() in special_ids for id in input_ids]
|
||||
|
||||
# Target: keep rate * (1-rate) tokens (rate=0.5 means keep half)
|
||||
n_to_keep = max(1, int(n_tokens * (1 - rate)))
|
||||
|
||||
# Mask: force-keep special tokens, then keep top-k by probability
|
||||
keep = np.zeros(n_tokens, dtype=bool)
|
||||
for i in range(n_tokens):
|
||||
if is_special[i]:
|
||||
keep[i] = True
|
||||
|
||||
n_kept_special = keep.sum()
|
||||
n_remaining = n_to_keep - n_kept_special
|
||||
|
||||
if n_remaining > 0:
|
||||
# Get indices of non-special tokens sorted by keep probability
|
||||
non_special_idx = [i for i in range(n_tokens) if not is_special[i]]
|
||||
sorted_idx = sorted(non_special_idx, key=lambda i: probs[i], reverse=True)
|
||||
for i in sorted_idx[:n_remaining]:
|
||||
keep[i] = True
|
||||
|
||||
# Decode kept tokens
|
||||
tokens = tokenizer.convert_ids_to_tokens(input_ids)
|
||||
compressed = _merge_token_spans(text, keep, tokens)
|
||||
|
||||
return {
|
||||
"compressed": compressed,
|
||||
"original_tokens": int(n_tokens),
|
||||
"compressed_tokens": int(keep.sum()),
|
||||
"ratio": float(keep.sum() / n_tokens),
|
||||
}
|
||||
|
||||
|
||||
def compress_messages(messages, rate=0.5):
|
||||
"""Compress an array of chat messages. Drops low-info tokens from each."""
|
||||
compressed = []
|
||||
for msg in messages:
|
||||
role = msg.get("role", "user")
|
||||
content = msg.get("content", "")
|
||||
if content:
|
||||
result = compress_text(content, rate)
|
||||
compressed.append({"role": role, "content": result["compressed"]})
|
||||
else:
|
||||
compressed.append(msg)
|
||||
return compressed
|
||||
|
||||
|
||||
# ── HTTP handler ─────────────────────────────────────────────────
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
server_version = "compress-server/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, default=str).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 == "/compress":
|
||||
body = self._read_body()
|
||||
messages = body.get("messages", [])
|
||||
rate = float(body.get("rate", 0.5))
|
||||
if not messages:
|
||||
return self._json_response(400, {"error": "missing messages"})
|
||||
t0 = time.time()
|
||||
compressed = compress_messages(messages, rate)
|
||||
elapsed = time.time() - t0
|
||||
log(f"compress {len(messages)} msgs rate={rate} in {elapsed:.2f}s")
|
||||
return self._json_response(200, {
|
||||
"compressed": compressed,
|
||||
"rate": rate,
|
||||
"elapsed_s": round(elapsed, 2),
|
||||
})
|
||||
self._json_response(404, {"error": "not found"})
|
||||
|
||||
def log_message(self, fmt, *args):
|
||||
return
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
log(f"compress-server on :{LISTEN_PORT}")
|
||||
http.server.ThreadingHTTPServer(("0.0.0.0", LISTEN_PORT), Handler).serve_forever()
|
||||
Executable
+22
@@ -0,0 +1,22 @@
|
||||
#!/bin/bash
|
||||
set -e
|
||||
cd ~/npu
|
||||
URLS=(
|
||||
"Codestral-22B-v0.1-Q4_K_M.gguf|https://huggingface.co/bartowski/Codestral-22B-v0.1-GGUF/resolve/main/Codestral-22B-v0.1-Q4_K_M.gguf"
|
||||
"gpt-oss-20b-mxfp4.gguf|https://huggingface.co/ggml-org/gpt-oss-20b-GGUF/resolve/main/gpt-oss-20b-mxfp4.gguf"
|
||||
"DeepSeek-Coder-V2-Lite-Instruct-Q4_K_M.gguf|https://huggingface.co/bartowski/DeepSeek-Coder-V2-Lite-Instruct-GGUF/resolve/main/DeepSeek-Coder-V2-Lite-Instruct-Q4_K_M.gguf"
|
||||
"starcoder2-15b-instruct-v0.1-Q4_K_M.gguf|https://huggingface.co/bartowski/starcoder2-15b-instruct-v0.1-GGUF/resolve/main/starcoder2-15b-instruct-v0.1-Q4_K_M.gguf"
|
||||
)
|
||||
for entry in "${URLS[@]}"; do
|
||||
fname="${entry%%|*}"
|
||||
url="${entry##*|}"
|
||||
if [[ -f "$fname" && $(stat -c%s "$fname") -gt 1000000000 ]]; then
|
||||
echo "[$(date +%H:%M:%S)] SKIP $fname (already $(du -h "$fname" | cut -f1))"
|
||||
continue
|
||||
fi
|
||||
echo "[$(date +%H:%M:%S)] START $fname"
|
||||
wget -c -O "$fname" "$url" 2>&1 | tail -5
|
||||
echo "[$(date +%H:%M:%S)] DONE $fname ($(du -h "$fname" | cut -f1))"
|
||||
done
|
||||
echo "[$(date +%H:%M:%S)] ALL DONE"
|
||||
ls -lh ~/npu/*.gguf
|
||||
Executable
+12
@@ -0,0 +1,12 @@
|
||||
#!/bin/bash
|
||||
# Manual launcher for the NPU Qwen 3.5 9B llama-server backend.
|
||||
# Used for ad-hoc testing; the systemd user unit
|
||||
# llama-server-qwen35-npu.service is the production path.
|
||||
set -e
|
||||
ulimit -n 65536
|
||||
exec /home/mfritsche/src/rk-llama.cpp/build/bin/llama-server \
|
||||
-m /home/mfritsche/models/Qwen3.5-9B-Q8_0.gguf \
|
||||
--alias qwen3.5-9b-npu \
|
||||
--reasoning off \
|
||||
-c 4096 -t 4 \
|
||||
--host 0.0.0.0 --port 8086
|
||||
Executable
+379
@@ -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()
|
||||
Executable
+6
@@ -0,0 +1,6 @@
|
||||
#!/bin/bash
|
||||
pkill -f llm-proxy.py 2>/dev/null
|
||||
sleep 1
|
||||
cd /home/mfritsche/npu
|
||||
nohup python3 llm-proxy.py > proxy.log 2>&1 &
|
||||
echo "proxy started PID $!"
|
||||
Reference in New Issue
Block a user