app.py

import gradio as gr
import asyncio
import os
import tempfile
import shutil
from pathlib import Path
import logging
import sys

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Add the current directory to Python path to import raganything
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

try:
    from raganything import RAGAnything, RAGAnythingConfig
    from lightrag.utils import EmbeddingFunc
    from lightrag.llm.openai import openai_complete_if_cache, openai_embed
    RAG_AVAILABLE = True
except ImportError as e:
    logger.error(f"RAGAnything import failed: {e}")
    RAG_AVAILABLE = False

# Global variables
rag_instance = None
processed_files = []

# Get API keys from environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
OPENAI_BASE_URL = os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")

def get_llm_model_func():
    """Get LLM model function"""
    if not OPENAI_API_KEY:
        logger.warning("No OpenAI API key found, using mock function")
        def mock_llm_func(prompt, system_prompt=None, history_messages=[], **kwargs):
            return f"Mock response to: {prompt[:100]}... (Add OpenAI API key to get real responses)"
        return mock_llm_func
    
    def llm_model_func(prompt, system_prompt=None, history_messages=[], **kwargs):
        return openai_complete_if_cache(
            "gpt-4o-mini",
            prompt,
            system_prompt=system_prompt,
            history_messages=history_messages,
            api_key=OPENAI_API_KEY,
            base_url=OPENAI_BASE_URL,
            **kwargs,
        )
    return llm_model_func

def get_vision_model_func():
    """Get vision model function for image processing"""
    if not OPENAI_API_KEY:
        def mock_vision_func(prompt, system_prompt=None, history_messages=[], image_data=None, **kwargs):
            return f"Mock vision response to: {prompt[:50]}..."
        return mock_vision_func
    
    def vision_model_func(prompt, system_prompt=None, history_messages=[], image_data=None, **kwargs):
        if image_data:
            return openai_complete_if_cache(
                "gpt-4o",
                "",
                system_prompt=None,
                history_messages=[],
                messages=[
                    {"role": "system", "content": system_prompt} if system_prompt else None,
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {
                                "type": "image_url",
                                "image_url": {"url": f"data:image/jpeg;base64,{image_data}"},
                            },
                        ],
                    } if image_data else {"role": "user", "content": prompt},
                ],
                api_key=OPENAI_API_KEY,
                base_url=OPENAI_BASE_URL,
                **kwargs,
            )
        else:
            return get_llm_model_func()(prompt, system_prompt, history_messages, **kwargs)
    
    return vision_model_func

def get_embedding_func():
    """Get embedding function"""
    if not OPENAI_API_KEY:
        def mock_embedding_func(texts):
            import numpy as np
            if isinstance(texts, str):
                texts = [texts]
            return np.random.rand(len(texts), 1536).tolist()
        
        return EmbeddingFunc(
            embedding_dim=1536,
            max_token_size=8192,
            func=mock_embedding_func
        )
    
    return EmbeddingFunc(
        embedding_dim=3072,
        max_token_size=8192,
        func=lambda texts: openai_embed(
            texts,
            model="text-embedding-3-large",
            api_key=OPENAI_API_KEY,
            base_url=OPENAI_BASE_URL,
        ),
    )

async def initialize_rag():
    """Initialize the RAG system"""
    global rag_instance
    
    if not RAG_AVAILABLE:
        return "❌ RAGAnything not installed. Please check requirements."
    
    try:
        # Create working directory
        working_dir = "./rag_storage"
        os.makedirs(working_dir, exist_ok=True)
        
        # Create configuration
        config = RAGAnythingConfig(
            working_dir=working_dir,
            mineru_parse_method="auto",
            enable_image_processing=True,
            enable_table_processing=True,
            enable_equation_processing=True,
        )
        
        # Get model functions
        llm_func = get_llm_model_func()
        vision_func = get_vision_model_func()
        embedding_func = get_embedding_func()
        
        # Initialize RAGAnything
        rag_instance = RAGAnything(
            config=config,
            llm_model_func=llm_func,
            vision_model_func=vision_func,
            embedding_func=embedding_func,
        )
        
        api_status = "with OpenAI API" if OPENAI_API_KEY else "in demo mode (add OPENAI_API_KEY for full functionality)"
        return f"✅ RAG-Anything initialized successfully {api_status}!"
        
    except Exception as e:
        logger.error(f"RAG initialization error: {e}")
        return f"❌ RAG initialization failed: {str(e)}"

