diff --git a/chatbot_ui/app/run.py b/chatbot_ui/app/run.py
index 52ce50c3..ae2ff037 100644
--- a/chatbot_ui/app/run.py
+++ b/chatbot_ui/app/run.py
@@ -58,22 +58,21 @@ def predict(message, state):
     # extract arch_state from metadata and store it in gradio session state
     # this state must be passed back to the gateway in the next request
     response_json = json.loads(raw_response.text)
-    arch_state = {}
     if response_json:
-        metadata = response_json.get("metadata", {})
-        if metadata:
-            arch_state_str = metadata.get(ARCH_STATE_HEADER, "{}")
-            arch_state = json.loads(arch_state_str)
-    if arch_state:
-        state["arch_state"] = arch_state
-
-    arch_messages_str = arch_state.get("messages", "[]")
-    arch_messages = json.loads(arch_messages_str)
+        # load arch_state from metadata
+        arch_state_str = response_json.get("metadata", {}).get(ARCH_STATE_HEADER, "{}")
+        # parse arch_state into json object
+        arch_state = json.loads(arch_state_str)
+        # load messages from arch_state
+        arch_messages_str = arch_state.get("messages", "[]")
+        # parse messages into json object
+        arch_messages = json.loads(arch_messages_str)
+        # append messages from arch gateway to history
+        for message in arch_messages:
+            history.append(message)
 
     content = response.choices[0].message.content
 
-    for message in arch_messages:
-        history.append(message)
     history.append({"role": "assistant", "content": content, "model": response.model})
     history_view = [h for h in history if h["role"] != "tool" and "content" in h]
     messages = [