Coverage for src/local_deep_research/web_search_engines/engines/search_engine_tavily.py: 99%
65 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +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: Retained for backward compatibility; full-content
62 retrieval is handled via the factory wrapper.
63 max_filtered_results: Maximum number of results to keep after filtering
64 search_depth: "basic" or "advanced" - controls search quality vs speed
65 include_domains: List of domains to include in search
66 exclude_domains: List of domains to exclude from search
67 settings_snapshot: Settings snapshot for thread context
68 **kwargs: Additional parameters (ignored but accepted for compatibility)
69 """
70 # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
71 super().__init__(
72 llm=llm,
73 max_filtered_results=max_filtered_results,
74 max_results=max_results,
75 include_full_content=include_full_content,
76 settings_snapshot=settings_snapshot,
77 )
78 self.search_depth = search_depth
79 self.time_period = time_period
80 self.include_domains = include_domains or []
81 self.exclude_domains = exclude_domains or []
83 # Get API key - check params, settings, or env vars
84 tavily_api_key = self._resolve_api_key(
85 api_key,
86 "search.engine.web.tavily.api_key",
87 engine_name="Tavily",
88 settings_snapshot=settings_snapshot,
89 )
91 self.api_key = tavily_api_key
92 self.base_url = "https://api.tavily.com"
94 # Full-content retrieval is the factory wrapper's job — the
95 # wrapper fetches pages via ``batch_fetch_and_extract`` and
96 # populates ``full_content`` on each result; the inner engine
97 # only emits snippets from ``_get_previews``.
99 def _get_previews(self, query: str) -> List[Dict[str, Any]]:
100 """
101 Get preview information from Tavily Search.
103 Args:
104 query: The search query
106 Returns:
107 List of preview dictionaries
108 """
109 logger.info("Getting search results from Tavily")
111 try:
112 # Prepare the request payload
113 payload = {
114 "api_key": self.api_key,
115 "query": query[:400], # Limit query length
116 "search_depth": self.search_depth,
117 "max_results": min(
118 20, self.max_results
119 ), # Tavily has a max limit
120 "include_answer": False, # We don't need the AI answer
121 "include_images": False, # We don't need images
122 "include_raw_content": self.include_full_content, # Get content if requested
123 }
125 # Add domain filters if specified
126 if self.include_domains:
127 payload["include_domains"] = self.include_domains
128 if self.exclude_domains:
129 payload["exclude_domains"] = self.exclude_domains
131 # Apply time-period filter (Tavily time_range). "all" / unknown
132 # map to None and are omitted so no filter is applied. Non-string
133 # values (possible via programmatic construction) are treated as
134 # "no filter" rather than raising on the unhashable dict lookup.
135 time_range = (
136 _TIME_RANGE_MAP.get(self.time_period)
137 if isinstance(self.time_period, str)
138 else None
139 )
140 if time_range:
141 payload["time_range"] = time_range
143 # Apply rate limiting before request
144 self._last_wait_time = self.rate_tracker.apply_rate_limit(
145 self.engine_type
146 )
148 # Make the API request
149 response = safe_post(
150 f"{self.base_url}/search",
151 json=payload,
152 headers={"Content-Type": "application/json"},
153 timeout=30,
154 )
156 # Check for rate limits
157 self._raise_if_rate_limit(response.status_code)
159 response.raise_for_status()
161 # Parse the response
162 data = response.json()
163 results = data.get("results", [])
165 # Format results as previews
166 previews = []
167 for i, result in enumerate(results):
168 url = self._clean_result_url(result.get("url"))
169 preview = {
170 "id": url or str(i), # Use URL as ID
171 "title": result.get("title", ""),
172 "link": url,
173 "snippet": result.get(
174 "content", ""
175 ), # Tavily calls it "content"
176 "displayed_link": url,
177 "position": i,
178 }
180 # Store full Tavily result for later
181 preview["_full_result"] = result
183 previews.append(preview)
185 # Store the previews for potential full content retrieval
186 self._search_results = previews
188 return previews
190 except RateLimitError:
191 raise # Re-raise rate limit errors
192 except requests.exceptions.RequestException as e:
193 safe_msg = self._scrub_error(e)
194 logger.warning(f"Error getting Tavily results: {safe_msg}")
195 self._raise_if_rate_limit(e)
196 return []
197 except Exception as e:
198 safe_msg = self._scrub_error(e)
199 logger.warning(
200 f"Unexpected error getting Tavily results: {safe_msg}"
201 )
202 return []
204 def _get_full_content(
205 self, relevant_items: List[Dict[str, Any]]
206 ) -> List[Dict[str, Any]]:
207 """
208 Get full content for the relevant search results.
209 Extends base implementation to include Tavily's raw_content.
211 Args:
212 relevant_items: List of relevant preview dictionaries
214 Returns:
215 List of result dictionaries with full content if available
216 """
217 results = super()._get_full_content(relevant_items)
219 # If Tavily provided raw_content and full content is requested, use it
220 if self.include_full_content:
221 for result in results:
222 if "raw_content" in result: 222 ↛ 221line 222 didn't jump to line 221 because the condition on line 222 was always true
223 result["content"] = result.get(
224 "raw_content", result.get("content", "")
225 )
227 return results