Spaces:
Running
Running
File size: 2,660 Bytes
893c5d9 c150a4d 893c5d9 c150a4d 893c5d9 c150a4d 893c5d9 c150a4d 893c5d9 |
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 |
import requests
import json
import re
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"
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': 10,
'corpusKey': corpora_key_list,
'context_config': {
'sentences_before': 2,
'sentences_after': 2,
'start_tag': "%START_SNIPPET%",
'end_tag': "%END_SNIPPET%",
},
'rerankingConfig': { 'rerankerId': 272725719 }
}
]
}
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 = "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()
top_k = 3
responses = res['responseSet'][0]['response'][:top_k]
documents = res['responseSet'][0]['document']
metadatas = []
for x in responses:
md = {m["name"]: m["value"] for m in x["metadata"]}
doc_num = x["documentIndex"]
doc_md = {f'doc_{m["name"]}': m["value"] for m in documents[doc_num]["metadata"]}
md.update(doc_md)
metadatas.append(md)
movie_title = metadatas[0].get("doc_title", None)
snippet_url = metadatas[0].get("url", None)
score = responses[0]["score"]
return movie_title, snippet_url, score
|