feat(nvidia): add NIM billing origin header — port to extracted modules

Original commit 13c3d4b4e by kchantharuan touched __init__ and
_apply_client_headers_for_base_url in pre-refactor run_agent.py. Re-applied to:

  - __init__: agent/agent_init.py (3 hunks — NVIDIA branch + _custom_headers
    fallback in routed-client and fallback-client paths)
  - _apply_client_headers_for_base_url: still in run_agent.py (1 hunk)

build_nvidia_nim_headers was already present in agent/auxiliary_client.py
from the prior merge — no additional port needed.

Co-authored-by: kchantharuan <kchantharuan@nvidia.com>
This commit is contained in:
teknium1
2026-05-16 23:25:11 -07:00
parent b07524e53a
commit 27df249564
2 changed files with 24 additions and 6 deletions
+17 -5
View File
@@ -617,6 +617,9 @@ def init_agent(
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"] = _ra()._routermint_headers()
elif base_url_host_matches(effective_base, "api.githubcopilot.com"):
@@ -655,9 +658,15 @@ def init_agent(
}
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
@@ -706,8 +715,11 @@ def init_agent(
}
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:
+7 -1
View File
@@ -2644,12 +2644,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"):