From 4546033b8e81d335c88e220e93df2258fac35fde Mon Sep 17 00:00:00 2001 From: quevedoSteven Date: Fri, 31 Jul 2026 15:51:30 +0000 Subject: [PATCH 1/4] 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" From c2408889d2840ab9cd5eb87760a234bd62824038 Mon Sep 17 00:00:00 2001 From: UnaiAlias <52742669+UnaiAlias@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:53:57 +0000 Subject: [PATCH 2/4] fix: remove duplicate litellm stream call in fetch_response_litellm_openai The streaming path called litellm.acompletion twice with the same kwargs. The first result was assigned to a variable that is never read, and only the second call's stream was returned. Every streamed request therefore consumed two provider requests (e.g. two NVIDIA NIM requests) with the same API key, halving effective per-key rate limits and leaking the first stream. Keep a single acompletion call for the streamed path; same fix in the tool_call_id truncation retry branch. --- src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py b/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py index 795d65a0..8533bb68 100644 --- a/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py +++ b/src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py @@ -68,7 +68,6 @@ async def fetch_response_litellm_openai( """ try: if stream: - ret = await litellm.acompletion(**kwargs) stream_obj = await litellm.acompletion(**kwargs) return _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls), stream_obj else: @@ -102,7 +101,6 @@ async def fetch_response_litellm_openai( kwargs["messages"] = messages if stream: - ret = await litellm.acompletion(**kwargs) stream_obj = await litellm.acompletion(**kwargs) return _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls), stream_obj else: From 929c2ea6c395cf7bc91fe0f37a735fea69fdacb5 Mon Sep 17 00:00:00 2001 From: UnaiAlias <52742669+UnaiAlias@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:07:37 +0000 Subject: [PATCH 3/4] fix: only apply NIM rotation to known NIM models The NIM key rotation was incorrectly applying to all models when NVIDIA_NIM_API_KEY_N are configured. This broke non-NIM models like z-ai/glm-5.2 which should route through the Alias gateway, not NVIDIA NIM. Add _is_nim_model() check with known NIM provider prefixes to gate rotation. Fix applies to both LiteLLM path and direct httpx path. --- .../agents/models/openai_chatcompletions.py | 45 ++++++++++++++----- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 6324fd6e..48d15683 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -3668,16 +3668,27 @@ class OpenAIChatCompletionsModel(Model): # 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. + # their own endpoint + key. Also only applies to known NIM models. + def _is_nim_model(model: str) -> bool: + """Check if model is a known NIM model (provider is on NVIDIA NIM).""" + provider = model.split("/")[0].lower() if "/" in model else model.lower() + nim_providers = { + "nvidia", "meta", "google", "microsoft", "mistralai", + "llama", "nemotron", "codellama", "mixtral", "phi", + "qwen", "yi", "deepseek", "gemma", "stable-diffusion", + "sdxl", "cosxl", "proteus", "realistic-vision", + } + return provider in nim_providers + 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("/") - ) + model_str = str(kwargs.get("model") or os.getenv("CAI_MODEL") or "") + if _is_nim_model(model_str): + kwargs["api_key"] = get_next_nim_key() + kwargs["custom_llm_provider"] = "openai" + kwargs["api_base"] = ( + resolve_llm_openai_compatible_base(model_str).rstrip("/") + or os.getenv("OPENAI_API_BASE", "").rstrip("/") + ) # Filter out NotGiven values to avoid JSON serialization issues filtered_kwargs = {} @@ -3785,7 +3796,21 @@ 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(): + model_str = str(kwargs.get("model") or os.getenv("CAI_MODEL") or "") + def _is_nim_model(model: str) -> bool: + provider = model.split("/")[0].lower() if "/" in model else model.lower() + nim_providers = { + "nvidia", "meta", "google", "microsoft", "mistralai", + "llama", "nemotron", "codellama", "mixtral", "phi", + "qwen", "yi", "deepseek", "gemma", "stable-diffusion", + "sdxl", "cosxl", "proteus", "realistic-vision", + } + return provider in nim_providers + if ( + "api.nvidia.com" in openai_api_base.lower() + and is_nim_rotation_configured() + and _is_nim_model(model_str) + ): direct_api_key = get_next_nim_key() or "sk-placeholder" else: direct_api_key = ( From f9e886f7923342157a9d3ea42d25806bc325221c Mon Sep 17 00:00:00 2001 From: UnaiAlias <52742669+UnaiAlias@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:10:50 +0000 Subject: [PATCH 4/4] fix: add z-ai to NIM model providers z-ai/glm-5.2 is hosted on NVIDIA NIM (build.nvidia.com/z-ai/glm-5.2). The NIM rotation was incorrectly excluding it. Add "z-ai" to the known NIM provider list in both LiteLLM and direct httpx paths. --- src/cai/sdk/agents/models/openai_chatcompletions.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 48d15683..60602da0 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -3677,6 +3677,7 @@ class OpenAIChatCompletionsModel(Model): "llama", "nemotron", "codellama", "mixtral", "phi", "qwen", "yi", "deepseek", "gemma", "stable-diffusion", "sdxl", "cosxl", "proteus", "realistic-vision", + "z-ai", } return provider in nim_providers @@ -3804,6 +3805,7 @@ class OpenAIChatCompletionsModel(Model): "llama", "nemotron", "codellama", "mixtral", "phi", "qwen", "yi", "deepseek", "gemma", "stable-diffusion", "sdxl", "cosxl", "proteus", "realistic-vision", + "z-ai", } return provider in nim_providers if (