app.py

"""
Intelligent Search Engine with RAG and OSINT capabilities.
"""
import os
import asyncio
import gradio as gr
from engines.search import SearchEngine
from engines.osint import OSINTEngine
from engines.image import ImageEngine
import markdown2
from typing import Dict, Any, List

# Initialize engines
search_engine = SearchEngine()
osint_engine = OSINTEngine()
image_engine = ImageEngine()

def format_search_results(results: Dict[str, Any]) -> str:
    """Format search results with markdown."""
    if not results or "answer" not in results:
        return "No results found."
        
    formatted = f"### Answer\n{results['answer']}\n\n"
    
    if results.get("sources"):
        formatted += "\n### Sources\n"
        for i, source in enumerate(results["sources"], 1):
            formatted += f"{i}. [{source}]({source})\n"
    
    return formatted

def format_osint_results(results: Dict[str, Any]) -> str:
    """Format OSINT results with markdown."""
    formatted = "### OSINT Results\n\n"
    
    if "error" in results:
        return f"Error: {results['error']}"
    
    if "found_on" in results:
        formatted += "#### Social Media Presence\n"
        for platform in results["found_on"]:
            formatted += f"- {platform['platform']}: [{platform['url']}]({platform['url']})\n"
    
    if "person_info" in results:
        person = results["person_info"]
        formatted += f"\n#### Personal Information\n"
        formatted += f"- Name: {person.get('name', 'N/A')}\n"
        if person.get("age"):
            formatted += f"- Age: {person['age']}\n"
        if person.get("location"):
            formatted += f"- Location: {person['location']}\n"
        if person.get("gender"):
            formatted += f"- Gender: {person['gender']}\n"
    
    return formatted

async def search_query(query: str) -> str:
    """Handle search queries."""
    try:
        results = await search_engine.search(query)
        return format_search_results(results)
    except Exception as e:
        return f"Error: {str(e)}"

async def search_username(username: str) -> str:
    """Search for username across platforms."""
    try:
        results = await osint_engine.search_username(username)
        return format_osint_results(results)
    except Exception as e:
        return f"Error: {str(e)}"

async def search_person(name: str, location: str = "", age: str = "", gender: str = "") -> str:
    """Search for person information."""
    try:
        age_int = int(age) if age.strip() else None
        person = await osint_engine.search_person(
            name=name,
            location=location if location.strip() else None,
            age=age_int,
            gender=gender if gender.strip() else None
        )
        return format_osint_results({"person_info": person.to_dict()})
    except Exception as e:
        return f"Error: {str(e)}"

async def analyze_image_file(image) -> str:
    """Analyze uploaded image."""
    try:
        if not image:
            return "No image provided."
            
        # Read image data
        with open(image.name, "rb") as f:
            image_data = f.read()
        
        # Analyze image
        results = await image_engine.analyze_image(image_data)
        
        if "error" in results:
            return f"Error analyzing image: {results['error']}"
        
        # Format results
        formatted = "### Image Analysis Results\n\n"
        
        # Add predictions
        formatted += "#### Content Detection\n"
        for pred in results["predictions"]:
            confidence = pred["confidence"] * 100
            formatted += f"- {pred['label']}: {confidence:.1f}%\n"
        
        # Add face detection results
        formatted += f"\n#### Face Detection\n"
        formatted += f"- Found {len(results['faces'])} faces\n"
        
        # Add metadata
        formatted += f"\n#### Image Metadata\n"
        metadata = results["metadata"]
        formatted += f"- Size: {metadata['width']}x{metadata['height']}\n"
        formatted += f"- Format: {metadata['format']}\n"
        formatted += f"- Mode: {metadata['mode']}\n"
        
        return formatted
        
    except Exception as e:
        return f"Error: {str(e)}"

def create_ui() -> gr.Blocks:
    """Create the Gradio interface."""
    with gr.Blocks(title="Intelligent Search Engine", theme=gr.themes.Soft()) as app:
        gr.Markdown("""
        # 🔍 Intelligent Search Engine
        
        Advanced search engine with RAG and OSINT capabilities.
        """)
        
        with gr.Tabs():
            # Intelligent Search Tab
            with gr.Tab("🌐 Search"):
                with gr.Column():
                    search_input = gr.Textbox(
                        label="Enter your search query",
                        placeholder="What would you like to know?"
                    )
                    search_button = gr.Button("Search", variant="primary")
                    search_output = gr.Markdown(label="Results")
                    
                search_button.click(
                    fn=search_query,
                    inputs=search_input,
                    outputs=search_output
                )
            
            # Username Search Tab
            with gr.Tab("👤 Username Search"):
                with gr.Column():
                    username_input = gr.Textbox(
                        label="Enter username",
                        placeholder="Username to search across platforms"
                    )
                    username_button = gr.Button("Search Username", variant="primary")
                    username_output = gr.Markdown(label="Results")
                    
                username_button.click(
                    fn=search_username,
                    inputs=username_input,
                    outputs=username_output
                )
            
            # Person Search Tab
            with gr.Tab("👥 Person Search"):
                with gr.Column():
                    name_input = gr.Textbox(
                        label="Full Name",
                        placeholder="Enter person's name"
                    )
                    location_input = gr.Textbox(
                        label="Location (optional)",
                        placeholder="City, Country"
                    )
                    age_input = gr.Textbox(
                        label="Age (optional)",
                        placeholder="Enter age"
                    )
                    gender_input = gr.Dropdown(
                        label="Gender (optional)",
                        choices=["", "Male", "Female", "Other"]
                    )
                    person_button = gr.Button("Search Person", variant="primary")
                    person_output = gr.Markdown(label="Results")
                    
                person_button.click(
                    fn=search_person,
                    inputs=[name_input, location_input, age_input, gender_input],
                    outputs=person_output
                )
            
            # Image Analysis Tab
            with gr.Tab("🖼️ Image Analysis"):
                with gr.Column():
                    image_input = gr.File(
                        label="Upload Image",
                        file_types=["image"]
                    )
                    image_button = gr.Button("Analyze Image", variant="primary")
                    image_output = gr.Markdown(label="Results")
                    
                image_button.click(
                    fn=analyze_image_file,
                    inputs=image_input,
                    outputs=image_output
                )
        
        gr.Markdown("""
        ### 📝 Notes
        - The search engine uses RAG (Retrieval-Augmented Generation) for intelligent answers
        - OSINT capabilities include social media presence, personal information, and image analysis
        - All searches are conducted using publicly available information
        """)
    
    return app

if __name__ == "__main__":
    app = create_ui()
    app.launch(share=True)