Coverage for open_webui/retrieval/web/jina_search.py: 35%

1from __future__ import annotations 

2 

3import logging 

4 

5import requests 

6from open_webui.retrieval.web.main import SearchResult 

7from yarl import URL 

8 

9log = logging.getLogger(__name__) 

10 

11 

12def search_jina(api_key: str, query: str, count: int, base_url: str = '') -> list[SearchResult]: 

13 """ 

14 Search using Jina's Search API and return the results as a list of SearchResult objects. 

15 Args: 

16 api_key (str): The Jina API key 

17 query (str): The query to search for 

18 count (int): The number of results to return 

19 base_url (str): Optional custom base URL for the Jina API 

20 

21 Returns: 

22 A list of SearchResult objects. 

23 """ 

24 jina_search_endpoint = base_url or 'https://s.jina.ai/' 

25 

26 headers = { 

27 'Accept': 'application/json', 

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

29 'Authorization': api_key, 

30 'X-Retain-Images': 'none', 

31 } 

32 

33 payload = {'q': query, 'count': count if count <= 10 else 10} 

34 

35 url = str(URL(jina_search_endpoint)) 

36 response = requests.post(url, headers=headers, json=payload) 

37 response.raise_for_status() 

38 data = response.json() 

39 

40 results = [] 

41 for result in data['data']: 

42 results.append( 

43 SearchResult( 

44 link=result['url'], 

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

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

47 ) 

48 ) 

49 

50 return results