llmock 0.1.0__py3-none-any.whl

llmock/__init__.py

@@ -0,0 +1,5 @@
+"""LLMock - OpenAI-compatible mock server for LLM API resilience testing."""
+
+__version__ = "0.1.0"
+
+__all__ = ["__version__"]

llmock/chaos.py

@@ -0,0 +1,187 @@
+"""Chaos Engineering middleware for LLMock."""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import random
+
+from fastapi import Request
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from llmock.simulation import (
+    ERROR_RATE_ENV_PREFIX,
+    SUPPORTED_ERROR_STATUS_CODES,
+    build_error_response,
+)
+
+
+class ChaosSettings:
+    """Runtime settings for latency and injected HTTP errors."""
+
+    def __init__(
+        self,
+        *,
+        latency_ms: int = 0,
+        error_rates: dict[int, float] | None = None,
+        error_rate_429: float = 0.0,
+        error_rate_500: float = 0.0,
+        error_rate_503: float = 0.0,
+        **named_error_rates: float,
+    ) -> None:
+        object.__setattr__(self, "latency_ms", latency_ms)
+        object.__setattr__(self, "_error_rates", {})
+
+        combined = dict(error_rates or {})
+        if error_rate_429 > 0.0:
+            combined[429] = error_rate_429
+        if error_rate_500 > 0.0:
+            combined[500] = error_rate_500
+        if error_rate_503 > 0.0:
+            combined[503] = error_rate_503
+
+        for key, value in named_error_rates.items():
+            status_code = self._status_code_from_attr(key)
+            if status_code is None:
+                raise TypeError(f"Unexpected chaos setting '{key}'.")
+            combined[status_code] = value
+
+        self.error_rates = combined
+
+    def __getattr__(self, name: str) -> float:
+        status_code = self._status_code_from_attr(name)
+        if status_code is not None:
+            return self._error_rates.get(status_code, 0.0)
+        raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: int | float) -> None:
+        status_code = self._status_code_from_attr(name)
+        if status_code is not None:
+            self._error_rates[status_code] = float(value)
+            return
+        object.__setattr__(self, name, value)
+
+    @property
+    def error_rates(self) -> dict[int, float]:
+        return dict(self._error_rates)
+
+    @error_rates.setter
+    def error_rates(self, rates: dict[int, float]) -> None:
+        normalized: dict[int, float] = {}
+        for status, rate in rates.items():
+            normalized[int(status)] = float(rate)
+        object.__setattr__(self, "_error_rates", normalized)
+
+    @staticmethod
+    def _status_code_from_attr(name: str) -> int | None:
+        if not name.startswith("error_rate_"):
+            return None
+        suffix = name.removeprefix("error_rate_")
+        if not suffix.isdigit():
+            return None
+        return int(suffix)
+
+    @classmethod
+    def from_env(cls) -> "ChaosSettings":
+        parsed_rates: dict[int, float] = {}
+        for key, value in os.environ.items():
+            if not key.startswith(ERROR_RATE_ENV_PREFIX):
+                continue
+            suffix = key[len(ERROR_RATE_ENV_PREFIX) :]
+            if not suffix.isdigit():
+                continue
+            parsed_rates[int(suffix)] = float(value)
+        return cls(
+            latency_ms=int(os.getenv("LLMOCK_LATENCY_MS", "0")),
+            error_rates=parsed_rates,
+        ).validated()
+
+    def with_overrides(
+        self,
+        *,
+        latency_ms: int | None = None,
+        error_rates: dict[int, float] | None = None,
+        error_rate_429: float | None = None,
+        error_rate_500: float | None = None,
+        error_rate_503: float | None = None,
+    ) -> "ChaosSettings":
+        combined = dict(self.error_rates)
+        if error_rates:
+            combined.update(error_rates)
+        if error_rate_429 is not None:
+            combined[429] = error_rate_429
+        if error_rate_500 is not None:
+            combined[500] = error_rate_500
+        if error_rate_503 is not None:
+            combined[503] = error_rate_503
+        return type(self)(
+            latency_ms=self.latency_ms if latency_ms is None else latency_ms,
+            error_rates=combined,
+        ).validated()
+
+    def validated(self) -> "ChaosSettings":
+        if self.latency_ms < 0:
+            raise ValueError("LLMOCK_LATENCY_MS must be greater than or equal to 0.")
+
+        total_probability = 0.0
+        for status, rate in self.error_rates.items():
+            if status not in SUPPORTED_ERROR_STATUS_CODES:
+                raise ValueError(
+                    f"LLMOCK_ERROR_RATE_{status} must target an HTTP error code between 400 and 599."
+                )
+            if not 0.0 <= rate <= 1.0:
+                raise ValueError(f"LLMOCK_ERROR_RATE_{status} must be between 0.0 and 1.0.")
+            total_probability += rate
+
+        if total_probability > 1.0:
+            raise ValueError("The sum of simulated error probabilities must be <= 1.0.")
+
+        return self
+
+    def as_env(self) -> dict[str, str]:
+        env = {"LLMOCK_LATENCY_MS": str(self.latency_ms)}
+        for status, rate in self.error_rates.items():
+            env[f"{ERROR_RATE_ENV_PREFIX}{status}"] = str(rate)
+        return env
+
+
+chaos_settings = ChaosSettings.from_env()
+
+
+class ChaosMiddleware(BaseHTTPMiddleware):
+    """Middleware that injects latency and random errors based on ChaosSettings."""
+
+    def __init__(self, app, settings: ChaosSettings | None = None) -> None:
+        super().__init__(app)
+        self.settings = (settings or ChaosSettings.from_env()).validated()
+
+    async def dispatch(self, request: Request, call_next):
+        cfg = self.settings
+
+        if request.url.path == "/health":
+            return await call_next(request)
+
+        forced_status = request.headers.get("x-llmock-force-status")
+        if forced_status:
+            status_code = int(forced_status)
+            if 400 <= status_code < 600:
+                return build_error_response(request.url.path, status_code)
+
+        if cfg.latency_ms > 0:
+            await asyncio.sleep(cfg.latency_ms / 1000.0)
+
+        sampled_status = _sample_error_status(cfg)
+        if sampled_status is not None:
+            return build_error_response(request.url.path, sampled_status)
+
+        return await call_next(request)
+
+
+def _sample_error_status(settings: ChaosSettings) -> int | None:
+    r = random.random()
+    cumulative = 0.0
+    for status, rate in sorted(settings.error_rates.items()):
+        cumulative += rate
+        if r < cumulative:
+            return status
+    return None

