diff --git a/docs/content/docs/documentation/API.mdx b/docs/content/docs/documentation/API.mdx index 9d93a58f8..2188e7394 100644 --- a/docs/content/docs/documentation/API.mdx +++ b/docs/content/docs/documentation/API.mdx @@ -715,7 +715,7 @@ OpenAI-compatible text completion endpoint. | `websearch` | `bool` | `false` | Augments the RAG context with live web search results. When used with a partition (`openrag-{partition}`), document and web results are combined. When used without a partition (direct LLM mode), web results are the sole context. Requires `WEBSEARCH_API_TOKEN` to be configured. See [web search configuration](/openrag/documentation/env_vars/#web-search-configuration). | | `spoken_style_answer` | `bool` | `false` | Generates a succinct spoken-style conversational answer based on the retrieved documents. | | `use_map_reduce` | `bool` | `false` | Uses a map-reduce strategy to aggregate information from multiple documents. See [map-reduce configuration](/openrag/documentation/env_vars/#map--reduce-configuration). | -| `llm_override` | `object` | `null` | Overrides only the downstream LLM model name while still using OpenRAG's configured LLM endpoint and credentials. Accepts: `model` (string). Endpoint URL and API key are server-side configuration and cannot be changed by a client request. | +| `llm_override` | `object` | `null` | Overrides the downstream LLM for this request. Accepts `model` (string), always honored. Also accepts `base_url` and `api_key`, which are honored **only** when the deployment sets `LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT` — otherwise they are ignored and the request goes to the server's configured endpoint. See [custom LLM endpoints](/openrag/documentation/env_vars/#client-supplied-llm-endpoints). | | `attachments` | `list[{"id": string}]` | `null` | Scopes RAG retrieval to a specific list of file IDs within the target partition, instead of searching the whole partition. Unknown or unindexed IDs are silently dropped; duplicates are deduplicated. The response's `extra.attachments` reports which IDs were actually used. | Examples: @@ -803,6 +803,30 @@ curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ }' ``` +With `LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT=true`, the same object may also carry the +endpoint and its credential, sending the request to a provider of the client's +choosing instead of the configured one: + +```json title="metadata.llm_override with a client-supplied endpoint" +{ + "llm_override": { + "base_url": "https://api.openai.com/v1", + "api_key": "sk-...", + "model": "gpt-4o" + } +} +``` + +`base_url` must be `https`, must carry no query string, fragment or `..` path +segment (percent-encoded or not), and is always requested as +`{base_url}/chat/completions` — `{base_url}/completions` when the request comes +in on the legacy `/v1/completions` route; anything else is rejected with a +**400**. The server's own API key is never forwarded — an override without +`api_key` sends no `Authorization` header at all. When the flag +is off, `base_url`/`api_key` are ignored (a warning is logged) and only `model` +applies, which typically surfaces as an "unknown model" error from the configured +provider. + ```bash title="Scoping retrieval to specific attachments" curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ -H 'accept: application/json' \ diff --git a/docs/content/docs/documentation/env_vars.md b/docs/content/docs/documentation/env_vars.md index b8c55bbd8..7602bb34c 100644 --- a/docs/content/docs/documentation/env_vars.md +++ b/docs/content/docs/documentation/env_vars.md @@ -289,10 +289,49 @@ These are external services to provide !!! | `API_KEY` | str | _(unset)_ | API key for authenticating with the LLM service | | `LLM_ENABLE_THINKING` | bool | _(unset)_ | Optional chat-template control for models that support `enable_thinking`; leave unset for Mistral tokenizers, set `false` to suppress Qwen-style reasoning traces | | `LLM_SEMAPHORE` | int | 10 | Maximum number of concurrent requests to allow for the LLM service | +| `LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT` | bool | `false` | Honor a client-supplied `base_url`/`api_key` in `metadata.llm_override`. Off by default; read the trade-off below before enabling. | | `MAX_LLM_CONTEXT_SIZE` | `int` | `8192` | Fallback maximum token limit for chat/completion requests. At startup, the `/v1/models` endpoint is queried for the model's `max_model_len`; if that query fails this value is used instead. Requests whose total token count (prompt + `max_tokens`) exceeds the limit are rejected with a **413** error. | | `MAX_OUTPUT_TOKENS` | `int` | `1024` | Default output-token budget (`max_tokens`) applied to chat completions when the request doesn't set one explicitly. | +#### Client-supplied LLM endpoints + +A client can always override the **model name** for a single request via +`metadata.llm_override` (see the [API reference](/openrag/documentation/api/#extra-arguments)). +`LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT=true` additionally honors `base_url` and +`api_key` from that object, so the request is served by a provider of the +client's choosing rather than the configured one. + +It exists for deployments whose clients already send the full object and would +otherwise break. Prefer registering a named endpoint under `/model-endpoints` +and binding it to the partition (`chat_llm`) — same outcome, none of the +trade-off below. + +**What enabling it means.** Any caller who can reach `/v1/chat/completions` can +make the server issue https POSTs from inside your network. The request shape is +pinned, which is what bounds the exposure: + +- `https` only — plaintext internal services are unreachable. +- The path is always `{base_url}/chat/completions` (`{base_url}/completions` for + the legacy `/v1/completions` route); a query string, fragment or `..` segment + — percent-encoded or not — is rejected with a **400**, so the override cannot + be aimed at an arbitrary internal path. +- Redirects are not followed, so a target cannot bounce the server elsewhere. +- The server's own API key is never forwarded; an override without `api_key` + sends no `Authorization` header at all. + +What remains reachable is therefore essentially *other https LLM gateways* — +including an internal one that trusts its network rather than a credential, which +such a caller could then use without holding its key. + +**What it does not change.** It grants no read access a caller does not already +have: `/search` returns the same partition content directly. What changes is the +way data leaves — as outbound LLM traffic from the server's egress rather than as +a user read. That matters against a DLP or approved-subprocessor constraint, not +against a caller who was never authorized in the first place. + +Enable only where every API caller is already trusted with both. + #### VLM Configuration | Variable | Type | Default | Description | diff --git a/infra/compose/.env.example b/infra/compose/.env.example index 92f990722..ac10d2854 100644 --- a/infra/compose/.env.example +++ b/infra/compose/.env.example @@ -14,6 +14,9 @@ BASE_URL= API_KEY= MODEL= LLM_SEMAPHORE=10 +# Allow clients to supply their own LLM endpoint through `metadata.llm_override`. +# LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT=true + # ── VLM (vision model, used for image understanding) ──────────────────────── # Can reuse the LLM values above if that model accepts images. diff --git a/openrag/api/routers/user/chat.py b/openrag/api/routers/user/chat.py index a31bcd328..5dd1c7ab6 100644 --- a/openrag/api/routers/user/chat.py +++ b/openrag/api/routers/user/chat.py @@ -29,6 +29,7 @@ from api.routers.user.source_links import build_document_source_link from api.schemas.user.chat import OpenAIChatCompletionRequest, OpenAICompletionRequest from core.config import load_config +from core.config.endpoints import client_llm_override, custom_endpoint_override_enabled from core.models.preset import resolve_partition_chat_llm from core.utils.exceptions import OpenRAGError from core.utils.logging import get_logger @@ -401,9 +402,18 @@ def _apply_default_max_tokens( consistent with the endpoint that serves the request. An explicit client-supplied value is always honoured. + + Skipped for a client-supplied endpoint: this budget describes the *server's* + endpoint, and the client's provider may reject ``max_tokens`` outright (newer + OpenAI models want ``max_completion_tokens``). Left unset it drops from the + payload; ``validate_tokens_limit`` still falls back to the configured default. """ - if request.max_tokens is None: - request.max_tokens = _effective_max_output_tokens(config, partitions) + if request.max_tokens is not None: + return + llm_override = client_llm_override(getattr(request, "metadata", None)) + if llm_override.get("base_url") and custom_endpoint_override_enabled(): + return + request.max_tokens = _effective_max_output_tokens(config, partitions) def check_tokens_limit( diff --git a/openrag/api/schemas/user/chat.py b/openrag/api/schemas/user/chat.py index 91606388a..3ea48754c 100644 --- a/openrag/api/schemas/user/chat.py +++ b/openrag/api/schemas/user/chat.py @@ -52,8 +52,14 @@ class OpenAIChatCompletionRequest(BaseModel): "llm_override": None, "include_all_retrieved_sources": False, }, - description="Extra custom parameters. Supports 'llm_override' object with an optional 'model' to override the downstream model name. The LLM endpoint and credentials are fixed by server configuration and cannot be overridden by the client. " - "'include_all_retrieved_sources' (default false) adds the full, unfiltered retrieval set to the response's extra.all_retrieved_sources — off by default since it can be large; opt in only for debugging/evaluation.", + description=( + "Extra custom parameters. Supports an 'llm_override' object with an optional 'model' " + "to override the downstream model name; its 'base_url' and 'api_key' are honored only " + "when the deployment sets LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT, and ignored otherwise. " + "'include_all_retrieved_sources' (default false) adds the full, unfiltered retrieval " + "set to the response's extra.all_retrieved_sources — off by default since it can be " + "large; opt in only for debugging/evaluation." + ), ) @model_validator(mode="after") @@ -102,6 +108,12 @@ class OpenAICompletionRequest(BaseModel): "llm_override": None, "include_all_retrieved_sources": False, }, - description="Extra custom parameters. Supports 'llm_override' object with an optional 'model' to override the downstream model name. The LLM endpoint and credentials are fixed by server configuration and cannot be overridden by the client. " - "'include_all_retrieved_sources' (default false) adds the full, unfiltered retrieval set to the response's extra.all_retrieved_sources — off by default since it can be large; opt in only for debugging/evaluation.", + description=( + "Extra custom parameters. Supports an 'llm_override' object with an optional 'model' " + "to override the downstream model name; its 'base_url' and 'api_key' are honored only " + "when the deployment sets LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT, and ignored otherwise. " + "'include_all_retrieved_sources' (default false) adds the full, unfiltered retrieval " + "set to the response's extra.all_retrieved_sources — off by default since it can be " + "large; opt in only for debugging/evaluation." + ), ) diff --git a/openrag/core/config/endpoints.py b/openrag/core/config/endpoints.py index 52da0a7cc..b54260462 100644 --- a/openrag/core/config/endpoints.py +++ b/openrag/core/config/endpoints.py @@ -2,10 +2,42 @@ from __future__ import annotations +import os +from collections.abc import Mapping + from pydantic import Field from .base import ConfigMixin +LLM_OVERRIDE_ENDPOINT_ENV = "LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT" + + +def custom_endpoint_override_enabled() -> bool: + """Is a client-supplied ``llm_override.base_url`` honored at all? + + Off by default: enabling it lets any authenticated caller make the server POST + to an arbitrary host (SSRF). Only the request *shape* is constrained — see + ``VLLMClient._resolve_endpoint_override``. + + Lives here rather than in ``services.inference`` because the API layer reads it + too and ``api -> services`` is a forbidden import direction. Read on demand so + a test or a reloaded worker sees the current environment. + """ + return os.getenv(LLM_OVERRIDE_ENDPOINT_ENV, "false").strip().lower() == "true" + + +def client_llm_override(metadata: object) -> Mapping[str, object]: + """Read ``metadata.llm_override`` as a mapping, or ``{}``. + + Only the *outer* ``metadata`` is schema-validated, so ``{"llm_override": + "gpt-4o"}`` is a valid request whose ``.get(...)`` would raise + ``AttributeError`` — a 500. A non-mapping override counts as absent. + """ + if not isinstance(metadata, Mapping): + return {} + override = metadata.get("llm_override") + return override if isinstance(override, Mapping) else {} + class LLMParamsConfig(ConfigMixin): """Shared parameters for LLM/VLM endpoints.""" diff --git a/openrag/services/inference/_circuit_breaker.py b/openrag/services/inference/_circuit_breaker.py index 79b1afc97..33a890d44 100644 --- a/openrag/services/inference/_circuit_breaker.py +++ b/openrag/services/inference/_circuit_breaker.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from datetime import timedelta from functools import wraps @@ -70,10 +71,26 @@ def get_breaker(name: str, fail_max: int = 50, timeout_duration: float = 60.0) - return _breakers[name] -def with_circuit_breaker(name: str, fail_max: int = 50, timeout_duration: float = 60.0): +def with_circuit_breaker( + name: str, + fail_max: int = 50, + timeout_duration: float = 60.0, + *, + skip_if: Callable[..., bool] | None = None, +): + """Guard *fn* with the shared breaker registered under *name*. + + *skip_if* receives the wrapped call's own arguments; returning True runs *fn* + outside the breaker entirely. For calls that don't reach the endpoint this + breaker describes — folding a second dependency into one health signal makes + it wrong in both directions. + """ + def decorator(fn): @wraps(fn) async def wrapper(*args, **kwargs): + if skip_if is not None and skip_if(*args, **kwargs): + return await fn(*args, **kwargs) breaker = get_breaker(name, fail_max, timeout_duration) try: return await breaker.call_async(fn, *args, **kwargs) diff --git a/openrag/services/inference/vllm_client.py b/openrag/services/inference/vllm_client.py index 10669e78c..f4dc58593 100644 --- a/openrag/services/inference/vllm_client.py +++ b/openrag/services/inference/vllm_client.py @@ -16,8 +16,14 @@ import base64 import re from collections.abc import AsyncIterator, Mapping +from urllib.parse import unquote, urlsplit import httpx +from core.config.endpoints import ( + LLM_OVERRIDE_ENDPOINT_ENV, + client_llm_override, + custom_endpoint_override_enabled, +) from core.embeddings import Embedder, embedder_registry from core.llm import LLM, llm_registry from core.utils.exceptions import ( @@ -140,6 +146,17 @@ def _strip_falsy_logprobs(payload: dict) -> dict: return payload +def _targets_client_endpoint(client: VLLMClient, *_args, **kwargs) -> bool: + """Breaker predicate: is this call routed to a client-supplied endpoint? + + The ``"llm"`` breaker is one process-wide instance describing the *configured* + endpoint, and it counts connection errors and timeouts. Without this any + caller could aim the override at an unresolvable host, repeat until + ``fail_max``, and open the breaker for every tenant. + """ + return client._has_endpoint_override(kwargs) + + @llm_registry.register("vllm") class VLLMClient(LLM): """OpenAI-compatible LLM client backed by vLLM. @@ -173,10 +190,12 @@ def __init__( self._api_key = api_key self._enable_thinking = enable_thinking self._defaults: dict = kwargs - headers: dict[str, str] = {"Content-Type": "application/json"} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + self._allow_custom_endpoint = custom_endpoint_override_enabled() + # Authorization per request, not on the client: httpx merges client-level + # headers into every request with no way to drop one, so the server's key + # would ride along to an overridden endpoint. + self._auth_headers: dict[str, str] = {"Authorization": f"Bearer {api_key}"} if api_key else {} + self._client = httpx.AsyncClient(timeout=timeout, headers={"Content-Type": "application/json"}) # Same construction breadcrumb as VLLMEmbedder: the component factories # cache instances per endpoint name, so this fires once per configured # endpoint and shows which base URL/model a preset name resolved to. @@ -187,32 +206,138 @@ def __init__( enable_thinking=self._enable_thinking, ).debug(f"{type(self).__name__} ready") - def _resolve_overrides(self, kwargs: dict) -> tuple[str, str, dict[str, str] | None]: + def _resolve_overrides(self, kwargs: dict) -> tuple[str, str, dict[str, str], bool]: """Read ``metadata.llm_override`` from *kwargs* without mutating caller data. Pure read: ``kwargs`` is untouched so retries see the original override on every attempt. The caller strips ``metadata`` from the outbound payload via ``_payload_kwargs`` — every ``metadata`` key is OpenRAG-internal and never belongs on the wire. + + Returns ``(base_url, model, headers, overridden)``. ``overridden`` says the + endpoint is the client's; callers need it to suppress the server's sampling + defaults, and reporting it here keeps the override parsed once. """ base_url = self._endpoint model = self._model - override_headers: dict[str, str] | None = None - - # Only `model` may be overridden by the client. `base_url` / `api_key` - # are deliberately NOT read from the request: honoring a client-supplied - # endpoint enables SSRF (the server would issue requests to an arbitrary - # host, e.g. cloud metadata) and would leak the server's API key to that - # host. The endpoint and credentials always come from server config. - llm_override = (kwargs.get("metadata") or {}).get("llm_override") or {} + headers = self._auth_headers + overridden = False + + # `model` is always client-overridable; `base_url` / `api_key` only when + # the operator sets LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT. + llm_override = client_llm_override(kwargs.get("metadata")) if llm_override.get("model"): model = llm_override["model"] - return base_url, model, override_headers + if llm_override.get("base_url"): + if self._allow_custom_endpoint: + base_url, headers = self._resolve_endpoint_override(llm_override) + overridden = True + else: + # Otherwise the failure is opaque: `model` applies while + # `base_url` is dropped, so the request hits the *server's* + # endpoint with a third party's model name and comes back as + # "invalid model name". + logger.bind( + requested_endpoint=llm_override.get("base_url"), + used_endpoint=base_url, + model=model, + ).warning( + f"Ignoring llm_override.base_url — {LLM_OVERRIDE_ENDPOINT_ENV} is not enabled. " + f"The request goes to the configured endpoint with model={model!r}." + ) + elif llm_override: + logger.bind(keys=sorted(llm_override), model=model).warning( + "llm_override carries no base_url; only the model name is overridden" + ) + + return base_url, model, headers, overridden + + def _resolve_endpoint_override(self, llm_override: Mapping) -> tuple[str, dict[str, str]]: + """Honor a client-supplied ``llm_override`` endpoint (opt-in, host-unrestricted). + + Restores the pre-refactor contract (``base_url`` + ``api_key`` + ``model``) + for deployments whose clients rely on it. The *host* is deliberately + unconstrained; the request *shape* is what keeps this from being a read + primitive against internal HTTP APIs: + + * **https only** — the plaintext internal services worth reaching (Milvus, + admin panels) stay unreachable. + * **No client-controlled path** — always ``{base_url}/chat/completions``, + or ``{base_url}/completions`` from ``generate``. A ``#`` or ``?`` would + truncate that suffix once concatenated and a ``..`` would traverse out of + it, making the override a path-picker. Tested on the raw string, not + ``urlsplit`` parts: a trailing ``#`` parses as an empty fragment and only + bites after concatenation. + * The server's own API key is never forwarded — the override's key, or no + ``Authorization`` at all. - def _chat_payload_kwargs(self, kwargs: dict) -> dict: - payload_kwargs = {**self._defaults, **kwargs} - enable_thinking = payload_kwargs.pop("enable_thinking", self._enable_thinking) + Rejections are 4xx so ``with_retry`` (429/502/503/504) does not re-attempt. + """ + candidate = str(llm_override["base_url"]).strip() + + if "#" in candidate or "?" in candidate: + raise InferenceError( + "llm_override.base_url must not carry a query string or fragment", + code="LLM_OVERRIDE_REJECTED", + status_code=400, + ) + try: + parts = urlsplit(candidate) + except ValueError as exc: + # Malformed input, e.g. the unclosed IPv6 literal `https://[::1/v1`. + # Uncaught it escapes the httpx handlers and error_handlers.py's + # OpenRAGError mapping, surfacing as an unstructured 500. + raise InferenceError( + "llm_override.base_url is not a valid URL", + code="LLM_OVERRIDE_REJECTED", + status_code=400, + ) from exc + scheme = parts.scheme.lower() + if scheme != "https": + raise InferenceError( + f"llm_override.base_url scheme {scheme!r} is not allowed (https only)", + code="LLM_OVERRIDE_REJECTED", + status_code=400, + ) + # Decoded too: `urlsplit` leaves `%XX` alone and httpx forwards it, so + # `%2e%2e` arrives as `..` at a target that decodes before routing. + if ".." in parts.path.split("/") or ".." in unquote(parts.path).split("/"): + raise InferenceError( + "llm_override.base_url must not contain a '..' path segment", + code="LLM_OVERRIDE_REJECTED", + status_code=400, + ) + + candidate = candidate.rstrip("/") + api_key = llm_override.get("api_key") + headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} + logger.bind(endpoint=candidate, configured=self._endpoint).debug( + "Honoring client-supplied llm_override endpoint" + ) + return candidate, headers + + def _has_endpoint_override(self, kwargs: dict) -> bool: + """Is this request routed to a client-supplied endpoint? + + Reads raw ``kwargs`` because the breaker predicate runs before the method + body. Inside the body, use the flag from ``_resolve_overrides``. + """ + if not self._allow_custom_endpoint: + return False + return bool(client_llm_override(kwargs.get("metadata")).get("base_url")) + + def _chat_payload_kwargs(self, kwargs: dict, *, use_defaults: bool = True) -> dict: + """Merge server sampling defaults under the request's own params. + + ``use_defaults=False`` for a client-supplied endpoint: the defaults describe + the *server's* model and another provider may reject them outright (Gemini + 400s on an unsolicited ``logprobs``). Dropping them is what made the restored + override actually work. + """ + payload_kwargs = {**self._defaults, **kwargs} if use_defaults else dict(kwargs) + fallback_thinking = self._enable_thinking if use_defaults else None + enable_thinking = payload_kwargs.pop("enable_thinking", fallback_thinking) if enable_thinking is not None and enable_thinking is True: chat_template_kwargs = dict(payload_kwargs.get("chat_template_kwargs") or {}) chat_template_kwargs.setdefault("enable_thinking", enable_thinking) @@ -220,12 +345,12 @@ def _chat_payload_kwargs(self, kwargs: dict) -> dict: payload_kwargs = _strip_falsy_logprobs(payload_kwargs) return payload_kwargs - @with_circuit_breaker("llm") + @with_circuit_breaker("llm", skip_if=_targets_client_endpoint) @with_retry(max_attempts=3) async def generate(self, prompt: str, **kwargs) -> dict: - base_url, model, headers = self._resolve_overrides(kwargs) + base_url, model, headers, overridden = self._resolve_overrides(kwargs) kwargs.pop("metadata", None) - payload = {**self._defaults, **kwargs, "model": model, "prompt": prompt} + payload = {**({} if overridden else self._defaults), **kwargs, "model": model, "prompt": prompt} payload = _strip_falsy_logprobs(payload) log_llm_call(caller="VLLMClient.generate", model=model, endpoint=base_url, prompt=prompt) try: @@ -242,12 +367,17 @@ async def generate(self, prompt: str, **kwargs) -> dict: ) from exc return _parse_response(resp) - @with_circuit_breaker("llm") + @with_circuit_breaker("llm", skip_if=_targets_client_endpoint) @with_retry(max_attempts=3) async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: - base_url, model, headers = self._resolve_overrides(kwargs) + base_url, model, headers, overridden = self._resolve_overrides(kwargs) kwargs.pop("metadata", None) - payload = {**self._chat_payload_kwargs(kwargs), "model": model, "messages": messages, "stream": False} + payload = { + **self._chat_payload_kwargs(kwargs, use_defaults=not overridden), + "model": model, + "messages": messages, + "stream": False, + } log_llm_call(caller="VLLMClient.chat", model=model, endpoint=base_url, messages=messages) try: resp = await self._client.post(f"{base_url}/chat/completions", json=payload, headers=headers) @@ -264,9 +394,14 @@ async def chat(self, messages: list[dict[str, str]], **kwargs) -> dict: return _parse_response(resp) async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: - base_url, model, headers = self._resolve_overrides(kwargs) + base_url, model, headers, overridden = self._resolve_overrides(kwargs) kwargs.pop("metadata", None) - payload = {**self._chat_payload_kwargs(kwargs), "model": model, "messages": messages, "stream": True} + payload = { + **self._chat_payload_kwargs(kwargs, use_defaults=not overridden), + "model": model, + "messages": messages, + "stream": True, + } log_llm_call(caller="VLLMClient.stream_chat", model=model, endpoint=base_url, messages=messages, stream=True) try: async with self._client.stream( @@ -541,6 +676,9 @@ async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> resp = await self._client.post( f"{self._endpoint}/chat/completions", json=payload, + # Explicit since Authorization moved off the shared httpx client; + # captioning always uses the configured endpoint. + headers=self._auth_headers, ) resp.raise_for_status() except httpx.ConnectError as exc: diff --git a/tests/unit/services/inference/test_vllm_client.py b/tests/unit/services/inference/test_vllm_client.py index 26e6d39d2..c021a6636 100644 --- a/tests/unit/services/inference/test_vllm_client.py +++ b/tests/unit/services/inference/test_vllm_client.py @@ -15,6 +15,7 @@ InferenceTimeoutError, ) from services.inference._circuit_breaker import _breakers +from services.inference._retry import _is_retryable from services.inference.vllm_client import ( _SUSPECT_UNICODE_ESCAPE, VLLMClient, @@ -353,10 +354,10 @@ def _make_client(self): def test_no_override_uses_defaults(self): client = self._make_client() kwargs: dict = {} - base_url, model, headers = client._resolve_overrides(kwargs) + base_url, model, headers, _ = client._resolve_overrides(kwargs) assert base_url == "http://default:8000/v1" assert model == "default-model" - assert headers is None + assert headers == {"Authorization": "Bearer default-key"} def test_llm_override_model_applied_endpoint_and_key_ignored(self): client = self._make_client() @@ -368,11 +369,11 @@ def test_llm_override_model_applied_endpoint_and_key_ignored(self): }, } kwargs: dict = {"metadata": original_metadata} - base_url, model, headers = client._resolve_overrides(kwargs) + base_url, model, headers, _ = client._resolve_overrides(kwargs) # Only `model` is honored; endpoint and credentials stay server-side. assert model == "custom-model" assert base_url == "http://default:8000/v1" - assert headers is None + assert headers == {"Authorization": "Bearer default-key"} # kwargs must not be mutated — retries depend on llm_override surviving. assert kwargs["metadata"] is original_metadata assert "llm_override" in kwargs["metadata"] @@ -384,10 +385,10 @@ def test_llm_override_partial(self): "use_map_reduce": True, } kwargs: dict = {"metadata": original_metadata} - base_url, model, headers = client._resolve_overrides(kwargs) + base_url, model, headers, _ = client._resolve_overrides(kwargs) assert base_url == "http://default:8000/v1" assert model == "override-model" - assert headers is None + assert headers == {"Authorization": "Bearer default-key"} assert kwargs["metadata"] is original_metadata assert kwargs["metadata"] == { "llm_override": {"model": "override-model"}, @@ -407,10 +408,277 @@ def test_client_base_url_and_api_key_override_ignored(self): } } } - base_url, model, headers = client._resolve_overrides(kwargs) + base_url, model, headers, _ = client._resolve_overrides(kwargs) assert model == "custom-model" assert base_url == "http://default:8000/v1" - assert headers is None + assert headers == {"Authorization": "Bearer default-key"} + + +class TestCustomEndpointOverride: + """LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT restores the full llm_override contract. + + The *host* is unrestricted; the request shape is not (https, fixed path), and + the server's own credential must never leak. + """ + + def _make_client(self, monkeypatch, enabled: bool, **kwargs): + if enabled: + monkeypatch.setenv("LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT", "true") + else: + monkeypatch.delenv("LLM_OVERRIDE_ALLOW_CUSTOM_ENDPOINT", raising=False) + return VLLMClient( + endpoint="http://default:8000/v1", + model_name="default-model", + api_key="default-key", + **kwargs, + ) + + def _kwargs(self, **override): + return {"metadata": {"llm_override": override}} + + @staticmethod + def _refusing_transport() -> httpx.AsyncClient: + def refuse(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused") + + return httpx.AsyncClient(transport=_make_transport(refuse)) + + @pytest.mark.asyncio + async def test_client_endpoint_failure_does_not_touch_the_shared_breaker(self, monkeypatch): + """A failing client endpoint must not count against the shared "llm" breaker. + + It is one global instance guarding the *server's* endpoint, so otherwise any + caller could aim the override at an unresolvable host and open it for every + tenant. + """ + client = self._make_client(monkeypatch, enabled=True) + client._client = self._refusing_transport() + + with pytest.raises(InferenceConnectionError): + await client.chat( + [{"role": "user", "content": "hi"}], + **self._kwargs(base_url="https://third-party.tld/v1", api_key="k"), + ) + + assert "llm" not in _breakers + + @pytest.mark.asyncio + async def test_configured_endpoint_failure_still_counts_against_the_shared_breaker(self, monkeypatch): + """Skipping the breaker must be scoped to overridden endpoints, not a + blanket disable of the server's own.""" + client = self._make_client(monkeypatch, enabled=True) + client._client = self._refusing_transport() + + with pytest.raises(InferenceConnectionError): + await client.chat([{"role": "user", "content": "hi"}]) + + assert _breakers["llm"].fail_counter == 1 + + def test_arbitrary_endpoint_is_honored_with_client_key(self, monkeypatch): + client = self._make_client(monkeypatch, enabled=True) + base_url, model, headers, overridden = client._resolve_overrides( + self._kwargs(base_url="https://api.openai.com/v1/", api_key="sk-client", model="gpt-5.1") + ) + + assert base_url == "https://api.openai.com/v1" + assert model == "gpt-5.1" + assert headers == {"Authorization": "Bearer sk-client"} + assert overridden is True + + @pytest.mark.parametrize( + "url", + [ + "https://169.254.169.254/latest/meta-data", + "https://internal-service.corp/v1", + "https://api.openai.com@evil.tld/v1", + ], + ) + def test_no_host_allowlist_when_enabled(self, monkeypatch, url): + """Explicit: with the flag on there is no host allowlist. That residual SSRF + (host, not path) is what the flag trades away — asserted so it cannot + regress into a false sense of safety. + """ + client = self._make_client(monkeypatch, enabled=True) + base_url, _, _, _ = client._resolve_overrides(self._kwargs(base_url=url, api_key="k", model="m")) + + assert base_url == url.rstrip("/") + + def test_server_api_key_never_reaches_an_overridden_endpoint(self, monkeypatch): + """An override with no api_key sends no Authorization at all rather than + falling back to the server's credential.""" + client = self._make_client(monkeypatch, enabled=True) + _, _, headers, _ = client._resolve_overrides(self._kwargs(base_url="https://evil.tld/v1", model="m")) + + assert headers == {} + + @pytest.mark.parametrize("base_url", ["http://milvus:19530/v2", "file:///etc/passwd"]) + def test_non_https_scheme_is_rejected_without_retry(self, monkeypatch, base_url): + """https only: plaintext http is where the internal-infra SSRF payoff lives. + Failing here turns it into a 4xx `with_retry` will not re-attempt.""" + client = self._make_client(monkeypatch, enabled=True) + + with pytest.raises(InferenceError) as exc_info: + client._resolve_overrides(self._kwargs(base_url=base_url, api_key="k")) + + assert exc_info.value.status_code == 400 + assert not _is_retryable(exc_info.value) + + @pytest.mark.parametrize( + "base_url", + [ + "https://milvus:19530/v2/vectordb/collections/list#", + "https://milvus:19530/v2/vectordb/collections/list?x=1", + "https://milvus:19530/v2/vectordb/../collections", + # Percent-encoded: httpx sends the octets verbatim, so a target that + # decodes before routing resolves these to the `..` form above. + "https://milvus:19530/v2/vectordb/%2e%2e/collections", + "https://milvus:19530/v2/vectordb/%2E%2E/collections", + ], + ) + def test_client_controlled_path_is_rejected_without_retry(self, monkeypatch, base_url): + """The request is always issued as ``{base_url}/chat/completions``, so a + fragment or query truncates that suffix and a ``..`` traverses out of it — + either aims the override at an arbitrary internal path (here Milvus' REST + API, reading cross-partition vectors). + """ + client = self._make_client(monkeypatch, enabled=True) + + with pytest.raises(InferenceError) as exc_info: + client._resolve_overrides(self._kwargs(base_url=base_url, api_key="k")) + + assert exc_info.value.status_code == 400 + assert not _is_retryable(exc_info.value) + + def test_unparseable_base_url_is_rejected_without_retry(self, monkeypatch): + """``urlsplit`` raises on an unclosed IPv6 literal; uncaught, that + ``ValueError`` surfaces as an unstructured 500.""" + client = self._make_client(monkeypatch, enabled=True) + + with pytest.raises(InferenceError) as exc_info: + client._resolve_overrides(self._kwargs(base_url="https://[::1/v1", api_key="k")) + + assert exc_info.value.status_code == 400 + assert not _is_retryable(exc_info.value) + + @pytest.mark.parametrize("llm_override", ["gpt-5.1", ["gpt-5.1"], 7]) + def test_non_mapping_llm_override_is_ignored(self, monkeypatch, llm_override): + """Only the outer ``metadata`` is schema-validated, so this is a well-formed + request whose ``.get`` would raise ``AttributeError`` — a 500. The breaker + predicate reads it first, before the method body runs. + """ + client = self._make_client(monkeypatch, enabled=True) + kwargs = {"metadata": {"llm_override": llm_override}} + + assert client._has_endpoint_override(kwargs) is False + base_url, model, headers, overridden = client._resolve_overrides(kwargs) + + assert (base_url, model, overridden) == ("http://default:8000/v1", "default-model", False) + assert headers == {"Authorization": "Bearer default-key"} + + def test_disabled_by_default_still_ignores_base_url(self, monkeypatch): + """Unset env keeps the hardened post-refactor behaviour, so an upgrade never + turns an untouched deployment into an SSRF surface.""" + client = self._make_client(monkeypatch, enabled=False) + base_url, model, headers, overridden = client._resolve_overrides( + self._kwargs(base_url="https://api.openai.com/v1", api_key="sk-client", model="gpt-5.1") + ) + + assert base_url == "http://default:8000/v1" + assert model == "gpt-5.1" + assert headers == {"Authorization": "Bearer default-key"} + assert overridden is False + + def test_override_does_not_mutate_caller_metadata(self, monkeypatch): + """Retries re-read the same kwargs, so the override must survive intact.""" + client = self._make_client(monkeypatch, enabled=True) + kwargs = self._kwargs(base_url="https://api.openai.com/v1", api_key="sk-client", model="gpt-5.1") + original = kwargs["metadata"]["llm_override"].copy() + + client._resolve_overrides(kwargs) + + assert kwargs["metadata"]["llm_override"] == original + + @pytest.mark.asyncio + async def test_override_reaches_the_wire_on_chat(self, monkeypatch): + """End-to-end: the request is issued to the overridden host with the + client's key, not the server's.""" + client = self._make_client(monkeypatch, enabled=True) + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["url"] = str(request.url) + seen["auth"] = request.headers.get("Authorization") + return _chat_response("ok") + + client._client = httpx.AsyncClient(transport=_make_transport(handler)) + await client.chat( + [{"role": "user", "content": "hi"}], + metadata={ + "llm_override": { + "base_url": "https://api.openai.com/v1", + "api_key": "sk-client", + "model": "gpt-5.1", + } + }, + ) + + assert seen["url"] == "https://api.openai.com/v1/chat/completions" + assert seen["auth"] == "Bearer sk-client" + + @pytest.mark.asyncio + async def test_server_sampling_defaults_are_not_sent_to_a_client_endpoint(self, monkeypatch): + """The server's sampling defaults must not ride along to another provider, + which may reject them outright — this is what made an otherwise-restored + override still fail in production.""" + client = self._make_client(monkeypatch, enabled=True, temperature=0.3) + captured: dict = {} + + def capture(request: httpx.Request) -> httpx.Response: + captured.update(json.loads(request.content)) + return _chat_response("ok") + + client._client = httpx.AsyncClient(transport=_make_transport(capture)) + await client.chat( + [{"role": "user", "content": "hi"}], + metadata={"llm_override": {"base_url": "https://third-party.tld/v1", "api_key": "k", "model": "m"}}, + ) + + assert "temperature" not in captured + assert captured["model"] == "m" + + @pytest.mark.asyncio + async def test_keyless_override_sends_no_authorization_on_the_wire(self, monkeypatch): + """httpx merges client-level headers into every request, so this only holds + because Authorization is set per request.""" + client = self._make_client(monkeypatch, enabled=True) + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["auth"] = request.headers.get("Authorization") + return _chat_response("ok") + + client._client = httpx.AsyncClient(transport=_make_transport(handler)) + await client.chat( + [{"role": "user", "content": "hi"}], + metadata={"llm_override": {"base_url": "https://no-auth.internal/v1", "model": "m"}}, + ) + + assert seen["auth"] is None + + @pytest.mark.asyncio + async def test_server_key_is_sent_when_there_is_no_override(self, monkeypatch): + """Moving Authorization off the client must not silently drop it.""" + client = self._make_client(monkeypatch, enabled=True) + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["auth"] = request.headers.get("Authorization") + return _chat_response("ok") + + client._client = httpx.AsyncClient(transport=_make_transport(handler)) + await client.chat([{"role": "user", "content": "hi"}]) + + assert seen["auth"] == "Bearer default-key" # --------------------------------------------------------------------------- @@ -621,6 +889,23 @@ def handler(request: httpx.Request) -> httpx.Response: await self._make_vision(handler).caption_image(b"\x89PNG\r\n\x1a\n") + @pytest.mark.asyncio + async def test_caption_image_sends_the_configured_api_key(self): + """VLLMClient sets Authorization per request, not on the shared httpx + client, so an overridden llm endpoint never receives the server's key. + Captioning inherits that client but always calls the configured endpoint, + so it must pass the header explicitly or authenticate as nobody. + """ + seen: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + seen["auth"] = request.headers.get("Authorization") + return _chat_response("ok") + + await self._make_vision(handler).caption_image(b"img") + + assert seen["auth"] == "Bearer test-key" + @pytest.mark.asyncio async def test_caption_images_batch(self): call_count = 0 diff --git a/tests/unit/test_token_validation.py b/tests/unit/test_token_validation.py index 783225654..9fac27e12 100644 --- a/tests/unit/test_token_validation.py +++ b/tests/unit/test_token_validation.py @@ -649,3 +649,18 @@ def test_explicit_client_value_is_honoured(self): req = OpenAIChatCompletionRequest(messages=[{"role": "user", "content": "hello"}], max_tokens=7) chat._apply_default_max_tokens(req, s, None) assert req.max_tokens == 7 + + @pytest.mark.parametrize("llm_override", ["gpt-5.1", ["gpt-5.1"], 7]) + def test_non_mapping_llm_override_falls_back_to_the_server_default(self, llm_override): + """``metadata`` is validated only as ``dict[str, Any]``, so this reaches the + read; ``.get`` on it would raise ``AttributeError`` — a 500.""" + import api.routers.user.chat as chat + from core.config.model_endpoints import LLM_OUTPUT_TOKENS_KEY + + s = _settings_with_default_llm(**{LLM_OUTPUT_TOKENS_KEY: 1024}) + req = OpenAIChatCompletionRequest( + messages=[{"role": "user", "content": "hello"}], + metadata={"llm_override": llm_override}, + ) + chat._apply_default_max_tokens(req, s, None) + assert req.max_tokens == 1024