Merge pull request #91 from brnaba-aws/main

Fixed agentId not populated in history section of the prompt for classifier when using in memory storage
awslabs · Nov 20, 2024 · 8a9283f · 8a9283f
2 parents bb317dd + 3156653
commit 8a9283f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git 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
@@ -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,