fix memoryConfiguration KeyError

weni-ai · Feb 7, 2025 · 274d7cf · 274d7cf
1 parent d5529f5
commit 274d7cf
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/nexus/task_managers/file_database/bedrock.py b/nexus/task_managers/file_database/bedrock.py
@@ -407,6 +407,15 @@ def bedrock_agent_to_supervisor(self, agent_id: str, to_supervisor: bool = True)
         agent_to_update = self.bedrock_agent.get_agent(agentId=agent_id)
         agent_to_update = agent_to_update['agent']
 
+        try:
+            memory_configuration = agent_to_update["memoryConfiguration"]
+        except KeyError:
+            memory_configuration = {
+                "enabledMemoryTypes": ["SESSION_SUMMARY"],
+                "sessionSummaryConfiguration": {"maxRecentSessions": 5},
+                "storageDays": 30
+            }
+
         agent_collaboration = {
             True: "SUPERVISOR",
             False: "DISABLED"
@@ -419,7 +428,7 @@ def bedrock_agent_to_supervisor(self, agent_id: str, to_supervisor: bool = True)
             agentCollaboration=agent_collaboration.get(to_supervisor),
             instruction=agent_to_update["instruction"],
             foundationModel=agent_to_update['foundationModel'],
-            memoryConfiguration=agent_to_update["memoryConfiguration"]
+            memoryConfiguration=memory_configuration,
         )
 
     def invoke_supervisor_stream(