app.py

import streamlit as st
import os
import json
from datetime import datetime
from pathlib import Path
import tempfile
import zipfile
import tkinter as tk
from tkinter import filedialog

from utils.llm_handler import LLMHandler
from utils.rag_pipeline import RAGPipeline
from utils.file_processor import FileProcessor
from utils.web_search import WebSearcher
from config import Config

# Page configuration
st.set_page_config(
    page_title="DeepResearcher",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize session state
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'rag_pipeline' not in st.session_state:
    st.session_state.rag_pipeline = None
if 'llm_handler' not in st.session_state:
    st.session_state.llm_handler = None
if 'file_processor' not in st.session_state:
    st.session_state.file_processor = FileProcessor()
if 'web_searcher' not in st.session_state:
    st.session_state.web_searcher = WebSearcher()
if 'selected_model_path' not in st.session_state:
    st.session_state.selected_model_path = ""

def browse_for_model():
    """Open file browser to select model file"""
    try:
        # Hide the main tkinter window
        root = tk.Tk()
        root.withdraw()
        root.attributes('-topmost', True)
        
        # Open file dialog
        file_path = filedialog.askopenfilename(
            title="Select Model File",
            filetypes=[
                ("Model files", "*.bin *.gguf *.ggml"),
                ("BIN files", "*.bin"),
                ("GGUF files", "*.gguf"),
                ("GGML files", "*.ggml"),
                ("All files", "*.*")
            ]
        )
        
        root.destroy()
        
        if file_path:
            return file_path
        return None
        
    except Exception as e:
        st.error(f"Error opening file browser: {str(e)}")
        return None

def initialize_components():
    """Initialize LLM and RAG components"""
    if not st.session_state.selected_model_path:
        st.warning("Please select a model first.")
        return False
        
    if st.session_state.llm_handler is None:
        try:
            with st.spinner("Loading LLM model..."):
                st.session_state.llm_handler = LLMHandler(
                    model_path=st.session_state.selected_model_path,
                    temperature=st.session_state.temperature,
                    max_tokens=st.session_state.max_tokens,
                    context_length=st.session_state.context_length
                )
        except Exception as e:
            st.error(f"Failed to load model: {str(e)}")
            return False
    
    if st.session_state.rag_pipeline is None:
        try:
            st.session_state.rag_pipeline = RAGPipeline(
                llm_handler=st.session_state.llm_handler
            )
        except Exception as e:
            st.error(f"Failed to initialize RAG pipeline: {str(e)}")
            return False
    
    return True

def sidebar_settings():
    """Render sidebar with settings"""
    st.sidebar.title("⚙️ Settings")
    
    # Model settings
    st.sidebar.subheader("Model Configuration")
    
    # Model selection
    available_models = Config.get_available_models()
    model_names = list(available_models.keys())
    
    if 'selected_model' not in st.session_state:
        st.session_state.selected_model = model_names[0] if model_names else "No models available"
    
    selected_model = st.sidebar.selectbox(
        "Select Model",
        model_names,
        index=model_names.index(st.session_state.selected_model) if st.session_state.selected_model in model_names else 0,
        help="Choose a local LLM model"
    )
    
    # Handle model selection
    if selected_model != st.session_state.get('selected_model'):
        st.session_state.selected_model = selected_model
        
        if selected_model == "📁 Browse for model file...":
            # Open file browser
            model_path = browse_for_model()
            if model_path:
                st.session_state.selected_model_path = model_path
                st.session_state.selected_model = f"Custom: {Path(model_path).name}"
                st.success(f"Selected model: {Path(model_path).name}")
            else:
                st.session_state.selected_model = model_names[0] if len(model_names) > 1 else "No models available"
        else:
            model_path = available_models.get(selected_model, "")
            if model_path and model_path != "BROWSE":
                st.session_state.selected_model_path = model_path
            
        # Reset LLM handler to reload model
        st.session_state.llm_handler = None
        st.session_state.rag_pipeline = None
    
    # Display current model info
    if st.session_state.selected_model_path:
        model_path = Path(st.session_state.selected_model_path)
        if model_path.exists():
            file_size = model_path.stat().st_size / (1024**3)  # GB
            st.sidebar.info(f"📁 **Current Model:**\n{model_path.name}\n📊 Size: {file_size:.1f} GB")
        else:
            st.sidebar.error("❌ Selected model file not found!")
    
    # Add directory scanner
    st.sidebar.subheader("Scan Directory for Models")
    scan_directory = st.sidebar.text_input(
        "Directory Path",
        placeholder="Enter path to scan for models...",
        help="Scan a specific directory for model files"
    )
    
    if st.sidebar.button("🔍 Scan Directory"):
        if scan_directory and os.path.exists(scan_directory):
            found_models = Config.scan_directory_for_models(scan_directory)
            if found_models:
                st.sidebar.success(f"Found {len(found_models)} models!")
                for name, path in found_models.items():
                    st.sidebar.write(f"• {name}")
            else:
                st.sidebar.warning("No models found in directory.")
        else:
            st.sidebar.error("Invalid directory path.")
    
    # Generation parameters
    st.sidebar.subheader("Generation Parameters")
    
    st.session_state.temperature = st.sidebar.slider(
        "Temperature", 0.0, 2.0, 0.7, 0.1,
        help="Controls randomness in generation"
    )
    
    st.session_state.max_tokens = st.sidebar.slider(
        "Max Tokens", 50, 2048, 512, 50,
        help="Maximum tokens to generate"
    )
    
    st.session_state.context_length = st.sidebar.slider(
        "Context Length", 512, 8192, 2048, 256,
        help="Maximum context window size"
    )
    
    # RAG settings
    st.sidebar.subheader("RAG Configuration")
    
    st.session_state.use_web_search = st.sidebar.checkbox(
        "Enable Web Search", True,
        help="Include web search results in responses"
    )
    
    st.session_state.max_web_results = st.sidebar.slider(
        "Max Web Results", 1, 10, 3,
        help="Number of web search results to include"
    )
    
    st.session_state.similarity_threshold = st.sidebar.slider(
        "Similarity Threshold", 0.0, 1.0, 0.7, 0.05,
        help="Minimum similarity for document retrieval"
    )
    
    # File processing settings
    st.sidebar.subheader("File Processing")
    
    st.session_state.chunk_size = st.sidebar.slider(
        "Chunk Size", 100, 2000, 500, 100,
        help="Size of text chunks for processing"
    )
    
    st.session_state.chunk_overlap = st.sidebar.slider(
        "Chunk Overlap", 0, 500, 50, 25,
        help="Overlap between text chunks"
    )

def file_upload_section():
    """Handle file uploads and processing"""
    st.subheader("📁 Document Upload")
    
    # Display processing capabilities
    processor = st.session_state.file_processor
    
    col1, col2 = st.columns(2)
    with col1:
        st.write("**📄 Supported Formats:**")
        st.write("• Documents: PDF, DOCX, TXT")
        st.write("• Images: PNG, JPG, JPEG, GIF")
        st.write("• Audio: MP3, WAV, M4A, FLAC")
        st.write("• Video: MP4, AVI, MOV, MKV")
    
    with col2:
        st.write("**🔧 Processing Status:**")
        capabilities = {
            "PDF/DOCX Processing": "✅" if hasattr(processor, '_process_pdf') else "❌",
            "Image OCR": "✅" if hasattr(processor, '_process_image') else "❌",
            "Audio Transcription": "✅" if processor.whisper_model else "❌",
            "Video Processing": "✅" if hasattr(processor, '_process_video') else "❌"
        }
        
        for capability, status in capabilities.items():
            st.write(f"{status} {capability}")
    
    uploaded_files = st.file_uploader(
        "Upload documents for analysis",
        type=['pdf', 'docx', 'txt', 'png', 'jpg', 'jpeg', 'mp3', 'wav', 'mp4', 'avi'],
        accept_multiple_files=True,
        help="Drag and drop files here or click to browse"
    )
    
    if uploaded_files:
        col1, col2 = st.columns([3, 1])
        
        with col1:
            st.write(f"**{len(uploaded_files)} file(s) uploaded:**")
            for file in uploaded_files:
                file_size = file.size / 1024  # KB
                st.write(f"• {file.name} ({file_size:.1f} KB)")
        
        with col2:
            if st.button("🔄 Process Files", type="primary"):
                process_uploaded_files(uploaded_files)

def process_uploaded_files(uploaded_files):
    """Process uploaded files and add to vector store"""
    if not initialize_components():
        return
        
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    total_files = len(uploaded_files)
    processed_texts = []
    
    for i, uploaded_file in enumerate(uploaded_files):
        status_text.text(f"Processing {uploaded_file.name}...")
        progress_bar.progress((i + 1) / total_files)
        
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=f"_{uploaded_file.name}") as tmp_file:
            tmp_file.write(uploaded_file.getvalue())
            tmp_path = tmp_file.name
        
        try:
            # Process file based on type
            text_content = st.session_state.file_processor.process_file(tmp_path)
            if text_content and text_content.strip():
                processed_texts.append({
                    'content': text_content,
                    'source': uploaded_file.name,
                    'type': uploaded_file.type or 'unknown'
                })
            else:
                st.warning(f"No content extracted from {uploaded_file.name}")
                
        except Exception as e:
            st.error(f"Error processing {uploaded_file.name}: {str(e)}")
        finally:
            # Clean up temporary file
            try:
                os.unlink(tmp_path)
            except:
                pass
    
    # Add processed texts to vector store
    if processed_texts:
        try:
            st.session_state.rag_pipeline.add_documents(processed_texts)
            st.success(f"✅ Successfully processed {len(processed_texts)} files!")
        except Exception as e:
            st.error(f"Error adding documents to knowledge base: {str(e)}")
    else:
        st.warning("⚠️ No content could be extracted from the uploaded files.")
    
    progress_bar.empty()
    status_text.empty()

def chat_interface():
    """Main chat interface"""
    st.subheader("💬 Research Assistant")
    
    # Check if model is loaded
    if not st.session_state.selected_model_path:
        st.warning("⚠️ Please select a model in the sidebar to start chatting.")
        return
    
    # Display chat history
    for i, (query, response, timestamp) in enumerate(st.session_state.chat_history):
        with st.container():
            st.markdown(f"**🧑 You ({timestamp}):**")
            st.markdown(query)
            st.markdown(f"**🤖 DeepResearcher:**")
            st.markdown(response)
            
            # Add re-run button for each query
            col1, col2, col3 = st.columns([1, 1, 8])
            with col1:
                if st.button(f"🔄", key=f"rerun_{i}", help="Re-run this query"):
                    rerun_query(query)
            with col2:
                if st.button(f"📋", key=f"copy_{i}", help="Copy response"):
                    st.code(response)
            
            st.divider()
    
    # Query input
    query = st.text_area(
        "Enter your research question:",
        height=100,
        placeholder="Ask me anything about your uploaded documents or any topic...",
        key="query_input"
    )
    
    col1, col2, col3 = st.columns([2, 1, 1])
    
    with col1:
        if st.button("🔍 Research", type="primary", disabled=not query.strip()):
            if query.strip():
                process_query(query.strip())
    
    with col2:
        if st.button("🗑️ Clear History"):
            st.session_state.chat_history = []
            st.rerun()
    
    with col3:
        if st.button("📥 Export"):
            export_chat_history()

def process_query(query):
    """Process a research query"""
    if not initialize_components():
        return
    
    with st.spinner("🔍 Researching your question..."):
        try:
            # Get response from RAG pipeline
            response = st.session_state.rag_pipeline.query(
                query,
                use_web_search=st.session_state.use_web_search,
                max_web_results=st.session_state.max_web_results,
                similarity_threshold=st.session_state.similarity_threshold
            )
            
            # Add to chat history
            timestamp = datetime.now().strftime("%H:%M:%S")
            st.session_state.chat_history.append((query, response, timestamp))
            
            st.rerun()
            
        except Exception as e:
            st.error(f"❌ Error processing query: {str(e)}")

def rerun_query(query):
    """Re-run a previous query"""
    process_query(query)

def export_chat_history():
    """Export chat history to file"""
    if not st.session_state.chat_history:
        st.warning("No chat history to export.")
        return
    
    # Create export content
    export_content = "# DeepResearcher Chat History\n\n"
    export_content += f"Exported on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
    
    for i, (query, response, timestamp) in enumerate(st.session_state.chat_history, 1):
        export_content += f"## Query {i} ({timestamp})\n\n"
        export_content += f"**Question:** {query}\n\n"
        export_content += f"**Answer:** {response}\n\n"
        export_content += "---\n\n"
    
    # Create download button
    st.download_button(
        label="📄 Download as Markdown",
        data=export_content,
        file_name=f"deepresearcher_history_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md",
        mime="text/markdown"
    )

def vector_store_management():
    """Manage vector store"""
    st.subheader("🗄️ Knowledge Base Management")
    
    if st.session_state.rag_pipeline and st.session_state.rag_pipeline.vector_store:
        # Display vector store stats
        doc_count = st.session_state.rag_pipeline.get_document_count()
        
        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric("📚 Documents", doc_count)
        with col2:
            if st.session_state.rag_pipeline.vector_store:
                stats = st.session_state.rag_pipeline.vector_store.get_statistics()
                st.metric("🔍 Vector Dimension", stats.get('embedding_dimension', 0))
        with col3:
            st.metric("💾 Storage", "Active" if doc_count > 0 else "Empty")
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            if st.button("🗑️ Clear Knowledge Base", type="secondary"):
                st.session_state.rag_pipeline.clear_vector_store()
                st.success("✅ Knowledge base cleared!")
                st.rerun()
        
        with col2:
            if st.button("💾 Save Knowledge Base"):
                # Save vector store to file
                save_path = f"knowledge_base_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pkl"
                try:
                    st.session_state.rag_pipeline.save_vector_store(save_path)
                    st.success(f"✅ Knowledge base saved to {save_path}")
                except Exception as e:
                    st.error(f"❌ Error saving: {str(e)}")
        
        with col3:
            uploaded_kb = st.file_uploader(
                "Load Knowledge Base",
                type=['pkl'],
                help="Upload a previously saved knowledge base"
            )
            if uploaded_kb:
                try:
                    # Save uploaded file temporarily and load
                    with tempfile.NamedTemporaryFile(delete=False, suffix='.pkl') as tmp_file:
                        tmp_file.write(uploaded_kb.getvalue())
                        st.session_state.rag_pipeline.load_vector_store(tmp_file.name)
                        os.unlink(tmp_file.name)
                    st.success("✅ Knowledge base loaded!")
                    st.rerun()
                except Exception as e:
                    st.error(f"❌ Error loading: {str(e)}")
    else:
        st.info("📝 No documents in knowledge base yet. Upload some files to get started!")

def main():
    """Main application"""
    st.title("🔬 DeepResearcher")
    st.markdown("*Advanced AI Research Assistant with Local LLM and RAG*")
    
    # Create necessary directories
    Config.create_directories()
    
    # Sidebar settings
    sidebar_settings()
    
    # Main content tabs
    tab1, tab2, tab3, tab4 = st.tabs(["💬 Chat", "📁 Upload", "🗄️ Knowledge Base", "ℹ️ About"])
    
    with tab1:
        chat_interface()
    
    with tab2:
        file_upload_section()
    
    with tab3:
        vector_store_management()
    
    with tab4:
        st.markdown("""
        ## About DeepResearcher
        
        DeepResearcher is an advanced AI research assistant that combines:
        
        - **🤖 Local LLM Integration**: Uses local language models for privacy and control
        - **📚 RAG Pipeline**: Retrieval-Augmented Generation for accurate, contextual responses
        - **🎯 Multi-modal Processing**: Supports PDF, DOCX, images, audio, and video files
        - **🌐 Web Search**: Integrates live web search for up-to-date information
        - **🗄️ Vector Database**: Efficient similarity search using FAISS
        
        ### 🔧 System Requirements:
        - **Python**: 3.8 or higher
        - **RAM**: 8GB+ recommended for 7B models
        - **Storage**: 5-10GB for model files
        - **Dependencies**: Tesseract OCR, FFmpeg
        
        ### 📋 Supported Model Formats:
        - **GGUF**: Latest format (recommended)
        - **BIN**: Legacy GGML format
        - **GGML**: Original format
        
        ### 🎯 Processing Capabilities:
        - **Documents**: Text extraction from PDF, DOCX, TXT
        - **Images**: OCR text extraction from PNG, JPG, etc.
        - **Audio**: Speech-to-text transcription
        - **Video**: Audio transcription + frame OCR
        
        ### 🚀 Performance Tips:
        - Use quantized models (Q4_0, Q5_0) for faster inference
        - Adjust context length based on your hardware
        - Enable GPU acceleration if available
        """)

if __name__ == "__main__":
    main()