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

1from typing import Any, Dict, List, Optional 

2 

3from langchain_community.tools import BraveSearch 

4from langchain_core.language_models import BaseLLM 

5 

6from ...security.secure_logging import logger 

7 

8from ..rate_limiting import RateLimitError 

9from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity 

10 

11 

12class BraveSearchEngine(BaseSearchEngine): 

13 """Brave 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 # secrets to redact from error messages (see BaseSearchEngine._scrub_error) 

22 _secret_attrs = ("_brave_api_key",) 

23 

24 def __init__( 

25 self, 

26 max_results: int = 10, 

27 region: str = "US", 

28 time_period: str = "y", 

29 safe_search: bool = True, 

30 search_language: str = "English", 

31 api_key: Optional[str] = None, 

32 language_code_mapping: Optional[Dict[str, str]] = None, 

33 llm: Optional[BaseLLM] = None, 

34 include_full_content: bool = True, 

35 max_filtered_results: Optional[int] = None, 

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

37 **kwargs, 

38 ): 

39 """ 

40 Initialize the Brave search engine. 

41 

42 Args: 

43 max_results: Maximum number of search results 

44 region: Region code for search results 

45 time_period: Time period for search results 

46 safe_search: Whether to enable safe search 

47 search_language: Language for search results 

48 api_key: Brave Search API key (can also be set via LDR_SEARCH_ENGINE_WEB_BRAVE_API_KEY env var or in UI settings) 

49 language_code_mapping: Mapping from language names to codes 

50 llm: Language model for relevance filtering 

51 include_full_content: Whether to include full webpage content in results 

52 max_filtered_results: Maximum number of results to keep after filtering 

53 settings_snapshot: Settings snapshot for thread context 

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

55 """ 

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

57 super().__init__( 

58 llm=llm, 

59 max_filtered_results=max_filtered_results, 

60 max_results=max_results, 

61 include_full_content=include_full_content, 

62 settings_snapshot=settings_snapshot, 

63 ) 

64 

65 # Set up language code mapping 

66 if language_code_mapping is None: 

67 from ...utilities.search_utilities import LANGUAGE_CODE_MAP 

68 

69 language_code_mapping = LANGUAGE_CODE_MAP 

70 

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

72 brave_api_key = self._resolve_api_key( 

73 api_key, 

74 "search.engine.web.brave.api_key", 

75 engine_name="Brave Search", 

76 settings_snapshot=settings_snapshot, 

77 ) 

78 self._brave_api_key = brave_api_key 

79 

80 # Get language code 

81 language_code = language_code_mapping.get(search_language.lower(), "en") 

82 

83 # Convert time period format to Brave's format 

84 brave_time_period = f"p{time_period}" 

85 

86 # Convert safe search to Brave's format 

87 brave_safe_search = "moderate" if safe_search else "off" 

88 

89 # Initialize Brave Search 

90 self.engine = BraveSearch.from_api_key( 

91 api_key=brave_api_key, 

92 search_kwargs={ 

93 "count": min(20, max_results), 

94 "country": region.upper(), 

95 "search_lang": language_code, 

96 "safesearch": brave_safe_search, 

97 "freshness": brave_time_period, 

98 }, 

99 ) 

100 

101 # User agent is not needed for Brave Search API 

102 

103 # If full content is requested, initialize FullSearchResults 

104 self._init_full_search( 

105 web_search=self.engine, 

106 language=search_language, 

107 max_results=max_results, 

108 region=region, 

109 time_period=time_period, 

110 safe_search=brave_safe_search, 

111 ) 

112 

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

114 """ 

115 Get preview information from Brave Search. 

116 

117 Args: 

118 query: The search query 

119 

120 Returns: 

121 List of preview dictionaries 

122 """ 

123 logger.info("Getting search results from Brave Search") 

124 

125 try: 

126 # Get search results from Brave Search 

127 raw_results = self.engine.run(query[:400]) 

128 

129 # Parse results if they're in string format 

130 if isinstance(raw_results, str): 

131 try: 

132 import json 

133 

134 raw_results = json.loads(raw_results) 

135 except json.JSONDecodeError as e: 

136 safe_msg = self._scrub_error(e) 

137 logger.warning( 

138 f"Unable to parse BraveSearch response as JSON: {safe_msg}" 

139 ) 

140 return [] 

141 

142 # Format results as previews 

143 previews = [] 

144 for i, result in enumerate(raw_results): 

145 preview = { 

146 "id": i, # Use index as ID 

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

148 "link": result.get("link", ""), 

149 "snippet": result.get("snippet", ""), 

150 "displayed_link": result.get("link", ""), 

151 "position": i, 

152 } 

153 

154 # Store full Brave result for later 

155 preview["_full_result"] = result 

156 

157 previews.append(preview) 

158 

159 # Store the previews for potential full content retrieval 

160 self._search_results = previews 

161 

162 return previews 

163 

164 except RateLimitError: 

165 raise 

166 except Exception as e: 

167 safe_msg = self._scrub_error(e) 

168 logger.warning(f"Error getting Brave Search results: {safe_msg}") 

169 self._raise_if_rate_limit(e) 

170 return []