llmstack-cli 0.1.0__py3-none-any.whl

llmstack/gateway/middleware/metrics.py

@@ -0,0 +1,115 @@
+"""Metrics middleware — tracks request counts, latencies, tokens in Prometheus format."""
+
+from __future__ import annotations
+
+import time
+from collections import defaultdict
+from threading import Lock
+
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+
+_lock = Lock()
+_request_count: dict[str, int] = defaultdict(int)
+_error_count: dict[str, int] = defaultdict(int)
+_tokens_in: int = 0
+_tokens_out: int = 0
+
+# Histogram buckets for latency
+_BUCKETS = [0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0]
+_latency_buckets: dict[str, list[int]] = defaultdict(lambda: [0] * (len(_BUCKETS) + 1))
+_latency_sum: dict[str, float] = defaultdict(float)
+_latency_count: dict[str, int] = defaultdict(int)
+
+
+class MetricsMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request: Request, call_next):
+        path = request.url.path
+        if path in ("/metrics", "/healthz"):
+            return await call_next(request)
+
+        start = time.monotonic()
+        response = await call_next(request)
+        duration = time.monotonic() - start
+
+        with _lock:
+            _request_count[path] += 1
+            _latency_sum[path] += duration
+            _latency_count[path] += 1
+
+            # Bucket assignment
+            buckets = _latency_buckets[path]
+            for i, bound in enumerate(_BUCKETS):
+                if duration <= bound:
+                    buckets[i] += 1
+            buckets[-1] += 1  # +Inf
+
+            if response.status_code >= 400:
+                _error_count[path] += 1
+
+        return response
+
+
+def record_tokens(input_tokens: int = 0, output_tokens: int = 0) -> None:
+    """Record token usage from a chat completion response."""
+    global _tokens_in, _tokens_out
+    with _lock:
+        _tokens_in += input_tokens
+        _tokens_out += output_tokens
+
+
+def get_metrics() -> dict:
+    """Return metrics as JSON (for /metrics JSON endpoint)."""
+    with _lock:
+        result = {}
+        for path in _request_count:
+            result[path] = {
+                "requests": _request_count[path],
+                "errors": _error_count.get(path, 0),
+                "latency_avg_ms": round(
+                    (_latency_sum[path] / _latency_count[path]) * 1000, 1
+                ) if _latency_count[path] else 0,
+            }
+        result["tokens"] = {"input": _tokens_in, "output": _tokens_out}
+        return result
+
+
+def get_prometheus_metrics() -> str:
+    """Return metrics in Prometheus exposition format."""
+    lines: list[str] = []
+
+    with _lock:
+        # Request counter
+        lines.append("# HELP llmstack_requests_total Total HTTP requests")
+        lines.append("# TYPE llmstack_requests_total counter")
+        for path, count in _request_count.items():
+            lines.append(f'llmstack_requests_total{{path="{path}"}} {count}')
+
+        # Error counter
+        lines.append("# HELP llmstack_errors_total Total HTTP errors (4xx/5xx)")
+        lines.append("# TYPE llmstack_errors_total counter")
+        for path, count in _error_count.items():
+            lines.append(f'llmstack_errors_total{{path="{path}"}} {count}')
+
+        # Latency histogram
+        lines.append("# HELP llmstack_request_duration_seconds Request latency histogram")
+        lines.append("# TYPE llmstack_request_duration_seconds histogram")
+        for path in _latency_count:
+            buckets = _latency_buckets[path]
+            cumulative = 0
+            for i, bound in enumerate(_BUCKETS):
+                cumulative += buckets[i]
+                lines.append(f'llmstack_request_duration_seconds_bucket{{path="{path}",le="{bound}"}} {cumulative}')
+            cumulative += buckets[-1]
+            lines.append(f'llmstack_request_duration_seconds_bucket{{path="{path}",le="+Inf"}} {cumulative}')
+            lines.append(f'llmstack_request_duration_seconds_sum{{path="{path}"}} {_latency_sum[path]:.4f}')
+            lines.append(f'llmstack_request_duration_seconds_count{{path="{path}"}} {_latency_count[path]}')
+
+        # Token counter
+        lines.append("# HELP llmstack_tokens_total Total tokens processed")
+        lines.append("# TYPE llmstack_tokens_total counter")
+        lines.append(f'llmstack_tokens_total{{type="input"}} {_tokens_in}')
+        lines.append(f'llmstack_tokens_total{{type="output"}} {_tokens_out}')
+
+    lines.append("")
+    return "\n".join(lines)

