#!/usr/bin/env python3
|
"""
|
Tavily AI Search API Client
|
Usage: python tavily_search.py "your search query" [--max-results 5] [--depth basic|advanced]
|
"""
|
|
import os
|
import sys
|
import json
|
import argparse
|
from typing import List, Dict, Any, Optional
|
|
|
def get_api_key(api_key: Optional[str] = None) -> str:
|
"""Get Tavily API key from environment variable or parameter."""
|
# Priority: parameter > environment variable
|
if api_key:
|
return api_key
|
|
env_key = os.environ.get("TAVILY_API_KEY")
|
if env_key:
|
return env_key
|
|
raise ValueError(
|
"Tavily API key required. Set via:\n"
|
" 1. Environment variable: export TAVILY_API_KEY=your_key\n"
|
" (Add to ~/.bashrc or ~/.zshrc for persistence)\n"
|
" 2. Direct parameter: pass api_key when calling the function\n"
|
"\nGet your API key at: https://tavily.com"
|
)
|
|
|
def tavily_search(
|
query: str,
|
max_results: int = 5,
|
search_depth: str = "basic",
|
include_answer: bool = False,
|
include_images: bool = False,
|
api_key: Optional[str] = None
|
) -> Dict[str, Any]:
|
"""
|
Search using Tavily AI Search API.
|
|
Args:
|
query: Search query string
|
max_results: Number of results to return (1-20)
|
search_depth: "basic" or "advanced"
|
include_answer: Include AI-generated answer
|
include_images: Include image URLs
|
api_key: Tavily API key (optional, auto-detected from env/config)
|
|
Returns:
|
Dictionary containing search results
|
"""
|
api_key = get_api_key(api_key)
|
|
try:
|
import requests
|
except ImportError:
|
raise ImportError("requests package required. Install with: pip install requests")
|
|
url = "https://api.tavily.com/search"
|
|
payload = {
|
"api_key": api_key,
|
"query": query,
|
"max_results": min(max(max_results, 1), 20),
|
"search_depth": search_depth,
|
"include_answer": include_answer,
|
"include_images": include_images,
|
}
|
|
response = requests.post(url, json=payload, timeout=30)
|
response.raise_for_status()
|
|
return response.json()
|
|
|
def format_results(results: Dict[str, Any]) -> str:
|
"""Format search results for display."""
|
output = []
|
|
if "answer" in results and results["answer"]:
|
output.append("=" * 60)
|
output.append("AI ANSWER")
|
output.append("=" * 60)
|
output.append(results["answer"])
|
output.append("")
|
|
output.append("=" * 60)
|
output.append("SEARCH RESULTS")
|
output.append("=" * 60)
|
output.append(f"Query: {results.get('query', 'N/A')}")
|
output.append("")
|
|
for i, result in enumerate(results.get("results", []), 1):
|
output.append(f"{i}. {result.get('title', 'No title')}")
|
output.append(f" URL: {result.get('url', 'N/A')}")
|
|
if result.get('published_date'):
|
output.append(f" Published: {result['published_date']}")
|
|
if result.get('score'):
|
output.append(f" Relevance: {result['score']:.2f}")
|
|
content = result.get('content', '')
|
if content:
|
# Truncate long content
|
if len(content) > 300:
|
content = content[:297] + "..."
|
output.append(f" {content}")
|
|
output.append("")
|
|
return "\n".join(output)
|
|
|
def main():
|
parser = argparse.ArgumentParser(description="Tavily AI Search")
|
parser.add_argument("query", help="Search query")
|
parser.add_argument("--max-results", type=int, default=5, help="Number of results (1-20)")
|
parser.add_argument("--depth", choices=["basic", "advanced"], default="basic", help="Search depth")
|
parser.add_argument("--answer", action="store_true", help="Include AI-generated answer")
|
parser.add_argument("--images", action="store_true", help="Include images")
|
parser.add_argument("--json", action="store_true", help="Output raw JSON")
|
|
args = parser.parse_args()
|
|
try:
|
results = tavily_search(
|
query=args.query,
|
max_results=args.max_results,
|
search_depth=args.depth,
|
include_answer=args.answer,
|
include_images=args.images
|
)
|
|
if args.json:
|
print(json.dumps(results, indent=2, ensure_ascii=False))
|
else:
|
print(format_results(results))
|
|
except ValueError as e:
|
print(f"Error: {e}", file=sys.stderr)
|
sys.exit(1)
|
except Exception as e:
|
print(f"Search failed: {e}", file=sys.stderr)
|
sys.exit(1)
|
|
|
if __name__ == "__main__":
|
main()
|
