app.py

import asyncio
import aiohttp
import gradio as gr
import json
import re
import time
from datetime import datetime
from typing import List, Dict, Optional, Tuple
from urllib.parse import quote_plus, urljoin
from dataclasses import dataclass
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer
import requests
from bs4 import BeautifulSoup
import newspaper
from newspaper import Article
import logging
import warnings

# Suppress warnings
warnings.filterwarnings("ignore")
logging.getLogger().setLevel(logging.ERROR)

@dataclass
class SearchResult:
    """Data class for search results"""
    title: str
    url: str
    snippet: str
    content: str = ""
    publication_date: Optional[str] = None
    relevance_score: float = 0.0

class QueryEnhancer:
    """Enhance user queries with search operators and entity quoting"""
    
    def __init__(self):
        # Common named entity patterns
        self.entity_patterns = [
            r'\b[A-Z][a-z]+ [A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b',  # Proper names
            r'\b[A-Z]{2,}(?:\s+[A-Z][a-z]+)*\b',  # Acronyms + words
            r'\b[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\s+(?:Inc|Corp|LLC|Ltd|Co|Company|Trust|Group|Holdings)\b'  # Companies
        ]
    
    def enhance_query(self, query: str) -> str:
        """Enhance query by quoting named entities and adding operators"""
        enhanced = query
        
        # Find and quote named entities
        for pattern in self.entity_patterns:
            matches = re.findall(pattern, enhanced)
            for match in matches:
                if len(match.split()) > 1:  # Only quote multi-word entities
                    enhanced = enhanced.replace(match, f'"{match}"')
        
        return enhanced

class SearchEngineInterface:
    """Interface for different search engines"""
    
    def __init__(self):
        self.session = None
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
            'Accept-Language': 'en-US,en;q=0.5',
            'Accept-Encoding': 'gzip, deflate',
            'Connection': 'keep-alive',
        }
    
    async def get_session(self):
        """Get or create aiohttp session"""
        if self.session is None:
            connector = aiohttp.TCPConnector(limit=10)
            timeout = aiohttp.ClientTimeout(total=30)
            self.session = aiohttp.ClientSession(
                headers=self.headers,
                connector=connector,
                timeout=timeout
            )
        return self.session
    
    async def search_google(self, query: str, num_results: int = 10) -> List[SearchResult]:
        """Search Google and parse results"""
        try:
            session = await self.get_session()
            url = f"https://www.google.com/search?q={quote_plus(query)}&num={num_results}"
            
            async with session.get(url) as response:
                if response.status != 200:
                    return []
                
                html = await response.text()
                soup = BeautifulSoup(html, 'html.parser')
                results = []
                
                # Parse Google search results
                for g in soup.find_all('div', class_='g')[:num_results]:
                    try:
                        title_elem = g.find('h3')
                        if not title_elem:
                            continue
                            
                        title = title_elem.get_text()
                        
                        # Get URL
                        link_elem = g.find('a')
                        if not link_elem or not link_elem.get('href'):
                            continue
                        url = link_elem['href']
                        
                        # Get snippet
                        snippet_elem = g.find('span', class_=['st', 'aCOpRe'])
                        if not snippet_elem:
                            snippet_elem = g.find('div', class_=['s', 'st'])
                        snippet = snippet_elem.get_text() if snippet_elem else ""
                        
                        if title and url.startswith('http'):
                            results.append(SearchResult(title=title, url=url, snippet=snippet))
                    except Exception as e:
                        continue
                
                return results
        except Exception as e:
            print(f"Google search error: {e}")
            return []
    
    async def search_bing(self, query: str, num_results: int = 10) -> List[SearchResult]:
        """Search Bing and parse results"""
        try:
            session = await self.get_session()
            url = f"https://www.bing.com/search?q={quote_plus(query)}&count={num_results}"
            
            async with session.get(url) as response:
                if response.status != 200:
                    return []
                
                html = await response.text()
                soup = BeautifulSoup(html, 'html.parser')
                results = []
                
                # Parse Bing search results
                for result in soup.find_all('li', class_='b_algo')[:num_results]:
                    try:
                        title_elem = result.find('h2')
                        if not title_elem:
                            continue
                        
                        link_elem = title_elem.find('a')
                        if not link_elem:
                            continue
                            
                        title = link_elem.get_text()
                        url = link_elem.get('href', '')
                        
                        snippet_elem = result.find('p', class_='b_paractl') or result.find('div', class_='b_caption')
                        snippet = snippet_elem.get_text() if snippet_elem else ""
                        
                        if title and url.startswith('http'):
                            results.append(SearchResult(title=title, url=url, snippet=snippet))
                    except Exception as e:
                        continue
                
                return results
        except Exception as e:
            print(f"Bing search error: {e}")
            return []
    
    async def search_yahoo(self, query: str, num_results: int = 10) -> List[SearchResult]:
        """Search Yahoo and parse results"""
        try:
            session = await self.get_session()
            url = f"https://search.yahoo.com/search?p={quote_plus(query)}&n={num_results}"
            
            async with session.get(url) as response:
                if response.status != 200:
                    return []
                
                html = await response.text()
                soup = BeautifulSoup(html, 'html.parser')
                results = []
                
                # Parse Yahoo search results
                for result in soup.find_all('div', class_='dd')[:num_results]:
                    try:
                        title_elem = result.find('h3', class_='title')
                        if not title_elem:
                            continue
                        
                        link_elem = title_elem.find('a')
                        if not link_elem:
                            continue
                            
                        title = link_elem.get_text()
                        url = link_elem.get('href', '')
                        
                        snippet_elem = result.find('div', class_='compText')
                        snippet = snippet_elem.get_text() if snippet_elem else ""
                        
                        if title and url.startswith('http'):
                            results.append(SearchResult(title=title, url=url, snippet=snippet))
                    except Exception as e:
                        continue
                
                return results
        except Exception as e:
            print(f"Yahoo search error: {e}")
            return []
    
    async def close(self):
        """Close the session"""
        if self.session:
            await self.session.close()

