Spaces:
Running
Running
File size: 5,362 Bytes
7f46a81 d5124c2 7f46a81 2a89f81 7f46a81 d5124c2 f26592e 7f46a81 d5124c2 7f46a81 9b851fc 7f46a81 d5124c2 7f46a81 6541511 7f46a81 f26592e d5124c2 7f46a81 6541511 9b851fc 5fc81fd 7f46a81 9b851fc 2a89f81 f26592e 5cebf82 55a083a 7f46a81 b4ea488 7f46a81 55a083a d5124c2 7f46a81 7ff5239 7f46a81 d5124c2 39e2176 d5124c2 39e2176 d5124c2 7f46a81 d5124c2 7f46a81 55a083a d5124c2 f26592e d5124c2 57fd6e4 d5124c2 55a083a d5124c2 a452991 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
import requests
import json
class VectaraQuery():
def __init__(self, api_key: str, customer_id: str, corpus_ids: list[str], prompt_name: str = None):
self.customer_id = customer_id
self.corpus_ids = corpus_ids
self.api_key = api_key
self.prompt_name = prompt_name if prompt_name else "vectara-experimental-summary-ext-2023-12-11-sml"
self.conv_id = None
def get_body(self, query_str: str):
corpora_key_list = [{
'customer_id': self.customer_id, 'corpus_id': corpus_id, 'lexical_interpolation_config': {'lambda': 0.005}
} for corpus_id in self.corpus_ids
]
return {
'query': [
{
'query': query_str,
'start': 0,
'numResults': 50,
'corpusKey': corpora_key_list,
'context_config': {
'sentences_before': 2,
'sentences_after': 2,
'start_tag': "%START_SNIPPET%",
'end_tag': "%END_SNIPPET%",
},
'rerankingConfig':
{
'rerankerId': 272725719,
},
'summary': [
{
'responseLang': 'eng',
'maxSummarizedResults': 10,
'summarizerPromptName': self.prompt_name,
'chat': {
'store': True,
'conversationId': self.conv_id
},
'citationParams': {
"style": "NONE",
}
}
]
}
]
}
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=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."
res = response.json()
summary = res['responseSet'][0]['summary'][0]['text']
chat = res['responseSet'][0]['summary'][0].get('chat', None)
if chat and chat['status'] is not 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.'
self.conv_id = chat['conversationId'] if chat else None
return summary
def submit_query_streaming(self, query_str: str):
endpoint = "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 in my brain. Please try again later."
chunks = []
for line in response.iter_lines():
if line: # filter out keep-alive new lines
data = json.loads(line.decode('utf-8'))
res = data['result']
response_set = res['responseSet']
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}")
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
chunk = summary['text']
chunks.append(chunk)
yield chunk
if summary['done']:
break
return ''.join(chunks)
|