diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..b22c043 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,33 @@ +name: CI + +on: + pull_request: + push: + branches: [main] + +permissions: + contents: read + +concurrency: + group: ci-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + unit: + name: Unit and static analysis + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.13" + cache: pip + - run: python -m pip install -e ".[dev]" + - run: ruff format --check . + - run: ruff check . + - run: mypy + - run: pytest tests/unit --cov=llm_router --cov-report=term-missing --cov-report=xml + - uses: actions/upload-artifact@v4 + with: + name: coverage + path: coverage.xml diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b924909 --- /dev/null +++ b/.gitignore @@ -0,0 +1,15 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.coverage +.mypy_cache/ +.pytest_cache/ +.ruff_cache/ +.venv/ +coverage.xml +htmlcov/ +node_modules/ +playwright-report/ +test-results/ +.env + diff --git a/README.md b/README.md new file mode 100644 index 0000000..cc7b847 --- /dev/null +++ b/README.md @@ -0,0 +1,22 @@ +# Production Local-LLM Inference & Routing Platform + +Policy-aware routing components for a production local-model inference platform. + +The complete architecture and design targets are documented in +[`02-production-local-llm-inference-routing-platform.md`](02-production-local-llm-inference-routing-platform.md). + +## Development + +Requires Python 3.11 or newer. + +```bash +python -m venv .venv +python -m pip install -e ".[dev]" +ruff format --check . +ruff check . +mypy +pytest tests/unit --cov=llm_router --cov-report=term-missing +``` + +This first delivery slice contains deterministic routing, privacy restrictions, quotas, +and bounded admission. API ingress and model-serving adapters are delivered separately. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..8fe78a7 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,56 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "local-llm-router" +version = "0.1.0" +description = "OpenAI-compatible gateway and policy router for local LLM inference." +readme = "README.md" +requires-python = ">=3.11" +dependencies = [ + "fastapi>=0.116,<1", + "pydantic-settings>=2.10,<3", + "uvicorn[standard]>=0.35,<1", +] + +[project.optional-dependencies] +dev = [ + "httpx>=0.28,<1", + "mypy>=1.17,<2", + "pytest>=8.4,<9", + "pytest-asyncio>=1.1,<2", + "pytest-cov>=6.2,<7", + "ruff>=0.12,<1", +] + +[tool.hatch.build.targets.wheel] +packages = ["src/llm_router"] + +[tool.pytest.ini_options] +addopts = "-ra --strict-config --strict-markers" +testpaths = ["tests"] +markers = [ + "integration: tests that exercise multiple application components", +] + +[tool.coverage.run] +branch = true +source = ["llm_router"] + +[tool.coverage.report] +fail_under = 90 +show_missing = true + +[tool.ruff] +line-length = 100 +target-version = "py311" + +[tool.ruff.lint] +select = ["E", "F", "I", "UP", "B", "ASYNC", "RUF"] + +[tool.mypy] +python_version = "3.11" +strict = true +packages = ["llm_router"] + diff --git a/src/llm_router/__init__.py b/src/llm_router/__init__.py new file mode 100644 index 0000000..6838d70 --- /dev/null +++ b/src/llm_router/__init__.py @@ -0,0 +1,3 @@ +"""Local LLM inference gateway and routing policy package.""" + +__version__ = "0.1.0" diff --git a/src/llm_router/admission.py b/src/llm_router/admission.py new file mode 100644 index 0000000..d881c28 --- /dev/null +++ b/src/llm_router/admission.py @@ -0,0 +1,48 @@ +import asyncio +import time +from collections import defaultdict, deque +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + + +class AdmissionRejectedError(RuntimeError): + """Raised when the bounded request queue cannot admit work in time.""" + + +class QuotaExceededError(RuntimeError): + """Raised when a caller exceeds its configured sliding-window quota.""" + + +class AdmissionController: + def __init__(self, max_concurrency: int, timeout_seconds: float) -> None: + self._semaphore = asyncio.Semaphore(max_concurrency) + self._timeout_seconds = timeout_seconds + + @asynccontextmanager + async def slot(self) -> AsyncIterator[None]: + try: + await asyncio.wait_for(self._semaphore.acquire(), timeout=self._timeout_seconds) + except TimeoutError as error: + raise AdmissionRejectedError("inference capacity is saturated") from error + try: + yield + finally: + self._semaphore.release() + + +class SlidingWindowQuota: + def __init__(self, requests_per_minute: int) -> None: + self._limit = requests_per_minute + self._events: dict[str, deque[float]] = defaultdict(deque) + self._lock = asyncio.Lock() + + async def consume(self, subject: str, *, now: float | None = None) -> None: + timestamp = time.monotonic() if now is None else now + cutoff = timestamp - 60 + async with self._lock: + events = self._events[subject] + while events and events[0] <= cutoff: + events.popleft() + if len(events) >= self._limit: + raise QuotaExceededError("request quota exceeded") + events.append(timestamp) diff --git a/src/llm_router/config.py b/src/llm_router/config.py new file mode 100644 index 0000000..a0822fb --- /dev/null +++ b/src/llm_router/config.py @@ -0,0 +1,32 @@ +from functools import lru_cache + +from pydantic import Field, model_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + + +class Settings(BaseSettings): + """Runtime settings loaded from environment variables prefixed with ROUTER_.""" + + model_config = SettingsConfigDict(env_prefix="ROUTER_", extra="ignore") + + environment: str = "development" + api_keys: str = "dev-key" + max_concurrency: int = Field(default=32, ge=1) + admission_timeout_seconds: float = Field(default=0.25, gt=0) + quota_requests_per_minute: int = Field(default=120, ge=1) + external_fallback_enabled: bool = False + + @model_validator(mode="after") + def reject_development_key_in_shared_environments(self) -> "Settings": + if self.environment not in {"development", "test"} and "dev-key" in self.accepted_api_keys: + raise ValueError("ROUTER_API_KEYS must be set outside development and test") + return self + + @property + def accepted_api_keys(self) -> frozenset[str]: + return frozenset(key.strip() for key in self.api_keys.split(",") if key.strip()) + + +@lru_cache +def get_settings() -> Settings: + return Settings() diff --git a/src/llm_router/models.py b/src/llm_router/models.py new file mode 100644 index 0000000..05b6f6a --- /dev/null +++ b/src/llm_router/models.py @@ -0,0 +1,94 @@ +from enum import StrEnum +from typing import Any, Literal + +from pydantic import BaseModel, Field, model_validator + + +class TaskClass(StrEnum): + EXTRACTION = "extraction" + CLASSIFICATION = "classification" + RAG = "rag" + SUMMARIZATION = "summarization" + REASONING = "reasoning" + CRITIQUE = "critique" + GENERAL = "general" + + +class PrivacyClass(StrEnum): + PUBLIC = "public" + PRIVATE = "private" + RESTRICTED = "restricted" + + +class ChatMessage(BaseModel): + role: Literal["system", "user", "assistant", "tool"] + content: str + + +class RoutingOptions(BaseModel): + task: TaskClass | None = None + privacy: PrivacyClass = PrivacyClass.PRIVATE + latency_tier: Literal["interactive", "standard", "batch"] = "standard" + quality_floor: float = Field(default=0.0, ge=0.0, le=1.0) + allow_external_fallback: bool = False + + +class ChatCompletionRequest(BaseModel): + model: str = "auto" + messages: list[ChatMessage] = Field(min_length=1) + max_tokens: int = Field(default=256, ge=1, le=8192) + temperature: float = Field(default=0.0, ge=0.0, le=2.0) + stream: bool = False + routing: RoutingOptions = Field(default_factory=RoutingOptions) + + @model_validator(mode="after") + def reject_streaming_for_initial_slice(self) -> "ChatCompletionRequest": + if self.stream: + raise ValueError("streaming is not available in the initial control-plane slice") + return self + + @property + def prompt(self) -> str: + return "\n".join(message.content for message in self.messages) + + +class ModelProfile(BaseModel): + id: str + revision: str + local: bool + healthy: bool = True + context_limit: int + supported_tasks: frozenset[TaskClass] + quality: float = Field(ge=0.0, le=1.0) + estimated_queue_ms: int = Field(default=0, ge=0) + cost_weight: float = Field(default=0.0, ge=0.0) + + +class RouteDecision(BaseModel): + profile: ModelProfile + task: TaskClass + reason: str + score: float + candidate_count: int + + +class ChatCompletionChoice(BaseModel): + index: int = 0 + message: ChatMessage + finish_reason: Literal["stop", "length"] = "stop" + + +class Usage(BaseModel): + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +class ChatCompletionResponse(BaseModel): + id: str + object: Literal["chat.completion"] = "chat.completion" + created: int + model: str + choices: list[ChatCompletionChoice] + usage: Usage + routing: dict[str, Any] diff --git a/src/llm_router/py.typed b/src/llm_router/py.typed new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/src/llm_router/py.typed @@ -0,0 +1 @@ + diff --git a/src/llm_router/routing.py b/src/llm_router/routing.py new file mode 100644 index 0000000..ad3c409 --- /dev/null +++ b/src/llm_router/routing.py @@ -0,0 +1,150 @@ +from dataclasses import dataclass + +from llm_router.models import ( + ChatCompletionRequest, + ModelProfile, + PrivacyClass, + RouteDecision, + TaskClass, +) + + +class NoEligibleModelError(RuntimeError): + """Raised when policy removes every model candidate.""" + + +def default_model_profiles() -> tuple[ModelProfile, ...]: + return ( + ModelProfile( + id="small-specialist", + revision="mock-small@sha256:dev", + local=True, + context_limit=8192, + supported_tasks=frozenset({TaskClass.EXTRACTION, TaskClass.CLASSIFICATION}), + quality=0.82, + estimated_queue_ms=12, + cost_weight=0.1, + ), + ModelProfile( + id="general-local", + revision="mock-general@sha256:dev", + local=True, + context_limit=32768, + supported_tasks=frozenset( + { + TaskClass.EXTRACTION, + TaskClass.CLASSIFICATION, + TaskClass.RAG, + TaskClass.SUMMARIZATION, + TaskClass.GENERAL, + } + ), + quality=0.89, + estimated_queue_ms=35, + cost_weight=0.35, + ), + ModelProfile( + id="high-capability", + revision="mock-high@sha256:dev", + local=True, + context_limit=65536, + supported_tasks=frozenset(TaskClass), + quality=0.96, + estimated_queue_ms=90, + cost_weight=0.9, + ), + ModelProfile( + id="approved-external-fallback", + revision="external-policy-v1", + local=False, + context_limit=128000, + supported_tasks=frozenset(TaskClass), + quality=0.98, + estimated_queue_ms=45, + cost_weight=1.5, + ), + ) + + +@dataclass(frozen=True) +class Router: + profiles: tuple[ModelProfile, ...] + external_fallback_enabled: bool = False + + def classify_task(self, request: ChatCompletionRequest) -> TaskClass: + if request.routing.task is not None: + return request.routing.task + + prompt = request.prompt.lower() + keywords = ( + (TaskClass.EXTRACTION, ("extract", "json schema", "fields from")), + (TaskClass.CLASSIFICATION, ("classify", "choose one label", "category")), + (TaskClass.SUMMARIZATION, ("summarize", "summary")), + (TaskClass.CRITIQUE, ("critique", "find flaws")), + (TaskClass.REASONING, ("reason step", "prove", "analyze deeply")), + (TaskClass.RAG, ("provided context", "according to the documents")), + ) + return next( + (task for task, terms in keywords if any(term in prompt for term in terms)), + TaskClass.GENERAL, + ) + + def select(self, request: ChatCompletionRequest) -> RouteDecision: + task = self.classify_task(request) + estimated_tokens = max(1, len(request.prompt) // 4) + request.max_tokens + + candidates = [ + profile + for profile in self.profiles + if profile.healthy + and task in profile.supported_tasks + and estimated_tokens <= profile.context_limit + and profile.quality >= request.routing.quality_floor + and self._privacy_allows(profile, request.routing.privacy) + and self._external_allows(profile, request) + ] + + if request.model != "auto": + candidates = [profile for profile in candidates if profile.id == request.model] + + if not candidates: + raise NoEligibleModelError( + "no healthy model satisfies capability, context, quality, " + "privacy, and fallback policy" + ) + + def score(profile: ModelProfile) -> float: + latency_penalty = profile.estimated_queue_ms / 100 + if request.routing.latency_tier == "interactive": + latency_penalty *= 2 + elif request.routing.latency_tier == "batch": + latency_penalty *= 0.5 + specialization_bonus = 10 if len(profile.supported_tasks) <= 2 else 0 + return ( + (profile.quality * 100) + + specialization_bonus + - latency_penalty + - (profile.cost_weight * 10) + ) + + selected = max(candidates, key=score) + return RouteDecision( + profile=selected, + task=task, + reason=( + f"selected highest policy score among {len(candidates)} eligible model(s); " + f"task={task.value}, privacy={request.routing.privacy.value}, " + f"latency_tier={request.routing.latency_tier}" + ), + score=round(score(selected), 3), + candidate_count=len(candidates), + ) + + @staticmethod + def _privacy_allows(profile: ModelProfile, privacy: PrivacyClass) -> bool: + return profile.local or privacy == PrivacyClass.PUBLIC + + def _external_allows(self, profile: ModelProfile, request: ChatCompletionRequest) -> bool: + return profile.local or ( + self.external_fallback_enabled and request.routing.allow_external_fallback + ) diff --git a/tests/unit/test_admission.py b/tests/unit/test_admission.py new file mode 100644 index 0000000..83e6caa --- /dev/null +++ b/tests/unit/test_admission.py @@ -0,0 +1,38 @@ +import asyncio + +import pytest + +from llm_router.admission import ( + AdmissionController, + AdmissionRejectedError, + QuotaExceededError, + SlidingWindowQuota, +) + + +@pytest.mark.asyncio +async def test_quota_uses_sliding_window() -> None: + quota = SlidingWindowQuota(requests_per_minute=2) + await quota.consume("tenant-a", now=100) + await quota.consume("tenant-a", now=101) + with pytest.raises(QuotaExceededError): + await quota.consume("tenant-a", now=102) + await quota.consume("tenant-a", now=161) + + +@pytest.mark.asyncio +async def test_quotas_are_isolated_by_subject() -> None: + quota = SlidingWindowQuota(requests_per_minute=1) + await quota.consume("tenant-a", now=100) + await quota.consume("tenant-b", now=100) + + +@pytest.mark.asyncio +async def test_admission_is_bounded() -> None: + admission = AdmissionController(max_concurrency=1, timeout_seconds=0.01) + async with admission.slot(): + with pytest.raises(AdmissionRejectedError): + async with admission.slot(): + pytest.fail("saturated request must not be admitted") + async with admission.slot(): + await asyncio.sleep(0) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py new file mode 100644 index 0000000..92f00f7 --- /dev/null +++ b/tests/unit/test_config.py @@ -0,0 +1,19 @@ +import pytest +from pydantic import ValidationError + +from llm_router.config import Settings + + +def test_api_keys_are_normalized() -> None: + settings = Settings(environment="test", api_keys=" first, second ,, ") + assert settings.accepted_api_keys == frozenset({"first", "second"}) + + +def test_production_rejects_default_development_key() -> None: + with pytest.raises(ValidationError, match="ROUTER_API_KEYS"): + Settings(environment="production") + + +def test_production_accepts_explicit_key() -> None: + settings = Settings(environment="production", api_keys="production-key") + assert settings.accepted_api_keys == frozenset({"production-key"}) diff --git a/tests/unit/test_routing.py b/tests/unit/test_routing.py new file mode 100644 index 0000000..ca98b65 --- /dev/null +++ b/tests/unit/test_routing.py @@ -0,0 +1,85 @@ +import pytest + +from llm_router.models import ( + ChatCompletionRequest, + ChatMessage, + PrivacyClass, + RoutingOptions, + TaskClass, +) +from llm_router.routing import NoEligibleModelError, Router, default_model_profiles + + +def request_for( + prompt: str, + *, + task: TaskClass | None = None, + privacy: PrivacyClass = PrivacyClass.PRIVATE, + model: str = "auto", + quality_floor: float = 0, + allow_external: bool = False, +) -> ChatCompletionRequest: + return ChatCompletionRequest( + model=model, + messages=[ChatMessage(role="user", content=prompt)], + routing=RoutingOptions( + task=task, + privacy=privacy, + quality_floor=quality_floor, + allow_external_fallback=allow_external, + ), + ) + + +def test_routes_extraction_to_small_specialist() -> None: + router = Router(default_model_profiles()) + decision = router.select(request_for("Extract the invoice fields as JSON")) + assert decision.task == TaskClass.EXTRACTION + assert decision.profile.id == "small-specialist" + assert "privacy=private" in decision.reason + assert decision.candidate_count == 3 + + +def test_quality_floor_promotes_request_to_high_capability() -> None: + router = Router(default_model_profiles()) + decision = router.select(request_for("Summarize this report", quality_floor=0.95)) + assert decision.profile.id == "high-capability" + + +def test_private_request_cannot_select_external_model() -> None: + router = Router(default_model_profiles(), external_fallback_enabled=True) + with pytest.raises(NoEligibleModelError, match="privacy"): + router.select( + request_for( + "Analyze deeply", + model="approved-external-fallback", + privacy=PrivacyClass.RESTRICTED, + allow_external=True, + ) + ) + + +def test_external_route_needs_operator_and_request_opt_in() -> None: + payload = request_for( + "Analyze deeply", + model="approved-external-fallback", + privacy=PrivacyClass.PUBLIC, + allow_external=True, + ) + with pytest.raises(NoEligibleModelError): + Router(default_model_profiles()).select(payload) + decision = Router(default_model_profiles(), external_fallback_enabled=True).select(payload) + assert decision.profile.id == "approved-external-fallback" + + +@pytest.mark.parametrize( + ("prompt", "expected"), + [ + ("Choose one label for this item", TaskClass.CLASSIFICATION), + ("Critique this architecture", TaskClass.CRITIQUE), + ("Hello there", TaskClass.GENERAL), + ], +) +def test_task_classifier(prompt: str, expected: TaskClass) -> None: + router = Router(default_model_profiles()) + assert router.classify_task(request_for(prompt)) == expected