added cost calculator

app.py

@@ -1,9 +1,11 @@
 from openai import OpenAI
 import streamlit as st
-from utils import im_2_b64
+from utils import im_2_b64, calculate_cost
 import pickle
 from upload import upload_file, get_file
 
+share_keys = ["messages", "cost"]
+
 
 st.title("ChatGPT with Vision")
 
@@ -13,12 +15,16 @@ if "messages" not in st.session_state:
     st.session_state.messages = []
 if "uploader_key" not in st.session_state:
     st.session_state["uploader_key"] = 0
+if "cost" not in st.session_state:
+    st.session_state.cost = []
 
 if len(st.session_state.messages) == 0 and "id" in st.query_params:
     with st.spinner("Loading chat..."):
         id = st.query_params["id"]
         data = get_file(id, 'chatgpt-vision-007')
-        st.session_state.messages = pickle.loads(data)
+        obj = pickle.loads(data)
+        for k, v in obj.items():
+            st.session_state[k] = v
 
 def clear_uploader():
     st.session_state["uploader_key"] += 1
@@ -31,7 +37,11 @@ def undo():
         st.rerun()
 
 def share():
-    data = pickle.dumps(st.session_state.messages)
+    obj = {}
+    for k in share_keys:
+        if k in st.session_state:
+            obj[k] = st.session_state[k]
+    data = pickle.dumps(obj)
     id = upload_file(data, 'chatgpt-vision-007')
     url = f"https://umbc-nlp-chatgpt-vision.hf.space/?id={id}"
     st.success(f"Share URL: {url}")
@@ -66,6 +76,16 @@ with st.sidebar:
             label_visibility="collapsed",
         )
 
+    with st.expander(f"Total Cost: ${sum(st.session_state.cost):.10f}"):
+        if len(st.session_state.cost) > 0:
+            st.subheader("Cost Breakdown")
+            for i, c in enumerate(st.session_state.cost):
+                st.write(f"Message {i+1}: ${c:.10f}")
+            st.markdown("---")
+            st.write(f"Total: ${sum(st.session_state.cost):.10f}")
+        else:
+            st.write("No cost incurred yet")
+
 
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
@@ -131,6 +151,8 @@ if prompt := st.chat_input("Type a message", key="chat_input", disabled=chat_inp
             max_tokens=max_tokens,
         )
         response = st.write_stream(stream)
+
     push_message("assistant", response)
+    calculate_cost()
     chat_input_disabled = False
     clear_uploader()

utils.py

@@ -1,6 +1,13 @@
 from PIL import Image
 from io import BytesIO
 import base64
+import streamlit as st
+
+
+UNIT_COST = {
+    "user": 0.01,
+    "assistant": 0.03,
+}
 
 
 # Convert Image to Base64 
@@ -13,3 +20,26 @@ def im_2_b64(image):
     img_str = base64.b64encode(buff.getvalue())
     return img_str
 
+
+def calculate_cost():
+    def get_text_cost(text, unit_cost):
+        num_of_words = len(text.split(" "))
+        tokens = max(1000.0 * num_of_words / 750.0, 0.0)
+        tokens = tokens / 1000.0
+        cost = tokens * unit_cost
+        return cost
+    def get_image_cost(unit_cost=0.01):
+        cost = 0.00255 # 512x512 image: https://openai.com/pricing
+        return cost
+
+    messages = st.session_state.messages
+    total_cost = 0
+    for message in messages:
+        role = message["role"]
+        for content in message["content"]:
+            if content["type"] == "image_url":
+                total_cost += get_image_cost(UNIT_COST[role])
+            else:
+                total_cost += get_text_cost(content["text"], UNIT_COST[role])
+    
+    st.session_state.cost.append(total_cost)