llmstack/gateway/proxy.py

@@ -0,0 +1,58 @@
+"""Proxy layer — forwards requests to inference and embedding backends."""
+
+from __future__ import annotations
+
+import os
+from typing import AsyncIterator
+
+import httpx
+
+INFERENCE_URL = os.getenv("LLMSTACK_INFERENCE_URL", "http://llmstack-ollama:11434/v1")
+EMBEDDINGS_URL = os.getenv("LLMSTACK_EMBEDDINGS_URL", "")
+
+# Timeout for inference (can be long for large models)
+REQUEST_TIMEOUT = int(os.getenv("LLMSTACK_REQUEST_TIMEOUT", "120"))
+
+
+async def proxy_chat_completion(payload: dict, stream: bool = False) -> dict | AsyncIterator[bytes]:
+    """Forward a chat completion request to the inference backend."""
+    url = f"{INFERENCE_URL}/chat/completions"
+    timeout = httpx.Timeout(REQUEST_TIMEOUT, connect=10)
+
+    if stream:
+        return _stream_response(url, payload, timeout)
+    else:
+        async with httpx.AsyncClient(timeout=timeout) as client:
+            resp = await client.post(url, json=payload)
+            resp.raise_for_status()
+            return resp.json()
+
+
+async def _stream_response(url: str, payload: dict, timeout: httpx.Timeout) -> AsyncIterator[bytes]:
+    """Stream SSE chunks from the inference backend."""
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        async with client.stream("POST", url, json=payload) as resp:
+            resp.raise_for_status()
+            async for chunk in resp.aiter_bytes():
+                yield chunk
+
+
+async def proxy_embeddings(payload: dict) -> dict:
+    """Forward an embeddings request to the embedding backend."""
+    url = EMBEDDINGS_URL or INFERENCE_URL
+    if not url.endswith("/embeddings"):
+        url = f"{url}/embeddings"
+
+    async with httpx.AsyncClient(timeout=30) as client:
+        resp = await client.post(url, json=payload)
+        resp.raise_for_status()
+        return resp.json()
+
+
+async def proxy_models() -> dict:
+    """List available models from the inference backend."""
+    url = f"{INFERENCE_URL}/models"
+    async with httpx.AsyncClient(timeout=10) as client:
+        resp = await client.get(url)
+        resp.raise_for_status()
+        return resp.json()

llmstack/gateway/routes/chat.py

@@ -0,0 +1,27 @@
+"""POST /v1/chat/completions — OpenAI-compatible chat endpoint."""
+
+from __future__ import annotations
+
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from llmstack.gateway.proxy import proxy_chat_completion
+
+router = APIRouter()
+
+
+@router.post("/chat/completions")
+async def chat_completions(request: Request):
+    payload = await request.json()
+    stream = payload.get("stream", False)
+
+    if stream:
+        chunks = await proxy_chat_completion(payload, stream=True)
+        return StreamingResponse(
+            chunks,
+            media_type="text/event-stream",
+            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
+        )
+    else:
+        result = await proxy_chat_completion(payload, stream=False)
+        return JSONResponse(content=result)

llmstack/gateway/routes/embeddings.py

@@ -0,0 +1,17 @@
+"""POST /v1/embeddings — OpenAI-compatible embeddings endpoint."""
+
+from __future__ import annotations
+
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
+
+from llmstack.gateway.proxy import proxy_embeddings
+
+router = APIRouter()
+
+
+@router.post("/embeddings")
+async def embeddings(request: Request):
+    payload = await request.json()
+    result = await proxy_embeddings(payload)
+    return JSONResponse(content=result)

llmstack/gateway/routes/health.py

