llmstack-cli 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

llmstack/__init__.py

@@ -1,3 +1,3 @@
 """LLMStack — One command. Full LLM stack. Zero config."""
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"

llmstack/core/health.py

@@ -16,7 +16,7 @@ async def wait_healthy(url: str, timeout: int = 120, interval: float = 2.0) -> b
                 resp = await client.get(url)
                 if resp.status_code == 200:
                     return True
-            except (httpx.ConnectError, httpx.ReadTimeout, httpx.ConnectTimeout):
+            except httpx.HTTPError:
                 pass
             await asyncio.sleep(interval)
             elapsed += interval

llmstack/gateway/Dockerfile

@@ -6,7 +6,9 @@ RUN pip install --no-cache-dir \
     fastapi>=0.115 \
     uvicorn[standard]>=0.30 \
     httpx>=0.27 \
-    starlette>=0.40
+    starlette>=0.40 \
+    redis>=5.0 \
+    pydantic>=2.0
 
 COPY src/llmstack/__init__.py /app/llmstack/__init__.py
 COPY src/llmstack/gateway/ /app/llmstack/gateway/

llmstack/gateway/cache.py

@@ -0,0 +1,187 @@
+"""Semantic response cache backed by Redis.
+
+Caches LLM completions keyed by a hash of (model + messages).
+Supports TTL-based expiration and cache hit/miss metrics.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import time
+from dataclasses import dataclass, field
+
+import redis.asyncio as aioredis
+
+REDIS_URL = os.getenv("LLMSTACK_REDIS_URL", "")
+CACHE_TTL = int(os.getenv("LLMSTACK_CACHE_TTL", "3600"))  # 1 hour default
+CACHE_ENABLED = os.getenv("LLMSTACK_CACHE_ENABLED", "true").lower() == "true"
+
+# Prefix to avoid key collisions
+_KEY_PREFIX = "llmstack:cache:"
+_STATS_KEY = "llmstack:cache:stats"
+
+
+@dataclass
+class CacheStats:
+    hits: int = 0
+    misses: int = 0
+    evictions: int = 0
+    avg_hit_latency_ms: float = 0.0
+    _hit_latencies: list[float] = field(default_factory=list, repr=False)
+
+    def record_hit(self, latency_ms: float) -> None:
+        self.hits += 1
+        self._hit_latencies.append(latency_ms)
+        self.avg_hit_latency_ms = sum(self._hit_latencies) / len(self._hit_latencies)
+
+    def record_miss(self) -> None:
+        self.misses += 1
+
+    @property
+    def hit_rate(self) -> float:
+        total = self.hits + self.misses
+        return self.hits / total if total > 0 else 0.0
+
+    def to_dict(self) -> dict:
+        return {
+            "hits": self.hits,
+            "misses": self.misses,
+            "hit_rate": round(self.hit_rate, 4),
+            "avg_hit_latency_ms": round(self.avg_hit_latency_ms, 2),
+        }
+
+
+class ResponseCache:
+    """Redis-backed LLM response cache with semantic key hashing."""
+
+    def __init__(self, redis_url: str = "", ttl: int = CACHE_TTL):
+        self._url = redis_url or REDIS_URL
+        self._ttl = ttl
+        self._redis: aioredis.Redis | None = None
+        self._stats = CacheStats()
+        self._connected = False
+
+    async def connect(self) -> None:
+        """Lazily connect to Redis."""
+        if self._connected or not self._url:
+            return
+        try:
+            self._redis = aioredis.from_url(
+                self._url,
+                decode_responses=True,
+                socket_connect_timeout=3,
+                socket_timeout=5,
+                retry_on_timeout=True,
+            )
+            await self._redis.ping()
+            self._connected = True
+        except Exception:
+            self._redis = None
+            self._connected = False
+
+    async def close(self) -> None:
+        if self._redis:
+            await self._redis.aclose()
+            self._connected = False
+
+    @staticmethod
+    def _build_cache_key(model: str, messages: list[dict], temperature: float = 1.0) -> str:
+        """Deterministic hash of the request for cache lookup.
+
+        Only caches when temperature <= 0.1 (near-deterministic outputs).
+        For higher temperatures, returns empty string (skip cache).
+        """
+        if temperature > 0.1:
+            return ""
+
+        # Normalize messages to a stable representation
+        normalized = []
+        for msg in messages:
+            normalized.append({
+                "role": msg.get("role", ""),
+                "content": msg.get("content", ""),
+            })
+
+        payload = json.dumps({"model": model, "messages": normalized}, sort_keys=True)
+        digest = hashlib.sha256(payload.encode()).hexdigest()
+        return f"{_KEY_PREFIX}{digest}"
+
+    async def get(self, model: str, messages: list[dict], temperature: float = 1.0) -> dict | None:
+        """Look up a cached response. Returns None on miss."""
+        if not self._connected or not CACHE_ENABLED:
+            return None
+
+        key = self._build_cache_key(model, messages, temperature)
+        if not key:
+            self._stats.record_miss()
+            return None
+
+        start = time.monotonic()
+        try:
+            raw = await self._redis.get(key)  # type: ignore[union-attr]
+            if raw is None:
+                self._stats.record_miss()
+                return None
+
+            latency_ms = (time.monotonic() - start) * 1000
+            self._stats.record_hit(latency_ms)
+
+            cached = json.loads(raw)
+            # Mark response as cached
+            cached["_cached"] = True
+            cached["_cache_age_s"] = int(time.time() - cached.get("_cached_at", 0))
+            return cached
+        except Exception:
+            self._stats.record_miss()
+            return None
+
+    async def put(self, model: str, messages: list[dict], response: dict,
+                  temperature: float = 1.0) -> None:
+        """Store a response in the cache."""
+        if not self._connected or not CACHE_ENABLED:
+            return
+
+        key = self._build_cache_key(model, messages, temperature)
+        if not key:
+            return
+
+        try:
+            response["_cached_at"] = int(time.time())
+            raw = json.dumps(response)
+            await self._redis.set(key, raw, ex=self._ttl)  # type: ignore[union-attr]
+        except Exception:
+            pass  # Cache write failure is non-fatal
+
+    async def invalidate(self, pattern: str = "*") -> int:
+        """Delete cache entries matching a pattern. Returns count deleted."""
+        if not self._connected:
+            return 0
+        try:
+            keys = []
+            async for key in self._redis.scan_iter(f"{_KEY_PREFIX}{pattern}"):  # type: ignore[union-attr]
+                keys.append(key)
+            if keys:
+                return await self._redis.delete(*keys)  # type: ignore[union-attr]
+            return 0
+        except Exception:
+            return 0
+
+    @property
+    def stats(self) -> CacheStats:
+        return self._stats
+
+
+# Module-level singleton
+_cache: ResponseCache | None = None
+
+
+async def get_cache() -> ResponseCache:
+    """Get or create the global cache instance."""
+    global _cache
+    if _cache is None:
+        _cache = ResponseCache()
+    if not _cache._connected:
+        await _cache.connect()
+    return _cache

llmstack/gateway/circuit_breaker.py

@@ -0,0 +1,177 @@
+"""Circuit breaker for inference backend resilience.
+
+Implements the three-state circuit breaker pattern:
+  CLOSED → OPEN → HALF_OPEN → CLOSED
+
+Prevents cascading failures when the inference backend is down
+by failing fast instead of timing out on every request.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from enum import Enum
+
+
+class CircuitState(str, Enum):
+    CLOSED = "closed"        # Normal operation — requests flow through
+    OPEN = "open"            # Backend is down — fail fast
+    HALF_OPEN = "half_open"  # Probing — allow one request to test recovery
+
+
+class CircuitBreakerError(Exception):
+    """Raised when the circuit is open and requests are rejected."""
+
+    def __init__(self, retry_after: float):
+        self.retry_after = retry_after
+        super().__init__(f"Circuit breaker is open. Retry after {retry_after:.0f}s")
+
+
+class CircuitBreaker:
+    """Circuit breaker with exponential backoff and half-open probing.
+
+    Usage:
+        breaker = CircuitBreaker()
+
+        async def call_inference(payload):
+            breaker.check()  # raises CircuitBreakerError if open
+            try:
+                result = await make_request(payload)
+                breaker.record_success()
+                return result
+            except Exception as e:
+                breaker.record_failure()
+                raise
+    """
+
+    def __init__(
+        self,
+        failure_threshold: int = 5,
+        recovery_timeout: float = 30.0,
+        half_open_max_calls: int = 1,
+        max_recovery_timeout: float = 300.0,
+        backoff_multiplier: float = 2.0,
+    ):
+        self.failure_threshold = failure_threshold
+        self.base_recovery_timeout = recovery_timeout
+        self.half_open_max_calls = half_open_max_calls
+        self.max_recovery_timeout = max_recovery_timeout
+        self.backoff_multiplier = backoff_multiplier
+
+        self._state = CircuitState.CLOSED
+        self._failure_count = 0
+        self._success_count = 0
+        self._last_failure_time = 0.0
+        self._half_open_calls = 0
+        self._consecutive_opens = 0
+        self._lock = asyncio.Lock()
+
+        # Metrics
+        self._total_failures = 0
+        self._total_successes = 0
+        self._total_rejections = 0
+        self._last_state_change = time.monotonic()
+
+    @property
+    def state(self) -> CircuitState:
+        return self._state
+
+    @property
+    def current_recovery_timeout(self) -> float:
+        """Exponential backoff on recovery timeout."""
+        timeout = self.base_recovery_timeout * (
+            self.backoff_multiplier ** self._consecutive_opens
+        )
+        return min(timeout, self.max_recovery_timeout)
+
+    def check(self) -> None:
+        """Check if a request is allowed. Raises CircuitBreakerError if not."""
+        if self._state == CircuitState.CLOSED:
+            return
+
+        if self._state == CircuitState.OPEN:
+            elapsed = time.monotonic() - self._last_failure_time
+            if elapsed >= self.current_recovery_timeout:
+                # Transition to half-open
+                self._state = CircuitState.HALF_OPEN
+                self._half_open_calls = 0
+                self._last_state_change = time.monotonic()
+                return
+
+            retry_after = self.current_recovery_timeout - elapsed
+            self._total_rejections += 1
+            raise CircuitBreakerError(retry_after=retry_after)
+
+        if self._state == CircuitState.HALF_OPEN:
+            if self._half_open_calls >= self.half_open_max_calls:
+                # Too many half-open calls pending, reject
+                self._total_rejections += 1
+                raise CircuitBreakerError(retry_after=5.0)
+            self._half_open_calls += 1
+
+    def record_success(self) -> None:
+        """Record a successful request."""
+        self._total_successes += 1
+
+        if self._state == CircuitState.HALF_OPEN:
+            self._success_count += 1
+            # Recovered — close the circuit
+            self._state = CircuitState.CLOSED
+            self._failure_count = 0
+            self._success_count = 0
+            self._consecutive_opens = 0
+            self._last_state_change = time.monotonic()
+        elif self._state == CircuitState.CLOSED:
+            # Reset failure count on success
+            self._failure_count = 0
+
+    def record_failure(self) -> None:
+        """Record a failed request."""
+        self._total_failures += 1
+        self._failure_count += 1
+        self._last_failure_time = time.monotonic()
+
+        if self._state == CircuitState.HALF_OPEN:
+            # Probe failed — back to open with longer timeout
+            self._state = CircuitState.OPEN
+            self._consecutive_opens += 1
+            self._last_state_change = time.monotonic()
+        elif self._state == CircuitState.CLOSED:
+            if self._failure_count >= self.failure_threshold:
+                self._state = CircuitState.OPEN
+                self._consecutive_opens += 1
+                self._last_state_change = time.monotonic()
+
+    def reset(self) -> None:
+        """Manually reset the circuit breaker to closed state."""
+        self._state = CircuitState.CLOSED
+        self._failure_count = 0
+        self._success_count = 0
+        self._half_open_calls = 0
+        self._consecutive_opens = 0
+        self._last_state_change = time.monotonic()
+
+    def metrics(self) -> dict:
+        """Return circuit breaker metrics."""
+        return {
+            "state": self._state.value,
+            "failure_count": self._failure_count,
+            "consecutive_opens": self._consecutive_opens,
+            "current_recovery_timeout_s": round(self.current_recovery_timeout, 1),
+            "total_successes": self._total_successes,
+            "total_failures": self._total_failures,
+            "total_rejections": self._total_rejections,
+            "time_in_state_s": round(time.monotonic() - self._last_state_change, 1),
+        }
+
+
+# Module-level singleton for the inference backend
+_inference_breaker: CircuitBreaker | None = None
+
+
+def get_inference_breaker() -> CircuitBreaker:
+    global _inference_breaker
+    if _inference_breaker is None:
+        _inference_breaker = CircuitBreaker()
+    return _inference_breaker

