Update query.py

query.py

@@ -20,7 +20,11 @@ class VectaraQuery():
         [
             {{
                 "role": "system",
-                "content": "You are an assistant that provides information about drink names based on a given corpus."
+                "content": "You are an assistant that provides information about drink names based on a given corpus. \
+                Format the response in the following way:\n\
+                Reason: <reason why the name cannot be used>\n\
+                Alternative: <alternative name>\n\
+                Notes: <additional notes>"
             }},
             {{
                 "role": "user",
@@ -88,25 +92,29 @@ class VectaraQuery():
                         text = result['text']
                         print(f"Processing text: {text}")  # Debugging line
 
-                        # Adjusting regex patterns to be more flexible
-                        reason_match = re.search(r"Reason Why it Can't be Used:\s*(.*?)(?:\n|$)", text, re.DOTALL)
-                        alternative_match = re.search(r"Alternative:\s*(.*?)(?:\n|$)", text, re.DOTALL)
-                        notes_match = re.search(r"Notes:\s*(.*?)(?:\n|$)", text, re.DOTALL)
-                        
-                        # Improved regex to capture multiline fields
-                        if not reason_match:
-                            reason_match = re.search(r"DISCUSSION\s*-\s*(.*?)(?=\n|\r\n)", text, re.DOTALL)
-                        if not alternative_match:
-                            alternative_match = re.search(r"Alternative\s*:\s*(.*?)(?=\n|\r\n)", text, re.DOTALL)
-                        if not notes_match:
-                            notes_match = re.search(r"Notes\s*:\s*(.*?)(?=\n|\r\n)", text, re.DOTALL)
+                        # 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:", "")
 
-                        reason = reason_match.group(1).strip() if reason_match else "Not available"
-                        alternative = alternative_match.group(1).strip() if alternative_match else "Not available"
-                        notes = notes_match.group(1).strip() if notes_match else "Not available"
-                        
-                        response = f"Reason: {reason}\nAlternative: {alternative}\nNotes: {notes}"
+                        response = f"Reason: {reason.strip()}\nAlternative: {alternative.strip()}\nNotes: {notes.strip()}"
                         print(f"Generated response: {response}")  # Debugging line
                         return response
 
         return "No relevant information found."
+
+    def extract_between_keywords(self, text, start_keyword, end_keyword):
+        start_idx = text.find(start_keyword)
+        if start_idx == -1:
+            return "Not available"
+        
+        start_idx += len(start_keyword)
+        if end_keyword:
+            end_idx = text.find(end_keyword, start_idx)
+            if end_idx == -1:
+                end_idx = len(text)
+        else:
+            end_idx = len(text)
+
+        return text[start_idx:end_idx].strip()
+