Coverage for open_webui/retrieval/web/perplexity_search.py: 32%

1import logging 

2from typing import Literal, Optional 

3 

4import requests 

5from open_webui.env import VERSION 

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

7from open_webui.utils.headers import include_user_info_headers 

8 

9log = logging.getLogger(__name__) 

10 

11 

12def search_perplexity_search( 

13 api_key: str, 

14 query: str, 

15 count: int, 

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

17 api_url: str = 'https://api.perplexity.ai/search', 

18 user=None, 

19) -> list[SearchResult]: 

20 """Search using Perplexity API and return the results as a list of SearchResult objects. 

21 

22 Args: 

23 api_key (str): A Perplexity API key 

24 query (str): The query to search for 

25 count (int): Maximum number of results to return 

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

27 api_url (str): Custom API URL (defaults to https://api.perplexity.ai/search) 

28 user: Optional user object for forwarding user info headers 

29 

30 """ 

31 

32 api_key = str(api_key) 

33 api_url = str(api_url) 

34 

35 try: 

36 url = api_url 

37 

38 # Create payload for the API call 

39 payload = { 

40 'query': query, 

41 'max_results': count, 

42 } 

43 

44 headers = { 

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

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

47 'X-Pplx-Integration': f'open-webui/{VERSION}', 

48 } 

49 

50 # Forward user info headers if user is provided 

51 if user is not None: 

52 headers = include_user_info_headers(headers, user) 

53 

54 # Make the API request 

55 response = requests.request('POST', url, json=payload, headers=headers) 

56 # Parse the JSON response 

57 json_response = response.json() 

58 

59 # Extract citations from the response 

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

61 

62 return [ 

63 SearchResult(link=result['url'], title=result['title'], snippet=result['snippet']) for result in results 

64 ] 

65 

66 except Exception as e: 

67 log.error(f'Error searching with Perplexity Search API: {e}') 

68 return []