Update query.py

query.py

@@ -8,6 +8,47 @@ def extract_between_tags(text, start_tag, end_tag):
     end_index = text.find(end_tag, start_index)
     return text[start_index+len(start_tag):end_index-len(end_tag)]
 
+class CitationNormalizer():
+
+    def __init__(self, responses, docs):
+        self.docs = docs
+        self.responses = responses
+        self.refs = []
+
+    def normalize_citations(self, summary):
+        start_tag = "%START_SNIPPET%"
+        end_tag = "%END_SNIPPET%"
+
+        # find all references in the summary
+        pattern = r'\[\d{1,2}\]'
+        matches = [match.span() for match in re.finditer(pattern, summary)]
+
+        # figure out unique list of references
+        for match in matches:
+            start, end = match
+            response_num = int(summary[start+1:end-1])
+            doc_num = self.responses[response_num-1]['documentIndex']
+            metadata = {item['name']: item['value'] for item in self.docs[doc_num]['metadata']}
+            text = extract_between_tags(self.responses[response_num-1]['text'], start_tag, end_tag)
+            if 'url' in metadata.keys():
+                url = f"{metadata['url']}#:~:text={quote(text)}"
+                if url not in self.refs:
+                    self.refs.append(url)
+
+        # replace references with markdown links
+        refs_dict = {url:(inx+1) for inx,url in enumerate(self.refs)}
+        for match in reversed(matches):
+            start, end = match
+            response_num = int(summary[start+1:end-1])
+            doc_num = self.responses[response_num-1]['documentIndex']
+            metadata = {item['name']: item['value'] for item in self.docs[doc_num]['metadata']}
+            text = extract_between_tags(self.responses[response_num-1]['text'], start_tag, end_tag)
+            url = f"{metadata['url']}#:~:text={quote(text)}"
+            citation_inx = refs_dict[url]
+            summary = summary[:start] + f'[\[{citation_inx}\]]({url})' + summary[end:]
+
+        return summary
+
 class VectaraQuery():
     def __init__(self, api_key: str, customer_id: str, corpus_ids: list[str], prompt_name: str = None):
         self.customer_id = customer_id
@@ -16,23 +57,13 @@ class VectaraQuery():
         self.prompt_name = prompt_name if prompt_name else "vectara-summary-ext-v1.2.0"
         self.conv_id = None
 
-    def submit_query(self, query_str: str):
+    def get_body(self, query_str: str):
         corpora_key_list = [{
                 'customer_id': self.customer_id, 'corpus_id': corpus_id, 'lexical_interpolation_config': {'lambda': 0.025}
             } for corpus_id in self.corpus_ids
         ]
 
-        endpoint = f"https://api.vectara.io/v1/query"
-        start_tag = "%START_SNIPPET%"
-        end_tag = "%END_SNIPPET%"
-        headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json",
-            "customer-id": self.customer_id,
-            "x-api-key": self.api_key,
-            "grpc-timeout": "60S"
-        }
-        body = {
+        return {
             'query': [
                 { 
                     'query': query_str,
@@ -42,8 +73,8 @@ class VectaraQuery():
                     'context_config': {
                         'sentences_before': 2,
                         'sentences_after': 2,
-                        'start_tag': start_tag,
-                        'end_tag': end_tag,
+                        'start_tag': "%START_SNIPPET%",
+                        'end_tag': "%END_SNIPPET%",
                     },
                     'rerankingConfig':
                     {
@@ -61,14 +92,27 @@ class VectaraQuery():
                                 'store': True,
                                 'conversationId': self.conv_id
                             },
-#                            'debug': True,
                         }
                     ]
                 } 
             ]
         }
+
+    def get_headers(self):
+        return {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+            "customer-id": self.customer_id,
+            "x-api-key": self.api_key,
+            "grpc-timeout": "60S"
+        }
+
+    def submit_query(self, query_str: str):
+
+        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=headers)    
+        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}")
             return "Sorry, something went wrong in my brain. Please try again later."
@@ -90,33 +134,68 @@ class VectaraQuery():
             return 'Sorry, something went wrong in my brain. Please try again later.'
         
         self.conv_id = res['responseSet'][0]['summary'][0]['chat']['conversationId']
+
+        summary = CitationNormalizer().normalize_citations(summary, responses, docs)        
+        return summary
+
+    def submit_query_streaming(self, query_str: str):
+
+        endpoint = f"https://api.vectara.io/v1/stream-query"
+        body = self.get_body(query_str)
         
-        pattern = r'\[\d{1,2}\]'
-        matches = [match.span() for match in re.finditer(pattern, summary)]
+        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."
 
-        # figure out unique list of references
-        refs = []
-        for match in matches:
-            start, end = match
-            response_num = int(summary[start+1:end-1])
-            doc_num = responses[response_num-1]['documentIndex']
-            metadata = {item['name']: item['value'] for item in docs[doc_num]['metadata']}
-            text = extract_between_tags(responses[response_num-1]['text'], start_tag, end_tag)
-            if 'url' in metadata.keys():
-                url = f"{metadata['url']}#:~:text={quote(text)}"
-                if url not in refs:
-                    refs.append(url)
+        chunks = []
+        accumulated_text = ""  # Initialize text accumulation
+        pattern_max_length = 50  # Example heuristic
+        for line in response.iter_lines():
+            if line:  # filter out keep-alive new lines
+                data = json.loads(line.decode('utf-8'))
+                res = data['result']
 
-        # replace references with markdown links
-        refs_dict = {url:(inx+1) for inx,url in enumerate(refs)}
-        for match in reversed(matches):
-            start, end = match
-            response_num = int(summary[start+1:end-1])
-            doc_num = responses[response_num-1]['documentIndex']
-            metadata = {item['name']: item['value'] for item in docs[doc_num]['metadata']}
-            text = extract_between_tags(responses[response_num-1]['text'], start_tag, end_tag)
-            url = f"{metadata['url']}#:~:text={quote(text)}"
-            citation_inx = refs_dict[url]
-            summary = summary[:start] + f'[\[{citation_inx}\]]({url})' + summary[end:]
+                # 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 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
 
-        return summary
+                else:
+                    # grab next chunk and yield it as output
+                    summary = res.get('summary', None)
+                    if summary is None:
+                        continue
+                    chunk = data['result']['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
+
+        # 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)