wehub-resource-sync
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
"""Built-in tool: Perplexity web search.
|
|
|
|
Requires environment variable:
|
|
- ``PERPLEXITY_API_KEY``: API key from Perplexity.
|
|
|
|
Uses the Perplexity Chat Completions API with an online model
|
|
to perform grounded web search. Returns the answer with citations.
|
|
|
|
See https://docs.perplexity.ai/
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
|
|
# Any: the OpenAI tool schema is a heterogeneous dict with string
|
|
# keys and mixed value types (str, dict, list).
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
_DEFAULT_PERPLEXITY_URL = "https://api.perplexity.ai/chat/completions"
|
|
|
|
# Perplexity model optimized for web search with citations.
|
|
_PERPLEXITY_MODEL = "sonar"
|
|
|
|
|
|
def _perplexity_url() -> str:
|
|
"""Resolve the Perplexity URL; ``OMNIGENT_PERPLEXITY_BASE_URL`` overrides for tests."""
|
|
return os.environ.get("OMNIGENT_PERPLEXITY_BASE_URL", _DEFAULT_PERPLEXITY_URL)
|
|
|
|
|
|
def _search_perplexity(
|
|
query: str,
|
|
config: dict[str, str],
|
|
) -> str:
|
|
"""
|
|
Call the Perplexity Chat Completions API with an online model.
|
|
|
|
:param query: The search query string.
|
|
:param config: Spec-level config; checked for ``api_key``
|
|
before falling back to the env var.
|
|
:returns: The answer text with citations, or an error message.
|
|
"""
|
|
api_key = config.get("api_key")
|
|
if not api_key:
|
|
return "Error: api_key must be provided in the web_search config in config.yaml."
|
|
try:
|
|
resp = httpx.post(
|
|
_perplexity_url(),
|
|
headers={
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json={
|
|
"model": _PERPLEXITY_MODEL,
|
|
"messages": [
|
|
{"role": "user", "content": query},
|
|
],
|
|
},
|
|
timeout=30.0,
|
|
)
|
|
resp.raise_for_status()
|
|
except httpx.HTTPStatusError as exc:
|
|
return f"Perplexity search error: HTTP {exc.response.status_code}"
|
|
except (httpx.ConnectError, httpx.TimeoutException) as exc:
|
|
return f"Perplexity search error: {exc}"
|
|
|
|
return _format_response(resp.json())
|
|
|
|
|
|
def _format_response(data: dict[str, Any]) -> str:
|
|
"""
|
|
Extract the answer and citations from Perplexity's response.
|
|
|
|
:param data: The parsed JSON response from Perplexity.
|
|
:returns: The answer text followed by numbered citations.
|
|
"""
|
|
choices = data.get("choices", [])
|
|
if not choices:
|
|
return "No answer returned."
|
|
message = choices[0].get("message", {})
|
|
content: str = str(message.get("content", ""))
|
|
# Perplexity includes citation URLs in the response metadata.
|
|
citations = data.get("citations", [])
|
|
if citations:
|
|
citation_lines = [f"[{i + 1}] {url}" for i, url in enumerate(citations)]
|
|
content += "\n\nSources:\n" + "\n".join(citation_lines)
|
|
return content
|