Update unified_document_processor.py

unified_document_processor.py

@@ -446,64 +446,64 @@ class UnifiedDocumentProcessor:
                     'error': str(e)
                 }
     
-        def get_available_files(self) -> Dict[str, List[str]]:
-            """Get list of all files in the database"""
-            try:
-                all_entries = self.collection.get(
-                    include=['metadatas']
-                )
-    
-                files = {
-                    'pdf': set(),
-                    'xml': set()
-                }
-    
-                for metadata in all_entries['metadatas']:
-                    file_type = metadata['content_type']
-                    file_name = metadata['source_file']
-                    files[file_type].add(file_name)
-    
-                return {
-                    'pdf': sorted(list(files['pdf'])),
-                    'xml': sorted(list(files['xml']))
-                }
-            except Exception as e:
-                print(f"Error getting available files: {str(e)}")
-                return {'pdf': [], 'xml': []}
-    
-        def ask_question_selective(self, question: str, selected_files: List[str], n_results: int = 5) -> str:
-            """Ask a question using only the selected files"""
-            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"]
-                )
-    
-                if not results['documents'][0]:
-                    return "No relevant content found in the selected files."
-    
-                context = "\n\n".join(results['documents'][0])
-    
-                prompt = f"""Based on the following content from the selected files, please answer this question: {question}
-    
-                Content:
-                {context}
-    
-                Please provide a direct answer based only on the information provided above."""
-    
-                response = self.groq_client.chat.completions.create(
-                    messages=[{"role": "user", "content": prompt}],
-                    model="llama3-8b-8192",
-                    temperature=0.2
-                )
-    
-                return response.choices[0].message.content
-    
-            except Exception as e:
-                return f"Error processing your question: {str(e)}"
+    def get_available_files(self) -> Dict[str, List[str]]:
+        """Get list of all files in the database"""
+        try:
+            all_entries = self.collection.get(
+                include=['metadatas']
+            )
+
+            files = {
+                'pdf': set(),
+                'xml': set()
+            }
+
+            for metadata in all_entries['metadatas']:
+                file_type = metadata['content_type']
+                file_name = metadata['source_file']
+                files[file_type].add(file_name)
+
+            return {
+                'pdf': sorted(list(files['pdf'])),
+                'xml': sorted(list(files['xml']))
+            }
+        except Exception as e:
+            print(f"Error getting available files: {str(e)}")
+            return {'pdf': [], 'xml': []}
+
+    def ask_question_selective(self, question: str, selected_files: List[str], n_results: int = 5) -> str:
+        """Ask a question using only the selected files"""
+        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"]
+            )
+
+            if not results['documents'][0]:
+                return "No relevant content found in the selected files."
+
+            context = "\n\n".join(results['documents'][0])
+
+            prompt = f"""Based on the following content from the selected files, please answer this question: {question}
+
+            Content:
+            {context}
+
+            Please provide a direct answer based only on the information provided above."""
+
+            response = self.groq_client.chat.completions.create(
+                messages=[{"role": "user", "content": prompt}],
+                model="llama3-8b-8192",
+                temperature=0.2
+            )
+
+            return response.choices[0].message.content
+
+        except Exception as e:
+            return f"Error processing your question: {str(e)}"