class ContentScraper:
    """Scrape and parse article content using newspaper3k"""
    
    def __init__(self):
        self.session = None
    
    async def get_session(self):
        """Get or create aiohttp session"""
        if self.session is None:
            connector = aiohttp.TCPConnector(limit=20)
            timeout = aiohttp.ClientTimeout(total=30)
            self.session = aiohttp.ClientSession(
                connector=connector,
                timeout=timeout
            )
        return self.session
    
    async def scrape_article(self, url: str) -> Tuple[str, Optional[str]]:
        """Scrape article content and publication date"""
        try:
            # Use newspaper3k for article extraction
            article = Article(url)
            article.download()
            article.parse()
            
            content = article.text
            pub_date = article.publish_date.isoformat() if article.publish_date else None
            
            return content, pub_date
        except Exception as e:
            print(f"Error scraping {url}: {e}")
            return "", None
    
    async def scrape_multiple(self, search_results: List[SearchResult]) -> List[SearchResult]:
        """Scrape multiple articles in parallel"""
        tasks = []
        for result in search_results:
            tasks.append(self.scrape_article(result.url))
        
        scraped_data = await asyncio.gather(*tasks, return_exceptions=True)
        
        for i, (content, pub_date) in enumerate(scraped_data):
            if not isinstance(content, Exception):
                search_results[i].content = content
                search_results[i].publication_date = pub_date
        
        return search_results
    
    async def close(self):
        """Close the session"""
        if self.session:
            await self.session.close()

