Update app.py

app.py

@@ -14,7 +14,7 @@ def initialize_session_state():
         os.makedirs(st.session_state.CHROMADB_DIR, exist_ok=True)
     
     if 'processed_files' not in st.session_state:
-        st.session_state.processed_files = set()
+        st.session_state.processed_files = dict(pdf=[], xml=[])
     
     if 'processor' not in st.session_state:
         try:
@@ -71,19 +71,18 @@ class StreamlitDocProcessor:
             persist_dir=st.session_state.CHROMADB_DIR
         )
 
-    def get_processed_files(self) -> set:
+    def get_processed_files(self) -> dict:
         """Get list of processed files from ChromaDB"""
         try:
             if st.session_state.processor:
-                available_files = st.session_state.processor.get_available_files()
-                return set(available_files['pdf'] + available_files['xml'])
-            return set()
+                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 set()
+            return dict(pdf=[], xml=[])
 
     def run(self):
-        st.title("AAS Assistant")
+        st.title("Document Assistant")
         
         # Create sidebar for navigation
         page = st.sidebar.selectbox(
@@ -112,7 +111,9 @@ class StreamlitDocProcessor:
                 progress_bar = st.progress(0)
                 status_text = st.empty()
                 
-                if uploaded_file.name not in st.session_state.processed_files:
+                # 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:
@@ -127,7 +128,9 @@ class StreamlitDocProcessor:
                         progress_bar.progress(75)
                         
                         if result['success']:
-                            st.session_state.processed_files.add(uploaded_file.name)
+                            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(f"Successfully processed {uploaded_file.name}")
                         else:
@@ -147,50 +150,54 @@ class StreamlitDocProcessor:
                     progress_bar.progress(100)
 
         # Display processed files
-        if st.session_state.processed_files:
+        if any(st.session_state.processed_files.values()):
             st.subheader("Processed Files")
-            for file in sorted(st.session_state.processed_files):
-                st.text(f"✓ {file}")
+            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}")
+            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 our database")
+        st.header("Query Documents")
 
         try:
             # Refresh available files
             st.session_state.processed_files = self.get_processed_files()
             
-            if not st.session_state.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
 
-            # Enhanced file selection with type indicators
-            available_files = self.get_processed_files()
-            xml_files = [f"📱 {f}" for f in available_files['xml']]
-            pdf_files = [f"📄 {f}" for f in available_files['pdf']]
-            all_files = sorted(xml_files + pdf_files)
+            # 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",
-                all_files,
-                default=all_files,
-                help="📱 = XML files, 📄 = PDF files"
+                sorted(all_files),
+                default=all_files
             )
 
-            # Clean up the file names (remove emojis) for processing
-            selected_files = [f[2:] for f in selected_files]  # Remove emoji prefix
+            # 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 with suggested prompts for XML
-            xml_selected = any(f.endswith('.xml') for f in selected_files)
-            if xml_selected:
-                st.info("Suggested questions for XML content:\n" +
-                    "• What are the main components and their relationships?\n" +
-                    "• What data types and properties are defined?\n" +
-                    "• How are the elements structured and organized?")
-
+            # Question input
             question = st.text_input("Enter your question:")
             
             if st.button("Ask Question") and question:
@@ -200,28 +207,12 @@ class StreamlitDocProcessor:
                             question,
                             selected_files
                         )
-                        
-                        # Display the answer in a structured way
                         st.write("Answer:", answer)
-                        
-                        # If XML files were queried, show additional metadata
-                        if xml_selected:
-                            with st.expander("Show XML Structure Details"):
-                                st.write("Related XML Elements:")
-                                # Get the structure information from the processor
-                                xml_details = st.session_state.processor.get_xml_structure_info(
-                                    selected_files,
-                                    question
-                                )
-                                for detail in xml_details:
-                                    st.code(detail, language="xml")
-                        
                 except Exception as e:
                     st.error(f"Error getting answer: {str(e)}")
                     
         except Exception as e:
             st.error(f"Error in Q&A interface: {str(e)}")
-            
 
 def main():
     # Initialize session state