diff options
| -rw-r--r-- | gn3/llms/process.py | 31 |
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", [] |