button feature

app.py

@@ -2,6 +2,9 @@ import streamlit as st
 import os
 import datetime as DT
 import pytz
+import time
+import json
+import re
 from transformers import AutoTokenizer
 
 from dotenv import load_dotenv
@@ -12,7 +15,10 @@ from groq import Groq
 client = Groq(
     api_key=os.environ.get("GROQ_API_KEY"),
 )
+
 MODEL = "llama-3.1-70b-versatile"
+JSON_SEPARATOR = ">>>>"
+
 
 tokenizer = AutoTokenizer.from_pretrained("Xenova/Meta-Llama-3.1-Tokenizer")
 
@@ -24,9 +30,18 @@ def countTokens(text):
     return len(tokens)
 
 
-SYSTEM_MSG = """
+SYSTEM_MSG = f"""
 You're an storytelling assistant who guides users through four phases of narrative development, helping them craft compelling personal or professional stories. The story created should be in simple language, yet evoke great emotions.
 Ask one question at a time, give the options in a well formatted manner in different lines
+If your response has number of options to choose from, only then append your final response with this exact keyword "{JSON_SEPARATOR}", and only after this, append with the JSON of options to choose from. The JSON should be of the format:
+{{
+    "options": [
+        {{ "id": "1", "label": "Option 1"}},
+        {{ "id": "2", "label": "Option 2"}},
+    ]
+}}
+Do not write "Choose one of the options below:"
+Keep options to less than 9
 
 # Tier 1: Story Creation
 You initiate the storytelling process through a series of engaging prompts:
@@ -61,13 +76,13 @@ How we're different: A marketing story
 
 Guided Storytelling Framework:
 You then lead users through a structured narrative development via the following prompts:
-Describe the day it happened
-What was the Call to Action / Invitation
-Describing the obstacles (up to three) in 4 lines
-Exploring emotions/fears experienced during the incident
-Recognize the helpers / any objects of help in the incident
-Detailing the resolution / Reaching the final goal
-Reflecting on personal growth or lessons learned (What did you do that changed your life forever?)
+- Describe the day it happened
+- What was the Call to Action / Invitation
+- Describing the obstacles (up to three) in 4 lines
+- Exploring emotions/fears experienced during the incident
+- Recognize the helpers / any objects of help in the incident
+- Detailing the resolution / Reaching the final goal
+- Reflecting on personal growth or lessons learned (What did you do that changed your life forever?)
 
 Now, show the story created so far, and ask for confirmation before proceeding to the next tier.
 
@@ -99,11 +114,13 @@ Tips for maximising emotional resonance and memorability
 By guiding users through these three tiers, you aim to cater to novice storytellers, offering a comprehensive platform for narrative skill development through its adaptive approach.
 You end it with the final story and seeking any suggestions from the user to refine the story further.
 Once the user confirms, you congratulate them with emojis on completing the story and provide the final story in a beatifully formatted manner.
+Note that the final story should include twist, turns and events that make it really engaging and enjoyable to read.
 
 """
 
 USER_ICON = "man.png"
 AI_ICON = "Kommuneity.png"
+START_MSG = "I want to create a story 😊"
 
 st.set_page_config(
     page_title="Kommuneity Story Creator",
@@ -113,7 +130,7 @@ st.set_page_config(
 ipAddress = st.context.headers.get("x-forwarded-for")
 
 
-def __nowInIST():
+def __nowInIST() -> DT.datetime:
     return DT.datetime.now(pytz.timezone("Asia/Kolkata"))
 
 
@@ -126,6 +143,32 @@ def pprint(log: str):
 pprint("\n")
 
 
+def __isInvalidResponse(response: str):
+    if len(re.findall(r'\n[a-z]', response)) > 3:
+        return True
+
+    if ('\n{\n    "options"' in response) and (JSON_SEPARATOR not in response):
+        return True
+
+
+def __resetButtonState():
+    st.session_state["buttonValue"] = ""
+
+
+def __setStartMsg(msg):
+    st.session_state.startMsg = msg
+
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+if "buttonValue" not in st.session_state:
+    __resetButtonState()
+
+if "startMsg" not in st.session_state:
+    st.session_state.startMsg = ""
+
+
 def predict(prompt):
     historyFormatted = [{"role": "system", "content": SYSTEM_MSG}]
     historyFormatted.extend(st.session_state.messages)
@@ -149,31 +192,9 @@ def predict(prompt):
             yield chunkContent
 
 
-st.title("Let's create your story 📖")
-with st.chat_message("user", avatar=AI_ICON):
-    st.write("Type 'Hi' to start")
-
-# st.markdown(
-#     """
-#     <style>
-#         :root {
-#             --bg-color: white;
-#             --text-color: black;
-#         }
-
-#         .main {
-#             background-color: var(--bg-color);
-#             color: var(--text-color);
-#             transition: background-color 0.3s, color 0.3s;
-#         }
-#     </style>
-#     """,
-#     unsafe_allow_html=True
-# )
-
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
+st.title("Kommuneity Story Creator 📖")
+if not st.session_state.startMsg:
+    st.button(START_MSG, on_click=lambda: __setStartMsg(START_MSG))
 
 for message in st.session_state.messages:
     role = message["role"]
@@ -182,17 +203,61 @@ for message in st.session_state.messages:
     with st.chat_message(role, avatar=avatar):
         st.markdown(content)
 
-if prompt := st.chat_input():
+if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_state.startMsg):
+    __resetButtonState()
+    __setStartMsg("")
+
     with st.chat_message("user", avatar=USER_ICON):
         st.markdown(prompt)
     pprint(f"{prompt=}")
     st.session_state.messages.append({"role": "user", "content": prompt })
 
     with st.chat_message("assistant", avatar=AI_ICON):
-        responseGenerator = predict(prompt)
-        response = st.write_stream(responseGenerator)
-        if not response:
+        placeholder = st.empty()
+
+        def getResponse():
+            response = ""
             responseGenerator = predict(prompt)
-            response = st.write_stream(responseGenerator)
-    pprint(f"{response=}")
+
+            for chunk in responseGenerator:
+                response += chunk
+                if __isInvalidResponse(response):
+                    return
+
+                if JSON_SEPARATOR not in response:
+                    placeholder.markdown(response)
+
+            return response
+
+        response = getResponse()
+        while not response:
+            pprint("Empty response. Retrying..")
+            time.sleep(0.5)
+            response = getResponse()
+
+        pprint(f"{response=}")
+
+        def selectButton(optionLabel):
+            st.session_state["buttonValue"] = optionLabel
+            pprint(f"Selected: {optionLabel}")
+
+        responseParts = response.split(JSON_SEPARATOR)
+        if len(responseParts) > 1:
+            [response, jsonStr] = responseParts
+
+            try:
+                json.loads(jsonStr)
+                jsonObj = json.loads(jsonStr)
+                options = jsonObj["options"]
+
+                for option in options:
+                    st.button(
+                        option["label"],
+                        key=option["id"],
+                        on_click=lambda label=option["label"]: selectButton(label)
+                    )
+                # st.code(jsonStr, language="json")
+            except Exception as e:
+                pprint(e)
+
     st.session_state.messages.append({"role": "assistant", "content": response})