Add MemMachine memory integration for NeMo Agent toolkit #1460
+3,789
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
- Add MemMachine memory plugin integration - Rename package from memmachine to nvidia_nat_memmachine to match NAT naming convention - Support both hosted API and self-hosted server modes - Make memmachine-server an optional dependency - Add httpx as required dependency for hosted API support - Update all references in notebook and documentation - Update tests to reflect package rename and memory consolidation behavior - Fix capitalization: 'NeMo Agent toolkit' (lowercase 't') in all text - Fix pyproject.toml to use workspace source for nvidia-nat dependency
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
WalkthroughAdds a new nvidia-nat-memmachine subpackage: packaging and metadata, README and example notebook, a MemMachineEditor and memmachine_memory_client factory integrating with the MemMachine SDK, extensive unit and integration tests, and adds the package to the meta-package dependencies. Changes
Sequence Diagram(s)sequenceDiagram
participant Agent as Agent/Application
participant Builder as WorkflowBuilder
participant Factory as memmachine_memory_client
participant Editor as MemMachineEditor
participant SDK as MemMachine SDK
participant Server as MemMachine Server
Agent->>Builder: register memory(config)
Builder->>Factory: instantiate with config
Factory->>SDK: create MemMachineClient(base_url, timeout, retries)
Factory->>SDK: get_or_create_project(org_id, project_id) (optional)
SDK->>Server: API request (project)
Server-->>SDK: project response
Factory->>Editor: initialize with client/project
Factory-->>Builder: yield memory_editor
Agent->>Editor: add_items(conversation/direct)
Editor->>SDK: memory.add (episodic and/or semantic)
SDK->>Server: store memories
Server-->>SDK: ack/indexed
SDK-->>Editor: add responses
Agent->>Editor: search(query, top_k)
Editor->>SDK: memory.search(query, limit=top_k)
SDK->>Server: execute query
Server-->>SDK: episodic + semantic results
SDK-->>Editor: results
Editor->>Agent: return list[MemoryItem]
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Tip 🧪 Unit Test Generation v2 is now available!We have significantly improved our unit test generation capabilities. To enable: Add this to your reviews:
finishing_touches:
unit_tests:
enabled: trueTry it out by using the Have feedback? Share your thoughts on our Discord thread! Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 16
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/pyproject.toml`:
- Around line 16-25: Update package dependency lists: in
packages/nvidia_nat_all/pyproject.toml add "nvidia-nat-memmachine" (no version
constraint) to the dependencies array, and in
packages/nvidia_nat_memmachine/pyproject.toml change the memmachine-server entry
from "memmachine-server~=0.2.2" to "memmachine-server~=0.2.1" so the declared
dependency matches the published version; ensure the dependencies remain sorted
per the file comments.
In
`@packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py`:
- Around line 1-9: Add the standard SPDX Apache-2.0 license header to the top of
memmachine_editor.py (above all imports); ensure the exact SPDX short-form line
"SPDX-License-Identifier: Apache-2.0" is present and include any required
copyright owner/year line per project policy so the file containing imports like
asyncio, requests, MemoryType, MemoryEditor, and MemoryItem is properly
licensed.
- Around line 185-196: The add_memory() closure captures outer loop variables
(memory, memory_text, metadata, memory_types) by reference so all scheduled
tasks use the last values; fix by binding those values into the closure before
scheduling (e.g., define add_memory(memory=memory, memory_text=memory_text,
metadata=metadata, memory_types=memory_types): or use functools.partial to bind
arguments) so the memory.add(...) call inside add_memory uses the per-iteration
values when you call task = asyncio.to_thread(add_memory).
- Around line 299-304: The bare except in the episodic_by_conversation loop
around conv_episodes.sort hides real errors; replace it with a targeted
exception handler that catches likely errors (e.g., TypeError, AttributeError,
ValueError) and log the failure including conv_key and the exception before
continuing. Specifically, wrap the conv_episodes.sort(key=lambda e:
e.get("created_at") or e.get("timestamp") or "") call in an except (TypeError,
AttributeError, ValueError) as e block and emit a warning or error via the
module's logger (or logging.warning) referencing conv_key and e so issues aren't
silently swallowed.
In `@packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py`:
- Around line 85-88: The except block in the project-creation fallback (inside
memory.py where the try/except around client/project creation is located)
currently swallows all exceptions with "pass"; replace that with a
logger.exception(...) call to record the full stack trace and context (e.g.,
"Failed to create project, falling back to client usage") so debugging info is
preserved; keep the existing fallback behavior (do not re-raise) but ensure the
exception is logged via logger.exception in the except Exception as e: handler.
In `@packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py`:
- Line 1: Add the standard SPDX Apache-2.0 license header at the top of this
file (register.py) before any code — ensure the file begins with the SPDX
license identifier (e.g., a comment containing "SPDX-License-Identifier:
Apache-2.0") so the header appears above the existing import line "from . import
memory"; update only the file header area to include the required license
comment.
In `@packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md`:
- Around line 1-5: Add the standard SPDX Apache-2.0 license header as the very
first line of the markdown file (API_CALL_VERIFICATION.md) by inserting
"SPDX-License-Identifier: Apache-2.0" at the top of the file so the file now
begins with that license identifier.
In `@packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py`:
- Around line 1-24: Add the missing SPDX Apache-2.0 license header at the top of
the file (above the docstring or immediately after the shebang) to satisfy
licensing requirements, and mark the integration test with pytest by importing
pytest and decorating the test function (e.g., test_add_and_retrieve) with
`@pytest.mark.integration` so external-service tests are properly classified;
update imports to include "import pytest" and ensure the decorator is applied to
the test function in the script.
In `@packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py`:
- Around line 1-2: Update the SPDX header years in the file's top comment lines:
change the copyright year range in the SPDX-FileCopyrightText line (currently
"2024-2025") to include 2026 (e.g., "2024-2026") and ensure the
SPDX-License-Identifier line remains unchanged; edit the two header comment
lines at the top of the file to reflect the new year range.
- Line 366: The assignment to results from calling editor_with_spy.search is
unused and should be removed to avoid lint warnings; simply call await
editor_with_spy.search(...) without assigning its return, or if you prefer
preserving a placeholder, assign to _ (underscore) instead—update the call site
where results is assigned in the test (the await editor_with_spy.search
invocation) to drop the unused variable.
- Line 50: Update the signature of get_calls to use explicit PEP 604 union
typing for the optional parameter: replace the parameter declaration
method_name: str = None with method_name: str | None, and ensure any related
type hints or usages in the get_calls function body or its callers remain valid
with the new annotation.
In `@packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py`:
- Around line 173-184: The for-loop in test_memmachine_integration.py uses an
unused loop variable named "attempt" which triggers a lint warning; rename it to
"_attempt" in both occurrences (the loop that calls memory_client.search with
query="morning work allergy peanuts" and the similar loop around lines ~403-414)
so the intent is preserved and the linter is satisfied, keeping the loop logic
and the await asyncio.sleep(2) retry behavior unchanged.
- Around line 1-2: Update the SPDX header year range in the file's top comment
to include 2026: locate the SPDX lines at the top of the file (the
SPDX-FileCopyrightText and SPDX-License-Identifier header) and change the
copyright year range "2024-2025" to "2024-2026" so the notice accurately
reflects the file modification year.
- Around line 22-25: Update the module docstring in the
tests/test_memmachine_integration.py file to wrap code entities in backticks to
satisfy Vale and docstring rules: surround MemMachine server with `MemMachine`,
the environment variable name with `MEMMACHINE_BASE_URL`, and the test command
with backticks like `pytest tests/test_memmachine_integration.py -v`; ensure any
other inline code-like tokens in that docstring are also backticked.
In `@packages/nvidia_nat_memmachine/tests/test_memory.py`:
- Around line 1-2: Update the SPDX header year range in the file header comment
to include 2026; specifically modify the top two comment lines (the SPDX
copyright and SPDX-License-Identifier block) so the year range reads "2024-2026"
instead of "2024-2025" to reflect the file modification in 2026.
- Around line 155-159: The test test_memmachine_memory_client_config_validation
currently catches a broad Exception when constructing
MemMachineMemoryClientConfig; change the assertion to catch the specific
pydantic validation error by replacing pytest.raises(Exception) with
pytest.raises(ValidationError) and add the necessary import for ValidationError
from pydantic so the test asserts the concrete validation failure from
MemMachineMemoryClientConfig.
🧹 Nitpick comments (7)
packages/nvidia_nat_memmachine/src/nat/meta/pypi.md (1)
20-24: Capitalize "Toolkit" in the heading.Per coding guidelines, titles/headings should use "Toolkit" (capital T) while body text uses "toolkit" (lowercase t). Also ensure the file ends with a single newline.
📝 Suggested fix
-# NVIDIA NeMo Agent toolkit Subpackage +# NVIDIA NeMo Agent Toolkit Subpackage This is a subpackage for MemMachine memory integration in NeMo Agent toolkit. For more information about the NVIDIA NeMo Agent toolkit, please visit the [NeMo Agent toolkit GitHub Repo](https://github.com/NVIDIA/NeMo-Agent-Toolkit). +packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py (3)
41-42: Add return type annotation and consider documenting unusedbuilderparameter.Per coding guidelines, public APIs require type hints on return values. The function should declare its return type. The
builderparameter appears unused (flagged by static analysis) — if it's required by the@register_memorydecorator contract, consider adding a comment or using_prefix to indicate intentional non-use.📝 Suggested fix
+from collections.abc import AsyncGenerator +from typing import Any + `@register_memory`(config_type=MemMachineMemoryClientConfig) -async def memmachine_memory_client(config: MemMachineMemoryClientConfig, builder: Builder): +async def memmachine_memory_client( + config: MemMachineMemoryClientConfig, + builder: Builder, # Required by `@register_memory` contract, unused here +) -> AsyncGenerator[Any, None]:
34-38: Clarify the distinction betweenmax_retriesandnum_retriesfromRetryMixin.The class defines
max_retriesfor SDK-level retries, while inheritingnum_retriesfromRetryMixin(used at line 95 forpatch_with_retry). This dual retry configuration may confuse users. Consider adding a note in the docstring explaining when each is used.📝 Suggested docstring addition
""" Configuration for MemMachine memory client. Based on the MemMachine Python SDK as documented at: https://github.com/MemMachine/MemMachine/blob/main/docs/examples/python.mdx Note: This integration is for local/self-hosted MemMachine instances. LLM API keys (e.g., OpenAI) are configured in the MemMachine cfg.yml file, not in this client configuration. + + Retry behavior: + - `max_retries`: SDK-level retries for individual MemMachine API calls. + - `num_retries` (inherited from RetryMixin): Application-level retries with + status code filtering, applied to the memory editor wrapper. """
92-98: Redundantisinstancecheck —configalways inherits fromRetryMixin.Since
MemMachineMemoryClientConfiginherits fromRetryMixin, this check is alwaysTrue. Consider removing the check for clarity, or add a comment if this is intentional defensive coding for future flexibility.📝 Suggested simplification
- if isinstance(config, RetryMixin): - memory_editor = patch_with_retry( - memory_editor, - retries=config.num_retries, - retry_codes=config.retry_on_status_codes, - retry_on_messages=config.retry_on_errors - ) + memory_editor = patch_with_retry( + memory_editor, + retries=config.num_retries, + retry_codes=config.retry_on_status_codes, + retry_on_messages=config.retry_on_errors + )packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py (1)
471-490: Unused variabledelete_semanticand incomplete bulk deletion path.The variable
delete_semanticis assigned but never used. If bulk deletion should respect this flag in the future, consider either implementing partial support or removing the variable to avoid confusion.Proposed fix - remove unused variable
project_id = kwargs.pop("project_id", None) org_id = kwargs.pop("org_id", None) - delete_semantic = kwargs.pop("delete_semantic_memory", False) # Note: MemMachine SDK doesn't have a delete_all methodpackages/nvidia_nat_memmachine/memmachine_memory_example.ipynb (1)
386-403: Consider usinghttpxfor consistency with coding guidelines.The coding guidelines prefer
httpxoverrequests. Whilerequestsworks fine here for a simple health check, consider aligning with the codebase convention.packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py (1)
179-183: Uselogger.exception()when catching and logging exceptions without re-raising.Per coding guidelines, when catching and logging exceptions without re-raising, use
logger.exception()to capture full stack trace. The current approach usingtraceback.print_exc()works but doesn't follow the project's logging conventions.
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py
Outdated
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py
Outdated
Show resolved
Hide resolved
willkill07
requested changes
Jan 23, 2026
Member
willkill07
left a comment
willkill07
left a comment
There was a problem hiding this comment.
Thank you for your contribution!
This needs some work, and I'm going to start with the low-hanging fruit and a quick review before addressing the files that GitHub hid by default due to size.
Overall, all comments by coderabbit and additional ones from me need to be addressed.
Additionally, copyright headers for new files should include the current year rather than 2024-2025
I will admit I don't know how many commits behind develop you are, but you should assume a base as being near or close to top of tree.
This is my own opinion here, but I also don't like the idea of mandating a local server. Is there not an API/microservice that can be used instead?
Finally, if you have any questions don't hesitate to ask!
| # version when adding a new package. If unsure, default to using `~=` instead of `==`. Does not apply to nvidia-nat packages. | ||
| # Keep sorted!!! | ||
| "nvidia-nat~=1.4", | ||
| "memmachine-server~=0.2.2", |
Member
There was a problem hiding this comment.
Is there not just a client version? It seems weird to force a locally running instance
Author
There was a problem hiding this comment.
Hey @willkill07 , thanks for looking over this. We are working on a hosted REST API platform where users can go and just make an API key but it hasn't been rolled out yet. Once that is ready, we'll make a new PR to give user the option to use that.
Member
There was a problem hiding this comment.
What I mean is this PyPI package exists: https://pypi.org/project/memmachine-client/
For the sake of the library component, is installing the server package necessary?
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
Actionable comments posted: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/nvidia_nat_all/pyproject.toml (1)
23-33: Addnvidia-nat-memmachineto the workspace sources (and syncuv.lock).Since this dependency is a workspace package, it should be listed in
[tool.uv.sources]to ensure local resolution. Also make sureuv.lockis updated to reflect the new dependency.🧩 Proposed update
[tool.uv.sources] nvidia-nat = { workspace = true } nvidia-nat-agno = { workspace = true } nvidia-nat-autogen = { workspace = true } nvidia-nat-crewai = { workspace = true } nvidia-nat-langchain = { workspace = true } nvidia-nat-llama-index = { workspace = true } nvidia-nat-mcp = { workspace = true } nvidia-nat-mem0ai = { workspace = true } +nvidia-nat-memmachine = { workspace = true } nvidia-nat-mysql = { workspace = true }Based on learnings.
🤖 Fix all issues with AI agents
In
`@packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py`:
- Around line 324-331: In the try/except around conv_episodes.sort in
memmachine_editor (the block catching TypeError/AttributeError/ValueError for
conv_key), replace the logger.warning call with logger.exception so the full
stack trace is preserved when the sort fails; keep the same descriptive message
referencing conv_key and that we are "continuing without sorting" but call
logger.exception(...) instead of logger.warning(...).
- Around line 29-45: The docstring for the MemMachineEditor class contains
several metadata keys and code entities that are not formatted as inline code;
update the MemMachineEditor docstring so all code identifiers (e.g., MemoryItem,
and metadata keys `session_id`, `agent_id`, `group_id`, `project_id`, `org_id`,
and the class name `MemMachineClient` reference) are wrapped in backticks to
satisfy Vale/coding guidelines and ensure clarity; keep the explanatory text
unchanged and only modify the docstring formatting around these identifiers.
- Around line 179-190: The closure add_memory currently captures the outer
variable memory by reference, causing deferred asyncio.to_thread tasks to all
use the last memory instance; fix this by binding memory into the closure
signature as a default parameter (e.g., add a parameter like memory=memory) so
each add_memory captures its own memory instance before scheduling; update the
add_memory definition that calls memory.add(...) to include memory as a default
arg consistent with the identical pattern used in the other branch.
In `@packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py`:
- Around line 48-51: The function signature for memmachine_memory_client
currently defines an unused parameter named builder which ruff flags; silence it
by renaming builder to _builder (or alternatively append a per-argument lint
suppression like "# noqa: ARG001") in the memmachine_memory_client definition so
the parameter remains for the `@register_memory` contract but no lint error is
raised.
- Around line 55-62: The project incorrectly requires the server package while
code imports the client: update the dependency in pyproject.toml from
memmachine-server to memmachine-client, and update the ImportError text in
memory.py (the block that catches ImportError for "from memmachine import
MemMachineClient") to instruct installing memmachine-client (e.g., "pip install
memmachine-client") and adjust any package name references/links so they point
to the client package documentation; keep the rest of the error message
structure (include the caught exception e) intact.
In `@packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py`:
- Around line 199-202: Replace the logger.exception call in the except block
that captures "Exception as e" with logger.error(..., exc_info=True) so the
error message is logged with the traceback without duplicating stack traces on
re-raise; specifically, locate the except Exception as e block (where
logger.exception("Error during test execution") is called) and change it to log
the same message using logger.error("Error during test execution",
exc_info=True) before re-raising the exception.
In `@packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py`:
- Around line 68-85: The two pytest fixture functions lack return type
annotations; update test_config_fixture to declare a return type of
MemMachineMemoryClientConfig (i.e., def test_config_fixture(...) ->
MemMachineMemoryClientConfig:) and update test_user_id_fixture to declare a
return type of str (def test_user_id_fixture(...) -> str:); ensure the
MemMachineMemoryClientConfig symbol is imported or referenced correctly in the
test module if not already.
- Around line 88-156: The async test function
test_add_and_retrieve_conversation_memory is missing the pytest-asyncio marker;
add `@pytest.mark.asyncio` above the function (place it directly after
`@pytest.mark.slow`) so pytest-asyncio runs the async test correctly.
- Around line 49-65: Change the broad except Exception in the MemMachineClient
health-check block to catch the specific requests exception type raised by the
client (e.g., requests.exceptions.RequestException or requests.RequestException)
so real errors aren't swallowed; import the exception from requests at top if
needed and keep the same behavior of raising RuntimeError when fail_missing is
true or calling pytest.skip otherwise while still preserving the existing
ImportError handling for the memmachine import and the use of
MemMachineClient.health_check().
♻️ Duplicate comments (2)
packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py (1)
46-209: Avoid running live MemMachine calls in the default test suite.
test_add_and_retrieveis still a discovered test and will run without the integration marker, even though it requires a live MemMachine instance. Make it a helper (non-test) or mark it as integration to keep it out of default runs.✅ Suggested restructuring
-async def test_add_and_retrieve(): +async def _run_add_and_retrieve(): """Test adding memories and retrieving them.""" @@ -@pytest.mark.integration -async def test_add_and_retrieve_integration(): +@pytest.mark.integration +async def test_add_and_retrieve_integration(): """Integration test for adding and retrieving memories.""" - await test_add_and_retrieve() + await _run_add_and_retrieve() @@ -if __name__ == "__main__": - asyncio.run(test_add_and_retrieve()) +if __name__ == "__main__": + asyncio.run(_run_add_and_retrieve())As per coding guidelines.
packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py (1)
16-24: Backtick code-like tokens in the module/fixture docstrings.The docstrings still leave code entities and the localhost URL unformatted, which can trigger Vale. Please wrap them in backticks for consistency. As per coding guidelines, please wrap code entities in backticks.
✍️ Proposed fix
-Integration tests for MemMachine memory integration. +Integration tests for `MemMachine` memory integration. @@ -Set `MEMMACHINE_BASE_URL` environment variable to override default (http://localhost:8080). +Set `MEMMACHINE_BASE_URL` environment variable to override default (`http://localhost:8080`). @@ -Set MEMMACHINE_BASE_URL environment variable to override default (http://localhost:8080). +Set `MEMMACHINE_BASE_URL` environment variable to override default (`http://localhost:8080`).Also applies to: 39-44
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py
Show resolved
Hide resolved
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py
Show resolved
Hide resolved
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Charlie-Yi-2002 <charlie.yi@memverge.com>
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py`:
- Around line 205-208: The async integration test function
test_add_and_retrieve_integration is missing the `@pytest.mark.asyncio` decorator;
add `@pytest.mark.asyncio` above the async def test_add_and_retrieve_integration()
declaration so pytest runs it as an async test (keep the existing
`@pytest.mark.integration` marker and the call to await test_add_and_retrieve()).
In `@packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py`:
- Around line 1-2: Update the SPDX copyright header lines to include 2026 by
changing the year range in the two SPDX comment lines: adjust the copyright
comment line that currently reads "2024-2025" to "2024-2026" and keep the
SPDX-License-Identifier line as-is; ensure these changes are applied at the top
of the file where the SPDX header is defined.
🧹 Nitpick comments (2)
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py (1)
322-329: Remove redundant exception object fromlogger.exceptioncall.
logger.exception()automatically includes the exception information, so including{e}in the message is redundant.Proposed fix
except (TypeError, AttributeError, ValueError) as e: # Skip sorting if timestamps are missing or incompatible logger.exception( - f"Failed to sort episodes for conversation '{conv_key}': {e}. " + f"Failed to sort episodes for conversation '{conv_key}'. " "Continuing without sorting." )packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py (1)
116-118: Consider movingasyncioimport to module level.The
asynciomodule is imported inside multiple test functions. Moving it to the top of the file with other imports would be cleaner.Proposed fix
Add to imports at top of file:
import asyncioThen remove the inline imports in each test function.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Charlie-Yi-2002 <charlie.yi@memverge.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py`:
- Around line 492-496: The mock return for mock_memory_instance.search is
inconsistent with other tests: instead of assigning a plain dict set
mock_memory_instance.search.return_value to a Mock (or MagicMock) instance whose
.content attribute contains the dict used in other tests (e.g.,
{"episodic_memory": [], "semantic_memory": [], "episode_summary": []}), so that
MemMachineEditor.search (and any code accessing result.content) behaves
consistently; update the test to wrap the dict in a Mock with .content and
adjust any assertions if they expect the Mock shape.
🧹 Nitpick comments (1)
packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py (1)
178-202: Consider removing unused fixture parameters.The
mock_clientandmock_projectparameters are injected but not used in the test body. Pytest automatically instantiates the fixture dependency chain whenmemmachine_editor_with_clientis requested, so you only need to explicitly request fixtures you directly interact with.This pattern appears in several tests:
test_add_items_with_direct_memory,test_add_items_with_memory_text_only,test_search_success,test_remove_items_by_memory_id_episodic,test_remove_items_by_memory_id_semantic,test_add_items_with_custom_project_and_org, andtest_search_with_custom_project_and_org.♻️ Example fix for this test
async def test_add_items_with_direct_memory( memmachine_editor_with_client: MemMachineEditor, - mock_client: Mock, - mock_project: Mock, mock_memory_instance: Mock, sample_direct_memory_item: MemoryItem ):
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py`:
- Around line 111-193: Add assertions after each search to validate results: for
Test 1 (first retrieval) assert retrieved is non-empty and that
retrieved[0].memory or retrieved[0].conversation contains an expected
substring/tag; for Test 2 (direct_memory) assert the search result length > 0
and that at least one result's memory includes "allergic to peanuts" or has tags
containing "allergy"/"preference"; for Test 3 (multiple_memories) assert
len(all_memories) >= 3 and that tags include "fact_1"/"fact_2"/"fact_3". Place
these assertions immediately after the corresponding memory_client.search calls
that populate retrieved and all_memories (use the MemoryItem,
memory_client.search, memory_client.add_items, user_id/session_id/agent_id
variables to locate the sections).
🧹 Nitpick comments (1)
packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py (1)
98-100: Replace fixed sleeps with polling + timeout to reduce flakiness.A hardcoded 2‑second wait can be too short or unnecessarily long depending on load. Consider polling the search until results appear (bounded by a timeout) to make the test deterministic.
Also applies to: 135-136, 175-176
willkill07
requested changes
Jan 27, 2026
…t running Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb`:
- Around line 108-113: Update the markdown text to remove possessive
constructions for inanimate objects: change "MemMachine's LLM provider API keys
(e.g., OpenAI) are configured in the MemMachine `cfg.yml` file, not here." to a
non-possessive form such as "LLM provider API keys for MemMachine (e.g., OpenAI)
are configured in the MemMachine `cfg.yml` file, not here." and apply the same
change to the other occurrence referenced (around lines 575-579) so all
instances use "for MemMachine" or "the MemMachine" instead of "MemMachine’s".
In `@packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py`:
- Around line 59-67: The helper function currently named test_add_and_retrieve
(which pytest will collect) should be renamed to _run_add_and_retrieve to avoid
double test collection; update the function definition name and every call site
(e.g., change calls in test_add_and_retrieve_integration or any other tests that
invoke test_add_and_retrieve to call _run_add_and_retrieve instead) so tests
still execute the helper but pytest won't treat it as a standalone test.
In `@packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py`:
- Around line 178-184: The test function test_add_items_with_direct_memory
currently declares fixture parameters that are only used for setup and trigger
ruff ARG001; update the function signature to silence unused-fixture warnings by
prefixing unused fixtures with an underscore (e.g., change mock_client,
mock_project, mock_memory_instance to _mock_client, _mock_project,
_mock_memory_instance) or alternatively append "# noqa: ARG001" if you prefer a
comment; ensure you only rename the fixtures in the test function signature
(leave the fixture definitions intact) so references to
test_add_items_with_direct_memory still run correctly.
In `@packages/nvidia_nat_memmachine/tests/test_memory.py`:
- Around line 19-23: This module contains multiple async test functions but
lacks the pytest-asyncio marker; add a module-level marker by inserting
pytestmark = pytest.mark.asyncio near the top (after the imports) so all async
tests run under pytest-asyncio, or alternatively decorate each async test with
`@pytest.mark.asyncio`; reference the existing imports (pytest, Builder,
MemMachineMemoryClientConfig, memmachine_memory_client) and place the
module-level pytestmark immediately after those imports.
🧹 Nitpick comments (2)
packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py (1)
29-51: Add Google-style docstrings for the public config and factory.
memmachine_memory_clienthas no docstring and the config docstring isn’t Google‑style (Args/Returns). Please add Google‑style docstrings and wrap code entities in backticks. As per coding guidelines.packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py (1)
26-36: Addasyncioto module-level imports.
asynciois imported inside multiple test functions (lines 117, 187, 269, 330, 411). Move it to the module-level imports for cleaner code.✍️ Proposed fix
import os import uuid +import asyncio import pytest import requestsThen remove the
import asynciostatements from within the test functions.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Charlie-Yi-2002 <charlie.yi@memverge.com>
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/tests/test_memory.py`:
- Around line 105-123: The test's import_side_effect raises ImportError with a
custom message which triggers Ruff TRY003; change the raise in
import_side_effect inside test_memmachine_memory_client_import_error to raise
ImportError (no custom string) so the patch of builtins.__import__ simply raises
the exception, leaving pytest.raises to assert the expected "Could not import
MemMachineClient" from memmachine_memory_client; update import_side_effect
accordingly and keep the rest of the test (patch and async with
memmachine_memory_client) unchanged.
🧹 Nitpick comments (1)
packages/nvidia_nat_memmachine/tests/test_memory.py (1)
26-66: Align fixture function names and return type hints with test conventions.
Fixture names should use thefixture_prefix and return type hints for readability and linting consistency.♻️ Proposed refactor
`@pytest.fixture`(name="mock_builder") -def mock_builder_fixture(): +def fixture_mock_builder() -> Mock: """Fixture to provide a mocked Builder instance.""" return Mock(spec=Builder) `@pytest.fixture`(name="config") -def config_fixture(): +def fixture_config() -> MemMachineMemoryClientConfig: """Fixture to provide a MemMachineMemoryClientConfig instance.""" return MemMachineMemoryClientConfig( base_url="http://localhost:8095", org_id="test_org", project_id="test_project", timeout=30, max_retries=3 ) `@pytest.fixture`(name="config_minimal") -def config_minimal_fixture(): +def fixture_config_minimal() -> MemMachineMemoryClientConfig: """Fixture to provide a minimal MemMachineMemoryClientConfig instance.""" return MemMachineMemoryClientConfig( base_url="http://localhost:8095" ) `@pytest.fixture`(name="mock_memmachine_client") -def mock_memmachine_client_fixture(): +def fixture_mock_memmachine_client() -> Mock: """Fixture to provide a mocked MemMachineClient.""" mock_client = Mock() mock_client.base_url = "http://localhost:8095" return mock_client `@pytest.fixture`(name="mock_project") -def mock_project_fixture(): +def fixture_mock_project() -> Mock: """Fixture to provide a mocked Project instance.""" mock_project = Mock() mock_project.org_id = "test_org" mock_project.project_id = "test_project" return mock_projectAs per coding guidelines, ...
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py`:
- Around line 17-24: Update the module docstring at the top of
test_add_and_retrieve.py to remove the claim that results are verified—replace
"Verifies they match" with wording that accurately states the test prints or
demonstrates API interaction without asserting retrieval correctness (e.g.,
"Prints retrieved memories to demonstrate API interaction; no correctness
assertions are performed"). Ensure the change is made in the module-level
docstring that describes the script's three-step flow so it accurately reflects
intended behavior.
- Around line 224-226: The integration test function
test_add_and_retrieve_integration should be marked as slow; add the pytest
marker by annotating the function with `@pytest.mark.slow` alongside the existing
`@pytest.mark.integration` and `@pytest.mark.asyncio` decorators so pytest can skip
it in default runs; locate the test_add_and_retrieve_integration definition in
tests/test_add_and_retrieve.py and insert `@pytest.mark.slow` above the async def.
🧹 Nitpick comments (1)
packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py (1)
51-56: Narrow exception handling in health check.Catching
Exceptionwill hide unexpected errors and can cause silent skips. Prefer arequests.RequestExceptioncatch so genuine bugs still surface.♻️ Proposed change
- except Exception: + except requests.RequestException: return False
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Charlie-Yi-2002 <charlie.yi@memverge.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py`:
- Around line 38-56: The test is missing the pytest slow marker and
_memmachine_available() is blocking the event loop by using requests.get; change
_memmachine_available to be async (async def _memmachine_available(base_url:
str) -> bool), switch the synchronous requests call to an async httpx client
(use httpx.AsyncClient and await the GET to f"{base_url}/api/v2/health" with
timeout), narrow the caught exception to httpx.RequestError (or httpx.HTTPError)
instead of Exception, and add `@pytest.mark.slow` to the test function that relies
on this helper so the test is correctly marked as slow.
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/pyproject.toml`:
- Around line 1-3: Add the SPDX Apache-2.0 header as the very first lines of the
file and update the copyright year to the current year; specifically insert the
standard SPDX header above the existing [build-system] section so the file
begins with the license header followed by the existing build-backend and
requires lines.
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Fix all issues with AI agents
In `@packages/nvidia_nat_memmachine/README.md`:
- Line 69: Update the heading text "Usage in NeMo Agent toolkit" to use a
capital T for "Toolkit" so it reads "Usage in NeMo Agent Toolkit"; locate the
heading in packages/nvidia_nat_memmachine/README.md (the line with the string
"Usage in NeMo Agent toolkit") and modify only that heading text to follow the
project's MD casing convention.
- Line 1: The title line "# NVIDIA NeMo Agent toolkit - MemMachine Integration"
should use "Toolkit" with a capital T; update the heading text (the top-level
markdown heading) to "# NVIDIA NeMo Agent Toolkit - MemMachine Integration" so
the title follows the capitalization guideline while leaving the rest of the
README content unchanged.
- Line 1: Add the required SPDX Apache-2.0 header to the top of the README.md
file: insert the SPDX header comment (e.g. "SPDX-FileCopyrightText: <copyright
holder>" and "SPDX-License-Identifier: Apache-2.0") as the first lines of
packages/nvidia_nat_memmachine/README.md so the file begins with the required
license metadata for md/mdx files.
- Around line 50-65: Update the incorrect default port statement in the README:
change the sentence that currently reads "The server will start on
`http://localhost:8095` by default" to state the actual default
`http://localhost:8080` (or explicitly say "`http://localhost:8080` by default,
or the port you configured") and, if keeping the recommendation, clarify that
`8095` is only a suggested custom port; locate and edit the paragraph following
the memmachine-server usage example to make this correction.
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
Member
|
Thanks for going through my feedback. We have a bit of a backlog and some churn on work right now, so it may take me a couple days to give this another pass. |
…memmachine-server as optional to support a local deployment Signed-off-by: Charlie Yi <charlie.yi@memverge.com>
Author
No worries, let me know whenever you have time which other things need to be addressed! |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Add MemMachine as a memory plugin within NAT and also add a notebook example to demonstrate how to use MemMachine in NAT
By Submitting this PR I confirm:
Summary by CodeRabbit
New Features
Documentation
Examples
Tests
Chores
✏️ Tip: You can customize this high-level summary in your review settings.