From b8e1c9975e9e02cf992060ec5e336f62182e4632 Mon Sep 17 00:00:00 2001 From: Devon Kelley Date: Mon, 6 Apr 2026 22:51:51 -0700 Subject: [PATCH] feat: add Tavily search provider support --- CHANGELOG.md | 4 +++ kalibr/router.py | 76 ++++++++++++++++++++++++++++++++++++++++++++ tests/test_router.py | 46 +++++++++++++++++++++++++++ 3 files changed, 126 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7fc1649..133e73b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### 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. + ## [1.9.3] - 2026-03-29 ### Fixed diff --git a/kalibr/router.py b/kalibr/router.py index 73afc15..a0ad1e9 100644 --- a/kalibr/router.py +++ b/kalibr/router.py @@ -55,6 +55,22 @@ class TraceHandle: path: Optional[str] +def _make_chat_completion(content: str, model: str = "shim") -> Any: + """Create an OpenAI-compatible ChatCompletion shim for non-LLM providers (e.g. search APIs).""" + from types import SimpleNamespace + import uuid as _uuid + message = SimpleNamespace(content=content, role="assistant", tool_calls=None) + choice = SimpleNamespace(message=message, finish_reason="stop", index=0) + usage = SimpleNamespace(prompt_tokens=0, completion_tokens=0, total_tokens=0) + return SimpleNamespace( + choices=[choice], + model=model, + usage=usage, + id=f"shim-{_uuid.uuid4().hex[:8]}", + object="chat.completion", + ) + + class Router: """ Routes LLM requests to the best model based on learned outcomes. @@ -914,6 +930,7 @@ def _call_deepgram_stt( "anthropic/": "anthropic", "google/": "google", "deepseek/": "deepseek", + "tavily/": "tavily", } def _dispatch( @@ -943,6 +960,8 @@ def _dispatch( return self._call_google(bare_model, messages, tools, **kwargs) elif vendor == "deepseek": return self._call_deepseek(bare_model, messages, tools, **kwargs) + elif vendor == "tavily": + return self._call_tavily(bare_model, messages, tools, **kwargs) # Standard prefix checks for bare model IDs if model_id.startswith(("gpt-", "o1-", "o3-")): @@ -1109,6 +1128,63 @@ def _huggingface_to_openai_response(self, response: Any, model: str) -> Any: ), ) + def _call_tavily(self, model: str, messages: List[Dict], tools: Any, **kwargs) -> Any: + """Call Tavily Search API and return an OpenAI-compatible ChatCompletion shim. + + Supports tavily/basic and tavily/advanced search depths. + Thompson Sampling can compete Tavily against LLMs on web_scraping/research goals. + """ + try: + import httpx + except ImportError: + raise ImportError("Install 'httpx' package: pip install httpx") + + api_key = os.environ.get("TAVILY_API_KEY") + if not api_key: + raise EnvironmentError( + "TAVILY_API_KEY environment variable not set.\n" + "Get your API key from: https://app.tavily.com" + ) + + # Extract query from last user message + query = "" + for msg in reversed(messages): + if msg.get("role") == "user": + query = msg.get("content", "") + break + if not query: + return _make_chat_completion("No user message found to search.", model=f"tavily/{model}") + + # model suffix: "basic" (default) or "advanced" + search_depth = "advanced" if model == "advanced" else "basic" + + resp = httpx.post( + "https://api.tavily.com/search", + json={ + "api_key": api_key, + "query": query, + "search_depth": search_depth, + "include_answer": True, + "max_results": 5, + }, + timeout=30.0, + ) + resp.raise_for_status() + data = resp.json() + + parts = [] + answer = data.get("answer") + if answer: + parts.append(answer) + for r in data.get("results", [])[:5]: + title = r.get("title", "") + url = r.get("url", "") + snippet = r.get("content", "") + parts.append(f"- [{title}]({url}): {snippet}") + + content = "\n\n".join(parts) if parts else "No results found." + return _make_chat_completion(content, model=f"tavily/{model}") + def _anthropic_to_openai_response(self, response: Any, model: str) -> Any: """Convert Anthropic response to OpenAI format.""" from types import SimpleNamespace diff --git a/tests/test_router.py b/tests/test_router.py index 37cb3ce..84691b9 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -61,6 +61,52 @@ def test_routes_to_anthropic(self, mock_anthropic): mock_anthropic.assert_called_once() +class TestTavilyRouting: + @patch("kalibr.router.Router._call_tavily") + def test_routes_tavily_basic(self, mock_tavily): + mock_tavily.return_value = MagicMock() + router = Router(goal="test", paths=["tavily/basic"], auto_register=False) + router._dispatch("tavily/basic", [{"role": "user", "content": "AI news"}], None) + mock_tavily.assert_called_once() + assert mock_tavily.call_args[0][0] == "basic" + + @patch("kalibr.router.Router._call_tavily") + def test_routes_tavily_advanced(self, mock_tavily): + mock_tavily.return_value = MagicMock() + router = Router(goal="test", paths=["tavily/advanced"], auto_register=False) + router._dispatch("tavily/advanced", [{"role": "user", "content": "AI news"}], None) + mock_tavily.assert_called_once() + assert mock_tavily.call_args[0][0] == "advanced" + + def test_tavily_missing_api_key(self): + import os + router = Router(goal="test", paths=["tavily/basic"], auto_register=False) + env_backup = os.environ.pop("TAVILY_API_KEY", None) + try: + with pytest.raises(EnvironmentError, match="TAVILY_API_KEY"): + router._call_tavily("basic", [{"role": "user", "content": "test"}], None) + finally: + if env_backup: + os.environ["TAVILY_API_KEY"] = env_backup + + @patch("httpx.post") + def test_tavily_response_shape(self, mock_post): + import os + os.environ["TAVILY_API_KEY"] = "tvly-test" + mock_post.return_value = MagicMock( + json=lambda: { + "answer": "AI is advancing rapidly.", + "results": [{"title": "AI News", "url": "https://example.com", "content": "details"}] + } + ) + mock_post.return_value.raise_for_status = lambda: None + router = Router(goal="test", paths=["tavily/basic"], auto_register=False) + resp = router._call_tavily("basic", [{"role": "user", "content": "AI news"}], None) + assert hasattr(resp, "choices") + assert resp.choices[0].message.content.startswith("AI is advancing rapidly.") + os.environ.pop("TAVILY_API_KEY", None) + + class TestRouterReport: def test_double_report_warning(self): router = Router(goal="test", auto_register=False)