commit

src/streamlit_app.py

@@ -1,428 +1,659 @@
-# app.py - Main Hugging Face Spaces Application
-import gradio as gr
-import PyPDF2
-import pdfplumber
-import fitz  # PyMuPDF
-import pandas as pd
-import re
-import logging
+# Enterprise AI Assistant with RAG using IBM Granite Models
+# Complete implementation with Streamlit interface
+
+import streamlit as st
 import os
+import pandas as pd
+import numpy as np
+from typing import List, Dict, Any, Optional
 import tempfile
-from typing import Dict, List, Tuple, Optional
-from pathlib import Path
 import json
+from datetime import datetime
+import logging
+from pathlib import Path
 
-# Set up logging
+# Core libraries for RAG
+import chromadb
+from sentence_transformers import SentenceTransformer
+import torch
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    pipeline,
+    BitsAndBytesConfig
+)
+
+# Document processing
+import PyPDF2
+from docx import Document
+import openpyxl
+from bs4 import BeautifulSoup
+import email
+from email.mime.text import MIMEText
+import chardet
+
+# Additional utilities
+import re
+from urllib.parse import urlparse
+import hashlib
+import pickle
+
+# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-class PDFProcessorError(Exception):
-    """Custom exception for PDF processing errors"""
-    pass
+# Configuration
+class Config:
+    GRANITE_MODEL_NAME = "ibm-granite/granite-3.1-8b-instruct"
+    GRANITE_GUARDIAN_MODEL = "ibm-granite/granite-guardian-3.2-5b"
+    EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+    CHUNK_SIZE = 1000
+    CHUNK_OVERLAP = 200
+    MAX_CONTEXT_LENGTH = 4000
+    TEMPERATURE = 0.7
+    MAX_NEW_TOKENS = 512
+    TOP_K = 5
+    SUPPORTED_FORMATS = ['.pdf', '.docx', '.xlsx', '.txt', '.csv', '.html', '.json', '.md']
 
-def enhanced_pdf_processor(file_path: str) -> Dict:
-    """
-    Enhanced PDF processor for Hugging Face deployment
-    """
-    results = {
-        'text': '',
-        'tables': [],
-        'metadata': {},
-        'extraction_method': 'unknown',
-        'success': False,
-        'error': None,
-        'file_info': {},
-        'summary': ''
-    }
+class DocumentProcessor:
+    """Handles document processing and text extraction"""
     
-    try:
-        # Validate file
-        if not os.path.exists(file_path):
-            results['error'] = f"File does not exist: {file_path}"
-            return results
-        
-        # Get file info
-        results['file_info'] = get_file_info(file_path)
-        
-        # Try different extraction methods
-        extraction_methods = [
-            ('PyMuPDF', extract_with_pymupdf),
-            ('pdfplumber', extract_with_pdfplumber), 
-            ('PyPDF2', extract_with_pypdf2)
-        ]
-        
-        for method_name, method_func in extraction_methods:
-            try:
-                logger.info(f"Trying extraction method: {method_name}")
-                
-                if method_name == 'pdfplumber':
-                    text_result, tables = method_func(file_path)
-                    if text_result and len(text_result.strip()) > 10:
-                        results['text'] = text_result
-                        results['tables'] = tables
-                        results['extraction_method'] = method_name
-                        results['success'] = True
-                        break
-                        
-                elif method_name == 'PyMuPDF':
-                    text_result, metadata = method_func(file_path)
-                    if text_result and len(text_result.strip()) > 10:
-                        results['text'] = text_result
-                        results['metadata'] = metadata
-                        results['extraction_method'] = method_name
-                        results['success'] = True
-                        break
-                        
-                else:  # PyPDF2
-                    text_result = method_func(file_path)
-                    if text_result and len(text_result.strip()) > 10:
-                        results['text'] = text_result
-                        results['extraction_method'] = method_name
-                        results['success'] = True
-                        break
-                        
-            except Exception as e:
-                logger.warning(f"{method_name} failed: {str(e)}")
-                continue
-        
-        # Generate summary if successful
-        if results['success']:
-            results['summary'] = generate_document_summary(results['text'])
-        else:
-            results['error'] = "All extraction methods failed"
-            
-    except Exception as e:
-        results['error'] = f"Processing error: {str(e)}"
-        logger.error(f"PDF processing error: {e}")
+    @staticmethod
+    def extract_text_from_pdf(file_path: str) -> str:
+        """Extract text from PDF files"""
+        try:
+            with open(file_path, 'rb') as file:
+                reader = PyPDF2.PdfReader(file)
+                text = ""
+                for page in reader.pages:
+                    text += page.extract_text() + "\n"
+                return text
+        except Exception as e:
+            logger.error(f"Error extracting text from PDF: {e}")
+            return ""
     
