app.py

import gradio as gr
import requests
import json
from typing import Optional, List

def tavily_search(
    api_key: str,
    query: str,
    topic: str = "general",
    search_depth: str = "basic",
    chunks_per_source: int = 3,
    max_results: int = 5,
    time_range: Optional[str] = None,
    days: int = 7,
    include_answer: bool = True,
    include_raw_content: str = "false",
    include_images: bool = False,
    include_image_descriptions: bool = False,
    include_domains: str = "",
    exclude_domains: str = "",
    country: Optional[str] = None
):
    """
    Perform a Tavily search with the given parameters.
    """
    if not api_key.strip():
        return "❌ Error: Please provide a valid Tavily API key."
    
    if not query.strip():
        return "❌ Error: Please provide a search query."
    
    url = "https://api.tavily.com/search"
    
    # Prepare the payload
    payload = {
        "query": query,
        "topic": topic,
        "search_depth": search_depth,
        "max_results": max_results,
        "include_answer": include_answer,
        "include_images": include_images,
        "include_image_descriptions": include_image_descriptions,
    }
    
    # Add optional parameters
    if search_depth == "advanced":
        payload["chunks_per_source"] = chunks_per_source
    
    if time_range and time_range != "None":
        payload["time_range"] = time_range
    
    if topic == "news" and days > 0:
        payload["days"] = days
    
    if include_raw_content != "false":
        payload["include_raw_content"] = include_raw_content if include_raw_content != "true" else True
    
    if include_domains.strip():
        payload["include_domains"] = [domain.strip() for domain in include_domains.split(",")]
    
    if exclude_domains.strip():
        payload["exclude_domains"] = [domain.strip() for domain in exclude_domains.split(",")]
    
    if country and country != "None":
        payload["country"] = country
    
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.post(url, json=payload, headers=headers, timeout=30)
        
        if response.status_code == 200:
            result = response.json()
            
            # Format the response for better readability
            formatted_result = "🔍 **Tavily Search Results**\n\n"
            
            if "answer" in result and result["answer"]:
                formatted_result += f"📋 **Answer:**\n{result['answer']}\n\n"
            
            if "results" in result:
                formatted_result += f"📊 **Found {len(result['results'])} results:**\n\n"
                
                for i, item in enumerate(result["results"], 1):
                    formatted_result += f"**{i}. {item.get('title', 'No Title')}**\n"
                    formatted_result += f"🔗 URL: {item.get('url', 'No URL')}\n"
                    formatted_result += f"📝 Content: {item.get('content', 'No content available')[:500]}...\n"
                    if item.get('score'):
                        formatted_result += f"⭐ Score: {item['score']}\n"
                    formatted_result += "\n" + "─" * 50 + "\n\n"
            
            # Add raw JSON for debugging
            formatted_result += "\n\n🔧 **Raw JSON Response:**\n```json\n" + json.dumps(result, indent=2) + "\n```"
            
            return formatted_result
            
        else:
            return f"❌ Error: HTTP {response.status_code}\n{response.text}"
            
    except requests.exceptions.Timeout:
        return "❌ Error: Request timed out. Please try again."
    except requests.exceptions.RequestException as e:
        return f"❌ Error: {str(e)}"
    except Exception as e:
        return f"❌ Unexpected error: {str(e)}"

# Define the country options
country_options = [
    "None", "afghanistan", "albania", "algeria", "andorra", "angola", "argentina", 
    "armenia", "australia", "austria", "azerbaijan", "bahamas", "bahrain", "bangladesh", 
    "barbados", "belarus", "belgium", "belize", "benin", "bhutan", "bolivia", 
    "bosnia and herzegovina", "botswana", "brazil", "brunei", "bulgaria", "burkina faso", 
    "burundi", "cambodia", "cameroon", "canada", "cape verde", "central african republic", 
    "chad", "chile", "china", "colombia", "comoros", "congo", "costa rica", "croatia", 
    "cuba", "cyprus", "czech republic", "denmark", "djibouti", "dominican republic", 
    "ecuador", "egypt", "el salvador", "equatorial guinea", "eritrea", "estonia", 
    "ethiopia", "fiji", "finland", "france", "gabon", "gambia", "georgia", "germany", 
    "ghana", "greece", "guatemala", "guinea", "haiti", "honduras", "hungary", "iceland", 
    "india", "indonesia", "iran", "iraq", "ireland", "israel", "italy", "jamaica", 
    "japan", "jordan", "kazakhstan", "kenya", "kuwait", "kyrgyzstan", "latvia", "lebanon", 
    "lesotho", "liberia", "libya", "liechtenstein", "lithuania", "luxembourg", "madagascar", 
    "malawi", "malaysia", "maldives", "mali", "malta", "mauritania", "mauritius", "mexico", 
    "moldova", "monaco", "mongolia", "montenegro", "morocco", "mozambique", "myanmar", 
    "namibia", "nepal", "netherlands", "new zealand", "nicaragua", "niger", "nigeria", 
    "north korea", "north macedonia", "norway", "oman", "pakistan", "panama", 
    "papua new guinea", "paraguay", "peru", "philippines", "poland", "portugal", "qatar", 
    "romania", "russia", "rwanda", "saudi arabia", "senegal", "serbia", "singapore", 
    "slovakia", "slovenia", "somalia", "south africa", "south korea", "south sudan", 
    "spain", "sri lanka", "sudan", "sweden", "switzerland", "syria", "taiwan", 
    "tajikistan", "tanzania", "thailand", "togo", "trinidad and tobago", "tunisia", 
    "turkey", "turkmenistan", "uganda", "ukraine", "united arab emirates", 
    "united kingdom", "united states", "uruguay", "uzbekistan", "venezuela", "vietnam", 
    "yemen", "zambia", "zimbabwe"
]

# Create the Gradio interface
with gr.Blocks(title="Tavily Search API", theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🔍 Tavily Search API Interface")
    gr.Markdown("Search the web using Tavily's powerful search API with customizable parameters.")
    
    with gr.Row():
        with gr.Column(scale=1):
            # API Configuration
            gr.Markdown("## 🔑 API Configuration")
            api_key = gr.Textbox(
                label="Tavily API Key",
                placeholder="Enter your Tavily API key here...",
                type="password",
                info="Your API key will be used to authenticate requests to Tavily."
            )
            
            # Search Parameters
            gr.Markdown("## 🎯 Search Parameters")
            query = gr.Textbox(
                label="Search Query",
                placeholder="e.g., 'who is Leo Messi?'",
                info="The search query you want to execute."
            )
            
            with gr.Row():
                topic = gr.Dropdown(
                    choices=["general", "news"],
                    value="general",
                    label="Topic",
                    info="general: broad searches, news: real-time updates"
                )
                
                search_depth = gr.Dropdown(
                    choices=["basic", "advanced"],
                    value="basic",
                    label="Search Depth",
                    info="basic: 1 credit, advanced: 2 credits"
                )
            
            with gr.Row():
                max_results = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=5,
                    step=1,
                    label="Max Results",
                    info="Maximum number of search results to return"
                )
                
                chunks_per_source = gr.Slider(
                    minimum=1,
                    maximum=3,
                    value=3,
                    step=1,
                    label="Chunks per Source",
                    info="Only applies to advanced search"
                )
            
            # Advanced Options
            gr.Markdown("## ⚙️ Advanced Options")
            
            with gr.Row():
                time_range = gr.Dropdown(
                    choices=["None", "day", "week", "month", "year"],
                    value="None",
                    label="Time Range",
                    info="Filter results by time period"
                )
                
                days = gr.Number(
                    value=7,
                    minimum=1,
                    label="Days (News only)",
                    info="Number of days back for news searches"
                )
            
            with gr.Row():
                include_answer = gr.Checkbox(
                    value=True,
                    label="Include Answer",
                    info="Include LLM-generated answer"
                )
                
                include_images = gr.Checkbox(
                    value=False,
                    label="Include Images",
                    info="Perform image search"
                )
                
                include_image_descriptions = gr.Checkbox(
                    value=False,
                    label="Include Image Descriptions",
                    info="Add descriptions for images"
                )
            
            include_raw_content = gr.Dropdown(
                choices=["false", "true", "markdown", "text"],
                value="false",
                label="Include Raw Content",
                info="Include cleaned HTML content"
            )
            
            with gr.Row():
                include_domains = gr.Textbox(
                    label="Include Domains",
                    placeholder="example.com, another.com",
                    info="Comma-separated list of domains to include"
                )
                
                exclude_domains = gr.Textbox(
                    label="Exclude Domains",
                    placeholder="example.com, another.com",
                    info="Comma-separated list of domains to exclude"
                )
            
            country = gr.Dropdown(
                choices=country_options,
                value="None",
                label="Country",
                info="Boost results from specific country (general topic only)"
            )
            
            # Search Button
            search_btn = gr.Button("🔍 Search", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            # Results
            gr.Markdown("## 📊 Search Results")
            results = gr.Textbox(
                label="Results",
                lines=25,
                max_lines=50,
                show_copy_button=True,
                placeholder="Results will appear here after searching..."
            )
    
    # Examples
    gr.Markdown("## 💡 Examples")
    gr.Examples(
        [
            ["Who is Leo Messi?", "general", "basic", 5, True],
            ["Latest news about artificial intelligence", "news", "advanced", 3, True],
            ["Python programming tutorials", "general", "basic", 10, True],
            ["Climate change recent developments", "news", "basic", 5, True],
        ],
        inputs=[query, topic, search_depth, max_results, include_answer],
        label="Click on an example to load it"
    )
    
    # Connect the search button
    search_btn.click(
        fn=tavily_search,
        inputs=[
            api_key, query, topic, search_depth, chunks_per_source, max_results,
            time_range, days, include_answer, include_raw_content, include_images,
            include_image_descriptions, include_domains, exclude_domains, country
        ],
        outputs=results
    )
    
    # Footer
    gr.Markdown("---")
    gr.Markdown("🔗 **Get your Tavily API key at:** [tavily.com](https://tavily.com)")

# Launch the app
if __name__ == "__main__":
    app.launch(
        share=True
    )