Separate caching logic per page

aws-ia · Sep 27, 2024 · 5e971ec · 5e971ec
1 parent 6072434
commit 5e971ec
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git a/cdk/examples/generative_ai_rag/web-app/pages/foundation_model.py b/cdk/examples/generative_ai_rag/web-app/pages/foundation_model.py
@@ -17,15 +17,15 @@ def get_parameter(name):
 
 st.title("💬 Chat")
 
-if "messages" not in st.session_state:
-    st.session_state["messages"] = [{"role": "assistant", "content": "Ask me something about reinvent 2024?"}]
+if "foundation_messages" not in st.session_state:
+    st.session_state["foundation_messages"] = [{"role": "assistant", "content": "Ask me something about reinvent 2024?"}]
 
-for msg in st.session_state.messages:
+for msg in st.session_state.foundation_messages:
     st.chat_message(msg["role"]).write(msg["content"])
 
 if prompt := st.chat_input():
     st.chat_message("user").write(prompt)
-    st.session_state.messages.append(
+    st.session_state.foundation_messages.append(
         {
             "role": "user",
             "content": prompt
@@ -62,7 +62,7 @@ def get_parameter(name):
 
     except (ClientError, Exception) as e:
         msg = f"ERROR: Can't invoke '{model_id}'. Reason: {e}"
-        st.session_state.messages.append(
+        st.session_state.foundation_messages.append(
             {
                 "role": "assistant",
                 "content": msg
@@ -75,7 +75,7 @@ def get_parameter(name):
 
     # Extract and print the response text.
     msg = model_response["content"][0]["text"]
-    st.session_state.messages.append(
+    st.session_state.foundation_messages.append(
         {
             "role": "assistant",
             "content": msg

diff --git a/cdk/examples/generative_ai_rag/web-app/pages/rag_integration.py b/cdk/examples/generative_ai_rag/web-app/pages/rag_integration.py
@@ -16,15 +16,15 @@ def get_parameter(name):
 
 st.title("💬 Chat")
 
-if "messages" not in st.session_state:
-    st.session_state["messages"] = [{"role": "assistant", "content": "Ask me something about reinvent 2024?"}]
+if "rag_messages" not in st.session_state:
+    st.session_state["rag_messages"] = [{"role": "assistant", "content": "Ask me something about reinvent 2024?"}]
 
-for msg in st.session_state.messages:
+for msg in st.session_state.rag_messages:
     st.chat_message(msg["role"]).write(msg["content"])
 
 if prompt := st.chat_input():
     st.chat_message("user").write(prompt)
-    st.session_state.messages.append(
+    st.session_state.rag_messages.append(
         {
             "role": "user",
             "content": prompt
@@ -49,7 +49,7 @@ def get_parameter(name):
             chunks = event.get('chunk').get('attribution').get('citations')
             for chunk in chunks:
                 msg = chunk['generatedResponsePart']['textResponsePart']['text']
-                st.session_state.messages.append(
+                st.session_state.rag_messages.append(
                     {
                         "role": "assistant",
                         "content": msg
@@ -58,7 +58,7 @@ def get_parameter(name):
                 st.chat_message("assistant").write(msg)
     except:
             msg = "Please give me a better prompt"
-            st.session_state.messages.append(
+            st.session_state.rag_messages.append(
                 {
                     "role": "assistant",
                     "content": msg