Coverage for open_webui/retrieval/web/perplexity.py: 23%

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 

7 

8MODELS = Literal[ 

9 'sonar', 

10 'sonar-pro', 

11 'sonar-reasoning', 

12 'sonar-reasoning-pro', 

13 'sonar-deep-research', 

14] 

15SEARCH_CONTEXT_USAGE_LEVELS = Literal['low', 'medium', 'high'] 

16 

17 

18log = logging.getLogger(__name__) 

19 

20 

21def search_perplexity( 

22 api_key: str, 

23 query: str, 

24 count: int, 

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

26 model: MODELS = 'sonar', 

27 search_context_usage: SEARCH_CONTEXT_USAGE_LEVELS = 'medium', 

28) -> list[SearchResult]: 

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

30 

31 Args: 

32 api_key (str): A Perplexity API key 

33 query (str): The query to search for 

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

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

36 model (str): The Perplexity model to use (sonar, sonar-pro) 

37 search_context_usage (str): Search context usage level (low, medium, high) 

38 

39 """ 

40 

41 api_key = str(api_key) 

42 

43 try: 

44 url = 'https://api.perplexity.ai/chat/completions' 

45 

46 # Create payload for the API call 

47 payload = { 

48 'model': model, 

49 'messages': [ 

50 { 

51 'role': 'system', 

52 'content': 'You are a search assistant. Provide factual information with citations.', 

53 }, 

54 {'role': 'user', 'content': query}, 

55 ], 

56 'temperature': 0.2, # Lower temperature for more factual responses 

57 'stream': False, 

58 'web_search_options': { 

59 'search_context_usage': search_context_usage, 

60 }, 

61 } 

62 

63 headers = { 

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

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

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

67 } 

68 

69 # Make the API request 

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

71 

72 # Parse the JSON response 

73 json_response = response.json() 

74 

75 # Extract citations from the response 

76 citations = json_response.get('citations', []) 

77 

78 # Create search results from citations 

79 results = [] 

80 for i, citation in enumerate(citations[:count]): 

81 # Extract content from the response to use as snippet 

82 content = '' 

83 if 'choices' in json_response and json_response['choices']: 

84 if i == 0: 

85 content = json_response['choices'][0]['message']['content'] 

86 

87 result = {'link': citation, 'title': f'Source {i + 1}', 'snippet': content} 

88 results.append(result) 

89 

90 if filter_list: 

91 results = get_filtered_results(results, filter_list) 

92 

93 return [ 

94 SearchResult(link=result['link'], title=result['title'], snippet=result['snippet']) 

95 for result in results[:count] 

96 ] 

97 

98 except Exception as e: 

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

100 return []