#!/usr/bin/env python3
"""
Perplexity Search via LitLLM and OpenRouter
This script performs AI-powered web searches using Perplexity models through
LiteLLM and OpenRouter. It provides real-time, grounded answers with source citations.
Usage:
python perplexity_search.py "search query" [options]
Requirements:
- OpenRouter API key set in OPENROUTER_API_KEY environment variable
- LiteLLM installed: uv pip install litellm
Author: Scientific Skills
License: MIT
"""
import os
import sys
import json
import argparse
from typing import Optional, Dict, Any, List
def check_dependencies():
"""Check if required packages are installed."""
try:
import litellm
return True
except ImportError:
print("Error: LiteLLM is not installed.", file=sys.stderr)
print("Install it with: uv pip install litellm", file=sys.stderr)
return False
def check_api_key() -> Optional[str]:
"""Check if OpenRouter API key is configured."""
api_key = os.environ.get("OPENROUTER_API_KEY")
if not api_key:
print("Error: OPENROUTER_API_KEY environment variable is not set.", file=sys.stderr)
print("\nTo set up your API key:", file=sys.stderr)
print("1. Get an API key from https://openrouter.ai/keys", file=sys.stderr)
print("2. Set the environment variable:", file=sys.stderr)
print(" export OPENROUTER_API_KEY='your-api-key-here'", file=sys.stderr)
print("\nOr create a .env file with:", file=sys.stderr)
print(" OPENROUTER_API_KEY=your-api-key-here", file=sys.stderr)
return None
return api_key
def search_with_perplexity(
query: str,
model: str = "openrouter/perplexity/sonar-pro",
max_tokens: int = 4000,
temperature: float = 0.2,
verbose: bool = False
) -> Dict[str, Any]:
"""
Perform a search using Perplexity models via LiteLLM and OpenRouter.
Args:
query: The search query
model: Model to use (default: sonar-pro)
max_tokens: Maximum tokens in response
temperature: Response temperature (0.0-1.0)
verbose: Print detailed information
Returns:
Dictionary containing the search results and metadata
"""
try:
from litellm import completion
except ImportError:
return {
"success": False,
"error": "LiteLLM not installed. Run: uv pip install litellm"
}
# Check API key
api_key = check_api_key()
if not api_key:
return {
"success": False,
"error": "OpenRouter API key not configured"
}
if verbose:
print(f"Model: {model}", file=sys.stderr)
print(f"Query: {query}", file=sys.stderr)
print(f"Max tokens: {max_tokens}", file=sys.stderr)
print(f"Temperature: {temperature}", file=sys.stderr)
print("", file=sys.stderr)
try:
# Perform the search using LiteLLM
response = completion(
model=model,
messages=[{
"role": "user",
"content": query
}],
max_tokens=max_tokens,
temperature=temperature
)
# Extract the response
result = {
"success": True,
"query": query,
"model": model,
"answer": response.choices[0].message.content,
"usage": {
"prompt_tokens": response.usage.prompt_tokens,
"completion_tokens": response.usage.completion_tokens,
"total_tokens": response.usage.total_tokens
}
}
# Check if citations are available in the response
if hasattr(response.choices[0].message, 'citations'):
result["citations"] = response.choices[0].message.citations
return result
except Exception as e:
return {
"success": False,
"error": str(e),
"query": query,
"model": model
}
def main():
"""Main entry point for the script."""
parser = argparse.ArgumentParser(
description="Perform AI-powered web searches using Perplexity via LiteLLM and OpenRouter",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Basic search
python perplexity_search.py "What are the latest developments in CRISPR?"
# Use Sonar Pro Search for deeper analysis
python perplexity_search.py "Compare mRNA and viral vector vaccines" --model sonar-pro-search
# Use Sonar Reasoning for complex queries
python perplexity_search.py "Explain quantum entanglement" --model sonar-reasoning-pro
# Save output to file
python perplexity_search.py "COVID-19 vaccine efficacy studies" --output results.json
# Verbose mode
python perplexity_search.py "Machine learning trends 2024" --verbose
Available Models:
- sonar-pro (default): General-purpose search with good balance
- sonar-pro-search: Most advanced agentic search with multi-step reasoning
- sonar: Standard model for basic searches
- sonar-reasoning-pro: Advanced reasoning capabilities
- sonar-reasoning: Basic reasoning model
"""
)
parser.add_argument(
"query",
help="The search query"
)
parser.add_argument(
"--model",
default="sonar-pro",
choices=[
"sonar-pro",
"sonar-pro-search",
"sonar",
"sonar-reasoning-pro",
"sonar-reasoning"
],
help="Perplexity model to use (default: sonar-pro)"
)
parser.add_argument(
"--max-tokens",
type=int,
default=4000,
help="Maximum tokens in response (default: 4000)"
)
parser.add_argument(
"--temperature",
type=float,
default=0.2,
help="Response temperature 0.0-1.0 (default: 0.2)"
)
parser.add_argument(
"--output",
help="Save results to JSON file"
)
parser.add_argument(
"--verbose",
action="store_true",
help="Print detailed information"
)
parser.add_argument(
"--check-setup",
action="store_true",
help="Check if dependencies and API key are configured"
)
args = parser.parse_args()
# Check setup if requested
if args.check_setup:
print("Checking setup...")
deps_ok = check_dependencies()
api_key_ok = check_api_key() is not None
if deps_ok and api_key_ok:
print("\nā Setup complete! Ready to search.")
return 0
else:
print("\nā Setup incomplete. Please fix the issues above.")
return 1
# Check dependencies
if not check_dependencies():
return 1
# Prepend openrouter/ to model name if not already present
model = args.model
if not model.startswith("openrouter/"):
model = f"openrouter/perplexity/{model}"
# Perform the search
result = search_with_perplexity(
query=args.query,
model=model,
max_tokens=args.max_tokens,
temperature=args.temperature,
verbose=args.verbose
)
# Handle results
if not result["success"]:
print(f"Error: {result['error']}", file=sys.stderr)
return 1
# Print answer
print("\n" + "="*80)
print("ANSWER")
print("="*80)
print(result["answer"])
print("="*80)
# Print usage stats if verbose
if args.verbose:
print(f"\nUsage:", file=sys.stderr)
print(f" Prompt tokens: {result['usage']['prompt_tokens']}", file=sys.stderr)
print(f" Completion tokens: {result['usage']['completion_tokens']}", file=sys.stderr)
print(f" Total tokens: {result['usage']['total_tokens']}", file=sys.stderr)
# Save to file if requested
if args.output:
with open(args.output, 'w') as f:
json.dump(result, f, indent=2)
print(f"\nā Results saved to {args.output}", file=sys.stderr)
return 0
if __name__ == "__main__":
sys.exit(main())