423 lines
16 KiB
Python
Executable File
423 lines
16 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""AI Profile Router – OpenAI-kompatibler Proxy vor llama.cpp.
|
||
|
||
Leitet OpenAI-kompatible Requests transparent an den lokalen llama.cpp-Server
|
||
weiter (Streaming, Tool Calls, JSON) und schaltet zwischen drei festen
|
||
Profilen um:
|
||
|
||
Profil Kontext
|
||
------ --------
|
||
fast 73728
|
||
medium 94208
|
||
long 131072
|
||
|
||
Virtuelle Modelle: qwen-fast, qwen-medium, qwen-long
|
||
Kommandos: POST /fast, /medium, /long (Profilwechsel)
|
||
GET /status (Zustand)
|
||
|
||
Ein Profilwechsel führt PROFILE_SCRIPT <profil> aus (ohne Shell, feste
|
||
Argumente → keine Injection), wartet dann, bis llama.cpp wieder erreichbar
|
||
ist, und erst dann wird eine erfolgreiche Antwort geliefert bzw. der
|
||
Request weitergeleitet.
|
||
|
||
Nur Python-Standardbibliothek. Logging nach stdout (journald).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import threading
|
||
import time
|
||
import http.client
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Konfiguration (über Umgebungsvariablen, vgl. systemd-Unit)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
HOST = os.environ.get("ROUTER_HOST", "0.0.0.0")
|
||
PORT = int(os.environ.get("ROUTER_PORT", "8081"))
|
||
UPSTREAM_URL = os.environ.get("UPSTREAM_URL", "http://127.0.0.1:8080").rstrip("/")
|
||
PROFILE_SCRIPT = os.environ.get("PROFILE_SCRIPT", "/usr/local/bin/llama-profile")
|
||
PROFILE_DIR = os.environ.get(
|
||
"PROFILE_DIR", "/etc/systemd/system/mike-ai-llama-ui.service.d")
|
||
|
||
SWITCH_TIMEOUT = float(os.environ.get("SWITCH_TIMEOUT", "600")) # s, Warten auf llama.cpp
|
||
REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "600")) # s, Read-Timeout Upstream
|
||
CONNECT_TIMEOUT = float(os.environ.get("CONNECT_TIMEOUT", "10")) # s, Connect-Timeout
|
||
POLL_INTERVAL = float(os.environ.get("POLL_INTERVAL", "2")) # s, Polling-Intervall
|
||
|
||
PROFILES = {"fast": 73728, "medium": 94208, "long": 131072}
|
||
VIRTUAL_MODELS = {f"qwen-{name}": name for name in PROFILES}
|
||
|
||
log = logging.getLogger("ai-profile-router")
|
||
|
||
# Hop-by-hop-Header, die nicht an Upstream/Client weitergereicht werden.
|
||
HOP_BY_HOP = {
|
||
"host", "connection", "keep-alive", "proxy-authenticate",
|
||
"proxy-authorization", "te", "trailer", "transfer-encoding",
|
||
"upgrade", "content-length",
|
||
}
|
||
|
||
|
||
def _parse_upstream(url: str) -> tuple[str, int]:
|
||
"""'http://127.0.0.1:8080' -> ('127.0.0.1', 8080)"""
|
||
hostport = url.split("://", 1)[-1]
|
||
host, _, port = hostport.partition(":")
|
||
return host, int(port) if port else 80
|
||
|
||
|
||
UPSTREAM_HOST, UPSTREAM_PORT = _parse_upstream(UPSTREAM_URL)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Zustand
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class _State:
|
||
"""Gemeinsamer, thread-sicherer Zustand."""
|
||
lock = threading.Lock() # serialisiert Profilwechsel
|
||
switching: str | None = None # Profil, das gerade gewechselt wird
|
||
started = time.time()
|
||
|
||
|
||
STATE = _State()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Upstream (llama.cpp)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def upstream_status() -> dict:
|
||
"""Prüft llama.cpp: erreichbar? welches Modell? welcher Kontext?"""
|
||
try:
|
||
conn = http.client.HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT,
|
||
timeout=CONNECT_TIMEOUT)
|
||
conn.request("GET", "/v1/models")
|
||
resp = conn.getresponse()
|
||
data = json.loads(resp.read())
|
||
conn.close()
|
||
except (OSError, ValueError) as e:
|
||
return {"reachable": False, "error": str(e)}
|
||
models = data.get("data") or []
|
||
if not models:
|
||
return {"reachable": True, "model": None, "ctx": None}
|
||
m = models[0]
|
||
return {"reachable": True,
|
||
"model": m.get("id"),
|
||
"ctx": (m.get("meta") or {}).get("n_ctx")}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Profile
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _read(path: str) -> str:
|
||
with open(path, encoding="utf-8") as f:
|
||
return f.read().strip()
|
||
|
||
|
||
def current_profile() -> str | None:
|
||
"""Aktives Profil, ermittelt durch Vergleich der override.conf."""
|
||
try:
|
||
override = _read(os.path.join(PROFILE_DIR, "override.conf"))
|
||
except OSError:
|
||
return None
|
||
for name in PROFILES:
|
||
try:
|
||
ref = _read(os.path.join(PROFILE_DIR, f"profile-{name}.conf.disabled"))
|
||
except OSError:
|
||
continue
|
||
if override == ref:
|
||
return name
|
||
return None
|
||
|
||
|
||
def switch_profile(profile: str, implicit: bool = False) -> None:
|
||
"""Führt einen Profilwechsel aus und wartet, bis llama.cpp bereit ist.
|
||
|
||
Wirft RuntimeError, wenn der Wechsel nicht erfolgreich war.
|
||
|
||
implicit=True (ausgelöst durch ein virtuelles Modell in einem Chat-Request):
|
||
Wenn das Profil bereits aktiv ist, aber llama.cpp down ist, wird sofort
|
||
eine RuntimeError geworfen (kein stiller Neustart). Der Nutzer kann den
|
||
Neustart explizit über /<profil> anstoßen.
|
||
"""
|
||
if profile not in PROFILES:
|
||
raise ValueError(f"unbekanntes Profil: {profile!r} "
|
||
f"(erlaubt: {', '.join(PROFILES)})")
|
||
with STATE.lock:
|
||
STATE.switching = profile
|
||
try:
|
||
cur = current_profile()
|
||
up = upstream_status()
|
||
if cur == profile and up["reachable"]:
|
||
log.info("Profil %s ist bereits aktiv", profile)
|
||
return
|
||
if cur == profile and not up["reachable"] and implicit:
|
||
raise RuntimeError(
|
||
f"llama.cpp nicht erreichbar (Profil {profile} ist bereits "
|
||
f"aktiv; Neustart über /{profile})")
|
||
log.info("Profilwechsel: %s -> %s", cur, profile)
|
||
try:
|
||
proc = subprocess.run(
|
||
[PROFILE_SCRIPT, profile],
|
||
stdin=subprocess.DEVNULL,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
timeout=120,
|
||
)
|
||
out = proc.stdout.decode(errors="replace").strip()
|
||
if out:
|
||
log.info("llama-profile: %s", out[-500:])
|
||
if proc.returncode != 0:
|
||
# whiptail bricht das Skript ohne TTY ab – der Wechsel
|
||
# selbst (cp + systemctl restart) ist dann aber erledigt.
|
||
log.warning("llama-profile Exit-Code %d (ohne TTY erwartet)",
|
||
proc.returncode)
|
||
except subprocess.TimeoutExpired:
|
||
log.error("llama-profile hat 120 s überschritten")
|
||
if current_profile() != profile:
|
||
raise RuntimeError(
|
||
f"Profildatei wurde nicht gesetzt (erwartet: {profile})")
|
||
log.info("Warte, bis llama.cpp wieder erreichbar ist ...")
|
||
deadline = time.monotonic() + SWITCH_TIMEOUT
|
||
while True:
|
||
status = upstream_status()
|
||
if status["reachable"]:
|
||
log.info("llama.cpp bereit: Profil=%s Modell=%s ctx=%s",
|
||
profile, status.get("model"), status.get("ctx"))
|
||
return
|
||
if time.monotonic() > deadline:
|
||
raise RuntimeError(
|
||
f"llama.cpp nach {SWITCH_TIMEOUT:.0f} s nicht erreichbar")
|
||
time.sleep(POLL_INTERVAL)
|
||
finally:
|
||
STATE.switching = None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# HTTP-Handler
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
server_version = "AIProfileRouter/1.0"
|
||
timeout = 60 # Socket-Timeout für Client-Requests (s)
|
||
|
||
# ---------- Routing ----------
|
||
|
||
def do_GET(self):
|
||
self._route()
|
||
|
||
def do_POST(self):
|
||
self._route()
|
||
|
||
def _route(self):
|
||
path = self.path.split("?", 1)[0]
|
||
started = time.monotonic()
|
||
try:
|
||
if path == "/v1/models" and self.command == "GET":
|
||
self._send_json(200, self._models_payload())
|
||
elif path == "/status":
|
||
self._send_json(200, self._status_payload())
|
||
elif path in ("/fast", "/medium", "/long"):
|
||
self._switch(path[1:])
|
||
elif (self.command == "POST" and path.startswith("/")
|
||
and path.count("/") == 1):
|
||
# Kommandonamensraum: unbekanntes Profil
|
||
self._send_error(400, f"unbekanntes Profil: {path[1:]}",
|
||
"invalid_request_error", "invalid_profile")
|
||
else:
|
||
self._forward()
|
||
except BrokenPipeError:
|
||
log.warning("Client getrennt: %s %s", self.command, path)
|
||
except Exception:
|
||
log.exception("Fehler bei %s %s", self.command, path)
|
||
self._safe_error(500, "interner Router-Fehler")
|
||
finally:
|
||
log.info("%s %s -> %s in %.3f s", self.command, path,
|
||
getattr(self, "_last_code", "-"), time.monotonic() - started)
|
||
|
||
# ---------- Router-eigene Endpunkte ----------
|
||
|
||
@staticmethod
|
||
def _models_payload() -> dict:
|
||
return {
|
||
"object": "list",
|
||
"data": [
|
||
{
|
||
"id": f"qwen-{name}",
|
||
"object": "model",
|
||
"created": 0,
|
||
"owned_by": "ai-profile-router",
|
||
"context_length": ctx,
|
||
"context_window": ctx,
|
||
}
|
||
for name, ctx in PROFILES.items()
|
||
],
|
||
}
|
||
|
||
def _status_payload(self) -> dict:
|
||
up = upstream_status()
|
||
return {
|
||
"router": "ai-profile-router",
|
||
"uptime_seconds": round(time.time() - STATE.started, 1),
|
||
"current_profile": current_profile(),
|
||
"switching": STATE.switching,
|
||
"profiles": PROFILES,
|
||
"upstream": {
|
||
"url": UPSTREAM_URL,
|
||
"reachable": up["reachable"],
|
||
"model": up.get("model"),
|
||
"ctx": up.get("ctx"),
|
||
},
|
||
}
|
||
|
||
def _switch(self, profile: str) -> None:
|
||
if profile not in PROFILES:
|
||
self._send_error(400, f"unbekanntes Profil: {profile}",
|
||
"invalid_request_error", "invalid_profile")
|
||
return
|
||
try:
|
||
switch_profile(profile)
|
||
except (ValueError, RuntimeError) as e:
|
||
self._send_error(503, str(e), "server_error", "profile_switch_failed")
|
||
return
|
||
up = upstream_status()
|
||
self._send_json(200, {
|
||
"status": "ok",
|
||
"profile": profile,
|
||
"context_length": PROFILES[profile],
|
||
"model": up.get("model"),
|
||
})
|
||
|
||
# ---------- Transparentes Forwarding ----------
|
||
|
||
def _forward(self) -> None:
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
body = self.rfile.read(length) if length > 0 else None
|
||
|
||
# Virtuelles Modell? -> Profil sicherstellen, dann Modell ersetzen.
|
||
if body is not None and self.path.startswith("/v1/"):
|
||
try:
|
||
data = json.loads(body)
|
||
except ValueError:
|
||
data = None
|
||
model = data.get("model") if isinstance(data, dict) else None
|
||
if isinstance(model, str) and model in VIRTUAL_MODELS:
|
||
profile = VIRTUAL_MODELS[model]
|
||
try:
|
||
switch_profile(profile, implicit=True)
|
||
except (ValueError, RuntimeError) as e:
|
||
self._send_error(502, str(e), "server_error",
|
||
"upstream_unavailable")
|
||
return
|
||
up = upstream_status()
|
||
if not up["reachable"] or not up.get("model"):
|
||
self._send_error(502, "llama.cpp nicht erreichbar",
|
||
"server_error", "upstream_unavailable")
|
||
return
|
||
data["model"] = up["model"]
|
||
body = json.dumps(data).encode()
|
||
elif isinstance(model, str) and model.startswith("qwen-"):
|
||
# qwen-* ist der Namensraum des Routers
|
||
self._send_error(400, f"unbekanntes virtuelles Modell: {model}",
|
||
"invalid_request_error", "unknown_model")
|
||
return
|
||
|
||
# An llama.cpp weiterleiten (Streaming bleibt erhalten).
|
||
try:
|
||
conn = http.client.HTTPConnection(UPSTREAM_HOST, UPSTREAM_PORT,
|
||
timeout=CONNECT_TIMEOUT)
|
||
conn.connect()
|
||
conn.sock.settimeout(REQUEST_TIMEOUT)
|
||
headers = {k: v for k, v in self.headers.items()
|
||
if k.lower() not in HOP_BY_HOP}
|
||
conn.request(self.command, self.path, body=body, headers=headers)
|
||
resp = conn.getresponse()
|
||
except (OSError, http.client.HTTPException) as e:
|
||
self._send_error(502, f"llama.cpp nicht erreichbar: {e}",
|
||
"server_error", "upstream_unavailable")
|
||
return
|
||
|
||
self._last_code = resp.status
|
||
self.send_response(resp.status)
|
||
for k, v in resp.getheaders():
|
||
if k.lower() not in HOP_BY_HOP:
|
||
self.send_header(k, v)
|
||
self.send_header("Connection", "close")
|
||
self.end_headers()
|
||
try:
|
||
while True:
|
||
chunk = resp.read(16384)
|
||
if not chunk:
|
||
break
|
||
self.wfile.write(chunk)
|
||
self.wfile.flush()
|
||
except (OSError, http.client.HTTPException) as e:
|
||
log.warning("Upstream-Stream abgebrochen: %s", e)
|
||
finally:
|
||
conn.close()
|
||
|
||
# ---------- Antworten ----------
|
||
|
||
def _send_json(self, code: int, payload: dict) -> None:
|
||
body = json.dumps(payload).encode()
|
||
self._last_code = code
|
||
self.send_response(code)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.send_header("Connection", "close")
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
|
||
def _send_error(self, code: int, message: str, etype: str, ecode: str) -> None:
|
||
# OpenAI-kompatibles Fehlerformat
|
||
self._send_json(code, {"error": {"message": message,
|
||
"type": etype,
|
||
"code": ecode}})
|
||
|
||
def _safe_error(self, code: int, message: str) -> None:
|
||
try:
|
||
self._send_error(code, message, "server_error", "internal_error")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class _FlushHandler(logging.StreamHandler):
|
||
"""StreamHandler, der nach jedem Record flusht (journald)."""
|
||
|
||
def emit(self, record):
|
||
super().emit(record)
|
||
self.flush()
|
||
|
||
|
||
def main() -> None:
|
||
handler = _FlushHandler(sys.stdout)
|
||
handler.setFormatter(logging.Formatter(
|
||
"%(asctime)s %(levelname)s %(message)s"))
|
||
logging.basicConfig(level=os.environ.get("LOG_LEVEL", "INFO"),
|
||
handlers=[handler])
|
||
log.info("AI Profile Router startet: %s:%s -> %s (Profile: %s)",
|
||
HOST, PORT, UPSTREAM_URL, ", ".join(PROFILES))
|
||
server = ThreadingHTTPServer((HOST, PORT), Handler)
|
||
server.daemon_threads = True
|
||
try:
|
||
server.serve_forever()
|
||
except KeyboardInterrupt:
|
||
pass
|
||
finally:
|
||
server.server_close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|