|
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('"', '\\"') |
|
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: |
|
data = json.loads(line.decode('utf-8')) |
|
print(f"Received data chunk: {json.dumps(data, indent=2)}") |
|
|
|
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}") |
|
|
|
|
|
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}") |
|
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() |
|
|
|
|