query.py

import requests
import json
import re

class VectaraQuery():
    def __init__(self, api_key: str, customer_id: str, corpus_id: str, prompt_name: str = None):
        self.customer_id = customer_id
        self.corpus_id = corpus_id
        self.api_key = api_key
        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):
        corpora_key_list = [{
            'customer_id': self.customer_id, 'corpus_id': self.corpus_id, 'lexical_interpolation_config': {'lambda': 0.025}
        }]

        user_response = user_response.replace('"', '\\"')  # Escape double quotes
        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>"
            }},
            {{
                "role": "user",
                "content": "{user_response}"
            }}
        ]
        '''

        return {
            'query': [
                { 
                    'query': user_response,
                    'start': 0,
                    'numResults': 10,
                    'corpusKey': corpora_key_list,
                    'context_config': {
                        'sentences_before': 2,
                        'sentences_after': 2,
                        'start_tag': "%START_SNIPPET%",
                        'end_tag': "%END_SNIPPET%",
                    }
                } 
            ]
        }

    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/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}")
            return "Sorry, something went wrong. Please try again later."

        accumulated_text = ""
        for line in response.iter_lines():
            if line:  # filter out keep-alive new lines
                data = json.loads(line.decode('utf-8'))
                print(f"Received data chunk: {json.dumps(data, indent=2)}")  # Debugging line

                if 'result' not in data:
                    print("No 'result' in data")
                    continue

                res = data['result']
                if 'responseSet' not in res:
                    print("No 'responseSet' in result")
                    continue
                
                response_set = res['responseSet']
                
                if response_set:
                    for result in response_set['response']:
                        if 'text' not in result:
                            print("No 'text' in result")
                            continue
                        text = result['text']
                        print(f"Processing text: {text}")  # Debugging line

                        # Instead of using regex, split the text by specific keywords
                        reason = self.extract_between_keywords(text, "Reason:", "Alternative:")
                        alternative = self.extract_between_keywords(text, "Alternative:", "Notes:")
                        notes = self.extract_between_keywords(text, "Notes:", "")

                        response = f"Reason: {reason.strip()}\nAlternative: {alternative.strip()}\nNotes: {notes.strip()}"
                        print(f"Generated response: {response}")  # Debugging line
                        return response

        return "No relevant information found."

    def extract_between_keywords(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()