Automatic Memory Wrapper for NeMo Agent Toolkit Agents

Complexity: 🟨 Intermediate

The auto_memory_agent wraps any agent to provide automatic memory capture and retrieval without requiring the LLM to invoke memory tools.

Why Use This?

Traditional tool-based memory:

• LLMs may forget to call memory tools
• Memory capture is inconsistent
• Requires explicit memory tool configuration

Automatic memory wrapper:

• Guaranteed capture: User messages and agent responses are automatically stored
• Automatic retrieval: Relevant context is injected before each agent call
• Memory backend agnostic: Works with Zep, Mem0, Redis, or any MemoryEditor
• Universal compatibility: Wraps any agent type (ReAct, ReWOO, Tool Calling, etc.)

Quick Start

Basic Configuration

memory:
  zep_memory:
    _type: nat.plugins.zep_cloud/zep_memory

functions:
  my_react_agent:
    _type: react_agent
    llm_name: nim_llm
    tool_names: [calculator]

workflow:
  _type: auto_memory_agent
  inner_agent_name: my_react_agent
  memory_name: zep_memory
  llm_name: nim_llm

Install this Workflow

From the root directory of the NeMo Agent Toolkit repository, run the following commands:

uv pip install -e examples/agents

Set Up API Keys

If you have not already done so, follow the Obtaining API Keys instructions to obtain an NVIDIA API key. You need to set your NVIDIA API key as an environment variable to access NVIDIA AI services:

export NVIDIA_API_KEY=<YOUR_API_KEY>

# Set Zep credentials
export ZEP_API_KEY=<YOUR_ZEP_API_KEY>

If you do not have access to a Zep API key, you can use config_mem0.yml with a Mem0 API key instead:

export MEM0_API_KEY=<YOUR_MEM0_API_KEY>

Running the Example

# Run the agent with Zep
nat run --config_file examples/agents/auto_memory_wrapper/configs/config_zep.yml

# Or with Mem0
nat run --config_file examples/agents/auto_memory_wrapper/configs/config_mem0.yml

Configuration Reference

Required Parameters

Parameter	Description
inner_agent_name	Name of the agent to wrap with automatic memory
memory_name	Name of the memory backend (from memory: section)
llm_name	LLM to use (required by AgentBaseConfig)

Optional Feature Flags

All default to true. Set to false to disable specific behaviors:

Parameter	Default	Description
save_user_messages_to_memory	true	Automatically save user messages before agent processing
retrieve_memory_for_every_response	true	Automatically retrieve and inject memory context
save_ai_messages_to_memory	true	Automatically save agent responses after generation

Memory Backend Parameters

search_params - Passed to memory_editor.search():

search_params:
  mode: "summary"      # Zep: "basic" or "summary"
  top_k: 10           # Maximum memories to retrieve

add_params - Passed to memory_editor.add_items():

add_params:
  ignore_roles: ["assistant"]  # Zep: Exclude roles from graph memory

See config_zep.yml for comprehensive parameter examples.

Multi-Tenant Memory Isolation

User ID is extracted at runtime for memory isolation. Configure it through the front end or session runtime, not the auto_memory_agent workflow block.

User ID Extraction Priority

1. SessionManager.session(user_id=...) - For production with custom auth middleware (recommended)
2. X-User-ID HTTP header - For testing without middleware
3. Console front end user_id - Defaults to "nat_run_user_id" for nat run

Conversation-aware memory backends can also use conversation_id to isolate separate conversations for the same user. For Zep Cloud, if no conversation ID is supplied, the integration uses a deterministic per-user default thread.

Production: Custom Middleware

Create middleware that extracts user ID from your authentication system:

from nat.runtime.session import SessionManager

class AuthenticatedUserManager:
    def __init__(self, user_id: str):
        self._user_id = user_id

    def get_id(self) -> str:
        return self._user_id

# In your request handler
async def handle_request(request):
    # Extract from JWT, OAuth, API key, etc.
    user_id = extract_user_from_jwt(request.headers["authorization"])

    async with session_manager.session(user_id=user_id, http_connection=request) as session:
        result = await session.run(user_input)
    return result

Testing: X-User-ID Header

For quick testing without custom middleware:

curl -X POST http://localhost:8000/chat \
  -H "Content-Type: application/json" \
  -H "X-User-ID: test_user_123" \
  -H "conversation-id: test_conv_001" \
  -d '{"messages": [{"role": "user", "content": "Hello!"}]}'

Local Development: Console User and Conversation IDs

For nat run, set --user_id to control memory isolation and --conversation_id to isolate a specific conversation:

nat run --config_file examples/agents/auto_memory_wrapper/configs/config_zep.yml \
  --user_id test_user_123 \
  --conversation_id test_conv_001

Advanced Example

See config_zep.yml for a fully-commented configuration with all available parameters.

workflow:
  _type: auto_memory_agent
  inner_agent_name: my_react_agent
  memory_name: zep_memory
  llm_name: nim_llm

  # Feature flags (optional - all default to true)
  save_user_messages_to_memory: true
  retrieve_memory_for_every_response: true
  save_ai_messages_to_memory: true

  # Memory retrieval configuration (optional)
  search_params:
    mode: "summary"  # Zep: "basic" (fast) or "summary" (comprehensive)
    top_k: 5         # Maximum number of memories to retrieve

  # Memory storage configuration (optional)
  add_params:
    ignore_roles: ["assistant"]  # Zep: Exclude assistant messages from graph

Important Notes

1. User ID is runtime/front-end scoped - Set via SessionManager.session(user_id=...), X-User-ID, or nat run --user_id
2. Memory backends are interchangeable - Works with any implementation of MemoryEditor interface

Examples

See configs/ directory:

• config_zep.yml - Zep Cloud memory backend with all parameters documented
• config_mem0.yml - Mem0 memory backend (alternative if Zep is unavailable)