Update query.py

query.py

@@ -1,6 +1,5 @@
 import requests
 import json
-import re
 
 class VectaraQuery():
     def __init__(self, api_key: str, customer_id: str, corpus_id: str, prompt_name: str = None):
@@ -92,18 +91,18 @@ class VectaraQuery():
                         text = result['text']
                         print(f"Processing text: {text}")  # Debugging line
 
-                        # Instead of using regex, split the text by specific keywords
-                        reason = self.extract_between_keywords(text, "Reason:", "Alternative:")
-                        alternative = self.extract_between_keywords(text, "Alternative:", "Notes:")
-                        notes = self.extract_between_keywords(text, "Notes:", "")
+                        # Improved keyword-based extraction
+                        reason = self.extract_text(text, "Reason:", "Alternative:")
+                        alternative = self.extract_text(text, "Alternative:", "Notes:")
+                        notes = self.extract_text(text, "Notes:", "")
 
-                        response = f"Reason: {reason.strip()}\nAlternative: {alternative.strip()}\nNotes: {notes.strip()}"
+                        response = f"Reason: {reason}\nAlternative: {alternative}\nNotes: {notes}"
                         print(f"Generated response: {response}")  # Debugging line
                         return response
 
         return "No relevant information found."
 
-    def extract_between_keywords(self, text, start_keyword, end_keyword):
+    def extract_text(self, text, start_keyword, end_keyword):
         start_idx = text.find(start_keyword)
         if start_idx == -1:
             return "Not available"