@@ -0,0 +1,55 @@
+"""GET /healthz — gateway health check."""
+
+from __future__ import annotations
+
+import os
+
+import httpx
+from fastapi import APIRouter
+from fastapi.responses import JSONResponse
+
+router = APIRouter()
+
+INFERENCE_URL = os.getenv("LLMSTACK_INFERENCE_URL", "")
+QDRANT_URL = os.getenv("LLMSTACK_QDRANT_URL", "")
+REDIS_URL = os.getenv("LLMSTACK_REDIS_URL", "")
+
+
+async def _check_url(url: str) -> bool:
+    if not url:
+        return False
+    try:
+        async with httpx.AsyncClient(timeout=5) as client:
+            resp = await client.get(url)
+            return resp.status_code == 200
+    except Exception:
+        return False
+
+
+@router.get("/healthz")
+async def healthz():
+    checks = {}
+
+    if INFERENCE_URL:
+        checks["inference"] = await _check_url(INFERENCE_URL.replace("/v1", "/health"))
+
+    if QDRANT_URL:
+        checks["qdrant"] = await _check_url(f"{QDRANT_URL}/healthz")
+
+    all_ok = all(checks.values()) if checks else True
+    status_code = 200 if all_ok else 503
+
+    return JSONResponse(
+        content={"status": "ok" if all_ok else "degraded", "services": checks},
+        status_code=status_code,
+    )
+
+
+@router.get("/metrics")
+async def metrics():
+    from llmstack.gateway.middleware.metrics import get_prometheus_metrics
+    from fastapi.responses import PlainTextResponse
+    return PlainTextResponse(
+        content=get_prometheus_metrics(),
+        media_type="text/plain; version=0.0.4; charset=utf-8",
+    )

llmstack/gateway/routes/models.py

@@ -0,0 +1,16 @@
+"""GET /v1/models — list available models."""
+
+from __future__ import annotations
+
+from fastapi import APIRouter
+from fastapi.responses import JSONResponse
+
+from llmstack.gateway.proxy import proxy_models
+
+router = APIRouter()
+
+
+@router.get("/models")
+async def list_models():
+    result = await proxy_models()
+    return JSONResponse(content=result)

llmstack/plugins/loader.py

@@ -0,0 +1,5 @@
+"""Plugin loader — re-exports the registry for convenience."""
+
+from llmstack.services.registry import ServiceRegistry
+
+__all__ = ["ServiceRegistry"]

llmstack/plugins/spec.py

@@ -0,0 +1,20 @@
+"""Plugin specification.
+
+To create an llmstack plugin:
+
+1. Create a class that extends ServiceBase
+2. Set `name` and `category` class attributes
+3. Implement: container_spec(), health_url()
+4. Optionally implement: post_start(), openai_base_url()
+5. Register via entry_points in pyproject.toml:
+
+    [project.entry-points."llmstack.services"]
+    my_service = "my_package:MyService"
+
+6. Publish to PyPI. Users install with:
+    pip install llmstack-plugin-myservice
+"""
+
+from llmstack.services.base import ServiceBase
+
+__all__ = ["ServiceBase"]

llmstack/services/base.py

