From 08b391797586b410bd3e66e30f7df268e73299dd Mon Sep 17 00:00:00 2001 From: Tavily PR Agent Date: Mon, 4 May 2026 15:07:37 +0000 Subject: [PATCH] feat: add Tavily as parallel search provider --- backend/app/services/search_service.py | 55 +++++++++++++++--- backend/requirements.txt | 3 +- backend/tests/test_search_service.py | 79 +++++++++++++++++++++++++- readme.md | 1 + 4 files changed, 129 insertions(+), 9 deletions(-) diff --git a/backend/app/services/search_service.py b/backend/app/services/search_service.py index 8fc144f..0ddd110 100644 --- a/backend/app/services/search_service.py +++ b/backend/app/services/search_service.py @@ -10,6 +10,7 @@ from app.models.search_model import SearchResult from app.utils.rate_limter import rate_limit from app.services.youtube_service import YouTubeAPIError +from tavily import TavilyClient # Constants BING_ENDPOINT = "https://api.bing.microsoft.com/v7.0/search" @@ -240,25 +241,65 @@ def search_youtube(query: str) -> List[SearchResult]: except Exception as e: raise YouTubeAPIError(f"YouTube search failed: {str(e)}") +@rate_limit(calls=CALLS_PER_MINUTE, period=60) +def search_tavily(query: str) -> List[SearchResult]: + """Perform a Tavily search for the given query. + + Args: + query: The search query to perform + + Returns: + List of SearchResult objects + + Raises: + SearchAPIError: If the Tavily API request fails + """ + api_key = os.getenv('TAVILY_API_KEY') + if not api_key: + raise SearchAPIError("TAVILY_API_KEY environment variable not set") + + try: + client = TavilyClient(api_key=api_key) + response = client.search(query, max_results=RESULTS_PER_ENGINE) + + results = [] + for item in response.get("results", []): + search_result = SearchResult( + question=query, + title=item.get("title", ""), + url=item.get("url", ""), + snippet=item.get("content", ""), + search_content=item.get("raw_content") or item.get("content", ""), + source="tavily" + ) + results.append(search_result) + + return results + + except Exception as e: + logger.error(f"Tavily search error: {str(e)}") + raise SearchAPIError(f"Tavily search failed: {str(e)}") + def perform_search(query: str) -> List[SearchResult]: - """Perform parallel searches on Google, Bing, and YouTube APIs. - + """Perform parallel searches on Google, Bing, YouTube, and Tavily APIs. + Args: query: The search query to run - + Returns: Combined list of unique SearchResult objects """ try: - with ThreadPoolExecutor(max_workers=3) as executor: + with ThreadPoolExecutor(max_workers=4) as executor: bing_future = executor.submit(search_bing, query) google_future = executor.submit(search_google, query) youtube_future = executor.submit(search_youtube, query) - + tavily_future = executor.submit(search_tavily, query) + results = [] - + # Gather results, handling potential failures - for future in [bing_future, google_future, youtube_future]: + for future in [bing_future, google_future, youtube_future, tavily_future]: try: results.extend(future.result()) except (SearchAPIError, YouTubeAPIError) as e: diff --git a/backend/requirements.txt b/backend/requirements.txt index 9e75875..29e82bb 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -6,4 +6,5 @@ python-decouple pydantic-settings # Add any other dependencies your project needs gunicorn -bs4 \ No newline at end of file +bs4 +tavily-python \ No newline at end of file diff --git a/backend/tests/test_search_service.py b/backend/tests/test_search_service.py index 174c4c4..38d8201 100644 --- a/backend/tests/test_search_service.py +++ b/backend/tests/test_search_service.py @@ -1,7 +1,9 @@ import pytest -from app.services.search_service import perform_search, fetch_content_from_custom_url +from unittest.mock import patch, MagicMock +from app.services.search_service import perform_search, fetch_content_from_custom_url, search_tavily, SearchAPIError from app.models.search_model import SearchResult + @pytest.mark.asyncio async def test_custom_url_fetch(): print("Running test_custom_url_fetch") @@ -14,3 +16,78 @@ async def test_custom_url_fetch(): url = "https://example.com" with pytest.raises(Exception): await fetch_content_from_custom_url(url) + + +def test_search_tavily_missing_key(): + """Test that search_tavily raises SearchAPIError when TAVILY_API_KEY is not set.""" + with patch.dict("os.environ", {}, clear=True): + with pytest.raises(SearchAPIError, match="TAVILY_API_KEY"): + search_tavily("test query") + + +def test_search_tavily_success(): + """Test that search_tavily returns correctly mapped SearchResult objects.""" + mock_response = { + "results": [ + { + "title": "Tavily Result 1", + "url": "https://example.com/1", + "content": "Snippet for result 1", + "raw_content": "Full content for result 1", + }, + { + "title": "Tavily Result 2", + "url": "https://example.com/2", + "content": "Snippet for result 2", + "raw_content": None, + }, + ] + } + mock_client = MagicMock() + mock_client.search.return_value = mock_response + + with patch.dict("os.environ", {"TAVILY_API_KEY": "tvly-test-key"}): + with patch("app.services.search_service.TavilyClient", return_value=mock_client): + results = search_tavily("test query") + + assert len(results) == 2 + assert results[0].source == "tavily" + assert results[0].title == "Tavily Result 1" + assert results[0].url == "https://example.com/1" + assert results[0].snippet == "Snippet for result 1" + assert results[0].search_content == "Full content for result 1" + # When raw_content is None, falls back to content + assert results[1].search_content == "Snippet for result 2" + + +def test_search_tavily_api_error(): + """Test that search_tavily wraps API errors as SearchAPIError.""" + mock_client = MagicMock() + mock_client.search.side_effect = Exception("API rate limit exceeded") + + with patch.dict("os.environ", {"TAVILY_API_KEY": "tvly-test-key"}): + with patch("app.services.search_service.TavilyClient", return_value=mock_client): + with pytest.raises(SearchAPIError, match="Tavily search failed"): + search_tavily("test query") + + +def test_perform_search_includes_tavily(): + """Test that perform_search includes Tavily results in combined output.""" + tavily_result = SearchResult( + question="test", + title="Tavily Result", + url="https://tavily.example.com", + snippet="tavily snippet", + search_content="tavily content", + source="tavily", + ) + + with patch("app.services.search_service.search_bing", side_effect=SearchAPIError("no key")), \ + patch("app.services.search_service.search_google", side_effect=SearchAPIError("no key")), \ + patch("app.services.search_service.search_youtube", side_effect=SearchAPIError("no key")), \ + patch("app.services.search_service.search_tavily", return_value=[tavily_result]): + results = perform_search("test") + + assert len(results) == 1 + assert results[0].source == "tavily" + assert results[0].url == "https://tavily.example.com" diff --git a/readme.md b/readme.md index aff7c2f..22515ac 100644 --- a/readme.md +++ b/readme.md @@ -228,6 +228,7 @@ Open [http://localhost:5173](http://localhost:5173) and start asking. | `GOOGLE_API_KEY` | ✅ | Google Custom Search API key | | `GOOGLE_SEARCH_CX` | ✅ | Custom Search Engine ID | | `BING_API_KEY` | ✅ | Bing Web Search v7 key | +| `TAVILY_API_KEY` | | Tavily Search API key (optional, additive provider) | | `CLERK_SECRET_KEY` | | Server-side Clerk verification | | `ALLOWED_ORIGINS` | | Comma-separated CORS origins |