Coverage for open_webui/retrieval/web/ollama.py: 29%

1import logging 

2from dataclasses import dataclass 

3from typing import Optional 

4 

5import requests 

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

7 

8log = logging.getLogger(__name__) 

9 

10 

11def search_ollama_cloud( 

12 url: str, 

13 api_key: str, 

14 query: str, 

15 count: int, 

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

17) -> list[SearchResult]: 

18 """Search using Ollama Search API and return the results as a list of SearchResult objects. 

19 

20 Args: 

21 api_key (str): A Ollama Search API key 

22 query (str): The query to search for 

23 count (int): Number of results to return 

24 filter_list (Optional[list[str]]): List of domains to filter results by 

25 """ 

26 log.info('Searching with Ollama for query: %s', query) 

27 

28 headers = {'Authorization': f'Bearer {api_key}', 'Content-Type': 'application/json'} 

29 payload = {'query': query, 'max_results': count} 

30 

31 try: 

32 response = requests.post(f'{url}/api/web_search', headers=headers, json=payload) 

33 response.raise_for_status() 

34 data = response.json() 

35 

36 results = data.get('results', []) 

37 log.info('Found %s results', len(results)) 

38 

39 if filter_list: 

40 results = get_filtered_results(results, filter_list) 

41 

42 return [ 

43 SearchResult( 

44 link=result.get('url', ''), 

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

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

47 ) 

48 for result in results 

49 ] 

50 except Exception as e: 

51 log.error(f'Error searching Ollama: {e}') 

52 return []