diff --git a/python/src/multi_agent_orchestrator/agents/amazon_bedrock_agent.py b/python/src/multi_agent_orchestrator/agents/amazon_bedrock_agent.py
index 0230b79..b208a49 100644
--- a/python/src/multi_agent_orchestrator/agents/amazon_bedrock_agent.py
+++ b/python/src/multi_agent_orchestrator/agents/amazon_bedrock_agent.py
@@ -59,12 +59,18 @@ async def process_request(
         :param additional_params: Optional additional parameters as key-value pairs.
         :return: A ConversationMessage object containing the agent's response.
         """
+
+        session_state = {}
+        if (additional_params and 'sessionState' in additional_params):
+            session_state = additional_params['sessionState']
         try:
+            print(session_state)
             response = self.client.invoke_agent(
                 agentId=self.agent_id,
                 agentAliasId=self.agent_alias_id,
                 sessionId=session_id,
-                inputText=input_text
+                inputText=input_text,
+                sessionState=session_state
             )
 
             completion = ""
diff --git a/python/src/multi_agent_orchestrator/storage/in_memory_chat_storage.py b/python/src/multi_agent_orchestrator/storage/in_memory_chat_storage.py
index 293bf0e..f24308d 100644
--- a/python/src/multi_agent_orchestrator/storage/in_memory_chat_storage.py
+++ b/python/src/multi_agent_orchestrator/storage/in_memory_chat_storage.py
@@ -59,6 +59,8 @@ async def fetch_all_chats(
             if stored_user_id == user_id and stored_session_id == session_id:
                 for message in messages:
                     new_content = message.content if message.content else []
+                    if len(new_content) > 0 and message.role == "assistant":
+                        new_content[0]['text'] = f"[{agent_id}] {new_content[0]['text']}"
                     all_messages.append(TimestampedMessage(
                         role=message.role,
                         content=new_content,