async def process_document(file_path, file_name):
    """Process uploaded document"""
    global rag_instance, processed_files
    
    if not rag_instance:
        init_result = await initialize_rag()
        if "❌" in init_result:
            return init_result, processed_files
    
    try:
        # Create output directory
        output_dir = "./rag_output"
        os.makedirs(output_dir, exist_ok=True)
        
        # Process document with RAG-Anything
        logger.info(f"Processing document: {file_name}")
        
        await rag_instance.process_document_complete(
            file_path=file_path,
            output_dir=output_dir,
            parse_method="auto"
        )
        
        processed_files.append(file_name)
        return f"✅ Successfully processed: {file_name}\n\nDocument has been parsed and added to the knowledge base. You can now ask questions about its content.", processed_files
        
    except Exception as e:
        logger.error(f"Document processing error: {e}")
        return f"❌ Failed to process {file_name}: {str(e)}", processed_files

async def query_documents(question, mode="hybrid"):
    """Query processed documents"""
    global rag_instance
    
    if not rag_instance:
        return "❌ Please initialize the system and process documents first."
    
    if not processed_files:
        return "❌ No documents processed yet. Please upload and process documents first."
    
    try:
        # Use RAG-Anything query
        result = await rag_instance.aquery(question, mode=mode)
        return f"📝 **Answer:**\n\n{result}\n\n---\n*Based on analysis of: {', '.join(processed_files)}*"
        
    except Exception as e:
        logger.error(f"Query error: {e}")
        return f"❌ Query failed: {str(e)}"

# Gradio interface functions
def upload_and_process(file):
    """Handle file upload and processing"""
    if file is None:
        return "❌ Please upload a file first.", processed_files
    
    try:
        # Copy file to temp location
        temp_path = f"./temp_{Path(file.name).name}"
        shutil.copy2(file.name, temp_path)
        
        # Process asynchronously
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        result, files = loop.run_until_complete(
            process_document(temp_path, Path(file.name).name)
        )
        loop.close()
        
        # Cleanup
        if os.path.exists(temp_path):
            os.remove(temp_path)
            
        return result, "\n".join(files) if files else "No files processed"
        
    except Exception as e:
        logger.error(f"Upload error: {e}")
        return f"❌ Upload failed: {str(e)}", processed_files

def ask_question(question, mode):
    """Handle question asking"""
    if not question.strip():
        return "❌ Please enter a question."
    
    try:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        result = loop.run_until_complete(query_documents(question, mode))
        loop.close()
        return result
    except Exception as e:
        logger.error(f"Query error: {e}")
        return f"❌ Query failed: {str(e)}"

def initialize_system():
    """Initialize the RAG system"""
    try:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        result = loop.run_until_complete(initialize_rag())
        loop.close()
        return result
    except Exception as e:
        logger.error(f"Initialization error: {e}")
        return f"❌ Initialization failed: {str(e)}"

