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/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: diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 47d871b0..60602da0 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,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: 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"