clossers/didi-lot1-ai
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Livrare LOT 1 - Didi

Browse source
This commit is contained in:
Dezvoltari Evotech 2026-06-25 14:13:25 -07:00
commit 5380c3fc63
990 changed files with 133308 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

16
ai_platform/modules/llm-inference/src/llm_inference/__init__.py Normal file
View file

@ -0,0 +1,16 @@
"""LLM Inference - Unified LLM inference with multiple backends."""
from llm_inference.client import LLMClient
from llm_inference.config import LLMSettings, get_settings
from llm_inference.types import BackendType, ChatMessage
__version__ = "0.1.0"
__all__ = [
"BackendType",
"ChatMessage",
"LLMClient",
"LLMSettings",
"__version__",
"get_settings",
]

5
ai_platform/modules/llm-inference/src/llm_inference/api/__init__.py Normal file
View file

@ -0,0 +1,5 @@
"""FastAPI application for LLM inference."""
from llm_inference.api.app import create_app
__all__ = ["create_app"]

165
ai_platform/modules/llm-inference/src/llm_inference/api/app.py Normal file
View file

@ -0,0 +1,165 @@
"""FastAPI application factory."""
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from fastapi import FastAPI
from llm_inference.api.dependencies import init_concurrency_limiter
from llm_inference.api.middleware import RateLimitMiddleware, RequestIdMiddleware
from llm_inference.api.routes import completions, health, info, models
from llm_inference.backends.llamacpp_backend import LlamaCppBackend
from llm_inference.client import LLMClient
from llm_inference.config import SettingsCache
from llm_inference.image_processing import close_http_client
from llm_inference.logging import configure_logging, get_logger
from llm_inference.runtime_config import RuntimeConfigClient
from llm_inference.types import BackendType
logger = get_logger("app")
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
"""Application lifespan manager.
Initializes the LLM client, settings, logging, and concurrency limiter
on startup, and performs cleanup on shutdown.
Args:
app: FastAPI application instance.
Yields:
None: Control to the application.
"""
settings = SettingsCache.get()
configure_logging(settings.log_level, settings.log_json)
logger.info("Starting LLM Inference API")
logger.info("Default backend: %s", settings.default_backend)
logger.info("vLLM enabled: %s", settings.enable_vllm)
logger.info("llama.cpp enabled: %s", settings.enable_llamacpp)
init_concurrency_limiter(settings.max_concurrent_completions)
app.state.settings = settings
app.state.client = LLMClient(settings)
# Start the runtime config client created in create_app()
if hasattr(app.state, "runtime_config"):
await app.state.runtime_config.start()
if settings.enable_llamacpp:
try:
backend = app.state.client.registry.get(BackendType.LLAMACPP)
if isinstance(backend, LlamaCppBackend):
backend.start_health_checks()
except Exception:
pass
logger.info("LLM Inference API started successfully")
yield
logger.info("Shutting down LLM Inference API")
if hasattr(app.state, "runtime_config"):
await app.state.runtime_config.stop()
await close_http_client()
def create_app() -> FastAPI:
"""Create and configure the FastAPI application.
Returns:
FastAPI: Configured FastAPI application.
"""
settings = SettingsCache.get()
app = FastAPI(
title="LLM Inference API",
description=(
"Unified LLM inference with multiple backends (LiteLLM, vLLM, llama.cpp)"
),
version="0.1.0",
lifespan=lifespan,
servers=[{"url": settings.external_url, "description": "LLM Inference API"}],
)
# ---------------------------------------------------------------- observability
# Prometheus /metrics + OTel tracing (no-op if deps missing or OTEL endpoint unset)
try:
from prometheus_fastapi_instrumentator import Instrumentator as _Inst # type: ignore
_Inst(should_group_status_codes=True).instrument(app).expose(app, endpoint="/metrics", include_in_schema=False)
except ImportError:
pass
import os as _os # noqa: E402
_otel_ep = _os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
if _otel_ep:
try:
from opentelemetry import trace as _trace # type: ignore
from opentelemetry.sdk.resources import Resource as _R # type: ignore
from opentelemetry.sdk.trace import TracerProvider as _TP # type: ignore
from opentelemetry.sdk.trace.export import BatchSpanProcessor as _BSP # type: ignore
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as _Exp # type: ignore
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor as _FInst # type: ignore
_provider = _TP(resource=_R.create({"service.name": _os.environ.get("OTEL_SERVICE_NAME", "didiAI-llm-api")}))
_provider.add_span_processor(_BSP(_Exp(endpoint=_otel_ep, insecure=True)))
_trace.set_tracer_provider(_provider)
_FInst.instrument_app(app)
print(f"[otel] didiAI-llm-api instrumented -> {_otel_ep}")
except ImportError as _e:
print(f"[otel] skip: {_e}")
# Runtime config client — instantiated here so middleware can capture it.
# Started/stopped inside lifespan().
runtime_config = RuntimeConfigClient(
dashboard_url=settings.dashboard_url,
live_log_logger_name="llm_inference",
live_log_key="llm.log.level",
)
app.state.runtime_config = runtime_config
# Add middleware (order matters - first added is outermost)
app.add_middleware(
RateLimitMiddleware,
rate=settings.rate_limit_rps,
burst=settings.rate_limit_burst,
exclude_paths=["/health", "/ready"],
runtime_config=runtime_config,
rate_key="llm.rate_limit.rps",
burst_key="llm.rate_limit.burst",
)
app.add_middleware(RequestIdMiddleware)
app.include_router(health.router, tags=["Health"])
app.include_router(completions.router, prefix="/v1", tags=["Completions"])
app.include_router(models.router, prefix="/v1", tags=["Models"])
app.include_router(info.router, tags=["Catalog"])
return app

232
ai_platform/modules/llm-inference/src/llm_inference/api/dependencies.py Normal file
View file

@ -0,0 +1,232 @@
"""FastAPI dependencies for LLM inference."""
import asyncio
import hmac
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from fastapi import Header, HTTPException, Request
from llm_inference.client import LLMClient
from llm_inference.config import LLMSettings
from llm_inference.logging import get_logger
logger = get_logger("dependencies")
def get_client(request: Request) -> LLMClient:
"""Get the LLM client from application state.
Args:
request: FastAPI request object.
Returns:
LLMClient: The LLM client instance.
"""
return request.app.state.client
def get_settings(request: Request) -> LLMSettings:
"""Get settings from application state.
Args:
request: FastAPI request object.
Returns:
LLMSettings: Application settings.
"""
return request.app.state.settings
def verify_bearer_token(
request: Request,
authorization: str | None = Header(default=None, alias="Authorization"),
) -> str | None:
"""Verify Bearer token authentication.
This dependency checks the Authorization header for a valid Bearer token.
If authentication is disabled (no tokens configured), returns None.
If authentication is enabled, validates the token and returns it.
Args:
request: FastAPI request object.
authorization: Authorization header value.
Returns:
The validated token if auth enabled, None if auth disabled.
Raises:
HTTPException: 401 if auth enabled and token is invalid/missing.
"""
settings = request.app.state.settings
# Auth disabled - allow all requests
if not settings.auth_enabled:
return None
# Auth enabled - validate token
if not authorization:
raise HTTPException(
status_code=401,
detail={
"error": "Authentication required",
"message": "Missing Authorization header",
},
headers={"WWW-Authenticate": "Bearer"},
)
# Parse Bearer token
parts = authorization.split()
if len(parts) != 2 or parts[0].lower() != "bearer":
raise HTTPException(
status_code=401,
detail={
"error": "Authentication required",
"message": "Invalid Authorization header format. Use: Bearer <token>",
},
headers={"WWW-Authenticate": "Bearer"},
)
token = parts[1]
# Constant-time comparison against all valid tokens
is_valid = any(
hmac.compare_digest(token.encode(), valid_token.encode())
for valid_token in settings.api_tokens
)
if not is_valid:
raise HTTPException(
status_code=401,
detail={
"error": "Authentication failed",
"message": "Invalid API token",
},
headers={"WWW-Authenticate": "Bearer"},
)
return token
class ConcurrencyLimiter:
"""Limits concurrent completions to prevent resource exhaustion.
Uses a semaphore to limit the number of concurrent completion requests.
Returns 503 Service Unavailable when limit is exceeded.
"""
def __init__(self, max_concurrent: int) -> None:
"""Initialize concurrency limiter.
Args:
max_concurrent: Maximum number of concurrent completions.
"""
self.max_concurrent = max_concurrent
self._semaphore = asyncio.Semaphore(max_concurrent)
self._counter_lock = asyncio.Lock()
self._current = 0
@property
def current_count(self) -> int:
"""Get current number of active requests."""
return self._current
@property
def available(self) -> int:
"""Get number of available slots."""
return self.max_concurrent - self._current
@asynccontextmanager
async def acquire(self, blocking: bool = False) -> AsyncIterator[None]:
"""Acquire a completion slot.
Args:
blocking: If True, wait for a slot. If False (default), return 503 immediately.
Yields:
None when slot is acquired.
Raises:
HTTPException: 503 if no slots available and blocking=False.
"""
if not blocking and self._semaphore.locked():
# Non-blocking mode and no slots available - reject immediately
# Note: small race window exists but semaphore still enforces limit
logger.warning(
"Concurrency limit reached: %d/%d active",
self._current,
self.max_concurrent,
)
raise HTTPException(
status_code=503,
detail={
"error": "Service temporarily unavailable",
"reason": "Too many concurrent requests",
"max_concurrent": self.max_concurrent,
"retry_after": 5,
},
headers={"Retry-After": "5"},
)
# Acquire semaphore (blocks if blocking=True and no slots, or immediate if available)
await self._semaphore.acquire()
# Update counter with lock for thread safety
async with self._counter_lock:
self._current += 1
try:
yield
finally:
async with self._counter_lock:
self._current -= 1
self._semaphore.release()
# Global concurrency limiter instance (initialized in app startup)
_concurrency_limiter: ConcurrencyLimiter | None = None
def init_concurrency_limiter(max_concurrent: int) -> ConcurrencyLimiter:
"""Initialize the global concurrency limiter.
Args:
max_concurrent: Maximum concurrent completions.
Returns:
ConcurrencyLimiter: The initialized limiter.
"""
global _concurrency_limiter
_concurrency_limiter = ConcurrencyLimiter(max_concurrent)
logger.info("Concurrency limiter initialized: max_concurrent=%d", max_concurrent)
return _concurrency_limiter
def get_concurrency_limiter() -> ConcurrencyLimiter:
"""Get the global concurrency limiter.
Returns:
ConcurrencyLimiter: The global limiter instance.
Raises:
RuntimeError: If limiter not initialized.
"""
if _concurrency_limiter is None:
raise RuntimeError("Concurrency limiter not initialized")
return _concurrency_limiter
async def require_completion_slot(request: Request) -> AsyncIterator[None]:
"""FastAPI dependency that acquires a completion slot.
Args:
request: FastAPI request.
Yields:
None when slot is acquired.
Raises:
HTTPException: 503 if no slots available.
"""
limiter = get_concurrency_limiter()
async with limiter.acquire():
yield

230
ai_platform/modules/llm-inference/src/llm_inference/api/middleware.py Normal file
View file

@ -0,0 +1,230 @@
"""FastAPI middleware for request handling and rate limiting.
Provides:
- RequestIdMiddleware: Extracts/generates request IDs and propagates via context
- RateLimitMiddleware: Token bucket rate limiting with 429 responses
"""
import time
import uuid
from collections.abc import Awaitable, Callable
from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from llm_inference.logging import get_logger, set_request_id
logger = get_logger("middleware")
class RequestIdMiddleware(BaseHTTPMiddleware):
"""Middleware to extract or generate request IDs.
Extracts X-Request-ID from incoming headers or generates a new UUID.
Sets the request ID in context for logging and returns it in response headers.
"""
async def dispatch(
self,
request: Request,
call_next: Callable[[Request], Awaitable[Response]],
) -> Response:
"""Process request with request ID tracking.
Args:
request: Incoming HTTP request.
call_next: Next middleware/handler in chain.
Returns:
Response: HTTP response with X-Request-ID header.
"""
# Extract from header or generate new
request_id = request.headers.get("X-Request-ID")
if not request_id:
request_id = str(uuid.uuid4())
# Set in context for logging
set_request_id(request_id)
# Store in request state for handlers
request.state.request_id = request_id
try:
response = await call_next(request)
response.headers["X-Request-ID"] = request_id
return response
finally:
# Clear context after request
set_request_id(None)
class TokenBucket:
"""Token bucket for rate limiting with optional live tuning.
Implements a simple token bucket algorithm where tokens are added
at a fixed rate up to a maximum burst size. If a runtime_config and
rate/burst keys are supplied, the bucket reads the latest override
on every acquire() — making the limit live-tunable from dashboard
without restart.
"""
def __init__(
self,
rate: float,
burst: int,
runtime_config: object | None = None,
rate_key: str | None = None,
burst_key: str | None = None,
) -> None:
"""Initialize token bucket.
Args:
rate: Tokens added per second (fallback / initial).
burst: Maximum tokens (bucket capacity, fallback / initial).
runtime_config: Optional RuntimeConfigClient. When set together
with rate_key / burst_key, the bucket pulls live values on
each acquire() and falls back to the constructor values.
rate_key: Dashboard config key for live rate (e.g., "llm.rate_limit.rps").
burst_key: Dashboard config key for live burst.
"""
self._fallback_rate = rate
self._fallback_burst = burst
self.rate = rate
self.burst = burst
self.tokens = float(burst)
self.last_update = time.monotonic()
self._runtime_config = runtime_config
self._rate_key = rate_key
self._burst_key = burst_key
def _refresh_from_config(self) -> None:
"""Pull latest rate/burst from runtime_config if wired."""
if self._runtime_config is None:
return
try:
new_rate = self._runtime_config.get_float(self._rate_key, self._fallback_rate) if self._rate_key else self._fallback_rate
new_burst = self._runtime_config.get_int(self._burst_key, self._fallback_burst) if self._burst_key else self._fallback_burst
except Exception:
return
if new_rate != self.rate or new_burst != self.burst:
# On burst increase, top up; on decrease, clamp tokens to new burst
self.rate = new_rate
self.burst = new_burst
self.tokens = min(self.tokens, float(new_burst))
def acquire(self) -> bool:
"""Try to acquire a token.
Returns:
bool: True if token acquired, False if rate limited.
"""
self._refresh_from_config()
now = time.monotonic()
elapsed = now - self.last_update
self.last_update = now
# Add tokens based on elapsed time
self.tokens = min(self.burst, self.tokens + elapsed * self.rate)
if self.tokens >= 1:
self.tokens -= 1
return True
return False
def retry_after(self) -> float:
"""Calculate seconds until a token is available.
Returns:
float: Seconds to wait before retrying.
"""
if self.tokens >= 1:
return 0.0
tokens_needed = 1 - self.tokens
return tokens_needed / self.rate
class RateLimitMiddleware(BaseHTTPMiddleware):
"""Middleware for global rate limiting using token bucket algorithm.
Returns 429 Too Many Requests with Retry-After header when exceeded.
WARNING: This is per-process rate limiting. In multi-replica deployments
(e.g., Kubernetes with multiple pods), each replica maintains its own
independent rate limit. A configured limit of 10 RPS with 5 replicas
effectively allows 50 RPS total.
For distributed rate limiting in production, use an external solution:
- Redis-based rate limiting (e.g., redis-rate-limiter)
- API gateway rate limiting (e.g., Kong, nginx)
- Cloud provider rate limiting (e.g., AWS API Gateway)
"""
def __init__(
self,
app: object,
rate: float = 10.0,
burst: int = 20,
exclude_paths: list[str] | None = None,
runtime_config: object | None = None,
rate_key: str | None = None,
burst_key: str | None = None,
) -> None:
"""Initialize rate limiter.
Args:
app: FastAPI application.
rate: Requests per second limit (fallback when no runtime_config).
burst: Maximum burst size (fallback).
exclude_paths: Paths to exclude from rate limiting.
runtime_config: Optional RuntimeConfigClient to enable live rate tuning.
rate_key: Dashboard config key (e.g., "llm.rate_limit.rps").
burst_key: Dashboard config key (e.g., "llm.rate_limit.burst").
"""
super().__init__(app)
self.bucket = TokenBucket(
rate=rate,
burst=burst,
runtime_config=runtime_config,
rate_key=rate_key,
burst_key=burst_key,
)
self.exclude_paths = exclude_paths or ["/health", "/ready"]
async def dispatch(
self,
request: Request,
call_next: Callable[[Request], Awaitable[Response]],
) -> Response:
"""Process request with rate limiting.
Args:
request: Incoming HTTP request.
call_next: Next middleware/handler in chain.
Returns:
Response: HTTP response or 429 if rate limited.
"""
# Skip rate limiting for excluded paths
if request.url.path in self.exclude_paths:
return await call_next(request)
# Try to acquire token
if not self.bucket.acquire():
retry_after = self.bucket.retry_after()
logger.warning(
"Rate limit exceeded for %s %s, retry_after=%.2f",
request.method,
request.url.path,
retry_after,
)
return JSONResponse(
status_code=429,
content={
"detail": "Too many requests",
"retry_after": retry_after,
},
headers={"Retry-After": str(int(retry_after) + 1)},
)
return await call_next(request)