-    return results
-
-def extract_with_pypdf2(file_path: str) -> str:
-    """Extract text using PyPDF2"""
-    text = ""
-    try:
-        with open(file_path, 'rb') as file:
-            reader = PyPDF2.PdfReader(file)
-            
-            if reader.is_encrypted:
-                try:
-                    reader.decrypt("")
-                except:
-                    raise PDFProcessorError("PDF is encrypted")
+    @staticmethod
+    def extract_text_from_docx(file_path: str) -> str:
+        """Extract text from DOCX files"""
+        try:
+            doc = Document(file_path)
+            text = ""
+            for paragraph in doc.paragraphs:
+                text += paragraph.text + "\n"
+            return text
+        except Exception as e:
+            logger.error(f"Error extracting text from DOCX: {e}")
+            return ""
+    
+    @staticmethod
+    def extract_text_from_xlsx(file_path: str) -> str:
+        """Extract text from Excel files"""
+        try:
+            workbook = openpyxl.load_workbook(file_path)
+            text = ""
+            for sheet_name in workbook.sheetnames:
+                sheet = workbook[sheet_name]
+                text += f"Sheet: {sheet_name}\n"
+                for row in sheet.iter_rows(values_only=True):
+                    row_text = " | ".join([str(cell) if cell else "" for cell in row])
+                    if row_text.strip():
+                        text += row_text + "\n"
+                text += "\n"
+            return text
+        except Exception as e:
+            logger.error(f"Error extracting text from XLSX: {e}")
+            return ""
+    
+    @staticmethod
+    def extract_text_from_csv(file_path: str) -> str:
+        """Extract text from CSV files"""
+        try:
+            df = pd.read_csv(file_path)
+            return df.to_string()
+        except Exception as e:
+            logger.error(f"Error extracting text from CSV: {e}")
+            return ""
+    
+    @staticmethod
+    def extract_text_from_html(file_path: str) -> str:
+        """Extract text from HTML files"""
+        try:
+            with open(file_path, 'r', encoding='utf-8') as file:
+                soup = BeautifulSoup(file.read(), 'html.parser')
+                return soup.get_text()
+        except Exception as e:
+            logger.error(f"Error extracting text from HTML: {e}")
+            return ""
+    
+    @staticmethod
+    def extract_text_from_txt(file_path: str) -> str:
+        """Extract text from TXT files"""
+        try:
+            # Detect encoding
+            with open(file_path, 'rb') as file:
+                raw_data = file.read()
+                encoding = chardet.detect(raw_data)['encoding']
             
-            for page_num, page in enumerate(reader.pages):
-                try:
-                    page_text = page.extract_text()
-                    if page_text:
-                        text += f"\n--- Page {page_num + 1} ---\n{page_text}\n"
-                except Exception as e:
-                    logger.warning(f"Failed to extract page {page_num + 1}: {e}")
-        
-        return clean_text(text)
-        
-    except Exception as e:
-        raise PDFProcessorError(f"PyPDF2 extraction failed: {e}")
-
-def extract_with_pdfplumber(file_path: str) -> Tuple[str, List[Dict]]:
-    """Extract text and tables using pdfplumber"""
-    text = ""
-    tables = []
+            with open(file_path, 'r', encoding=encoding) as file:
+                return file.read()
+        except Exception as e:
+            logger.error(f"Error extracting text from TXT: {e}")
+            return ""
     