@@ -0,0 +1,65 @@
+"""Base class for all managed services."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+
+class ServiceState(str, Enum):
+    STOPPED = "stopped"
+    STARTING = "starting"
+    RUNNING = "running"
+    UNHEALTHY = "unhealthy"
+    ERROR = "error"
+
+
+@dataclass
+class ServiceStatus:
+    name: str
+    state: ServiceState
+    port: int | None = None
+    container_id: str | None = None
+    message: str = ""
+
+
+class ServiceBase(ABC):
+    """Every llmstack service (inference, vectordb, cache, etc.) implements this."""
+
+    name: str
+    category: str  # inference, vectordb, cache, embeddings, gateway, observe
+
+    @abstractmethod
+    def container_spec(self) -> dict[str, Any]:
+        """Return kwargs for docker.containers.run().
+
+        Must include at least: image, ports, environment.
+        May include: volumes, device_requests, healthcheck, command.
+        """
+
+    @abstractmethod
+    def health_url(self) -> str:
+        """HTTP URL to GET for health checks (from the host)."""
+
+    async def post_start(self) -> None:
+        """Hook called after the container is healthy.
+
+        Override for actions like pulling a model.
+        """
+
+    def openai_base_url(self) -> str | None:
+        """If this service exposes an OpenAI-compatible API, return its internal Docker URL."""
+        return None
+
+    def internal_url(self) -> str:
+        """Return the URL reachable from other containers on the Docker network."""
+        spec = self.container_spec()
+        ports = spec.get("ports", {})
+        # Get the first container port
+        if ports:
+            container_port = list(ports.values())[0] if isinstance(ports, dict) else None
+            if container_port:
+                return f"http://{self.name}:{container_port}"
+        return f"http://{self.name}"

llmstack/services/cache/redis.py

@@ -0,0 +1,33 @@
+"""Redis cache service."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from llmstack.config.schema import CacheConfig
+from llmstack.services.base import ServiceBase
+
+
+class RedisService(ServiceBase):
+    name = "redis"
+    category = "cache"
+
+    def __init__(self, config: CacheConfig):
+        self.config = config
+
+    def container_spec(self) -> dict[str, Any]:
+        return {
+            "image": "redis:7-alpine",
+            "name": "llmstack-redis",
+            "ports": {"6379/tcp": self.config.port},
+            "command": [
+                "redis-server",
+                "--maxmemory", self.config.max_memory,
+                "--maxmemory-policy", "allkeys-lru",
+            ],
+            "environment": {},
+        }
+
+    def health_url(self) -> str:
+        # Redis doesn't have HTTP health, we check via TCP
+        return f"http://localhost:{self.config.port}"

llmstack/services/embeddings/tei.py

@@ -0,0 +1,49 @@
+"""HuggingFace Text Embeddings Inference (TEI) service."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from llmstack.config.schema import EmbeddingSpec
+from llmstack.core.hardware import HardwareProfile
+from llmstack.services.base import ServiceBase
+
+
+class TEIService(ServiceBase):
+    name = "tei"
+    category = "embeddings"
+
+    def __init__(self, spec: EmbeddingSpec, hw: HardwareProfile):
+        self.spec = spec
+        self.hw = hw
+        self.host_port = 8002
+
+    def container_spec(self) -> dict[str, Any]:
+        cmd = ["--model-id", self.spec.name, "--port", "80"]
+
+        spec: dict[str, Any] = {
+            "image": "ghcr.io/huggingface/text-embeddings-inference:cpu-latest",
+            "name": "llmstack-tei",
+            "ports": {"80/tcp": self.host_port},
+            "command": cmd,
+            "volumes": {
+                "llmstack_tei_cache": {"bind": "/data", "mode": "rw"},
+            },
+            "environment": {},
+        }
+
+        # Use GPU image if NVIDIA available
+        if self.hw.gpu_vendor == "nvidia":
+            import docker
+            spec["image"] = "ghcr.io/huggingface/text-embeddings-inference:latest"
+            spec["device_requests"] = [
+                docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])
+            ]
+
+        return spec
+
+    def health_url(self) -> str:
+        return f"http://localhost:{self.host_port}/health"
+
+    def openai_base_url(self) -> str:
+        return "http://llmstack-tei:80/v1"

llmstack/services/gateway/service.py

@@ -0,0 +1,47 @@
+"""Gateway service — runs the FastAPI proxy as a Docker container."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from llmstack.config.schema import GatewayConfig
+from llmstack.services.base import ServiceBase
+
+
+class GatewayService(ServiceBase):
+    name = "gateway"
+    category = "gateway"
+
+    def __init__(
+        self,
+        config: GatewayConfig,
+        inference_url: str,
+        embeddings_url: str,
+        qdrant_url: str = "",
+        redis_url: str = "",
+    ):
+        self.config = config
+        self.inference_url = inference_url
+        self.embeddings_url = embeddings_url
+        self.qdrant_url = qdrant_url
+        self.redis_url = redis_url
+
+    def container_spec(self) -> dict[str, Any]:
+        return {
+            "image": "ghcr.io/mara-werils/llmstack-gateway:latest",
+            "name": "llmstack-gateway",
+            "ports": {"8000/tcp": self.config.port},
+            "environment": {
+                "LLMSTACK_INFERENCE_URL": self.inference_url,
+                "LLMSTACK_EMBEDDINGS_URL": self.embeddings_url,
+                "LLMSTACK_QDRANT_URL": self.qdrant_url,
+                "LLMSTACK_REDIS_URL": self.redis_url,
+                "LLMSTACK_API_KEYS": ",".join(self.config.api_keys),
+                "LLMSTACK_CORS_ORIGINS": ",".join(self.config.cors),
+                "LLMSTACK_REQUEST_TIMEOUT": str(self.config.request_timeout),
+                "LLMSTACK_RATE_LIMIT": self.config.rate_limit,
+            },
+        }
+
+    def health_url(self) -> str:
+        return f"http://localhost:{self.config.port}/healthz"

