Merge f9e886f792 into 62871b6f5a

.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/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:

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,33 @@ 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. 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",
+                "z-ai",
+            }
+            return provider in nim_providers
+
+        if is_nim_rotation_configured() and not kwargs.get("api_base"):
+            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 = {}
         for key, value in kwargs.items():
@@ -3763,9 +3797,34 @@ 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"
+                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",
+                        "z-ai",
+                    }
+                    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 = (
+                        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"