fixed history context

app.py

@@ -283,12 +283,12 @@ def __setStartMsg(msg):
     st.session_state.startMsg = msg
 
 
+if "chatHistory" not in st.session_state:
+    st.session_state.chatHistory = []
+
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-if "history" not in st.session_state:
-    st.session_state.history = []
-
 if "buttonValue" not in st.session_state:
     __resetButtonState()
 
@@ -296,27 +296,22 @@ if "startMsg" not in st.session_state:
     st.session_state.startMsg = ""
 
 
-def __getChatMessages(prompt: str):
-    st.session_state.history.append({
-       "role": "user",
-       "content": prompt
-    })
-
+def __getMessages():
     def getContextSize():
-        currContextSize = countTokens(SYSTEM_MSG) + countTokens(st.session_state.history) + 100
+        currContextSize = countTokens(SYSTEM_MSG) + countTokens(st.session_state.messages) + 100
         pprint(f"{currContextSize=}")
         return currContextSize
 
     while getContextSize() > MAX_CONTEXT:
         pprint("Context size exceeded, removing first message")
-        st.session_state.history.pop(0)
+        st.session_state.messages.pop(0)
 
-    return st.session_state.history
+    return st.session_state.messages
 
 
 def predict(prompt):
     messagesFormatted = [{"role": "system", "content": SYSTEM_MSG}]
-    messagesFormatted.extend(__getChatMessages(prompt))
+    messagesFormatted.extend(__getMessages())
     contextSize = countTokens(messagesFormatted)
     pprint(f"{contextSize=} | {MODEL}")
 
@@ -356,10 +351,10 @@ st.title("Kommuneity Story Creator 📖")
 if not (st.session_state["buttonValue"] or st.session_state["startMsg"]):
     st.button(START_MSG, on_click=lambda: __setStartMsg(START_MSG))
 
-for message in st.session_state.messages:
-    role = message["role"]
-    content = message["content"]
-    imagePath = message.get("image")
+for chat in st.session_state.chatHistory:
+    role = chat["role"]
+    content = chat["content"]
+    imagePath = chat.get("image")
     avatar = AI_ICON if role == "assistant" else USER_ICON
     with st.chat_message(role, avatar=avatar):
         st.markdown(content)
@@ -373,7 +368,8 @@ if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_s
     with st.chat_message("user", avatar=USER_ICON):
         st.markdown(prompt)
     pprint(f"{prompt=}")
-    st.session_state.messages.append({"role": "user", "content": prompt })
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.session_state.chatHistory.append({"role": "user", "content": prompt })
 
     with st.chat_message("assistant", avatar=AI_ICON):
         responseContainer = st.empty()
@@ -454,5 +450,9 @@ if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_s
         st.session_state.messages.append({
             "role": "assistant",
             "content": response,
+        })
+        st.session_state.chatHistory.append({
+            "role": "assistant",
+            "content": response,
             "image": imagePath,
         })