class EmbeddingFilter:
    """Filter search results using embedding-based similarity"""
    
    def __init__(self):
        self.vectorizer = TfidfVectorizer(max_features=1000, stop_words='english')
    
    def filter_by_relevance(self, query: str, search_results: List[SearchResult], 
                          threshold: float = 0.1) -> List[SearchResult]:
        """Filter results by cosine similarity with query"""
        if not search_results:
            return search_results
        
        # Combine title, snippet, and content for each result
        result_texts = []
        for result in search_results:
            combined_text = f"{result.title} {result.snippet} {result.content[:1000]}"
            result_texts.append(combined_text)
        
        if not result_texts:
            return search_results
        
        try:
            # Add query to the corpus for vectorization
            all_texts = [query] + result_texts
            
            # Vectorize texts
            tfidf_matrix = self.vectorizer.fit_transform(all_texts)
            
            # Calculate cosine similarity between query and each result
            query_vector = tfidf_matrix[0:1]
            result_vectors = tfidf_matrix[1:]
            
            similarities = cosine_similarity(query_vector, result_vectors)[0]
            
            # Add relevance scores and filter
            filtered_results = []
            for i, result in enumerate(search_results):
                result.relevance_score = similarities[i]
                if similarities[i] >= threshold:
                    filtered_results.append(result)
            
            # Sort by relevance score
            filtered_results.sort(key=lambda x: x.relevance_score, reverse=True)
            return filtered_results
            
        except Exception as e:
            print(f"Embedding filter error: {e}")
            return search_results

class LLMSummarizer:
    """Summarize search results using Groq or OpenRouter APIs"""
    
    def __init__(self, groq_api_key: str = "", openrouter_api_key: str = ""):
        self.groq_api_key = groq_api_key
        self.openrouter_api_key = openrouter_api_key
        self.groq_model = "meta-llama/llama-4-maverick-17b-128e-instruct"
        self.openrouter_model = "deepseek/deepseek-r1:free"
    
    def create_system_prompt(self) -> str:
        """Create system prompt for summarization"""
        return """You are an expert summarizer. Your task is to analyze search results and provide a comprehensive, accurate summary that directly answers the user's query.

Instructions:
1. Focus only on information relevant to the user's query
2. Filter out noise, advertisements, and unrelated content
3. Synthesize information from multiple sources when possible
4. Maintain factual accuracy and cite sources when appropriate
5. If information is contradictory, note the discrepancies
6. Provide a clear, concise summary that directly addresses the query
7. Include relevant dates, numbers, and specific details when available

Format your response as a comprehensive summary, not bullet points."""
    
    async def summarize_with_groq(self, query: str, search_results: List[SearchResult], 
                                temperature: float = 0.3, max_tokens: int = 2000) -> str:
        """Summarize using Groq API"""
        if not self.groq_api_key:
            return "Groq API key not provided"
        
        try:
            # Prepare the content for summarization
            content_json = {
                "user_query": query,
                "search_results": []
            }
            
            for result in search_results:
                content_json["search_results"].append({
                    "title": result.title,
                    "url": result.url,
                    "snippet": result.snippet,
                    "content": result.content[:2000],  # Limit content length
                    "publication_date": result.publication_date,
                    "relevance_score": result.relevance_score
                })
            
            user_prompt = f"""Please summarize the following search results for the query: "{query}"

Search Results Data:
{json.dumps(content_json, indent=2)}

Provide a comprehensive summary that directly answers the user's query based on the most relevant and recent information available."""
            
            headers = {
                "Authorization": f"Bearer {self.groq_api_key}",
                "Content-Type": "application/json"
            }
            
            payload = {
                "model": self.groq_model,
                "messages": [
                    {"role": "system", "content": self.create_system_prompt()},
                    {"role": "user", "content": user_prompt}
                ],
                "temperature": temperature,
                "max_tokens": max_tokens
            }
            
            async with aiohttp.ClientSession() as session:
                async with session.post("https://api.groq.com/openai/v1/chat/completions", 
                                      headers=headers, json=payload) as response:
                    if response.status == 200:
                        result = await response.json()
                        return result["choices"][0]["message"]["content"]
                    else:
                        error_text = await response.text()
                        return f"Groq API error: {response.status} - {error_text}"
        
        except Exception as e:
            return f"Error with Groq summarization: {str(e)}"
    
    async def summarize_with_openrouter(self, query: str, search_results: List[SearchResult], 
                                      temperature: float = 0.3, max_tokens: int = 2000) -> str:
        """Summarize using OpenRouter API"""
        if not self.openrouter_api_key:
            return "OpenRouter API key not provided"
        
        try:
            # Prepare the content for summarization
            content_json = {
                "user_query": query,
                "search_results": []
            }
            
            for result in search_results:
                content_json["search_results"].append({
                    "title": result.title,
                    "url": result.url,
                    "snippet": result.snippet,
                    "content": result.content[:2000],  # Limit content length
                    "publication_date": result.publication_date,
                    "relevance_score": result.relevance_score
                })
            
            user_prompt = f"""Please summarize the following search results for the query: "{query}"

Search Results Data:
{json.dumps(content_json, indent=2)}

Provide a comprehensive summary that directly answers the user's query based on the most relevant and recent information available."""
            
            headers = {
                "Authorization": f"Bearer {self.openrouter_api_key}",
                "Content-Type": "application/json",
                "HTTP-Referer": "https://huggingface.co/spaces",
                "X-Title": "AI Search Engine"
            }
            
            payload = {
                "model": self.openrouter_model,
                "messages": [
                    {"role": "system", "content": self.create_system_prompt()},
                    {"role": "user", "content": user_prompt}
                ],
                "temperature": temperature,
                "max_tokens": max_tokens
            }
            
            async with aiohttp.ClientSession() as session:
                async with session.post("https://openrouter.ai/api/v1/chat/completions", 
                                      headers=headers, json=payload) as response:
                    if response.status == 200:
                        result = await response.json()
                        return result["choices"][0]["message"]["content"]
                    else:
                        error_text = await response.text()
                        return f"OpenRouter API error: {response.status} - {error_text}"
        
        except Exception as e:
            return f"Error with OpenRouter summarization: {str(e)}"

