| from elasticsearch import Elasticsearch |
| import os |
| import json |
| import requests |
|
|
| ES_URL = os.environ["ES_URL"] |
| ES_USER = os.environ["ES_USER"] |
| ES_PASS = os.environ["ES_PASS"] |
| ES_CA_CERT = os.environ["ES_CA_CERT"] |
|
|
|
|
| class ESGPT: |
| def __init__(self, index_name): |
| self.es = Elasticsearch(ES_URL, http_auth=(ES_USER, ES_PASS), |
| ca_certs=ES_CA_CERT, verify_certs=True) |
| self.index_name = index_name |
| self.model_engine = os.environ["OPENAI_GPT_ENGINE"] |
| self.api_key = os.environ["OPENAI_API_KEY"] |
|
|
| def index(self, doc_id, doc): |
| self.es.index(index=self.index_name, |
| id=doc_id, |
| document=doc) |
|
|
| def search(self, query): |
| body = { |
| "query": { |
| "query_string": {"query": query} |
| } |
| } |
|
|
| results = self.es.search(index=self.index_name, body=body) |
| return results['hits']['hits'] |
|
|
| def _paper_results_to_text(self, results): |
| text_result = "" |
| for paper in results: |
| title = "" |
| if "title" in paper["_source"]: |
| title = paper["_source"]["title"] |
| |
| abstract = "" |
| if "abctract" in paper["_source"]: |
| abstract = paper["_source"]["abstract"] |
| |
| paper_str = f"{title}:\n{abstract[:100]}\n\n" |
| text_result += paper_str |
| return text_result |
|
|
| def summarize(self, query, results): |
| |
| result_json_str = self._paper_results_to_text(results) |
| if result_json_str == "": |
| result_json_str = "No results found" |
|
|
| print(result_json_str[:500]) |
|
|
| body = { |
| "model": self.model_engine, |
| "prompt": f"Please summarize the following search results for query: {query}:\n{result_json_str[:1000]}", |
| "max_tokens": 1000, |
| "n": 1, |
| "stop": None, |
| "temperature": 0.5, |
| "stream": True, |
| } |
|
|
| headers = {"Content-Type": "application/json", |
| "Authorization": f"Bearer {self.api_key}"} |
|
|
| resp = requests.post("https://api.openai.com/v1/completions", |
| headers=headers, |
| data=json.dumps(body), |
| stream=True) |
| return resp |
|
|