5
ai_platform/modules/llm-inference/src/llm_inference/api/routes/__init__.py Normal file
View file

@ -0,0 +1,5 @@
"""API routes for LLM inference."""
from llm_inference.api.routes import completions, health, info, models
__all__ = ["completions", "health", "info", "models"]

221
ai_platform/modules/llm-inference/src/llm_inference/api/routes/completions.py Normal file
View file

@ -0,0 +1,221 @@
"""Chat completion routes with streaming support."""
import json
from collections.abc import AsyncGenerator
from fastapi import APIRouter, Depends, HTTPException
from sse_starlette.sse import EventSourceResponse
from llm_inference.api.dependencies import (
ConcurrencyLimiter,
get_client,
get_concurrency_limiter,
verify_bearer_token,
)
from llm_inference.client import LLMClient
from llm_inference.exceptions import (
BackendNotAvailableError,
BackendNotEnabledError,
CompletionError,
)
from llm_inference.schemas import (
CompletionRequest,
CompletionResponse,
TextCompletionRequest,
TextCompletionResponse,
)
router = APIRouter(dependencies=[Depends(verify_bearer_token)])
@router.post(
"/chat/completions",
response_model=None, # Disable auto-generation due to Union with EventSourceResponse
)
async def chat_completions(
request: CompletionRequest,
client: LLMClient = Depends(get_client),
) -> CompletionResponse | EventSourceResponse:
"""OpenAI-compatible chat completions endpoint.
Supports both streaming (SSE) and non-streaming responses.
Backend can be overridden per-request via the `backend` field.
Concurrency is limited per-process. For streaming requests, the slot is
held for the entire duration of the stream.
Args:
request: Completion request with messages, model, and options.
client: LLM client instance.
Returns:
CompletionResponse or EventSourceResponse for streaming.
Raises:
HTTPException: If the request fails or concurrency limit exceeded (503).
"""
limiter = get_concurrency_limiter()
if request.stream:
# For streaming: wrap generator to hold slot throughout entire stream
# This is critical - the slot must be held until streaming completes
return EventSourceResponse(
_stream_with_slot(_stream_generator(client, request), limiter),
media_type="text/event-stream",
ping=15, # Send ping every 15s to detect dead connections
)
# For non-streaming: acquire slot, run completion, release slot
async with limiter.acquire():
try:
kwargs = _build_completion_kwargs(request)
response = await client.complete(
messages=request.messages,
model=request.model,
backend=request.backend,
**kwargs,
)
return response
except (BackendNotAvailableError, BackendNotEnabledError) as e:
raise HTTPException(status_code=400, detail=str(e)) from None
except CompletionError as e:
raise HTTPException(status_code=500, detail=str(e)) from None
@router.post("/completions", response_model=TextCompletionResponse)
async def text_completions(
request: TextCompletionRequest,
client: LLMClient = Depends(get_client),
) -> TextCompletionResponse:
"""OpenAI-compatible legacy text completions endpoint (/v1/completions).
Non-streaming. Routes to the model's backend (vLLM primary, cloud via
LiteLLM). Backends that do not support text completion return 501.
"""
limiter = get_concurrency_limiter()
async with limiter.acquire():
try:
kwargs: dict[str, object] = {"temperature": request.temperature}
if request.max_tokens is not None:
kwargs["max_tokens"] = request.max_tokens
if request.top_p is not None:
kwargs["top_p"] = request.top_p
if request.frequency_penalty is not None:
kwargs["frequency_penalty"] = request.frequency_penalty
if request.presence_penalty is not None:
kwargs["presence_penalty"] = request.presence_penalty
if request.stop is not None:
stop_seq = (
[request.stop]
if isinstance(request.stop, str)
else list(request.stop)
)
stop_seq = [s for s in stop_seq if s]
if stop_seq:
kwargs["stop"] = stop_seq
return await client.text_complete(
prompt=request.prompt,
model=request.model,
backend=request.backend,
**kwargs,
)
except NotImplementedError as e:
raise HTTPException(status_code=501, detail=str(e)) from None
except (BackendNotAvailableError, BackendNotEnabledError) as e:
raise HTTPException(status_code=400, detail=str(e)) from None
except CompletionError as e:
raise HTTPException(status_code=500, detail=str(e)) from None
async def _stream_generator(
client: LLMClient,
request: CompletionRequest,
) -> AsyncGenerator[dict[str, str], None]:
"""Generate SSE events for streaming response.
Args:
client: LLM client instance.
request: Completion request.
Yields:
SSE event dictionaries with data field.
"""
try:
kwargs = _build_completion_kwargs(request)
async for chunk in client.stream(
messages=request.messages,
model=request.model,
backend=request.backend,
**kwargs,
):
yield {"data": json.dumps(chunk.model_dump())}
# Send [DONE] marker to signal end of stream
yield {"data": "[DONE]"}
except (BackendNotAvailableError, BackendNotEnabledError, CompletionError) as e:
# Send error event followed by [DONE] to signal stream termination
yield {
"event": "error",
"data": json.dumps({"error": str(e)}),
}
yield {"data": "[DONE]"}
async def _stream_with_slot(
generator: AsyncGenerator[dict[str, str], None],
limiter: ConcurrencyLimiter,
) -> AsyncGenerator[dict[str, str], None]:
"""Wrap a stream generator to hold a concurrency slot throughout streaming.
This ensures the concurrency limiter slot is held for the entire duration
of the SSE stream, not just until the EventSourceResponse is returned.
Args:
generator: The underlying stream generator.
limiter: Concurrency limiter to acquire slot from.
Yields:
SSE event dictionaries from the wrapped generator.
"""
async with limiter.acquire():
async for item in generator:
yield item
def _build_completion_kwargs(request: CompletionRequest) -> dict[str, object]:
"""Build kwargs dict from completion request.
Args:
request: Completion request.
Returns:
dict: Keyword arguments for completion call.
"""
kwargs: dict[str, object] = {
"temperature": request.temperature,
}
if request.max_tokens is not None:
kwargs["max_tokens"] = request.max_tokens
if request.top_p is not None:
kwargs["top_p"] = request.top_p
if request.frequency_penalty is not None:
kwargs["frequency_penalty"] = request.frequency_penalty
if request.presence_penalty is not None:
kwargs["presence_penalty"] = request.presence_penalty
if request.stop is not None:
# Normalize to list and filter out empty strings
# Some backends choke on empty stop sequences
stop_seq = (
[request.stop] if isinstance(request.stop, str) else list(request.stop)
)
stop_seq = [s for s in stop_seq if s] # Filter empty strings
if stop_seq:
kwargs["stop"] = stop_seq
return kwargs

78
ai_platform/modules/llm-inference/src/llm_inference/api/routes/health.py Normal file
View file

@ -0,0 +1,78 @@
"""Health check routes."""
from fastapi import APIRouter, Depends
from llm_inference.api.dependencies import get_client
from llm_inference.client import LLMClient
from llm_inference.logging import get_logger
from llm_inference.schemas import BackendHealth, HealthResponse, ReadinessResponse
router = APIRouter()
logger = get_logger("routes.health")
@router.get("/health", response_model=HealthResponse)
async def health_check(
client: LLMClient = Depends(get_client),
) -> HealthResponse:
"""Health check endpoint.
Returns the overall health status and per-backend health status.
Returns:
HealthResponse: Overall and per-backend health status.
"""
backend_health: list[BackendHealth] = []
for backend_type in client.list_backends():
backend = client.registry.get(backend_type)
is_healthy = await backend.health_check()
backend_health.append(
BackendHealth(
name=backend_type.value,
healthy=is_healthy,
)
)
# Determine overall status
all_healthy = all(b.healthy for b in backend_health)
any_healthy = any(b.healthy for b in backend_health)
if all_healthy:
status = "healthy"
elif any_healthy:
status = "degraded"
else:
status = "unhealthy"
return HealthResponse(status=status, backends=backend_health)
@router.get("/ready", response_model=ReadinessResponse)
async def readiness_check(
client: LLMClient = Depends(get_client),
) -> ReadinessResponse:
"""Readiness probe for Kubernetes.
Verifies that the default backend is healthy and able to serve requests.
This is the critical check for load balancer routing.
Returns:
ReadinessResponse: Readiness status (ready if default backend is healthy).
"""
try:
# Get the default backend and check its health
default_backend = client.registry.get()
is_healthy = await default_backend.health_check()
if not is_healthy:
logger.warning(
"Readiness check failed: default backend '%s' is unhealthy",
default_backend.name,
)
return ReadinessResponse(ready=is_healthy)
except Exception as e:
logger.error("Readiness check failed with error: %s", str(e))
return ReadinessResponse(ready=False)

468
ai_platform/modules/llm-inference/src/llm_inference/api/routes/info.py Normal file
View file

@ -0,0 +1,468 @@
"""Component information endpoint for service catalog."""
from fastapi import APIRouter, Depends
from llm_inference.api.dependencies import get_client
from llm_inference.client import LLMClient
from llm_inference.config import SettingsCache
router = APIRouter()
@router.get("/v1/info")
async def get_component_info(
client: LLMClient = Depends(get_client),
) -> dict:
"""Get component information for service catalog.
Returns complete metadata about this service including:
- Resource information (component metadata)
- Available models (from all enabled backends)
- Available functions (API endpoints)
This endpoint is used by the catalog-api to aggregate service information
and by backend systems to populate the catalog database.
Returns:
dict: Component information matching catalog.resources, catalog.models,
and catalog.functions schemas.
"""
settings = SettingsCache.get()
# Determine which backends are available
backends = ["litellm"] # Always available
if settings.enable_vllm:
backends.append("vllm")
if settings.enable_llamacpp:
backends.append("llamacpp")
# Build resource information (maps to catalog.resources)
resource = {
"name": "LLM Inference Gateway",
"slug": "llm-inference",
"resource_type": "api_service",
"provider": "internal",
"base_url": f"http://didiAI-llm-api:{settings.port}",
"configuration": {
"version": "1.0.0",
"port": settings.port,
"external_url": settings.external_url,
"default_backend": settings.default_backend,
"backends": backends,
"default_model": settings.default_model,
},
"authentication": {
"type": "bearer",
"required": bool(settings.api_tokens),
"env_var": "LLM_API_TOKENS",
},
"headers": {
"Content-Type": "application/json",
"Accept": "application/json",
},
"rate_limits": {
"requests_per_second": settings.rate_limit_rps,
"burst": settings.rate_limit_burst,
"concurrent": settings.max_concurrent_completions,
},
"cost_tracking": {
"enabled": False,
},
"tags": ["llm", "inference", "openai-compatible", "gateway", "nlp"],
"is_active": True,
"metadata": {
"category": "nlp",
"gpu_required": settings.enable_vllm,
"status": "healthy",
"vllm_base_url": settings.vllm_base_url if settings.enable_vllm else None,
"llamacpp_base_url": (
settings.llamacpp_base_url if settings.enable_llamacpp else None
),
},
}
# Collect models from all backends
models = []
# Get models from enabled backends
try:
all_models = await client.list_models()
for model_info in all_models:
model_entry = {
"name": model_info.id,
"slug": model_info.id.lower().replace("/", "-").replace("_", "-"),
"provider": _get_provider_from_model_id(model_info.id),
"model_type": _get_model_type(model_info),
"capabilities": model_info.capabilities or [],
"configuration": {
"backend": model_info.backend,
"context_length": model_info.context_length,
"loaded": model_info.loaded,
},
"endpoint": _get_model_endpoint(model_info, settings),
"api_key_ref": None,
"tags": _get_model_tags(model_info),
"is_active": model_info.loaded,
"metadata": {
"backend": model_info.backend,
"model_id": model_info.id,
},
}
models.append(model_entry)
except Exception:
# If model listing fails, continue with empty models list
pass
# Define available functions (maps to catalog.functions)
functions = [
{
"name": "Chat Completions",
"slug": "llm-chat-completions",
"category": "completion",
"description": (
"Create a chat completion with optional streaming support. "
"OpenAI-compatible endpoint supporting multiple backends."
),
"input_schema": {
"type": "object",
"properties": {
"messages": {
"type": "array",
"description": "List of messages (1-1000)",
"items": {
"type": "object",
"properties": {
"role": {
"type": "string",
"enum": ["system", "user", "assistant", "function", "tool"],
},
"content": {"type": "string"},
},
"required": ["role"],
},
},
"model": {"type": "string", "description": "Model identifier"},
"temperature": {
"type": "number",
"minimum": 0.0,
"maximum": 2.0,
"default": 0.7,
},
"max_tokens": {"type": "integer", "minimum": 1, "maximum": 1000000},
"stream": {"type": "boolean", "default": False},
"backend": {
"type": "string",
"enum": backends,
"description": "Backend override",
},
"top_p": {"type": "number", "minimum": 0.0, "maximum": 1.0},
"frequency_penalty": {"type": "number", "minimum": -2.0, "maximum": 2.0},
"presence_penalty": {"type": "number", "minimum": -2.0, "maximum": 2.0},
"stop": {
"oneOf": [
{"type": "string"},
{"type": "array", "items": {"type": "string"}},
]
},
},
"required": ["messages", "model"],
},
"output_schema": {
"type": "object",
"properties": {
"id": {"type": "string"},
"object": {"type": "string", "const": "chat.completion"},
"created": {"type": "integer"},
"model": {"type": "string"},
"choices": {
"type": "array",
"items": {
"type": "object",
"properties": {
"index": {"type": "integer"},
"message": {
"type": "object",
"properties": {
"role": {"type": "string"},
"content": {"type": "string"},
},
},
"finish_reason": {"type": "string"},
},
},
},
"usage": {
"type": "object",
"properties": {
"prompt_tokens": {"type": "integer"},
"completion_tokens": {"type": "integer"},
"total_tokens": {"type": "integer"},
},
},
"backend": {"type": "string"},
},
},
"implementation": {
"method": "POST",
"path": "/v1/chat/completions",
"content_type": "application/json",
"timeout": 120,
"supports_streaming": True,
},
"endpoint": f"http://localhost:{settings.port}/v1/chat/completions",
"tags": ["llm", "chat", "openai", "streaming"],
"is_active": True,
"metadata": {
"rate_limited": True,
"auth_required": bool(settings.api_tokens),
},
},
{
"name": "List Models",
"slug": "llm-list-models",
"category": "discovery",
"description": "List all available models across all enabled backends",
"input_schema": {
"type": "object",
"properties": {
"backend": {
"type": "string",
"enum": backends,
"description": "Optional backend filter",
}
},
},
"output_schema": {
"type": "object",
"properties": {
"object": {"type": "string", "const": "list"},
"data": {
"type": "array",
"items": {
"type": "object",
"properties": {
"id": {"type": "string"},
"backend": {"type": "string"},
"loaded": {"type": "boolean"},
"context_length": {"type": "integer"},
"capabilities": {
"type": "array",
"items": {"type": "string"},
},
},
},
},
},
},
"implementation": {
"method": "GET",
"path": "/v1/models",
"timeout": 10,
},
"endpoint": f"http://localhost:{settings.port}/v1/models",
"tags": ["discovery", "models"],
"is_active": True,
"metadata": {},
},
{
"name": "List Backends",
"slug": "llm-list-backends",
"category": "discovery",
"description": "List all available backends",
"input_schema": {"type": "object", "properties": {}},
"output_schema": {
"type": "object",
"properties": {
"backends": {
"type": "array",
"items": {"type": "string"},
}
},
},
"implementation": {
"method": "GET",
"path": "/v1/backends",
"timeout": 5,
},
"endpoint": f"http://localhost:{settings.port}/v1/backends",
"tags": ["discovery"],
"is_active": True,
"metadata": {},
},
{
"name": "Load Model",
"slug": "llm-load-model",
"category": "management",
"description": "Load a model on a local backend (vLLM or llama.cpp only)",
"input_schema": {
"type": "object",
"properties": {
"model": {"type": "string"},
"backend": {"type": "string", "enum": ["vllm", "llamacpp"]},
},
"required": ["model", "backend"],
},
"output_schema": {
"type": "object",
"properties": {
"success": {"type": "boolean"},
"model": {"type": "string"},
"backend": {"type": "string"},
"message": {"type": "string"},
},
},
"implementation": {
"method": "POST",
"path": "/v1/models/load",
"content_type": "application/json",
"timeout": 60,
},
"endpoint": f"http://localhost:{settings.port}/v1/models/load",
"tags": ["management", "models"],
"is_active": settings.enable_vllm or settings.enable_llamacpp,
"metadata": {},
},
{
"name": "Unload Model",
"slug": "llm-unload-model",
"category": "management",
"description": "Unload a model from a local backend (vLLM or llama.cpp only)",
"input_schema": {
"type": "object",
"properties": {
"model": {"type": "string"},
"backend": {"type": "string", "enum": ["vllm", "llamacpp"]},
},
"required": ["model", "backend"],
},
"output_schema": {
"type": "object",
"properties": {
"success": {"type": "boolean"},
"model": {"type": "string"},
"backend": {"type": "string"},
"message": {"type": "string"},
},
},
"implementation": {
"method": "POST",
"path": "/v1/models/unload",
"content_type": "application/json",
"timeout": 30,
},
"endpoint": f"http://localhost:{settings.port}/v1/models/unload",
"tags": ["management", "models"],
"is_active": settings.enable_vllm or settings.enable_llamacpp,
"metadata": {},
},
]
return {
"resource": resource,
"models": models,
"functions": functions,
}
def _get_provider_from_model_id(model_id: str) -> str:
"""Extract provider name from model ID.
Args:
model_id: Model identifier (e.g., "openai/gpt-4", "meta-llama/Llama-2-7b")
Returns:
str: Provider name (e.g., "openai", "meta-llama", "unknown")
"""
if "/" in model_id:
return model_id.split("/")[0]
if "gpt" in model_id.lower():
return "openai"
if "claude" in model_id.lower():
return "anthropic"
if "llama" in model_id.lower():
return "meta"
if "qwen" in model_id.lower():
return "qwen"
if "mistral" in model_id.lower():
return "mistralai"
return "unknown"
def _get_model_type(model_info) -> str:
"""Determine model type from capabilities.
Args:
model_info: Model information object
Returns:
str: Model type (llm, vision, audio, etc.)
"""
capabilities = model_info.capabilities or []
if "vision" in capabilities or "multimodal" in capabilities:
return "vision"
if "audio" in capabilities:
return "audio"
if "embedding" in capabilities:
return "embedding"
return "llm"
def _get_model_tags(model_info) -> list[str]:
"""Generate tags for a model.
Args:
model_info: Model information object
Returns:
list[str]: List of tags
"""
tags = []
# Add model type tags
if "gpt" in model_info.id.lower():
tags.extend(["gpt", "openai"])
if "claude" in model_info.id.lower():
tags.extend(["claude", "anthropic"])
if "llama" in model_info.id.lower():
tags.extend(["llama", "meta"])
if "qwen" in model_info.id.lower():
tags.extend(["qwen"])
if "mistral" in model_info.id.lower():
tags.extend(["mistral"])
# Add capability tags
capabilities = model_info.capabilities or []
tags.extend(capabilities)
# Add backend tag
tags.append(model_info.backend)
# Add size tags if detectable
model_id_lower = model_info.id.lower()
if "7b" in model_id_lower:
tags.append("7b")
elif "13b" in model_id_lower:
tags.append("13b")
elif "30b" in model_id_lower:
tags.append("30b")
elif "70b" in model_id_lower:
tags.append("70b")
elif "120b" in model_id_lower:
tags.append("120b")
return list(set(tags)) # Remove duplicates
def _get_model_endpoint(model_info, settings) -> str:
"""Get the endpoint URL for a model.
Args:
model_info: Model information object
settings: Application settings
Returns:
str: Endpoint URL
"""
# Default to gateway endpoint
return f"http://localhost:{settings.port}/v1/chat/completions"

