Coverage for src/local_deep_research/web_search_engines/engines/search_engine_tavily.py: 98%
72 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1from typing import Any, Dict, List, Optional
3import requests
4from langchain_core.language_models import BaseLLM
5from ...security.secure_logging import logger
7from ...security.safe_requests import safe_post
8from ..rate_limiting import RateLimitError
9from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity
11# Map LDR time_period codes to Tavily API time_range values.
12# "all" and any unrecognized value (None, "", wrong case, whitespace,
13# unknown strings) omit time_range so no filter is applied. The settings
14# UI only emits the canonical d/w/m/y/all codes; programmatic callers
15# must match exactly — we deliberately do NOT .lower()/strip() so a
16# mistyped value fails visibly as "unfiltered" rather than silently
17# being coerced.
18_TIME_RANGE_MAP = {"d": "day", "w": "week", "m": "month", "y": "year"}
21class TavilySearchEngine(BaseSearchEngine):
22 """Tavily search engine implementation with two-phase approach"""
24 # Mark as public search engine
25 is_public = True
26 egress_sensitivity = Sensitivity.NON_SENSITIVE
27 egress_exposure = Exposure.EXPOSING
28 # Mark as generic search engine (general web search)
29 is_generic = True
31 def __init__(
32 self,
33 max_results: int = 10,
34 region: str = "US",
35 time_period: str = "all",
36 safe_search: bool = True,
37 search_language: str = "English",
38 api_key: Optional[str] = None,
39 llm: Optional[BaseLLM] = None,
40 include_full_content: bool = True,
41 max_filtered_results: Optional[int] = None,
42 search_depth: str = "basic",
43 include_domains: Optional[List[str]] = None,
44 exclude_domains: Optional[List[str]] = None,
45 settings_snapshot: Optional[Dict[str, Any]] = None,
46 **kwargs,
47 ):
48 """
49 Initialize the Tavily search engine.
51 Args:
52 max_results: Maximum number of search results
53 region: Region code for search results (not used by Tavily currently)
54 time_period: Time period filter (d/w/m/y/all); forwarded to the
55 Tavily ``time_range`` param as day/week/month/year. ``all``
56 (or any unrecognized value) omits the filter.
57 safe_search: Whether to enable safe search (not used by Tavily currently)
58 search_language: Language for search results (not used by Tavily currently)
59 api_key: Tavily API key (can also be set via LDR_SEARCH_ENGINE_WEB_TAVILY_API_KEY env var or in UI settings)
60 llm: Language model for relevance filtering
61 include_full_content: Whether to include full webpage content in results
62 max_filtered_results: Maximum number of results to keep after filtering
63 search_depth: "basic" or "advanced" - controls search quality vs speed
64 include_domains: List of domains to include in search
65 exclude_domains: List of domains to exclude from search
66 settings_snapshot: Settings snapshot for thread context
67 **kwargs: Additional parameters (ignored but accepted for compatibility)
68 """
69 # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
70 super().__init__(
71 llm=llm,
72 max_filtered_results=max_filtered_results,
73 max_results=max_results,
74 include_full_content=include_full_content,
75 settings_snapshot=settings_snapshot,
76 )
77 self.search_depth = search_depth
78 self.time_period = time_period
79 self.include_domains = include_domains or []
80 self.exclude_domains = exclude_domains or []
82 # Get API key - check params, settings, or env vars
83 tavily_api_key = self._resolve_api_key(
84 api_key,
85 "search.engine.web.tavily.api_key",
86 engine_name="Tavily",
87 settings_snapshot=settings_snapshot,
88 )
90 self.api_key = tavily_api_key
91 self.base_url = "https://api.tavily.com"
93 # If full content is requested, initialize FullSearchResults
94 if include_full_content:
95 # Create a simple wrapper for Tavily API calls
96 class TavilyWrapper:
97 def __init__(self, parent):
98 self.parent = parent
100 def run(self, query):
101 return self.parent._get_previews(query)
103 self._init_full_search(
104 web_search=TavilyWrapper(self),
105 language=search_language,
106 max_results=max_results,
107 region=region,
108 time_period=time_period,
109 safe_search="moderate" if safe_search else "off",
110 )
112 def _get_previews(self, query: str) -> List[Dict[str, Any]]:
113 """
114 Get preview information from Tavily Search.
116 Args:
117 query: The search query
119 Returns:
120 List of preview dictionaries
121 """
122 logger.info("Getting search results from Tavily")
124 try:
125 # Prepare the request payload
126 payload = {
127 "api_key": self.api_key,
128 "query": query[:400], # Limit query length
129 "search_depth": self.search_depth,
130 "max_results": min(
131 20, self.max_results
132 ), # Tavily has a max limit
133 "include_answer": False, # We don't need the AI answer
134 "include_images": False, # We don't need images
135 "include_raw_content": self.include_full_content, # Get content if requested
136 }
138 # Add domain filters if specified
139 if self.include_domains:
140 payload["include_domains"] = self.include_domains
141 if self.exclude_domains:
142 payload["exclude_domains"] = self.exclude_domains
144 # Apply time-period filter (Tavily time_range). "all" / unknown
145 # map to None and are omitted so no filter is applied. Non-string
146 # values (possible via programmatic construction) are treated as
147 # "no filter" rather than raising on the unhashable dict lookup.
148 time_range = (
149 _TIME_RANGE_MAP.get(self.time_period)
150 if isinstance(self.time_period, str)
151 else None
152 )
153 if time_range:
154 payload["time_range"] = time_range
156 # Apply rate limiting before request
157 self._last_wait_time = self.rate_tracker.apply_rate_limit(
158 self.engine_type
159 )
161 # Make the API request
162 response = safe_post(
163 f"{self.base_url}/search",
164 json=payload,
165 headers={"Content-Type": "application/json"},
166 timeout=30,
167 )
169 # Check for rate limits
170 self._raise_if_rate_limit(response.status_code)
172 response.raise_for_status()
174 # Parse the response
175 data = response.json()
176 results = data.get("results", [])
178 # Format results as previews
179 previews = []
180 for i, result in enumerate(results):
181 url = self._clean_result_url(result.get("url"))
182 preview = {
183 "id": url or str(i), # Use URL as ID
184 "title": result.get("title", ""),
185 "link": url,
186 "snippet": result.get(
187 "content", ""
188 ), # Tavily calls it "content"
189 "displayed_link": url,
190 "position": i,
191 }
193 # Store full Tavily result for later
194 preview["_full_result"] = result
196 previews.append(preview)
198 # Store the previews for potential full content retrieval
199 self._search_results = previews
201 return previews
203 except RateLimitError:
204 raise # Re-raise rate limit errors
205 except requests.exceptions.RequestException as e:
206 safe_msg = self._scrub_error(e)
207 logger.warning(f"Error getting Tavily results: {safe_msg}")
208 self._raise_if_rate_limit(e)
209 return []
210 except Exception as e:
211 safe_msg = self._scrub_error(e)
212 logger.warning(
213 f"Unexpected error getting Tavily results: {safe_msg}"
214 )
215 return []
217 def _get_full_content(
218 self, relevant_items: List[Dict[str, Any]]
219 ) -> List[Dict[str, Any]]:
220 """
221 Get full content for the relevant search results.
222 Extends base implementation to include Tavily's raw_content.
224 Args:
225 relevant_items: List of relevant preview dictionaries
227 Returns:
228 List of result dictionaries with full content if available
229 """
230 results = super()._get_full_content(relevant_items)
232 # If Tavily provided raw_content and full content is requested, use it
233 if self.include_full_content:
234 for result in results:
235 if "raw_content" in result: 235 ↛ 234line 235 didn't jump to line 234 because the condition on line 235 was always true
236 result["content"] = result.get(
237 "raw_content", result.get("content", "")
238 )
240 return results