app.py

import streamlit as st
import os
import tempfile
from typing import List
from unified_document_processor import UnifiedDocumentProcessor, CustomEmbeddingFunction
import chromadb
from chromadb.config import Settings
from groq import Groq

def initialize_session_state():
    """Initialize all session state variables"""
    if 'CHROMADB_DIR' not in st.session_state:
        st.session_state.CHROMADB_DIR = os.path.join(os.getcwd(), 'chromadb_data')
        os.makedirs(st.session_state.CHROMADB_DIR, exist_ok=True)
    
    if 'processed_files' not in st.session_state:
        st.session_state.processed_files = dict(pdf=[], xml=[])
    
    if 'processor' not in st.session_state:
        try:
            st.session_state.processor = None  # Will be initialized in StreamlitDocProcessor
        except Exception as e:
            st.error(f"Error initializing processor: {str(e)}")

class StreamlitDocProcessor:
    def __init__(self):
        if st.session_state.processor is None:
            try:
                groq_api_key = st.secrets["GROQ_API_KEY"]
                # Initialize processor with persistent ChromaDB
                st.session_state.processor = self.initialize_processor(groq_api_key)
                # Update processed files after initializing processor
                st.session_state.processed_files = self.get_processed_files()
            except Exception as e:
                st.error(f"Error initializing processor: {str(e)}")
                return

    def initialize_processor(self, groq_api_key):
        """Initialize the processor with persistent ChromaDB"""
        class PersistentUnifiedDocumentProcessor(UnifiedDocumentProcessor):
            def __init__(self, api_key, collection_name="unified_content", persist_dir=None):
                self.groq_client = Groq(api_key=api_key)
                self.max_elements_per_chunk = 50
                self.pdf_chunk_size = 500
                self.pdf_overlap = 50
                self._initialize_nltk()
                
                # Initialize persistent ChromaDB
                self.chroma_client = chromadb.PersistentClient(
                    path=persist_dir,
                    settings=Settings(
                        allow_reset=True,
                        is_persistent=True
                    )
                )
                
                # Get or create collection
                try:
                    self.collection = self.chroma_client.get_collection(
                        name=collection_name,
                        embedding_function=CustomEmbeddingFunction()
                    )
                except:
                    self.collection = self.chroma_client.create_collection(
                        name=collection_name,
                        embedding_function=CustomEmbeddingFunction()
                    )

        return PersistentUnifiedDocumentProcessor(
            groq_api_key, 
            persist_dir=st.session_state.CHROMADB_DIR
        )

    def get_processed_files(self) -> dict:
        """Get list of processed files from ChromaDB"""
        try:
            if st.session_state.processor:
                return st.session_state.processor.get_available_files()
            return dict(pdf=[], xml=[])
        except Exception as e:
            st.error(f"Error getting processed files: {str(e)}")
            return dict(pdf=[], xml=[])

    def run(self):
        st.title("Document Assistant")
        
        # Create sidebar for navigation
        st.sidebar.title("Navigation")
        page = st.sidebar.selectbox(
            "Choose a page",
            ["Upload & Process", "Query"]
        )
        
        # Add sidebar information
        with st.sidebar.expander("About"):
            st.write("""
            This application allows you to:
            - Upload PDF and XML documents
            - Process them for semantic search
            - Query the documents with different levels of detail
            """)
            
        if page == "Upload & Process":
            self.upload_and_process_page()
        else:
            self.qa_page()

    def upload_and_process_page(self):
        st.header("Upload and Process Documents")
        
        # Add instructions
        with st.expander("Instructions", expanded=True):
            st.write("""
            1. Click 'Browse files' to select documents
            2. You can select multiple files at once
            3. Supported formats: PDF and XML
            4. Wait for processing to complete
            5. Processed files will be listed below
            """)
        
        # File uploader
        uploaded_files = st.file_uploader(
            "Upload PDF or XML files",
            type=['pdf', 'xml'],
            accept_multiple_files=True,
            help="Select one or more PDF or XML files to upload"
        )

        if uploaded_files:
            for uploaded_file in uploaded_files:
                # Create progress bar and status container
                col1, col2 = st.columns([3, 1])
                with col1:
                    progress_bar = st.progress(0)
                with col2:
                    status_text = st.empty()
                
                # Check if file is already processed
                file_ext = os.path.splitext(uploaded_file.name)[1][1:]  # Get extension without dot
                if uploaded_file.name not in st.session_state.processed_files.get(file_ext, []):
                    try:
                        # Create a temporary file
                        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
                            tmp_file.write(uploaded_file.getbuffer())
                            temp_path = tmp_file.name

                        # Process the file
                        status_text.info('Processing...')
                        progress_bar.progress(25)
                        
                        result = st.session_state.processor.process_file(temp_path)
                        progress_bar.progress(75)
                        
                        if result['success']:
                            if file_ext not in st.session_state.processed_files:
                                st.session_state.processed_files[file_ext] = []
                            st.session_state.processed_files[file_ext].append(uploaded_file.name)
                            progress_bar.progress(100)
                            status_text.success("✓ Success")
                        else:
                            progress_bar.progress(100)
                            status_text.error("✗ Failed")
                            st.error(f"Failed to process {uploaded_file.name}: {result['error']}")

                    except Exception as e:
                        status_text.error("✗ Error")
                        st.error(f"Error processing {uploaded_file.name}: {str(e)}")
                    finally:
                        # Clean up temporary file
                        try:
                            os.unlink(temp_path)
                        except:
                            pass
                else:
                    status_text.info("Already processed")
                    progress_bar.progress(100)

        # Display processed files
        if any(st.session_state.processed_files.values()):
            st.subheader("Processed Files")
            
            col1, col2 = st.columns(2)
            
            with col1:
                if st.session_state.processed_files.get('xml'):
                    st.write("📱 XML Files:")
                    for file in sorted(st.session_state.processed_files['xml']):
                        st.text(f"  • {file}")
                        
            with col2:
                if st.session_state.processed_files.get('pdf'):
                    st.write("📄 PDF Files:")
                    for file in sorted(st.session_state.processed_files['pdf']):
                        st.text(f"  • {file}")

    def qa_page(self):
        st.header("Query Documents")

        try:
            # Refresh available files
            st.session_state.processed_files = self.get_processed_files()
            
            if not any(st.session_state.processed_files.values()):
                st.warning("No processed files available. Please upload and process some files first.")
                return

            # Create combined list of files with icons
            all_files = []
            for file in st.session_state.processed_files.get('xml', []):
                all_files.append(f"📱 {file}")
            for file in st.session_state.processed_files.get('pdf', []):
                all_files.append(f"📄 {file}")

            if not all_files:
                st.warning("No processed files available. Please upload and process some files first.")
                return

            # Add query instructions
            with st.expander("Query Instructions", expanded=True):
                st.write("""
                Choose your query type:
                - **Quick Answer**: Basic response with essential information
                - **Detailed Answer**: Shows sources and relevance with expandable details
                - **Complete Analysis**: Provides summary and full breakdown with XML hierarchies
                """)

            # File selection
            selected_files = st.multiselect(
                "Select files to search through",
                sorted(all_files),
                default=all_files,
                help="Choose which files to include in your search"
            )

            # Remove icons from selected files
            selected_files = [f.split(' ', 1)[1] for f in selected_files]

            if not selected_files:
                st.warning("Please select at least one file to search through.")
                return

            # Question input
            question = st.text_input(
                "Enter your question:",
                help="Type your question here and choose a query type below"
            )
            
            if question:
                col1, col2, col3 = st.columns(3)
                
                with col1:
                    if st.button("Quick Answer", help="Get a concise answer quickly"):
                        try:
                            with st.spinner("Getting quick answer..."):
                                answer = st.session_state.processor.ask_question_selective(
                                    question,
                                    selected_files
                                )
                                st.write("Answer:", answer)
                        except Exception as e:
                            st.error(f"Error getting answer: {str(e)}")

                with col2:
                    if st.button("Detailed Answer", help="Get answer with sources and relevance scores"):
                        try:
                            with st.spinner("Getting detailed answer..."):
                                result = st.session_state.processor.get_detailed_context(
                                    question,
                                    selected_files
                                )
                                if result['success']:
                                    st.write("### Relevant Information")
                                    for item in result['results']:
                                        with st.expander(f"Source: {item['metadata']['source_file']} ({item['metadata']['content_type'].upper()})"):
                                            st.write(f"Relevance Score: {item['relevance_score']:.2f}")
                                            if item['metadata']['content_type'] == 'xml':
                                                st.write(f"XML Path: {item['source_info']['path']}")
                                            st.write("Content:", item['content'])
                                else:
                                    st.error(result['error'])
                        except Exception as e:
                            st.error(f"Error getting detailed answer: {str(e)}")

                with col3:
                    if st.button("Complete Analysis", help="Get comprehensive analysis with XML hierarchy"):
                        try:
                            with st.spinner("Performing complete analysis..."):
                                result = st.session_state.processor.get_summary_and_details(
                                    question,
                                    selected_files
                                )
                                if result['success']:
                                    st.write("### Summary")
                                    st.write(result['summary'])
                                    
                                    st.write("### Detailed Information")
                                    for item in result['details']:
                                        with st.expander(f"Source: {item['metadata']['source_file']} ({item['metadata']['content_type'].upper()})"):
                                            st.write(f"Relevance Score: {item['relevance_score']:.2f}")
                                            if item['metadata']['content_type'] == 'xml':
                                                st.write(f"XML Path: {item['source_info']['path']}")
                                                if 'parent_info' in item:
                                                    st.write("Parent Element:", item['parent_info']['content'])
                                                if 'children_info' in item:
                                                    st.write("Related Elements:")
                                                    for child in item['children_info']:
                                                        st.write(f"- {child['content']}")
                                            st.write("Content:", item['content'])
                                else:
                                    st.error(result['error'])
                        except Exception as e:
                            st.error(f"Error getting complete analysis: {str(e)}")
                    
        except Exception as e:
            st.error(f"Error in Q&A interface: {str(e)}")

def main():
    # Set page config
    st.set_page_config(
        page_title="Document Assistant",
        page_icon="📚",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Initialize session state
    initialize_session_state()
    
    # Create and run app
    app = StreamlitDocProcessor()
    app.run()

if __name__ == "__main__":
    main()