From 3d3a4400002a12ce27b12e77cf62ac8a25c1091b Mon Sep 17 00:00:00 2001
From: Mohammadreza Pourreza
 <71866535+MohammadrezaPourreza@users.noreply.github.com>
Date: Mon, 6 Nov 2023 14:01:18 -0500
Subject: [PATCH] Dh 4953/fix five hundered error for question endpoint (#238)

* DH-4953/fixing the question 500 error

* DH-4953/raising 400 when API key is invalid
---
 dataherald/api/fastapi.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/dataherald/api/fastapi.py b/dataherald/api/fastapi.py
index f98cda7f..0e6c870f 100644
--- a/dataherald/api/fastapi.py
+++ b/dataherald/api/fastapi.py
@@ -5,6 +5,7 @@
 import time
 from typing import List
 
+import openai
 from bson import json_util
 from bson.objectid import InvalidId, ObjectId
 from fastapi import BackgroundTasks, HTTPException
@@ -146,9 +147,9 @@ def answer_question(
                     "error_message": "Connections doesn't exist",
                 },
             )
-        context = context_store.retrieve_context_for_question(user_question)
-        start_generated_answer = time.time()
         try:
+            context = context_store.retrieve_context_for_question(user_question)
+            start_generated_answer = time.time()
             generated_answer = sql_generation.generate_response(
                 user_question, database_connection, context[0]
             )
@@ -368,6 +369,8 @@ def update_response(self, response_id: str) -> Response:
             response = generates_nl_answer.execute(response)
             response.exec_time = time.time() - start_generated_answer
             response_repository.update(response)
+        except openai.error.AuthenticationError as e:
+            raise HTTPException(status_code=400, detail=str(e)) from e
         except ValueError as e:
             raise HTTPException(status_code=404, detail=str(e)) from e
         except SQLInjectionError as e:
@@ -454,6 +457,8 @@ def create_response(
                 response.confidence_score = confidence_score
             response.exec_time = time.time() - start_generated_answer
             response_repository.update(response)
+        except openai.error.AuthenticationError as e:
+            raise HTTPException(status_code=400, detail=str(e)) from e
         except ValueError as e:
             raise HTTPException(status_code=404, detail=str(e)) from e
         except SQLInjectionError as e: