394 lines
16 KiB
Python
394 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""Authenticated LAN bridge from CodexPi to a PC-local Ollama server."""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import hmac
|
||
import json
|
||
import os
|
||
import platform
|
||
import secrets
|
||
import shutil
|
||
import socket
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
from pathlib import Path
|
||
|
||
|
||
MAX_REQUEST_BYTES = 64 * 1024 * 1024
|
||
ALLOWED_GET = {"/api/tags", "/v1/models", "/api/experimental/model-recommendations"}
|
||
ALLOWED_POST = {"/api/show", "/api/chat", "/api/generate", "/api/pull", "/v1/chat/completions", "/v1/responses"}
|
||
|
||
|
||
def checked_url(value: str, *, local_only: bool = False) -> str:
|
||
parsed = urllib.parse.urlsplit(value)
|
||
if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password:
|
||
raise ValueError("Adresse HTTP invalide")
|
||
if local_only and parsed.hostname not in {"127.0.0.1", "localhost", "::1"}:
|
||
raise ValueError("Ollama doit rester attaché à l’adresse locale du PC")
|
||
return value.rstrip("/")
|
||
|
||
|
||
def json_request(server: str, path: str, body: dict) -> dict:
|
||
data = json.dumps(body, ensure_ascii=False).encode("utf-8")
|
||
request = urllib.request.Request(
|
||
checked_url(server) + path,
|
||
data=data,
|
||
method="POST",
|
||
headers={"Accept": "application/json", "Content-Type": "application/json", "User-Agent": "CodexPi-AI-Bridge/1"},
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=30) as response:
|
||
raw = response.read()
|
||
except urllib.error.HTTPError as exc:
|
||
try:
|
||
detail = json.loads(exc.read()).get("detail")
|
||
except Exception:
|
||
detail = None
|
||
raise RuntimeError(detail or f"Erreur CodexPi HTTP {exc.code}") from exc
|
||
value = json.loads(raw or b"{}")
|
||
if not isinstance(value, dict):
|
||
raise RuntimeError("Réponse CodexPi invalide")
|
||
return value
|
||
|
||
|
||
def ollama_models(upstream: str) -> list[dict]:
|
||
try:
|
||
with urllib.request.urlopen(checked_url(upstream, local_only=True) + "/api/tags", timeout=5) as response:
|
||
value = json.loads(response.read(2_000_000) or b"{}")
|
||
except (OSError, urllib.error.URLError, ValueError) as exc:
|
||
raise RuntimeError("Ollama ne répond pas sur ce PC. Ouvrez Ollama puis réessayez.") from exc
|
||
models = value.get("models") if isinstance(value, dict) else None
|
||
if not isinstance(models, list) or not models:
|
||
raise RuntimeError("Ollama fonctionne mais aucun modèle n’est installé")
|
||
return models
|
||
|
||
|
||
def hardware_info() -> dict:
|
||
"""Return only coarse sizing data needed to recommend a local model."""
|
||
memory_bytes = 0
|
||
try:
|
||
if platform.system() == "Windows":
|
||
import ctypes
|
||
|
||
class MemoryStatus(ctypes.Structure):
|
||
_fields_ = [
|
||
("length", ctypes.c_ulong),
|
||
("memory_load", ctypes.c_ulong),
|
||
("total_physical", ctypes.c_ulonglong),
|
||
("available_physical", ctypes.c_ulonglong),
|
||
("total_page_file", ctypes.c_ulonglong),
|
||
("available_page_file", ctypes.c_ulonglong),
|
||
("total_virtual", ctypes.c_ulonglong),
|
||
("available_virtual", ctypes.c_ulonglong),
|
||
("available_extended_virtual", ctypes.c_ulonglong),
|
||
]
|
||
|
||
status = MemoryStatus()
|
||
status.length = ctypes.sizeof(status)
|
||
if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
|
||
memory_bytes = int(status.total_physical)
|
||
elif hasattr(os, "sysconf"):
|
||
memory_bytes = int(os.sysconf("SC_PAGE_SIZE")) * int(os.sysconf("SC_PHYS_PAGES"))
|
||
except (AttributeError, OSError, TypeError, ValueError):
|
||
memory_bytes = 0
|
||
|
||
gpu_name = ""
|
||
gpu_vram_bytes = 0
|
||
executable = shutil.which("nvidia-smi")
|
||
if executable:
|
||
try:
|
||
result = subprocess.run(
|
||
[executable, "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
|
||
check=True,
|
||
capture_output=True,
|
||
text=True,
|
||
timeout=5,
|
||
)
|
||
candidates: list[tuple[int, str]] = []
|
||
for line in result.stdout.splitlines():
|
||
name, separator, memory = line.rpartition(",")
|
||
if separator:
|
||
candidates.append((int(memory.strip()) * 1024 * 1024, name.strip()[:120]))
|
||
if candidates:
|
||
gpu_vram_bytes, gpu_name = max(candidates)
|
||
except (OSError, subprocess.SubprocessError, ValueError):
|
||
pass
|
||
|
||
try:
|
||
disk = shutil.disk_usage(Path.home())
|
||
disk_free_bytes = int(disk.free)
|
||
except OSError:
|
||
disk_free_bytes = 0
|
||
return {
|
||
"platform": platform.system(),
|
||
"architecture": platform.machine(),
|
||
"memory_bytes": memory_bytes,
|
||
"disk_free_bytes": disk_free_bytes,
|
||
"gpu_name": gpu_name,
|
||
"gpu_vram_bytes": gpu_vram_bytes,
|
||
}
|
||
|
||
|
||
def save_config(path: Path, value: dict) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
descriptor, temporary = tempfile.mkstemp(prefix="ai-bridge-", suffix=".json", dir=path.parent)
|
||
try:
|
||
with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
|
||
json.dump(value, handle, ensure_ascii=False, indent=2)
|
||
handle.write("\n")
|
||
try:
|
||
os.chmod(temporary, 0o600)
|
||
except OSError:
|
||
pass
|
||
os.replace(temporary, path)
|
||
finally:
|
||
if os.path.exists(temporary):
|
||
os.unlink(temporary)
|
||
|
||
|
||
def discover_address(server: str) -> str:
|
||
parsed = urllib.parse.urlsplit(checked_url(server))
|
||
port = parsed.port or (443 if parsed.scheme == "https" else 80)
|
||
addresses = socket.getaddrinfo(parsed.hostname, port, type=socket.SOCK_DGRAM)
|
||
last_error: OSError | None = None
|
||
for family, socket_type, protocol, _, target in addresses:
|
||
probe = socket.socket(family, socket_type, protocol)
|
||
try:
|
||
probe.connect(target)
|
||
return str(probe.getsockname()[0])
|
||
except OSError as exc:
|
||
last_error = exc
|
||
finally:
|
||
probe.close()
|
||
raise RuntimeError(f"Impossible de déterminer l’adresse locale du PC : {last_error}")
|
||
|
||
|
||
def handler_for(upstream: str, token: str):
|
||
class BridgeHandler(BaseHTTPRequestHandler):
|
||
protocol_version = "HTTP/1.0"
|
||
server_version = "CodexPiAIBridge/1"
|
||
|
||
def log_message(self, fmt: str, *args) -> None:
|
||
print(f"{self.client_address[0]} · {fmt % args}")
|
||
|
||
def _authorized(self) -> bool:
|
||
scheme, _, supplied = self.headers.get("Authorization", "").partition(" ")
|
||
return scheme.lower() == "bearer" and hmac.compare_digest(supplied, token)
|
||
|
||
def _json_error(self, status: int, message: str) -> None:
|
||
body = json.dumps({"error": message}, ensure_ascii=False).encode("utf-8")
|
||
self.send_response(status)
|
||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.send_header("Cache-Control", "no-store")
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
|
||
def _path(self) -> str:
|
||
return urllib.parse.urlsplit(self.path).path
|
||
|
||
def _proxy(self, method: str) -> None:
|
||
if not self._authorized():
|
||
self._json_error(401, "Jeton du bridge IA invalide")
|
||
return
|
||
path = self._path()
|
||
allowed = ALLOWED_GET if method == "GET" else ALLOWED_POST
|
||
if path not in allowed:
|
||
self._json_error(404, "Route Ollama non autorisée")
|
||
return
|
||
length = int(self.headers.get("Content-Length") or 0)
|
||
if length < 0 or length > MAX_REQUEST_BYTES:
|
||
self._json_error(413, "Requête trop volumineuse")
|
||
return
|
||
data = self.rfile.read(length) if length else None
|
||
headers = {"Accept": self.headers.get("Accept", "application/json"), "User-Agent": "CodexPi-AI-Bridge/1"}
|
||
if data is not None:
|
||
headers["Content-Type"] = self.headers.get("Content-Type", "application/json")
|
||
request = urllib.request.Request(
|
||
upstream + self.path,
|
||
data=data,
|
||
method=method,
|
||
headers=headers,
|
||
)
|
||
try:
|
||
with urllib.request.urlopen(request, timeout=900) as response:
|
||
self.send_response(response.status)
|
||
self.send_header("Content-Type", response.headers.get("Content-Type", "application/json"))
|
||
self.send_header("Cache-Control", "no-store")
|
||
self.send_header("Connection", "close")
|
||
self.end_headers()
|
||
while True:
|
||
chunk = response.read(64 * 1024)
|
||
if not chunk:
|
||
break
|
||
self.wfile.write(chunk)
|
||
self.wfile.flush()
|
||
except urllib.error.HTTPError as exc:
|
||
body = exc.read(2_000_000)
|
||
self.send_response(exc.code)
|
||
self.send_header("Content-Type", exc.headers.get("Content-Type", "application/json"))
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
except (BrokenPipeError, ConnectionResetError):
|
||
return
|
||
except (OSError, urllib.error.URLError) as exc:
|
||
self._json_error(502, f"Ollama est indisponible : {exc}")
|
||
|
||
def do_GET(self) -> None:
|
||
if self._path() == "/codexpi/info":
|
||
if not self._authorized():
|
||
self._json_error(401, "Jeton du bridge IA invalide")
|
||
return
|
||
body = json.dumps(hardware_info(), ensure_ascii=False).encode("utf-8")
|
||
self.send_response(200)
|
||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.send_header("Cache-Control", "no-store")
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
return
|
||
if self._path() == "/health":
|
||
if not self._authorized():
|
||
self._json_error(401, "Jeton du bridge IA invalide")
|
||
return
|
||
try:
|
||
models = ollama_models(upstream)
|
||
except RuntimeError as exc:
|
||
self._json_error(503, str(exc))
|
||
return
|
||
body = json.dumps({"ok": True, "runtime": "ollama", "models": len(models)}).encode("utf-8")
|
||
self.send_response(200)
|
||
self.send_header("Content-Type", "application/json")
|
||
self.send_header("Content-Length", str(len(body)))
|
||
self.send_header("Cache-Control", "no-store")
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
return
|
||
self._proxy("GET")
|
||
|
||
def do_POST(self) -> None:
|
||
self._proxy("POST")
|
||
|
||
return BridgeHandler
|
||
|
||
|
||
def create_server(listen: str, port: int, upstream: str, token: str) -> ThreadingHTTPServer:
|
||
return ThreadingHTTPServer((listen, port), handler_for(checked_url(upstream, local_only=True), token))
|
||
|
||
|
||
def pair(
|
||
pairing_path: Path,
|
||
config_path: Path,
|
||
listen: str,
|
||
port: int,
|
||
upstream: str,
|
||
name: str,
|
||
advertise_address: str | None,
|
||
) -> tuple[dict, ThreadingHTTPServer]:
|
||
pairing = json.loads(pairing_path.read_text(encoding="utf-8"))
|
||
server_url = checked_url(str(pairing["server"]))
|
||
ollama_models(upstream)
|
||
token = secrets.token_urlsafe(48)
|
||
server = create_server(listen, port, upstream, token)
|
||
thread = threading.Thread(target=server.serve_forever, name="codexpi-ai-bridge", daemon=True)
|
||
thread.start()
|
||
address = advertise_address or discover_address(server_url)
|
||
try:
|
||
result = json_request(server_url, "/api/local-ai/agent/pair", {
|
||
"pairing_id": pairing["pairing_id"],
|
||
"code": pairing["code"],
|
||
"name": name,
|
||
"platform": platform.platform()[:160],
|
||
"address": address,
|
||
"port": port,
|
||
"token": token,
|
||
})
|
||
except Exception:
|
||
server.shutdown()
|
||
server.server_close()
|
||
raise
|
||
config = {
|
||
"listen": listen,
|
||
"port": port,
|
||
"upstream": checked_url(upstream, local_only=True),
|
||
"token": token,
|
||
"bridge_id": result.get("id"),
|
||
"codexpi": server_url,
|
||
}
|
||
save_config(config_path, config)
|
||
print(f"PC associé à CodexPi : {name}")
|
||
print(f"Adresse privée annoncée : http://{address}:{port}")
|
||
print(f"Modèle par défaut : {result.get('default_model') or 'à choisir dans CodexPi'}")
|
||
print(f"Configuration privée : {config_path}")
|
||
return config, server
|
||
|
||
|
||
def main() -> int:
|
||
default_config = Path.home() / ".codexpi" / "ai-bridge.json"
|
||
parser = argparse.ArgumentParser(description="Relie CodexPi à Ollama sur ce PC sans exposer l’API Ollama brute.")
|
||
parser.add_argument("--pairing", type=Path, help="fichier .codexpi-ai téléchargé depuis Comptes IA")
|
||
parser.add_argument("--config", type=Path, default=default_config, help="configuration privée du bridge")
|
||
parser.add_argument("--listen", default="0.0.0.0", help="adresse d’écoute du bridge")
|
||
parser.add_argument("--port", type=int, default=11435, help="port privé du bridge")
|
||
parser.add_argument("--upstream", default="http://127.0.0.1:11434", help="adresse locale d’Ollama")
|
||
parser.add_argument("--name", default=socket.gethostname(), help="nom affiché dans CodexPi")
|
||
parser.add_argument("--advertise-address", help="adresse IP privée du PC visible depuis le Pi")
|
||
args = parser.parse_args()
|
||
if not 1024 <= args.port <= 65535:
|
||
parser.error("--port doit être compris entre 1024 et 65535")
|
||
active_server: ThreadingHTTPServer | None = None
|
||
background_server = False
|
||
try:
|
||
if args.pairing:
|
||
_, active_server = pair(
|
||
args.pairing,
|
||
args.config,
|
||
args.listen,
|
||
args.port,
|
||
args.upstream,
|
||
args.name,
|
||
args.advertise_address,
|
||
)
|
||
background_server = True
|
||
else:
|
||
config = json.loads(args.config.read_text(encoding="utf-8"))
|
||
ollama_models(str(config["upstream"]))
|
||
active_server = create_server(
|
||
str(config.get("listen") or args.listen),
|
||
int(config.get("port") or args.port),
|
||
str(config["upstream"]),
|
||
str(config["token"]),
|
||
)
|
||
print(f"Bridge IA CodexPi prêt sur le port {config.get('port') or args.port}")
|
||
print("Laissez cette fenêtre ouverte. Ctrl+C arrête le bridge.")
|
||
if background_server:
|
||
stop = threading.Event()
|
||
while not stop.wait(3600):
|
||
pass
|
||
else:
|
||
active_server.serve_forever()
|
||
return 0
|
||
except KeyboardInterrupt:
|
||
print("\nBridge IA arrêté.")
|
||
return 0
|
||
except Exception as exc:
|
||
print(f"Erreur : {exc}", file=sys.stderr)
|
||
return 1
|
||
finally:
|
||
if active_server is not None:
|
||
if background_server:
|
||
active_server.shutdown()
|
||
active_server.server_close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|