backend/open_webui/retrieval/web/perplexity.py

import logging
from typing import Literal, Optional

import requests
from open_webui.env import VERSION
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 ConfigVar 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',
            'X-Pplx-Integration': f'open-webui/{VERSION}',
        }

        # 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 []