feat(nvidia): add NIM billing origin header

This commit is contained in:
kchantharuan 2026-05-13 12:46:07 -07:00 committed by Teknium
parent 4e89c53082
commit 13c3d4b4ef
5 changed files with 162 additions and 6 deletions

View file

@ -1664,6 +1664,9 @@ class AIAgent:
if base_url_host_matches(effective_base, "openrouter.ai"):
from agent.auxiliary_client import build_or_headers
client_kwargs["default_headers"] = build_or_headers()
elif base_url_host_matches(effective_base, "integrate.api.nvidia.com"):
from agent.auxiliary_client import build_nvidia_nim_headers
client_kwargs["default_headers"] = build_nvidia_nim_headers(effective_base)
elif base_url_host_matches(effective_base, "api.routermint.com"):
client_kwargs["default_headers"] = _routermint_headers()
elif base_url_host_matches(effective_base, "api.githubcopilot.com"):
@ -1702,9 +1705,15 @@ class AIAgent:
}
if _provider_timeout is not None:
client_kwargs["timeout"] = _provider_timeout
# Preserve any default_headers the router set
if hasattr(_routed_client, '_default_headers') and _routed_client._default_headers:
client_kwargs["default_headers"] = dict(_routed_client._default_headers)
# Preserve provider-specific headers the router set. The
# OpenAI SDK stores caller-provided default_headers in
# _custom_headers; older/mocked clients may expose
# _default_headers instead.
_routed_headers = getattr(_routed_client, "_custom_headers", None)
if not _routed_headers:
_routed_headers = getattr(_routed_client, "_default_headers", None)
if _routed_headers:
client_kwargs["default_headers"] = dict(_routed_headers)
else:
# When the user explicitly chose a non-OpenRouter provider
# but no credentials were found, fail fast with a clear
@ -1753,8 +1762,11 @@ class AIAgent:
}
if _provider_timeout is not None:
client_kwargs["timeout"] = _provider_timeout
if hasattr(_fb_client, "_default_headers") and _fb_client._default_headers:
client_kwargs["default_headers"] = dict(_fb_client._default_headers)
_fb_headers = getattr(_fb_client, "_custom_headers", None)
if not _fb_headers:
_fb_headers = getattr(_fb_client, "_default_headers", None)
if _fb_headers:
client_kwargs["default_headers"] = dict(_fb_headers)
_fb_resolved = True
break
if not _fb_resolved:
@ -7334,12 +7346,18 @@ class AIAgent:
return True
def _apply_client_headers_for_base_url(self, base_url: str) -> None:
from agent.auxiliary_client import _AI_GATEWAY_HEADERS, build_or_headers
from agent.auxiliary_client import (
_AI_GATEWAY_HEADERS,
build_nvidia_nim_headers,
build_or_headers,
)
if base_url_host_matches(base_url, "openrouter.ai"):
self._client_kwargs["default_headers"] = build_or_headers()
elif base_url_host_matches(base_url, "ai-gateway.vercel.sh"):
self._client_kwargs["default_headers"] = dict(_AI_GATEWAY_HEADERS)
elif base_url_host_matches(base_url, "integrate.api.nvidia.com"):
self._client_kwargs["default_headers"] = build_nvidia_nim_headers(base_url)
elif base_url_host_matches(base_url, "api.routermint.com"):
self._client_kwargs["default_headers"] = _routermint_headers()
elif base_url_host_matches(base_url, "api.githubcopilot.com"):