llmock/cli.py

@@ -0,0 +1,275 @@
+"""Typer CLI entry point for LLMock."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+import typer
+import uvicorn
+import yaml
+
+from llmock.chaos import ChaosSettings
+from llmock.simulation import ERROR_RATE_ENV_PREFIX, MockResponseSettings, SUPPORTED_ERROR_STATUS_CODES
+
+app = typer.Typer(
+    name="llmock",
+    help="OpenAI-compatible mock server for LLM API resilience testing.",
+    no_args_is_help=True,
+)
+
+ERROR_RATE_OPTION_HELP = (
+    "Inject any HTTP error with STATUS=PROBABILITY. Repeat the option for as many 4xx/5xx "
+    "statuses as needed, for example: --error-rate 400=0.05 --error-rate 429=0.2 --error-rate 503=0.1."
+)
+
+
+def _resolve_config_path(*, config: str | None) -> str | None:
+    return config or os.getenv("LLMOCK_CONFIG")
+
+
+def _load_config_file(path: str | None) -> dict[str, Any]:
+    if not path:
+        return {}
+
+    try:
+        with open(path, "r", encoding="utf-8") as handle:
+            raw = handle.read()
+
+        lower_path = path.lower()
+        if lower_path.endswith(".json"):
+            data = json.loads(raw)
+        elif lower_path.endswith((".yaml", ".yml")):
+            data = yaml.safe_load(raw) or {}
+        else:
+            raise typer.BadParameter(
+                "Unsupported config file format. Use a .json, .yaml, or .yml file.",
+                param_hint="--config",
+            )
+    except OSError as exc:
+        raise typer.BadParameter(
+            f"Could not read config file: {exc}",
+            param_hint="--config",
+        ) from exc
+    except (json.JSONDecodeError, yaml.YAMLError) as exc:
+        raise typer.BadParameter(
+            f"Could not parse config file: {exc}",
+            param_hint="--config",
+        ) from exc
+
+    if not isinstance(data, dict):
+        raise typer.BadParameter(
+            "Config file root must be a JSON/YAML object.",
+            param_hint="--config",
+        )
+    return data
+
+
+def _config_scalar(config: dict[str, Any], key: str, *, section: str | None = None) -> Any:
+    if section:
+        nested = config.get(section)
+        if isinstance(nested, dict) and key in nested:
+            return nested[key]
+    return config.get(key)
+
+
+def _config_error_rates(config: dict[str, Any]) -> dict[int, float]:
+    parsed: dict[int, float] = {}
+
+    for source in (config, config.get("chaos", {})):
+        if not isinstance(source, dict):
+            continue
+
+        mapping = source.get("error_rates")
+        if mapping is not None:
+            if not isinstance(mapping, dict):
+                raise typer.BadParameter(
+                    "Config key 'error_rates' must be an object mapping status codes to probabilities.",
+                    param_hint="--config",
+                )
+            for status, rate in mapping.items():
+                parsed[int(status)] = float(rate)
+
+        for key, value in source.items():
+            if not isinstance(key, str) or not key.startswith("error_rate_"):
+                continue
+            suffix = key.removeprefix("error_rate_")
+            if suffix.isdigit():
+                parsed[int(suffix)] = float(value)
+
+    return parsed
+
+
+def _parse_error_rate_options(values: list[str] | None) -> dict[int, float]:
+    parsed: dict[int, float] = {}
+    for value in values or []:
+        try:
+            status_str, rate_str = value.split("=", 1)
+            status_code = int(status_str)
+            rate = float(rate_str)
+        except ValueError as exc:
+            raise typer.BadParameter(
+                "Each --error-rate value must use the form STATUS=PROBABILITY, e.g. 429=0.4."
+            ) from exc
+
+        if status_code not in SUPPORTED_ERROR_STATUS_CODES:
+            raise typer.BadParameter(
+                f"Unsupported status code {status_code}. Use an HTTP error code between 400 and 599."
+            )
+
+        parsed[status_code] = rate
+    return parsed
+
+
+def _set_server_env(*, chaos: ChaosSettings, responses: MockResponseSettings) -> None:
+    for key in list(os.environ):
+        if key.startswith(ERROR_RATE_ENV_PREFIX):
+            del os.environ[key]
+
+    for key, value in chaos.as_env().items():
+        os.environ[key] = value
+    for key, value in responses.as_env().items():
+        os.environ[key] = value
+
+
+def _resolve_chaos_settings(
+    *,
+    config: dict[str, Any],
+    latency_ms: int | None,
+    error_rates: list[str] | None = None,
+    error_rate_429: float | None,
+    error_rate_500: float | None,
+    error_rate_503: float | None,
+) -> ChaosSettings:
+    base = ChaosSettings(
+        latency_ms=int(_config_scalar(config, "latency_ms", section="chaos") or 0),
+        error_rates=_config_error_rates(config),
+    ).validated()
+
+    env_settings = ChaosSettings.from_env()
+    env_latency = os.getenv("LLMOCK_LATENCY_MS")
+    return base.with_overrides(
+        latency_ms=int(env_latency) if env_latency is not None else None,
+        error_rates=env_settings.error_rates,
+    ).with_overrides(
+        latency_ms=latency_ms,
+        error_rates=_parse_error_rate_options(error_rates),
+        error_rate_429=error_rate_429,
+        error_rate_500=error_rate_500,
+        error_rate_503=error_rate_503,
+    )
+
+
+def _resolve_mock_response_settings(*, config: dict[str, Any], response_style: str | None) -> MockResponseSettings:
+    settings = MockResponseSettings(
+        response_style=str(
+            _config_scalar(config, "response_style")
+            or _config_scalar(config, "response_style", section="responses")
+            or _config_scalar(config, "style", section="responses")
+            or "varied"
+        )
+    ).validated()
+    env_style = os.getenv("LLMOCK_RESPONSE_STYLE")
+    if env_style is not None:
+        settings.response_style = env_style
+    if response_style is not None:
+        settings.response_style = response_style
+    return settings.validated()
+
+
+def _resolve_server_host(*, config: dict[str, Any], host: str | None) -> str:
+    return str(host or os.getenv("LLMOCK_HOST") or _config_scalar(config, "host", section="server") or "127.0.0.1")
+
+
+def _resolve_server_port(*, config: dict[str, Any], port: int | None) -> int:
+    if port is not None:
+        return port
+    return int(os.getenv("LLMOCK_PORT") or _config_scalar(config, "port", section="server") or "8000")
+
+
+def _format_error_rates(settings: ChaosSettings) -> str:
+    return ", ".join(
+        f"{status}={rate:.0%}" for status, rate in sorted(settings.error_rates.items()) if rate
+    )
+
+
+@app.callback()
+def main() -> None:
+    """LLMock command group."""
+
+
+@app.command()
+def serve(
+    config: str | None = typer.Option(
+        None,
+        "--config",
+        help="Load startup settings from a JSON or YAML file. CLI flags override env vars, and env vars override the config file.",
+    ),
+    host: str | None = typer.Option(None, "--host", "-h", help="Override LLMOCK_HOST."),
+    port: int | None = typer.Option(None, "--port", "-p", help="Override LLMOCK_PORT."),
+    reload: bool = typer.Option(False, "--reload", help="Enable auto-reload (dev mode)"),
+    log_level: str = typer.Option("info", "--log-level", help="Uvicorn log level"),
+    latency_ms: int | None = typer.Option(None, "--latency-ms", help="Override LLMOCK_LATENCY_MS."),
+    error_rates: list[str] | None = typer.Option(
+        None,
+        "--error-rate",
+        help=ERROR_RATE_OPTION_HELP,
+    ),
+    error_rate_429: float | None = typer.Option(
+        None,
+        "--error-rate-429",
+        help="Shortcut for --error-rate 429=RATE.",
+    ),
+    error_rate_500: float | None = typer.Option(
+        None,
+        "--error-rate-500",
+        help="Shortcut for --error-rate 500=RATE.",
+    ),
+    error_rate_503: float | None = typer.Option(
+        None,
+        "--error-rate-503",
+        help="Shortcut for --error-rate 503=RATE.",
+    ),
+    response_style: str | None = typer.Option(
+        None,
+        "--response-style",
+        help="Override LLMOCK_RESPONSE_STYLE (static, hello, echo, varied).",
+    ),
+) -> None:
+    """Start the LLMock server."""
+    config_path = _resolve_config_path(config=config)
+    loaded_config = _load_config_file(config_path)
+    resolved_host = _resolve_server_host(config=loaded_config, host=host)
+    resolved_port = _resolve_server_port(config=loaded_config, port=port)
+    chaos = _resolve_chaos_settings(
+        config=loaded_config,
+        latency_ms=latency_ms,
+        error_rates=error_rates,
+        error_rate_429=error_rate_429,
+        error_rate_500=error_rate_500,
+        error_rate_503=error_rate_503,
+    )
+    responses = _resolve_mock_response_settings(config=loaded_config, response_style=response_style)
+    _set_server_env(chaos=chaos, responses=responses)
+
+    configured_rates = _format_error_rates(chaos)
+    if chaos.latency_ms or configured_rates:
+        typer.echo(f"Chaos: latency={chaos.latency_ms}ms  errors=[{configured_rates or 'none'}]")
+
+    if config_path:
+        typer.echo(f"Config: {config_path}")
+    typer.echo(f"Responses: style={responses.response_style}")
+    typer.echo(f"Starting LLMock on http://{resolved_host}:{resolved_port}")
+    uvicorn.run(
+        "llmock.main:create_app",
+        factory=True,
+        host=resolved_host,
+        port=resolved_port,
+        reload=reload,
+        log_level=log_level,
+    )
+
+
+if __name__ == "__main__":
+    app()

llmock/errors.py

@@ -0,0 +1,91 @@
+"""Provider-specific exception handlers for LLMock.
+
+Overrides FastAPI's default `{"detail": "..."}` shape with the real
+error envelope each provider uses so that clients can parse them normally.
+"""
+
+from __future__ import annotations
+
+from fastapi import FastAPI, Request
+from fastapi.exceptions import RequestValidationError
+from starlette.exceptions import HTTPException
+from starlette.responses import JSONResponse
+
+from llmock.simulation import _build_error_content, provider_from_path
+
+
+def _message_from_exc(exc: HTTPException) -> str:
+    detail = exc.detail
+    if isinstance(detail, str):
+        return detail
+    if isinstance(detail, dict):
+        return str(detail.get("msg") or detail.get("message") or exc.detail)
+    return str(detail)
+
+
+def _message_from_validation(exc: RequestValidationError) -> str:
+    errors = exc.errors()
+    if not errors:
+        return "Request validation error."
+    first = errors[0]
+    loc = " → ".join(str(p) for p in first.get("loc", []))
+    msg = first.get("msg", "validation error")
+    return f"{loc}: {msg}" if loc else msg
+
+
+async def _http_exception_handler(request: Request, exc: HTTPException) -> JSONResponse:
+    provider = provider_from_path(request.url.path)
+    message = _message_from_exc(exc)
+    content = _build_error_content(provider=provider, status_code=exc.status_code)
+
+    # Inject the actual message instead of the generic label.
+    _inject_message(content, message, provider)
+
+    headers = dict(exc.headers or {})
+    if exc.status_code in {429, 503, 504, 529} and "retry-after" not in {k.lower() for k in headers}:
+        headers["retry-after"] = "1"
+
+    return JSONResponse(status_code=exc.status_code, content=content, headers=headers or None)
+
+
+async def _validation_exception_handler(
+    request: Request, exc: RequestValidationError
+) -> JSONResponse:
+    provider = provider_from_path(request.url.path)
+    status_code = 422
+    message = _message_from_validation(exc)
+    content = _build_error_content(provider=provider, status_code=status_code)
+    _inject_message(content, message, provider)
+    return JSONResponse(status_code=status_code, content=content)
+
+
+def _inject_message(content: dict, message: str, provider: str) -> None:
+    """Mutate *content* in-place to use *message* instead of the generic label."""
+    if provider == "anthropic":
+        # {"type": "error", "error": {"type": "...", "message": "..."}}
+        if isinstance(content.get("error"), dict):
+            content["error"]["message"] = message
+    elif provider == "gemini":
+        # {"error": {"code": ..., "message": "...", "status": "..."}}
+        if isinstance(content.get("error"), dict):
+            content["error"]["message"] = message
+    elif provider == "cohere":
+        # {"message": "..."}
+        content["message"] = message
+    elif provider == "mistral":
+        # {"object": "error", "message": "...", "type": "...", "param": null, "code": "..."}
+        content["message"] = message
+    elif provider == "xai":
+        # {"error": {"message": "...", "type": "...", "code": "..."}}
+        if isinstance(content.get("error"), dict):
+            content["error"]["message"] = message
+    else:
+        # OpenAI-compatible: {"error": {"message": "...", "type": "...", ...}}
+        if isinstance(content.get("error"), dict):
+            content["error"]["message"] = message
+
+
+def register_error_handlers(app: FastAPI) -> None:
+    """Register provider-aware error handlers on *app*."""
+    app.add_exception_handler(HTTPException, _http_exception_handler)  # type: ignore[arg-type]
+    app.add_exception_handler(RequestValidationError, _validation_exception_handler)  # type: ignore[arg-type]

llmock/main.py

@@ -0,0 +1,53 @@
+"""FastAPI application factory for LLMock."""
+
+from fastapi import FastAPI
+
+from llmock import __version__
+from llmock.chaos import ChaosMiddleware, ChaosSettings, chaos_settings
+from llmock.errors import register_error_handlers
+
+# Import every router module so their registry.register() calls fire at import time.
+import llmock.routers.ai21  # noqa: F401
+import llmock.routers.anthropic  # noqa: F401
+import llmock.routers.cohere  # noqa: F401
+import llmock.routers.gemini  # noqa: F401
+import llmock.routers.groq  # noqa: F401
+import llmock.routers.mistral  # noqa: F401
+import llmock.routers.openai  # noqa: F401
+import llmock.routers.perplexity  # noqa: F401
+import llmock.routers.together  # noqa: F401
+import llmock.routers.xai  # noqa: F401
+
+from llmock.routers.registry import get_all_routers
+from llmock.simulation import MockResponseSettings
+
+
+def create_app(
+    chaos: ChaosSettings | None = None,
+    responses: MockResponseSettings | None = None,
+) -> FastAPI:
+    settings = (chaos or ChaosSettings.from_env()).validated()
+    response_settings = (responses or MockResponseSettings.from_env()).validated()
+
+    app = FastAPI(
+        title="LLMock",
+        description="OpenAI-compatible mock server for LLM API resilience testing",
+        version=__version__,
+    )
+
+    app.state.chaos_settings = settings
+    app.state.mock_response_settings = response_settings
+    register_error_handlers(app)
+    app.add_middleware(ChaosMiddleware, settings=settings)
+
+    for router in get_all_routers():
+        app.include_router(router)
+
+    @app.get("/health")
+    def health() -> dict:
+        return {"status": "ok", "version": __version__}
+
+    return app
+
+
+app = create_app(chaos=chaos_settings)