class AISearchEngine:
    """Main AI-powered search engine class"""
    
    def __init__(self, groq_api_key: str = "", openrouter_api_key: str = ""):
        self.query_enhancer = QueryEnhancer()
        self.search_interface = SearchEngineInterface()
        self.content_scraper = ContentScraper()
        self.embedding_filter = EmbeddingFilter()
        self.llm_summarizer = LLMSummarizer(groq_api_key, openrouter_api_key)
    
    async def search_and_summarize(self, 
                                 query: str,
                                 search_engines: List[str],
                                 model: str,
                                 use_embeddings: bool,
                                 temperature: float,
                                 max_results: int,
                                 max_tokens: int) -> Tuple[str, str]:
        """Main search and summarization pipeline"""
        
        start_time = time.time()
        status_updates = []
        
        try:
            # Step 1: Query Enhancement
            status_updates.append("🔍 Enhancing search query...")
            enhanced_query = self.query_enhancer.enhance_query(query)
            status_updates.append(f"Enhanced query: {enhanced_query}")
            
            # Step 2: Parallel Search across engines
            status_updates.append("🌐 Searching across multiple engines...")
            search_tasks = []
            
            if "Google" in search_engines:
                search_tasks.append(self.search_interface.search_google(enhanced_query, max_results))
            if "Bing" in search_engines:
                search_tasks.append(self.search_interface.search_bing(enhanced_query, max_results))
            if "Yahoo" in search_engines:
                search_tasks.append(self.search_interface.search_yahoo(enhanced_query, max_results))
            
            if not search_tasks:
                return "No search engines selected", "\n".join(status_updates)
            
            search_results_lists = await asyncio.gather(*search_tasks)
            
            # Combine and deduplicate results
            all_results = []
            seen_urls = set()
            
            for results_list in search_results_lists:
                for result in results_list:
                    if result.url not in seen_urls:
                        all_results.append(result)
                        seen_urls.add(result.url)
            
            status_updates.append(f"Found {len(all_results)} unique results")
            
            if not all_results:
                return "No search results found", "\n".join(status_updates)
            
            # Step 3: Content Scraping
            status_updates.append("📄 Scraping article content...")
            scraped_results = await self.content_scraper.scrape_multiple(all_results[:max_results])
            
            # Filter results with content
            results_with_content = [r for r in scraped_results if r.content.strip()]
            status_updates.append(f"Successfully scraped {len(results_with_content)} articles")
            
            # Step 4: Optional Embedding-based Filtering
            if use_embeddings and results_with_content:
                status_updates.append("🧠 Filtering results using embeddings...")
                filtered_results = self.embedding_filter.filter_by_relevance(query, results_with_content)
                status_updates.append(f"Filtered to {len(filtered_results)} most relevant results")
            else:
                filtered_results = results_with_content
            
            if not filtered_results:
                return "No relevant results found after filtering", "\n".join(status_updates)
            
            # Step 5: LLM Summarization
            status_updates.append(f"🤖 Generating summary using {model}...")
            
            if model.startswith("Groq"):
                summary = await self.llm_summarizer.summarize_with_groq(
                    query, filtered_results, temperature, max_tokens
                )
            else:  # OpenRouter
                summary = await self.llm_summarizer.summarize_with_openrouter(
                    query, filtered_results, temperature, max_tokens
                )
            
            # Add metadata
            end_time = time.time()
            processing_time = end_time - start_time
            
            metadata = f"\n\n---\n**Search Metadata:**\n"
            metadata += f"- Processing time: {processing_time:.2f} seconds\n"
            metadata += f"- Results found: {len(all_results)}\n"
            metadata += f"- Articles scraped: {len(results_with_content)}\n"
            metadata += f"- Results used for summary: {len(filtered_results)}\n"
            metadata += f"- Search engines: {', '.join(search_engines)}\n"
            metadata += f"- Model: {model}\n"
            metadata += f"- Embeddings used: {use_embeddings}\n"
            
            final_summary = summary + metadata
            status_updates.append(f"✅ Summary generated in {processing_time:.2f}s")
            
            return final_summary, "\n".join(status_updates)
            
        except Exception as e:
            error_msg = f"Error in search pipeline: {str(e)}"
            status_updates.append(f"❌ {error_msg}")
            return error_msg, "\n".join(status_updates)
        
        finally:
            # Cleanup
            await self.search_interface.close()
            await self.content_scraper.close()

