diff --git a/ci/.nim_models_used.json b/ci/.nim_models_used.json index 2055ebf564..48d369eb21 100644 --- a/ci/.nim_models_used.json +++ b/ci/.nim_models_used.json @@ -11,11 +11,11 @@ }, { "model": "meta/llama-3.3-70b-instruct", - "num_configs": 27 + "num_configs": 26 }, { "model": "nvidia/nemotron-3-super-120b-a12b", - "num_configs": 16 + "num_configs": 17 }, { "model": "meta/llama-3.1-8b-instruct", @@ -53,7 +53,11 @@ "embedders": [ { "model": "nvidia/nv-embedqa-e5-v5", - "num_configs": 23 + "num_configs": 22 + }, + { + "model": "nvidia/nemotron-3-embed-1b", + "num_configs": 1 }, { "model": "nvidia/llama-nemotron-embed-1b-v2", diff --git a/ci/scripts/path_checks.py b/ci/scripts/path_checks.py index 6a03e65615..36d0367272 100644 --- a/ci/scripts/path_checks.py +++ b/ci/scripts/path_checks.py @@ -158,6 +158,7 @@ "file/console", "files/functions", "I/O", + "Illustrative/testing", "include/exclude", "Input/Observation", "input/output", diff --git a/docs/source/build-workflows/memory.md b/docs/source/build-workflows/memory.md index 3b0761a950..575d25e5cb 100644 --- a/docs/source/build-workflows/memory.md +++ b/docs/source/build-workflows/memory.md @@ -21,6 +21,18 @@ The NeMo Agent Toolkit Memory subsystem is designed to store and retrieve a user The memory module is designed to be extensible, allowing developers to create custom memory back-ends, providers in NeMo Agent Toolkit terminology. +## User Identity for Memory Tools + +The built-in `add_memory`, `get_memory`, and `delete_memory` tools bind every operation to an identity. Set an optional `user_id` in the tool configuration for a fixed service or single-user identity. + +```yaml +functions: + get_memory: + _type: get_memory + memory: user_memory + user_id: service_user +``` + ## Included Memory Modules The NeMo Agent Toolkit includes four memory module providers, all of which are available as plugins: * [Mem0](https://mem0.ai/) which is provided by the [`nvidia-nat-mem0ai`](https://pypi.org/project/nvidia-nat-mem0ai/) plugin. @@ -32,6 +44,44 @@ The NeMo Agent Toolkit includes four memory module providers, all of which are a Additional memory backends are available as community plugins: * [Synap](https://maximem.ai) — managed memory layer with user and customer scoping, provided by the [`maximem-synap-nemo-agent-toolkit`](https://pypi.org/project/maximem-synap-nemo-agent-toolkit/) plugin. See `examples/memory/synap/` for usage. ([Open source integration package](https://github.com/maximem-ai/maximem_synap_sdk/tree/main/packages/integrations)) +## Authenticating Memory Tool Users + +Each of the built-in `add_memory`, `get_memory`, and `delete_memory` tools require configuring an identity source: + +- Use `user_id` for a fixed, single-user memory namespace. +- Use `user_id_resolver` for a multi-user application. Its value is the import path of a trusted, zero-argument Python callable that returns the current authenticated user's stable ID. The callable can be synchronous or asynchronous and is invoked for every memory operation. + +For example, application code can obtain a user that authentication middleware has already verified: + +```python +from my_application.request_context import get_authenticated_user + + +def resolve_memory_user_id() -> str: + user = get_authenticated_user() + if user is None: + raise RuntimeError("An authenticated user is required") + return user.user_id +``` + +Reference that callable from each memory tool: + +```yaml +functions: + add_memory: + _type: add_memory + memory: user_memory + user_id_resolver: my_application.auth.resolve_memory_user_id + get_memory: + _type: get_memory + memory: user_memory + user_id_resolver: my_application.auth.resolve_memory_user_id + delete_memory: + _type: delete_memory + memory: user_memory + user_id_resolver: my_application.auth.resolve_memory_user_id +``` + ## Automatic Memory Wrapper Agent The NeMo Agent Toolkit provides an [`auto_memory_agent`](../components/agents/auto-memory-wrapper/index.md) wrapper that adds automatic memory capture and retrieval to any agent without requiring the LLM to invoke memory tools explicitly. @@ -93,12 +143,14 @@ The automatic memory wrapper agent supports several configuration parameters: User ID is automatically extracted at runtime for memory isolation via: 1. `SessionManager.session(user_id=...)` - For production with custom auth middleware (recommended) -2. `X-User-ID` HTTP header - For testing without middleware +2. `X-User-ID` HTTP header - Illustrative/testing only; assumes a trusted upstream proxy authenticates the request and injects this header 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 `nat run`, pass `--conversation_id` when testing independent memory conversations from the CLI. +> Never treat a client-supplied `X-User-ID` header as authentication. + For detailed configuration and usage examples, refer to the `examples/agents/auto_memory_wrapper/README.md` guide. ## Examples diff --git a/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md b/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md index 1f2d3f0459..8c3c4274c9 100644 --- a/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md +++ b/docs/source/components/agents/auto-memory-wrapper/auto-memory-wrapper.md @@ -166,7 +166,7 @@ through the front end or session runtime, not the `auto_memory_agent` workflow b ### 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 +2. **`X-User-ID` HTTP header** - Illustrative/testing only; assumes a hypothetical trusted upstream proxy authenticates the request and injects the header 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. @@ -327,10 +327,14 @@ workflow: ## Important Notes -1. **User ID is runtime/front-end scoped** - Set via `SessionManager.session(user_id=...)`, `X-User-ID`, or `nat run --user_id` +1. **User ID is runtime/front-end scoped** - Set via `SessionManager.session(user_id=...)`, `X-User-ID`, or + `nat run --user_id`. When no runtime identity or header is available, the wrapper falls back to `"default_user"`. + Use this fallback for development and testing only. Production deployments must require an authenticated runtime + identity. 2. **Memory backends are interchangeable** - Works with any implementation of `MemoryEditor` interface 3. **No memory tools needed** - The wrapped agent does not need explicit memory tools configured 4. **Transparent to inner agent** - The wrapped agent is unaware of memory operations +5. **X-User-ID** - We used this header for illustrative purposes only. Do not rely on it for authentication in production. --- diff --git a/docs/source/extend/custom-components/memory.md b/docs/source/extend/custom-components/memory.md index 8d85aec4f9..80e837be73 100644 --- a/docs/source/extend/custom-components/memory.md +++ b/docs/source/extend/custom-components/memory.md @@ -19,6 +19,10 @@ limitations under the License. This documentation presumes familiarity with the NeMo Agent Toolkit [memory module](../../build-workflows/memory.md), [plugin architecture](../plugins.md), the concept of "function registration" using `@register_function`, and how we define [tool](../../build-workflows/functions-and-function-groups/functions.md#agents-and-tools) and workflow configurations in the NeMo Agent Toolkit config described in the [Creating a New Tool and Workflow](../../get-started/tutorials/create-a-new-workflow.md) tutorial. +For applications that expose the built-in memory tools to multiple authenticated users, see +[Authenticating Memory Tool Users](../../build-workflows/memory.md#authenticating-memory-tool-users). Configure a trusted +`user_id_resolver`. + ## Key Memory Module Components * **Memory Data Models** @@ -190,12 +194,14 @@ functions: add_memory: _type: add_memory memory: saas_memory + user_id: user_12 description: | Add any facts about user preferences to long term memory. Always use this if users mention a preference. The input to this tool should be a string that describes the user's preference, not the question or answer. get_memory: _type: get_memory memory: saas_memory + user_id: user_12 description: | Always call this tool before calling any other tools, even if the user does not mention to use it. The question should be about user preferences which will help you format your response. @@ -214,6 +220,7 @@ Explanation: - We define a memory entry named `saas_memory` with `_type: mem0_memory`, using the [Mem0](https://mem0.ai/) provider included in the [`nvidia-nat-mem0ai`](https://pypi.org/project/nvidia-nat-mem0ai/) plugin. - Then we define two tools (functions in NeMo Agent Toolkit terminology) that reference `saas_memory`: `add_memory` and `get_memory`. +- The optional `user_id` is a fixed identity for these tools, alternately `user_id_resolver` can be used to dynamically resolve the user identity at runtime. - Finally, the `agent_memory` workflow references these two tool names. ### Automatic Memory with the Auto-Memory Wrapper diff --git a/examples/RAG/simple_rag/configs/milvus_memory_rag_config.yml b/examples/RAG/simple_rag/configs/milvus_memory_rag_config.yml index f955b09668..a06c812d27 100644 --- a/examples/RAG/simple_rag/configs/milvus_memory_rag_config.yml +++ b/examples/RAG/simple_rag/configs/milvus_memory_rag_config.yml @@ -44,12 +44,14 @@ functions: add_memory: _type: add_memory memory: saas_memory + user_id: user_12 description: | Add any facts about user preferences to long term memory. Always use this if users mention a preference. The input to this tool should be a string that describes the user's preference, not the question or answer. get_memory: _type: get_memory memory: saas_memory + user_id: user_12 description: | Always call this tool before calling any other tools, even if the user does not mention to use it. The question should be about user preferences which will help you format your response. @@ -86,16 +88,14 @@ workflow: IMPORTANT MEMORY TOOL REQUIREMENTS: 1. You MUST call get_memory tool FIRST, before calling any other tools - 2. You MUST use user_id "user_12" for all memory operations - 3. You MUST include ALL required parameters when calling memory tools - 4. When calling add_memory or get_memory, you MUST use the exact format as below, don't include any other content, + 2. You MUST include ALL required parameters when calling memory tools + 3. When calling add_memory or get_memory, you MUST use the exact format as below, don't include any other content, and make sure the input is a valid JSON object. For get_memory tool, you MUST use this exact format: {{ "query": "user preferences", - "top_k": 1, - "user_id": "user_12" + "top_k": 1 }} For add_memory tool, you MUST use this exact format: @@ -110,7 +110,6 @@ workflow: "content": "Hello Alex! I've noted you are looking for a trip to New York." }} ], - "user_id": "user_12", "metadata": {{ "key_value_pairs": {{ "type": "travel", diff --git a/examples/RAG/simple_rag/configs/milvus_memory_rag_tools_config.yml b/examples/RAG/simple_rag/configs/milvus_memory_rag_tools_config.yml index f1953b90c2..0c488348a0 100644 --- a/examples/RAG/simple_rag/configs/milvus_memory_rag_tools_config.yml +++ b/examples/RAG/simple_rag/configs/milvus_memory_rag_tools_config.yml @@ -44,12 +44,14 @@ functions: add_memory: _type: add_memory memory: saas_memory + user_id: user_12 description: | Add any facts about user preferences to long term memory. Always use this if users mention a preference. The input to this tool should be a string that describes the user's preference, not the question or answer. get_memory: _type: get_memory memory: saas_memory + user_id: user_12 description: | Always call this tool before calling any other tools, even if the user does not mention to use it. The question should be about user preferences which will help you format your response. diff --git a/examples/agents/auto_memory_wrapper/README.md b/examples/agents/auto_memory_wrapper/README.md index c7e07dc6d8..aece1e59b0 100644 --- a/examples/agents/auto_memory_wrapper/README.md +++ b/examples/agents/auto_memory_wrapper/README.md @@ -133,7 +133,7 @@ User ID is extracted at runtime for memory isolation. Configure it through the f ### 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 +2. **`X-User-ID` HTTP header** - Illustrative/testing only; assumes a hypothetical trusted upstream proxy authenticates the request and injects the header 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. @@ -175,6 +175,9 @@ curl -X POST http://localhost:8000/chat \ -d '{"messages": [{"role": "user", "content": "Hello!"}]}' ``` +The example usage of the `X-User-ID` header is for illustrative purposes only; do not accept this header directly from untrusted clients. + + ### 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: @@ -215,6 +218,7 @@ workflow: 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 +3. `X-User-ID` HTTP header - Illustrative/testing only; assumes a trusted upstream proxy authenticates the request and injects this header ## Examples diff --git a/examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/configs/config.yml b/examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/configs/config.yml index 72a68fca10..1d193bdb62 100644 --- a/examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/configs/config.yml +++ b/examples/frameworks/semantic_kernel_demo/src/nat_semantic_kernel_demo/configs/config.yml @@ -34,6 +34,7 @@ functions: add_memory: _type: add_memory memory: saas_memory + user_id: user_1 description: | Add any facts about user preferences to long term memory. Always use this if users mention a preference. The input to this tool should be a string that describes the user's preference, not the question or answer. @@ -41,6 +42,7 @@ functions: get_memory: _type: get_memory memory: saas_memory + user_id: user_1 description: | Always call this tool before calling any other tools, even if the user does not mention to use it. The question should be about user preferences which will help you format your response. @@ -82,16 +84,14 @@ workflow: You have access to long term memory. IMPORTANT MEMORY TOOL REQUIREMENTS: 1. You MUST call get_memory tool FIRST, before calling any other tools - 2. You MUST use user_id "user_1" for all memory operations - 3. You MUST include ALL required parameters when calling memory tools - 4. When calling add_memory or get_memory, you MUST use the exact format as below, don't include any other content, + 2. You MUST include ALL required parameters when calling memory tools + 3. When calling add_memory or get_memory, you MUST use the exact format as below, don't include any other content, and make sure the input is a valid JSON object. For get_memory tool, you MUST use this exact format: { "query": "user preferences", - "top_k": 1, - "user_id": "user_1" + "top_k": 1 } For add_memory tool, you MUST use this exact format: @@ -106,7 +106,6 @@ workflow: "content": "Hello Alex! I've noted you are looking for a trip to New York." } ], - "user_id": "user_1", "metadata": { "key_value_pairs": { "type": "travel", diff --git a/examples/memory/memmachine/memmachine_memory_example.ipynb b/examples/memory/memmachine/memmachine_memory_example.ipynb index 5a5a080a3e..9762043ffa 100644 --- a/examples/memory/memmachine/memmachine_memory_example.ipynb +++ b/examples/memory/memmachine/memmachine_memory_example.ipynb @@ -435,17 +435,19 @@ " get_memory:\n", " _type: get_memory\n", " memory: memmachine_memory\n", + " user_id: \"{user_id}\"\n", " description: |\n", " Retrieve memories relevant to a query. Always call this tool first to check\n", " for existing user preferences or facts.\n", - " Use the exact JSON format with user_id, query, and top_k parameters.\n", + " Use the exact JSON format with query and top_k parameters.\n", "\n", " add_memory:\n", " _type: add_memory\n", " memory: memmachine_memory\n", + " user_id: \"{user_id}\"\n", " description: |\n", " Add facts about user preferences or information to long-term memory.\n", - " Use the exact JSON format with user_id, memory, conversation (optional), metadata, and tags.\n", + " Use the exact JSON format with memory, conversation (optional), metadata, and tags.\n", "\n", "workflow:\n", " _type: react_agent\n", diff --git a/examples/memory/redis/configs/config.yml b/examples/memory/redis/configs/config.yml index b2c853e4c8..d2a32ffb7e 100644 --- a/examples/memory/redis/configs/config.yml +++ b/examples/memory/redis/configs/config.yml @@ -26,14 +26,15 @@ general: llms: nim_llm: _type: nim - model_name: meta/llama-3.3-70b-instruct + model_name: nvidia/nemotron-3-super-120b-a12b temperature: 0.7 - max_tokens: 1024 + chat_template_kwargs: + enable_thinking: false embedders: - nv-embedqa-e5-v5: + nemotron-3-embed-1b: _type: nim - model_name: nvidia/nv-embedqa-e5-v5 + model_name: nvidia/nemotron-3-embed-1b memory: redis_memory: @@ -42,29 +43,30 @@ memory: db: 0 port: 6379 key_prefix: nat - embedder: nv-embedqa-e5-v5 + embedder: nemotron-3-embed-1b functions: get_memory: _type: get_memory memory: redis_memory + user_id: redis description: | Always call this tool before calling any other tools, even if the user does not mention to use it. The question should be about user preferences which will help you format your response. For example: "How does the user like responses formatted?". - Use "redis" for the user_id memory_add: _type: add_memory memory: redis_memory + user_id: redis description: | Add any facts about user preferences to long term memory. Always use this if users mention a preference. The input to this tool should be a string that describes the user's preference, not the question or answer. - Use "redis" for the user_id. Be sure to include any relevant tags for the memory as a list of strings. Also include key value pairs for metadata + Include any relevant tags for the memory as a list of strings and key-value pairs for metadata. workflow: - _type: react_agent + _type: tool_calling_agent tool_names: [memory_add, get_memory] description: "A chat agent that can make memories and also recall memories" llm_name: nim_llm @@ -75,14 +77,13 @@ workflow: IMPORTANT MEMORY TOOL REQUIREMENTS: 1. You MUST use get_memory tool with the exact JSON format below - 2. You MUST include ALL required parameters (query, top_k, user_id) + 2. You MUST include ALL required parameters (query, top_k) 3. The input MUST be a valid JSON object with no extra text or formatting For get_memory tool, you MUST use this exact format: {{ "query": "your search query here", - "top_k": 5, - "user_id": "redis" + "top_k": 5 }} For memory_add tool, you MUST use this exact format: @@ -97,7 +98,6 @@ workflow: "content": "Hello Alex! I've noted you are looking for a trip to New York." }} ], - "user_id": "redis", "metadata": {{ "key_value_pairs": {{ "type": "travel", diff --git a/packages/nvidia_nat_core/src/nat/tool/memory_tools/add_memory_tool.py b/packages/nvidia_nat_core/src/nat/tool/memory_tools/add_memory_tool.py index d43f95440e..282291531a 100644 --- a/packages/nvidia_nat_core/src/nat/tool/memory_tools/add_memory_tool.py +++ b/packages/nvidia_nat_core/src/nat/tool/memory_tools/add_memory_tool.py @@ -20,22 +20,20 @@ from nat.builder.builder import Builder from nat.builder.function_info import FunctionInfo from nat.cli.register_workflow import register_function -from nat.data_models.component_ref import MemoryRef -from nat.data_models.function import FunctionBaseConfig from nat.memory.models import MemoryItem +from nat.tool.memory_tools.common import AddMemoryInput +from nat.tool.memory_tools.common import MemoryToolConfigBase +from nat.tool.memory_tools.common import resolve_memory_user_id logger = logging.getLogger(__name__) -class AddToolConfig(FunctionBaseConfig, name="add_memory"): +class AddToolConfig(MemoryToolConfigBase, name="add_memory"): """Function to add memory to a hosted memory platform.""" description: str = Field(default=("Tool to add a memory about a user's interactions to a system " "for retrieval later."), description="The description of this function's use for tool calling agents.") - memory: MemoryRef = Field(default=MemoryRef("saas_memory"), - description=("Instance name of the memory client instance from the workflow " - "configuration object.")) @register_function(config_type=AddToolConfig) @@ -48,14 +46,13 @@ async def add_memory_tool(config: AddToolConfig, builder: Builder): # First, retrieve the memory client memory_editor = await builder.get_memory_client(config.memory) - async def _arun(item: MemoryItem) -> str: + async def _arun(item: AddMemoryInput) -> str: """ Asynchronous execution of addition of memories. Args: - item (MemoryItem): The memory item to add. Must include: + item (AddMemoryInput): The memory item to add. May include: - conversation: List of dicts with "role" and "content" keys - - user_id: String identifier for the user - metadata: Dict of metadata (can be empty) - tags: Optional list of tags - memory: Optional memory string @@ -64,16 +61,18 @@ async def _arun(item: MemoryItem) -> str: if available, otherwise an error will be raised. """ try: + memory_item = MemoryItem(**item.model_dump(), user_id=await resolve_memory_user_id(config)) + # If conversation is not provided but memory is, create a conversation - if not item.conversation and item.memory: - item.conversation = [{"role": "user", "content": item.memory}] - elif not item.conversation: + if not memory_item.conversation and memory_item.memory: + memory_item.conversation = [{"role": "user", "content": memory_item.memory}] + elif not memory_item.conversation: raise ToolException("Either conversation or memory must be provided") - await memory_editor.add_items([item]) + await memory_editor.add_items([memory_item]) return "Memory added successfully. You can continue. Please respond to the user." except Exception as e: raise ToolException(f"Error adding memory: {e}") from e - yield FunctionInfo.from_fn(_arun, description=config.description) + yield FunctionInfo.from_fn(_arun, description=config.description, input_schema=AddMemoryInput) diff --git a/packages/nvidia_nat_core/src/nat/tool/memory_tools/common.py b/packages/nvidia_nat_core/src/nat/tool/memory_tools/common.py new file mode 100644 index 0000000000..27cc05fbb6 --- /dev/null +++ b/packages/nvidia_nat_core/src/nat/tool/memory_tools/common.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import typing +from collections.abc import Awaitable +from collections.abc import Callable + +from pydantic import BaseModel +from pydantic import ConfigDict +from pydantic import Field +from pydantic import ImportString +from pydantic import field_validator +from pydantic import model_validator + +from nat.data_models.component_ref import MemoryRef +from nat.data_models.function import FunctionBaseConfig + +UserIdResolver = Callable[[], str | Awaitable[str]] + + +class MemoryToolConfigBase(FunctionBaseConfig): + """Shared configuration for memory tools.""" + + memory: MemoryRef = Field( + default=MemoryRef("saas_memory"), + description="Instance name of the memory client from the workflow configuration.", + ) + user_id: str | None = Field( + default=None, + description=( + "Optional fixed user identity for all memory operations. Configure either user_id or user_id_resolver. " + "This value is never exposed to the LLM."), + ) + user_id_resolver: ImportString[UserIdResolver] | None = Field( + default=None, + description=( + "Import path to a trusted zero-argument callable that returns the authenticated user's ID. The callable " + "is invoked for every memory operation and may be synchronous or asynchronous. Configure either " + "user_id or user_id_resolver. This value is never exposed to the LLM."), + ) + + @field_validator("user_id") + @classmethod + def validate_user_id(cls, value: str | None) -> str | None: + if value is None: + return None + + value = value.strip() + if not value: + raise ValueError("user_id must not be empty") + + return value + + @model_validator(mode="after") + def validate_user_id_source(self) -> typing.Self: + if self.user_id is not None and self.user_id_resolver is not None: + raise ValueError("Configure only one of user_id or user_id_resolver") + + return self + + +async def resolve_memory_user_id(config: MemoryToolConfigBase) -> str: + """Resolve a memory identity from trusted configuration or application code.""" + if config.user_id is not None: + return config.user_id + + if config.user_id_resolver is None: + raise ValueError("No user identity is available. Configure user_id or user_id_resolver.") + + resolved_user_id = config.user_id_resolver() + if inspect.isawaitable(resolved_user_id): + resolved_user_id = await resolved_user_id + + if not isinstance(resolved_user_id, str) or not resolved_user_id.strip(): + raise ValueError("The configured user_id_resolver must return a non-empty string.") + + return resolved_user_id.strip() + + +class AddMemoryInput(BaseModel): + """LLM-controlled input for adding a memory.""" + + model_config = ConfigDict(extra="forbid") + + conversation: list[dict[str, str]] | None = Field( + default=None, + description=( + "List of conversation messages. Each message must have a role key (user or assistant) and a content key."), + ) + tags: list[str] = Field(default_factory=list, description="List of tags applied to the item.") + metadata: dict[str, typing.Any] = Field(default_factory=dict, description="Metadata about the memory item.") + memory: str | None = Field(default=None, description="A memory to store.") + + +class GetMemoryInput(BaseModel): + """LLM-controlled input for retrieving memories.""" + + model_config = ConfigDict(extra="forbid") + + query: str = Field(description="Search query for which to retrieve memory.") + top_k: int = Field(description="Maximum number of memories to return.", gt=0) + + +class DeleteMemoryInput(BaseModel): + """LLM-controlled input for deleting a user's memories.""" + + model_config = ConfigDict(extra="forbid") diff --git a/packages/nvidia_nat_core/src/nat/tool/memory_tools/delete_memory_tool.py b/packages/nvidia_nat_core/src/nat/tool/memory_tools/delete_memory_tool.py index aeabbd53fe..7c6ecae607 100644 --- a/packages/nvidia_nat_core/src/nat/tool/memory_tools/delete_memory_tool.py +++ b/packages/nvidia_nat_core/src/nat/tool/memory_tools/delete_memory_tool.py @@ -20,21 +20,19 @@ from nat.builder.builder import Builder from nat.builder.function_info import FunctionInfo from nat.cli.register_workflow import register_function -from nat.data_models.component_ref import MemoryRef -from nat.data_models.function import FunctionBaseConfig -from nat.memory.models import DeleteMemoryInput + +from .common import DeleteMemoryInput +from .common import MemoryToolConfigBase +from .common import resolve_memory_user_id logger = logging.getLogger(__name__) -class DeleteToolConfig(FunctionBaseConfig, name="delete_memory"): +class DeleteToolConfig(MemoryToolConfigBase, name="delete_memory"): """Function to delete memory from a hosted memory platform.""" description: str = Field(default="Tool to delete a memory from a hosted memory platform.", description="The description of this function's use for tool calling agents.") - memory: MemoryRef = Field(default=MemoryRef("saas_memory"), - description=("Instance name of the memory client instance from the workflow " - "configuration object.")) @register_function(config_type=DeleteToolConfig) @@ -48,14 +46,15 @@ async def delete_memory_tool(config: DeleteToolConfig, builder: Builder): # First, retrieve the memory client memory_editor = await builder.get_memory_client(config.memory) - async def _arun(user_id: str) -> str: + async def _arun(delete_input: DeleteMemoryInput) -> str: """ Asynchronous execution of deletion of memories. """ try: + del delete_input - await memory_editor.remove_items(user_id=user_id, ) + await memory_editor.remove_items(user_id=await resolve_memory_user_id(config)) return "Memories deleted!" diff --git a/packages/nvidia_nat_core/src/nat/tool/memory_tools/get_memory_tool.py b/packages/nvidia_nat_core/src/nat/tool/memory_tools/get_memory_tool.py index 7de09c48d7..4f8d284b4c 100644 --- a/packages/nvidia_nat_core/src/nat/tool/memory_tools/get_memory_tool.py +++ b/packages/nvidia_nat_core/src/nat/tool/memory_tools/get_memory_tool.py @@ -20,22 +20,20 @@ from nat.builder.builder import Builder from nat.builder.function_info import FunctionInfo from nat.cli.register_workflow import register_function -from nat.data_models.component_ref import MemoryRef -from nat.data_models.function import FunctionBaseConfig -from nat.memory.models import SearchMemoryInput + +from .common import GetMemoryInput +from .common import MemoryToolConfigBase +from .common import resolve_memory_user_id logger = logging.getLogger(__name__) -class GetToolConfig(FunctionBaseConfig, name="get_memory"): +class GetToolConfig(MemoryToolConfigBase, name="get_memory"): """Function to get memory to a hosted memory platform.""" description: str = Field(default=("Tool to retrieve a memory about a user's " "interactions to help answer questions in a personalized way."), description="The description of this function's use for tool calling agents.") - memory: MemoryRef = Field(default=MemoryRef("saas_memory"), - description=("Instance name of the memory client instance from the workflow " - "configuration object.")) @register_function(config_type=GetToolConfig) @@ -51,7 +49,7 @@ async def get_memory_tool(config: GetToolConfig, builder: Builder): # First, retrieve the memory client memory_editor = await builder.get_memory_client(config.memory) - async def _arun(search_input: SearchMemoryInput) -> str: + async def _arun(search_input: GetMemoryInput) -> str: """ Asynchronous execution of collection of memories. """ @@ -59,7 +57,7 @@ async def _arun(search_input: SearchMemoryInput) -> str: memories = await memory_editor.search( query=search_input.query, top_k=search_input.top_k, - user_id=search_input.user_id, + user_id=await resolve_memory_user_id(config), ) memory_str = f"Memories as a JSON: \n{json.dumps([mem.model_dump(mode='json') for mem in memories])}" diff --git a/packages/nvidia_nat_core/tests/nat/tools/test_memory_tools.py b/packages/nvidia_nat_core/tests/nat/tools/test_memory_tools.py new file mode 100644 index 0000000000..a201c44c70 --- /dev/null +++ b/packages/nvidia_nat_core/tests/nat/tools/test_memory_tools.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import AsyncExitStack + +import pytest +from pydantic import ValidationError + +from nat.tool.memory_tools.common import AddMemoryInput +from nat.tool.memory_tools.common import DeleteMemoryInput +from nat.tool.memory_tools.common import GetMemoryInput +from nat.tool.memory_tools.common import resolve_memory_user_id + + +class _MemoryEditor: + + def __init__(self): + self.added_user_ids: list[str] = [] + self.searched_user_ids: list[str] = [] + self.deleted_user_ids: list[str] = [] + + async def add_items(self, items): + self.added_user_ids.extend(item.user_id for item in items) + + async def search(self, *, query, top_k, user_id): + del query, top_k + self.searched_user_ids.append(user_id) + return [] + + async def remove_items(self, *, user_id): + self.deleted_user_ids.append(user_id) + + +class _Builder: + + def __init__(self, memory_editor): + self.memory_editor = memory_editor + + async def get_memory_client(self, memory): + del memory + return self.memory_editor + + +@pytest.mark.parametrize( + "input_type, input_value", + [ + (AddMemoryInput, { + "memory": "hello", "user_id": "attacker" + }), + (GetMemoryInput, { + "query": "hello", "top_k": 1, "user_id": "attacker" + }), + (DeleteMemoryInput, { + "user_id": "attacker" + }), + ], +) +def test_memory_tool_inputs_reject_llm_supplied_user_id(input_type, input_value): + with pytest.raises(ValidationError): + input_type.model_validate(input_value) + + +@pytest.mark.parametrize("input_type", [AddMemoryInput, GetMemoryInput, DeleteMemoryInput]) +def test_memory_tool_schemas_do_not_publish_user_id(input_type): + assert "user_id" not in input_type.model_json_schema().get("properties", {}) + + +@pytest.mark.asyncio +async def test_resolve_memory_user_id_uses_fixed_identity(): + from nat.tool.memory_tools.add_memory_tool import AddToolConfig + + assert await resolve_memory_user_id(AddToolConfig(user_id=" fixed-user ")) == "fixed-user" + + +@pytest.mark.asyncio +async def test_resolve_memory_user_id_uses_synchronous_resolver(): + from nat.tool.memory_tools.add_memory_tool import AddToolConfig + + config = AddToolConfig(user_id_resolver=lambda: " resolved-user ") + assert await resolve_memory_user_id(config) == "resolved-user" + + +@pytest.mark.asyncio +async def test_resolve_memory_user_id_uses_asynchronous_resolver(): + from nat.tool.memory_tools.add_memory_tool import AddToolConfig + + async def resolver(): + return "resolved-user" + + assert await resolve_memory_user_id(AddToolConfig(user_id_resolver=resolver)) == "resolved-user" + + +@pytest.mark.parametrize("resolved_value", [None, "", " ", 123]) +@pytest.mark.asyncio +async def test_resolve_memory_user_id_rejects_invalid_resolver_result(resolved_value): + from nat.tool.memory_tools.add_memory_tool import AddToolConfig + + config = AddToolConfig(user_id_resolver=lambda: resolved_value) + with pytest.raises(ValueError, match="non-empty string"): + await resolve_memory_user_id(config) + + +@pytest.mark.asyncio +async def test_resolve_memory_user_id_requires_an_explicit_source(): + from nat.tool.memory_tools.add_memory_tool import AddToolConfig + + with pytest.raises(ValueError, match="user_id or user_id_resolver"): + await resolve_memory_user_id(AddToolConfig()) + + +def test_memory_tool_config_rejects_blank_fixed_user_id(): + from nat.tool.memory_tools.add_memory_tool import AddToolConfig + + with pytest.raises(ValidationError, match="user_id must not be empty"): + AddToolConfig(user_id=" ") + + +def test_memory_tool_config_rejects_multiple_identity_sources(): + from nat.tool.memory_tools.add_memory_tool import AddToolConfig + + with pytest.raises(ValidationError, match="only one"): + AddToolConfig(user_id="fixed-user", user_id_resolver=lambda: "resolved-user") + + +@pytest.mark.asyncio +async def test_memory_tools_use_resolved_identity_for_every_operation(): + from nat.tool.memory_tools.add_memory_tool import AddToolConfig + from nat.tool.memory_tools.add_memory_tool import add_memory_tool + from nat.tool.memory_tools.delete_memory_tool import DeleteToolConfig + from nat.tool.memory_tools.delete_memory_tool import delete_memory_tool + from nat.tool.memory_tools.get_memory_tool import GetToolConfig + from nat.tool.memory_tools.get_memory_tool import get_memory_tool + + memory_editor = _MemoryEditor() + builder = _Builder(memory_editor) + resolver_calls = 0 + + def resolver(): + nonlocal resolver_calls + resolver_calls += 1 + return "authenticated-user" + + async with AsyncExitStack() as stack: + add_tool = await stack.enter_async_context(add_memory_tool(AddToolConfig(user_id_resolver=resolver), builder)) + get_tool = await stack.enter_async_context(get_memory_tool(GetToolConfig(user_id_resolver=resolver), builder)) + delete_tool = await stack.enter_async_context( + delete_memory_tool(DeleteToolConfig(user_id_resolver=resolver), builder)) + + await add_tool.single_fn(AddMemoryInput(memory="strawberry")) + await get_tool.single_fn(GetMemoryInput(query="favorite flavor", top_k=5)) + await delete_tool.single_fn(DeleteMemoryInput()) + + assert resolver_calls == 3 + assert memory_editor.added_user_ids == ["authenticated-user"] + assert memory_editor.searched_user_ids == ["authenticated-user"] + assert memory_editor.deleted_user_ids == ["authenticated-user"] diff --git a/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/token_verifier.py b/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/token_verifier.py index 3b9e8c2c9d..3516bd8511 100644 --- a/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/token_verifier.py +++ b/packages/nvidia_nat_fastmcp/src/nat/plugins/fastmcp/server/token_verifier.py @@ -24,6 +24,7 @@ class NATFastMCPTokenVerifier(TokenVerifier): """FastMCP token verifier that delegates validation to BearerTokenValidator.""" def __init__(self, config: OAuth2ResourceServerConfig, *, base_url: str): + """Initialize the verifier with OAuth2 configuration and a public base URL.""" super().__init__(base_url=base_url, required_scopes=config.scopes or []) self._bearer_token_validator = BearerTokenValidator( issuer=config.issuer_url, diff --git a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py index 3bfdfb5c5c..62562b59a5 100644 --- a/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py +++ b/packages/nvidia_nat_langchain/src/nat/plugins/langchain/agent/auto_memory_wrapper/agent.py @@ -72,9 +72,11 @@ def _get_user_id_from_context(self) -> str: Extract user_id from runtime context. Priority order: + 1. Context.user_id - For authenticated sessions (set via SessionManager.session()) 2. user_manager.get_id() - Legacy/custom context compatibility - 3. X-User-ID HTTP header - For testing/simple auth without middleware + 3. X-User-ID HTTP header - Illustrative/testing only; assumes a trusted upstream proxy + has authenticated the request and injected the header 4. "default_user" - Fallback for development/testing without authentication Returns: @@ -100,7 +102,8 @@ def _get_user_id_from_context(self) -> str: except Exception as e: logger.debug(f"Failed to get user_id from user_manager: {e}") - # Priority 3: Extract from X-User-ID HTTP header (temporary workaround for testing) + # Priority 3: Extract an identity header injected by a trusted upstream proxy. This is illustrative/testing + # support only; an application must not accept a client-supplied X-User-ID header as authentication. metadata = getattr(self._context, "metadata", None) headers = getattr(metadata, "headers", None) if metadata else None if headers: