Update src/app.py

src/app.py

@@ -1,224 +1,397 @@
 """
-Main Streamlit Application for Multimodal RAG LLM System
+Multimodal RAG LLM System - Streamlit App
+Complete working version with VISUAL image analysis using gpt-4o
 """
+
 import streamlit as st
 import os
-import tempfile
 from pathlib import Path
+
+# Import optimized versions
 from pdf_parser import PDFParser
 from vector_store import VectorStore
-from rag_system import MultimodalRAG
+from rag_system import VisualMultimodalRAG  # NEW - Vision model
 from config import UPLOAD_FOLDER, MAX_PDF_SIZE_MB
-import json
 
 
-def initialize_session_state():
-    """Initialize Streamlit session state"""
-    if 'vector_store' not in st.session_state:
-        st.session_state.vector_store = VectorStore()
-    
-    if 'pdf_parser' not in st.session_state:
-        st.session_state.pdf_parser = PDFParser()
-    
-    if 'rag_system' not in st.session_state:
-        st.session_state.rag_system = None
-    
-    if 'api_key_set' not in st.session_state:
-        st.session_state.api_key_set = False
-    
-    if 'processed_documents' not in st.session_state:
-        st.session_state.processed_documents = st.session_state.pdf_parser.processed_files
+# ============================================================================
+# PAGE CONFIGURATION
+# ============================================================================
 
+st.set_page_config(
+    page_title="📄 Multimodal RAG LLM System",
+    page_icon="🤖",
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
 
-def load_existing_documents():
-    """Load documents that were processed in previous runs"""
-    docs = st.session_state.pdf_parser.get_all_documents()
-    
-    for doc_id, doc_data in docs.items():
-        # Check if this document is already in vector store
-        st.session_state.vector_store.add_documents(doc_data, doc_id)
-
-
-def process_uploaded_pdf(pdf_file):
-    """Process uploaded PDF file"""
-    try:
-        # Save temporary file
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
-            tmp_file.write(pdf_file.getbuffer())
-            tmp_path = tmp_file.name
-        
-        # Check file size
-        file_size_mb = os.path.getsize(tmp_path) / (1024 * 1024)
-        if file_size_mb > MAX_PDF_SIZE_MB:
-            st.error(f"File size {file_size_mb:.2f}MB exceeds maximum {MAX_PDF_SIZE_MB}MB")
-            os.remove(tmp_path)
-            return False
-        
-        # Parse PDF
-        with st.spinner(f"Processing {pdf_file.name}..."):
-            text, images, tables = st.session_state.pdf_parser.parse_pdf(tmp_path)
-            
-            # Add to vector store
-            doc_data = {
-                'text': text,
-                'images': images,
-                'tables': tables
-            }
-            doc_id = Path(pdf_file.name).stem
-            st.session_state.vector_store.add_documents(doc_data, doc_id)
-            st.session_state.vector_store.persist()
-            
-            # Store document info
-            st.session_state.processed_documents[doc_id] = st.session_state.pdf_parser.processed_files.get(doc_id, '')
-        
-        # Clean up
-        os.remove(tmp_path)
-        
-        st.success(f"Successfully processed: {pdf_file.name}")
-        st.info(f"Extracted: {len(text)} characters, {len(images)} images, {len(tables)} tables")
-        return True
-    
-    except Exception as e:
-        st.error(f"Error processing PDF: {str(e)}")
-        return False
+# ============================================================================
+# SESSION STATE INITIALIZATION
+# ============================================================================
+
+if 'api_key_set' not in st.session_state:
+    st.session_state.api_key_set = False
+
+if 'api_key' not in st.session_state:
+    st.session_state.api_key = None
+
+if 'visual_rag_system' not in st.session_state:  # NEW - Vision model
+    st.session_state.visual_rag_system = None
+
+if 'vector_store' not in st.session_state:
+    st.session_state.vector_store = None
+
+if 'parser' not in st.session_state:
+    st.session_state.parser = None
 
+if 'current_document' not in st.session_state:
+    st.session_state.current_document = None
 
-def main():
-    st.set_page_config(page_title="Multimodal RAG LLM System", layout="wide")
-    st.title("📄 Multimodal RAG LLM System")
-    st.markdown("Extract, analyze, and query PDF documents with text, images, and tables using AI")
+if 'current_text' not in st.session_state:
+    st.session_state.current_text = None
+
+if 'current_images' not in st.session_state:
+    st.session_state.current_images = None
+
+if 'current_tables' not in st.session_state:
+    st.session_state.current_tables = None
+
+if 'processing_results' not in st.session_state:  # NEW
+    st.session_state.processing_results = None
+
+
+# ============================================================================
+# MAIN HEADER
+# ============================================================================
+
+st.title("📄 Multimodal RAG LLM System")
+st.markdown("""
+Process PDF documents with visual image analysis:
+- **PDF Parser** with OCR for Russian & English
+- **Visual Analysis** (gpt-4o) for image understanding
+- **Vector Store** (ChromaDB) for semantic search
+- **Individual Component** summarization and storage
+""")
+
+
+# ============================================================================
+# SIDEBAR - CONFIGURATION
+# ============================================================================
+
+with st.sidebar:
+    st.header("⚙️ Configuration")
     
-    # Initialize session state
-    initialize_session_state()
+    # API Key Section
+    st.subheader("🔑 OpenAI API Key")
     
-    # Sidebar for settings
-    with st.sidebar:
-        st.header("⚙️ Settings")
-        
-        # API Key input (masked)
-        api_key = st.text_input(
-            "Enter OpenAI API Key",
-            type="password",
-            help="Your OpenAI API key for GPT-4o access"
-        )
-        
-        if api_key:
-            st.session_state.rag_system = MultimodalRAG(api_key=api_key)
-            st.session_state.api_key_set = True
-            st.success("✅ API Key set")
-        
-        st.divider()
-        
-        # Vector store info
-        st.subheader("📊 Vector Store Status")
-        collection_info = st.session_state.vector_store.get_collection_info()
-        st.metric("Documents in Store", collection_info.get('count', 0))
-        
-        st.divider()
-        
-        # Document management
-        st.subheader("📁 Document Management")
-        if st.button("Refresh Document List"):
-            load_existing_documents()
-            st.success("Document list refreshed")
-        
-        if st.session_state.processed_documents:
-            st.write("**Processed Documents:**")
-            for doc_id in st.session_state.processed_documents.keys():
-                col1, col2 = st.columns([3, 1])
-                col1.write(f"📌 {doc_id}")
-                if col2.button("🗑️", key=f"delete_{doc_id}"):
-                    st.session_state.vector_store.delete_by_doc_id(doc_id)
-                    del st.session_state.processed_documents[doc_id]
-                    st.rerun()
+    api_key = st.text_input(
+        "Enter your OpenAI API key:",
+        type="password",
+        key="api_key_input"
+    )
     
-    # Main area
-    st.header("📤 Upload PDF Document")
-    col1, col2 = st.columns([2, 1])
-
-    if st.button("⚙️ Summarize & Store Each Component"):
-        results = rag_system.process_and_store_document(
-            text=st.session_state.current_text,
-            images=st.session_state.current_images,
-            tables=st.session_state.current_tables,
-            vector_store=vector_store,
-            doc_id=st.session_state.current_document
-        )
+    if api_key:
+        st.session_state.api_key = api_key
+        st.session_state.api_key_set = True
+        
+        # Initialize RAG systems if not already done
+        if st.session_state.visual_rag_system is None:
+            try:
+                st.session_state.visual_rag_system = VisualMultimodalRAG(api_key=api_key, debug=True)  # NEW
+                st.session_state.vector_store = VectorStore()
+                st.session_state.parser = PDFParser(debug=True)
+                st.success("✅ API Key set & systems initialized")
+            except Exception as e:
+                st.error(f"Error initializing systems: {e}")
+    else:
+        st.session_state.api_key_set = False
+        st.warning("⚠️ Please enter your API key to continue")
     
-    with col1:
-        uploaded_file = st.file_uploader(
-            "Choose a PDF file",
-            type="pdf",
-            help=f"Maximum file size: {MAX_PDF_SIZE_MB}MB"
-        )
+    st.divider()
     
-    if uploaded_file:
-        if process_uploaded_pdf(uploaded_file):
-            st.rerun()
+    # Vector Store Status
+    st.subheader("📊 Vector Store Status")
+    if st.session_state.vector_store:
+        try:
+            info = st.session_state.vector_store.get_collection_info()
+            st.metric("Items in Store", info['count'])
+            st.metric("Status", info['status'])
+            st.caption(f"Path: {info['persist_path']}")
+        except Exception as e:
+            st.error(f"Error getting store info: {e}")
+    else:
+        st.info("Set API key to initialize vector store")
     
     st.divider()
     
-    # Q&A Section
-    if st.session_state.api_key_set:
-        st.header("❓ Ask Questions About Your Documents")
-        
-        # Check if there are documents in vector store
-        collection_info = st.session_state.vector_store.get_collection_info()
-        if collection_info.get('count', 0) == 0:
-            st.warning("No documents in vector store. Please upload a PDF first.")
+    # Document Management
+    st.subheader("📁 Document Management")
+    if st.button("🔄 Clear Vector Store"):
+        if st.session_state.vector_store:
+            try:
+                st.session_state.vector_store.clear_all()
+                st.success("✅ Vector store cleared")
+            except Exception as e:
+                st.error(f"Error clearing store: {e}")
+
+
+# ============================================================================
+# MAIN CONTENT
+# ============================================================================
+
+# Upload Section
+st.header("📤 Upload PDF Document")
+
+uploaded_file = st.file_uploader(
+    "Choose a PDF file",
+    type=['pdf'],
+    help="PDF with text, images, and tables"
+)
+
+if uploaded_file is not None:
+    # Save uploaded file
+    upload_path = Path(UPLOAD_FOLDER)
+    upload_path.mkdir(exist_ok=True)
+    
+    file_path = upload_path / uploaded_file.name
+    with open(file_path, 'wb') as f:
+        f.write(uploaded_file.getbuffer())
+    
+    st.success(f"✅ File saved: {uploaded_file.name}")
+    
+    # Parse PDF
+    if st.button("🔍 Parse PDF"):
+        if not st.session_state.api_key_set:
+            st.error("❌ Please set OpenAI API key first")
         else:
-            # Question input
-            question = st.text_area(
-                "Enter your question",
-                placeholder="Ask anything about your PDF documents...",
-                height=100
-            )
-            
-            col1, col2 = st.columns([1, 4])
-            
-            with col1:
-                if st.button("🔍 Search & Answer", type="primary"):
-                    if question.strip():
-                        with st.spinner("Searching and generating answer..."):
-                            # Search vector store
-                            search_results = st.session_state.vector_store.search(question, n_results=5)
+            try:
+                with st.spinner("📄 Parsing PDF..."):
+                    print(f"\n{'='*70}")
+                    print(f"PARSING: {uploaded_file.name}")
+                    print(f"{'='*70}")
+                    
+                    # Parse PDF - returns text, images, tables
+                    parser = st.session_state.parser
+                    text, images, tables = parser.parse_pdf(str(file_path))
+                    
+                    # Store in session state
+                    st.session_state.current_document = uploaded_file.name
+                    st.session_state.current_text = text
+                    st.session_state.current_images = images
+                    st.session_state.current_tables = tables
+                    
+                    # Display results
+                    col1, col2, col3 = st.columns(3)
+                    with col1:
+                        st.metric("📝 Text", f"{len(text):,} chars")
+                    with col2:
+                        st.metric("🖼️ Images", len(images))
+                    with col3:
+                        st.metric("📋 Tables", len(tables))
+                    
+                    # Show image OCR details
+                    if images:
+                        st.subheader("🖼️ Extracted Images")
+                        for idx, img in enumerate(images):
+                            ocr_text = img.get('ocr_text', '')
+                            ocr_len = len(ocr_text)
                             
-                            if search_results:
-                                # Generate answer
-                                answer = st.session_state.rag_system.answer_question(question, search_results)
-                                
-                                # Display results
-                                st.subheader("🤖 Answer")
-                                st.write(answer)
-                                
-                                # Display search results
-                                with st.expander("📚 Source Documents"):
-                                    for idx, result in enumerate(search_results, 1):
-                                        st.write(f"**[{result.get('type', 'unknown').upper()}]**")
-                                        st.write(result.get('content', ''))
-                                        st.divider()
+                            if ocr_len > 0:
+                                st.success(f"✅ Image {idx}: {ocr_len} characters (OCR)")
                             else:
-                                st.warning("No relevant documents found. Try different search terms.")
-                    else:
-                        st.warning("Please enter a question first.")
+                                st.warning(f"⚠️ Image {idx}: No OCR text (will use visual analysis)")
+                    
+                    st.success("✅ PDF parsing complete!")
+                    
+            except Exception as e:
+                st.error(f"❌ Error parsing PDF: {e}")
+                print(f"Error: {e}")
+
+
+# ============================================================================
+# VISUAL IMAGE ANALYSIS & COMPONENT STORAGE
+# ============================================================================
+
+st.divider()
+st.header("🖼️ Visual Analysis & Storage")
+
+st.info("""
+**How it works:**
+1. Images are sent to gpt-4o for visual analysis (not just text OCR)
+2. Text is split into chunks and each chunk is summarized
+3. Tables are analyzed individually
+4. ALL summaries are stored in the vector store for semantic search
+""")
+
+if st.button("🖼️ Analyze Images Visually & Store Components"):
+    if not st.session_state.api_key_set:
+        st.error("❌ Please set OpenAI API key first")
+    elif st.session_state.current_text is None:
+        st.error("❌ Please parse a PDF document first")
     else:
-        st.warning("⚠️ Please enter your OpenAI API Key in the sidebar to use the Q&A feature.")
-    
-    st.divider()
-    
-    # Footer
-    st.markdown("---")
-    st.markdown(
-        """
-        **About this system:**
-        - Extracts text, images, and tables from PDF documents
-        - Uses CLIP-ViT embeddings for multimodal understanding
-        - Stores documents locally in ChromaDB vector store
-        - Analyzes queries using OpenAI GPT-4o
-        - Supports Russian language documents
-        """
-    )
+        try:
+            with st.spinner("🖼️ Analyzing images visually with gpt-4o..."):
+                print(f"\n{'='*70}")
+                print(f"VISUAL IMAGE ANALYSIS")
+                print(f"{'='*70}")
+                
+                # Process with visual analysis
+                visual_rag = st.session_state.visual_rag_system
+                vector_store = st.session_state.vector_store
+                
+                results = visual_rag.process_and_store_document(
+                    text=st.session_state.current_text,
+                    images=st.session_state.current_images,    # Actual images sent to gpt-4o
+                    tables=st.session_state.current_tables,
+                    vector_store=vector_store,
+                    doc_id=st.session_state.current_document or "current_doc"
+                )
+                
+                st.session_state.processing_results = results
+                
+                # Display results
+                st.success("✅ Visual analysis complete & stored!")
+                
+                col1, col2, col3 = st.columns(3)
+                with col1:
+                    st.metric("🖼️ Images Analyzed", len(results['image_visual_analyses']))
+                with col2:
+                    st.metric("📝 Text Chunks", len(results['text_summaries']))
+                with col3:
+                    st.metric("📋 Tables Analyzed", len(results['table_summaries']))
+                
+                st.metric("📊 Total Stored in Vector", results['total_stored'])
+                
+                # Show image visual analyses
+                if results['image_visual_analyses']:
+                    st.subheader("🖼️ Visual Image Analyses (gpt-4o)")
+                    for img_analysis in results['image_visual_analyses']:
+                        with st.expander(f"Image {img_analysis['image_index']} - Visual Analysis"):
+                            st.write("**Visual Analysis by gpt-4o:**")
+                            st.write(img_analysis['visual_analysis'])
+                            
+                            st.write("**Image Path:**")
+                            st.code(img_analysis['image_path'])
+                            
+                            if img_analysis['ocr_text']:
+                                st.write("**OCR Text (backup):**")
+                                st.text(img_analysis['ocr_text'][:500])
+                
+                # Show text chunk summaries
+                if results['text_summaries']:
+                    st.subheader("📝 Text Chunk Summaries")
+                    for chunk_summary in results['text_summaries']:
+                        with st.expander(
+                            f"Chunk {chunk_summary['chunk_index']} "
+                            f"({chunk_summary['chunk_length']} chars)"
+                        ):
+                            st.write("**Summary:**")
+                            st.write(chunk_summary['summary'])
+                            st.write("**Original Text (first 500 chars):**")
+                            st.text(chunk_summary['original_text'])
+                
+                # Show table analyses
+                if results['table_summaries']:
+                    st.subheader("📋 Table Analyses")
+                    for table_summary in results['table_summaries']:
+                        with st.expander(
+                            f"Table {table_summary['table_index']} "
+                            f"({table_summary['table_length']} chars)"
+                        ):
+                            st.write("**Analysis:**")
+                            st.write(table_summary['summary'])
+                            st.write("**Original Content (first 500 chars):**")
+                            st.text(table_summary['original_content'])
+                
+                print(f"\n✅ Visual analysis processing complete!")
+                
+        except Exception as e:
+            st.error(f"❌ Error during visual analysis: {e}")
+            print(f"Error: {e}")
+
+
+# ============================================================================
+# QUESTION & ANSWERING
+# ============================================================================
+
+st.divider()
+st.header("❓ Ask Questions About Document")
+
+question = st.text_area(
+    "Enter your question:",
+    height=100,
+    placeholder="What does the document say about...?"
+)
+
+if st.button("🔍 Search & Answer"):
+    if not st.session_state.api_key_set:
+        st.error("❌ Please set OpenAI API key first")
+    elif st.session_state.current_text is None:
+        st.error("❌ Please parse a PDF document first")
+    elif not question:
+        st.error("❌ Please enter a question")
+    else:
+        try:
+            with st.spinner("🔄 Searching and generating answer..."):
+                print(f"\n{'='*70}")
+                print(f"QUESTION: {question}")
+                print(f"{'='*70}")
+                
+                # Search vector store
+                store = st.session_state.vector_store
+                
+                # Add documents to store first if not already added
+                doc_name = st.session_state.current_document or "current_doc"
+                doc_data = {
+                    'text': st.session_state.current_text,
+                    'images': [],  # Images already stored via visual analysis
+                    'tables': []   # Tables already stored via visual analysis
+                }
+                store.add_documents(doc_data, doc_name)
+                
+                # Search
+                search_results = store.search(question, n_results=5)
+                
+                # Display results
+                st.write("### Search Results")
+                for idx, result in enumerate(search_results, 1):
+                    content_type = result.get('type', 'unknown')
+                    distance = result.get('distance', 0)
+                    content = result.get('content', '')
+                    
+                    with st.expander(
+                        f"Result {idx} - {content_type.upper()} "
+                        f"(relevance: {1-distance:.2%})"
+                    ):
+                        st.write(content)
+                
+                st.success("✅ Search complete! Use results above to understand the document.")
+                
+        except Exception as e:
+            st.error(f"❌ Error processing question: {e}")
+            print(f"Error: {e}")
+
+
+# ============================================================================
+# FOOTER
+# ============================================================================
+
+st.divider()
+
+col1, col2, col3 = st.columns(3)
+
+with col1:
+    st.info("📖 **Text Processing**: PyPDF2 extraction with UTF-8 support")
+
+with col2:
+    st.info("🖼️ **Visual Analysis**: GPT-4o vision for image understanding")
 
+with col3:
+    st.info("📊 **Vector Storage**: ChromaDB with auto-persist")
 
-if __name__ == "__main__":
-    main()
+st.caption(
+    "Multimodal RAG System | "
+    "Visual Image Analysis | "
+    "Russian Language Support | "
+    "Individual Component Summarization"
+)