# Global search engine instance
search_engine = None

async def initialize_search_engine(groq_key: str, openrouter_key: str):
    """Initialize the search engine with API keys"""
    global search_engine
    search_engine = AISearchEngine(groq_key, openrouter_key)
    return search_engine

async def perform_search(query: str, 
                        search_engines: List[str],
                        model: str,
                        use_embeddings: bool,
                        temperature: float,
                        max_results: int,
                        max_tokens: int,
                        groq_key: str,
                        openrouter_key: str):
    """Perform search with given parameters"""
    global search_engine
    
    if search_engine is None:
        search_engine = await initialize_search_engine(groq_key, openrouter_key)
    
    return await search_engine.search_and_summarize(
        query, search_engines, model, use_embeddings, 
        temperature, max_results, max_tokens
    )

async def chat_inference(message, history, groq_key, openrouter_key, model_choice, search_engines, use_embeddings, temperature, max_results, max_tokens):
    """Main chat inference function for ChatInterface with additional inputs"""
    try:
        if not message.strip():
            yield "Please enter a search query."
            return
        
        if not groq_key and not openrouter_key:
            yield "❌ Please provide at least one API key (Groq or OpenRouter) to use the AI summarization features."
            return
        
        if not search_engines:
            yield "❌ Please select at least one search engine."
            return
        
        # Initialize search engine
        global search_engine
        if search_engine is None:
            search_engine = await initialize_search_engine(groq_key, openrouter_key)
        else:
            # Update API keys if they changed
            search_engine.llm_summarizer.groq_api_key = groq_key
            search_engine.llm_summarizer.openrouter_api_key = openrouter_key
        
        # Start with status updates
        yield "🔍 Enhancing query and searching across multiple engines..."
        
        # Small delay to show the initial status
        await asyncio.sleep(0.1)
        
        # Update status
        yield "🌐 Fetching results from search engines..."
        await asyncio.sleep(0.1)
        
        # Update status  
        yield "📄 Scraping article content..."
        await asyncio.sleep(0.1)
        
        if use_embeddings:
            yield "🧠 Filtering results using embeddings..."
            await asyncio.sleep(0.1)
        
        yield "🤖 Generating AI-powered summary..."
        await asyncio.sleep(0.1)
        
        # Perform the actual search and summarization
        summary, status = await search_engine.search_and_summarize(
            message,
            search_engines,
            model_choice,
            use_embeddings,
            temperature,
            max_results,
            max_tokens
        )
        
        # Stream the final result
        yield summary
        
    except Exception as e:
        yield f"❌ Search failed: {str(e)}\n\nPlease check your API keys and try again."

