Coverage for open_webui/retrieval/web/firecrawl.py: 15%

1from __future__ import annotations 

2 

3import logging 

4import time 

5from typing import TYPE_CHECKING, Any 

6 

7import requests 

8from langchain_core.documents import Document 

9 

10if TYPE_CHECKING: 10 ↛ 11line 10 didn't jump to line 11 because the condition on line 10 was never true

11 from open_webui.retrieval.web.main import SearchResult 

12 

13log = logging.getLogger(__name__) 

14 

15DEFAULT_FIRECRAWL_API_BASE_URL = 'https://api.firecrawl.dev' 

16FIRECRAWL_RETRY_STATUS_CODES = {429, 500, 502, 503, 504} 

17FIRECRAWL_MAX_RETRIES = 2 

18 

19 

20def build_firecrawl_url(base_url: str | None, path: str) -> str: 

21 base_url = (base_url or DEFAULT_FIRECRAWL_API_BASE_URL).rstrip('/') 

22 path = path.lstrip('/') 

23 

24 if base_url.endswith('/v2'): 

25 return f'{base_url}/{path}' 

26 

27 return f'{base_url}/v2/{path}' 

28 

29 

30def build_firecrawl_headers(api_key: str | None) -> dict[str, str]: 

31 headers = {'Content-Type': 'application/json'} 

32 if api_key: 

33 headers['Authorization'] = f'Bearer {api_key}' 

34 return headers 

35 

36 

37def get_firecrawl_timeout_seconds(timeout: Any) -> float | None: 

38 if timeout in (None, ''): 

39 return None 

40 

41 try: 

42 timeout = float(timeout) 

43 except (TypeError, ValueError): 

44 return None 

45 

46 return timeout if timeout > 0 else None 

47 

48 

49def get_firecrawl_scrape_timeout_ms(timeout: Any) -> int | None: 

50 timeout_seconds = get_firecrawl_timeout_seconds(timeout) 

51 if timeout_seconds is None: 

52 return None 

53 

54 # Firecrawl v2 expects scrape timeouts in milliseconds. 

55 return min(300000, max(1000, int(timeout_seconds * 1000))) 

56 

57 

58def get_firecrawl_client_timeout_seconds(timeout: Any, fallback: float = 60) -> float: 

59 # Keep the local HTTP timeout slightly above Firecrawl's scrape timeout. 

60 return (get_firecrawl_timeout_seconds(timeout) or fallback) + 10 

61 

62 

63def get_firecrawl_retry_delay(headers: Any, attempt: int) -> float: 

64 retry_after = headers.get('Retry-After') if headers else None 

65 if retry_after: 

66 try: 

67 return min(10.0, max(0.0, float(retry_after))) 

68 except (TypeError, ValueError): 

69 pass 

70 

71 return min(8.0, float(2**attempt)) 

72 

73 

74def request_firecrawl_json( 

75 method: str, 

76 url: str, 

77 *, 

78 headers: dict[str, str], 

79 json: dict[str, Any] | None = None, 

80 timeout: float | None = None, 

81 verify: bool = True, 

82) -> dict[str, Any]: 

83 last_error = None 

84 

85 for attempt in range(FIRECRAWL_MAX_RETRIES + 1): 

86 try: 

87 response = requests.request( 

88 method, 

89 url, 

90 headers=headers, 

91 json=json, 

92 timeout=timeout, 

93 verify=verify, 

94 ) 

95 

96 if response.status_code in FIRECRAWL_RETRY_STATUS_CODES and attempt < FIRECRAWL_MAX_RETRIES: 

97 delay = get_firecrawl_retry_delay(response.headers, attempt) 

98 log.warning( 

99 'Firecrawl %s %s returned HTTP %s; retrying in %.1fs', 

100 method, 

101 url, 

102 response.status_code, 

103 delay, 

104 ) 

105 time.sleep(delay) 

106 continue 

107 

108 response.raise_for_status() 

109 return response.json() 

110 except (requests.ConnectionError, requests.Timeout) as e: 

111 last_error = e 

112 if attempt >= FIRECRAWL_MAX_RETRIES: 

113 break 

114 

115 delay = get_firecrawl_retry_delay(None, attempt) 

