Update query.py

query.py

@@ -111,7 +111,7 @@ class VectaraQuery():
 
         endpoint = f"https://api.vectara.io/v1/query"
         body = self.get_body(query_str)
-        
+
         response = requests.post(endpoint, data=json.dumps(body), verify=True, headers=self.get_headers())    
         if response.status_code != 200:
             print(f"Query failed with code {response.status_code}, reason {response.reason}, text {response.text}")
@@ -141,7 +141,7 @@ class VectaraQuery():
 
         endpoint = f"https://api.vectara.io/v1/stream-query"
         body = self.get_body(query_str)
-        
+
         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}")
@@ -154,14 +154,14 @@ class VectaraQuery():
             if line:  # filter out keep-alive new lines
                 data = json.loads(line.decode('utf-8'))
                 res = data['result']
-
-                # capture responses and docs if we get that first
-                response_set = res['responseSet']
-                if response_set is not None:
-                    # do we have chat conv_id to update?
-                    summary = response_set.get('summary', [])
-                    if len(summary)>0:
-                        chat = summary[0].get('chat', None)
+                
+                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}")
@@ -172,13 +172,8 @@ class VectaraQuery():
                         conv_id = chat.get('conversationId', None) if chat else None
                         if conv_id:
                             self.conv_id = conv_id
-
-                else:
-                    # grab next chunk and yield it as output
-                    summary = res.get('summary', None)
-                    if summary is None:
-                        continue
-                    chunk = data['result']['summary']['text']
+                        
+                    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)
@@ -197,4 +192,4 @@ class VectaraQuery():
             chunks.append(accumulated_text)
             yield accumulated_text        
         
-        return ''.join(chunks)
+        return ''.join(chunks)