diff --git a/packages/nvidia_nat_selfmemory/LICENSE.md b/packages/nvidia_nat_selfmemory/LICENSE.md new file mode 100644 index 0000000000..f49a4e16e6 --- /dev/null +++ b/packages/nvidia_nat_selfmemory/LICENSE.md @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. \ No newline at end of file diff --git a/packages/nvidia_nat_selfmemory/pyproject.toml b/packages/nvidia_nat_selfmemory/pyproject.toml new file mode 100644 index 0000000000..2ab67f491e --- /dev/null +++ b/packages/nvidia_nat_selfmemory/pyproject.toml @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools>=64", "setuptools-scm>=8", "setuptools_dynamic_dependencies>=1.0.0"] + + +[tool.setuptools.packages.find] +where = ["src"] +include = ["nat.*"] + + +[tool.setuptools_scm] +git_describe_command = "git describe --long --first-parent" +root = "../.." + + +[project] +name = "nvidia-nat-selfmemory" +dynamic = ["version", "dependencies", "optional-dependencies"] +requires-python = ">=3.11,<3.14" +description = "Subpackage for SelfMemory integration in NeMo Agent Toolkit" +readme = "src/nat/meta/pypi.md" +keywords = ["ai", "agents", "memory", "selfmemory"] +license = { text = "Apache-2.0" } +authors = [{ name = "NVIDIA Corporation" }] +maintainers = [{ name = "NVIDIA Corporation" }] +classifiers = [ + "Programming Language :: Python", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", +] + +[project.urls] +documentation = "https://docs.nvidia.com/nemo/agent-toolkit/latest/" +source = "https://github.com/NVIDIA/NeMo-Agent-Toolkit" + +[tool.setuptools_dynamic_dependencies] +dependencies = [ + # Keep package version constraints as open as possible to avoid conflicts with other packages. Always define a minimum + # version when adding a new package. If unsure, default to using `~=` instead of `==`. Does not apply to nvidia-nat packages. + # Keep sorted!!! + "nvidia-nat-core == {version}", + "selfmemory>=0.9.4,<2.0.0", +] + +[tool.setuptools_dynamic_dependencies.optional-dependencies] +test = [ + "nvidia-nat-test == {version}", +] + +[tool.uv] +build-constraint-dependencies = ["setuptools>=64", "setuptools-scm>=8", "setuptools_dynamic_dependencies>=1.0.0"] +managed = true +config-settings = { editable_mode = "compat" } + +[tool.uv.sources] +nvidia-nat-core = { path = "../nvidia_nat_core", editable = true } +nvidia-nat-test = { path = "../nvidia_nat_test", editable = true } + +[project.entry-points.'nat.components'] +nat_selfmemory = "nat.plugins.selfmemory.register" diff --git a/packages/nvidia_nat_selfmemory/src/nat/meta/pypi.md b/packages/nvidia_nat_selfmemory/src/nat/meta/pypi.md new file mode 100644 index 0000000000..c5900ee62e --- /dev/null +++ b/packages/nvidia_nat_selfmemory/src/nat/meta/pypi.md @@ -0,0 +1,47 @@ + + +![NVIDIA NeMo Agent Toolkit](https://media.githubusercontent.com/media/NVIDIA/NeMo-Agent-Toolkit/refs/heads/main/docs/source/_static/banner.png "NeMo Agent Toolkit banner image") + +# nvidia-nat-selfmemory + +SelfMemory memory provider plugin for the NVIDIA NeMo Agent Toolkit. + +This package provides a `MemoryEditor` implementation that uses [SelfMemory](https://github.com/selfmemory/selfmemory) as the memory backend, enabling AI agents built with NeMo Agent Toolkit to store and retrieve long-term memories through SelfMemory's multi-tenant, vector-based memory system. + +## Features + +- 29+ vector store backends (Qdrant, ChromaDB, Pinecone, Milvus, etc.) +- 15+ embedding providers (OpenAI, Ollama, HuggingFace, etc.) +- Multi-tenant user isolation +- Optional LLM-based intelligent fact extraction +- Built-in encryption support + +## Usage + +```yaml +memory: + user_store: + _type: selfmemory + vector_store_provider: qdrant + vector_store_config: + host: localhost + port: 6333 + embedding_provider: openai + embedding_config: + model: text-embedding-3-small +``` diff --git a/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/__init__.py b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/__init__.py new file mode 100644 index 0000000000..3bcc1c39bb --- /dev/null +++ b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. diff --git a/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/memory.py b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/memory.py new file mode 100644 index 0000000000..2d6e7e876a --- /dev/null +++ b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/memory.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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 nat.builder.builder import Builder +from nat.cli.register_workflow import register_memory +from nat.data_models.memory import MemoryBaseConfig + + +class SelfMemoryProviderConfig(MemoryBaseConfig, name="selfmemory"): + vector_store_provider: str = "qdrant" + vector_store_config: dict = {} + embedding_provider: str = "openai" + embedding_config: dict = {} + llm_provider: str | None = None + llm_config: dict = {} + encryption_key: str | None = None + + +@register_memory(config_type=SelfMemoryProviderConfig) +async def selfmemory_provider(config: SelfMemoryProviderConfig, builder: Builder): + import os + + from selfmemory import SelfMemory + + from nat.plugins.selfmemory.selfmemory_editor import SelfMemoryEditor + + config_dict = { + "vector_store": { + "provider": config.vector_store_provider, + "config": config.vector_store_config, + }, + "embedding": { + "provider": config.embedding_provider, + "config": config.embedding_config, + }, + } + + if config.llm_provider: + config_dict["llm"] = { + "provider": config.llm_provider, + "config": config.llm_config, + } + + encryption_key = config.encryption_key or os.environ.get("MASTER_ENCRYPTION_KEY") + + if encryption_key: + os.environ["MASTER_ENCRYPTION_KEY"] = encryption_key + + memory = SelfMemory(config=config_dict) + + try: + yield SelfMemoryEditor(memory) + finally: + memory.close() diff --git a/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/register.py b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/register.py new file mode 100644 index 0000000000..58e1f1e710 --- /dev/null +++ b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/register.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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. + +# flake8: noqa +# isort:skip_file + +# Import any providers which need to be automatically registered here + +from . import memory diff --git a/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/selfmemory_editor.py b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/selfmemory_editor.py new file mode 100644 index 0000000000..e125f9e2ce --- /dev/null +++ b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/selfmemory_editor.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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 asyncio + +from selfmemory import SelfMemory + +from nat.memory.interfaces import MemoryEditor +from nat.memory.models import MemoryItem + +from .translator import memory_item_to_add_kwargs, search_results_to_memory_items + + +class SelfMemoryEditor(MemoryEditor): + """ + Wrapper class that implements NAT MemoryEditor for SelfMemory. + + Bridges SelfMemory's synchronous API to NeMo Agent Toolkit's + asynchronous MemoryEditor interface using asyncio.to_thread(). + """ + + def __init__(self, backend: SelfMemory): + self._backend = backend + + async def add_items(self, items: list[MemoryItem]) -> None: + """Insert multiple MemoryItems into SelfMemory.""" + coroutines = [ + asyncio.to_thread(self._backend.add, **memory_item_to_add_kwargs(item)) + for item in items + ] + await asyncio.gather(*coroutines) + + async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]: + """Retrieve items relevant to the given query.""" + user_id = kwargs.pop("user_id") + + result = await asyncio.to_thread( + self._backend.search, query, user_id=user_id, limit=top_k + ) + + return search_results_to_memory_items(result, user_id) + + async def remove_items(self, **kwargs) -> None: + """Remove items by memory_id or user_id.""" + if "memory_id" in kwargs: + memory_id = kwargs.pop("memory_id") + await asyncio.to_thread(self._backend.delete, memory_id) + + elif "user_id" in kwargs: + user_id = kwargs.pop("user_id") + await asyncio.to_thread(self._backend.delete_all, user_id=user_id) diff --git a/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/translator.py b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/translator.py new file mode 100644 index 0000000000..7451e8238e --- /dev/null +++ b/packages/nvidia_nat_selfmemory/src/nat/plugins/selfmemory/translator.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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 typing import Any + +from nat.memory.models import MemoryItem + + +def memory_item_to_add_kwargs(item: MemoryItem) -> dict[str, Any]: + """Convert a NeMo MemoryItem to SelfMemory add() keyword arguments.""" + messages = item.conversation if item.conversation else item.memory or "" + tags = ",".join(item.tags) if item.tags else None + + metadata = dict(item.metadata) if item.metadata else {} + people_mentioned = metadata.pop("people_mentioned", None) + topic_category = metadata.pop("topic_category", None) + project_id = metadata.pop("project_id", None) + organization_id = metadata.pop("organization_id", None) + + kwargs = { + "messages": messages, + "user_id": item.user_id or "default", + "tags": tags, + "people_mentioned": people_mentioned, + "topic_category": topic_category, + "project_id": project_id, + "organization_id": organization_id, + "metadata": metadata if metadata else None, + } + + return {k: v for k, v in kwargs.items() if v is not None} + + +def search_results_to_memory_items( + results: dict[str, Any], user_id: str +) -> list[MemoryItem]: + """Convert SelfMemory search results to a list of NeMo MemoryItems.""" + items = [] + + for result in results.get("results", []): + metadata = result.get("metadata", {}) + tags_str = metadata.get("tags", "") + tags = [t.strip() for t in tags_str.split(",") if t.strip()] if tags_str else [] + + items.append( + MemoryItem( + conversation=[], + user_id=user_id, + memory=result.get("content", ""), + tags=tags, + metadata=metadata, + ) + ) + + return items diff --git a/packages/nvidia_nat_selfmemory/tests/test_config.py b/packages/nvidia_nat_selfmemory/tests/test_config.py new file mode 100644 index 0000000000..1e31ab607e --- /dev/null +++ b/packages/nvidia_nat_selfmemory/tests/test_config.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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 nat.plugins.selfmemory.memory import SelfMemoryProviderConfig + + +class TestSelfMemoryProviderConfig: + + def test_default_values(self): + """Test config has correct defaults.""" + config = SelfMemoryProviderConfig() + + assert config.vector_store_provider == "qdrant" + assert config.vector_store_config == {} + assert config.embedding_provider == "openai" + assert config.embedding_config == {} + assert config.llm_provider is None + assert config.llm_config == {} + assert config.encryption_key is None + + def test_custom_values(self): + """Test config accepts custom values.""" + config = SelfMemoryProviderConfig( + vector_store_provider="chroma", + vector_store_config={"persist_directory": "/tmp/chroma"}, + embedding_provider="ollama", + embedding_config={"model": "nomic-embed-text"}, + llm_provider="openai", + llm_config={"model": "gpt-4o-mini"}, + encryption_key="my-secret-key", + ) + + assert config.vector_store_provider == "chroma" + assert config.vector_store_config["persist_directory"] == "/tmp/chroma" + assert config.embedding_provider == "ollama" + assert config.embedding_config["model"] == "nomic-embed-text" + assert config.llm_provider == "openai" + assert config.encryption_key == "my-secret-key" + + def test_name_attribute(self): + """Test the config registers with name 'selfmemory'.""" + config = SelfMemoryProviderConfig() + + assert config.__class__.__name__ == "SelfMemoryProviderConfig" diff --git a/packages/nvidia_nat_selfmemory/tests/test_selfmemory_editor.py b/packages/nvidia_nat_selfmemory/tests/test_selfmemory_editor.py new file mode 100644 index 0000000000..d679896cdb --- /dev/null +++ b/packages/nvidia_nat_selfmemory/tests/test_selfmemory_editor.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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 unittest.mock import MagicMock + +import pytest + +from nat.memory.models import MemoryItem +from nat.plugins.selfmemory.selfmemory_editor import SelfMemoryEditor + + +@pytest.fixture(name="mock_backend") +def mock_backend_fixture() -> MagicMock: + """Fixture to provide a mocked SelfMemory instance.""" + backend = MagicMock() + backend.add.return_value = {"id": "mem_123", "success": True} + backend.search.return_value = {"results": []} + backend.delete.return_value = {"success": True} + backend.delete_all.return_value = {"success": True} + return backend + + +@pytest.fixture(name="editor") +def editor_fixture(mock_backend: MagicMock): + """Fixture to provide a SelfMemoryEditor with a mocked backend.""" + return SelfMemoryEditor(backend=mock_backend) + + +@pytest.fixture(name="sample_memory_item") +def sample_memory_item_fixture(): + """Fixture to provide a sample MemoryItem.""" + conversation = [ + {"role": "user", "content": "I love Italian food, especially pizza."}, + {"role": "assistant", "content": "Noted! You love Italian food and pizza."}, + ] + + return MemoryItem( + conversation=conversation, + user_id="user123", + memory="Loves Italian food", + metadata={"topic_category": "preferences"}, + tags=["food", "preferences"], + ) + + +async def test_add_items_success(editor: SelfMemoryEditor, mock_backend: MagicMock, sample_memory_item: MemoryItem): + """Test adding multiple MemoryItem objects successfully.""" + items = [sample_memory_item, sample_memory_item] + await editor.add_items(items) + + assert mock_backend.add.call_count == len(items) + + +async def test_add_items_translates_tags(editor: SelfMemoryEditor, mock_backend: MagicMock, sample_memory_item: MemoryItem): + """Test that tags are converted from list to comma-separated string.""" + await editor.add_items([sample_memory_item]) + + call_kwargs = mock_backend.add.call_args[1] + assert call_kwargs["tags"] == "food,preferences" + + +async def test_add_items_translates_metadata(editor: SelfMemoryEditor, mock_backend: MagicMock, sample_memory_item: MemoryItem): + """Test that metadata fields are extracted to SelfMemory kwargs.""" + await editor.add_items([sample_memory_item]) + + call_kwargs = mock_backend.add.call_args[1] + assert call_kwargs["topic_category"] == "preferences" + assert call_kwargs["user_id"] == "user123" + + +async def test_add_items_empty_list(editor: SelfMemoryEditor, mock_backend: MagicMock): + """Test adding an empty list of MemoryItem objects.""" + await editor.add_items([]) + + mock_backend.add.assert_not_called() + + +async def test_search_success(editor: SelfMemoryEditor, mock_backend: MagicMock): + """Test searching with a valid query and user ID.""" + mock_backend.search.return_value = { + "results": [ + { + "id": "mem_1", + "content": "Loves Italian food", + "score": 0.95, + "metadata": { + "data": "Loves Italian food", + "user_id": "user123", + "tags": "food,preferences", + }, + } + ] + } + + result = await editor.search(query="food preferences", user_id="user123", top_k=5) + + assert len(result) == 1 + assert result[0].memory == "Loves Italian food" + assert result[0].user_id == "user123" + assert result[0].tags == ["food", "preferences"] + + mock_backend.search.assert_called_once_with("food preferences", user_id="user123", limit=5) + + +async def test_search_empty_results(editor: SelfMemoryEditor, mock_backend: MagicMock): + """Test searching with no results.""" + mock_backend.search.return_value = {"results": []} + + result = await editor.search(query="nonexistent", user_id="user123") + + assert len(result) == 0 + + +async def test_search_missing_user_id(editor: SelfMemoryEditor): + """Test searching without providing a user ID.""" + with pytest.raises(KeyError, match="user_id"): + await editor.search(query="test query") + + +async def test_remove_items_by_memory_id(editor: SelfMemoryEditor, mock_backend: MagicMock): + """Test removing items by memory ID.""" + await editor.remove_items(memory_id="mem_123") + + mock_backend.delete.assert_called_once_with("mem_123") + + +async def test_remove_items_by_user_id(editor: SelfMemoryEditor, mock_backend: MagicMock): + """Test removing all items for a specific user ID.""" + await editor.remove_items(user_id="user123") + + mock_backend.delete_all.assert_called_once_with(user_id="user123") + + +async def test_remove_items_missing_arguments(editor: SelfMemoryEditor): + """Test removing items with missing required arguments.""" + result = await editor.remove_items() + + assert result is None diff --git a/packages/nvidia_nat_selfmemory/tests/test_translator.py b/packages/nvidia_nat_selfmemory/tests/test_translator.py new file mode 100644 index 0000000000..46a54eb4e4 --- /dev/null +++ b/packages/nvidia_nat_selfmemory/tests/test_translator.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-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 nat.memory.models import MemoryItem +from nat.plugins.selfmemory.translator import ( + memory_item_to_add_kwargs, + search_results_to_memory_items, +) + + +class TestMemoryItemToAddKwargs: + + def test_basic_conversion(self): + """Test basic MemoryItem to add kwargs conversion.""" + item = MemoryItem( + conversation=[{"role": "user", "content": "I love pizza"}], + user_id="alice", + memory="Loves pizza", + tags=["food"], + metadata={}, + ) + + kwargs = memory_item_to_add_kwargs(item) + + assert kwargs["messages"] == [{"role": "user", "content": "I love pizza"}] + assert kwargs["user_id"] == "alice" + assert kwargs["tags"] == "food" + + def test_multiple_tags(self): + """Test that multiple tags are joined with commas.""" + item = MemoryItem( + conversation=[], + user_id="bob", + memory="Note", + tags=["work", "meeting", "important"], + metadata={}, + ) + + kwargs = memory_item_to_add_kwargs(item) + + assert kwargs["tags"] == "work,meeting,important" + + def test_empty_tags(self): + """Test that empty tags are omitted.""" + item = MemoryItem( + conversation=[], + user_id="charlie", + memory="Note", + tags=[], + metadata={}, + ) + + kwargs = memory_item_to_add_kwargs(item) + + assert "tags" not in kwargs + + def test_extracts_metadata_fields(self): + """Test that known metadata fields are extracted to top-level kwargs.""" + item = MemoryItem( + conversation=[], + user_id="alice", + memory="Meeting note", + tags=[], + metadata={ + "people_mentioned": "Sarah,Mike", + "topic_category": "work", + "project_id": "proj_1", + "organization_id": "org_1", + "custom_field": "value", + }, + ) + + kwargs = memory_item_to_add_kwargs(item) + + assert kwargs["people_mentioned"] == "Sarah,Mike" + assert kwargs["topic_category"] == "work" + assert kwargs["project_id"] == "proj_1" + assert kwargs["organization_id"] == "org_1" + assert kwargs["metadata"] == {"custom_field": "value"} + + def test_fallback_to_memory_string(self): + """Test that memory string is used when conversation is empty.""" + item = MemoryItem( + conversation=[], + user_id="alice", + memory="Direct memory text", + tags=[], + metadata={}, + ) + + kwargs = memory_item_to_add_kwargs(item) + + assert kwargs["messages"] == "Direct memory text" + + def test_default_user_id(self): + """Test that empty user_id defaults to 'default'.""" + item = MemoryItem( + conversation=[], + user_id="", + memory="Note", + tags=[], + metadata={}, + ) + + kwargs = memory_item_to_add_kwargs(item) + + assert kwargs["user_id"] == "default" + + +class TestSearchResultsToMemoryItems: + + def test_basic_conversion(self): + """Test basic search result to MemoryItem conversion.""" + results = { + "results": [ + { + "id": "mem_1", + "content": "Loves pizza", + "score": 0.95, + "metadata": { + "data": "Loves pizza", + "user_id": "alice", + "tags": "food,preferences", + }, + } + ] + } + + items = search_results_to_memory_items(results, "alice") + + assert len(items) == 1 + assert items[0].memory == "Loves pizza" + assert items[0].user_id == "alice" + assert items[0].tags == ["food", "preferences"] + + def test_empty_results(self): + """Test empty search results.""" + results = {"results": []} + + items = search_results_to_memory_items(results, "alice") + + assert len(items) == 0 + + def test_empty_tags(self): + """Test result with empty tags string.""" + results = { + "results": [ + { + "id": "mem_1", + "content": "Note", + "metadata": {"tags": ""}, + } + ] + } + + items = search_results_to_memory_items(results, "alice") + + assert items[0].tags == [] + + def test_missing_metadata(self): + """Test result with missing metadata.""" + results = { + "results": [ + { + "id": "mem_1", + "content": "Note", + } + ] + } + + items = search_results_to_memory_items(results, "alice") + + assert items[0].memory == "Note" + assert items[0].tags == [] + + def test_multiple_results(self): + """Test multiple search results.""" + results = { + "results": [ + {"id": "mem_1", "content": "First", "metadata": {"tags": "a"}}, + {"id": "mem_2", "content": "Second", "metadata": {"tags": "b"}}, + {"id": "mem_3", "content": "Third", "metadata": {"tags": "c"}}, + ] + } + + items = search_results_to_memory_items(results, "alice") + + assert len(items) == 3 + assert items[0].memory == "First" + assert items[2].memory == "Third" diff --git a/pyproject.toml b/pyproject.toml index b43e9dc10b..90cd6e6448 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,7 @@ openpipe-art = ["nvidia-nat-openpipe-art == {version}"] opentelemetry = ["nvidia-nat-opentelemetry == {version}"] phoenix = ["nvidia-nat-phoenix == {version}"] profiler = ["nvidia-nat-profiler == {version}"] +selfmemory = ["nvidia-nat-selfmemory == {version}"] rag = ["nvidia-nat-rag == {version}"] ragas = ["nvidia-nat-ragas == {version}"] ragaai = ["nvidia-nat-ragaai == {version}"]