Update query.py

query.py

@@ -1,12 +1,6 @@
 import requests
 import json
 import re
-from urllib.parse import quote
-
-def extract_between_tags(text, start_tag, end_tag):
-    start_index = text.find(start_tag)
-    end_index = text.find(end_tag, start_index)
-    return text[start_index+len(start_tag):end_index]
 
 class VectaraQuery():
     def __init__(self, api_key: str, customer_id: str, corpus_id: str, prompt_name: str = None):
@@ -68,7 +62,7 @@ class VectaraQuery():
         response = requests.post(endpoint, data=json.dumps(body), verify=True, headers=self.get_headers(), stream=True) 
         if response.status_code != 200:
             print(f"Query failed with code {response.status_code}, reason {response.reason}, text {response.text}")
-            return "Sorry, something went wrong in my brain. Please try again later."
+            return "Sorry, something went wrong. Please try again later."
 
         chunks = []
         accumulated_text = ""  # Initialize text accumulation
@@ -77,42 +71,17 @@ class VectaraQuery():
             if line:  # filter out keep-alive new lines
                 data = json.loads(line.decode('utf-8'))                
                 res = data['result']
-                response_set = res['responseSet']                
-                if response_set is None:
-                    # grab next chunk and yield it as output
-                    summary = res.get('summary', None)
-                    if summary is None or len(summary)==0:
-                        continue
-                    else:
-                        chat = summary.get('chat', None)
-                        if chat and chat.get('status', None):
-                            st_code = chat['status']
-                            print(f"Chat query failed with code {st_code}")
-                            if st_code == 'RESOURCE_EXHAUSTED':
-                                self.conv_id = None
-                                return 'Sorry, Vectara chat turns exceeds plan limit.'
-                            return 'Sorry, something went wrong in my brain. Please try again later.'
-                        conv_id = chat.get('conversationId', None) if chat else None
-                        if conv_id:
-                            self.conv_id = conv_id
+                response_set = res['responseSet']
+                
+                if response_set:
+                    for result in response_set:
+                        text = result['text']
+                        # Extract relevant information from the text
+                        reason = re.search(r"Reason Why it Can't be Used: (.*?)\n", text).group(1)
+                        alternative = re.search(r"Alternative: (.*?)\n", text).group(1)
+                        notes = re.search(r"Notes: (.*?)\n", text).group(1)
                         
-                    chunk = summary['text']
-                    accumulated_text += chunk  # Append current chunk to accumulation
-                    if len(accumulated_text) > pattern_max_length:
-                        accumulated_text = re.sub(r"\[\d+\]", "", accumulated_text)
-                        accumulated_text = re.sub(r"\s+\.", ".", accumulated_text)
-                        out_chunk = accumulated_text[:-pattern_max_length]
-                        chunks.append(out_chunk)
-                        yield out_chunk
-                        accumulated_text = accumulated_text[-pattern_max_length:]
-
-                    if summary['done']:
-                        break
+                        response = f"Reason: {reason}\nAlternative: {alternative}\nNotes: {notes}"
+                        return response
 
-        # yield the last piece
-        if len(accumulated_text) > 0:
-            accumulated_text = re.sub(r" \[\d+\]\.", ".", accumulated_text)
-            chunks.append(accumulated_text)
-            yield accumulated_text        
-        
-        return ''.join(chunks)
+        return "No relevant information found."