llmstack/gateway/main.py

@@ -1,8 +1,9 @@
-"""LLMStack Gateway — OpenAI-compatible API gateway."""
+"""LLMStack Gateway — OpenAI-compatible API gateway with caching, RAG, and resilience."""
 
 from __future__ import annotations
 
 import os
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
@@ -11,15 +12,28 @@ from llmstack.gateway.routes.chat import router as chat_router
 from llmstack.gateway.routes.embeddings import router as embeddings_router
 from llmstack.gateway.routes.models import router as models_router
 from llmstack.gateway.routes.health import router as health_router
+from llmstack.gateway.routes.rag import router as rag_router
 from llmstack.gateway.middleware.auth import AuthMiddleware
 from llmstack.gateway.middleware.metrics import MetricsMiddleware
+from llmstack.gateway.middleware.rate_limit import RateLimitMiddleware
+from llmstack.gateway.middleware.logging import LoggingMiddleware
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Startup: connect cache. Shutdown: close connections."""
+    from llmstack.gateway.cache import get_cache
+    cache = await get_cache()
+    yield
+    await cache.close()
 
 
 def create_app() -> FastAPI:
     app = FastAPI(
         title="LLMStack Gateway",
-        description="OpenAI-compatible API gateway for LLMStack",
-        version="0.1.0",
+        description="OpenAI-compatible API gateway with caching, RAG, and resilience",
+        version="0.3.0",
+        lifespan=lifespan,
     )
 
     # CORS
@@ -32,18 +46,27 @@ def create_app() -> FastAPI:
         allow_headers=["*"],
     )
 
-    # Auth
+    # Middleware stack (order matters: outermost first)
+    # 1. Logging (outermost — captures everything)
+    app.add_middleware(LoggingMiddleware)
+
+    # 2. Auth
     api_keys = os.getenv("LLMSTACK_API_KEYS", "")
     if api_keys:
         app.add_middleware(AuthMiddleware, api_keys=api_keys.split(","))
 
-    # Metrics
+    # 3. Rate limiting
+    rate_limit = os.getenv("LLMSTACK_RATE_LIMIT", "100/min")
+    app.add_middleware(RateLimitMiddleware, rate_limit=rate_limit)
+
+    # 4. Metrics (innermost — measures actual handler time)
     app.add_middleware(MetricsMiddleware)
 
     # Routes
     app.include_router(chat_router, prefix="/v1")
     app.include_router(embeddings_router, prefix="/v1")
     app.include_router(models_router, prefix="/v1")
+    app.include_router(rag_router, prefix="/v1")
     app.include_router(health_router)
 
     return app

llmstack/gateway/middleware/logging.py

@@ -0,0 +1,114 @@
+"""Structured request logging middleware with correlation IDs.
+
+Adds X-Request-ID to every request/response and logs structured JSON
+for each request including method, path, status, duration, and tokens.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import sys
+import time
+import uuid
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+
+
+LOG_LEVEL = os.getenv("LLMSTACK_LOG_LEVEL", "INFO").upper()
+LOG_FORMAT = os.getenv("LLMSTACK_LOG_FORMAT", "json")  # json | text
+
+# Configure structured logger
+logger = logging.getLogger("llmstack.gateway")
+
+
+def _setup_logger() -> None:
+    """Configure the gateway logger with structured JSON output."""
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(getattr(logging, LOG_LEVEL, logging.INFO))
+
+    if LOG_FORMAT == "json":
+        handler.setFormatter(_JsonFormatter())
+    else:
+        handler.setFormatter(logging.Formatter(
+            "%(asctime)s %(levelname)s [%(request_id)s] %(message)s",
+            datefmt="%Y-%m-%dT%H:%M:%S",
+        ))
+
+    logger.addHandler(handler)
+    logger.setLevel(getattr(logging, LOG_LEVEL, logging.INFO))
+    logger.propagate = False
+
+
+class _JsonFormatter(logging.Formatter):
+    """Emit log records as single-line JSON."""
+
+    def format(self, record: logging.LogRecord) -> str:
+        log_dict = {
+            "ts": self.formatTime(record, "%Y-%m-%dT%H:%M:%S"),
+            "level": record.levelname,
+            "msg": record.getMessage(),
+        }
+        # Merge extra fields
+        for key in ("request_id", "method", "path", "status", "duration_ms",
+                     "client_ip", "tokens_in", "tokens_out", "cache_hit", "model"):
+            val = getattr(record, key, None)
+            if val is not None:
+                log_dict[key] = val
+
+        return json.dumps(log_dict, default=str)
+
+
+class LoggingMiddleware(BaseHTTPMiddleware):
+    """Adds correlation ID and structured logging to every request."""
+
+    SKIP_PATHS = {"/metrics", "/healthz"}
+
+    def __init__(self, app):
+        super().__init__(app)
+        _setup_logger()
+
+    async def dispatch(self, request: Request, call_next):
+        if request.url.path in self.SKIP_PATHS:
+            return await call_next(request)
+
+        # Generate or reuse correlation ID
+        request_id = request.headers.get("X-Request-ID", str(uuid.uuid4())[:8])
+
+        # Store in request state for downstream use
+        request.state.request_id = request_id
+
+        start = time.monotonic()
+        response = await call_next(request)
+        duration_ms = round((time.monotonic() - start) * 1000, 1)
+
+        # Add correlation headers
+        response.headers["X-Request-ID"] = request_id
+
+        # Extract client IP
+        client = request.client
+        client_ip = client.host if client else "unknown"
+        forwarded = request.headers.get("X-Forwarded-For", "")
+        if forwarded:
+            client_ip = forwarded.split(",")[0].strip()
+
+        # Log the request
+        logger.info(
+            "%s %s %d %.1fms",
+            request.method,
+            request.url.path,
+            response.status_code,
+            duration_ms,
+            extra={
+                "request_id": request_id,
+                "method": request.method,
+                "path": request.url.path,
+                "status": response.status_code,
+                "duration_ms": duration_ms,
+                "client_ip": client_ip,
+            },
+        )
+
+        return response