added MongoDB support

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import model
 from config import app_config
+import mongo_utils as mongo
 
 
 def clear():
@@ -29,19 +30,34 @@ def create_interface():
                 api_name=False,
             )
         with gr.Row():
-            gr.Markdown(
-                """
-                # Storyteller
-                **This app can craft captivating narratives from captivating images, 
-                potentially surpassing even Shakespearean standards. Select an image 
-                that inspires a story, choose a story length (up to 100 words), and 
-                adjust the creativity index to enhance its creative flair.**  
-                <br>
-                ***Please exercise patience, as the models employed are extensive and may
-                require a few seconds to load. If you encounter an unrelated story, 
-                it is likely still loading; wait a moment and try again.***
-                """
-            )
+            with gr.Column(scale=5):
+                gr.Markdown(
+                    """
+                    # Storyteller
+                    **This app can craft captivating narratives from captivating images, 
+                    potentially surpassing even Shakespearean standards. Select an image 
+                    that inspires a story, choose a story length (up to 100 words), and 
+                    adjust the creativity index to enhance its creative flair.**  
+                    <br>
+                    ***Please exercise patience, as the models employed are extensive and may
+                    require a few seconds to load. If you encounter an unrelated story, 
+                    it is likely still loading; wait a moment and try again.***
+                    """
+                )
+            with gr.Column(scale=2):
+                max_count = gr.Textbox(
+                    label="Max allowed OpenAI requests:",
+                    value=app_config.openai_max_access_count,
+                )
+                curr_count = gr.Textbox(
+                    label="Used up OpenAI requests:",
+                    value=app_config.openai_curr_access_count,
+                )
+                available_count = gr.Textbox(
+                    label="Available OpenAI requests:",
+                    value=app_config.openai_max_access_count
+                    - app_config.openai_curr_access_count,
+                )
         with gr.Row():
             with gr.Column():
                 image = gr.Image(
@@ -89,7 +105,7 @@ def create_interface():
         submit_button.click(
             fn=model.generate_story,
             inputs=[image, word_count, creativity],
-            outputs=[story],
+            outputs=[story, max_count, curr_count, available_count],
         )
         clear_button.click(
             fn=clear, inputs=[], outputs=[image, word_count, creativity, story]
@@ -100,5 +116,6 @@ def create_interface():
 
 
 if __name__ == "__main__":
+    mongo.fetch_curr_access_count()
     app = create_interface()
     app.launch()

config.py

@@ -1,6 +1,5 @@
 import os
 from dataclasses import dataclass
-import gradio as gr
 
 
 @dataclass
@@ -8,9 +7,16 @@ class AppConfig:
     title = "Picture to Story Generator"
     theme = "freddyaboulton/dracula_revamped"
     css = "style.css"
-    HF_TOKEN = os.getenv("HF_TOKEN")
-    OPENAI_KEY = os.getenv("OPENAI_KEY")
-    I2T_API_URL = os.getenv("I2T_API_URL")
+    openai_max_access_count = 200
+    openai_curr_access_count = None
+    mongo_client = None
+    db = "mydb"
+    collection = "pic2story-openai-access-counter"
+    key = "current_count"
+    # HF_TOKEN = os.getenv("HF_TOKEN")
+    # OPENAI_KEY = os.getenv("OPENAI_KEY")
+    # I2T_API_URL = os.getenv("I2T_API_URL")
+    # MONGO_CONN_STR = os.getenv("MONGO_CONN_STR")
 
 
 app_config = AppConfig()

model.py

@@ -8,6 +8,7 @@ from langchain.prompts import (
     ChatPromptTemplate,
 )
 from config import app_config
+import mongo_utils as mongo
 
 
 def __image2text(image):
@@ -51,11 +52,18 @@ def __text2story(image_desc, word_count, creativity):
 
 def generate_story(image_file, word_count, creativity):
     """Generates a story given an image"""
-    # read image as bytes array
+    # read image as bytes arrayS
     with open(image_file, "rb") as f:
         input_image = f.read()
+    # generate caption for image
     image_desc = __image2text(image=input_image)
+    # generate story from caption
     story = __text2story(
         image_desc=image_desc, word_count=word_count, creativity=creativity
     )
-    return story
+    # increment the openai access counter and compute count stats
+    mongo.increment_curr_access_count()
+    max_count = app_config.openai_max_access_count
+    curr_count = app_config.openai_curr_access_count
+    available_count = max_count - curr_count
+    return story, max_count, curr_count, available_count

mongo_utils.py

@@ -0,0 +1,57 @@
+import pymongo
+from config import app_config
+
+
+# Generic functions
+def get_db_client():
+    """Returns MongoDB client object, connect to MongoDB Atlas instances if required"""
+    try:
+        if app_config.mongo_client == None:
+            client = pymongo.MongoClient(app_config.MONGO_CONN_STR)
+            app_config.mongo_client = client
+    except Exception as e:
+        print(e)
+    return app_config.mongo_client
+
+
+def fetch_document(client, db, collection):
+    """Get a single document from the provided db and collection"""
+    try:
+        document = client[db][collection].find_one()
+    except Exception as e:
+        print(e)
+    return document
+
+
+def update_document(client, db, collection, key, value):
+    """Update the passed key in the document for provided db and collection"""
+    try:
+        document = fetch_document(client, db, collection)
+        client[db][collection].update_one(
+            {"_id": document["_id"]},
+            {"$set": {key: value}},
+        )
+    except Exception as e:
+        print(e)
+
+
+# Use case specific functions
+def fetch_curr_access_count():
+    client = get_db_client()
+    curr_count = fetch_document(
+        client=client, db=app_config.db, collection=app_config.collection
+    )[app_config.key]
+    app_config.openai_curr_access_count = curr_count
+
+
+def increment_curr_access_count():
+    client = get_db_client()
+    updated_count = app_config.openai_curr_access_count + 1
+    update_document(
+        client=client,
+        db=app_config.db,
+        collection=app_config.collection,
+        key=app_config.key,
+        value=updated_count,
+    )
+    app_config.openai_curr_access_count = updated_count

requirements.txt

@@ -1,4 +1,5 @@
 gradio==3.45.*
 requests==2.29.*
 langchain==0.0.*
-openai==0.27.*
+openai==0.27.*
+pymongo[tls, srv]==4.4.*