Initial commit
This commit is contained in:
Zhongwei Li
277
skills/perplexity-search/scripts/perplexity_search.py
Executable file
277
skills/perplexity-search/scripts/perplexity_search.py
Executable file
@@ -0,0 +1,277 @@
|
||||
#!/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())
|
||||
171
skills/perplexity-search/scripts/setup_env.py
Executable file
171
skills/perplexity-search/scripts/setup_env.py
Executable file
@@ -0,0 +1,171 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Setup script for Perplexity Search environment configuration.
|
||||
|
||||
This script helps users configure their OpenRouter API key and validates the setup.
|
||||
|
||||
Usage:
|
||||
python setup_env.py [--api-key YOUR_KEY] [--env-file .env]
|
||||
|
||||
Author: Scientific Skills
|
||||
License: MIT
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def create_env_file(api_key: str, env_file: str = ".env") -> bool:
|
||||
"""
|
||||
Create or update .env file with OpenRouter API key.
|
||||
|
||||
Args:
|
||||
api_key: The OpenRouter API key
|
||||
env_file: Path to .env file (default: .env)
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
env_path = Path(env_file)
|
||||
|
||||
# Read existing content if file exists
|
||||
existing_content = []
|
||||
if env_path.exists():
|
||||
with open(env_path, 'r') as f:
|
||||
existing_content = [
|
||||
line for line in f.readlines()
|
||||
if not line.startswith('OPENROUTER_API_KEY=')
|
||||
]
|
||||
|
||||
# Write new content
|
||||
with open(env_path, 'w') as f:
|
||||
# Write existing content (excluding old OPENROUTER_API_KEY)
|
||||
f.writelines(existing_content)
|
||||
|
||||
# Add new API key
|
||||
if existing_content and not existing_content[-1].endswith('\n'):
|
||||
f.write('\n')
|
||||
f.write(f'OPENROUTER_API_KEY={api_key}\n')
|
||||
|
||||
print(f"✓ API key saved to {env_file}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error creating .env file: {e}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
|
||||
def validate_setup() -> bool:
|
||||
"""
|
||||
Validate that the environment is properly configured.
|
||||
|
||||
Returns:
|
||||
True if setup is valid, False otherwise
|
||||
"""
|
||||
print("Validating setup...")
|
||||
print()
|
||||
|
||||
# Check for API key
|
||||
api_key = os.environ.get("OPENROUTER_API_KEY")
|
||||
if not api_key:
|
||||
print("✗ OPENROUTER_API_KEY environment variable not set")
|
||||
print()
|
||||
print("To set up your API key:")
|
||||
print("1. Get an API key from https://openrouter.ai/keys")
|
||||
print("2. Run this script with --api-key flag:")
|
||||
print(" python setup_env.py --api-key YOUR_KEY")
|
||||
print()
|
||||
return False
|
||||
else:
|
||||
# Mask the key for display
|
||||
masked_key = api_key[:8] + "..." + api_key[-4:] if len(api_key) > 12 else "***"
|
||||
print(f"✓ OPENROUTER_API_KEY is set ({masked_key})")
|
||||
|
||||
# Check for LiteLLM
|
||||
try:
|
||||
import litellm
|
||||
print(f"✓ LiteLLM is installed (version {litellm.__version__})")
|
||||
except ImportError:
|
||||
print("✗ LiteLLM is not installed")
|
||||
print()
|
||||
print("Install LiteLLM with:")
|
||||
print(" uv pip install litellm")
|
||||
print()
|
||||
return False
|
||||
|
||||
print()
|
||||
print("✓ Setup is complete! You're ready to use Perplexity Search.")
|
||||
return True
|
||||
|
||||
|
||||
def main():
|
||||
"""Main entry point for the setup script."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Setup Perplexity Search environment configuration",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
# Set up API key
|
||||
python setup_env.py --api-key sk-or-v1-xxxxx
|
||||
|
||||
# Validate existing setup
|
||||
python setup_env.py --validate
|
||||
|
||||
# Use custom .env file location
|
||||
python setup_env.py --api-key sk-or-v1-xxxxx --env-file /path/to/.env
|
||||
|
||||
Get your OpenRouter API key from:
|
||||
https://openrouter.ai/keys
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--api-key",
|
||||
help="Your OpenRouter API key"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--env-file",
|
||||
default=".env",
|
||||
help="Path to .env file (default: .env)"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--validate",
|
||||
action="store_true",
|
||||
help="Validate existing setup"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# If no arguments, show validation
|
||||
if not args.api_key and not args.validate:
|
||||
args.validate = True
|
||||
|
||||
# Handle API key setup
|
||||
if args.api_key:
|
||||
print("Setting up OpenRouter API key...")
|
||||
if create_env_file(args.api_key, args.env_file):
|
||||
print()
|
||||
print("Next steps:")
|
||||
print(f"1. Load the environment variables:")
|
||||
print(f" source {args.env_file}")
|
||||
print("2. Or export directly:")
|
||||
print(f" export OPENROUTER_API_KEY={args.api_key}")
|
||||
print("3. Test the setup:")
|
||||
print(" python perplexity_search.py --check-setup")
|
||||
print()
|
||||
|
||||
# Validate setup
|
||||
if args.validate:
|
||||
if not validate_setup():
|
||||
return 1
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user