Update app.py

app.py

@@ -1,11 +1,220 @@
 import streamlit as st
 import os
 import tempfile
-from typing import List
+from typing import List, Dict, Optional, Tuple
 from unified_document_processor import UnifiedDocumentProcessor, CustomEmbeddingFunction
 import chromadb
 from chromadb.config import Settings
 from groq import Groq
+import json
+import xml.etree.ElementTree as ET
+
+class EnhancedXMLProcessor:
+    def __init__(self):
+        self.processed_nodes = set()
+        self.reference_map = {}
+        self.node_info = {}
+
+    def build_reference_map(self, root) -> None:
+        for element in root.findall('.//*'):
+            node_id = element.get('NodeId')
+            if node_id:
+                self.node_info[node_id] = {
+                    'tag': element.tag,
+                    'browse_name': element.get('BrowseName', ''),
+                    'display_name': self._get_display_name(element),
+                    'description': self._get_description(element),
+                    'data_type': element.get('DataType', ''),
+                    'references': []
+                }
+                
+                refs = element.find('References')
+                if refs is not None:
+                    for ref in refs.findall('Reference'):
+                        ref_type = ref.get('ReferenceType')
+                        is_forward = ref.get('IsForward', 'true') == 'true'
+                        target = ref.text
+                        
+                        if ref_type in ['HasComponent', 'HasProperty', 'HasTypeDefinition']:
+                            self.reference_map.setdefault(node_id, []).append({
+                                'type': ref_type,
+                                'target': target,
+                                'is_forward': is_forward
+                            })
+                            self.node_info[node_id]['references'].append({
+                                'type': ref_type,
+                                'target': target,
+                                'is_forward': is_forward
+                            })
+
+    def _get_display_name(self, element) -> str:
+        display_name = element.find('DisplayName')
+        if display_name is not None:
+            return display_name.text
+        return ''
+
+    def _get_description(self, element) -> str:
+        desc = element.find('Description')
+        if desc is not None:
+            return desc.text
+        return ''
+
+    def _get_value(self, element) -> Optional[str]:
+        value_elem = element.find('.//Value')
+        if value_elem is not None:
+            for child in value_elem:
+                if child.text:
+                    return child.text
+        return None
+
+    def generate_natural_language(self, node_id: str, depth: int = 0, visited: set = None) -> List[str]:
+        if visited is None:
+            visited = set()
+            
+        if node_id in visited:
+            return []
+            
+        visited.add(node_id)
+        descriptions = []
+        
+        node = self.node_info.get(node_id)
+        if not node:
+            return []
+
+        base_desc = self._build_base_description(node, depth)
+        if base_desc:
+            descriptions.append(base_desc)
+
+        if node_id in self.reference_map:
+            child_descriptions = self._process_forward_references(node_id, depth + 1, visited)
+            descriptions.extend(child_descriptions)
+
+        return descriptions
+
+    def _build_base_description(self, node: Dict, depth: int) -> str:
+        indentation = "  " * depth
+        desc_parts = []
+
+        if node['browse_name']:
+            browse_name = node['browse_name'].split(':')[-1]
+            desc_parts.append(f"a {browse_name}")
+
+        if node['display_name']:
+            desc_parts.append(f"(displayed as '{node['display_name']}')")
+
+        if node['data_type']:
+            desc_parts.append(f"of type {node['data_type']}")
+
+        if node['description']:
+            desc_parts.append(f"which {node['description']}")
+
+        if desc_parts:
+            return f"{indentation}Contains {' '.join(desc_parts)}"
+        return ""
+
+    def _process_forward_references(self, node_id: str, depth: int, visited: set) -> List[str]:
+        descriptions = []
+        
+        for ref in self.reference_map.get(node_id, []):
+            if ref['is_forward'] and ref['type'] in ['HasComponent', 'HasProperty']:
+                target_descriptions = self.generate_natural_language(ref['target'], depth, visited)
+                descriptions.extend(target_descriptions)
+
+        return descriptions
+
+    def generate_complete_description(self, root) -> str:
+        self.build_reference_map(root)
+        
+        root_descriptions = []
+        for node_id in self.node_info:
+            is_root = True
+            for ref_list in self.reference_map.values():
+                for ref in ref_list:
+                    if not ref['is_forward'] and ref['type'] == 'HasComponent' and ref['target'] == node_id:
+                        is_root = False
+                        break
+                if not is_root:
+                    break
+                    
+            if is_root:
+                descriptions = self.generate_natural_language(node_id)
+                root_descriptions.extend(descriptions)
+
+        return "\n".join(root_descriptions)
+
+    def flatten_xml_to_text(self, element, depth=0) -> str:
+        try:
+            return self.generate_complete_description(element)
+        except Exception as e:
+            print(f"Error in enhanced XML processing: {str(e)}")
+            return self._original_flatten_xml_to_text(element, depth)
+
+    def _original_flatten_xml_to_text(self, element, depth=0) -> str:
+        text_parts = []
+        
+        element_info = f"Element: {element.tag}"
+        if element.attrib:
+            element_info += f", Attributes: {json.dumps(element.attrib)}"
+        if element.text and element.text.strip():
+            element_info += f", Text: {element.text.strip()}"
+        text_parts.append(element_info)
+        
+        for child in element:
+            child_text = self._original_flatten_xml_to_text(child, depth + 1)
+            text_parts.append(child_text)
+            
+        return "\n".join(text_parts)
+
+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()
+        self.xml_processor = EnhancedXMLProcessor()
+        
+        self.chroma_client = chromadb.PersistentClient(
+            path=persist_dir,
+            settings=Settings(
+                allow_reset=True,
+                is_persistent=True
+            )
+        )
+        
+        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()
+            )
+
+    def flatten_xml_to_text(self, element, depth=0) -> str:
+        try:
+            return self.xml_processor.generate_complete_description(element)
+        except Exception as e:
+            st.error(f"Error in enhanced XML processing: {str(e)}")
+            return self._original_flatten_xml_to_text(element, depth)
+
+    def _original_flatten_xml_to_text(self, element, depth=0) -> str:
+        text_parts = []
+        
+        element_info = f"Element: {element.tag}"
+        if element.attrib:
+            element_info += f", Attributes: {json.dumps(element.attrib)}"
+        if element.text and element.text.strip():
+            element_info += f", Text: {element.text.strip()}"
+        text_parts.append(element_info)
+        
+        for child in element:
+            child_text = self._original_flatten_xml_to_text(child, depth + 1)
+            text_parts.append(child_text)
+            
+        return "\n".join(text_parts)
 
 def initialize_session_state():
     """Initialize all session state variables"""
