Update query.py

query.py

@@ -23,7 +23,11 @@ class VectaraQuery():
                 Format the response in the following way:\n\
                 Reason: <reason why the name cannot be used>\n\
                 Alternative: <alternative name>\n\
-                Notes: <additional notes>"
+                Notes: <additional notes>\n\n\
+                Example:\n\
+                Reason: The name 'Vodka Sunrise' cannot be used because it is trademarked.\n\
+                Alternative: Use 'Morning Delight' instead.\n\
+                Notes: Ensure the drink contains vodka to match the alternative name."
             }},
             {{
                 "role": "user",
@@ -44,7 +48,19 @@ class VectaraQuery():
                         'sentences_after': 2,
                         'start_tag': "%START_SNIPPET%",
                         'end_tag': "%END_SNIPPET%",
-                    }
+                    },
+                    'summary': [
+                        {
+                            'responseLang': 'eng',
+                            'maxSummarizedResults': 1,
+                            'summarizerPromptName': self.prompt_name,
+                            'promptText': prompt,
+                            'chat': {
+                                'store': True,
+                                'conversationId': self.conv_id
+                            },
+                        }
+                    ]
                 } 
             ]
         }
@@ -90,30 +106,60 @@ class VectaraQuery():
                             continue
                         text = result['text']
                         print(f"Processing text: {text}")  # Debugging line
+                        accumulated_text += text
 
-                        # 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}\nAlternative: {alternative}\nNotes: {notes}"
-                        print(f"Generated response: {response}")  # Debugging line
-                        return response
+        if accumulated_text:
+            return self.summarize_text(accumulated_text)
 
         return "No relevant information found."
 
-    def extract_text(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()
+    def summarize_text(self, text):
+        endpoint = f"https://api.vectara.io/v1/stream-summary"
+        prompt = f'''
+        [
+            {{
+                "role": "system",
+                "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>\n\n\
+                Example:\n\
+                Reason: The name 'Vodka Sunrise' cannot be used because it is trademarked.\n\
+                Alternative: Use 'Morning Delight' instead.\n\
+                Notes: Ensure the drink contains vodka to match the alternative name."
+            }},
+            {{
+                "role": "user",
+                "content": "{text}"
+            }}
+        ]
+        '''
+        body = {
+            'text': text,
+            'summary': {
+                'responseLang': 'eng',
+                'maxSummarizedResults': 1,
+                'summarizerPromptName': self.prompt_name,
+                'promptText': prompt,
+                'chat': {
+                    'store': True,
+                    'conversationId': self.conv_id
+                },
+            }
+        }
+        response = requests.post(endpoint, data=json.dumps(body), verify=True, headers=self.get_headers(), stream=True)
+        if response.status_code != 200:
+            print(f"Summary query failed with code {response.status_code}, reason {response.reason}, text {response.text}")
+            return "Sorry, something went wrong. Please try again later."
+
+        for line in response.iter_lines():
+            if line:  # filter out keep-alive new lines
+                data = json.loads(line.decode('utf-8'))
+                print(f"Received summary data chunk: {json.dumps(data, indent=2)}")  # Debugging line
+
+                if 'summary' in data:
+                    return data['summary']['text']
+
+        return "No relevant information found."