-    try:
-        with pdfplumber.open(file_path) as pdf:
-            for page_num, page in enumerate(pdf.pages):
-                try:
-                    # Extract text
-                    page_text = page.extract_text()
-                    if page_text:
-                        text += f"\n--- Page {page_num + 1} ---\n{page_text}\n"
-                    
-                    # Extract tables
-                    page_tables = page.extract_tables()
-                    for table_num, table in enumerate(page_tables):
-                        if table and len(table) > 1:
-                            tables.append({
-                                'page': page_num + 1,
-                                'table_number': table_num + 1,
-                                'data': table,
-                                'text_representation': table_to_text(table)
-                            })
-                            
-                except Exception as e:
-                    logger.warning(f"Failed to process page {page_num + 1}: {e}")
-        
-        return clean_text(text), tables
-        
-    except Exception as e:
-        raise PDFProcessorError(f"pdfplumber extraction failed: {e}")
+    @staticmethod
+    def extract_text_from_json(file_path: str) -> str:
+        """Extract text from JSON files"""
+        try:
+            with open(file_path, 'r') as file:
+                data = json.load(file)
+                return json.dumps(data, indent=2)
+        except Exception as e:
+            logger.error(f"Error extracting text from JSON: {e}")
+            return ""
+    
+    def process_document(self, file_path: str) -> str:
+        """Process document based on file extension"""
+        file_extension = Path(file_path).suffix.lower()
+        
+        extractors = {
+            '.pdf': self.extract_text_from_pdf,
+            '.docx': self.extract_text_from_docx,
+            '.xlsx': self.extract_text_from_xlsx,
+            '.csv': self.extract_text_from_csv,
+            '.html': self.extract_text_from_html,
+            '.txt': self.extract_text_from_txt,
+            '.md': self.extract_text_from_txt,
+            '.json': self.extract_text_from_json,
+        }
+        
+        extractor = extractors.get(file_extension)
+        if extractor:
+            return extractor(file_path)
+        else:
+            logger.warning(f"Unsupported file format: {file_extension}")
+            return ""
 
-def extract_with_pymupdf(file_path: str) -> Tuple[str, Dict]:
-    """Extract text using PyMuPDF"""
-    text = ""
-    metadata = {}
+class TextChunker:
+    """Handles text chunking for RAG"""
+    
+    def __init__(self, chunk_size: int = Config.CHUNK_SIZE, chunk_overlap: int = Config.CHUNK_OVERLAP):
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
     
-    try:
-        doc = fitz.open(file_path)
+    def chunk_text(self, text: str, document_name: str = "") -> List[Dict[str, Any]]:
+        """Split text into chunks with metadata"""
+        chunks = []
+        sentences = re.split(r'[.!?]+', text)
         
-        # Extract metadata
-        try:
-            doc_metadata = doc.metadata or {}
-            metadata = {
-                'page_count': doc.page_count,
-                'title': doc_metadata.get('title', ''),
-                'author': doc_metadata.get('author', ''),
-                'subject': doc_metadata.get('subject', ''),
-                'creator': doc_metadata.get('creator', ''),
-                'creation_date': doc_metadata.get('creationDate', '')
-            }
-        except Exception as e:
-            metadata = {'page_count': doc.page_count}
+        current_chunk = ""
+        current_length = 0
         
-        # Extract text
-        for page_num in range(doc.page_count):
-            try:
-                page = doc[page_num]
-                page_text = page.get_text()
-                if page_text:
-                    text += f"\n--- Page {page_num + 1} ---\n{page_text}\n"
-            except Exception as e:
-                logger.warning(f"Failed to extract page {page_num + 1}: {e}")
+        for sentence in sentences:
+            sentence = sentence.strip()
+            if not sentence:
+                continue
+                
+            sentence_length = len(sentence)
+            
+            if current_length + sentence_length > self.chunk_size and current_chunk:
+                # Save current chunk
+                chunks.append({
+                    'text': current_chunk.strip(),
+                    'metadata': {
+                        'document_name': document_name,
+                        'chunk_id': len(chunks),
+                        'timestamp': datetime.now().isoformat()
+                    }
+                })
+                
+                # Start new chunk with overlap
+                overlap_text = current_chunk[-self.chunk_overlap:] if len(current_chunk) > self.chunk_overlap else current_chunk
+                current_chunk = overlap_text + " " + sentence
+                current_length = len(current_chunk)
+            else:
+                current_chunk += " " + sentence
+                current_length += sentence_length
         
-        doc.close()
-        return clean_text(text), metadata
+        # Add the last chunk
+        if current_chunk.strip():
+            chunks.append({
+                'text': current_chunk.strip(),
+                'metadata': {
+                    'document_name': document_name,
+                    'chunk_id': len(chunks),
+                    'timestamp': datetime.now().isoformat()
+                }
+            })
         
-    except Exception as e:
-        raise PDFProcessorError(f"PyMuPDF extraction failed: {e}")
+        return chunks
 
