Spaces:
Build error
Build error
File size: 3,162 Bytes
87a665c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | import logging
from typing import Optional, Literal
import requests
from open_webui.retrieval.web.main import SearchResult, get_filtered_results
MODELS = Literal[
'sonar',
'sonar-pro',
'sonar-reasoning',
'sonar-reasoning-pro',
'sonar-deep-research',
]
SEARCH_CONTEXT_USAGE_LEVELS = Literal['low', 'medium', 'high']
log = logging.getLogger(__name__)
def search_perplexity(
api_key: str,
query: str,
count: int,
filter_list: Optional[list[str]] = None,
model: MODELS = 'sonar',
search_context_usage: SEARCH_CONTEXT_USAGE_LEVELS = 'medium',
) -> list[SearchResult]:
"""Search using Perplexity API and return the results as a list of SearchResult objects.
Args:
api_key (str): A Perplexity API key
query (str): The query to search for
count (int): Maximum number of results to return
filter_list (Optional[list[str]]): List of domains to filter results
model (str): The Perplexity model to use (sonar, sonar-pro)
search_context_usage (str): Search context usage level (low, medium, high)
"""
# Handle PersistentConfig object
if hasattr(api_key, '__str__'):
api_key = str(api_key)
try:
url = 'https://api.perplexity.ai/chat/completions'
# Create payload for the API call
payload = {
'model': model,
'messages': [
{
'role': 'system',
'content': 'You are a search assistant. Provide factual information with citations.',
},
{'role': 'user', 'content': query},
],
'temperature': 0.2, # Lower temperature for more factual responses
'stream': False,
'web_search_options': {
'search_context_usage': search_context_usage,
},
}
headers = {
'Authorization': f'Bearer {api_key}',
'Content-Type': 'application/json',
}
# Make the API request
response = requests.request('POST', url, json=payload, headers=headers)
# Parse the JSON response
json_response = response.json()
# Extract citations from the response
citations = json_response.get('citations', [])
# Create search results from citations
results = []
for i, citation in enumerate(citations[:count]):
# Extract content from the response to use as snippet
content = ''
if 'choices' in json_response and json_response['choices']:
if i == 0:
content = json_response['choices'][0]['message']['content']
result = {'link': citation, 'title': f'Source {i + 1}', 'snippet': content}
results.append(result)
if filter_list:
results = get_filtered_results(results, filter_list)
return [
SearchResult(link=result['link'], title=result['title'], snippet=result['snippet'])
for result in results[:count]
]
except Exception as e:
log.error(f'Error searching with Perplexity API: {e}')
return []
|