Create es_gpt.py

es_gpt.py

@@ -0,0 +1,75 @@
+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):
+        # Generate summaries for each search result
+        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