script.py

import os
import requests
import google.generativeai as genai
from dotenv import load_dotenv
from datetime import datetime

load_dotenv()

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
PERPLEXITY_API_KEY = os.getenv("PERPLEXITY_API_KEY")

if not GEMINI_API_KEY or not PERPLEXITY_API_KEY:
    raise ValueError(
        "API keys not found. Please create a .env file and add your "
        "GEMINI_API_KEY and PERPLEXITY_API_KEY."
    )
genai.configure(api_key=GEMINI_API_KEY)
gemini_model = genai.GenerativeModel('gemini-2.0-flash')

def local_quick_answer(query: str):
    q = query.strip().lower()
    if q in ("what is the date", "what's the date", "what is today's date", "today's date"):
        return {"text": datetime.now().strftime("%Y-%m-%d"), "search_results": []}
    if q in ("what time is it", "what's the time", "what is the time"):
        return {"text": datetime.now().strftime("%H:%M:%S %Z") or datetime.now().strftime("%H:%M:%S"), "search_results": []}
    return None

def search_perplexity(query: str, model: str = "sonar-pro", timeout: int = 30):
    """
    Perform a Perplexity chat completion (web grounded).
    Returns dict: {"text": str, "search_results": [ {title,url,date}, ... ], "raw": dict}
    """
    
    quick = local_quick_answer(query)
    if quick:
        return quick

    url = "https://api.perplexity.ai/chat/completions"
    payload = {
        "model": "sonar-pro",
        "messages": [
            {"role": "system", "content": "You are an AI assistant that provides concise, factual answers with citations."},
            {"role": "user", "content": query},
        ],
        "search_mode": "web",            
        "search_recency_filter": "month",
        "return_related_questions": False,
        "return_images": False,
        
    }
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "Authorization": f"Bearer {PERPLEXITY_API_KEY}",
    }

    try:
        resp = requests.post(url, json=payload, headers=headers, timeout=timeout)
        resp.raise_for_status()
    except requests.exceptions.RequestException as e:
        
        resp_text = getattr(e, "response", None)
        extra = resp_text.text if resp_text is not None else ""
        return {"text": f"Perplexity API request failed: {str(e)}. {extra}", "search_results": [], "raw": None}

    try:
        data = resp.json()
    except ValueError:
        return {"text": "Perplexity returned non-JSON response.", "search_results": [], "raw": resp.text}

    text = ""
    try:
        choices = data.get("choices") or []
        if len(choices) and isinstance(choices[0], dict):
            message = choices[0].get("message") or {}
            text = message.get("content") or ""
    except Exception:
        text = ""

    if not text:
        text = data.get("message") or data.get("text") or ""

    search_results = data.get("search_results") or []
    
    normalized = []
    for item in search_results:
        normalized.append({
            "title": item.get("title"),
            "url": item.get("url"),
            "date": item.get("date"),
        })

    return {"text": text, "search_results": normalized, "raw": data}

def main():
    print("🤖 Gemini Chatbot with Perplexity search (type '/search query' or 'exit')")
    print("-" * 60)

    while True:
        user_input = input("You: ").strip()
        if not user_input:
            continue

        if user_input.lower() in ("exit", "quit"):
            print("Goodbye! 👋")
            break

        if user_input.lower().startswith("/search"):
            query = user_input[len("/search"):].strip()
            if not query:
                print("Bot: Please provide a query after /search.")
                continue

            search_response = search_perplexity(query)
           
            if search_response.get("raw") is None and search_response.get("search_results") == [] and search_response.get("text", "").startswith("Perplexity API request failed"):
                print("Bot (search error):", search_response["text"])
                print("-" * 30)
                continue

            sources_text = ""
            if search_response["search_results"]:
                top = search_response["search_results"][:5]
                sources_text = "\n".join([f"- {s.get('title') or 'untitled'} ({s.get('date') or 'no-date'}): {s.get('url')}" for s in top])

            augmented_prompt = f"""
Based on the following web-grounded search results (from Perplexity), answer the user's question concisely and include source citations.

Web Search Context:
\"\"\"{search_response['text']}\"\"\"

Top sources:
{sources_text}

User question: "{query}"

Answer:
"""

            gemini_resp = gemini_model.generate_content(augmented_prompt)
           
            bot_text = getattr(gemini_resp, "text", None) or str(gemini_resp)
            print("Bot:", bot_text)
            
            if sources_text:
                print("\nSources:")
                print(sources_text)
        else:
            
            gemini_resp = gemini_model.generate_content(user_input)
            bot_text = getattr(gemini_resp, "text", None) or str(gemini_resp)
            print("Bot:", bot_text)

        print("-" * 60)

if __name__ == "__main__":
    main()