Fix tools previous history truncating (#19)

* Fix tools previous history trucating Signed-off-by: Sanket <[email protected]> * revert import change Signed-off-by: Sanket <[email protected]> --------- Signed-off-by: Sanket <[email protected]>
infracloudio · Aug 26, 2024 · ef23b7f · ef23b7f
1 parent 900907d
commit ef23b7f
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/multi_tenant_rag.py b/multi_tenant_rag.py
@@ -21,6 +21,11 @@
     setup_agent,
 )
 
+SYSTEM = "system"
+USER = "user"
+ASSISTANT = "assistant"
+
+
 logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO)
 
 
@@ -115,7 +120,7 @@ def main():
     )
 
     chat_history = st.session_state.get(
-        "chat_history", [{"role": "system", "content": system_instructions.content}]
+        "chat_history", [{"role": SYSTEM, "content": system_instructions.content}]
     )
 
     for message in chat_history[1:]:
@@ -158,7 +163,7 @@ def main():
         )
 
     if question:
-        st.chat_message("user").markdown(question)
+        st.chat_message(USER).markdown(question)
         with st.spinner():
             if use_tools:
                 answer = agent_executor.invoke(
@@ -167,8 +172,8 @@ def main():
                         "chat_history": chat_history,
                     }
                 )["output"]
-                with st.chat_message("assistant"):
-                    answer = st.write(answer)
+                with st.chat_message(ASSISTANT):
+                    st.write(answer)
                     logger.info(f"answer: {answer}")
             else:
                 answer = rag.query_docs(
@@ -179,12 +184,12 @@ def main():
                     chat_history=chat_history,
                     use_reranker=use_reranker,
                 )
-                with st.chat_message("assistant"):
+                with st.chat_message(ASSISTANT):
                     answer = st.write_stream(answer)
                     logger.info(f"answer: {answer}")
 
-            chat_history.append({"role": "user", "content": question})
-            chat_history.append({"role": "assistant", "content": answer})
+            chat_history.append({"role": USER, "content": question})
+            chat_history.append({"role": ASSISTANT, "content": answer})
             st.session_state["chat_history"] = chat_history