mirror of
https://github.com/open-webui/open-webui.git
synced 2026-05-03 18:59:38 -05:00
101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
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 []
|