Spaces:
Sleeping
Sleeping
| 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") | |
| uploaded_files = st.file_uploader( | |
| "Upload PDF or XML files", | |
| type=['pdf', 'xml'], | |
| accept_multiple_files=True | |
| ) | |
| if uploaded_files: | |
| for uploaded_file in uploaded_files: | |
| # Create progress containers | |
| progress_bar = st.progress(0) | |
| status_container = st.empty() | |
| file_ext = os.path.splitext(uploaded_file.name)[1][1:] | |
| if uploaded_file.name not in st.session_state.processed_files.get(file_ext, []): | |
| try: | |
| 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 | |
| status_container.text(f'Processing {uploaded_file.name}...') | |
| if file_ext == 'xml': | |
| # Add processing status updates | |
| status_container.text('Parsing XML...') | |
| progress_bar.progress(10) | |
| result = st.session_state.processor.process_file(temp_path) | |
| if result['success']: | |
| total_chunks = result['total_chunks'] | |
| for i, chunk_result in enumerate(result['results']): | |
| # Update progress for each batch | |
| progress = min(95, int(10 + (85 * (i / total_chunks)))) | |
| progress_bar.progress(progress) | |
| status_container.text(f'Processing chunk {i+1}/{total_chunks}...') | |
| else: | |
| # Regular PDF processing | |
| result = st.session_state.processor.process_file(temp_path) | |
| 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_container.success(f"Successfully processed {uploaded_file.name}") | |
| else: | |
| progress_bar.progress(100) | |
| status_container.error(f"Failed to process {uploaded_file.name}: {result['error']}") | |
| except Exception as e: | |
| status_container.error(f"Error processing {uploaded_file.name}: {str(e)}") | |
| finally: | |
| try: | |
| os.unlink(temp_path) | |
| except: | |
| pass | |
| else: | |
| status_container.info(f"{uploaded_file.name} has already been processed") | |
| progress_bar.progress(100) | |
| 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 | |
| # File selection | |
| selected_files = st.multiselect( | |
| "Select files to search through", | |
| sorted(all_files), | |
| default=all_files | |
| ) | |
| # 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:") | |
| if question: | |
| col1, col2, col3 = st.columns(3) | |
| with col1: | |
| if st.button("Quick Answer"): | |
| 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"): | |
| 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()})"): | |
| # Use similarity_score instead of relevance_score | |
| st.write(f"Similarity Score: {item['similarity_score']}%") | |
| 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"): | |
| 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()})"): | |
| # Use similarity_score instead of relevance_score | |
| st.write(f"Similarity Score: {item.get('similarity_score', 'N/A')}%") | |
| 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() |