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

@@ -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,

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()

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

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)