# Create Gradio interface
def create_interface():
    # Custom CSS for professional look
    custom_css = """
    .gradio-container {
        max-width: 1200px !important;
        margin: auto;
    }
    .main-header {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 2rem;
        border-radius: 15px;
        margin-bottom: 2rem;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
    }
    .status-box {
        background: #f8f9fa;
        border: 1px solid #dee2e6;
        border-radius: 8px;
        padding: 1rem;
    }
    .feature-grid {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));
        gap: 1rem;
        margin: 1rem 0;
    }
    .feature-card {
        background: white;
        padding: 1.5rem;
        border-radius: 10px;
        box-shadow: 0 2px 10px rgba(0,0,0,0.1);
        border-left: 4px solid #667eea;
    }
    """
    
    with gr.Blocks(
        title="RAG-Anything: Production System", 
        theme=gr.themes.Soft(),
        css=custom_css
    ) as demo:
        
        # Header
        gr.HTML("""
        <div class="main-header">
            <h1>🚀 RAG-Anything</h1>
            <h2>Production Multimodal Document AI System</h2>
            <p>Real document processing • Advanced AI understanding • Production ready</p>
        </div>
        """)
        
        # System status and initialization
        with gr.Row():
            with gr.Column():
                init_btn = gr.Button("🔧 Initialize RAG System", variant="primary", size="lg")
                init_status = gr.Textbox(
                    label="System Status", 
                    value="Click 'Initialize RAG System' to start the engine",
                    interactive=False,
                    elem_classes=["status-box"]
                )
        
        # Main processing area
        with gr.Row():
            # Document processing column
            with gr.Column(scale=1):
                gr.Markdown("## 📄 Document Processing")
                file_input = gr.File(
                    label="Upload Document",
                    file_types=[".pdf", ".docx", ".pptx", ".xlsx", ".jpg", ".png", ".txt", ".md"],
                    file_count="single"
                )
                process_btn = gr.Button("📤 Process with RAG-Anything", variant="secondary", size="lg")
                process_status = gr.Textbox(label="Processing Status", lines=4, interactive=False)
                processed_list = gr.Textbox(
                    label="Processed Documents", 
                    value="No documents processed yet",
                    lines=3,
                    interactive=False
                )
            
            # Query column
            with gr.Column(scale=1):
                gr.Markdown("## 🔍 Intelligent Query System")
                question_input = gr.Textbox(
                    label="Ask Questions About Your Documents",
                    placeholder="What are the main findings? Explain the methodology? Compare the data...",
                    lines=3
                )
                mode_dropdown = gr.Dropdown(
                    choices=["hybrid", "local", "global", "naive"],
                    value="hybrid",
                    label="Retrieval Mode",
                    info="Hybrid combines vector search + knowledge graph"
                )
                ask_btn = gr.Button("🤖 Get AI Answer", variant="primary", size="lg")
        
        # Results area
        answer_output = gr.Textbox(
            label="AI Response",
            lines=12,
            interactive=False,
            show_copy_button=True
        )
        
        # Example questions
        gr.Markdown("## 💡 Example Questions")
        examples = gr.Examples(
            examples=[
                ["What are the main findings discussed in this document?"],
                ["Summarize the key data points from tables and figures"],
                ["What methodology or approach is described?"],
                ["Compare the performance metrics or results shown"],
                ["Explain any mathematical formulas or equations present"],
                ["What are the conclusions and recommendations?"],
                ["How do the images and charts support the text content?"]
            ],
            inputs=[question_input],
            label="Click any example to try it"
        )
        
        # Technology showcase
        gr.HTML("""
        <div style="margin-top: 2rem; padding: 2rem; background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%); border-radius: 15px;">
            <h3 style="text-align: center; margin-bottom: 1.5rem;">🛠️ RAG-Anything Technology Stack</h3>
            <div class="feature-grid">
                <div class="feature-card">
                    <h4>🧠 LightRAG Engine</h4>
                    <p>Fast retrieval-augmented generation with knowledge graphs</p>
                </div>
                <div class="feature-card">
                    <h4>⚡ MinerU Parser</h4>
                    <p>High-fidelity document structure extraction and analysis</p>
                </div>
                <div class="feature-card">
                    <h4>🔗 Multimodal Processing</h4>
                    <p>Unified handling of text, images, tables, and equations</p>
                </div>
                <div class="feature-card">
                    <h4>🎯 Hybrid Retrieval</h4>
                    <p>Vector similarity + graph traversal for precise answers</p>
                </div>
            </div>
        </div>
        """)
        
        # API Configuration info
        gr.HTML(f"""
        <div style="margin-top: 1rem; padding: 1rem; background: {'#d4edda' if OPENAI_API_KEY else '#f8d7da'}; 
                    border-radius: 8px; border: 1px solid {'#c3e6cb' if OPENAI_API_KEY else '#f5c6cb'};">
            <h4>🔑 API Configuration</h4>
            <p><strong>Status:</strong> {'✅ OpenAI API configured' if OPENAI_API_KEY else '⚠️ OpenAI API key not found'}</p>
            <p><small>{'Full functionality enabled' if OPENAI_API_KEY else 'Add OPENAI_API_KEY environment variable for full functionality'}</small></p>
        </div>
        """)
        
        # Footer
        gr.HTML("""
        <div style="text-align: center; margin-top: 2rem; padding: 1.5rem; 
                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                    color: white; border-radius: 10px;">
            <h3>🚀 RAG-Anything: Production Ready</h3>
            <p>🌟 <a href="https://github.com/HKUDS/RAG-Anything" style="color: #ffd700;" target="_blank">
                View Source Code</a> | 
               📧 <strong>Enterprise Solutions Available</strong></p>
        </div>
        """)
        
        # Event handlers
        init_btn.click(
            fn=initialize_system,
            outputs=init_status
        )
        
        process_btn.click(
            fn=upload_and_process,
            inputs=file_input,
            outputs=[process_status, processed_list]
        )
        
        ask_btn.click(
            fn=ask_question,
            inputs=[question_input, mode_dropdown],
            outputs=answer_output
        )
    
    return demo

# Launch the application
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )