From 4546033b8e81d335c88e220e93df2258fac35fde Mon Sep 17 00:00:00 2001 From: quevedoSteven Date: Fri, 31 Jul 2026 15:51:30 +0000 Subject: [PATCH] feat: round-robin NVIDIA NIM API key rotation Adds cai.util.nim_rotation which cycles NVIDIA_NIM_API_KEY_1..N per request, keeping under NIM per-key rate limits (~40 req/min). Rotation applies in the OpenAI-compatible LLM path and the direct httpx completion path when the API base is a NIM endpoint (api.nvidia.com). Also defaults cache_write_tokens to 0 in the streamed usage object: litellm CustomResponseUsage now requires the field while NIM usage payloads omit it, which previously aborted streaming at end of response. --- .env.example | 5 ++ .../agents/models/openai_chatcompletions.py | 34 ++++++++++++- src/cai/util/nim_rotation.py | 49 +++++++++++++++++++ tests/util/test_nim_rotation.py | 45 +++++++++++++++++ 4 files changed, 132 insertions(+), 1 deletion(-) create mode 100644 src/cai/util/nim_rotation.py create mode 100644 tests/util/test_nim_rotation.py diff --git a/.env.example b/.env.example index e36532cb..39b1a4b2 100644 --- a/.env.example +++ b/.env.example @@ -4,6 +4,11 @@ OLLAMA="" PROMPT_TOOLKIT_NO_CPR=1 CAI_STREAM=false CAI_MODEL="alias1" +# NVIDIA NIM: set N keys as NVIDIA_NIM_API_KEY_1..N to round-robin through +# them per request (NIM enforces ~40 req/min per key). Requires +# OPENAI_API_BASE pointing at a NIM endpoint (e.g. https://integrate.api.nvidia.com/v1). +# NVIDIA_NIM_API_KEY_1="" +# NVIDIA_NIM_API_KEY_2="" # Model sampling parameters (optional - defaults shown) # CAI_TEMPERATURE=0.7 # CAI_TOP_P=1.0 diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 47d871b0..6324fd6e 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -153,6 +153,10 @@ from cai.util.llm_api_base import ( resolve_llm_openai_compatible_base, resolve_llm_openai_compatible_api_key, ) +from cai.util.nim_rotation import ( + get_next_nim_key, + is_nim_rotation_configured, +) from cai.errors import LLMEmptyAssistantError, LLMRateLimited, LLMTimeout from cai.util.gateway_rate_limiter import ( COMPLETION_BUDGET_TOKENS, @@ -2786,6 +2790,9 @@ class OpenAIChatCompletionsModel(Model): and hasattr(usage.prompt_tokens_details, "cached_tokens") and usage.prompt_tokens_details.cached_tokens else 0, + "cache_write_tokens": cache_creation + if cache_creation is not None + else 0, }, cache_creation_input_tokens=cache_creation, cache_read_input_tokens=cache_read, @@ -3657,6 +3664,21 @@ class OpenAIChatCompletionsModel(Model): if hasattr(model_settings, "reasoning_effort"): kwargs["reasoning_effort"] = model_settings.reasoning_effort + # NIM key rotation: when multiple NVIDIA_NIM_API_KEY_N are set, + # round-robin through them to stay under 40 req/min per key. + # Only applies when no explicit api_base was routed above (alias, + # ollama_cloud, or a custom provider fallback): those must keep + # their own endpoint + key. + if is_nim_rotation_configured() and not kwargs.get("api_base"): + kwargs["api_key"] = get_next_nim_key() + kwargs["custom_llm_provider"] = "openai" + kwargs["api_base"] = ( + resolve_llm_openai_compatible_base( + str(kwargs.get("model") or os.getenv("CAI_MODEL") or "") + ).rstrip("/") + or os.getenv("OPENAI_API_BASE", "").rstrip("/") + ) + # Filter out NotGiven values to avoid JSON serialization issues filtered_kwargs = {} for key, value in kwargs.items(): @@ -3763,9 +3785,19 @@ class OpenAIChatCompletionsModel(Model): request_body = {k: v for k, v in request_body.items() if v is not None} api_url = f"{openai_api_base.rstrip('/')}/chat/completions" + if "api.nvidia.com" in openai_api_base.lower() and is_nim_rotation_configured(): + direct_api_key = get_next_nim_key() or "sk-placeholder" + else: + direct_api_key = ( + get_config().openai_api_key + or resolve_llm_openai_compatible_api_key( + str(kwargs.get("model") or os.getenv("CAI_MODEL") or "") + ) + or "sk-placeholder" + ) headers = { "Content-Type": "application/json", - "Authorization": f"Bearer {get_config().openai_api_key or 'sk-placeholder'}", + "Authorization": f"Bearer {direct_api_key}", } if stream: diff --git a/src/cai/util/nim_rotation.py b/src/cai/util/nim_rotation.py new file mode 100644 index 00000000..ccdfe63e --- /dev/null +++ b/src/cai/util/nim_rotation.py @@ -0,0 +1,49 @@ +"""Round-robin rotation for NVIDIA NIM API keys. + +Reads ``NVIDIA_NIM_API_KEY_1``, ``NVIDIA_NIM_API_KEY_2``, … from the +environment and cycles through them on every call to ``get_next_nim_key()``. + +Set ``NVIDIA_NIM_API_KEY_1``, ``NVIDIA_NIM_API_KEY_2``, ``NVIDIA_NIM_API_KEY_3`` +with your NIM keys. Each request gets the next key in sequence, spreading +the load so no single key exceeds the NIM rate limit (40 req/min). +""" + +from __future__ import annotations + +import itertools +import os +import sys + +_NIM_CYCLE: itertools.cycle[str] | None = None + + +def _ensure_cycle() -> itertools.cycle[str] | None: + global _NIM_CYCLE + if _NIM_CYCLE is not None: + return _NIM_CYCLE + keys: list[str] = [] + for i in itertools.count(1): + raw = os.getenv(f"NVIDIA_NIM_API_KEY_{i}") + if raw: + keys.append(raw.strip()) + else: + break + if keys: + _NIM_CYCLE = itertools.cycle(keys) + n = len(keys) + print( + f"[NIM] Round-robin active ({n} key{'s' if n != 1 else ''})", + file=sys.stderr, + ) + return _NIM_CYCLE + + +def is_nim_rotation_configured() -> bool: + return _ensure_cycle() is not None + + +def get_next_nim_key() -> str | None: + cycle = _ensure_cycle() + if cycle is None: + return None + return next(cycle) diff --git a/tests/util/test_nim_rotation.py b/tests/util/test_nim_rotation.py new file mode 100644 index 00000000..04372672 --- /dev/null +++ b/tests/util/test_nim_rotation.py @@ -0,0 +1,45 @@ +"""Tests for NVIDIA NIM API key round-robin rotation.""" + +from __future__ import annotations + +import cai.util.nim_rotation as nim_rotation + + +def _reset() -> None: + nim_rotation._NIM_CYCLE = None + + +def test_no_keys_configured_returns_none(monkeypatch): + _reset() + for i in range(1, 5): + monkeypatch.delenv(f"NVIDIA_NIM_API_KEY_{i}", raising=False) + assert nim_rotation.is_nim_rotation_configured() is False + assert nim_rotation.get_next_nim_key() is None + + +def test_rotates_through_all_keys(monkeypatch, capsys): + _reset() + monkeypatch.setenv("NVIDIA_NIM_API_KEY_1", "key-one") + monkeypatch.setenv("NVIDIA_NIM_API_KEY_2", "key-two") + monkeypatch.setenv("NVIDIA_NIM_API_KEY_3", "key-three") + monkeypatch.delenv("NVIDIA_NIM_API_KEY_4", raising=False) + + assert nim_rotation.is_nim_rotation_configured() is True + seen = [ + nim_rotation.get_next_nim_key(), + nim_rotation.get_next_nim_key(), + nim_rotation.get_next_nim_key(), + nim_rotation.get_next_nim_key(), + ] + assert seen == ["key-one", "key-two", "key-three", "key-one"] + + err = capsys.readouterr().err + assert "Round-robin active (3 keys)" in err + + +def test_single_key_never_moves(monkeypatch): + _reset() + monkeypatch.setenv("NVIDIA_NIM_API_KEY_1", "only-key") + monkeypatch.delenv("NVIDIA_NIM_API_KEY_2", raising=False) + assert nim_rotation.get_next_nim_key() == "only-key" + assert nim_rotation.get_next_nim_key() == "only-key"