-def clean_text(text: str) -> str:
-    """Clean extracted text"""
-    if not text:
-        return ""
+class VectorStore:
+    """Handles vector storage and retrieval using ChromaDB"""
     
-    # Remove excessive whitespace
-    text = re.sub(r'\n\s*\n', '\n\n', text)
-    text = re.sub(r' +', ' ', text)
-    
-    # Remove problematic characters
-    text = text.replace('\ufffd', '')
-    text = text.replace('\x00', '')
-    text = text.replace('\u200b', '')
+    def __init__(self, collection_name: str = "enterprise_documents"):
+        self.client = chromadb.PersistentClient(path="./chroma_db")
+        self.collection_name = collection_name
+        self.embedding_model = SentenceTransformer(Config.EMBEDDING_MODEL_NAME)
+        
+        # Create or get collection
+        try:
+            self.collection = self.client.get_collection(collection_name)
+        except:
+            self.collection = self.client.create_collection(
+                name=collection_name,
+                metadata={"description": "Enterprise document embeddings"}
+            )
     
-    return text.strip()
-
-def table_to_text(table: List[List]) -> str:
-    """Convert table to text"""
-    if not table:
-        return ""
+    def add_documents(self, chunks: List[Dict[str, Any]]) -> None:
+        """Add document chunks to vector store"""
+        texts = [chunk['text'] for chunk in chunks]
+        metadatas = [chunk['metadata'] for chunk in chunks]
+        
+        # Generate embeddings
+        embeddings = self.embedding_model.encode(texts).tolist()
+        
+        # Generate IDs
+        ids = [f"doc_{i}_{hashlib.md5(text.encode()).hexdigest()[:8]}" 
+               for i, text in enumerate(texts)]
+        
+        # Add to collection
+        self.collection.add(
+            documents=texts,
+            embeddings=embeddings,
+            metadatas=metadatas,
+            ids=ids
+        )
+        
+        logger.info(f"Added {len(chunks)} chunks to vector store")
     
-    text_lines = []
-    for row in table:
-        if row:
-            clean_row = [str(cell).strip() if cell else "" for cell in row]
-            if any(clean_row):
-                text_lines.append(" | ".join(clean_row))
+    def similarity_search(self, query: str, k: int = Config.TOP_K) -> List[Dict[str, Any]]:
+        """Search for similar documents"""
+        query_embedding = self.embedding_model.encode([query]).tolist()
+        
+        results = self.collection.query(
+            query_embeddings=query_embedding,
+            n_results=k,
+            include=['documents', 'metadatas', 'distances']
+        )
+        
+        search_results = []
+        for i in range(len(results['documents'][0])):
+            search_results.append({
+                'text': results['documents'][0][i],
+                'metadata': results['metadatas'][0][i],
+                'distance': results['distances'][0][i] if results['distances'] else 0
+            })
+        
+        return search_results
     
-    return "\n".join(text_lines)
-
-def get_file_info(file_path: str) -> Dict:
-    """Get file information"""
-    try:
-        path = Path(file_path)
-        stat = path.stat()
+    def get_collection_stats(self) -> Dict[str, Any]:
+        """Get statistics about the collection"""
+        count = self.collection.count()
         return {
-            'name': path.name,
-            'size': stat.st_size,
-            'size_mb': round(stat.st_size / (1024 * 1024), 2)
+            'total_documents': count,
+            'collection_name': self.collection_name
         }
-    except Exception:
-        return {}
 
-def generate_document_summary(text: str) -> str:
-    """Generate a simple document summary"""
-    if not text:
-        return "No text extracted"
-    
-    # Basic statistics
-    words = len(text.split())
-    lines = len(text.split('\n'))
-    chars = len(text)
+class GraniteModel:
+    """Handles IBM Granite model loading and inference"""
     
