Coverage for open_webui/retrieval/web/tavily.py: 37%

1from __future__ import annotations 

2 

3import logging 

4 

5import requests 

6from open_webui.env import TAVILY_API_BASE_URL 

7from open_webui.retrieval.web.main import SearchResult, get_filtered_results 

8 

9log = logging.getLogger(__name__) 

10 

11 

12def search_tavily( 

13 api_key: str, 

14 query: str, 

15 count: int, 

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

17 # **kwargs, 

18) -> list[SearchResult]: 

19 """Search using Tavily's Search API and return the results as a list of SearchResult objects. 

20 

21 Args: 

22 api_key (str): A Tavily Search API key 

23 query (str): The query to search for 

24 count (int): The maximum number of results to return 

25 

26 Returns: 

27 A list of SearchResult objects. 

28 """ 

29 url = f'{TAVILY_API_BASE_URL}/search' 

30 headers = { 

31 'Content-Type': 'application/json', 

32 'Authorization': f'Bearer {api_key}', 

33 } 

34 data = {'query': query, 'max_results': count} 

35 response = requests.post(url, headers=headers, json=data) 

36 response.raise_for_status() 

37 

38 json_response = response.json() 

39 

40 results = json_response.get('results', []) 

41 if filter_list: 

42 results = get_filtered_results(results, filter_list) 

43 

44 return [ 

45 SearchResult( 

46 link=result['url'], 

47 title=result.get('title', ''), 

48 snippet=result.get('content'), 

49 ) 

50 for result in results 

51 ]