llmstack/services/inference/ollama.py

@@ -0,0 +1,60 @@
+"""Ollama inference service."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import httpx
+
+from llmstack.config.schema import ModelSpec
+from llmstack.core.hardware import HardwareProfile
+from llmstack.services.base import ServiceBase
+
+
+class OllamaService(ServiceBase):
+    name = "ollama"
+    category = "inference"
+
+    def __init__(self, model: ModelSpec, hw: HardwareProfile):
+        self.model = model
+        self.hw = hw
+        self.host_port = 11434
+
+    def container_spec(self) -> dict[str, Any]:
+        spec: dict[str, Any] = {
+            "image": "ollama/ollama:latest",
+            "name": "llmstack-ollama",
+            "ports": {"11434/tcp": self.host_port},
+            "volumes": {
+                "llmstack_ollama_data": {"bind": "/root/.ollama", "mode": "rw"},
+            },
+            "environment": {},
+        }
+
+        # GPU passthrough for NVIDIA
+        if self.hw.gpu_vendor == "nvidia":
+            import docker
+            spec["device_requests"] = [
+                docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])
+            ]
+
+        return spec
+
+    def health_url(self) -> str:
+        return f"http://localhost:{self.host_port}"
+
+    async def post_start(self) -> None:
+        """Pull the model after Ollama is healthy."""
+        model_name = self.model.name
+        if self.model.quantization:
+            model_name = f"{self.model.name}:{self.model.quantization}"
+
+        async with httpx.AsyncClient(timeout=600) as client:
+            resp = await client.post(
+                f"http://localhost:{self.host_port}/api/pull",
+                json={"name": model_name, "stream": False},
+            )
+            resp.raise_for_status()
+
+    def openai_base_url(self) -> str:
+        return f"http://llmstack-ollama:{self.host_port}/v1"

llmstack/services/inference/vllm.py

@@ -0,0 +1,57 @@
+"""vLLM inference service."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from llmstack.config.schema import ModelSpec
+from llmstack.core.hardware import HardwareProfile
+from llmstack.services.base import ServiceBase
+
+
+class VllmService(ServiceBase):
+    name = "vllm"
+    category = "inference"
+
+    def __init__(self, model: ModelSpec, hw: HardwareProfile):
+        self.model = model
+        self.hw = hw
+        self.host_port = 8001
+
+    def container_spec(self) -> dict[str, Any]:
+        import docker
+
+        cmd = [
+            "--model", self.model.name,
+            "--host", "0.0.0.0",
+            "--port", "8000",
+            "--max-model-len", str(self.model.context_length),
+        ]
+
+        if self.model.quantization:
+            cmd.extend(["--quantization", self.model.quantization])
+
+        spec: dict[str, Any] = {
+            "image": "vllm/vllm-openai:latest",
+            "name": "llmstack-vllm",
+            "ports": {"8000/tcp": self.host_port},
+            "command": cmd,
+            "environment": {
+                "HUGGING_FACE_HUB_TOKEN": "",
+            },
+            "volumes": {
+                "llmstack_vllm_cache": {"bind": "/root/.cache/huggingface", "mode": "rw"},
+            },
+            "device_requests": [
+                docker.types.DeviceRequest(count=-1, capabilities=[["gpu"]])
+            ],
+            "shm_size": "4g",
+        }
+
+        return spec
+
+    def health_url(self) -> str:
+        return f"http://localhost:{self.host_port}/health"
+
+    def openai_base_url(self) -> str:
+        return "http://llmstack-vllm:8000/v1"