app.py

import gradio as gr
import requests
import json
import os
import time
from collections import defaultdict
from PIL import Image
import io

BASE_URL = "https://api.jigsawstack.com/v1"
headers = {
    "x-api-key": os.getenv("JIGSAWSTACK_API_KEY")
}

# Rate limiting configuration
request_times = defaultdict(list)
MAX_REQUESTS = 20  # Maximum requests per time window
TIME_WINDOW = 3600   # Time window in seconds (1 hour)

def get_real_ip(request: gr.Request):
    """Extract real IP address using x-forwarded-for header or fallback"""
    if not request:
        return "unknown"
    
    forwarded = request.headers.get("x-forwarded-for")
    if forwarded:
        ip = forwarded.split(",")[0].strip()  # First IP in the list is the client's
    else:
        ip = request.client.host  # fallback
    return ip

def check_rate_limit(request: gr.Request):
    """Check if the current request exceeds rate limits"""
    if not request:
        return True, "Rate limit check failed - no request info"
    
    ip = get_real_ip(request)
    now = time.time()

    # Clean up old timestamps outside the time window
    request_times[ip] = [t for t in request_times[ip] if now - t < TIME_WINDOW]
    

    # Check if rate limit exceeded
    if len(request_times[ip]) >= MAX_REQUESTS:
        time_remaining = int(TIME_WINDOW - (now - request_times[ip][0]))
        time_remaining_minutes = round(time_remaining / 60, 1)      
        time_window_minutes = round(TIME_WINDOW / 60, 1)
        
        return False, f"Rate limit exceeded. You can make {MAX_REQUESTS} requests per {time_window_minutes} minutes. Try again in {time_remaining_minutes} minutes."
    
    # Add current request timestamp
    request_times[ip].append(now)
    return True, ""

# ----------------- JigsawStack API Wrappers ------------------

def web_ai_search(query, ai_overview, safe_search, spell_check, deep_research, max_depth, max_breadth, max_output_tokens, target_output_tokens, request: gr.Request):
    rate_limit_ok, rate_limit_msg = check_rate_limit(request)
    if not rate_limit_ok:
        return f"❌ {rate_limit_msg}", "", [], ""
    
    if not query or not query.strip():
        return "❌ Please enter a search query.", "", [], ""
    payload = {
        "query": query.strip(),
        "ai_overview": ai_overview,
        "safe_search": safe_search,
        "spell_check": spell_check
    }
    if deep_research:
        payload["deep_research"] = True
        config = {
            "max_depth": max_depth if max_depth is not None else 3,
            "max_breadth": max_breadth if max_breadth is not None else 3,
            "max_output_tokens": max_output_tokens if max_output_tokens is not None else 32000
        }
        if target_output_tokens is not None and target_output_tokens != "":
            config["target_output_tokens"] = target_output_tokens
        payload["deep_research_config"] = config
    try:
        response = requests.post(f"{BASE_URL}/web/search", headers=headers, json=payload)
        if response.status_code != 200:
            return f"❌ Error: {response.status_code} - {response.text}", "", [], ""
        result = response.json()
        if not result.get("success"):
            return "❌ Search failed.", "", [], ""
        status = "✅ Search successful!"
        overview = result.get("ai_overview", "")
        results = result.get("results", [])
        # Format results for display
        formatted_results = []
        for r in results:
            title = r.get("title", "")
            url = r.get("url", "")
            snippet = r.get("snippet", "")
            formatted_results.append(f"{title}\n{url}\n{snippet}")
        raw_json = json.dumps(result, indent=2)
        return status, overview, formatted_results, raw_json
    except Exception as e:
        return f"❌ Error: {str(e)}", "", [], ""

with gr.Blocks() as demo:
    gr.Markdown("""
        <div style='text-align: center; margin-bottom: 24px;'>
            <h1 style='font-size:2.2em; margin-bottom: 0.2em;'>🧩 AI Search</h1>
            <p style='font-size:1.2em; margin-top: 0;'>Effortlessly search the web and get high-quality results powered by AI.</p>
            <p style='font-size:1em; margin-top: 0.5em;'>For more details and API usage, see the <a href='https://jigsawstack.com/docs/api-reference/web/ai-search' target='_blank'>documentation</a>.</p>
        </div>
        """)
    with gr.Row():
            with gr.Column():
                search_query = gr.Textbox(label="Search Query", placeholder="Type your search here...")
                ai_overview = gr.Checkbox(label="AI Overview", value=True)
                safe_search = gr.Dropdown(choices=["moderate", "strict", "off"], value="moderate", label="Safe Search")
                spell_check = gr.Checkbox(label="Spell Check", value=True)
                deep_research = gr.Checkbox(label="Deep Research", value=False)
                with gr.Group(visible=False) as deep_research_group:
                    max_depth = gr.Number(label="Max Depth", value=3, precision=0)
                    max_breadth = gr.Number(label="Max Breadth", value=3, precision=0)
                    max_output_tokens = gr.Number(label="Max Output Tokens", value=32000, precision=0)
                    target_output_tokens = gr.Number(label="Target Output Tokens (optional)", value=None, precision=0)
                search_btn = gr.Button("🔍 Search")
                search_clear_btn = gr.Button("Clear")
            with gr.Column():
                search_status = gr.Textbox(label="Status", interactive=False)
                overview_box = gr.Textbox(label="AI Overview", lines=4, interactive=False)
                results_box = gr.Dataframe(headers=["Result"], label="Results", interactive=False)
                search_json_box = gr.Accordion("Raw JSON Response", open=False)
                with search_json_box:
                    search_json_output = gr.Textbox(show_label=False, lines=20, interactive=False)
    def toggle_deep_research(checked):
            return {deep_research_group: gr.update(visible=checked)}
    deep_research.change(fn=toggle_deep_research, inputs=deep_research, outputs=deep_research_group)
    def on_search(query, overview, safe, spell, deep, d_depth, d_breadth, d_tokens, d_target, request: gr.Request):
        return web_ai_search(query, overview, safe, spell, deep, d_depth, d_breadth, d_tokens, d_target, request)
    search_btn.click(fn=on_search, inputs=[search_query, ai_overview, safe_search, spell_check, deep_research, max_depth, max_breadth, max_output_tokens, target_output_tokens],
                        outputs=[search_status, overview_box, results_box, search_json_output])
    def clear_search():
            return "", "", [], ""
    search_clear_btn.click(fn=clear_search, inputs=[], outputs=[search_query, overview_box, results_box, search_json_output])
demo.launch()