Update app.py

app.py

@@ -5,6 +5,10 @@ import os
 # Set your OpenAI API key
 openai.api_key = os.getenv('OPENAI_API_KEY')
 
+# Pricing constants
+INPUT_COST_PER_TOKEN = 0.50 / 1_000_000
+OUTPUT_COST_PER_TOKEN = 1.50 / 1_000_000
+
 def print_like_dislike(x: gr.LikeData):
     print(x.index, x.value, x.liked)
 
@@ -19,16 +23,20 @@ def add_file(history, file):
     return history
 
 def initialize_chat():
-    # This function initializes the chat with a "Hello" message.
-    return [(None, "Hello, my name is <strong>Andrea</strong>, I'm a <em>Friendly Chatbot</em> and will help you with your learning journey. <br>Select a question from below to start!")]
-
-chat_history = initialize_chat()
+    # This function initializes the chat with an initial question.
+    initial_question = "I'm 14 years old female and want to become a graphic designer. I'm living in Uttar Pradesh in India. How can I start?"
+    chat_history = [(None, initial_question)]
+    response, follow_up_questions, input_tokens, output_tokens, cost = generate_response(initial_question)
+    chat_history.append((None, response))
+    return chat_history, follow_up_questions, input_tokens, output_tokens, cost
 
 def generate_response(selected_question):
-    global chat_history
     prompt = selected_question  # Ensure selected_question is a string
     
     try:
+        # Calculate input tokens
+        input_tokens = len(openai.Completion.create(engine="text-davinci-003", prompt=prompt, max_tokens=1).usage['total_tokens']) - 1
+        
         response = openai.ChatCompletion.create(
             model="gpt-3.5-turbo",
             messages=[
@@ -37,9 +45,12 @@ def generate_response(selected_question):
             ],
             max_tokens=150,
             temperature=0.7,
-        ).choices[0].message['content'].strip()
+        )
         
-        follow_up_prompt = f"Based on the following response, suggest three follow-up questions: {response}"
+        output_text = response.choices[0].message['content'].strip()
+        output_tokens = response.usage['completion_tokens']
+        
+        follow_up_prompt = f"Based on the following response, suggest three follow-up questions: {output_text}"
         follow_up_response = openai.ChatCompletion.create(
             model="gpt-3.5-turbo",
             messages=[
@@ -48,20 +59,30 @@ def generate_response(selected_question):
             ],
             max_tokens=100,
             temperature=0.7,
-        ).choices[0].message['content'].strip()
+        )
 
-        follow_up_questions = follow_up_response.split('\n')
+        follow_up_questions = follow_up_response.choices[0].message['content'].strip().split('\n')
         topics_str = "Topic analysis not available"
         
-        # Adjusted to return a list of tuples as expected by the Chatbot component
-        new_response = (None, response + "\n\nTopics: " + topics_str)
-        chat_history.append(new_response)
+        # Calculate the total tokens used
+        total_input_tokens = input_tokens + response.usage['prompt_tokens']
+        total_output_tokens = output_tokens + follow_up_response.usage['completion_tokens']
+        
+        # Calculate cost
+        input_cost = total_input_tokens * INPUT_COST_PER_TOKEN
+        output_cost = total_output_tokens * OUTPUT_COST_PER_TOKEN
+        total_cost = input_cost + output_cost
+
+        # Adjusted to return the response and follow-up questions
+        new_response = output_text + "\n\nTopics: " + topics_str
     except Exception as e:
-        new_response = (None, f"Error generating response: {e}")
-        chat_history.append(new_response)
+        new_response = f"Error generating response: {e}"
         follow_up_questions = []
+        total_input_tokens = 0
+        total_output_tokens = 0
+        total_cost = 0.0
     
-    return chat_history, follow_up_questions
+    return new_response, follow_up_questions, total_input_tokens, total_output_tokens, total_cost
 
 def update_suggested_questions(follow_up_questions):
     return gr.Markdown.update(value="\n".join(f"* {q}" for q in follow_up_questions))
@@ -88,6 +109,9 @@ css = """
 }
 """
 
+# Initialize the chat history and suggested questions
+chat_history, initial_suggested_questions, initial_input_tokens, initial_output_tokens, initial_cost = initialize_chat()
+
 with gr.Blocks(css=css) as demo:
     with gr.Row():
         with gr.Column(scale=1):
@@ -97,10 +121,15 @@ with gr.Blocks(css=css) as demo:
             In the realm of digital communication, the development of an advanced chatbot that incorporates topic modeling represents a significant leap towards enhancing user interaction and maintaining focus during conversations. This innovative chatbot design is specifically engineered to streamline discussions by guiding users to select from a curated list of suggested questions. This approach is crafted to mitigate the risk of diverging into off-topic dialogues, which are common pitfalls in conventional chatbot systems.
             """
             )
-            suggested_questions = gr.Markdown("### Suggested Questions:\n\n* No questions yet.")
+            suggested_questions = gr.Markdown(
+                value="### Suggested Questions:\n\n" + "\n".join(f"* {q}" for q in initial_suggested_questions)
+            )
+            token_info = gr.Markdown(
+                value=f"### Token Usage:\n\n* Input Tokens: {initial_input_tokens}\n* Output Tokens: {initial_output_tokens}\n* Total Cost: ${initial_cost:.4f}"
+            )
         with gr.Column(scale=1, elem_id="chat-container"):
             chatbot = gr.Chatbot(
-                initialize_chat(),
+                value=chat_history,
                 elem_id="chatbot",
                 bubble_full_width=False,
                 label="Safe Chatbot v1",
@@ -111,7 +140,7 @@ with gr.Blocks(css=css) as demo:
                 txt = gr.Textbox(scale=4, show_label=False, placeholder="Select Question", container=False, interactive=False)  # Adjust based on need
                 btn = gr.Button("Submit")
 
-            btn.click(fn=generate_response, inputs=[txt], outputs=[chatbot, suggested_questions])
+            btn.click(fn=generate_response, inputs=[txt], outputs=[chatbot, suggested_questions, token_info])
 
             chatbot.like(print_like_dislike, None, None)