134 lines
4.2 KiB
Python
134 lines
4.2 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
Perplexity API Query Interface
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
python3 query.py "What is the capital of France?"
|
||
|
|
python3 query.py "Latest AI news" --model sonar-pro --citations
|
||
|
|
"""
|
||
|
|
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import urllib.request
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
def load_config():
|
||
|
|
"""Load API configuration"""
|
||
|
|
config_path = Path(__file__).parent.parent / "config.json"
|
||
|
|
try:
|
||
|
|
with open(config_path) as f:
|
||
|
|
return json.load(f)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error loading config: {e}", file=sys.stderr)
|
||
|
|
return None
|
||
|
|
|
||
|
|
def query_perplexity(query, model=None, max_tokens=None, include_citations=False, search_context="low"):
|
||
|
|
"""
|
||
|
|
Query Perplexity API
|
||
|
|
|
||
|
|
Args:
|
||
|
|
query: The question/prompt to send
|
||
|
|
model: Model to use (sonar, sonar-pro, sonar-reasoning, sonar-deep-research)
|
||
|
|
max_tokens: Maximum tokens in response
|
||
|
|
include_citations: Whether to include source citations
|
||
|
|
search_context: Search depth (low, medium, high)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
dict with response text, citations, and usage info
|
||
|
|
"""
|
||
|
|
config = load_config()
|
||
|
|
if not config:
|
||
|
|
return {"error": "Failed to load configuration"}
|
||
|
|
|
||
|
|
model = model or config.get("default_model", "sonar")
|
||
|
|
max_tokens = max_tokens or config.get("default_max_tokens", 1000)
|
||
|
|
api_key = config.get("api_key")
|
||
|
|
base_url = config.get("base_url", "https://api.perplexity.ai")
|
||
|
|
|
||
|
|
if not api_key:
|
||
|
|
return {"error": "API key not configured"}
|
||
|
|
|
||
|
|
payload = {
|
||
|
|
"model": model,
|
||
|
|
"messages": [
|
||
|
|
{"role": "system", "content": "Be precise and concise."},
|
||
|
|
{"role": "user", "content": query}
|
||
|
|
],
|
||
|
|
"max_tokens": max_tokens,
|
||
|
|
"search_context_size": search_context
|
||
|
|
}
|
||
|
|
|
||
|
|
data = json.dumps(payload).encode()
|
||
|
|
|
||
|
|
req = urllib.request.Request(
|
||
|
|
f"{base_url}/chat/completions",
|
||
|
|
data=data,
|
||
|
|
headers={
|
||
|
|
"Authorization": f"Bearer {api_key}",
|
||
|
|
"Content-Type": "application/json"
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
try:
|
||
|
|
with urllib.request.urlopen(req, timeout=60) as response:
|
||
|
|
result = json.loads(response.read().decode())
|
||
|
|
|
||
|
|
output = {
|
||
|
|
"text": result["choices"][0]["message"]["content"],
|
||
|
|
"model": result.get("model"),
|
||
|
|
"usage": result.get("usage", {})
|
||
|
|
}
|
||
|
|
|
||
|
|
if include_citations:
|
||
|
|
output["citations"] = result.get("citations", [])
|
||
|
|
output["search_results"] = result.get("search_results", [])
|
||
|
|
|
||
|
|
return output
|
||
|
|
|
||
|
|
except urllib.error.HTTPError as e:
|
||
|
|
error_body = e.read().decode()
|
||
|
|
return {"error": f"HTTP {e.code}: {error_body}"}
|
||
|
|
except Exception as e:
|
||
|
|
return {"error": str(e)}
|
||
|
|
|
||
|
|
def main():
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser(description="Query Perplexity API")
|
||
|
|
parser.add_argument("query", help="The query to send")
|
||
|
|
parser.add_argument("--model", default="sonar",
|
||
|
|
choices=["sonar", "sonar-pro", "sonar-reasoning", "sonar-deep-research"],
|
||
|
|
help="Model to use")
|
||
|
|
parser.add_argument("--max-tokens", type=int, default=1000,
|
||
|
|
help="Maximum tokens in response")
|
||
|
|
parser.add_argument("--citations", action="store_true",
|
||
|
|
help="Include citations in output")
|
||
|
|
parser.add_argument("--search-context", default="low",
|
||
|
|
choices=["low", "medium", "high"],
|
||
|
|
help="Search context size")
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
result = query_perplexity(
|
||
|
|
args.query,
|
||
|
|
model=args.model,
|
||
|
|
max_tokens=args.max_tokens,
|
||
|
|
include_citations=args.citations,
|
||
|
|
search_context=args.search_context
|
||
|
|
)
|
||
|
|
|
||
|
|
if "error" in result:
|
||
|
|
print(f"Error: {result['error']}", file=sys.stderr)
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
print(result["text"])
|
||
|
|
|
||
|
|
if args.citations and result.get("citations"):
|
||
|
|
print("\n--- Sources ---")
|
||
|
|
for i, citation in enumerate(result["citations"][:5], 1):
|
||
|
|
print(f"[{i}] {citation}")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|