aboutsummaryrefslogtreecommitdiff
path: root/gn3/llms/process.py
diff options
context:
space:
mode:
Diffstat (limited to 'gn3/llms/process.py')
-rw-r--r--gn3/llms/process.py31
1 files changed, 29 insertions, 2 deletions
diff --git a/gn3/llms/process.py b/gn3/llms/process.py
index abd307e..549c7e6 100644
--- a/gn3/llms/process.py
+++ b/gn3/llms/process.py
@@ -1,6 +1,7 @@
"""this module contains code for processing response from fahamu client.py"""
+import os
import string
import json
@@ -10,6 +11,7 @@ import requests
from gn3.llms.client import GeneNetworkQAClient
from gn3.llms.response import DocIDs
+from gn3.settings import TMPDIR
BASE_URL = 'https://genenetwork.fahamuai.com/api/tasks'
@@ -67,13 +69,36 @@ def rate_document(task_id, doc_id, rating, auth_token):
raise RuntimeError(f"An error occurred: {str(error)}") from error
+def load_file(filename):
+ """function to open and load json file"""
+ file_path = os.path.join(TMPDIR, filename)
+ if not os.path.isfile(file_path):
+ raise FileNotFoundError(f"{filename} was not found or is a directory")
+ with open(file_path, "rb") as file_handler:
+ return json.load(file_handler)
+
+
+def fetch_pubmed(references, file_name):
+ """method to fetch and populate references with pubmed"""
+
+ try:
+ pubmed = load_file(file_name)
+ for reference in references:
+ if pubmed.get(reference["doc_id"]):
+ reference["pubmed"] = pubmed.get(reference["doc_id"])
+ return references
+
+ except FileNotFoundError:
+ return references
+
+
def get_gnqa(query, auth_token):
"""entry function for the gn3 api endpoint()"""
api_client = GeneNetworkQAClient(requests.Session(), api_key=auth_token)
res, task_id = api_client.ask('?ask=' + quote(query), auth_token)
- if task_id == 0 :
- raise RuntimeError(f"Error connecting to Fahamu Api: {str(res)}")
+ if task_id == 0:
+ raise RuntimeError(f"Error connecting to Fahamu Api: {str(res)}")
res, success = api_client.get_answer(task_id)
if success == 1:
resp_text = filter_response_text(res.text)
@@ -83,6 +108,8 @@ def get_gnqa(query, auth_token):
context = resp_text['data']['context']
references = parse_context(
context, DocIDs().getInfo, format_bibliography_info)
+ references = fetch_pubmed(references, "pubmed.json")
+
return task_id, answer, references
else:
return task_id, "Unfortunately, I have nothing on the query", []