116 log.warning('Firecrawl %s %s failed; retrying in %.1fs: %s', method, url, delay, e) 

117 time.sleep(delay) 

118 

119 if last_error: 

120 raise last_error 

121 

122 raise RuntimeError(f'Firecrawl {method} {url} failed without a response') 

123 

124 

125def get_firecrawl_result_url(result: dict[str, Any]) -> str: 

126 metadata = result.get('metadata') or {} 

127 return ( 

128 result.get('url') 

129 or result.get('link') 

130 or metadata.get('url') 

131 or metadata.get('sourceURL') 

132 or metadata.get('source_url') 

133 or '' 

134 ) 

135 

136 

137def scrape_firecrawl_url( 

138 firecrawl_url: str, 

139 firecrawl_api_key: str, 

140 url: str, 

141 *, 

142 verify_ssl: bool = True, 

143 timeout: Any = None, 

144 params: dict[str, Any] | None = None, 

145) -> Document | None: 

146 payload = { 

147 'url': url, 

148 'formats': ['markdown'], 

149 'skipTlsVerification': not verify_ssl, 

150 'removeBase64Images': True, 

151 **(params or {}), 

152 } 

153 scrape_timeout_ms = get_firecrawl_scrape_timeout_ms(timeout) 

154 if scrape_timeout_ms is not None: 

155 payload['timeout'] = scrape_timeout_ms 

156 

157 response = request_firecrawl_json( 

158 'POST', 

159 build_firecrawl_url(firecrawl_url, 'scrape'), 

160 headers=build_firecrawl_headers(firecrawl_api_key), 

161 json=payload, 

162 timeout=get_firecrawl_client_timeout_seconds(timeout), 

163 verify=verify_ssl, 

164 ) 

165 data = response.get('data') or {} 

166 content = data.get('markdown') or '' 

167 if not isinstance(content, str) or not content.strip(): 

168 return None 

169 

170 metadata = data.get('metadata') or {} 

171 document_metadata = {'source': get_firecrawl_result_url(data) or url} 

172 if metadata.get('title'): 

173 document_metadata['title'] = metadata['title'] 

174 if metadata.get('description'): 

175 document_metadata['description'] = metadata['description'] 

176 

177 return Document(page_content=content, metadata=document_metadata) 

178 

179 

180def search_firecrawl( 

181 firecrawl_url: str, 

182 firecrawl_api_key: str, 

183 query: str, 

184 count: int, 

185 filter_list: list[str] | None = None, 

186) -> list[SearchResult]: 

187 try: 

188 response = request_firecrawl_json( 

189 'POST', 

190 build_firecrawl_url(firecrawl_url, 'search'), 

191 headers=build_firecrawl_headers(firecrawl_api_key), 

192 json={ 

193 'query': query, 

194 'limit': count, 

195 'timeout': count * 3000, 

196 'ignoreInvalidURLs': True, 

197 }, 

198 timeout=count * 3 + 10, 

199 ) 

200 # Firecrawl /search has historically returned both `{"data": [...]}` 

201 # (flat list, what v1 did and what frost19k reported under #23966) 

202 # and `{"data": {"web": [...]}}` (current v2). Accept either. 

203 data = response.get('data') or {} 

204 results = data if isinstance(data, list) else (data.get('web') or []) 

205 

206 if filter_list: 

207 from open_webui.retrieval.web.main import get_filtered_results 

208 

209 results = get_filtered_results(results, filter_list) 

210 

211 from open_webui.retrieval.web.main import SearchResult 

212 

213 search_results = [] 

214 for result in results[:count]: 

215 url = get_firecrawl_result_url(result) 

216 if not url: 

217 continue 

218 

219 metadata = result.get('metadata') or {} 

220 search_results.append( 

221 SearchResult( 

222 link=url, 

223 title=result.get('title') or metadata.get('title'), 

224 snippet=result.get('description') or result.get('snippet') or metadata.get('description'), 

225 ) 

226 ) 

227 

228 log.info('FireCrawl search results: %s', search_results) 

229 return search_results 

230 except Exception as e: 

231 log.error(f'Error in FireCrawl search: {e}') 

232 return []