diff --git a/agent/deep_assistant.py b/agent/deep_assistant.py
index e973e67..49640e5 100644
--- a/agent/deep_assistant.py
+++ b/agent/deep_assistant.py
@@ -462,7 +462,7 @@ def create_custom_cli_agent(
         # Add memory middleware
         if enable_memory:
             agent_middleware.append(
-                AgentMemoryMiddleware(settings=settings, assistant_id=assistant_id)
+                CustomAgentMemoryMiddleware(settings=settings, assistant_id=assistant_id)
             )
 
         # Add skills middleware