@@ -14,96 +223,63 @@ 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 = dict(pdf=[], xml=[])
+        st.session_state.processed_files = set()
     
     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)}")
+        st.session_state.processor = None
 
 class StreamlitDocProcessor:
     def __init__(self):
+        self.ensure_processor_initialized()
+
+    def ensure_processor_initialized(self):
         if st.session_state.processor is None:
             try:
-                groq_api_key = st.secrets["GROQ_API_KEY"]
-                # Initialize processor with persistent ChromaDB
+                groq_api_key = os.getenv('GROQ_API_KEY')
+                if not groq_api_key:
+                    st.error("GROQ API key not found. Please set the GROQ_API_KEY environment variable.")
+                    return False
+                
                 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()
+                if st.session_state.processor:
+                    st.session_state.processed_files = self.get_processed_files()
+                    return True
             except Exception as e:
                 st.error(f"Error initializing processor: {str(e)}")
-                return
+                return False
+        return True
 
     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"""
+    def get_processed_files(self) -> set:
         try:
             if st.session_state.processor:
-                return st.session_state.processor.get_available_files()
-            return dict(pdf=[], xml=[])
+                available_files = st.session_state.processor.get_available_files()
+                return set(available_files['pdf'] + available_files['xml'])
+            return set()
         except Exception as e:
             st.error(f"Error getting processed files: {str(e)}")
-            return dict(pdf=[], xml=[])
+            return set()
 
     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"]
-        )
+        st.title("AAS Assistant")
         
-        # 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 self.ensure_processor_initialized():
+            page = st.sidebar.selectbox(
+                "Choose a page",
+                ["Upload & Process", "Query"]
+            )
             
-        if page == "Upload & Process":
-            self.upload_and_process_page()
+            if page == "Upload & Process":
+                self.upload_and_process_page()
+            else:
+                self.qa_page()
         else:
-            self.qa_page()
+            st.error("Please set the GROQ_API_KEY environment variable and restart the application.")
 
     def upload_and_process_page(self):
         st.header("Upload and Process Documents")
@@ -113,187 +289,86 @@ class StreamlitDocProcessor:
             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()
+                status_text = 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, []):
+                if uploaded_file.name not in st.session_state.processed_files:
                     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}...')
+
+                        status_text.text(f'Processing {uploaded_file.name}...')
+                        progress_bar.progress(25)
                         
-                        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)
+                        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)
+                            st.session_state.processed_files.add(uploaded_file.name)
                             progress_bar.progress(100)
-                            status_container.success(f"Successfully processed {uploaded_file.name}")
+                            status_text.success(f"Successfully processed {uploaded_file.name}")
                         else:
                             progress_bar.progress(100)
-                            status_container.error(f"Failed to process {uploaded_file.name}: {result['error']}")
-    
+                            status_text.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)}")
+                        status_text.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")
+                    status_text.info(f"{uploaded_file.name} has already been processed")
                     progress_bar.progress(100)
 
+        if st.session_state.processed_files:
+            st.subheader("Processed Files")
+            for file in sorted(st.session_state.processed_files):
+                st.text(f"✓ {file}")
 
     def qa_page(self):
-        st.header("Query Documents")
-    
+        st.header("Query our database")
+
         try:
-            # Refresh available files
             st.session_state.processed_files = self.get_processed_files()
             
-            if not any(st.session_state.processed_files.values()):
+            if not st.session_state.processed_files:
                 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
+                sorted(list(st.session_state.processed_files)),
+                default=list(st.session_state.processed_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)}")
-                        
+            if st.button("Ask Question") and question:
+                try:
+                    with st.spinner("Searching for 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)}")
+                    
         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()