134
ai_platform/modules/llm-inference/src/llm_inference/api/routes/models.py Normal file
View file

@ -0,0 +1,134 @@
"""Model management routes."""
from fastapi import APIRouter, Depends, HTTPException
from llm_inference.api.dependencies import get_client, verify_bearer_token
from llm_inference.client import LLMClient
from llm_inference.exceptions import (
BackendNotAvailableError,
BackendNotEnabledError,
ModelLoadError,
)
from llm_inference.schemas import (
BackendListResponse,
ModelListResponse,
ModelLoadRequest,
ModelLoadResponse,
)
from llm_inference.types import BackendType
router = APIRouter(dependencies=[Depends(verify_bearer_token)])
@router.get("/models", response_model=ModelListResponse)
async def list_models(
backend: str | None = None,
client: LLMClient = Depends(get_client),
) -> ModelListResponse:
"""List available models.
Args:
backend: Optional backend filter. If not specified, returns models
from all available backends.
client: LLM client instance.
Returns:
ModelListResponse: List of available models.
Raises:
HTTPException: If the specified backend is not available.
"""
try:
backend_type = BackendType(backend) if backend else None
models = await client.list_models(backend=backend_type)
return ModelListResponse(data=models)
except ValueError:
raise HTTPException(
status_code=400, detail=f"Invalid backend: {backend}"
) from None
except (BackendNotAvailableError, BackendNotEnabledError) as e:
raise HTTPException(status_code=400, detail=str(e)) from None
@router.post("/models/load", response_model=ModelLoadResponse)
async def load_model(
request: ModelLoadRequest,
client: LLMClient = Depends(get_client),
) -> ModelLoadResponse:
"""Load a model on the specified backend.
This is only supported by local backends (vLLM, llama.cpp).
Args:
request: Model load request.
client: LLM client instance.
Returns:
ModelLoadResponse: Result of the load operation.
Raises:
HTTPException: If loading fails or backend doesn't support it.
"""
try:
success = await client.load_model(request.model, request.backend)
return ModelLoadResponse(
success=success,
model=request.model,
backend=request.backend.value,
message="Model loaded successfully" if success else "Model load failed",
)
except NotImplementedError as e:
raise HTTPException(status_code=400, detail=str(e)) from None
except (BackendNotAvailableError, BackendNotEnabledError) as e:
raise HTTPException(status_code=400, detail=str(e)) from None
except ModelLoadError as e:
raise HTTPException(status_code=500, detail=str(e)) from None
@router.post("/models/unload", response_model=ModelLoadResponse)
async def unload_model(
request: ModelLoadRequest,
client: LLMClient = Depends(get_client),
) -> ModelLoadResponse:
"""Unload a model from the specified backend.
This is only supported by local backends (vLLM, llama.cpp).
Args:
request: Model unload request.
client: LLM client instance.
Returns:
ModelLoadResponse: Result of the unload operation.
Raises:
HTTPException: If unloading fails or backend doesn't support it.
"""
try:
success = await client.unload_model(request.model, request.backend)
return ModelLoadResponse(
success=success,
model=request.model,
backend=request.backend.value,
message="Model unloaded successfully" if success else "Model unload failed",
)
except NotImplementedError as e:
raise HTTPException(status_code=400, detail=str(e)) from None
except (BackendNotAvailableError, BackendNotEnabledError) as e:
raise HTTPException(status_code=400, detail=str(e)) from None
@router.get("/backends", response_model=BackendListResponse)
async def list_backends(
client: LLMClient = Depends(get_client),
) -> BackendListResponse:
"""List available backends.
Args:
client: LLM client instance.
Returns:
BackendListResponse: List of available backend names.
"""
backends = [b.value for b in client.list_backends()]
return BackendListResponse(backends=backends)

9
ai_platform/modules/llm-inference/src/llm_inference/backends/__init__.py Normal file
View file

@ -0,0 +1,9 @@
"""Backend implementations for LLM inference."""
from llm_inference.backends.base import LLMBackend
from llm_inference.backends.registry import BackendRegistry
__all__ = [
"BackendRegistry",
"LLMBackend",
]

146
ai_platform/modules/llm-inference/src/llm_inference/backends/base.py Normal file
View file