-    # Extract first few sentences for preview
-    sentences = re.split(r'[.!?]+', text)
-    preview = '. '.join(sentences[:3]).strip()
-    if len(preview) > 300:
-        preview = preview[:300] + "..."
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.guardian_pipeline = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+    @st.cache_resource
+    def load_model(_self):
+        """Load Granite model with caching"""
+        try:
+            # Configure for efficient loading
+            if _self.device == "cuda":
+                quantization_config = BitsAndBytesConfig(
+                    load_in_4bit=True,
+                    bnb_4bit_compute_dtype=torch.float16,
+                    bnb_4bit_use_double_quant=True,
+                    bnb_4bit_quant_type="nf4"
+                )
+            else:
+                quantization_config = None
+            
+            # Load tokenizer
+            _self.tokenizer = AutoTokenizer.from_pretrained(
+                Config.GRANITE_MODEL_NAME,
+                trust_remote_code=True
+            )
+            
+            # Load model
+            _self.model = AutoModelForCausalLM.from_pretrained(
+                Config.GRANITE_MODEL_NAME,
+                quantization_config=quantization_config,
+                device_map="auto" if _self.device == "cuda" else None,
+                torch_dtype=torch.float16 if _self.device == "cuda" else torch.float32,
+                trust_remote_code=True
+            )
+            
+            # Load Guardian model for safety
+            try:
+                _self.guardian_pipeline = pipeline(
+                    "text-classification",
+                    model=Config.GRANITE_GUARDIAN_MODEL,
+                    device=0 if _self.device == "cuda" else -1
+                )
+                logger.info("Granite Guardian model loaded successfully")
+            except Exception as e:
+                logger.warning(f"Could not load Guardian model: {e}")
+                _self.guardian_pipeline = None
+            
+            logger.info(f"Granite model loaded successfully on {_self.device}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error loading Granite model: {e}")
+            return False
     
-    return f"""
-Document Statistics:
-- Characters: {chars:,}
-- Words: {words:,}
-- Lines: {lines:,}
-
-Preview:
-{preview}
-"""
-
-def process_pdf_file(file) -> Tuple[str, str, str, str]:
-    """
-    Process uploaded PDF file for Gradio interface
-    """
-    if file is None:
-        return "No file uploaded", "", "", ""
+    def check_safety(self, text: str) -> bool:
+        """Check if text is safe using Guardian model"""
+        if not self.guardian_pipeline:
+            return True  # If no guardian model, assume safe
+        
+        try:
+            result = self.guardian_pipeline(text)
+            # Assuming Guardian returns safety classification
+            return result[0]['label'].lower() == 'safe'
+        except Exception as e:
+            logger.warning(f"Error in safety check: {e}")
+            return True  # Default to safe if error
     
-    try:
-        # Create temporary file
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
-            tmp_file.write(file.read())
-            tmp_file_path = tmp_file.name
+    def generate_response(self, prompt: str, context: str = "") -> str:
+        """Generate response using Granite model"""
+        if not self.model or not self.tokenizer:
+            if not self.load_model():
+                return "Error: Could not load the model. Please check your setup."
         
-        # Process the PDF
-        result = enhanced_pdf_processor(tmp_file_path)
+        # Safety check
+        if not self.check_safety(prompt):
+            return "I cannot provide a response to that query due to safety concerns."
         
-        # Clean up
-        os.unlink(tmp_file_path)
+        # Construct the full prompt
+        system_prompt = """You are an Enterprise AI Assistant with access to company documents and policies. 
+        Provide helpful, accurate, and professional responses based on the provided context.
+        If you cannot answer based on the context, say so clearly."""
         
-        if result['success']:
-            # Format results for display
-            status = f"✅ Successfully processed using {result['extraction_method']}"
-            
-            # File info
-            file_info = result.get('file_info', {})
-            info = f"""
-File: {file_info.get('name', 'Unknown')}
-Size: {file_info.get('size_mb', 0)} MB
-Pages: {result.get('metadata', {}).get('page_count', 'Unknown')}
-"""
-            
-            # Summary
-            summary = result.get('summary', 'No summary available')
+        if context:
+            full_prompt = f"{system_prompt}\n\nContext:\n{context}\n\nUser Question: {prompt}\n\nAssistant:"
+        else:
+            full_prompt = f"{system_prompt}\n\nUser Question: {prompt}\n\nAssistant:"
+        
+        try:
+            # Tokenize input
+            inputs = self.tokenizer.encode(full_prompt, return_tensors='pt')
             
-            # Full text (truncated for display)
-            full_text = result['text']
-            if len(full_text) > 5000:
-                display_text = full_text[:5000] + f"\n\n... (Text truncated. Total length: {len(full_text)} characters)"
-            else:
-                display_text = full_text
+            # Truncate if too long
+            if inputs.shape[1] > Config.MAX_CONTEXT_LENGTH:
+                inputs = inputs[:, -Config.MAX_CONTEXT_LENGTH:]
             
-            # Tables info
-            if result['tables']:
-                tables_info = f"\n\nTables found: {len(result['tables'])}"
-                for i, table in enumerate(result['tables'][:3]):  # Show first 3 tables
-                    tables_info += f"\n\nTable {i+1} (Page {table['page']}):\n"
-                    tables_info += table['text_representation'][:500]
-                    if len(table['text_representation']) > 500:
-                        tables_info += "..."
-                display_text += tables_info
+            inputs = inputs.to(self.device)
             
-            return status, info, summary, display_text
+            # Generate response
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs,
+                    max_new_tokens=Config.MAX_NEW_TOKENS,
+                    temperature=Config.TEMPERATURE,
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.eos_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    repetition_penalty=1.1
+                )
             
-        else:
-            error_msg = result.get('error', 'Unknown error')
-            return f"❌ Processing failed: {error_msg}", "", "", ""
+            # Decode response
+            response = self.tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
+            return response.strip()
             
-    except Exception as e:
-        return f"❌ Error: {str(e)}", "", "", ""
+        except Exception as e:
+            logger.error(f"Error generating response: {e}")
+            return f"I apologize, but I encountered an error while generating a response: {str(e)}"
 
-def answer_question(text: str, question: str) -> str:
-    """
-    Simple keyword-based question answering
-    """
-    if not text or not question:
-        return "Please provide both text and a question."
+class EnterpriseRAGAssistant:
+    """Main RAG Assistant class"""
     
-    # Convert to lowercase for searching
-    text_lower = text.lower()
-    question_lower = question.lower()
+    def __init__(self):
+        self.doc_processor = DocumentProcessor()
+        self.text_chunker = TextChunker()
+        self.vector_store = VectorStore()
+        self.granite_model = GraniteModel()
+        
+    def process_and_store_documents(self, uploaded_files) -> Dict[str, Any]:
+        """Process uploaded files and store in vector database"""
+        results = {
+            'processed_files': [],
+            'errors': [],
+            'total_chunks': 0
+        }
+        
+        for uploaded_file in uploaded_files:
+            try:
+                # Save uploaded file temporarily
+                with tempfile.NamedTemporaryFile(delete=False, suffix=Path(uploaded_file.name).suffix) as tmp_file:
+                    tmp_file.write(uploaded_file.read())
+                    tmp_file_path = tmp_file.name
+                
+                # Extract text
+                text = self.doc_processor.process_document(tmp_file_path)
+                
+                if text:
+                    # Chunk text
+                    chunks = self.text_chunker.chunk_text(text, uploaded_file.name)
+                    
+                    # Store in vector database
+                    self.vector_store.add_documents(chunks)
+                    
+                    results['processed_files'].append({
+                        'name': uploaded_file.name,
+                        'chunks': len(chunks),
+                        'text_length': len(text)
+                    })
+                    results['total_chunks'] += len(chunks)
+                else:
+                    results['errors'].append(f"Could not extract text from {uploaded_file.name}")
+                
+                # Clean up temporary file
+                os.unlink(tmp_file_path)
+                
+            except Exception as e:
+                results['errors'].append(f"Error processing {uploaded_file.name}: {str(e)}")
+        
+        return results
     
-    # Extract keywords from question
-    keywords = [word for word in question_lower.split() if len(word) > 3]
+    def answer_query(self, query: str) -> Dict[str, Any]:
+        """Answer user query using RAG"""
+        # Retrieve relevant documents
+        search_results = self.vector_store.similarity_search(query)
+        
+        # Prepare context
+        context = "\n\n".join([result['text'] for result in search_results])
+        
+        # Generate response
+        response = self.granite_model.generate_response(query, context)
+        
+        return {
+            'response': response,
+            'sources': search_results,
+            'context_used': bool(context)
+        }
+
+def main():
+    """Main Streamlit application"""
+    st.set_page_config(
+        page_title="Enterprise AI Assistant with RAG",
+        page_icon="🏢",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
     
-    # Find relevant sentences
-    sentences = re.split(r'[.!?]+', text)
-    relevant_sentences = []
+    # Custom CSS
+    st.markdown("""
+    <style>
+    .main-header {
+        font-size: 2.5rem;
+        color: #1f4e79;
+        text-align: center;
+        margin-bottom: 2rem;
+    }
+    .stButton > button {
+        background-color: #0f62fe;
+        color: white;
+        font-weight: bold;
+    }
+    .success-box {
+        padding: 1rem;
+        border-radius: 0.5rem;
+        background-color: #d4edda;
+        border: 1px solid #c3e6cb;
+        color: #155724;
+    }
+    .error-box {
+        padding: 1rem;
+        border-radius: 0.5rem;
+        background-color: #f8d7da;
+        border: 1px solid #f5c6cb;
+        color: #721c24;
+    }
+    </style>
+    """, unsafe_allow_html=True)
     
-    for sentence in sentences:
-        sentence_lower = sentence.lower()
-        score = sum(1 for keyword in keywords if keyword in sentence_lower)
-        if score > 0:
-            relevant_sentences.append((sentence.strip(), score))
+    # Initialize session state
+    if 'rag_assistant' not in st.session_state:
+        st.session_state.rag_assistant = EnterpriseRAGAssistant()
     
-    # Sort by relevance and take top 3
-    relevant_sentences.sort(key=lambda x: x[1], reverse=True)
-    top_sentences = [sent[0] for sent in relevant_sentences[:3]]
+    if 'chat_history' not in st.session_state:
+        st.session_state.chat_history = []
     
-    if top_sentences:
-        return f"Based on the document, here are the most relevant sections:\n\n" + "\n\n".join(top_sentences)
-    else:
-        return "I couldn't find information related to your question in the document."
-
-# Global variable to store extracted text
-extracted_text = ""
-
-def update_extracted_text(status, info, summary, full_text):
-    """Update global extracted text variable"""
-    global extracted_text
-    extracted_text = full_text
-    return status, info, summary, full_text
-
-def qa_interface(question):
-    """Interface for question answering"""
-    global extracted_text
-    return answer_question(extracted_text, question)
-
-# Create Gradio interface
-with gr.Blocks(title="PDF Processor & Q&A System") as app:
-    gr.Markdown("# 📄 PDF Processor & Question Answering System")
-    gr.Markdown("Upload a PDF file to extract text and ask questions about its content.")
+    # Header
+    st.markdown('<h1 class="main-header">🏢 Enterprise AI Assistant with RAG</h1>', unsafe_allow_html=True)
+    st.markdown("**Powered by IBM Granite Models | Intelligent Document Processing & Q&A**")
     
-    with gr.Tab("PDF Processing"):
-        with gr.Row():
-            with gr.Column():
-                file_input = gr.File(label="Upload PDF", file_types=[".pdf"])
-                process_btn = gr.Button("Process PDF", variant="primary")
+    # Sidebar
+    with st.sidebar:
+        st.header("📁 Document Management")
+        
+        # File upload
+        uploaded_files = st.file_uploader(
+            "Upload Enterprise Documents",
+            type=['pdf', 'docx', 'xlsx', 'txt', 'csv', 'html', 'json', 'md'],
+            accept_multiple_files=True,
+            help="Upload documents to build your knowledge base"
+        )
+        
+        if uploaded_files and st.button("Process Documents", type="primary"):
+            with st.spinner("Processing documents..."):
+                results = st.session_state.rag_assistant.process_and_store_documents(uploaded_files)
             
-            with gr.Column():
-                status_output = gr.Textbox(label="Status", lines=2)
-                info_output = gr.Textbox(label="File Information", lines=4)
+            if results['processed_files']:
+                st.markdown('<div class="success-box">', unsafe_allow_html=True)
+                st.success(f"Successfully processed {len(results['processed_files'])} files!")
+                st.write(f"Total chunks created: {results['total_chunks']}")
+                
+                for file_info in results['processed_files']:
+                    st.write(f"✓ {file_info['name']}: {file_info['chunks']} chunks")
+                st.markdown('</div>', unsafe_allow_html=True)
+            
+            if results['errors']:
+                st.markdown('<div class="error-box">', unsafe_allow_html=True)
+                st.error("Some files had errors:")
+                for error in results['errors']:
+                    st.write(f"✗ {error}")
+                st.markdown('</div>', unsafe_allow_html=True)
         
-        summary_output = gr.Textbox(label="Document Summary", lines=8)
-        text_output = gr.Textbox(label="Extracted Text", lines=15, max_lines=20)
-    
-    with gr.Tab("Question & Answer"):
-        gr.Markdown("Ask questions about the processed PDF content.")
-        with gr.Row():
-            question_input = gr.Textbox(label="Your Question", placeholder="What is this document about?")
-            ask_btn = gr.Button("Ask Question", variant="primary")
+        # Database stats
+        st.header("📊 Knowledge Base Stats")
+        try:
+            stats = st.session_state.rag_assistant.vector_store.get_collection_stats()
+            st.metric("Total Documents", stats['total_documents'])
+        except:
+            st.metric("Total Documents", 0)
         
-        answer_output = gr.Textbox(label="Answer", lines=8)
+        # Model info
+        st.header("🤖 Model Information")
+        st.info(f"**Main Model**: {Config.GRANITE_MODEL_NAME}")
+        st.info(f"**Safety Model**: {Config.GRANITE_GUARDIAN_MODEL}")
+        st.info(f"**Embedding Model**: {Config.EMBEDDING_MODEL_NAME}")
     
-    # Event handlers
-    process_btn.click(
-        fn=process_pdf_file,
-        inputs=[file_input],
-        outputs=[status_output, info_output, summary_output, text_output]
-    ).then(
-        fn=update_extracted_text,
-        inputs=[status_output, info_output, summary_output, text_output],
-        outputs=[status_output, info_output, summary_output, text_output]
-    )
+    # Main content area
+    col1, col2 = st.columns([2, 1])
     
-    ask_btn.click(
-        fn=qa_interface,
-        inputs=[question_input],
-        outputs=[answer_output]
-    )
+    with col1:
+        st.header("💬 Chat with Your Documents")
+        
+        # Chat interface
+        query = st.text_input(
+            "Ask a question about your documents:",
+            placeholder="e.g., What is our company's policy on remote work?",
+            key="user_query"
+        )
+        
+        if st.button("Send Query", type="primary") and query:
+            with st.spinner("Generating response..."):
+                result = st.session_state.rag_assistant.answer_query(query)
+            
+            # Add to chat history
+            st.session_state.chat_history.append({
+                'query': query,
+                'response': result['response'],
+                'sources': result['sources'],
+                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            })
+        
+        # Display chat history
+        if st.session_state.chat_history:
+            st.header("📜 Chat History")
+            
+            for i, chat in enumerate(reversed(st.session_state.chat_history)):
+                with st.expander(f"Q: {chat['query'][:50]}... ({chat['timestamp']})", expanded=i==0):
+                    st.markdown("**Question:**")
+                    st.write(chat['query'])
+                    
+                    st.markdown("**Answer:**")
+                    st.write(chat['response'])
+                    
+                    if chat['sources']:
+                        st.markdown("**Sources:**")
+                        for j, source in enumerate(chat['sources'][:3]):
+                            st.markdown(f"**Source {j+1}** (from {source['metadata']['document_name']}):")
+                            st.text(source['text'][:200] + "...")
+    
+    with col2:
+        st.header("🔍 Search Results")
+        
+        if st.session_state.chat_history:
+            latest_chat = st.session_state.chat_history[-1]
+            
+            st.subheader("Latest Query Sources")
+            for i, source in enumerate(latest_chat['sources']):
+                with st.expander(f"Source {i+1}: {source['metadata']['document_name']}"):
+                    st.write(f"**Relevance Score**: {1 - source['distance']:.3f}")
+                    st.write(f"**Document**: {source['metadata']['document_name']}")
+                    st.write(f"**Chunk ID**: {source['metadata']['chunk_id']}")
+                    st.text_area("Content", source['text'], height=150, disabled=True)
+        
+        # Quick actions
+        st.header("⚡ Quick Actions")
+        if st.button("Clear Chat History"):
+            st.session_state.chat_history = []
+            st.rerun()
+        
+        if st.button("Reset Knowledge Base"):
+            if st.confirm("Are you sure you want to reset the knowledge base? This cannot be undone."):
+                # This would require implementing a reset method
+                st.warning("Knowledge base reset functionality would be implemented here")
     
-    # Example
-    gr.Examples(
-        examples=[
-            ["What is the main topic of this document?"],
-            ["What are the key findings?"],
-            ["Who are the authors?"],
-            ["What is the conclusion?"]
-        ],
-        inputs=[question_input]
+    # Footer
+    st.markdown("---")
+    st.markdown(
+        "Built with ❤️ using IBM Granite Models, Streamlit, and ChromaDB | "
+        "Enterprise-grade AI Assistant for document processing and intelligent Q&A"
     )
 
 if __name__ == "__main__":
-    app.launch()
+    main()