Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
25 changes: 25 additions & 0 deletions kalibr/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -931,6 +931,7 @@ def _call_deepgram_stt(
"google/": "google",
"deepseek/": "deepseek",
"tavily/": "tavily",
"nebius/": "nebius",
}

def _dispatch(
Expand Down Expand Up @@ -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-")):
Expand Down Expand Up @@ -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:
Expand Down
29 changes: 29 additions & 0 deletions tests/test_router.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading