fix: handle rate limit errors on backend and frontend (#22)

* chore: check for rate limits

* fix: only show one error message at a time

* fix: show user error from server

src/backend/app.py

@@ -3,6 +3,7 @@
 import openai
 from flask import Flask, abort, request
 
+from backend.error_messages import ErrorMessages
 from backend.http_status_codes import StatusCodes
 
 app = Flask(__name__)
@@ -23,7 +24,7 @@ def get_models() -> list[str]:
 def get_gpt_3_5_response(dataset: str) -> str:
     openai.api_key = request.headers.get("Authorization")
     if not openai.api_key:
-        abort(StatusCodes.UNAUTHORISED.value, "API key not given")
+        abort(StatusCodes.UNAUTHORISED.value, ErrorMessages.NO_API_KEY.value)
 
     try:
         response = openai.ChatCompletion.create(
@@ -36,17 +37,24 @@ def get_gpt_3_5_response(dataset: str) -> str:
                 }
             ]
         )
-        # TODO: handle rate limit error
         return response.choices[0].message.content
+    except openai.error.RateLimitError:
+        abort(
+            StatusCodes.BAD_REQUEST.value,
+            ErrorMessages.RATE_LIMIT.value
+        )
     except openai.error.AuthenticationError:
-        abort(StatusCodes.UNAUTHORISED.value, "API key given is not valid")
+        abort(
+            StatusCodes.UNAUTHORISED.value,
+            ErrorMessages.INVALID_API_KEY.value
+        )
 
 
 @app.route("/davinci/<dataset>")
 def get_davinci_response(dataset: str) -> str:
     openai.api_key = request.headers.get("Authorization")
     if not openai.api_key:
-        abort(StatusCodes.UNAUTHORISED.value, "API key not given")
+        abort(StatusCodes.UNAUTHORISED.value, ErrorMessages.NO_API_KEY.value)
 
     try:
         response = openai.Completion.create(
@@ -55,10 +63,15 @@ def get_davinci_response(dataset: str) -> str:
             max_tokens=4000,
             temperature=0.2,
         )
-        # TODO: handle rate limit error
         return response.choices[0].text
+    except openai.error.RateLimitError:
+        abort(
+            StatusCodes.BAD_REQUEST.value,
+            ErrorMessages.RATE_LIMIT.value
+        )
     except openai.error.AuthenticationError:
-        abort(StatusCodes.UNAUTHORISED.value, "API key given is not valid")
+        abort(StatusCodes.UNAUTHORISED.value,
+              ErrorMessages.INVALID_API_KEY.value)
 
 
 if __name__ == '__main__':

src/backend/error_messages.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class ErrorMessages(Enum):
+    NO_API_KEY = "API key not given"
+    INVALID_API_KEY = "API key given is not valid"
+    RATE_LIMIT = "You have reached your rate limit. See here for more " \
+                 "information: " \
+                 "https://platform.openai.com/docs/guides/rate-limits/overview"

src/backend/http_status_codes.py

@@ -2,4 +2,5 @@
 
 
 class StatusCodes(Enum):
+    BAD_REQUEST = 400
     UNAUTHORISED = 403

src/frontend/api.py

@@ -1,13 +1,16 @@
 import ast
+import bs4
 import logging
 from enum import Enum
 from io import StringIO
+from requests.exceptions import HTTPError
 from typing import List
 
-import openai
 import pandas as pd
 import requests
 
+from frontend.server_exception import ServerException
+
 BASE_URL = "http://127.0.0.1:8000"
 
 logger = logging.getLogger(__name__)
@@ -41,8 +44,18 @@ def get_response(dataset_of: str, gpt_choice, api_key: str) -> pd.DataFrame:
         url = f"{BASE_URL}/davinci/{dataset_of}"
 
     headers = get_auth_header(api_key)
-    response = requests.get(url=url, headers=headers)
-    return str_to_df(response.text)
+
+    try:
+        response = requests.get(url=url, headers=headers)
+        response.raise_for_status()
+        status_code = response.status_code
+        if status_code == 200:
+            return str_to_df(response.text)
+    except HTTPError as e:
+        html = bs4.BeautifulSoup(e.response.text, features="html.parser")
+        server_message = html.p.text
+        logger.info(server_message)
+        raise ServerException(server_message)
 
 
 # TODO: add a function that checks if the input is valid markdown

src/frontend/server_exception.py

@@ -0,0 +1,3 @@
+class ServerException(Exception):
+    def __init__(self, message):
+        self.message = message

src/frontend/st_app.py

@@ -1,16 +1,18 @@
 import time
 
+import pandas as pd
 import streamlit as st
 
 from frontend.api import get_models, get_response
+from frontend.server_exception import ServerException
 from frontend.util import df_to_csv, df_to_json
 
 app_title = "Dataset Generator"
 st.set_page_config(page_title=app_title)
 st.title(app_title)
 
 
-def show_export_buttons():
+def show_export_buttons(df: pd.DataFrame):
     col1, col2 = st.columns(2, gap="small")
     with col1:
         st.download_button(
@@ -28,9 +30,9 @@ def show_export_buttons():
         )
 
 
-def show_result():
+def show_result(df: pd.DataFrame):
     st.dataframe(df, use_container_width=True)
-    show_export_buttons()
+    show_export_buttons(df)
 
 
 def get_gpt_radio():
@@ -61,31 +63,42 @@ def get_gpt_radio():
     generate_clicked = st.form_submit_button(label="Generate dataset",
                                              disabled=False)
 
-if generate_clicked:
+
+def generate_dataset():
     if not api_key:
         st.error(
             f"You did not enter in an API key. Please enter your OpenAI API "
             f"key and try again", icon="🥸"
         )
+        return
     if not dataset_entered:
         st.error(
             f"Please enter the dataset you would like us to generate (e.g. "
             f"'Harry Potter quotes') dataset.", icon="🥸"
         )
+        return
 
     # Show spinner until a DataFrame is returned
     with st.spinner(f"Generating a dataset of {dataset_entered}..."):
-        df = get_response(dataset_entered, gpt_choice, api_key)
-        # the truth value of a DataFrame is ambiguous, so cannot use
-        # 'while not df'
-        while df is None:
-            # check every 1 second
-            time.sleep(1)
+        try:
+            df = get_response(dataset_entered, gpt_choice, api_key)
+            # the truth value of a DataFrame is ambiguous, so cannot use
+            # 'while not df'
+            while df is None:
+                # check every 1 second
+                time.sleep(1)
+        except ServerException as e:
+            st.error(e.message)
+            return
 
     if df.empty:
         st.error(
             f"We could not generate a {dataset_entered} dataset. "
             f"Try generating a different dataset.", icon="🤔"
         )
     else:
-        show_result()
+        show_result(df)
+
+
+if generate_clicked:
+    generate_dataset()

src/requirements-dev.txt

@@ -1 +1,2 @@
 pre-commit==3.3.3
+pytest==7.4.0

src/requirements.txt

@@ -1,5 +1,6 @@
-streamlit==1.23.1
+streamlit==1.25.0
 openai==0.27.8
 python-dotenv==1.0.0
 pandas==2.0.2
 Flask==2.3.2
+beautifulsoup4==4.12.2