diff --git a/CHANGELOG.md b/CHANGELOG.md index 133e73b..ec0b659 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - **Tavily Search provider** — `tavily/basic` and `tavily/advanced` as Router paths. Returns web search results wrapped in an OpenAI-compatible ChatCompletion shim so Thompson Sampling can compete Tavily against LLMs on web research goals. Set `TAVILY_API_KEY` env var. +- **Nebius AI provider** — `nebius/` prefix routes to Nebius AI Studio (OpenAI-compatible). Set `NEBIUS_API_KEY` env var. Supported models: `nebius/meta-llama/Llama-3.3-70B-Instruct`, `nebius/Qwen/Qwen2.5-72B-Instruct`, `nebius/mistralai/Mistral-Nemo-Instruct-2407`. ## [1.9.3] - 2026-03-29 diff --git a/kalibr/router.py b/kalibr/router.py index a0ad1e9..1e7b97d 100644 --- a/kalibr/router.py +++ b/kalibr/router.py @@ -931,6 +931,7 @@ def _call_deepgram_stt( "google/": "google", "deepseek/": "deepseek", "tavily/": "tavily", + "nebius/": "nebius", } def _dispatch( @@ -962,6 +963,8 @@ def _dispatch( return self._call_deepseek(bare_model, messages, tools, **kwargs) elif vendor == "tavily": return self._call_tavily(bare_model, messages, tools, **kwargs) + elif vendor == "nebius": + return self._call_nebius(bare_model, messages, tools, **kwargs) # Standard prefix checks for bare model IDs if model_id.startswith(("gpt-", "o1-", "o3-")): @@ -1017,6 +1020,28 @@ def _call_deepseek(self, model: str, messages: List[Dict], tools: Any, **kwargs) call_kwargs = {"model": model, "messages": messages, **kwargs} return client.chat.completions.create(**call_kwargs) + def _call_nebius(self, model: str, messages: List[Dict], tools: Any, **kwargs) -> Any: + """Call Nebius AI API using OpenAI-compatible client with Nebius base URL.""" + try: + from openai import OpenAI + except ImportError: + raise ImportError("Install 'openai' package: pip install openai") + + api_key = os.environ.get("NEBIUS_API_KEY") + if not api_key: + raise EnvironmentError( + "NEBIUS_API_KEY environment variable not set.\n" + "Get your API key from: https://studio.nebius.ai" + ) + + client = OpenAI( + api_key=api_key, + base_url="https://api.studio.nebius.ai/v1/", + ) + + call_kwargs = {"model": model, "messages": messages, **kwargs} + return client.chat.completions.create(**call_kwargs) + def _call_anthropic(self, model: str, messages: List[Dict], tools: Any, **kwargs) -> Any: """Call Anthropic API and convert response to OpenAI format.""" try: diff --git a/tests/test_router.py b/tests/test_router.py index 84691b9..14b5dea 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -107,6 +107,35 @@ def test_tavily_response_shape(self, mock_post): os.environ.pop("TAVILY_API_KEY", None) +class TestNebiusRouting: + @patch("kalibr.router.Router._call_nebius") + def test_routes_nebius_prefix(self, mock_nebius): + mock_nebius.return_value = MagicMock() + router = Router(goal="test", paths=["nebius/meta-llama/Llama-3.3-70B-Instruct"], auto_register=False) + router._dispatch("nebius/meta-llama/Llama-3.3-70B-Instruct", [{"role": "user", "content": "hi"}], None) + mock_nebius.assert_called_once() + assert mock_nebius.call_args[0][0] == "meta-llama/Llama-3.3-70B-Instruct" + + @patch("kalibr.router.Router._call_nebius") + def test_routes_nebius_qwen(self, mock_nebius): + mock_nebius.return_value = MagicMock() + router = Router(goal="test", paths=["nebius/Qwen/Qwen2.5-72B-Instruct"], auto_register=False) + router._dispatch("nebius/Qwen/Qwen2.5-72B-Instruct", [{"role": "user", "content": "hi"}], None) + mock_nebius.assert_called_once() + assert mock_nebius.call_args[0][0] == "Qwen/Qwen2.5-72B-Instruct" + + def test_nebius_missing_api_key(self): + import os + router = Router(goal="test", paths=["nebius/meta-llama/Llama-3.3-70B-Instruct"], auto_register=False) + env_backup = os.environ.pop("NEBIUS_API_KEY", None) + try: + with pytest.raises(EnvironmentError, match="NEBIUS_API_KEY"): + router._call_nebius("meta-llama/Llama-3.3-70B-Instruct", [{"role": "user", "content": "hi"}], None) + finally: + if env_backup: + os.environ["NEBIUS_API_KEY"] = env_backup + + class TestRouterReport: def test_double_report_warning(self): router = Router(goal="test", auto_register=False)