def create_gradio_interface():
    """Create the modern Gradio ChatInterface"""
    
    # Define additional inputs for the accordion
    additional_inputs = [
        gr.Textbox(
            label="🔑 Groq API Key",
            type="password",
            placeholder="Enter your Groq API key (get from: https://console.groq.com/)",
            info="Required for Groq Llama-4 model"
        ),
        gr.Textbox(
            label="🔑 OpenRouter API Key", 
            type="password",
            placeholder="Enter your OpenRouter API key (get from: https://openrouter.ai/)",
            info="Required for OpenRouter DeepSeek-R1 model"
        ),
        gr.Dropdown(
            choices=["Groq (Llama-4)", "OpenRouter (DeepSeek-R1)"],
            value="Groq (Llama-4)",
            label="🤖 AI Model",
            info="Choose the AI model for summarization"
        ),
        gr.CheckboxGroup(
            choices=["Google", "Bing", "Yahoo"],
            value=["Google", "Bing"],
            label="🔍 Search Engines",
            info="Select which search engines to use (multiple recommended)"
        ),
        gr.Checkbox(
            value=True,
            label="🧠 Use Embedding-based Filtering",
            info="Filter results by relevance using TF-IDF similarity (recommended)"
        ),
        gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.3,
            step=0.1,
            label="🌡️ Temperature",
            info="Higher = more creative, Lower = more focused (0.1-0.3 recommended for factual queries)"
        ),
        gr.Slider(
            minimum=5,
            maximum=20,
            value=10,
            step=1,
            label="📊 Max Results per Engine",
            info="Number of search results to fetch from each engine"
        ),
        gr.Slider(
            minimum=500,
            maximum=4000,
            value=2000,
            step=100,
            label="📝 Max Tokens",
            info="Maximum length of the AI-generated summary"
        )
    ]
    
    # Create the main ChatInterface
    chat_interface = gr.ChatInterface(
        fn=chat_inference,
        additional_inputs=additional_inputs,
        additional_inputs_accordion=gr.Accordion("⚙️ Configuration & Advanced Parameters", open=True),
        title="🔍 AI-Powered Search Engine",
        description="""
        **Search across Google, Bing, and Yahoo, then get AI-powered summaries!**
        
        ✨ **Features:** Multi-engine search • Query enhancement • Parallel scraping • AI summarization • Embedding filtering
        
        📋 **Quick Start:** 1) Add your API key below 2) Select search engines 3) Ask any question!
        """,
        cache_examples=False,
        #retry_btn="🔄 Retry",
        #undo_btn="↩️ Undo", 
        #clear_btn="🗑️ Clear",
        submit_btn="🔍 Search & Summarize",
        stop_btn="⏹️ Stop",
        chatbot=gr.Chatbot(
            show_copy_button=True,
            #likeable=True,
            layout="bubble",
            height=600,
            placeholder="🚀 Ready to search! Configure your settings below and ask me anything.",
            show_share_button=True
        ),
        theme=gr.themes.Soft(),
        analytics_enabled=False,
        type="messages"  # Use the modern message format
    )
    
    return chat_interface

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=True)