@ -0,0 +1,146 @@
"""Abstract base class for LLM backends."""
from abc import ABC, abstractmethod
from collections.abc import AsyncIterator
from llm_inference.schemas import (
CompletionChunk,
CompletionResponse,
TextCompletionResponse,
)
from llm_inference.types import ChatMessage, ModelInfo
class LLMBackend(ABC):
"""Abstract base class for LLM backends.
All backend implementations must inherit from this class and implement
the required abstract methods.
"""
@property
@abstractmethod
def name(self) -> str:
"""Backend identifier name.
Returns:
str: Unique name for this backend (e.g., 'litellm', 'vllm').
"""
...
@abstractmethod
async def complete(
self,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> CompletionResponse:
"""Generate a chat completion.
Args:
messages: List of chat messages in the conversation.
model: Model identifier to use for completion.
**kwargs: Additional parameters passed to the backend.
Returns:
CompletionResponse: The completion response.
Raises:
CompletionError: If the completion request fails.
"""
...
@abstractmethod
async def stream(
self,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> AsyncIterator[CompletionChunk]:
"""Generate a streaming chat completion.
Args:
messages: List of chat messages in the conversation.
model: Model identifier to use for completion.
**kwargs: Additional parameters passed to the backend.
Yields:
CompletionChunk: Streaming chunks of the completion.
Raises:
CompletionError: If the completion request fails.
"""
...
yield # pragma: no cover (makes this a generator)
async def text_complete(
self,
prompt: str | list[str],
model: str,
**kwargs: object,
) -> TextCompletionResponse:
"""Generate a (legacy) text completion for ``prompt``.
Concrete by default so backends that cannot serve /v1/completions are
not forced to implement it; override in OpenAI-compatible backends.
Raises:
NotImplementedError: If the backend does not support text completion.
"""
raise NotImplementedError(
f"Backend '{self.name}' does not support text completions"
)
@abstractmethod
async def list_models(self) -> list[ModelInfo]:
"""List available models for this backend.
Returns:
list[ModelInfo]: List of available models.
"""
...
async def load_model(self, model: str) -> bool:
"""Load a model on this backend.
This is optional and only supported by local backends (vLLM, llama.cpp).
Args:
model: Model identifier to load.
Returns:
bool: True if the model was loaded successfully.
Raises:
NotImplementedError: If the backend doesn't support model loading.
ModelLoadError: If loading the model fails.
"""
raise NotImplementedError(
f"Backend '{self.name}' does not support dynamic model loading"
)
async def unload_model(self, model: str) -> bool:
"""Unload a model from this backend.
This is optional and only supported by local backends (vLLM, llama.cpp).
Args:
model: Model identifier to unload.
Returns:
bool: True if the model was unloaded successfully.
Raises:
NotImplementedError: If the backend doesn't support model unloading.
"""
raise NotImplementedError(
f"Backend '{self.name}' does not support dynamic model unloading"
)
async def health_check(self) -> bool:
"""Check if this backend is healthy and ready to serve requests.
Returns:
bool: True if the backend is healthy.
"""
return True

518
ai_platform/modules/llm-inference/src/llm_inference/backends/litellm_backend.py Normal file
View file

@ -0,0 +1,518 @@
"""LiteLLM backend implementation.
LiteLLM provides a unified interface for 100+ LLM providers including:
- OpenAI
- Anthropic
- OpenRouter
- Azure OpenAI
- Google AI (Gemini)
- AWS Bedrock
- And many more
"""
import logging
import time
from collections.abc import AsyncIterator
import httpx
import litellm
import litellm.exceptions
from litellm import acompletion, atext_completion
from llm_inference.backends.base import LLMBackend
from llm_inference.config import LLMSettings
from llm_inference.exceptions import (
CompletionError,
LLMRateLimitError,
LLMTimeoutError,
)
from llm_inference.retry import retry_with_backoff
from llm_inference.schemas import (
CompletionChunk,
CompletionResponse,
TextChoice,
TextCompletionResponse,
)
from llm_inference.types import (
ChatMessage,
Choice,
Delta,
ModelInfo,
StreamChoice,
Usage,
)
from llm_inference.utils import safe_close_stream
logger = logging.getLogger("llm_inference.backends.litellm")
class LiteLLMBackend(LLMBackend):
"""LiteLLM backend supporting 100+ LLM providers.
This is the default backend and handles most cloud-based LLM providers
including OpenRouter for accessing various models through a single API.
"""
# Default cache TTL for model list (1 hour)
_MODELS_CACHE_TTL: float = 3600.0
def __init__(self, settings: LLMSettings) -> None:
"""Initialize LiteLLM backend.
Args:
settings: Application settings with API keys.
"""
self._settings = settings
self._configure_litellm()
# Model list cache: (timestamp, models)
self._models_cache: tuple[float, list[ModelInfo]] | None = None
def _configure_litellm(self) -> None:
"""Configure LiteLLM with API keys from settings."""
# Set API keys if provided
if self._settings.openrouter_api_key:
litellm.openrouter_api_key = self._settings.openrouter_api_key
if self._settings.openai_api_key:
litellm.openai_api_key = self._settings.openai_api_key
if self._settings.anthropic_api_key:
litellm.anthropic_api_key = self._settings.anthropic_api_key
# Suppress verbose logging
litellm.suppress_debug_info = True
@property
def name(self) -> str:
"""Backend identifier name."""
return "litellm"
async def complete(
self,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> CompletionResponse:
"""Generate a chat completion via LiteLLM.
Args:
messages: List of chat messages.
model: Model identifier (e.g., 'gpt-4', 'claude-3-opus').
**kwargs: Additional parameters (temperature, max_tokens, etc.).
Returns:
CompletionResponse: The completion response.
Raises:
CompletionError: If the request fails after retries.
LLMRateLimitError: If rate limited.
LLMTimeoutError: If the request times out.
"""
async def _do_complete() -> CompletionResponse:
response = await acompletion(
model=model,
messages=[m.model_dump(exclude_none=True) for m in messages],
timeout=self._settings.request_timeout,
**kwargs,
)
# Convert response to our schema
choices = []
for choice in response.choices:
choices.append(
Choice(
index=choice.index,
message=ChatMessage(
role=choice.message.role,
content=choice.message.content,
),
finish_reason=choice.finish_reason,
)
)
usage = None
if response.usage:
usage = Usage(
prompt_tokens=response.usage.prompt_tokens,
completion_tokens=response.usage.completion_tokens,
total_tokens=response.usage.total_tokens,
)
return CompletionResponse(
id=response.id,
created=response.created,
model=response.model,
choices=choices,
usage=usage,
backend=self.name,
)
return await retry_with_backoff(
_do_complete,
max_retries=self._settings.max_retries,
min_wait=self._settings.retry_min_wait,
max_wait=self._settings.retry_max_wait,
backend=self.name,
model=model,
)
async def text_complete(
self,
prompt: str | list[str],
model: str,
**kwargs: object,
) -> TextCompletionResponse:
"""Generate a legacy text completion via LiteLLM (cloud providers)."""
async def _do_text() -> TextCompletionResponse:
response = await atext_completion(
model=model,
prompt=prompt,
timeout=self._settings.request_timeout,
**kwargs,
)
choices = [
TextChoice(
index=getattr(c, "index", i),
text=getattr(c, "text", ""),
finish_reason=getattr(c, "finish_reason", None),
)
for i, c in enumerate(response.choices)
]
usage = None
if response.usage:
usage = Usage(
prompt_tokens=response.usage.prompt_tokens,
completion_tokens=response.usage.completion_tokens,
total_tokens=response.usage.total_tokens,
)
return TextCompletionResponse(
id=response.id,
created=response.created,
model=response.model,
choices=choices,
usage=usage,
backend=self.name,
)
return await retry_with_backoff(
_do_text,
max_retries=self._settings.max_retries,
min_wait=self._settings.retry_min_wait,
max_wait=self._settings.retry_max_wait,
backend=self.name,
model=model,
)
async def stream(
self,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> AsyncIterator[CompletionChunk]:
"""Generate a streaming chat completion via LiteLLM.
Args:
messages: List of chat messages.
model: Model identifier.
**kwargs: Additional parameters.
Yields:
CompletionChunk: Streaming chunks of the completion.
Raises:
CompletionError: If the request fails.
LLMRateLimitError: If rate limited.
LLMTimeoutError: If the request times out.
"""
response = None
try:
response = await acompletion(
model=model,
messages=[m.model_dump(exclude_none=True) for m in messages],
stream=True,
timeout=self._settings.request_timeout,
**kwargs,
)
async for chunk in response:
choices = []
if chunk.choices:
for choice in chunk.choices:
delta = Delta(
role=getattr(choice.delta, "role", None),
content=getattr(choice.delta, "content", None),
)
choices.append(
StreamChoice(
index=choice.index,
delta=delta,
finish_reason=choice.finish_reason,
)
)
yield CompletionChunk(
id=chunk.id,
created=int(time.time()),
model=chunk.model or model,
choices=choices,
)
except litellm.exceptions.RateLimitError as e:
logger.warning("Rate limited during stream: %s", str(e))
raise LLMRateLimitError(str(e), backend=self.name) from e
except litellm.exceptions.Timeout as e:
logger.warning("Timeout during stream: %s", str(e))
raise LLMTimeoutError(str(e), backend=self.name) from e
except litellm.exceptions.AuthenticationError as e:
logger.error("Authentication failed: %s", str(e))
raise CompletionError(
f"Authentication failed: {e}", backend=self.name, model=model
) from e
except litellm.exceptions.APIError as e:
logger.error("API error during stream: %s", str(e))
raise CompletionError(str(e), backend=self.name, model=model) from e
except Exception as e:
logger.error("Unexpected error during stream: %s", str(e))
raise CompletionError(str(e), backend=self.name, model=model) from e
finally:
# Ensure stream is properly closed on cancellation or error
await safe_close_stream(response, logger)
async def list_models(self) -> list[ModelInfo]:
"""List available models via LiteLLM.
Uses a TTL-based cache to avoid hitting external APIs on every call.
Falls back to a curated list if dynamic fetching fails.
Returns:
list[ModelInfo]: List of available models.
"""
# Check cache first
now = time.time()
if self._models_cache is not None:
cached_time, cached_models = self._models_cache
if now - cached_time < self._MODELS_CACHE_TTL:
logger.debug(
"Returning cached model list (%d models)", len(cached_models)
)
return cached_models
models: list[ModelInfo] = []
async with httpx.AsyncClient(
timeout=httpx.Timeout(self._settings.connect_timeout)
) as client:
# Try to fetch OpenAI models dynamically
if self._settings.openai_api_key:
try:
response = await client.get(
"https://api.openai.com/v1/models",
headers={
"Authorization": f"Bearer {self._settings.openai_api_key}"
},
)
if response.status_code == 200:
data = response.json()
for m in data.get("data", []):
model_id = m.get("id", "")
# Filter to chat models (gpt-*)
if model_id.startswith("gpt-"):
models.append(
ModelInfo(
id=model_id,
backend=self.name,
capabilities=["chat"],
)
)
except Exception as e:
logger.debug("Failed to fetch OpenAI models dynamically: %s", e)
# Fall back to curated list for OpenAI
models.extend(self._get_fallback_openai_models())
# Try to fetch OpenRouter models dynamically
if self._settings.openrouter_api_key:
try:
response = await client.get(
"https://openrouter.ai/api/v1/models",
headers={
"Authorization": f"Bearer {self._settings.openrouter_api_key}"
},
)
if response.status_code == 200:
data = response.json()
for m in data.get("data", []):
model_id = m.get("id", "")
if model_id:
context_length = m.get("context_length")
models.append(
ModelInfo(
id=f"openrouter/{model_id}",
backend=self.name,
context_length=context_length,
capabilities=["chat"],
)
)
except Exception as e:
logger.debug("Failed to fetch OpenRouter models dynamically: %s", e)
# Fall back to curated list for OpenRouter
models.extend(self._get_fallback_openrouter_models())
# Anthropic doesn't have a public models endpoint, use curated list
if self._settings.anthropic_api_key:
models.extend(self._get_fallback_anthropic_models())
# Cache the result
self._models_cache = (time.time(), models)
logger.debug("Cached model list (%d models)", len(models))
return models
def _get_fallback_openai_models(self) -> list[ModelInfo]:
"""Get fallback list of OpenAI models."""
return [
ModelInfo(
id="gpt-4o",
backend=self.name,
context_length=128000,
capabilities=["chat"],
),
ModelInfo(
id="gpt-4o-mini",
backend=self.name,
context_length=128000,
capabilities=["chat"],
),
ModelInfo(
id="gpt-4-turbo",
backend=self.name,
context_length=128000,
capabilities=["chat"],
),
ModelInfo(
id="gpt-3.5-turbo",
backend=self.name,
context_length=16385,
capabilities=["chat"],
),
]
def _get_fallback_anthropic_models(self) -> list[ModelInfo]:
"""Get fallback list of Anthropic models."""
return [
ModelInfo(
id="claude-3-5-sonnet-20241022",
backend=self.name,
context_length=200000,
capabilities=["chat"],
),
ModelInfo(
id="claude-3-opus-20240229",
backend=self.name,
context_length=200000,
capabilities=["chat"],
),
ModelInfo(
id="claude-3-haiku-20240307",
backend=self.name,
context_length=200000,
capabilities=["chat"],
),
]
def _get_fallback_openrouter_models(self) -> list[ModelInfo]:
"""Get fallback list of OpenRouter models."""
return [
ModelInfo(
id="openrouter/anthropic/claude-3.5-sonnet",
backend=self.name,
capabilities=["chat"],
),
ModelInfo(
id="openrouter/openai/gpt-4o",
backend=self.name,
capabilities=["chat"],
),
ModelInfo(
id="openrouter/meta-llama/llama-3.1-405b-instruct",
backend=self.name,
capabilities=["chat"],
),
ModelInfo(
id="openrouter/google/gemini-pro-1.5",
backend=self.name,
capabilities=["chat"],
),
]
async def health_check(self) -> bool:
"""Check if LiteLLM backend can reach at least one provider.
Actually tests connectivity to configured providers instead of
always returning True.
Returns:
bool: True if at least one provider is reachable.
"""
async with httpx.AsyncClient(
timeout=httpx.Timeout(self._settings.connect_timeout)
) as client:
# Check OpenAI
if self._settings.openai_api_key:
try:
response = await client.get(
"https://api.openai.com/v1/models",
headers={
"Authorization": f"Bearer {self._settings.openai_api_key}"
},
)
# 200 = success, 401 = reachable but bad key (still healthy)
# Treating 401 as healthy supports key rotation scenarios
if response.status_code in (200, 401):
if response.status_code == 401:
logger.warning(
"OpenAI: API key invalid but endpoint reachable"
)
return True
except httpx.RequestError as e:
logger.debug("OpenAI health check failed: %s", e)
# Check Anthropic
if self._settings.anthropic_api_key:
try:
response = await client.get(
"https://api.anthropic.com/v1/models",
headers={
"x-api-key": self._settings.anthropic_api_key,
"anthropic-version": "2023-06-01",
},
)
if response.status_code in (200, 401):
if response.status_code == 401:
logger.warning(
"Anthropic: API key invalid but endpoint reachable"
)
return True
except httpx.RequestError as e:
logger.debug("Anthropic health check failed: %s", e)
# Check OpenRouter
if self._settings.openrouter_api_key:
try:
response = await client.get(
"https://openrouter.ai/api/v1/models",
headers={
"Authorization": f"Bearer {self._settings.openrouter_api_key}"
},
)
if response.status_code in (200, 401):
if response.status_code == 401:
logger.warning(
"OpenRouter: API key invalid but endpoint reachable"
)
return True
except httpx.RequestError as e:
logger.debug("OpenRouter health check failed: %s", e)
# No providers reachable or configured
logger.warning("No LLM providers reachable")
return False

477
ai_platform/modules/llm-inference/src/llm_inference/backends/llamacpp_backend.py Normal file
View file

@ -0,0 +1,477 @@
"""llama.cpp backend with load balancing across multiple servers.
Supports round-robin distribution with automatic failover and periodic
health checks. When a server goes down, traffic is routed to healthy
servers. When it recovers, it re-enters the rotation.
Configuration:
Single server (legacy):
LLM_LLAMACPP_BASE_URL=http://localhost:8102
Multiple servers (load balanced):
LLM_LLAMACPP_BASE_URLS=http://10.11.10.18:14001,http://10.11.10.19:14001
"""
import asyncio
import logging
import time
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
import httpx
from llm_inference.backends.base import LLMBackend
from llm_inference.config import LLMSettings
from llm_inference.exceptions import (
CompletionError,
LLMConnectionError,
LLMRateLimitError,
LLMTimeoutError,
ModelListError,
)
from llm_inference.retry import retry_with_backoff
from llm_inference.schemas import CompletionChunk, CompletionResponse
from llm_inference.types import (
ChatMessage,
Choice,
Delta,
ModelInfo,
StreamChoice,
Usage,
)
from llm_inference.utils import safe_close_stream
try:
import openai
from openai import AsyncOpenAI
OPENAI_AVAILABLE = True
except ImportError:
OPENAI_AVAILABLE = False
logger = logging.getLogger("llm_inference.backends.llamacpp")
@dataclass
class LlamaCppServer:
"""State for a single llama.cpp server."""
url: str
healthy: bool = True
request_count: int = 0
error_count: int = 0
last_check: float = 0
response_times: list[float] = field(default_factory=list)
@property
def avg_response_ms(self) -> float:
if not self.response_times:
return 0
recent = self.response_times[-50:]
return sum(recent) / len(recent)
@property
def short_name(self) -> str:
"""Extract host:port for logging."""
return self.url.replace("http://", "").replace("https://", "")
class LlamaCppBackend(LLMBackend):
"""llama.cpp backend with round-robin load balancing.
When multiple URLs are configured via LLM_LLAMACPP_BASE_URLS,
requests are distributed round-robin across healthy servers.
Unhealthy servers are skipped and periodically re-checked.
"""
def __init__(self, settings: LLMSettings) -> None:
if not OPENAI_AVAILABLE:
raise ImportError(
"openai package is required for llama.cpp backend. "
"Install with: pip install llm-inference[llamacpp]"
)
self._settings = settings
self._timeout = httpx.Timeout(
connect=settings.connect_timeout,
read=settings.request_timeout,
write=settings.request_timeout,
pool=settings.connect_timeout,
)
# Initialize servers and clients
self._servers: list[LlamaCppServer] = []
self._clients: dict[str, AsyncOpenAI] = {}
for url in settings.llamacpp_urls:
server = LlamaCppServer(url=url)
self._servers.append(server)
self._clients[url] = AsyncOpenAI(
base_url=f"{url}/v1",
api_key="not-needed",
timeout=self._timeout,
max_retries=0,
)
self._current_index = 0
self._lock = asyncio.Lock()
self._health_task: asyncio.Task[None] | None = None
self._health_check_interval = settings.llamacpp_health_check_interval
server_names = [s.short_name for s in self._servers]
logger.info(
"llama.cpp backend: %d server(s): %s", len(self._servers), server_names
)
# Start health check loop if multiple servers
if len(self._servers) > 1:
try:
loop = asyncio.get_running_loop()
self._health_task = loop.create_task(
self._periodic_health_check(self._health_check_interval)
)
except RuntimeError:
# No running loop yet (e.g., during import). Will be started later.
logger.debug("No running loop, deferring health check task")
def start_health_checks(self) -> None:
"""Start periodic health checks if not already running.
Call this from the application lifespan when the event loop is available.
"""
if self._health_task is None and len(self._servers) > 1:
loop = asyncio.get_running_loop()
self._health_task = loop.create_task(
self._periodic_health_check(self._health_check_interval)
)
logger.info("Started periodic health checks (interval=%ds)", self._health_check_interval)
@property
def name(self) -> str:
return "llamacpp"
async def _get_server(self) -> tuple[LlamaCppServer, AsyncOpenAI]:
"""Get next healthy server via round-robin."""
async with self._lock:
healthy = [s for s in self._servers if s.healthy]
if not healthy:
# All down — try all servers (failover attempt)
logger.warning("All llama.cpp servers unhealthy, trying all")
healthy = self._servers
server = healthy[self._current_index % len(healthy)]
self._current_index += 1
return server, self._clients[server.url]
def _mark_unhealthy(self, server: LlamaCppServer) -> None:
server.error_count += 1
if server.healthy:
server.healthy = False
logger.warning(
"llama.cpp server %s marked unhealthy (errors: %d)",
server.short_name,
server.error_count,
)
def _mark_healthy(self, server: LlamaCppServer) -> None:
if not server.healthy:
logger.info(
"llama.cpp server %s recovered (had %d errors)",
server.short_name,
server.error_count,
)
server.error_count = 0
server.healthy = True
async def _periodic_health_check(self, interval: int) -> None:
"""Periodically check health of all servers."""
while True:
await asyncio.sleep(interval)
for server in self._servers:
client = self._clients[server.url]
try:
await client.models.list()
self._mark_healthy(server)
except Exception:
self._mark_unhealthy(server)
server.last_check = time.time()
async def _complete_on_server(
self,
server: LlamaCppServer,
client: AsyncOpenAI,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> CompletionResponse:
"""Execute completion on a specific server."""
start = time.time()
response = await client.chat.completions.create(
model=model,
messages=[m.model_dump(exclude_none=True) for m in messages], # type: ignore[arg-type]
**kwargs,
)
elapsed_ms = (time.time() - start) * 1000
server.response_times.append(elapsed_ms)
if len(server.response_times) > 50:
server.response_times = server.response_times[-50:]
server.request_count += 1
self._mark_healthy(server)
choices = []
for choice in response.choices:
choices.append(
Choice(
index=choice.index,
message=ChatMessage(
role=choice.message.role,
content=choice.message.content,
),
finish_reason=choice.finish_reason,
)
)
usage = None
if response.usage:
usage = Usage(
prompt_tokens=response.usage.prompt_tokens,
completion_tokens=response.usage.completion_tokens,
total_tokens=response.usage.total_tokens,
)
return CompletionResponse(
id=response.id,
created=response.created,
model=response.model or model,
choices=choices,
usage=usage,
backend=f"{self.name}@{server.short_name}",
)
async def complete(
self,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> CompletionResponse:
"""Generate a chat completion with load balancing and failover.
Tries healthy servers round-robin. On failure, marks server unhealthy
and retries on the next available server.
Args:
messages: List of chat messages.
model: Model identifier (may be ignored by server).
**kwargs: Additional parameters.
Returns:
CompletionResponse: The completion response.
Raises:
CompletionError: If all servers fail.
"""
tried: set[str] = set()
last_error: Exception | None = None
for _attempt in range(len(self._servers)):
server, client = await self._get_server()
if server.url in tried:
continue
tried.add(server.url)
async def _do_complete(
_server: LlamaCppServer = server,
_client: AsyncOpenAI = client,
) -> CompletionResponse:
return await self._complete_on_server(
_server, _client, messages, model, **kwargs
)
try:
return await retry_with_backoff(
_do_complete,
max_retries=self._settings.max_retries,
min_wait=self._settings.retry_min_wait,
max_wait=self._settings.retry_max_wait,
backend=self.name,
model=model,
)
except (LLMConnectionError, LLMTimeoutError) as e:
self._mark_unhealthy(server)
last_error = e
logger.warning(
"llama.cpp server %s failed, trying next: %s",
server.short_name,
str(e),
)
continue
except Exception:
raise
raise CompletionError(
f"All llama.cpp servers failed. Last error: {last_error}",
backend=self.name,
model=model,
)
async def stream(
self,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> AsyncIterator[CompletionChunk]:
"""Generate a streaming chat completion with failover.
Args:
messages: List of chat messages.
model: Model identifier.
**kwargs: Additional parameters.
Yields:
CompletionChunk: Streaming chunks of the completion.
Raises:
CompletionError: If all servers fail.
"""
tried: set[str] = set()
last_error: Exception | None = None
for _attempt in range(len(self._servers)):
server, client = await self._get_server()
if server.url in tried:
continue
tried.add(server.url)
stream = None
try:
stream = await client.chat.completions.create(
model=model,
messages=[m.model_dump(exclude_none=True) for m in messages], # type: ignore[arg-type]
stream=True,
**kwargs,
)
server.request_count += 1
self._mark_healthy(server)
async for chunk in stream:
choices = []
if chunk.choices:
for choice in chunk.choices:
delta = Delta(
role=getattr(choice.delta, "role", None),
content=getattr(choice.delta, "content", None),
)
choices.append(
StreamChoice(
index=choice.index,
delta=delta,
finish_reason=choice.finish_reason,
)
)
yield CompletionChunk(
id=chunk.id,
created=int(time.time()),
model=chunk.model or model,
choices=choices,
)
# Stream completed successfully
return
except openai.APIConnectionError as e:
self._mark_unhealthy(server)
last_error = e
logger.warning(
"llama.cpp server %s stream failed: %s", server.short_name, e
)
continue
except openai.APITimeoutError as e:
self._mark_unhealthy(server)
last_error = e
logger.warning(
"llama.cpp server %s stream timeout: %s", server.short_name, e
)
continue
except openai.RateLimitError as e:
raise LLMRateLimitError(str(e), backend=self.name) from e
except openai.AuthenticationError as e:
raise CompletionError(
f"Authentication failed: {e}", backend=self.name, model=model
) from e
except openai.APIStatusError as e:
raise CompletionError(str(e), backend=self.name, model=model) from e
except Exception as e:
raise CompletionError(str(e), backend=self.name, model=model) from e
finally:
await safe_close_stream(stream, logger)
raise CompletionError(
f"All llama.cpp servers failed for stream. Last error: {last_error}",
backend=self.name,
model=model,
)
async def list_models(self) -> list[ModelInfo]:
"""List models from the first healthy server.
Returns:
list[ModelInfo]: List of available models.
Raises:
ModelListError: If models cannot be retrieved.
"""
for server in self._servers:
if not server.healthy and len(self._servers) > 1:
continue
client = self._clients[server.url]
try:
models = await client.models.list()
return [
ModelInfo(
id=m.id,
backend=self.name,
capabilities=["chat", "vision"],
)
for m in models.data
]
except Exception as e:
self._mark_unhealthy(server)
logger.debug("Failed to list models from %s: %s", server.short_name, e)
continue
raise ModelListError(self.name, "All llama.cpp servers unreachable")
async def health_check(self) -> bool:
"""Check if at least one llama.cpp server is healthy.
Returns:
bool: True if any server is reachable.
"""
for server in self._servers:
client = self._clients[server.url]
try:
await client.models.list()
self._mark_healthy(server)
return True
except Exception:
self._mark_unhealthy(server)
continue
return False
def get_servers_status(self) -> list[dict[str, object]]:
"""Get status of all servers (for health endpoint)."""
return [
{
"url": s.url,
"healthy": s.healthy,
"request_count": s.request_count,
"error_count": s.error_count,
"avg_response_ms": round(s.avg_response_ms, 1),
}
for s in self._servers
]

113
ai_platform/modules/llm-inference/src/llm_inference/backends/registry.py Normal file
View file

@ -0,0 +1,113 @@
"""Backend registry for managing LLM backend instances."""
from llm_inference.backends.base import LLMBackend
from llm_inference.config import LLMSettings
from llm_inference.exceptions import BackendNotAvailableError, BackendNotEnabledError
from llm_inference.types import BackendType
class BackendRegistry:
"""Registry for managing LLM backend instances.
The registry initializes and provides access to configured backends.
LiteLLM is always available; vLLM and llama.cpp are optional.
"""
def __init__(self, settings: LLMSettings) -> None:
"""Initialize the backend registry.
Args:
settings: Application settings.
"""
self._settings = settings
self._backends: dict[BackendType, LLMBackend] = {}
self._initialize_backends()
def _initialize_backends(self) -> None:
"""Initialize enabled backends."""
# LiteLLM is always available (core dependency)
from llm_inference.backends.litellm_backend import LiteLLMBackend
self._backends[BackendType.LITELLM] = LiteLLMBackend(self._settings)
# vLLM is optional but MUST work if enabled (fail-fast)
if self._settings.enable_vllm:
try:
from llm_inference.backends.vllm_backend import VLLMBackend
self._backends[BackendType.VLLM] = VLLMBackend(self._settings)
except ImportError as e:
raise BackendNotAvailableError(
"vllm",
f"vLLM backend is enabled but dependencies are not installed. "
f"Install with: pip install llm-inference[vllm]. "
f"Original error: {e}",
) from e
# llama.cpp is optional but MUST work if enabled (fail-fast)
if self._settings.enable_llamacpp:
try:
from llm_inference.backends.llamacpp_backend import LlamaCppBackend
self._backends[BackendType.LLAMACPP] = LlamaCppBackend(self._settings)
except ImportError as e:
raise BackendNotAvailableError(
"llamacpp",
f"llama.cpp backend is enabled but dependencies are not installed. "
f"Install with: pip install llm-inference[llamacpp]. "
f"Original error: {e}",
) from e
def get(self, backend_type: BackendType | str | None = None) -> LLMBackend:
"""Get a backend instance.
Args:
backend_type: Backend to retrieve. If None, uses default from settings.
Returns:
LLMBackend: The requested backend instance.
Raises:
BackendNotAvailableError: If the backend is not available.
BackendNotEnabledError: If the backend is not enabled.
"""
if backend_type is None:
backend_type = BackendType(self._settings.default_backend)
elif isinstance(backend_type, str):
backend_type = BackendType(backend_type)
if backend_type not in self._backends:
# Check if it's a valid backend that's just not enabled
if backend_type == BackendType.VLLM and not self._settings.enable_vllm:
raise BackendNotEnabledError("vllm")
if (
backend_type == BackendType.LLAMACPP
and not self._settings.enable_llamacpp
):
raise BackendNotEnabledError("llamacpp")
raise BackendNotAvailableError(
backend_type.value, "Backend not initialized"
)
return self._backends[backend_type]
def list_backends(self) -> list[BackendType]:
"""List all available backends.
Returns:
list[BackendType]: List of available backend types.
"""
return list(self._backends.keys())
def is_available(self, backend_type: BackendType | str) -> bool:
"""Check if a backend is available.
Args:
backend_type: Backend to check.
Returns:
bool: True if the backend is available.
"""
if isinstance(backend_type, str):
backend_type = BackendType(backend_type)
return backend_type in self._backends

348
ai_platform/modules/llm-inference/src/llm_inference/backends/vllm_backend.py Normal file
View file

@ -0,0 +1,348 @@
"""vLLM backend implementation.
vLLM is a high-throughput and memory-efficient inference engine for LLMs.
This backend connects to a vLLM server via its OpenAI-compatible API.
"""
import logging
import time
from collections.abc import AsyncIterator
import httpx
from llm_inference.backends.base import LLMBackend
from llm_inference.config import LLMSettings
from llm_inference.exceptions import (
CompletionError,
LLMConnectionError,
LLMRateLimitError,
LLMTimeoutError,
ModelListError,
)
from llm_inference.retry import retry_with_backoff
from llm_inference.schemas import (
CompletionChunk,
CompletionResponse,
TextChoice,
TextCompletionResponse,
)
from llm_inference.types import (
ChatMessage,
Choice,
Delta,
ModelInfo,
StreamChoice,
Usage,
)
from llm_inference.utils import safe_close_stream
try:
import openai
from openai import AsyncOpenAI
OPENAI_AVAILABLE = True
except ImportError:
OPENAI_AVAILABLE = False
logger = logging.getLogger("llm_inference.backends.vllm")
class VLLMBackend(LLMBackend):
"""vLLM backend using OpenAI-compatible API.
Requires the vLLM server to be running with the OpenAI-compatible
API enabled.
Example:
Start vLLM server:
```bash
python -m vllm.entrypoints.openai.api_server \
--model meta-llama/Llama-2-7b-chat-hf \
--host 0.0.0.0 --port 8101
```
Configure:
```bash
export LLM_ENABLE_VLLM=true
export LLM_VLLM_BASE_URL=http://localhost:8101
```
"""
def __init__(self, settings: LLMSettings) -> None:
"""Initialize vLLM backend.
Args:
settings: Application settings.
Raises:
ImportError: If openai package is not installed.
"""
if not OPENAI_AVAILABLE:
raise ImportError(
"openai package is required for vLLM backend. "
"Install with: pip install llm-inference[vllm]"
)
self._settings = settings
# Configure httpx timeout
timeout = httpx.Timeout(
connect=settings.connect_timeout,
read=settings.request_timeout,
write=settings.request_timeout,
pool=settings.connect_timeout,
)
self._client = AsyncOpenAI(
base_url=f"{settings.vllm_base_url}/v1",
api_key=settings.vllm_api_key or "not-needed",
timeout=timeout,
max_retries=0, # We handle retries ourselves
)
@property
def name(self) -> str:
"""Backend identifier name."""
return "vllm"
async def complete(
self,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> CompletionResponse:
"""Generate a chat completion via vLLM.
Args:
messages: List of chat messages.
model: Model identifier.
**kwargs: Additional parameters.
Returns:
CompletionResponse: The completion response.
Raises:
CompletionError: If the request fails after retries.
LLMRateLimitError: If rate limited.
LLMTimeoutError: If the request times out.
"""
async def _do_complete() -> CompletionResponse:
response = await self._client.chat.completions.create(
model=model,
messages=[m.model_dump(exclude_none=True) for m in messages], # type: ignore[arg-type]
**kwargs,
)
choices = []
for choice in response.choices:
choices.append(
Choice(
index=choice.index,
message=ChatMessage(
role=choice.message.role,
content=choice.message.content,
),
finish_reason=choice.finish_reason,
)
)
usage = None
if response.usage:
usage = Usage(
prompt_tokens=response.usage.prompt_tokens,
completion_tokens=response.usage.completion_tokens,
total_tokens=response.usage.total_tokens,
)
return CompletionResponse(
id=response.id,
created=response.created,
model=response.model,
choices=choices,
usage=usage,
backend=self.name,
)
return await retry_with_backoff(
_do_complete,
max_retries=self._settings.max_retries,
min_wait=self._settings.retry_min_wait,
max_wait=self._settings.retry_max_wait,
backend=self.name,
model=model,
)
async def text_complete(
self,
prompt: str | list[str],
model: str,
**kwargs: object,
) -> TextCompletionResponse:
"""Generate a legacy text completion via vLLM (/v1/completions)."""
async def _do_text() -> TextCompletionResponse:
response = await self._client.completions.create(
model=model,
prompt=prompt, # type: ignore[arg-type]
**kwargs,
)
choices = [
TextChoice(
index=c.index, text=c.text, finish_reason=c.finish_reason
)
for c in response.choices
]
usage = None
if response.usage:
usage = Usage(
prompt_tokens=response.usage.prompt_tokens,
completion_tokens=response.usage.completion_tokens,
total_tokens=response.usage.total_tokens,
)
return TextCompletionResponse(
id=response.id,
created=response.created,
model=response.model,
choices=choices,
usage=usage,
backend=self.name,
)
return await retry_with_backoff(
_do_text,
max_retries=self._settings.max_retries,
min_wait=self._settings.retry_min_wait,
max_wait=self._settings.retry_max_wait,
backend=self.name,
model=model,
)
async def stream(
self,
messages: list[ChatMessage],
model: str,
**kwargs: object,
) -> AsyncIterator[CompletionChunk]:
"""Generate a streaming chat completion via vLLM.
Args:
messages: List of chat messages.
model: Model identifier.
**kwargs: Additional parameters.
Yields:
CompletionChunk: Streaming chunks of the completion.
Raises:
CompletionError: If the request fails.
LLMRateLimitError: If rate limited.
LLMTimeoutError: If the request times out.
"""
stream = None
try:
stream = await self._client.chat.completions.create(
model=model,
messages=[m.model_dump(exclude_none=True) for m in messages], # type: ignore[arg-type]
stream=True,
**kwargs,
)
async for chunk in stream:
choices = []
if chunk.choices:
for choice in chunk.choices:
delta = Delta(
role=getattr(choice.delta, "role", None),
content=getattr(choice.delta, "content", None),
)
choices.append(
StreamChoice(
index=choice.index,
delta=delta,
finish_reason=choice.finish_reason,
)
)
yield CompletionChunk(
id=chunk.id,
created=int(time.time()),
model=chunk.model or model,
choices=choices,
)
except openai.RateLimitError as e:
logger.warning("Rate limited during stream: %s", str(e))
raise LLMRateLimitError(str(e), backend=self.name) from e
except openai.APITimeoutError as e:
logger.warning("Timeout during stream: %s", str(e))
raise LLMTimeoutError(str(e), backend=self.name) from e
except openai.APIConnectionError as e:
logger.error("Connection error during stream: %s", str(e))
raise LLMConnectionError(self.name, reason=str(e)) from e
except openai.AuthenticationError as e:
logger.error("Authentication failed: %s", str(e))
raise CompletionError(
f"Authentication failed: {e}", backend=self.name, model=model
) from e
except openai.APIStatusError as e:
logger.error("API error during stream: %s", str(e))
raise CompletionError(str(e), backend=self.name, model=model) from e
except Exception as e:
logger.error("Unexpected error during stream: %s", str(e))
raise CompletionError(str(e), backend=self.name, model=model) from e
finally:
# Ensure stream is properly closed on cancellation or error
await safe_close_stream(stream, logger)
async def list_models(self) -> list[ModelInfo]:
"""List models available on the vLLM server.
Returns:
list[ModelInfo]: List of available models.
Raises:
ModelListError: If models cannot be retrieved.
"""
try:
models = await self._client.models.list()
return [
ModelInfo(
id=m.id,
backend=self.name,
capabilities=["chat"],
)
for m in models.data
]
except openai.APIConnectionError as e:
logger.error("Failed to list models - connection error: %s", str(e))
raise ModelListError(
self.name, f"Cannot connect to vLLM server: {e}"
) from e
except openai.APITimeoutError as e:
logger.error("Failed to list models - timeout: %s", str(e))
raise ModelListError(
self.name, f"Timeout connecting to vLLM server: {e}"
) from e
except Exception as e:
logger.error("Failed to list models: %s", str(e))
raise ModelListError(self.name, str(e)) from e
async def health_check(self) -> bool:
"""Check if vLLM server is healthy.
Returns:
bool: True if the server is reachable and responding.
"""
try:
await self._client.models.list()
return True
except openai.APIConnectionError:
logger.debug("vLLM health check failed: connection error")
return False
except openai.APITimeoutError:
logger.debug("vLLM health check failed: timeout")
return False
except Exception as e:
logger.debug("vLLM health check failed: %s", str(e))
return False

140
ai_platform/modules/llm-inference/src/llm_inference/cli.py Normal file
View file

@ -0,0 +1,140 @@
"""CLI entry point for the LLM inference server."""
import argparse
import signal
import sys
from typing import Any
import uvicorn
from llm_inference.config import SettingsCache
from llm_inference.logging import configure_logging, get_logger
class GracefulShutdown:
"""Handles graceful shutdown on signals.
Allows in-flight requests to complete before exiting.
"""
def __init__(self) -> None:
"""Initialize shutdown handler."""
self._server: uvicorn.Server | None = None
self._shutdown_requested = False
self._logger = get_logger("cli.shutdown")
def register_signals(self) -> None:
"""Register signal handlers for graceful shutdown."""
signal.signal(signal.SIGTERM, self._signal_handler)
signal.signal(signal.SIGINT, self._signal_handler)
def _signal_handler(self, signum: int, frame: Any) -> None:
"""Handle shutdown signals.
Args:
signum: Signal number.
frame: Current stack frame (unused).
"""
if self._shutdown_requested:
# Second signal - force exit
self._logger.warning("Forced shutdown requested")
sys.exit(1)
self._shutdown_requested = True
signal_name = signal.Signals(signum).name
self._logger.info("Received %s, initiating graceful shutdown...", signal_name)
if self._server:
self._server.should_exit = True
def set_server(self, server: uvicorn.Server) -> None:
"""Set the uvicorn server for shutdown control.
Args:
server: Uvicorn server instance.
"""
self._server = server
def main() -> None:
"""Run the LLM inference server with graceful shutdown support."""
parser = argparse.ArgumentParser(
description="LLM Inference API Server",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
"--host",
type=str,
default=None,
help="Host to bind to (default: from LLM_HOST env or 0.0.0.0)",
)
parser.add_argument(
"--port",
type=int,
default=None,
help="Port to bind to (default: from LLM_PORT env or 14011)",
)
parser.add_argument(
"--workers",
type=int,
default=1,
help="Number of worker processes",
)
parser.add_argument(
"--reload",
action="store_true",
help="Enable auto-reload for development",
)
parser.add_argument(
"--graceful-timeout",
type=int,
default=30,
help="Timeout in seconds for graceful shutdown",
)
args = parser.parse_args()
# Get settings for defaults
settings = SettingsCache.get()
# Configure logging early
configure_logging(settings.log_level, settings.log_json)
logger = get_logger("cli")
host = args.host or settings.host
port = args.port or settings.port
logger.info("Starting LLM Inference API Server on %s:%d", host, port)
if args.reload:
# Development mode - use simple uvicorn.run
logger.info("Development mode - auto-reload enabled")
uvicorn.run(
"llm_inference.api.app:create_app",
factory=True,
host=host,
port=port,
workers=1,
reload=True,
)
else:
# Production mode - use server with graceful shutdown
shutdown_handler = GracefulShutdown()
shutdown_handler.register_signals()
config = uvicorn.Config(
"llm_inference.api.app:create_app",
factory=True,
host=host,
port=port,
workers=args.workers,
timeout_graceful_shutdown=args.graceful_timeout,
)
server = uvicorn.Server(config)
shutdown_handler.set_server(server)
server.run()
if __name__ == "__main__":
main()

357
ai_platform/modules/llm-inference/src/llm_inference/client.py Normal file
View file

@ -0,0 +1,357 @@
"""High-level LLM client for unified inference."""
from collections.abc import AsyncIterator
from time import perf_counter
from llm_inference import metrics
from llm_inference.backends import BackendRegistry
from llm_inference.config import LLMSettings, SettingsCache
from llm_inference.exceptions import (
BackendNotAvailableError,
BackendNotEnabledError,
CompletionError,
LLMRateLimitError,
LLMTimeoutError,
)
from llm_inference.image_processing import process_messages
from llm_inference.schemas import (
CompletionChunk,
CompletionResponse,
TextCompletionResponse,
)
from llm_inference.types import BackendType, ChatMessage, ModelInfo
class LLMClient:
"""High-level client for LLM inference.
This is the primary interface for using the module as a Python library.
It provides a unified interface across all configured backends.
Example:
```python
from llm_inference import LLMClient
client = LLMClient()
# Simple completion
response = await client.complete(
messages=[{"role": "user", "content": "Hello!"}],
model="gpt-3.5-turbo",
)
print(response.choices[0].message.content)
# Streaming
async for chunk in client.stream(
messages=[{"role": "user", "content": "Tell me a story"}],
model="gpt-4",
):
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="")
# Override backend per-request
response = await client.complete(
messages=[{"role": "user", "content": "Hello!"}],
model="meta-llama/Llama-2-7b-chat-hf",
backend=BackendType.VLLM,
)
```
"""
def __init__(self, settings: LLMSettings | None = None) -> None:
"""Initialize the LLM client.
Args:
settings: Optional settings instance. If not provided, loads from environment.
"""
self._settings = settings or SettingsCache.get()
self._registry = BackendRegistry(self._settings)
@property
def registry(self) -> BackendRegistry:
"""Get the backend registry.
Returns:
BackendRegistry: The backend registry instance.
"""
return self._registry
async def complete(
self,
messages: list[dict[str, str] | ChatMessage],
model: str | None = None,
backend: BackendType | str | None = None,
**kwargs: object,
) -> CompletionResponse:
"""Generate a chat completion.
Args:
messages: List of chat messages. Can be dicts or ChatMessage objects.
model: Model identifier. Uses default if not specified.
backend: Backend to use. If not specified, auto-detects from model
or uses default.
**kwargs: Additional parameters (temperature, max_tokens, etc.).
Returns:
CompletionResponse: The completion response.
Raises:
BackendNotAvailableError: If the requested backend is not available.
CompletionError: If the completion request fails.
"""
parsed_messages = self._parse_messages(messages)
parsed_messages = await process_messages(parsed_messages)
model = self._resolve_model(model or self._settings.default_model)
explicit = backend is not None
if backend is None:
backend = await self._resolve_backend_for_model(model)
cascade = self._build_cascade(backend, explicit)
last_error: Exception | None = None
for backend_type in cascade:
name = backend_type.value
start = perf_counter()
metrics.LLM_INFLIGHT.labels(backend=name).inc()
try:
response = await self._registry.get(backend_type).complete(
parsed_messages, model, **kwargs
)
metrics.observe_success(
model, response.backend or name, perf_counter() - start,
response.usage,
)
return response
except (
BackendNotAvailableError,
BackendNotEnabledError,
CompletionError,
LLMTimeoutError,
LLMRateLimitError,
) as exc:
metrics.observe_error(model, name)
last_error = exc
continue
finally:
metrics.LLM_INFLIGHT.labels(backend=name).dec()
# All backends in the cascade failed — surface the last error.
raise last_error if last_error else CompletionError("no backend available")
async def stream(
self,
messages: list[dict[str, str] | ChatMessage],
model: str | None = None,
backend: BackendType | str | None = None,
**kwargs: object,
) -> AsyncIterator[CompletionChunk]:
"""Generate a streaming chat completion.
Args:
messages: List of chat messages.
model: Model identifier. Uses default if not specified.
backend: Backend to use. Uses default if not specified.
**kwargs: Additional parameters.
Yields:
CompletionChunk: Streaming chunks of the completion.
Raises:
BackendNotAvailableError: If the requested backend is not available.
CompletionError: If the completion request fails.
"""
parsed_messages = self._parse_messages(messages)
parsed_messages = await process_messages(parsed_messages)
model = self._resolve_model(model or self._settings.default_model)
if backend is None:
backend = await self._resolve_backend_for_model(model)
backend_instance = self._registry.get(backend)
async for chunk in backend_instance.stream(parsed_messages, model, **kwargs):
yield chunk
async def text_complete(
self,
prompt: str | list[str],
model: str | None = None,
backend: BackendType | str | None = None,
**kwargs: object,
) -> TextCompletionResponse:
"""Generate a legacy text completion (/v1/completions).
Resolves the model alias and target backend, then delegates to the
backend's ``text_complete``. Backends that do not support text
completion raise NotImplementedError.
"""
model = self._resolve_model(model or self._settings.default_model)
if backend is None:
backend = await self._resolve_backend_for_model(model)
backend_instance = self._registry.get(backend)
return await backend_instance.text_complete(prompt, model, **kwargs)
async def list_models(
self,
backend: BackendType | str | None = None,
) -> list[ModelInfo]:
"""List available models.
Args:
backend: If specified, only list models for this backend.
Otherwise, lists models from all available backends.
Returns:
list[ModelInfo]: List of available models.
"""
if backend is not None:
return await self._registry.get(backend).list_models()
# Aggregate from all backends
all_models: list[ModelInfo] = []
for backend_type in self._registry.list_backends():
backend_instance = self._registry.get(backend_type)
models = await backend_instance.list_models()
all_models.extend(models)
return all_models
async def load_model(
self,
model: str,
backend: BackendType | str,
) -> bool:
"""Load a model on the specified backend.
This is only supported by local backends (vLLM, llama.cpp).
Args:
model: Model identifier to load.
backend: Backend to load the model on.
Returns:
bool: True if the model was loaded successfully.
Raises:
NotImplementedError: If the backend doesn't support model loading.
ModelLoadError: If loading fails.
"""
backend_instance = self._registry.get(backend)
return await backend_instance.load_model(model)
async def unload_model(
self,
model: str,
backend: BackendType | str,
) -> bool:
"""Unload a model from the specified backend.
This is only supported by local backends (vLLM, llama.cpp).
Args:
model: Model identifier to unload.
backend: Backend to unload the model from.
Returns:
bool: True if the model was unloaded successfully.
Raises:
NotImplementedError: If the backend doesn't support model unloading.
"""
backend_instance = self._registry.get(backend)
return await backend_instance.unload_model(model)
def list_backends(self) -> list[BackendType]:
"""List available backends.
Returns:
list[BackendType]: List of available backend types.
"""
return self._registry.list_backends()
async def health_check(self) -> dict[str, bool]:
"""Check health of all backends.
Returns:
dict[str, bool]: Mapping of backend name to health status.
"""
health: dict[str, bool] = {}
for backend_type in self._registry.list_backends():
backend_instance = self._registry.get(backend_type)
health[backend_type.value] = await backend_instance.health_check()
return health
def _resolve_model(self, model: str) -> str:
"""Resolve a friendly model alias to its real served id.
Returns the configured alias target, or the input unchanged when no
alias is defined (passthrough).
"""
return self._settings.model_aliases.get(model, model)
def _build_cascade(
self, primary: BackendType | str | None, explicit: bool
) -> list[BackendType]:
"""Build the ordered backend cascade for a request.
Returns ``[primary]`` when the caller pinned a backend or fallback is
disabled; otherwise ``[primary, *other enabled backends]`` in
``fallback_order`` (offer §1.6).
"""
if primary is None:
primary = BackendType(self._settings.default_backend)
elif isinstance(primary, str):
primary = BackendType(primary)
if explicit or not self._settings.enable_fallback:
return [primary]
enabled = set(self._registry.list_backends())
chain: list[BackendType] = [primary] if primary in enabled else []
for name in self._settings.fallback_order:
try:
bt = BackendType(name)
except ValueError:
continue
if bt in enabled and bt != primary:
chain.append(bt)
return chain or [primary]
async def _resolve_backend_for_model(self, model: str) -> BackendType | None:
"""Auto-detect the correct backend for a model.
Queries each enabled local backend's model list to find which one
serves the requested model. Returns None to use default.
Args:
model: Model identifier to look up.
Returns:
BackendType if found, None for default.
"""
for backend_type in self._registry.list_backends():
if backend_type == BackendType.LITELLM:
continue
try:
models = await self._registry.get(backend_type).list_models()
if any(m.id == model for m in models):
return backend_type
except Exception:
continue
return None
def _parse_messages(
self,
messages: list[dict[str, str] | ChatMessage],
) -> list[ChatMessage]:
"""Parse messages into ChatMessage objects.
Args:
messages: List of messages as dicts or ChatMessage objects.
Returns:
list[ChatMessage]: List of ChatMessage objects.
"""
parsed: list[ChatMessage] = []
for msg in messages:
if isinstance(msg, ChatMessage):
parsed.append(msg)
else:
parsed.append(ChatMessage.model_validate(msg))
return parsed

337
ai_platform/modules/llm-inference/src/llm_inference/config.py Normal file
View file

@ -0,0 +1,337 @@
"""Configuration management using Pydantic V2 Settings."""
import threading
import warnings
from typing import Literal
from pydantic import Field, field_validator, model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
class LLMSettings(BaseSettings):
"""LLM Inference module configuration.
All settings can be configured via environment variables with the LLM_ prefix.
Required environment variables (no defaults - fail-fast):
LLM_DEFAULT_BACKEND: Backend to use (litellm, vllm, llamacpp)
LLM_ENABLE_VLLM: Enable vLLM backend (true/false)
LLM_ENABLE_LLAMACPP: Enable llama.cpp backend (true/false)
Example:
LLM_DEFAULT_BACKEND=litellm
LLM_ENABLE_VLLM=false
LLM_ENABLE_LLAMACPP=false
LLM_PORT=14011
"""
model_config = SettingsConfigDict(
env_prefix="LLM_",
env_file=".env",
env_file_encoding="utf-8",
extra="forbid", # Reject unknown fields to catch typos
)
# ==========================================================================
# REQUIRED fields (no defaults - fail-fast)
# ==========================================================================
default_backend: Literal["litellm", "vllm", "llamacpp"] = Field(
description="Default backend to use for inference (REQUIRED)",
)
enable_vllm: bool = Field(
description="Enable vLLM backend (REQUIRED, set to true or false)",
)
enable_llamacpp: bool = Field(
description="Enable llama.cpp backend (REQUIRED, set to true or false)",
)
# ==========================================================================
# Optional fields with sensible defaults
# ==========================================================================
# Default model selection — the platform's primary model (offer: Qwen3.5).
default_model: str = Field(
default="qwen3.5",
description="Default model alias used when a request omits `model`",
)
# Friendly alias → real served model id. Lets callers use a stable name
# (e.g. "qwen3.5") regardless of the backend's served model id. Set via
# LLM_MODEL_ALIASES='{"qwen3.5":"Qwen/Qwen3.5-35B-A3B"}'. Empty = passthrough
# (works when vLLM is started with --served-model-name qwen3.5).
model_aliases: dict[str, str] = Field(
default_factory=dict,
description="Map of friendly model alias to the real served model id",
)
# Cross-backend fallback cascade (offer §1.6): on a backend failure, retry
# the request on the next enabled backend in this order. Disabled when the
# caller pins an explicit backend.
enable_fallback: bool = Field(
default=True,
description="Enable cross-backend fallback cascade on backend failure",
)
fallback_order: list[str] = Field(
default_factory=lambda: ["vllm", "llamacpp", "litellm"],
description="Backend order tried in the fallback cascade",
)
# API server settings
host: str = Field(
default="0.0.0.0",
description="Host to bind the server to",
)
port: int = Field(
default=14011,
description="Port to bind the server to (Prod LLM API: 14011)",
)
external_url: str = Field(
description="External URL for OpenAPI spec (e.g., http://10.11.10.42:14011). REQUIRED.",
)
# LiteLLM / Provider API keys (optional but at least one needed for litellm)
openrouter_api_key: str | None = Field(
default=None,
description="OpenRouter API key",
)
openai_api_key: str | None = Field(
default=None,
description="OpenAI API key",
)
anthropic_api_key: str | None = Field(
default=None,
description="Anthropic API key",
)
# vLLM settings
vllm_base_url: str = Field(
default="http://localhost:14001",
description="Base URL for vLLM server",
)
vllm_api_key: str | None = Field(
default=None,
description="API key for vLLM server (if required)",
)
# llama.cpp settings
llamacpp_base_url: str = Field(
default="http://localhost:8080",
description="Base URL for llama.cpp server (single server, legacy)",
)
llamacpp_base_urls: list[str] | None = Field(
default=None,
description=(
"Comma-separated list of llama.cpp server URLs for load balancing. "
"If set, overrides llamacpp_base_url. "
"Example: http://10.11.10.18:14001"
),
)
llamacpp_health_check_interval: int = Field(
default=30,
ge=5,
description="Health check interval in seconds for llama.cpp backends",
)
@field_validator("llamacpp_base_urls", mode="before")
@classmethod
def parse_llamacpp_base_urls(cls, v: str | list[str] | None) -> list[str] | None:
"""Parse comma-separated URLs into a list."""
if v is None or v == "":
return None
if isinstance(v, str):
urls = [u.strip().rstrip("/") for u in v.split(",") if u.strip()]
return urls if urls else None
return [u.strip().rstrip("/") for u in v if u.strip()]
@property
def llamacpp_urls(self) -> list[str]:
"""Get all llama.cpp server URLs (from list or single)."""
if self.llamacpp_base_urls:
return self.llamacpp_base_urls
return [self.llamacpp_base_url.rstrip("/")]
# Model storage
models_dir: str = Field(
default="/models",
description="Directory for local model storage",
)
# ==========================================================================
# Timeout and retry settings
# ==========================================================================
request_timeout: float = Field(
default=120.0,
ge=1.0,
description="Default timeout in seconds for completion requests",
)
connect_timeout: float = Field(
default=10.0,
ge=1.0,
description="Timeout in seconds for establishing connections",
)
max_retries: int = Field(
default=3,
ge=0,
description="Maximum number of retries for transient failures",
)
retry_min_wait: float = Field(
default=1.0,
ge=0.1,
description="Minimum wait time in seconds between retries",
)
retry_max_wait: float = Field(
default=60.0,
ge=1.0,
description="Maximum wait time in seconds between retries",
)
# ==========================================================================
# Observability settings
# ==========================================================================
log_level: str = Field(
default="INFO",
description="Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)",
)
log_json: bool = Field(
default=False,
description="Output logs in JSON format for production",
)
# ==========================================================================
# Rate limiting and concurrency
# ==========================================================================
rate_limit_rps: float = Field(
default=10.0,
ge=0.1,
description="Rate limit: requests per second",
)
rate_limit_burst: int = Field(
default=20,
ge=1,
description="Rate limit: burst capacity",
)
max_concurrent_completions: int = Field(
default=10,
ge=1,
description="Maximum concurrent completion requests",
)
# ==========================================================================
# Authentication settings
# ==========================================================================
api_tokens: frozenset[str] | None = Field(
default=None,
description=(
"API tokens for Bearer authentication (comma-separated). "
"If not set, authentication is disabled."
),
)
# ==========================================================================
# Runtime config (dashboard polling)
# ==========================================================================
dashboard_url: str | None = Field(
default=None,
description=(
"Optional dashboard base URL (e.g., http://didiAI-dashboard:51300). "
"When set, runtime_config polls /api/config every 30s for live "
"overrides. Unset disables polling — module uses env-only config."
),
)
@field_validator("api_tokens", mode="before")
@classmethod
def parse_api_tokens(cls, v: str | list[str] | None) -> frozenset[str] | None:
"""Parse comma-separated tokens into a frozenset."""
if v is None or v == "":
return None
if isinstance(v, str):
tokens = [t.strip() for t in v.split(",") if t.strip()]
return frozenset(tokens) if tokens else None
return frozenset(v) if v else None
@property
def auth_enabled(self) -> bool:
"""Check if authentication is enabled."""
return bool(self.api_tokens)
# ==========================================================================
# Validators
# ==========================================================================
@model_validator(mode="after")
def validate_api_keys(self) -> "LLMSettings":
"""Warn if no API keys are set when using litellm backend."""
if self.default_backend == "litellm":
has_any_key = any(
[
self.openrouter_api_key,
self.openai_api_key,
self.anthropic_api_key,
]
)
if not has_any_key:
warnings.warn(
"No API keys configured for litellm backend. "
"Set at least one of: OPENROUTER_API_KEY, "
"OPENAI_API_KEY, ANTHROPIC_API_KEY",
UserWarning,
stacklevel=2,
)
return self
class SettingsCache:
"""Thread-safe settings cache that can be cleared for testing."""
_instance: LLMSettings | None = None
_lock: threading.Lock = threading.Lock()
@classmethod
def get(cls) -> LLMSettings:
"""Get or create the settings instance.
Thread-safe singleton pattern. Always acquires lock for correctness.
Performance impact is negligible for singleton access patterns.
Returns:
LLMSettings: Application settings loaded from environment.
"""
with cls._lock:
if cls._instance is None:
cls._instance = LLMSettings()
return cls._instance
@classmethod
def clear(cls) -> None:
"""Clear the cached settings instance.
Use this in tests to reset settings between test cases.
"""
with cls._lock:
cls._instance = None
@classmethod
def set(cls, settings: LLMSettings) -> None:
"""Set a specific settings instance.
Use this in tests to inject mock settings.
Args:
settings: Settings instance to use.
"""
with cls._lock:
cls._instance = settings
def get_settings() -> LLMSettings:
"""Get cached settings instance.
Returns:
LLMSettings: Application settings loaded from environment.
"""
return SettingsCache.get()

137
ai_platform/modules/llm-inference/src/llm_inference/exceptions.py Normal file
View file

@ -0,0 +1,137 @@
"""Custom exceptions for LLM inference."""
class LLMInferenceError(Exception):
"""Base exception for LLM inference errors."""
class AuthenticationError(LLMInferenceError):
"""Raised when authentication fails."""
def __init__(self, message: str = "Authentication required") -> None:
super().__init__(message)
class BackendNotAvailableError(LLMInferenceError):
"""Raised when a requested backend is not available."""
def __init__(self, backend: str, reason: str | None = None) -> None:
self.backend = backend
self.reason = reason
message = f"Backend '{backend}' is not available"
if reason:
message += f": {reason}"
super().__init__(message)
class BackendNotEnabledError(LLMInferenceError):
"""Raised when a backend is not enabled in configuration."""
def __init__(self, backend: str) -> None:
self.backend = backend
super().__init__(
f"Backend '{backend}' is not enabled. "
f"Set LLM_ENABLE_{backend.upper()}=true to enable it."
)
class ModelNotFoundError(LLMInferenceError):
"""Raised when a requested model is not found."""
def __init__(self, model: str, backend: str | None = None) -> None:
self.model = model
self.backend = backend
message = f"Model '{model}' not found"
if backend:
message += f" on backend '{backend}'"
super().__init__(message)
class CompletionError(LLMInferenceError):
"""Raised when a completion request fails."""
def __init__(
self,
message: str,
backend: str | None = None,
model: str | None = None,
) -> None:
self.backend = backend
self.model = model
full_message = message
if backend or model:
details = []
if backend:
details.append(f"backend={backend}")
if model:
details.append(f"model={model}")
full_message = f"{message} ({', '.join(details)})"
super().__init__(full_message)
class ModelLoadError(LLMInferenceError):
"""Raised when loading a model fails."""
def __init__(self, model: str, backend: str, reason: str | None = None) -> None:
self.model = model
self.backend = backend
self.reason = reason
message = f"Failed to load model '{model}' on backend '{backend}'"
if reason:
message += f": {reason}"
super().__init__(message)
class ModelListError(LLMInferenceError):
"""Raised when listing models fails."""
def __init__(self, backend: str, reason: str | None = None) -> None:
self.backend = backend
self.reason = reason
message = f"Failed to list models on backend '{backend}'"
if reason:
message += f": {reason}"
super().__init__(message)
class LLMRateLimitError(LLMInferenceError):
"""Raised when a rate limit is hit."""
def __init__(
self,
message: str,
backend: str | None = None,
retry_after: float | None = None,
) -> None:
self.backend = backend
self.retry_after = retry_after
full_message = message
if retry_after:
full_message += f" (retry after {retry_after}s)"
super().__init__(full_message)
class LLMTimeoutError(LLMInferenceError):
"""Raised when a request times out."""
def __init__(
self,
message: str,
backend: str | None = None,
timeout: float | None = None,
) -> None:
self.backend = backend
self.timeout = timeout
super().__init__(message)
class LLMConnectionError(LLMInferenceError):
"""Raised when connection to backend fails."""
def __init__(self, backend: str, reason: str | None = None) -> None:
self.backend = backend
self.reason = reason
message = f"Connection to backend '{backend}' failed"
if reason:
message += f": {reason}"
super().__init__(message)

185
ai_platform/modules/llm-inference/src/llm_inference/image_processing.py Normal file
View file

@ -0,0 +1,185 @@
"""Image processing for multimodal requests.
Downloads image URLs and converts them to base64 data URIs so that all
backends (vLLM, llama.cpp, LiteLLM) receive a uniform format.
Supports:
1. OpenAI format: {"type": "image_url", "image_url": {"url": "https://..."}}
2. Already base64: {"type": "image_url", "image_url": {"url": "data:image/...;base64,..."}}
→ passed through unchanged.
"""
import base64
import logging
import mimetypes
import re
from urllib.parse import urlparse
import httpx
from llm_inference.types import ChatMessage
logger = logging.getLogger("llm_inference.image_processing")
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
IMAGE_URL_PATTERN = re.compile(
r"https?://[^\s<>\"]+\.(?:jpg|jpeg|png|gif|webp|bmp)", re.IGNORECASE
)
# Shared client for image downloads (reused across requests)
_http_client: httpx.AsyncClient | None = None
def _get_http_client() -> httpx.AsyncClient:
global _http_client
if _http_client is None:
_http_client = httpx.AsyncClient(
timeout=30,
follow_redirects=True,
headers={"User-Agent": "LLMInference/1.0"},
)
return _http_client
async def close_http_client() -> None:
"""Close the shared HTTP client. Call on application shutdown."""
global _http_client
if _http_client is not None:
await _http_client.aclose()
_http_client = None
def _guess_mime_type(url: str, content_type: str = "") -> str:
"""Guess MIME type from URL or Content-Type header."""
if "jpeg" in content_type or "jpg" in content_type:
return "image/jpeg"
if "png" in content_type:
return "image/png"
if "gif" in content_type:
return "image/gif"
if "webp" in content_type:
return "image/webp"
if "bmp" in content_type:
return "image/bmp"
ext = urlparse(url).path.rsplit(".", 1)[-1].lower() if "." in url else ""
mime = mimetypes.guess_type(f"file.{ext}")[0]
return mime or "image/jpeg"
async def _download_and_encode(url: str) -> str:
"""Download image URL and return as base64 data URI."""
client = _get_http_client()
resp = await client.get(url)
resp.raise_for_status()
content_type = resp.headers.get("content-type", "")
mime_type = _guess_mime_type(url, content_type)
b64 = base64.b64encode(resp.content).decode("utf-8")
logger.info(
"Downloaded image: %s (%d bytes, %s)", url, len(resp.content), mime_type
)
return f"data:{mime_type};base64,{b64}"
async def _process_content_item(item: dict[str, object]) -> dict[str, object]:
"""Process a single content item, downloading image URLs if needed."""
if not isinstance(item, dict):
return item
if item.get("type") != "image_url":
return item
image_url = item.get("image_url", {})
url = image_url.get("url", "") if isinstance(image_url, dict) else str(image_url)
# Already base64 — pass through
if url.startswith("data:"):
return item
# Not an HTTP URL — pass through
if not url.startswith(("http://", "https://")):
return item
try:
data_uri = await _download_and_encode(url)
return {
"type": "image_url",
"image_url": {"url": data_uri},
}
except Exception as e:
logger.warning("Failed to download image %s: %s", url, e)
return item # Return original on failure
async def _process_text_with_urls(text: str) -> str | list[dict[str, object]]:
"""Detect image URLs in plain text and convert to multimodal format."""
urls = IMAGE_URL_PATTERN.findall(text)
if not urls:
return text
# Build multimodal content array
content: list[dict[str, object]] = []
# Strip URLs from text
clean_text = text
for url in urls:
clean_text = clean_text.replace(url, "").strip()
if clean_text:
content.append({"type": "text", "text": clean_text})
for url in urls:
try:
data_uri = await _download_and_encode(url)
content.append({"type": "image_url", "image_url": {"url": data_uri}})
except Exception as e:
logger.warning("Failed to download image %s: %s", url, e)
# Put URL back in text
if content and content[0].get("type") == "text":
content[0]["text"] += f" {url}" # type: ignore[operator]
else:
content.insert(0, {"type": "text", "text": url})
return content
async def process_messages(messages: list[ChatMessage]) -> list[ChatMessage]:
"""Process all messages, downloading image URLs to base64.
This ensures uniform image format for all backends. Images arrive
as base64 data URIs regardless of whether the client sent URLs or
base64 originally.
Args:
messages: Original messages (may contain image URLs).
Returns:
Processed messages with images as base64 data URIs.
"""
processed: list[ChatMessage] = []
for msg in messages:
content = msg.content
# Already multimodal array (OpenAI format)
if isinstance(content, list):
new_content = []
for item in content:
new_content.append(await _process_content_item(item))
processed.append(
ChatMessage(role=msg.role, content=new_content, name=msg.name)
)
# Plain text — check for embedded image URLs
elif isinstance(content, str):
new_content = await _process_text_with_urls(content)
processed.append(
ChatMessage(role=msg.role, content=new_content, name=msg.name)
)
else:
processed.append(msg)
return processed

135
ai_platform/modules/llm-inference/src/llm_inference/logging.py Normal file
View file

@ -0,0 +1,135 @@
"""Structured logging with request_id context propagation.
This module provides:
- ContextVar-based request_id tracking
- Structured JSON logging option
- Request ID injection into all log records
"""
import logging
import sys
from contextvars import ContextVar
from typing import Any
# Context variable for request ID propagation
request_id_ctx: ContextVar[str | None] = ContextVar("request_id", default=None)
class RequestIdFilter(logging.Filter):
"""Logging filter that adds request_id to log records."""
def filter(self, record: logging.LogRecord) -> bool:
"""Add request_id to the log record.
Args:
record: The log record to modify.
Returns:
bool: Always True (never filters out records).
"""
record.request_id = request_id_ctx.get() or "-"
return True
class JsonFormatter(logging.Formatter):
"""JSON log formatter for structured logging."""
def format(self, record: logging.LogRecord) -> str:
"""Format log record as JSON.
Args:
record: The log record to format.
Returns:
str: JSON-formatted log string.
"""
import json
log_data: dict[str, Any] = {
"timestamp": self.formatTime(record, self.datefmt),
"level": record.levelname,
"logger": record.name,
"message": record.getMessage(),
"request_id": getattr(record, "request_id", "-"),
}
# Add exception info if present
if record.exc_info:
log_data["exception"] = self.formatException(record.exc_info)
# Add extra fields
if hasattr(record, "extra_data"):
log_data.update(record.extra_data)
return json.dumps(log_data)
def configure_logging(level: str = "INFO", json_format: bool = False) -> None:
"""Configure logging for the application.
Args:
level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL).
json_format: If True, use JSON formatting; otherwise use text.
"""
# Get the root logger for our package
root_logger = logging.getLogger("llm_inference")
root_logger.setLevel(getattr(logging, level.upper(), logging.INFO))
# Remove existing handlers
for handler in root_logger.handlers[:]:
root_logger.removeHandler(handler)
# Create console handler
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.DEBUG)
# Add request ID filter
handler.addFilter(RequestIdFilter())
# Set formatter based on preference
if json_format:
handler.setFormatter(JsonFormatter())
else:
handler.setFormatter(
logging.Formatter(
"%(asctime)s [%(levelname)s] [%(request_id)s] %(name)s: %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
)
root_logger.addHandler(handler)
# Prevent propagation to root logger
root_logger.propagate = False
def get_logger(name: str) -> logging.Logger:
"""Get a logger with the llm_inference prefix.
Args:
name: Logger name (will be prefixed with 'llm_inference.').
Returns:
logging.Logger: Configured logger instance.
"""
if name.startswith("llm_inference."):
return logging.getLogger(name)
return logging.getLogger(f"llm_inference.{name}")
def set_request_id(request_id: str | None) -> None:
"""Set the request ID for the current context.
Args:
request_id: The request ID to set, or None to clear.
"""
request_id_ctx.set(request_id)
def get_request_id() -> str | None:
"""Get the request ID for the current context.
Returns:
str | None: The current request ID, or None if not set.
"""
return request_id_ctx.get()

63
ai_platform/modules/llm-inference/src/llm_inference/metrics.py Normal file
View file

@ -0,0 +1,63 @@
"""Custom Prometheus metrics for LLM inference (offer §1.10).
Registered on the default prometheus_client registry, so they are exposed by the
existing /metrics endpoint. Gives per-model latency percentiles (p50/p90/p99 via
``histogram_quantile``), token throughput and request/error counts.
VRAM and loaded-model counts per backend are best scraped from the vLLM server's
own /metrics (it exports GPU cache usage natively); this module covers the
gateway-level signals.
"""
from __future__ import annotations
from prometheus_client import Counter, Gauge, Histogram
LLM_LATENCY = Histogram(
"llm_request_duration_seconds",
"End-to-end completion latency at the gateway, per model and backend.",
["model", "backend"],
buckets=(0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 20, 30, 60, 120),
)
LLM_TOKENS = Counter(
"llm_tokens_total",
"Tokens processed, per model, backend and kind (prompt/completion).",
["model", "backend", "kind"],
)
LLM_REQUESTS = Counter(
"llm_requests_total",
"Completion requests, per model, backend and status (success/error).",
["model", "backend", "status"],
)
LLM_INFLIGHT = Gauge(
"llm_inflight_requests",
"In-flight completion requests, per backend.",
["backend"],
)
def observe_success(
model: str, backend: str, duration_s: float, usage: object | None
) -> None:
"""Record a successful completion's latency, tokens and request count."""
LLM_LATENCY.labels(model=model, backend=backend).observe(duration_s)
LLM_REQUESTS.labels(model=model, backend=backend, status="success").inc()
if usage is not None:
prompt = getattr(usage, "prompt_tokens", 0) or 0
completion = getattr(usage, "completion_tokens", 0) or 0
if prompt:
LLM_TOKENS.labels(model=model, backend=backend, kind="prompt").inc(prompt)
if completion:
LLM_TOKENS.labels(
model=model, backend=backend, kind="completion"
).inc(completion)
def observe_error(model: str, backend: str) -> None:
"""Record a failed completion request."""
LLM_REQUESTS.labels(model=model, backend=backend, status="error").inc()

222
ai_platform/modules/llm-inference/src/llm_inference/retry.py Normal file
View file

@ -0,0 +1,222 @@
"""Retry utilities with exponential backoff for transient failures."""
import asyncio
import logging
import random
from collections.abc import Awaitable, Callable
from typing import TypeVar
import httpx
from llm_inference.exceptions import (
CompletionError,
LLMConnectionError,
LLMRateLimitError,
LLMTimeoutError,
)
logger = logging.getLogger("llm_inference.retry")
T = TypeVar("T")
# Import provider-specific exceptions conditionally
try:
import litellm.exceptions as litellm_exc
LITELLM_RETRYABLE = (
litellm_exc.RateLimitError,
litellm_exc.Timeout,
litellm_exc.ServiceUnavailableError,
)
LITELLM_NON_RETRYABLE = (
litellm_exc.AuthenticationError,
litellm_exc.BadRequestError,
)
except ImportError:
LITELLM_RETRYABLE = ()
LITELLM_NON_RETRYABLE = ()
try:
import openai
OPENAI_RETRYABLE = (
openai.RateLimitError,
openai.APIConnectionError,
openai.APITimeoutError,
)
OPENAI_NON_RETRYABLE = (
openai.AuthenticationError,
openai.BadRequestError,
openai.NotFoundError,
)
except ImportError:
OPENAI_RETRYABLE = ()
OPENAI_NON_RETRYABLE = ()
# Combine all retryable exceptions
RETRYABLE_EXCEPTIONS: tuple[type[Exception], ...] = (
*LITELLM_RETRYABLE,
*OPENAI_RETRYABLE,
httpx.TimeoutException,
httpx.ConnectError,
)
# Combine all non-retryable exceptions
NON_RETRYABLE_EXCEPTIONS: tuple[type[Exception], ...] = (
*LITELLM_NON_RETRYABLE,
*OPENAI_NON_RETRYABLE,
)
def is_retryable_exception(exc: Exception) -> bool:
"""Check if an exception is retryable.
Args:
exc: Exception to check.
Returns:
bool: True if the exception is a transient failure that can be retried.
"""
return isinstance(exc, RETRYABLE_EXCEPTIONS)
def extract_retry_after(exc: Exception) -> float | None:
"""Extract retry-after value from rate limit exceptions.
Args:
exc: Exception that may contain retry-after header.
Returns:
float | None: Retry-after value in seconds, or None if not available.
"""
if hasattr(exc, "response") and exc.response is not None:
headers = getattr(exc.response, "headers", {})
retry_after = headers.get("retry-after")
if retry_after:
try:
return float(retry_after)
except ValueError:
pass
return None
def translate_exception(
exc: Exception,
backend: str,
model: str | None = None,
) -> Exception:
"""Translate provider-specific exceptions to our exception types.
Args:
exc: Original exception from provider.
backend: Backend name for error context.
model: Model name for error context.
Returns:
Exception: Translated exception of appropriate type.
"""
# Rate limit errors
if LITELLM_RETRYABLE and isinstance(exc, litellm_exc.RateLimitError):
retry_after = extract_retry_after(exc)
return LLMRateLimitError(str(exc), backend=backend, retry_after=retry_after)
if OPENAI_RETRYABLE and isinstance(exc, openai.RateLimitError):
retry_after = extract_retry_after(exc)
return LLMRateLimitError(str(exc), backend=backend, retry_after=retry_after)
# Timeout errors
if LITELLM_RETRYABLE and isinstance(exc, litellm_exc.Timeout):
return LLMTimeoutError(str(exc), backend=backend)
if OPENAI_RETRYABLE and isinstance(exc, openai.APITimeoutError):
return LLMTimeoutError(str(exc), backend=backend)
if isinstance(exc, httpx.TimeoutException):
return LLMTimeoutError(str(exc), backend=backend)
# Connection errors
if OPENAI_RETRYABLE and isinstance(exc, openai.APIConnectionError):
return LLMConnectionError(backend, reason=str(exc))
if isinstance(exc, httpx.ConnectError):
return LLMConnectionError(backend, reason=str(exc))
# Default to CompletionError
return CompletionError(str(exc), backend=backend, model=model)
async def retry_with_backoff(
func: Callable[[], Awaitable[T]],
max_retries: int,
min_wait: float,
max_wait: float,
backend: str,
model: str | None = None,
) -> T:
"""Execute an async function with exponential backoff retry.
Args:
func: Async function to execute (takes no arguments).
max_retries: Maximum number of retry attempts.
min_wait: Minimum wait time between retries in seconds.
max_wait: Maximum wait time between retries in seconds.
backend: Backend name for error context.
model: Model name for error context.
Returns:
T: The result of the function.
Raises:
LLMRateLimitError: On rate limit with no retries left.
LLMTimeoutError: On timeout with no retries left.
LLMConnectionError: On connection failure with no retries left.
CompletionError: On other API errors.
"""
last_exception: Exception | None = None
for attempt in range(max_retries + 1):
try:
return await func()
except NON_RETRYABLE_EXCEPTIONS as e:
# Don't retry authentication or bad request errors
raise translate_exception(e, backend, model) from e
except RETRYABLE_EXCEPTIONS as e:
last_exception = e
if attempt >= max_retries:
# No more retries
raise translate_exception(e, backend, model) from e
# Calculate wait time with exponential backoff + jitter
retry_after = extract_retry_after(e)
if retry_after:
# Respect the API's retry-after header, but cap at max_wait
# to prevent malicious/buggy APIs from causing indefinite waits
wait_time = min(retry_after, max_wait)
else:
# Exponential backoff: min_wait * 2^attempt (capped at max_wait)
wait_time = min(min_wait * (2**attempt), max_wait)
# Add jitter (0-25% of wait time)
wait_time += random.uniform(0, wait_time * 0.25)
logger.warning(
"Retry %d/%d for %s after %.1fs: %s",
attempt + 1,
max_retries,
backend,
wait_time,
str(e),
)
await asyncio.sleep(wait_time)
except Exception as e:
# Unknown exception - translate and raise without retry
raise translate_exception(e, backend, model) from e
# Should not reach here, but just in case
if last_exception:
raise translate_exception(last_exception, backend, model)
raise CompletionError("Unknown error occurred", backend=backend, model=model)

157
ai_platform/modules/llm-inference/src/llm_inference/runtime_config.py Normal file
View file

@ -0,0 +1,157 @@
"""Runtime config client — fetches config overrides from dashboard.
Polls the dashboard /api/config endpoint periodically and caches values
in memory. All accessors fall back to caller-supplied defaults if the
dashboard is unreachable or the key is missing.
Single source of truth = dashboard `KNOWN_KEYS`. This client carries no
local default registry; consumers pass their own fallback (typically the
pydantic settings value).
"""
import asyncio
import contextlib
import logging
from typing import Any
import httpx
logger = logging.getLogger(__name__)
class RuntimeConfigClient:
"""Polls dashboard /api/config and caches values in-process."""
def __init__(
self,
dashboard_url: str | None,
poll_interval_seconds: int = 30,
request_timeout: float = 5.0,
live_log_logger_name: str | None = None,
live_log_key: str | None = None,
) -> None:
self.dashboard_url = dashboard_url.rstrip("/") if dashboard_url else None
self.poll_interval = poll_interval_seconds
self._timeout = request_timeout
self._cache: dict[str, Any] = {}
self._client: httpx.AsyncClient | None = None
self._task: asyncio.Task[None] | None = None
self._enabled = bool(dashboard_url)
# Optional auto-apply: re-set log level on the named logger when the
# configured key changes value (e.g., `llm.log.level`).
self._live_log_logger_name = live_log_logger_name
self._live_log_key = live_log_key
self._last_log_level: str | None = None
@property
def enabled(self) -> bool:
return self._enabled
def get(self, key: str, default: Any = None) -> Any:
v = self._cache.get(key)
return v if v is not None else default
def get_bool(self, key: str, default: bool = False) -> bool:
v = self._cache.get(key)
return bool(v) if v is not None else default
def get_int(self, key: str, default: int = 0) -> int:
v = self._cache.get(key)
if v is None:
return default
try:
return int(v)
except (TypeError, ValueError):
return default
def get_float(self, key: str, default: float = 0.0) -> float:
v = self._cache.get(key)
if v is None:
return default
try:
return float(v)
except (TypeError, ValueError):
return default
def get_str(self, key: str, default: str = "") -> str:
v = self._cache.get(key)
return str(v) if v is not None else default
async def start(self) -> None:
if not self._enabled:
logger.info("RuntimeConfigClient disabled (no dashboard URL)")
return
self._client = httpx.AsyncClient(
timeout=httpx.Timeout(
connect=2.0, read=self._timeout, write=2.0, pool=5.0
)
)
# Fetch once synchronously so first requests already have overrides
await self._refresh()
self._task = asyncio.create_task(self._loop())
logger.info(
"RuntimeConfigClient started (polling %s every %ds, %d keys cached)",
self.dashboard_url,
self.poll_interval,
len(self._cache),
)
async def stop(self) -> None:
if self._task is not None:
self._task.cancel()
with contextlib.suppress(asyncio.CancelledError, Exception):
await self._task
self._task = None
if self._client is not None and not self._client.is_closed:
await self._client.aclose()
self._client = None
async def _loop(self) -> None:
while True:
try:
await asyncio.sleep(self.poll_interval)
await self._refresh()
except asyncio.CancelledError:
raise
except Exception as e:
logger.debug("Config poll failed: %s", e)
async def _refresh(self) -> None:
if self._client is None or self.dashboard_url is None:
return
try:
resp = await self._client.get(f"{self.dashboard_url}/api/config")
resp.raise_for_status()
data = resp.json()
except Exception as e:
logger.debug("Config refresh failed: %s", e)
return
items = data.get("items", {})
new_cache: dict[str, Any] = {}
for key, entry in items.items():
new_cache[key] = entry.get("value")
self._cache = new_cache
self._maybe_apply_log_level()
def _maybe_apply_log_level(self) -> None:
"""Re-apply log level live if configured key changed."""
if not self._live_log_key or not self._live_log_logger_name:
return
new_level = self.get_str(self._live_log_key)
if not new_level:
return
if new_level == self._last_log_level:
return
try:
level_int = logging.getLevelName(new_level.upper())
if isinstance(level_int, int):
logging.getLogger(self._live_log_logger_name).setLevel(level_int)
self._last_log_level = new_level
logger.info(
"Log level for %s changed to %s (via runtime config)",
self._live_log_logger_name,
new_level,
)
except Exception as e:
logger.warning("Failed to apply log level %s: %s", new_level, e)

223
ai_platform/modules/llm-inference/src/llm_inference/schemas.py Normal file
View file

@ -0,0 +1,223 @@
"""Pydantic V2 API request/response schemas."""
from pydantic import BaseModel, ConfigDict, Field, field_validator
from llm_inference.types import (
BackendType,
ChatMessage,
Choice,
ModelInfo,
StreamChoice,
Usage,
)
# =============================================================================
# Completion Schemas
# =============================================================================
class CompletionRequest(BaseModel):
"""Request schema for chat completions."""
model_config = ConfigDict(extra="forbid")
messages: list[ChatMessage] = Field(
min_length=1,
max_length=1000,
description="List of messages in the conversation (1-1000 messages)",
)
model: str = Field(
description="Model to use for completion",
)
temperature: float = Field(
default=0.7,
ge=0.0,
le=2.0,
description="Sampling temperature",
)
max_tokens: int | None = Field(
default=None,
ge=1,
le=1000000,
description="Maximum tokens to generate (1-1000000)",
)
stream: bool = Field(
default=False,
description="Enable streaming response",
)
backend: BackendType | None = Field(
default=None,
description="Backend to use (overrides default)",
)
top_p: float | None = Field(
default=None,
ge=0.0,
le=1.0,
description="Top-p sampling parameter",
)
frequency_penalty: float | None = Field(
default=None,
ge=-2.0,
le=2.0,
description="Frequency penalty",
)
presence_penalty: float | None = Field(
default=None,
ge=-2.0,
le=2.0,
description="Presence penalty",
)
stop: list[str] | str | None = Field(
default=None,
description="Stop sequences",
)
@field_validator("messages")
@classmethod
def validate_messages_content(cls, v: list[ChatMessage]) -> list[ChatMessage]:
"""Validate that messages have content (except assistant messages)."""
for i, msg in enumerate(v):
# Assistant messages can have empty content (for function calls etc)
if msg.role != "assistant" and msg.content is None:
raise ValueError(
f"Message at index {i} with role '{msg.role}' cannot have empty content"
)
return v
class CompletionResponse(BaseModel):
"""Response schema for chat completions."""
model_config = ConfigDict(extra="forbid")
id: str = Field(description="Unique completion ID")
object: str = Field(default="chat.completion")
created: int = Field(description="Unix timestamp of creation")
model: str = Field(description="Model used for completion")
choices: list[Choice] = Field(description="Completion choices")
usage: Usage | None = Field(default=None, description="Token usage")
backend: str = Field(description="Backend that served the request")
class CompletionChunk(BaseModel):
"""Streaming chunk for chat completions."""
model_config = ConfigDict(extra="forbid")
id: str = Field(description="Unique completion ID")
object: str = Field(default="chat.completion.chunk")
created: int = Field(description="Unix timestamp of creation")
model: str = Field(description="Model used for completion")
choices: list[StreamChoice] = Field(description="Streaming choices")
# =============================================================================
# Text Completion Schemas (legacy OpenAI /v1/completions)
# =============================================================================
class TextChoice(BaseModel):
"""One text-completion choice."""
model_config = ConfigDict(extra="forbid")
index: int = Field(description="Choice index")
text: str = Field(description="Generated text")
finish_reason: str | None = Field(default=None, description="Finish reason")
class TextCompletionRequest(BaseModel):
"""Request schema for text completions (/v1/completions)."""
model_config = ConfigDict(extra="forbid")
prompt: str | list[str] = Field(description="Prompt(s) to complete")
model: str | None = Field(default=None, description="Model (alias) to use")
temperature: float = Field(default=0.7, ge=0.0, le=2.0)
max_tokens: int | None = Field(default=None, ge=1, le=1000000)
backend: BackendType | None = Field(
default=None, description="Backend to use (overrides default)"
)
top_p: float | None = Field(default=None, ge=0.0, le=1.0)
frequency_penalty: float | None = Field(default=None, ge=-2.0, le=2.0)
presence_penalty: float | None = Field(default=None, ge=-2.0, le=2.0)
stop: list[str] | str | None = Field(default=None, description="Stop sequences")
class TextCompletionResponse(BaseModel):
"""Response schema for text completions."""
model_config = ConfigDict(extra="forbid")
id: str = Field(description="Unique completion ID")
object: str = Field(default="text_completion")
created: int = Field(description="Unix timestamp of creation")
model: str = Field(description="Model used for completion")
choices: list[TextChoice] = Field(description="Completion choices")
usage: Usage | None = Field(default=None, description="Token usage")
backend: str = Field(description="Backend that served the request")
# =============================================================================
# Model Management Schemas
# =============================================================================
class ModelListResponse(BaseModel):
"""Response schema for listing models."""
object: str = Field(default="list")
data: list[ModelInfo] = Field(description="List of available models")
class ModelLoadRequest(BaseModel):
"""Request schema for loading a model."""
model: str = Field(description="Model identifier to load")
backend: BackendType = Field(description="Backend to load the model on")
class ModelLoadResponse(BaseModel):
"""Response schema for model load/unload operations."""
success: bool = Field(description="Whether the operation succeeded")
model: str = Field(description="Model identifier")
backend: str = Field(description="Backend the operation was performed on")
message: str | None = Field(default=None, description="Optional message")
# =============================================================================
# Health Schemas
# =============================================================================
class BackendHealth(BaseModel):
"""Health status for a single backend."""
name: str = Field(description="Backend name")
healthy: bool = Field(description="Whether the backend is healthy")
message: str | None = Field(default=None, description="Optional status message")
class HealthResponse(BaseModel):
"""Response schema for health check."""
status: str = Field(description="Overall status: healthy, degraded, unhealthy")
backends: list[BackendHealth] = Field(description="Per-backend health status")
class ReadinessResponse(BaseModel):
"""Response schema for readiness probe."""
ready: bool = Field(default=True)
# =============================================================================
# Backend Schemas
# =============================================================================
class BackendListResponse(BaseModel):
"""Response schema for listing backends."""
backends: list[str] = Field(description="List of available backend names")

87
ai_platform/modules/llm-inference/src/llm_inference/types.py Normal file
View file

@ -0,0 +1,87 @@
"""Core types and enums for LLM inference."""
from enum import Enum
from typing import Literal
from typing import Any
from pydantic import BaseModel, ConfigDict, Field
class BackendType(str, Enum):
"""Supported LLM backends."""
LITELLM = "litellm"
VLLM = "vllm"
LLAMACPP = "llamacpp"
class ChatMessage(BaseModel):
"""A chat message in a conversation."""
model_config = ConfigDict(extra="forbid")
role: Literal["system", "user", "assistant", "function", "tool"] = Field(
description="The role of the message author",
)
content: str | list[dict[str, Any]] | None = Field(
default=None,
description="The content of the message (string or multimodal content array)",
)
name: str | None = Field(
default=None,
description="Optional name for the message author",
)
class Usage(BaseModel):
"""Token usage information."""
prompt_tokens: int = Field(default=0)
completion_tokens: int = Field(default=0)
total_tokens: int = Field(default=0)
class Choice(BaseModel):
"""A completion choice."""
model_config = ConfigDict(extra="forbid")
index: int = Field(default=0)
message: ChatMessage | None = Field(default=None)
finish_reason: str | None = Field(default=None)
class Delta(BaseModel):
"""Streaming delta content."""
model_config = ConfigDict(extra="forbid")
role: str | None = Field(default=None)
content: str | None = Field(default=None)
class StreamChoice(BaseModel):
"""A streaming completion choice."""
model_config = ConfigDict(extra="forbid")
index: int = Field(default=0)
delta: Delta = Field(default_factory=Delta)
finish_reason: str | None = Field(default=None)
class ModelInfo(BaseModel):
"""Information about an available model."""
id: str = Field(description="Model identifier")
backend: str = Field(description="Backend serving this model")
loaded: bool = Field(default=True, description="Whether the model is loaded")
context_length: int | None = Field(
default=None,
description="Maximum context length in tokens",
)
capabilities: list[str] = Field(
default_factory=list,
description="Model capabilities (e.g., 'chat', 'completion')",
)

29
ai_platform/modules/llm-inference/src/llm_inference/utils.py Normal file
View file

@ -0,0 +1,29 @@
"""Shared utilities for LLM inference module."""
import logging
from typing import Any
async def safe_close_stream(
stream: Any | None,
logger: logging.Logger,
) -> None:
"""Safely close an async stream, handling both aclose() and close().
Attempts aclose() first (async close), falls back to close() if not available.
Silently handles any exceptions during cleanup.
Args:
stream: Stream object to close, or None.
logger: Logger for debug messages on cleanup errors.
"""
if stream is None:
return
try:
if hasattr(stream, "aclose"):
await stream.aclose()
elif hasattr(stream, "close"):
await stream.close()
except Exception as e:
logger.debug("Error during stream cleanup: %s", e)