Coverage for src/local_deep_research/web_search_engines/engines/search_engine_tavily.py: 98%

1from typing import Any, Dict, List, Optional 

2 

3import requests 

4from langchain_core.language_models import BaseLLM 

5from ...security.secure_logging import logger 

6 

7from ...security.safe_requests import safe_post 

8from ..rate_limiting import RateLimitError 

9from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity 

10 

11 

12class TavilySearchEngine(BaseSearchEngine): 

13 """Tavily search engine implementation with two-phase approach""" 

14 

15 # Mark as public search engine 

16 is_public = True 

17 egress_sensitivity = Sensitivity.NON_SENSITIVE 

18 egress_exposure = Exposure.EXPOSING 

19 # Mark as generic search engine (general web search) 

20 is_generic = True 

21 

22 def __init__( 

23 self, 

24 max_results: int = 10, 

25 region: str = "US", 

26 time_period: str = "y", 

27 safe_search: bool = True, 

28 search_language: str = "English", 

29 api_key: Optional[str] = None, 

30 llm: Optional[BaseLLM] = None, 

31 include_full_content: bool = True, 

32 max_filtered_results: Optional[int] = None, 

33 search_depth: str = "basic", 

34 include_domains: Optional[List[str]] = None, 

35 exclude_domains: Optional[List[str]] = None, 

36 settings_snapshot: Optional[Dict[str, Any]] = None, 

37 **kwargs, 

38 ): 

39 """ 

40 Initialize the Tavily search engine. 

41 

42 Args: 

43 max_results: Maximum number of search results 

44 region: Region code for search results (not used by Tavily currently) 

45 time_period: Time period for search results (not used by Tavily currently) 

46 safe_search: Whether to enable safe search (not used by Tavily currently) 

47 search_language: Language for search results (not used by Tavily currently) 

48 api_key: Tavily API key (can also be set via LDR_SEARCH_ENGINE_WEB_TAVILY_API_KEY env var or in UI settings) 

49 llm: Language model for relevance filtering 

50 include_full_content: Whether to include full webpage content in results 

51 max_filtered_results: Maximum number of results to keep after filtering 

52 search_depth: "basic" or "advanced" - controls search quality vs speed 

53 include_domains: List of domains to include in search 

54 exclude_domains: List of domains to exclude from search 

55 settings_snapshot: Settings snapshot for thread context 

56 **kwargs: Additional parameters (ignored but accepted for compatibility) 

57 """ 

58 # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results 

59 super().__init__( 

60 llm=llm, 

61 max_filtered_results=max_filtered_results, 

62 max_results=max_results, 

63 include_full_content=include_full_content, 

64 settings_snapshot=settings_snapshot, 

65 ) 

66 self.search_depth = search_depth 

67 self.include_domains = include_domains or [] 

68 self.exclude_domains = exclude_domains or [] 

69 

70 # Get API key - check params, settings, or env vars 

71 tavily_api_key = self._resolve_api_key( 

72 api_key, 

73 "search.engine.web.tavily.api_key", 

74 engine_name="Tavily", 

75 settings_snapshot=settings_snapshot, 

76 ) 

77 

78 self.api_key = tavily_api_key 

79 self.base_url = "https://api.tavily.com" 

80 

81 # If full content is requested, initialize FullSearchResults 

82 if include_full_content: 

83 # Create a simple wrapper for Tavily API calls 

84 class TavilyWrapper: 

85 def __init__(self, parent): 

86 self.parent = parent 

87 

88 def run(self, query): 

89 return self.parent._get_previews(query) 

90 

91 self._init_full_search( 

92 web_search=TavilyWrapper(self), 

93 language=search_language, 

94 max_results=max_results, 

95 region=region, 

96 time_period=time_period, 

97 safe_search="moderate" if safe_search else "off", 

98 ) 

99 

100 def _get_previews(self, query: str) -> List[Dict[str, Any]]: 

101 """ 

102 Get preview information from Tavily Search. 

103 

104 Args: 

105 query: The search query 

106 

107 Returns: 

108 List of preview dictionaries 

109 """ 

110 logger.info("Getting search results from Tavily") 

111 

112 try: 

113 # Prepare the request payload 

114 payload = { 

115 "api_key": self.api_key, 

116 "query": query[:400], # Limit query length 

117 "search_depth": self.search_depth, 

118 "max_results": min( 

119 20, self.max_results 

120 ), # Tavily has a max limit 

121 "include_answer": False, # We don't need the AI answer 

122 "include_images": False, # We don't need images 

123 "include_raw_content": self.include_full_content, # Get content if requested 

124 } 

125 

126 # Add domain filters if specified 

127 if self.include_domains: 

128 payload["include_domains"] = self.include_domains 

129 if self.exclude_domains: 

130 payload["exclude_domains"] = self.exclude_domains 

131 

132 # Apply rate limiting before request 

133 self._last_wait_time = self.rate_tracker.apply_rate_limit( 

134 self.engine_type 

135 ) 

136 

137 # Make the API request 

138 response = safe_post( 

139 f"{self.base_url}/search", 

140 json=payload, 

141 headers={"Content-Type": "application/json"}, 

142 timeout=30, 

143 ) 

144 

145 # Check for rate limits 

146 self._raise_if_rate_limit(response.status_code) 

147 

148 response.raise_for_status() 

149 

150 # Parse the response 

151 data = response.json() 

152 results = data.get("results", []) 

153 

154 # Format results as previews 

155 previews = [] 

156 for i, result in enumerate(results): 

157 url = self._clean_result_url(result.get("url")) 

158 preview = { 

159 "id": url or str(i), # Use URL as ID 

160 "title": result.get("title", ""), 

161 "link": url, 

162 "snippet": result.get( 

163 "content", "" 

164 ), # Tavily calls it "content" 

165 "displayed_link": url, 

166 "position": i, 

167 } 

168 

169 # Store full Tavily result for later 

170 preview["_full_result"] = result 

171 

172 previews.append(preview) 

173 

174 # Store the previews for potential full content retrieval 

175 self._search_results = previews 

176 

177 return previews 

178 

179 except RateLimitError: 

180 raise # Re-raise rate limit errors 

181 except requests.exceptions.RequestException as e: 

182 safe_msg = self._scrub_error(e) 

183 logger.warning(f"Error getting Tavily results: {safe_msg}") 

184 self._raise_if_rate_limit(e) 

185 return [] 

186 except Exception as e: 

187 safe_msg = self._scrub_error(e) 

188 logger.warning( 

189 f"Unexpected error getting Tavily results: {safe_msg}" 

190 ) 

191 return [] 

192 

193 def _get_full_content( 

194 self, relevant_items: List[Dict[str, Any]] 

195 ) -> List[Dict[str, Any]]: 

196 """ 

197 Get full content for the relevant search results. 

198 Extends base implementation to include Tavily's raw_content. 

199 

200 Args: 

201 relevant_items: List of relevant preview dictionaries 

202 

203 Returns: 

204 List of result dictionaries with full content if available 

205 """ 

206 results = super()._get_full_content(relevant_items) 

207 

208 # If Tavily provided raw_content and full content is requested, use it 

209 if self.include_full_content: 

210 for result in results: 

211 if "raw_content" in result: 211 ↛ 210line 211 didn't jump to line 210 because the condition on line 211 was always true

212 result["content"] = result.get( 

213 "raw_content", result.get("content", "") 

214 ) 

215 

216 return results