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

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

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:

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)

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"