Update query.py

query.py

@@ -6,7 +6,7 @@ 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-len(end_tag)]
+    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):
@@ -16,7 +16,7 @@ class VectaraQuery():
         self.prompt_name = prompt_name if prompt_name else "vectara-experimental-summary-ext-2023-12-11-large"
         self.conv_id = None
 
-    def get_body(self, user_response: str, role: str, topic: str, style: str):
+    def get_body(self, user_response: str):
         corpora_key_list = [{
             'customer_id': self.customer_id, 'corpus_id': self.corpus_id, 'lexical_interpolation_config': {'lambda': 0.025}
         }]
@@ -26,29 +26,11 @@ class VectaraQuery():
         [
             {{
                 "role": "system",
-                "content": "You are a professional debate bot. 
-                You specialize in the {style} debate style.
-                You are provided with search results related to {topic}. 
-                Follow these INSTRUCTIONS carefully:
-                1. Provide a thoughtful and convincing reply.
-                2. Do not base your response on information or knowledge that is not in the search results.
-                3. Respond with respect to your opponent. 
-                4. Limit your responses to not more than 2 paragraphs."
-            }},
-            {{
-                "role": "assistant",
-                "content": "
-                    #foreach ($qResult in $vectaraQueryResults)
-                        Search result $esc.java(${{foreach.index}}+1): $esc.java(${{qResult.getText()}})
-                    #end
-                "
+                "content": "You are an assistant that provides information about drink names based on a given corpus."
             }},
             {{
                 "role": "user",
-                "content": "Provide a convincing response {role} {topic}, to the question '$esc.java(${{vectaraQuery}})'. 
-                Consider the search results as relevant information with which to form your response, but do not mention the results in your response.
-                Consider the last argument from your opponent: '{user_response}'.
-                Use the {style} debate style to make your argument."
+                "content": "{user_response}"
             }}
         ]
         '''
@@ -56,35 +38,16 @@ class VectaraQuery():
         return {
             'query': [
                 { 
-                    'query': f"What is a good argument {role} {topic}",
+                    'query': user_response,
                     'start': 0,
-                    'numResults': 50,
+                    'numResults': 10,
                     'corpusKey': corpora_key_list,
                     'context_config': {
                         'sentences_before': 2,
                         'sentences_after': 2,
                         'start_tag': "%START_SNIPPET%",
                         'end_tag': "%END_SNIPPET%",
-                    },
-                    'rerankingConfig':
-                    {
-                        'rerankerId': 272725718,
-                        'mmrConfig': {
-                            'diversityBias': 0.3
-                        }
-                    },
-                    'summary': [
-                        {
-                            'responseLang': 'eng',
-                            'maxSummarizedResults': 7,
-                            'summarizerPromptName': self.prompt_name,
-                            'promptText': prompt,
-                            'chat': {
-                                'store': True,
-                                'conversationId': self.conv_id
-                            },
-                        }
-                    ]
+                    }
                 } 
             ]
         }
@@ -98,11 +61,10 @@ class VectaraQuery():
             "grpc-timeout": "60S"
         }
 
-    def submit_query(self, query_str: str, bot_role: str, topic: str, style: str):
+    def submit_query(self, query_str: str):
 
         endpoint = f"https://api.vectara.io/v1/stream-query"
-        body = self.get_body(query_str, bot_role, topic, style)
-
+        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}")
@@ -153,4 +115,4 @@ class VectaraQuery():
             chunks.append(accumulated_text)
             yield accumulated_text        
         
-        return ''.join(chunks)
+        return ''.join(chunks)