"""Pluggable inference client: vllm | ollama | mock backends. All backends implement the same contract: health() -> BackendStatus (never raises) generate(prompt, system, max_tokens) -> InferenceResult with token counts and wall-clock latency validate() -> deterministic test-prompt round trip The mock backend is the default for tests: fully deterministic, no network. The ollama backend follows the platform's existing local pattern (sensorflow/evaluation/copilot.py). The vllm backend targets any OpenAI-compatible endpoint and is only exercisable on CUDA/ROCm hosts — see sensorflow/retro/inference/vllm_server/. """ from __future__ import annotations import hashlib import os import time from abc import ABC, abstractmethod from typing import List, Optional import httpx from pydantic import BaseModel, Field VALIDATION_PROMPT = "Respond with exactly the token RETRO-OK and nothing else." VALIDATION_EXPECT = "RETRO-OK" DEFAULT_OLLAMA_URL = os.environ.get("RETRO_OLLAMA_URL", "http://localhost:11434") DEFAULT_VLLM_URL = os.environ.get("RETRO_VLLM_URL", "http://localhost:8001/v1") class InferenceResult(BaseModel): backend: str model: str text: str prompt_tokens: int completion_tokens: int latency_ms: float notes: str = "" class BackendStatus(BaseModel): backend: str available: bool endpoint: Optional[str] = None model: Optional[str] = None detail: str = "" health_latency_ms: Optional[float] = None class ValidationResult(BaseModel): backend: str passed: bool expected: str got: str latency_ms: float prompt_tokens: int completion_tokens: int detail: str = "" class InferenceBackend(ABC): name: str = "abstract" @abstractmethod def health(self) -> BackendStatus: ... @abstractmethod def generate(self, prompt: str, system: Optional[str] = None, max_tokens: int = 1024) -> InferenceResult: ... def validate(self) -> ValidationResult: """Deterministic test-prompt inference validation with latency.""" t0 = time.perf_counter() try: res = self.generate(VALIDATION_PROMPT, max_tokens=16) except Exception as exc: return ValidationResult(backend=self.name, passed=False, expected=VALIDATION_EXPECT, got="", latency_ms=(time.perf_counter() - t0) * 1000, prompt_tokens=0, completion_tokens=0, detail=f"generate failed: {exc}") passed = VALIDATION_EXPECT in res.text return ValidationResult(backend=self.name, passed=passed, expected=VALIDATION_EXPECT, got=res.text.strip()[:200], latency_ms=res.latency_ms, prompt_tokens=res.prompt_tokens, completion_tokens=res.completion_tokens, detail="" if passed else "model reply did not contain the expected token") def _approx_tokens(text: str) -> int: """Deterministic whitespace-ish token approximation (mock backend only).""" return max(1, len(text.split())) if text else 0 class MockBackend(InferenceBackend): """Deterministic, offline. Replies are pure functions of the prompt. The retrospective agent does NOT parse mock replies for analysis content — in mock mode the analysis itself is produced by the deterministic scripted path (sensorflow/retro/agent). The mock backend exists to satisfy the client contract (health/generate/validate) without any model. """ name = "mock" def health(self) -> BackendStatus: return BackendStatus(backend="mock", available=True, endpoint=None, model="deterministic-mock", detail="always available; replies are deterministic " "functions of the prompt", health_latency_ms=0.0) def generate(self, prompt: str, system: Optional[str] = None, max_tokens: int = 1024) -> InferenceResult: t0 = time.perf_counter() if VALIDATION_EXPECT in prompt or "RETRO-OK" in prompt: text = VALIDATION_EXPECT else: digest = hashlib.sha256(((system or "") + "\n" + prompt).encode()).hexdigest()[:12] text = f"[mock-response {digest}] Deterministic placeholder reply." return InferenceResult( backend="mock", model="deterministic-mock", text=text, prompt_tokens=_approx_tokens((system or "") + " " + prompt), completion_tokens=_approx_tokens(text), latency_ms=(time.perf_counter() - t0) * 1000, notes="deterministic mock; token counts are whitespace approximations") class OllamaBackend(InferenceBackend): """Local Ollama server (Metal/CPU on this machine — NOT vLLM numbers).""" name = "ollama" def __init__(self, base_url: str = DEFAULT_OLLAMA_URL, model: Optional[str] = None, timeout: float = 120.0): self.base_url = base_url.rstrip("/") self.model = model or os.environ.get("RETRO_OLLAMA_MODEL", "") self.timeout = timeout def _pick_model(self) -> Optional[str]: if self.model: return self.model try: res = httpx.get(f"{self.base_url}/api/tags", timeout=3.0) models = [m.get("name") for m in res.json().get("models", [])] self.model = models[0] if models else None except Exception: self.model = None return self.model def health(self) -> BackendStatus: t0 = time.perf_counter() try: res = httpx.get(f"{self.base_url}/api/tags", timeout=3.0) ms = (time.perf_counter() - t0) * 1000 if res.status_code == 200: models = [m.get("name") for m in res.json().get("models", [])] return BackendStatus(backend="ollama", available=True, endpoint=self.base_url, model=self.model or (models[0] if models else None), detail=f"{len(models)} model(s) installed", health_latency_ms=round(ms, 1)) return BackendStatus(backend="ollama", available=False, endpoint=self.base_url, detail=f"HTTP {res.status_code} from /api/tags") except Exception as exc: return BackendStatus(backend="ollama", available=False, endpoint=self.base_url, detail=f"unreachable: {exc.__class__.__name__}: {exc}") def generate(self, prompt: str, system: Optional[str] = None, max_tokens: int = 1024) -> InferenceResult: model = self._pick_model() if not model: raise RuntimeError(f"no Ollama model available at {self.base_url}") messages = ([{"role": "system", "content": system}] if system else []) + [ {"role": "user", "content": prompt}] t0 = time.perf_counter() res = httpx.post(f"{self.base_url}/api/chat", json={"model": model, "messages": messages, "stream": False, "options": {"num_predict": max_tokens}}, timeout=self.timeout) latency_ms = (time.perf_counter() - t0) * 1000 res.raise_for_status() body = res.json() return InferenceResult( backend="ollama", model=model, text=body.get("message", {}).get("content", ""), prompt_tokens=int(body.get("prompt_eval_count", 0)), completion_tokens=int(body.get("eval_count", 0)), latency_ms=round(latency_ms, 1), notes="Ollama on local CPU/Metal — these are NOT vLLM/GPU numbers") class VLLMBackend(InferenceBackend): """OpenAI-compatible vLLM endpoint. Only runnable against a CUDA/ROCm host (see vllm_server/); on this macOS machine health() honestly reports it unavailable unless a remote endpoint is configured.""" name = "vllm" def __init__(self, base_url: str = DEFAULT_VLLM_URL, model: Optional[str] = None, timeout: float = 120.0): self.base_url = base_url.rstrip("/") self.model = model or os.environ.get("RETRO_VLLM_MODEL", "") self.timeout = timeout def health(self) -> BackendStatus: t0 = time.perf_counter() try: res = httpx.get(f"{self.base_url}/models", timeout=3.0) ms = (time.perf_counter() - t0) * 1000 if res.status_code == 200: models = [m.get("id") for m in res.json().get("data", [])] return BackendStatus(backend="vllm", available=True, endpoint=self.base_url, model=self.model or (models[0] if models else None), detail=f"OpenAI-compatible server, models: {models}", health_latency_ms=round(ms, 1)) return BackendStatus(backend="vllm", available=False, endpoint=self.base_url, detail=f"HTTP {res.status_code} from /models") except Exception as exc: return BackendStatus(backend="vllm", available=False, endpoint=self.base_url, detail=f"unreachable ({exc.__class__.__name__}); vLLM " "cannot run on this macOS machine — point " "RETRO_VLLM_URL at a CUDA/ROCm host") def generate(self, prompt: str, system: Optional[str] = None, max_tokens: int = 1024) -> InferenceResult: model = self.model if not model: res = httpx.get(f"{self.base_url}/models", timeout=5.0) data = res.json().get("data", []) if not data: raise RuntimeError(f"no models served at {self.base_url}") model = data[0]["id"] messages = ([{"role": "system", "content": system}] if system else []) + [ {"role": "user", "content": prompt}] t0 = time.perf_counter() res = httpx.post(f"{self.base_url}/chat/completions", json={"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": 0.0}, timeout=self.timeout) latency_ms = (time.perf_counter() - t0) * 1000 res.raise_for_status() body = res.json() usage = body.get("usage", {}) return InferenceResult( backend="vllm", model=model, text=body["choices"][0]["message"]["content"], prompt_tokens=int(usage.get("prompt_tokens", 0)), completion_tokens=int(usage.get("completion_tokens", 0)), latency_ms=round(latency_ms, 1)) BACKENDS = {"mock": MockBackend, "ollama": OllamaBackend, "vllm": VLLMBackend} def get_backend(name: str = "mock", **kwargs) -> InferenceBackend: if name not in BACKENDS: raise ValueError(f"unknown backend {name!r}; expected one of {sorted(BACKENDS)}") return BACKENDS[name](**kwargs) def all_backend_statuses() -> List[BackendStatus]: return [get_backend(n).health() for n in ("mock", "ollama", "vllm")]