From 89fe81a799fd3ee636966e4b265257a2ce1ca89f Mon Sep 17 00:00:00 2001 From: Saruar Date: Sun, 1 Nov 2026 18:20:22 +0500 Subject: [PATCH] feat(nace): send X-Nace-Client attribution header Nace requests, including the Auto Router classifier, identify litellm and its version in X-Nace-Client for Nace's usage attribution. Other providers do not receive it --- litellm/llms/nace/decisions/transformation.py | 24 +++++++++++++ .../complexity_router/jev_classifier.py | 7 ++-- tests/unit/decisions/test_main.py | 34 +++++++++++++++++++ .../complexity_router/test_jev_classifier.py | 34 +++++++++++++++++++ 4 files changed, 97 insertions(+), 2 deletions(-) diff --git a/litellm/llms/nace/decisions/transformation.py b/litellm/llms/nace/decisions/transformation.py index 058f35731f6..0e5e0b24b96 100644 --- a/litellm/llms/nace/decisions/transformation.py +++ b/litellm/llms/nace/decisions/transformation.py @@ -1,7 +1,9 @@ +import re from collections.abc import Mapping from dataclasses import replace from typing import Final +from litellm._version import version as _litellm_version from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig from litellm.types.decisions import ( DecisionsBase64Image, @@ -11,6 +13,15 @@ from litellm.types.decisions import ( UnsupportedDecisionsRequest, ) +NACE_CLIENT_HEADER: Final = "X-Nace-Client" +_NACE_CLIENT_DISALLOWED: Final = re.compile(r"[^A-Za-z0-9 ._@+-]") +_NACE_CLIENT_MAX_LENGTH: Final = 64 + + +def nace_client_header() -> Mapping[str, str]: + version: Final = _NACE_CLIENT_DISALLOWED.sub("-", _litellm_version) + return {NACE_CLIENT_HEADER: f"litellm/{version}"[:_NACE_CLIENT_MAX_LENGTH]} + def _data_url(image: DecisionsImage) -> DecisionsImage: if isinstance(image, DecisionsBase64Image): @@ -35,6 +46,19 @@ class NaceDecisionsConfig(BaseDecisionsConfig): def get_default_api_base(self) -> str | None: return "https://console.nace.ai" + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + api_key: str | None, + litellm_params: Mapping[str, object], + ) -> Mapping[str, str]: + base: Final = super().validate_environment(headers, model, api_key, litellm_params) + return { + **{name: value for name, value in base.items() if name.lower() != NACE_CLIENT_HEADER.lower()}, + **nace_client_header(), + } + def transform_decisions_request( self, model: str, diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index 39877150cb5..623461dc55c 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -22,7 +22,7 @@ from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG, is_ai_decide_model from litellm.llms.laya.common_utils import laya_response_model -from litellm.llms.nace.decisions.transformation import NACE_DECISIONS_CONFIG +from litellm.llms.nace.decisions.transformation import NACE_DECISIONS_CONFIG, nace_client_header from litellm.llms.pass_through.typesafe_logging_handler import TypeSafePassthroughLoggingHandler from litellm.router_deprecated import renamed_router_param from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS as _DEFAULT_JEV_INSTRUCTIONS @@ -316,10 +316,13 @@ class HttpJevClassifierClient: authorization: Final[Mapping[str, str]] = ( MappingProxyType({"Authorization": f"Bearer {self._api_key}"}) if self._api_key else MappingProxyType({}) ) + attribution: Final[Mapping[str, str]] = ( + MappingProxyType(nace_client_header()) if self._provider == "nace" else MappingProxyType({}) + ) response: Final = await self._http_client.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler has a dynamic post signature self._request_url(request.model), json=request.model_dump(mode="json", exclude=self._excluded_fields(request.model)), - headers=MappingProxyType({**authorization, "Content-Type": "application/json"}), # pyright: ignore[reportArgumentType] # HTTP headers are not mutated by AsyncHTTPHandler + headers=MappingProxyType({**authorization, **attribution, "Content-Type": "application/json"}), # pyright: ignore[reportArgumentType] # HTTP headers are not mutated by AsyncHTTPHandler timeout=timeout_s, ) response.raise_for_status() diff --git a/tests/unit/decisions/test_main.py b/tests/unit/decisions/test_main.py index f509de0f93c..ef8ae383bc5 100644 --- a/tests/unit/decisions/test_main.py +++ b/tests/unit/decisions/test_main.py @@ -21,6 +21,7 @@ from botocore.credentials import Credentials from respx.models import Call import litellm +from litellm._version import version as litellm_version from litellm.constants import AWS_SIGNING_MAX_THREADS from litellm.cost_calculator import get_response_cost_from_hidden_params from litellm.integrations.custom_logger import CustomLogger @@ -4814,3 +4815,36 @@ def test_nace_refuses_a_ninth_image_before_http(respx_mock: respx.MockRouter) -> assert caught.value.status_code == 400 assert not route.called + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("nace_account") +async def test_nace_requests_identify_the_litellm_version(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post(_NACE_SYSTEMONE_URL).respond(json=_DREX_MAX_RESPONSE) + + await litellm.adecisions( + model="nace/drex-max", + input="Order 1234.", + questions=_PREDICATE_QUESTIONS, + extra_headers={"X-Nace-Client": "caller/9", "X-Other": "kept"}, + ) + + sent: Final = route.calls[0].request.headers + assert sent["x-nace-client"] == f"litellm/{litellm_version}" + assert re.fullmatch(r"litellm/[A-Za-z0-9 ._@+-]+", sent["x-nace-client"]) + assert len(sent["x-nace-client"]) <= 64 + assert sent["x-other"] == "kept" + + +@pytest.mark.asyncio +async def test_other_decisions_providers_do_not_get_the_nace_client_header( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("OPENROUTER_API_KEY", "or-key") + route: Final = respx_mock.post("https://openrouter.ai/api/alpha/decisions").respond(json=_RESPONSE) + + await litellm.adecisions( + model="openrouter/nace/drex-test-alias", input="Order 1234.", questions=_PREDICATE_QUESTIONS + ) + + assert "x-nace-client" not in route.calls[0].request.headers diff --git a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py index 17ed41ee8bc..62caa376f11 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -13,6 +13,7 @@ import tiktoken import litellm from litellm._logging import verbose_router_logger +from litellm._version import version as litellm_version from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -23,6 +24,7 @@ from litellm.router_strategy.complexity_router.complexity_router import Complexi from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, JevClassifierConfig from litellm.router_strategy.complexity_router.jev_classifier import ( DEFAULT_JEV_INSTRUCTIONS, + ClassifierProvider, HttpJevClassifierClient, JevChoiceAnswer, JevSystemOneRequest, @@ -1488,3 +1490,35 @@ async def test_nace_routes_through_its_console_with_its_own_key_and_accounts_the assert json.loads(sent.content)["model"] == "drex-max" assert len(recorder.calls) == 1 assert recorder.calls[0]["response_cost"] == pytest.approx(0.31) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("provider", ["nace", "typesafe"]) +async def test_only_the_nace_classifier_sends_the_client_header(provider: ClassifierProvider) -> None: + seen: Final[list[httpx.Headers]] = [] + + def respond(http_request: httpx.Request) -> httpx.Response: + seen.append(http_request.headers) + return httpx.Response( + 200, + json={ + "model": "drex-max", + "answers": { + "tier": { + "type": "choice", + "choice": "SIMPLE", + "probabilities": {"SIMPLE": 1.0}, + "confidence": 1.0, + } + }, + }, + ) + + handler: Final = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(respond)) + client: Final = HttpJevClassifierClient("secret", "https://classifier.test", handler, provider) + await client.evaluate( + build_jev_request("Hello", None, "drex-max", DEFAULT_JEV_INSTRUCTIONS, {"SIMPLE": "facts"}), 1.0 + ) + + assert seen[0].get("x-nace-client") == (f"litellm/{litellm_version}" if provider == "nace" else None)