Spaces:
Running
Running
| import requests | |
| import json | |
| class VectaraQuery(): | |
| def __init__(self, api_key: str, corpus_keys: list[str], prompt_name: str = None): | |
| self.corpus_keys = corpus_keys | |
| self.api_key = api_key | |
| self.prompt_name = prompt_name if prompt_name else "vectara-summary-ext-24-05-sml" | |
| self.conv_id = None | |
| def get_body(self, query_str: str, stream: False): | |
| corpora_list = [{ | |
| 'corpus_key': corpus_key, 'lexical_interpolation': 0.005 | |
| } for corpus_key in self.corpus_keys | |
| ] | |
| return { | |
| 'query': query_str, | |
| 'search': | |
| { | |
| 'corpora': corpora_list, | |
| 'offset': 0, | |
| 'limit': 50, | |
| 'context_configuration': | |
| { | |
| 'sentences_before': 2, | |
| 'sentences_after': 2, | |
| 'start_tag': "%START_SNIPPET%", | |
| 'end_tag': "%END_SNIPPET%", | |
| }, | |
| 'reranker': | |
| { | |
| 'type': 'customer_reranker', | |
| 'reranker_id': 'rnk_272725719' | |
| }, | |
| }, | |
| 'generation': | |
| { | |
| 'prompt_name': self.prompt_name, | |
| 'max_used_search_results': 10, | |
| 'response_language': 'eng', | |
| 'citations': | |
| { | |
| 'style': 'none' | |
| } | |
| }, | |
| 'chat': | |
| { | |
| 'store': True | |
| }, | |
| 'stream_response': stream | |
| } | |
| def get_headers(self): | |
| return { | |
| "Content-Type": "application/json", | |
| "Accept": "application/json", | |
| "x-api-key": self.api_key, | |
| "grpc-timeout": "60S" | |
| } | |
| def get_stream_headers(self): | |
| return { | |
| "Content-Type": "application/json", | |
| "Accept": "text/event-stream", | |
| "x-api-key": self.api_key, | |
| "grpc-timeout": "60S" | |
| } | |
| def submit_query(self, query_str: str): | |
| if self.conv_id: | |
| endpoint = f"https://api.vectara.io/v2/chats/{self.conv_id}/turns" | |
| else: | |
| endpoint = "https://api.vectara.io/v2/chats" | |
| body = self.get_body(query_str, stream=False) | |
| 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}") | |
| if response.status_code == 429: | |
| return "Sorry, Vectara chat turns exceeds plan limit." | |
| return "Sorry, something went wrong in my brain. Please try again later." | |
| res = response.json() | |
| if self.conv_id is None: | |
| self.conv_id = res['chat_id'] | |
| summary = res['answer'] | |
| return summary | |
| def submit_query_streaming(self, query_str: str): | |
| if self.conv_id: | |
| endpoint = f"https://api.vectara.io/v2/chats/{self.conv_id}/turns" | |
| else: | |
| endpoint = "https://api.vectara.io/v2/chats" | |
| body = self.get_body(query_str, stream=True) | |
| response = requests.post(endpoint, data=json.dumps(body), verify=True, headers=self.get_stream_headers(), stream=True) | |
| if response.status_code != 200: | |
| print(f"Query failed with code {response.status_code}, reason {response.reason}, text {response.text}") | |
| if response.status_code == 429: | |
| return "Sorry, Vectara chat turns exceeds plan limit." | |
| return "Sorry, something went wrong in my brain. Please try again later." | |
| chunks = [] | |
| for line in response.iter_lines(): | |
| line = line.decode('utf-8') | |
| if line: # filter out keep-alive new lines | |
| key, value = line.split(':', 1) | |
| if key == 'data': | |
| line = json.loads(value) | |
| if line['type'] == 'generation_chunk': | |
| chunk = line['generation_chunk'] | |
| chunks.append(chunk) | |
| yield chunk | |
| return ''.join(chunks) |