Update unified_document_processor.py

unified_document_processor.py

@@ -377,4 +377,170 @@ class UnifiedDocumentProcessor:
             return response.choices[0].message.content
 
         except Exception as e:
-            return f"Error processing your question: {str(e)}"
+            return f"Error processing your question: {str(e)}"
+    def get_detailed_context(self, question: str, selected_files: List[str], n_results: int = 5) -> Dict:
+    """Get detailed context including path and metadata information"""
+    try:
+        filter_dict = {
+            'source_file': {'$in': selected_files}
+        }
+
+        results = self.collection.query(
+            query_texts=[question],
+            n_results=n_results,
+            where=filter_dict,
+            include=["documents", "metadatas", "distances"]
+        )
+
+        if not results['documents'][0]:
+            return {
+                'success': False,
+                'error': "No relevant content found"
+            }
+
+        detailed_results = []
+        for doc, meta, distance in zip(results['documents'][0], results['metadatas'][0], results['distances'][0]):
+            result_info = {
+                'content': doc,
+                'metadata': meta,
+                'relevance_score': 1 - distance,  # Convert distance to similarity score
+                'source_info': {
+                    'file': meta['source_file'],
+                    'type': meta['content_type'],
+                    'path': meta.get('xml_path', 'N/A'),  # Only for XML files
+                    'context': json.loads(meta['context']) if meta.get('context') else {}
+                }
+            }
+            detailed_results.append(result_info)
+
+        return {
+            'success': True,
+            'results': detailed_results,
+            'query': question
+        }
+
+    except Exception as e:
+        return {
+            'success': False,
+            'error': str(e)
+        }
+
+def get_hierarchical_context(self, question: str, selected_files: List[str], n_results: int = 5) -> Dict:
+    """Get hierarchical context for XML files including parent-child relationships"""
+    try:
+        # Get initial results
+        initial_results = self.get_detailed_context(question, selected_files, n_results)
+        
+        if not initial_results['success']:
+            return initial_results
+
+        hierarchical_results = []
+        for result in initial_results['results']:
+            if result['metadata']['content_type'] == 'xml':
+                # Get parent elements
+                parent_path = '/'.join(result['source_info']['path'].split('/')[:-1])
+                if parent_path:
+                    parent_filter = {
+                        'source_file': result['metadata']['source_file'],
+                        'xml_path': parent_path
+                    }
+                    parent_results = self.collection.query(
+                        query_texts=[""],  # Empty query to get exact match
+                        where=parent_filter,
+                        include=["documents", "metadatas"],
+                        n_results=1
+                    )
+                    if parent_results['documents'][0]:
+                        result['parent_info'] = {
+                            'content': parent_results['documents'][0][0],
+                            'metadata': parent_results['metadatas'][0][0]
+                        }
+
+                # Get immediate children
+                child_path_prefix = result['source_info']['path'] + '/'
+                child_filter = {
+                    'source_file': result['metadata']['source_file'],
+                    'xml_path': {'$contains': child_path_prefix}
+                }
+                child_results = self.collection.query(
+                    query_texts=[""],  # Empty query to get exact matches
+                    where=child_filter,
+                    include=["documents", "metadatas"],
+                    n_results=5
+                )
+                if child_results['documents'][0]:
+                    result['children_info'] = [{
+                        'content': doc,
+                        'metadata': meta
+                    } for doc, meta in zip(child_results['documents'][0], child_results['metadatas'][0])]
+
+            hierarchical_results.append(result)
+
+        return {
+            'success': True,
+            'results': hierarchical_results,
+            'query': question
+        }
+
+    except Exception as e:
+        return {
+            'success': False,
+            'error': str(e)
+        }
+
+def get_summary_and_details(self, question: str, selected_files: List[str]) -> Dict:
+    """Get both a summary answer and detailed supporting information"""
+    try:
+        # Get hierarchical context first
+        detailed_results = self.get_hierarchical_context(question, selected_files)
+        
+        if not detailed_results['success']:
+            return detailed_results
+
+        # Create summary prompt
+        relevant_content = []
+        for result in detailed_results['results']:
+            if result['metadata']['content_type'] == 'xml':
+                content_info = [
+                    f"XML Path: {result['source_info']['path']}",
+                    f"Content: {result['content']}"
+                ]
+                if 'parent_info' in result:
+                    content_info.append(f"Parent: {result['parent_info']['content']}")
+                if 'children_info' in result:
+                    children_content = [child['content'] for child in result['children_info']]
+                    content_info.append(f"Related Elements: {', '.join(children_content)}")
+            else:
+                content_info = [f"Content: {result['content']}"]
+
+            relevant_content.append('\n'.join(content_info))
+
+        summary_prompt = f"""Based on the following content, please provide:
+        1. A concise answer to the question
+        2. Key supporting points
+        3. Related context if relevant
+
+        Question: {question}
+
+        Content:
+        {'\n\n'.join(relevant_content)}
+        """
+
+        response = self.groq_client.chat.completions.create(
+            messages=[{"role": "user", "content": summary_prompt}],
+            model="llama3-8b-8192",
+            temperature=0.2
+        )
+
+        return {
+            'success': True,
+            'summary': response.choices[0].message.content,
+            'details': detailed_results['results'],
+            'query': question
+        }
+
+    except Exception as e:
+        return {
+            'success': False,
+            'error': str(e)
+        }