From 19f7193f13b8584cd6e311edfce6e83a64371209 Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Thu, 22 Jan 2026 13:52:16 -0800 Subject: [PATCH 01/18] Add MemMachine memory integration for NeMo Agent toolkit - 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 --- packages/nvidia_nat_memmachine/README.md | 85 ++ .../memmachine_memory_example.ipynb | 937 ++++++++++++++++++ packages/nvidia_nat_memmachine/pyproject.toml | 56 ++ .../src/nat/meta/pypi.md | 24 + .../src/nat/plugins/__init__.py | 20 + .../src/nat/plugins/memmachine/__init__.py | 0 .../plugins/memmachine/memmachine_editor.py | 490 +++++++++ .../src/nat/plugins/memmachine/memory.py | 100 ++ .../src/nat/plugins/memmachine/register.py | 1 + .../tests/API_CALL_VERIFICATION.md | 129 +++ .../nvidia_nat_memmachine/tests/__init__.py | 14 + .../tests/test_add_and_retrieve.py | 187 ++++ .../tests/test_memmachine_api_calls.py | 619 ++++++++++++ .../tests/test_memmachine_editor.py | 511 ++++++++++ .../tests/test_memmachine_integration.py | 431 ++++++++ .../tests/test_memory.py | 193 ++++ 16 files changed, 3797 insertions(+) create mode 100644 packages/nvidia_nat_memmachine/README.md create mode 100644 packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb create mode 100644 packages/nvidia_nat_memmachine/pyproject.toml create mode 100644 packages/nvidia_nat_memmachine/src/nat/meta/pypi.md create mode 100644 packages/nvidia_nat_memmachine/src/nat/plugins/__init__.py create mode 100644 packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/__init__.py create mode 100644 packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py create mode 100644 packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py create mode 100644 packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py create mode 100644 packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md create mode 100644 packages/nvidia_nat_memmachine/tests/__init__.py create mode 100644 packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py create mode 100644 packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py create mode 100644 packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py create mode 100644 packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py create mode 100644 packages/nvidia_nat_memmachine/tests/test_memory.py diff --git a/packages/nvidia_nat_memmachine/README.md b/packages/nvidia_nat_memmachine/README.md new file mode 100644 index 0000000000..d0d032450f --- /dev/null +++ b/packages/nvidia_nat_memmachine/README.md @@ -0,0 +1,85 @@ +# NVIDIA NeMo Agent toolkit - MemMachine Integration + +This package provides integration with MemMachine for memory management in NeMo Agent toolkit. + +## Overview + +MemMachine is a unified memory management system that supports both episodic and semantic memory through a single interface. This integration allows you to use MemMachine as a memory backend for your NeMo Agent toolkit workflows. + +## Prerequisites + +- Python 3.11+ +- MemMachine server (install via `pip install memmachine-server`) + +## Installation + +Install the package: + +```bash +pip install nvidia-nat-memmachine +``` + +Or for development: + +```bash +uv pip install -e packages/nvidia_nat_memmachine +``` + +## MemMachine Server Setup + +### Step 1: Install MemMachine Server + +```bash +pip install memmachine-server +``` + +### Step 2: Run the Configuration Wizard + +Before starting the server, you need to configure MemMachine using the interactive configuration wizard: + +```bash +memmachine-configure +``` + +The wizard will guide you through setting up: + +- **Neo4j Database**: Option to install Neo4j automatically or provide connection details for an existing instance +- **Large Language Model (LLM) Provider**: Choose from supported providers like OpenAI, AWS Bedrock, or Ollama +- **Model Selection**: Select specific LLM and embedding models +- **API Keys and Credentials**: Input necessary API keys for your selected LLM provider +- **Server Settings**: Configure server host and port + +**Note**: +- The wizard installs Neo4j and Java automatically when you choose to install Neo4j (platform-specific: Windows uses ZIP, macOS uses brew, Linux uses tar.gz) +- The wizard uses Neo4j as the vector database and SQLite as the relational database by default +- The configuration file will be generated at `/.config/memmachine/cfg.yml` + +### Step 3: Start the MemMachine Server + +After completing the configuration wizard, start the server: + +```bash +memmachine-server +``` + +The server will start on `http://localhost:8080` by default (or the port you configured). + +For more details, see the [MemMachine Configuration Wizard Documentation](https://docs.memmachine.ai/open_source/configuration-wizard). + +## Usage in NeMo Agent toolkit + +Add MemMachine memory to your workflow configuration: + +```yaml +memory: + memmachine_memory: + base_url: "http://localhost:8080" # MemMachine server URL + org_id: "my_org" # Optional: default organization ID + project_id: "my_project" # Optional: default project ID +``` + +## Additional Resources + +- [MemMachine Documentation](https://docs.memmachine.ai/) +- [NeMo Agent toolkit Documentation](https://docs.nvidia.com/nemo/agent-toolkit/latest/) + diff --git a/packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb b/packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb new file mode 100644 index 0000000000..d6b9f8128b --- /dev/null +++ b/packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb @@ -0,0 +1,937 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# MemMachine Memory Integration with NeMo Agent toolkit\n", + "\n", + "This notebook demonstrates how to use MemMachine memory inside NeMo Agent toolkit end-to-end. MemMachine provides a unified memory management system where users can add conversations or memories directly.\n", + "\n", + "## What You'll Learn\n", + "\n", + "- How to set up and configure MemMachine server\n", + "- How to integrate MemMachine memory with NeMo Agent toolkit\n", + "- How to add and retrieve memories from conversations\n", + "- How to add and retrieve memories directly\n", + "- How to use memory in an agent workflow with tools" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Table of Contents\n", + "\n", + "- [0) Setup](#setup)\n", + " - [0.1) Prerequisites](#prereqs)\n", + " - [0.2) MemMachine Server Setup](#memmachine-setup)\n", + " - [0.3) API Keys](#api-keys)\n", + " - [0.4) Installing Dependencies](#installing-deps)\n", + " - [0.5) Verify MemMachine Server](#verify-server)\n", + "- [1) Basic Memory Operations](#basic-memory)\n", + " - [1.1) Programmatic Memory Usage](#programmatic)\n", + " - [1.2) Adding Memories from Conversations](#conversation-memory)\n", + " - [1.3) Adding Memories Directly](#direct-memory)\n", + " - [1.4) Searching Memories](#searching)\n", + "- [2) Agent Workflow with Memory](#agent-workflow)\n", + " - [2.1) Create Configuration](#config)\n", + " - [2.2) Run Agent with Memory](#run-agent)\n", + "- [3) Next Steps](#next-steps)\n", + "\n", + "Note: In Google Colab use the Table of Contents tab to navigate." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "## 0) Setup\n", + "\n", + "\n", + "### 0.1) Prerequisites\n", + "\n", + "- **Platform:** Linux, macOS, or Windows\n", + "- **Python:** version 3.11, 3.12, or 3.13\n", + "- **MemMachine Server:** Must be installed and running (see next section)\n", + "- **Neo4j Database:** Required by MemMachine (can be auto-installed)\n", + "- **LLM Provider:** OpenAI, AWS Bedrock, or Ollama (for memory processing)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "### 0.2) MemMachine Server Setup\n", + "\n", + "**IMPORTANT:** Before running this notebook, you must set up and start the MemMachine server.\n", + "\n", + "#### Step 1: Install MemMachine Server\n", + "```bash\n", + "pip install memmachine-server\n", + "```\n", + "\n", + "#### Step 2: Run Configuration Wizard\n", + "```bash\n", + "memmachine-configure\n", + "```\n", + "\n", + "The wizard will guide you through:\n", + "- **Neo4j Database**: Option to install Neo4j automatically or provide connection details\n", + "- **LLM Provider**: Choose from OpenAI, AWS Bedrock, Ollama, etc\n", + "- **Model Selection**: Select LLM and embedding models\n", + "- **API Keys**: Input necessary credentials\n", + "- **Server Settings**: Configure host and port (default: localhost:8080)\n", + "\n", + "Configuration file will be saved at: `~/.config/memmachine/cfg.yml`\n", + "\n", + "#### Step 3: Start MemMachine Server\n", + "In a separate terminal, run:\n", + "```bash\n", + "memmachine-server\n", + "```\n", + "\n", + "The server will start on `http://localhost:8080` by default.\n", + "\n", + "For more details, see the [MemMachine Documentation](https://docs.memmachine.ai/open_source/configuration-wizard)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "### 0.3) API Keys\n", + "\n", + "For this notebook, you will need the following API keys:\n", + "\n", + "- **NVIDIA Build:** You can obtain an NVIDIA Build API Key by creating an [NVIDIA Build](https://build.nvidia.com) account and generating a key at https://build.nvidia.com/settings/api-keys\n", + "\n", + "**Note:** MemMachine's LLM provider API keys (e.g., OpenAI) are configured in the MemMachine `cfg.yml` file, not here." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verify that the langchain plugin is properly installed and registered\n", + "import sys\n", + "from importlib.metadata import entry_points\n", + "\n", + "try:\n", + " # Check if the entry point exists\n", + " components = entry_points(group='nat.components')\n", + " langchain_entry = None\n", + " for ep in components:\n", + " if 'langchain' in ep.name.lower() and 'tools' not in ep.name.lower():\n", + " langchain_entry = ep\n", + " break\n", + " \n", + " if langchain_entry:\n", + " print(f\"Found langchain entry point: {langchain_entry.name} -> {langchain_entry.module}\")\n", + " \n", + " # Check if the module path is correct\n", + " expected_module = \"nat.plugins.langchain.register\"\n", + " if langchain_entry.module != expected_module:\n", + " print(f\"āŒ Entry point module path is incorrect!\")\n", + " print(f\" Expected: {expected_module}\")\n", + " print(f\" Found: {langchain_entry.module}\")\n", + " else:\n", + " print(f\"āœ… Entry point module path is correct: {langchain_entry.module}\")\n", + " try:\n", + " # Try to load the module to verify it works\n", + " langchain_entry.load()\n", + " print(\"āœ… LangChain plugin module loaded successfully!\")\n", + " except Exception as e:\n", + " print(f\"āŒ Failed to load langchain plugin module: {e}\")\n", + " print(\"\\nšŸ’” Try reinstalling the package:\")\n", + " print(\" uv pip install --force-reinstall \\\"nvidia-nat[langchain]\\\"\")\n", + " else:\n", + " print(\"āŒ langchain entry point not found!\")\n", + " print(\"\\nšŸ’” The langchain plugin may not be installed correctly.\")\n", + " print(\" Try reinstalling:\")\n", + " print(\" uv pip install --force-reinstall \\\"nvidia-nat[langchain]\\\"\")\n", + " print(\"\\n Or if you're in the repo directory:\")\n", + " print(\" uv pip install --force-reinstall -e packages/nvidia_nat_langchain\")\n", + "except Exception as e:\n", + " print(f\"āŒ Error checking entry points: {e}\")\n", + " print(\"\\nšŸ’” Try reinstalling the package:\")\n", + " print(\" uv pip install --force-reinstall \\\"nvidia-nat[langchain]\\\"\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import getpass\n", + "import os\n", + "\n", + "if \"NVIDIA_API_KEY\" not in os.environ:\n", + " nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n", + " os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verify that the memmachine plugin is properly installed and registered\n", + "import sys\n", + "from importlib.metadata import entry_points\n", + "\n", + "try:\n", + " # Check if the entry point exists\n", + " components = entry_points(group='nat.components')\n", + " memmachine_entry = None\n", + " for ep in components:\n", + " if 'memmachine' in ep.name.lower():\n", + " memmachine_entry = ep\n", + " break\n", + " \n", + " if memmachine_entry:\n", + " print(f\"Found memmachine entry point: {memmachine_entry.name} -> {memmachine_entry.module}\")\n", + " \n", + " # Check if the module path is correct\n", + " expected_module = \"nat.plugins.memmachine.register\"\n", + " if memmachine_entry.module != expected_module:\n", + " print(f\"āŒ Entry point module path is incorrect!\")\n", + " print(f\" Expected: {expected_module}\")\n", + " print(f\" Found: {memmachine_entry.module}\")\n", + " print(\"\\nšŸ’” The package needs to be reinstalled in editable mode.\")\n", + " print(\" Run this command in a terminal:\")\n", + " print(\" uv pip install --force-reinstall -e ../../packages/nvidia_nat_memmachine\")\n", + " print(\"\\n Or if you're in the repo root:\")\n", + " print(\" uv pip install --force-reinstall -e packages/nvidia_nat_memmachine\")\n", + " else:\n", + " print(f\"āœ… Entry point module path is correct: {memmachine_entry.module}\")\n", + " try:\n", + " # Try to load the module to verify it works\n", + " memmachine_entry.load()\n", + " print(\"āœ… Plugin module loaded successfully!\")\n", + " except Exception as e:\n", + " print(f\"āŒ Failed to load plugin module: {e}\")\n", + " print(\"\\nšŸ’” Try reinstalling the package in editable mode:\")\n", + " print(\" uv pip install --force-reinstall -e ../../packages/nvidia_nat_memmachine\")\n", + " else:\n", + " print(\"āŒ memmachine entry point not found!\")\n", + " print(\"\\nšŸ’” The plugin may not be installed correctly.\")\n", + " print(\" Try reinstalling in editable mode:\")\n", + " print(\" uv pip install --force-reinstall -e ../../packages/nvidia_nat_memmachine\")\n", + " print(\"\\n Or if you're not in the repo directory:\")\n", + " print(\" uv pip install --force-reinstall nvidia-nat-memmachine\")\n", + "except Exception as e:\n", + " print(f\"āŒ Error checking entry points: {e}\")\n", + " print(\"\\nšŸ’” Try reinstalling the package:\")\n", + " print(\" uv pip install --force-reinstall -e ../../packages/nvidia_nat_memmachine\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "### 0.4) Installing Dependencies\n", + "\n", + "We'll use `pip` to install packages directly into the Jupyter kernel's Python environment. This ensures that when we run agent commands from within the notebook, they have access to all installed packages." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# We'll use pip (which comes with Python) for package installation\n", + "# This ensures packages are installed into the same environment as the Jupyter kernel\n", + "print(\"Ready to install packages using pip\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now install NeMo Agent toolkit with langchain support and the MemMachine integration.\n", + "\n", + "**Note:** These installation cells use `sys.executable` to ensure packages are installed into the same Python environment as your Jupyter kernel. After installation, you may need to restart the kernel for the changes to take effect." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "import subprocess\n", + "\n", + "# Install nvidia-nat with langchain support using the same Python as Jupyter\n", + "# This is REQUIRED for react_agent workflow type\n", + "print(\"Checking if nvidia-nat-langchain is installed...\")\n", + "result = subprocess.run(\n", + " [sys.executable, \"-m\", \"pip\", \"show\", \"nvidia-nat-langchain\"],\n", + " capture_output=True,\n", + " text=True\n", + ")\n", + "\n", + "if result.returncode != 0:\n", + " print(\"Installing nvidia-nat[langchain]...\")\n", + " subprocess.run(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"nvidia-nat[langchain]\"],\n", + " check=True\n", + " )\n", + " print(\"āœ… Installation complete\")\n", + " print(\"\\nāš ļø Please restart the kernel (Kernel ā†’ Restart Kernel) for changes to take effect.\")\n", + "else:\n", + " print(\"āœ… nvidia-nat-langchain is already installed\")\n", + " print(\"\\nIf you encounter 'No module named nat.plugins.langchain' errors,\")\n", + " print(\"try running:\")\n", + " print(f\" !{sys.executable} -m pip install --force-reinstall 'nvidia-nat[langchain]'\")\n", + " print(\"Then restart the kernel.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "import subprocess\n", + "import os\n", + "\n", + "# Install nvidia-nat-memmachine package using the same Python as Jupyter\n", + "# IMPORTANT: For the plugin to work with YAML configs, it must be installed\n", + "# in editable mode if developing locally, or entry points won't be registered\n", + "print(\"Checking if nvidia-nat-memmachine is installed...\")\n", + "result = subprocess.run(\n", + " [sys.executable, \"-m\", \"pip\", \"show\", \"nvidia-nat-memmachine\"],\n", + " capture_output=True,\n", + " text=True\n", + ")\n", + "\n", + "if result.returncode != 0:\n", + " # Try installing from local source first (for development)\n", + " if os.path.isdir(\"../../packages/nvidia_nat_memmachine\"):\n", + " print(\"Installing nvidia-nat-memmachine in editable mode from local source...\")\n", + " subprocess.run(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"-e\", \"../../packages/nvidia_nat_memmachine\"],\n", + " check=True\n", + " )\n", + " else:\n", + " print(\"Installing nvidia-nat-memmachine from PyPI...\")\n", + " subprocess.run(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"nvidia-nat-memmachine\"],\n", + " check=True\n", + " )\n", + " print(\"āœ… Installation complete\")\n", + " print(\"\\nāš ļø Please restart the kernel for changes to take effect.\")\n", + "else:\n", + " print(\"āœ… nvidia-nat-memmachine is already installed\")\n", + " print(\"\\nIf you encounter plugin loading errors, try running:\")\n", + " print(f\" !{sys.executable} -m pip install --force-reinstall -e ../../packages/nvidia_nat_memmachine\")\n", + " print(\"Then restart the kernel.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sys\n", + "import subprocess\n", + "\n", + "# Install memmachine-server using the same Python as Jupyter\n", + "print(\"Checking if memmachine-server is installed...\")\n", + "result = subprocess.run(\n", + " [sys.executable, \"-m\", \"pip\", \"show\", \"memmachine-server\"],\n", + " capture_output=True,\n", + " text=True\n", + ")\n", + "\n", + "if result.returncode != 0:\n", + " print(\"Installing memmachine-server~=0.2.2...\")\n", + " subprocess.run(\n", + " [sys.executable, \"-m\", \"pip\", \"install\", \"memmachine-server~=0.2.2\"],\n", + " check=True\n", + " )\n", + " print(\"āœ… Installation complete\")\n", + "else:\n", + " print(\"āœ… memmachine-server is already installed\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "### 0.5) Verify MemMachine Server is Running\n", + "\n", + "Let's check if the MemMachine server is accessible:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "# MemMachine server URL (default)\n", + "MEMMACHINE_BASE_URL = os.environ.get(\"MEMMACHINE_BASE_URL\", \"http://localhost:8080\")\n", + "\n", + "try:\n", + " response = requests.get(f\"{MEMMACHINE_BASE_URL}/api/v2/health\", timeout=5)\n", + " if response.status_code == 200:\n", + " print(f\"āœ… MemMachine server is running at {MEMMACHINE_BASE_URL}\")\n", + " else:\n", + " print(f\"āš ļø MemMachine server responded with status {response.status_code}\")\n", + "except requests.exceptions.ConnectionError:\n", + " print(f\"āŒ Cannot connect to MemMachine server at {MEMMACHINE_BASE_URL}\")\n", + " print(\"Please ensure the server is running. Start it with: memmachine-server\")\n", + "except Exception as e:\n", + " print(f\"āŒ Error checking server: {e}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "## 1) Basic Memory Operations\n", + "\n", + "Let's explore how to use MemMachine memory programmatically with NeMo Agent toolkit.\n", + "\n", + "\n", + "### 1.1) Programmatic Memory Usage\n", + "\n", + "First, let's import the necessary modules:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import asyncio\n", + "import uuid\n", + "from nat.builder.workflow_builder import WorkflowBuilder\n", + "from nat.data_models.config import GeneralConfig\n", + "from nat.memory.models import MemoryItem\n", + "from nat.plugins.memmachine.memory import MemMachineMemoryClientConfig\n", + "\n", + "# Create a unique test ID for this session\n", + "test_id = str(uuid.uuid4())[:8]\n", + "print(f\"Test session ID: {test_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now, let's configure and create a MemMachine memory client:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Configure MemMachine memory client\n", + "memmachine_config = MemMachineMemoryClientConfig(\n", + " base_url=MEMMACHINE_BASE_URL,\n", + " org_id=f\"demo_org_{test_id}\",\n", + " project_id=f\"demo_project_{test_id}\",\n", + " timeout=30,\n", + " max_retries=3\n", + ")\n", + "\n", + "print(f\"āœ… MemMachine configuration created\")\n", + "print(f\" Base URL: {memmachine_config.base_url}\")\n", + "print(f\" Org ID: {memmachine_config.org_id}\")\n", + "print(f\" Project ID: {memmachine_config.project_id}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "### 1.2) Adding Memories from Conversations\n", + "\n", + "Memories can be added from conversations, preserving the full context of interactions. All memories are added to both episodic and semantic memory types. Let's add a memory from a conversation:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "async def add_memory_from_conversation():\n", + " \"\"\"Add a memory from a conversation\"\"\"\n", + " general_config = GeneralConfig()\n", + " \n", + " async with WorkflowBuilder(general_config=general_config) as builder:\n", + " # Add MemMachine memory client\n", + " await builder.add_memory_client(\"memmachine_memory\", memmachine_config)\n", + " memory_client = await builder.get_memory_client(\"memmachine_memory\")\n", + " \n", + " # Create a memory with conversation context\n", + " user_id = f\"demo_user_{test_id}\"\n", + " conversation = [\n", + " {\"role\": \"user\", \"content\": \"I love pizza and Italian food, especially margherita pizza.\"},\n", + " {\"role\": \"assistant\", \"content\": \"I'll remember that you love pizza and Italian food, especially margherita pizza.\"},\n", + " ]\n", + " \n", + " memory_item = MemoryItem(\n", + " conversation=conversation,\n", + " user_id=user_id,\n", + " memory=\"User loves pizza and Italian food, especially margherita pizza\",\n", + " metadata={\n", + " \"session_id\": f\"session_{test_id}\",\n", + " \"agent_id\": f\"agent_{test_id}\",\n", + " \"test_id\": \"conversation_demo\"\n", + " },\n", + " tags=[\"food\", \"preference\", \"italian\"]\n", + " )\n", + " \n", + " # Add the memory\n", + " await memory_client.add_items([memory_item])\n", + " print(f\"āœ… Added memory from conversation for user: {user_id}\")\n", + " \n", + " # Wait a moment for indexing\n", + " await asyncio.sleep(2)\n", + " \n", + " return user_id, memory_client\n", + "\n", + "# Run the async function\n", + "user_id, memory_client = await add_memory_from_conversation()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "### 1.3) Adding Memories Directly\n", + "\n", + "Memories can also be added directly without a conversation. All memories are added to both episodic and semantic memory types. These are great for storing long-term user preferences and facts:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "async def add_memory_directly():\n", + " \"\"\"Add a memory directly (without conversation) using the existing memory client\"\"\"\n", + " # Reuse the memory_client from the previous cell\n", + " # This avoids trying to create the project again\n", + " \n", + " # Create a memory directly (without conversation)\n", + " direct_memory = MemoryItem(\n", + " conversation=None, # No conversation for direct memory\n", + " user_id=user_id,\n", + " memory=\"User prefers working in the morning (9 AM - 12 PM) and is allergic to peanuts\",\n", + " metadata={\n", + " \"session_id\": f\"session_{test_id}\",\n", + " \"agent_id\": f\"agent_{test_id}\",\n", + " \"test_id\": \"direct_demo\"\n", + " },\n", + " tags=[\"preference\", \"allergy\", \"schedule\"]\n", + " )\n", + " \n", + " # Add the memory using the existing memory_client\n", + " await memory_client.add_items([direct_memory])\n", + " print(f\"āœ… Added memory directly for user: {user_id}\")\n", + " \n", + " # Direct memories are processed asynchronously\n", + " # Wait longer for background ingestion task\n", + " print(\"ā³ Waiting for memory ingestion (this may take 2-5 seconds)...\")\n", + " await asyncio.sleep(5)\n", + " \n", + " return memory_client\n", + "\n", + "memory_client = await add_memory_directly()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "### 1.4) Searching Memories\n", + "\n", + "Now let's search for the memories we just added. **Note:** MemMachine's search function returns all memories in a single search call, whether they were added from conversations or directly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "async def search_memories():\n", + " \"\"\"Search for memories - returns all memories in one call\"\"\"\n", + " # Reuse the memory_client from the previous cell\n", + " # This avoids trying to create the project again\n", + " \n", + " # Single search returns all memories\n", + " print(\"šŸ” Searching for memories (pizza/Italian food)...\")\n", + " print(\" Note: This search returns all memories (from conversations and direct)\\n\")\n", + " \n", + " all_results = await memory_client.search(\n", + " query=\"pizza Italian food margherita\",\n", + " top_k=10,\n", + " user_id=user_id,\n", + " session_id=f\"session_{test_id}\",\n", + " agent_id=f\"agent_{test_id}\"\n", + " )\n", + " \n", + " print(f\" Found {len(all_results)} total memories\\n\")\n", + " \n", + " # Display results - memories from conversations have conversation field, direct memories don't\n", + " for i, mem in enumerate(all_results, 1):\n", + " memory_type = \"From Conversation\" if mem.conversation else \"Direct\"\n", + " print(f\" {i}. [{memory_type}] {mem.memory}\")\n", + " if mem.conversation:\n", + " print(f\" Conversation: {len(mem.conversation)} messages\")\n", + " if mem.tags:\n", + " print(f\" Tags: {', '.join(mem.tags)}\")\n", + " print()\n", + " \n", + " # Now search for direct memory (may need retries due to async processing)\n", + " print(\"\\nšŸ” Searching for direct memory (morning work allergy)...\")\n", + " print(\" Note: Direct memories may take a few seconds to be searchable\\n\")\n", + " \n", + " direct_results = []\n", + " for attempt in range(3):\n", + " direct_results = await memory_client.search(\n", + " query=\"morning work schedule allergy peanuts\",\n", + " top_k=10,\n", + " user_id=user_id,\n", + " session_id=f\"session_{test_id}\",\n", + " agent_id=f\"agent_{test_id}\"\n", + " )\n", + " # Filter for direct memories (no conversation)\n", + " direct_only = [m for m in direct_results if not m.conversation]\n", + " if len(direct_only) > 0:\n", + " direct_results = direct_only\n", + " break\n", + " print(f\" Attempt {attempt + 1}: No direct memory results yet, waiting...\")\n", + " await asyncio.sleep(2)\n", + " \n", + " print(f\" Found {len(direct_results)} direct memories\")\n", + " for i, mem in enumerate(direct_results, 1):\n", + " print(f\" {i}. {mem.memory}\")\n", + " if mem.tags:\n", + " print(f\" Tags: {', '.join(mem.tags)}\")\n", + "\n", + "await search_memories()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "## 2) Agent Workflow with Memory\n", + "\n", + "Now let's create an agent workflow that can use memory tools to remember and recall information.\n", + "\n", + "\n", + "### 2.1) Create Configuration\n", + "\n", + "Let's create a YAML configuration file for an agent with memory capabilities:\n", + "\n", + "**Note:** The `react_agent` workflow type requires `nvidia-nat[langchain]` to be installed. This is because the ReAct agent uses LangChain/LangGraph for its agent framework. If you encounter errors about missing `langchain` modules (like `No module named 'langchain.schema'`), make sure you've run the installation cell above that installs `nvidia-nat[langchain]`. The ReAct agent is one of several agent types available in NeMo Agent toolkit - it requires langchain because it uses LangGraph for agent orchestration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Write the agent config file using the same test_id from Part 1\n", + "# This ensures the agent uses the same MemMachine project we created earlier\n", + "\n", + "agent_config = f'''general:\n", + " telemetry:\n", + " enabled: false\n", + "\n", + "llms:\n", + " nim_llm:\n", + " _type: nim\n", + " model_name: meta/llama-3.1-70b-instruct\n", + " temperature: 0.7\n", + " max_tokens: 1024\n", + "\n", + "memory:\n", + " memmachine_memory:\n", + " _type: memmachine_memory\n", + " base_url: \"http://localhost:8080\"\n", + " org_id: \"{memmachine_config.org_id}\"\n", + " project_id: \"{memmachine_config.project_id}\"\n", + "\n", + "functions:\n", + " get_memory:\n", + " _type: get_memory\n", + " memory: memmachine_memory\n", + " description: |\n", + " Retrieve memories relevant to a query. Always call this tool first to check for existing user preferences or facts.\n", + " Use the exact JSON format with user_id, query, and top_k parameters.\n", + "\n", + " add_memory:\n", + " _type: add_memory\n", + " memory: memmachine_memory\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", + "\n", + "workflow:\n", + " _type: react_agent\n", + " tool_names: [add_memory, get_memory]\n", + " description: \"A chat agent that can remember and recall user preferences using MemMachine memory\"\n", + " llm_name: nim_llm\n", + " verbose: true\n", + " max_tool_calls: 5\n", + " system_prompt: |\n", + " You are a helpful assistant with access to memory tools. Always use user_id \"{user_id}\" for memory operations.\n", + "\n", + " {{tools}}\n", + "\n", + " Use this format:\n", + "\n", + " Question: the input question you must answer\n", + " Thought: think about what to do\n", + " Action: the action to take, one of [{{tool_names}}]\n", + " Action Input: {{{{\"key\": \"value\"}}}}\n", + " Observation: the result of the action\n", + " ... (repeat Thought/Action/Action Input/Observation as needed)\n", + " Thought: I now know the final answer\n", + " Final Answer: your final answer\n", + "\n", + " CRITICAL: Action Input must be ONLY valid JSON on a single line. No extra text before or after the JSON.\n", + "'''\n", + "\n", + "# Write the config file\n", + "with open('memmachine_agent_config.yml', 'w') as f:\n", + " f.write(agent_config)\n", + "\n", + "print(f\"āœ… Created agent config with project: {memmachine_config.org_id}/{memmachine_config.project_id}\")\n", + "print(f\" Config file: memmachine_agent_config.yml\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verify that langchain plugin can be imported in the current Python environment\n", + "import sys\n", + "print(f\"Current Python: {sys.executable}\")\n", + "print(f\"Python version: {sys.version}\\n\")\n", + "\n", + "try:\n", + " # Try to import the langchain plugin module directly\n", + " from nat.plugins.langchain import register\n", + " print(\"āœ… Successfully imported nat.plugins.langchain.register\")\n", + " print(\"āœ… LangChain plugin is available in this Python environment\\n\")\n", + " \n", + " # Also verify entry points\n", + " from importlib.metadata import entry_points\n", + " components = entry_points(group='nat.components')\n", + " langchain_eps = [ep for ep in components if 'langchain' in ep.name.lower() and 'tools' not in ep.name.lower()]\n", + " \n", + " if langchain_eps:\n", + " print(f\"āœ… Found langchain entry point: {langchain_eps[0].name}\")\n", + " print(\"āœ… Ready to run agent workflows with langchain!\\n\")\n", + " else:\n", + " print(\"āš ļø Warning: LangChain entry point not found, but module can be imported\")\n", + " print(\" This may still work, but entry points should be registered.\\n\")\n", + " \n", + "except ImportError as e:\n", + " print(f\"āŒ Failed to import langchain plugin: {e}\")\n", + " print(\"\\nšŸ’” The langchain plugin is not available in this Python environment.\")\n", + " print(\" This means `!nat run` will fail when trying to use react_agent workflow.\")\n", + " print(\"\\n Solutions:\")\n", + " print(\" 1. Make sure you've run the installation cell above that installs nvidia-nat[langchain]\")\n", + " print(\" 2. Restart the Jupyter kernel after installing\")\n", + " print(\" 3. Verify you're using the correct Python environment (venv)\")\n", + " print(f\" Current Python: {sys.executable}\")\n", + " print(\"\\n If the issue persists, try reinstalling:\")\n", + " print(\" uv pip install --force-reinstall 'nvidia-nat[langchain]'\")\n", + " raise" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### 2.2) Verify LangChain Plugin and Run Agent\n", + "\n", + "**Important:** The `react_agent` workflow requires the langchain plugin. When running `!nat run` from Jupyter, it may use system Python instead of the venv, causing the langchain plugin to not be found.\n", + "\n", + "**Solution:** We use `python -m nat.cli.main run` instead of `!nat run` to ensure we use the venv's Python interpreter, which has access to the langchain plugin." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Run the agent using the venv's Python to ensure langchain plugin is accessible\n", + "# Using sys.executable ensures we use the same Python as the notebook kernel (venv)\n", + "import sys\n", + "\n", + "# Use python -m nat.cli.main instead of !nat to ensure correct Python environment\n", + "!{sys.executable} -m nat.cli.main run --config_file memmachine_agent_config.yml --input \"What is my favorite food?\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Verify that the memmachine plugin is properly installed and registered\n", + "import sys\n", + "from importlib.metadata import entry_points\n", + "\n", + "try:\n", + " # Check if the entry point exists\n", + " components = entry_points(group='nat.components')\n", + " memmachine_entry = None\n", + " for ep in components:\n", + " if 'memmachine' in ep.name.lower():\n", + " memmachine_entry = ep\n", + " break\n", + " \n", + " if memmachine_entry:\n", + " print(f\"āœ… Found memmachine entry point: {memmachine_entry.name} -> {memmachine_entry.module}\")\n", + " try:\n", + " # Try to load the module to verify it works\n", + " memmachine_entry.load()\n", + " print(\"āœ… Plugin module loaded successfully\")\n", + " except Exception as e:\n", + " print(f\"āŒ Failed to load plugin module: {e}\")\n", + " print(\"\\nšŸ’” Try reinstalling the package in editable mode:\")\n", + " print(\" uv pip install --force-reinstall -e ../../packages/nvidia_nat_memmachine\")\n", + " else:\n", + " print(\"āŒ memmachine entry point not found!\")\n", + " print(\"\\nšŸ’” The plugin may not be installed correctly.\")\n", + " print(\" Try reinstalling in editable mode:\")\n", + " print(\" uv pip install --force-reinstall -e ../../packages/nvidia_nat_memmachine\")\n", + " print(\"\\n Or if you're not in the repo directory:\")\n", + " print(\" uv pip install --force-reinstall nvidia-nat-memmachine\")\n", + "except Exception as e:\n", + " print(f\"āŒ Error checking entry points: {e}\")\n", + " print(\"\\nšŸ’” Try reinstalling the package:\")\n", + " print(\" uv pip install --force-reinstall -e ../../packages/nvidia_nat_memmachine\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Great! The agent should have retrieved the memory about Python. Now let's tell the agent about another preference so it can add it to memory:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Use Python's -m flag to ensure we use the venv's nat module\n", + "import sys\n", + "!{sys.executable} -m nat.cli.main run --config_file memmachine_agent_config.yml --input \"I love reading science fiction novels like Dune, can you recommend some other books in the genre?\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now let's verify the agent can recall our book preference from memory:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Use Python's -m flag to ensure we use the venv's nat module\n", + "import sys\n", + "!{sys.executable} -m nat.cli.main run --config_file memmachine_agent_config.yml --input \"What other books do you think I would like? Also recommend some movies in the genre.\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "## 3) Next Steps\n", + "\n", + "Congratulations! You've successfully integrated MemMachine memory with NeMo Agent toolkit. Here are some next steps to explore:\n", + "\n", + "1. **Explore Advanced Memory Features**:\n", + " - Use metadata and tags for better memory organization\n", + " - Experiment with different ways to add memories (from conversations vs directly)\n", + " - Try memory deletion and cleanup strategies\n", + "\n", + "2. **Integrate with Other Components**:\n", + " - Combine memory with RAG (Retrieval Augmented Generation)\n", + " - Use memory in multi-agent workflows\n", + " - Add memory to custom tools and functions\n", + "\n", + "3. **Production Considerations**:\n", + " - Set up proper Neo4j database management\n", + " - Configure memory retention policies\n", + " - Implement memory search optimization\n", + " - Add monitoring and observability\n", + "\n", + "4. **Additional Resources**:\n", + " - [MemMachine Documentation](https://docs.memmachine.ai/)\n", + " - [NeMo Agent toolkit Documentation](https://docs.nvidia.com/nemo/agent-toolkit/latest/)\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.9" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/packages/nvidia_nat_memmachine/pyproject.toml b/packages/nvidia_nat_memmachine/pyproject.toml new file mode 100644 index 0000000000..2e2b0c6ceb --- /dev/null +++ b/packages/nvidia_nat_memmachine/pyproject.toml @@ -0,0 +1,56 @@ +[build-system] +build-backend = "setuptools.build_meta" +requires = ["setuptools >= 64", "setuptools-scm>=8"] + + +[tool.setuptools.packages.find] +where = ["src"] +include = ["nat.*"] + + +[tool.setuptools_scm] +git_describe_command = "git describe --long --first-parent" +root = "../.." + + +[project] +name = "nvidia-nat-memmachine" +dynamic = ["version"] +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~=1.4", + "memmachine-server~=0.2.2", +] +requires-python = ">=3.11,<3.14" +description = "Subpackage for MemMachine integration in NeMo Agent toolkit. Requires a cfg.yml configuration file with database and AI model settings." +readme = "src/nat/meta/pypi.md" +keywords = ["ai", "agents", "memory"] +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.uv] +managed = true +config-settings = { editable_mode = "compat" } + + +[tool.uv.sources] +nvidia-nat = { workspace = true } + + +[project.entry-points.'nat.components'] +nat_memmachine = "nat.plugins.memmachine.register" + diff --git a/packages/nvidia_nat_memmachine/src/nat/meta/pypi.md b/packages/nvidia_nat_memmachine/src/nat/meta/pypi.md new file mode 100644 index 0000000000..9f602c9b63 --- /dev/null +++ b/packages/nvidia_nat_memmachine/src/nat/meta/pypi.md @@ -0,0 +1,24 @@ + + +![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 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). + diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/__init__.py b/packages/nvidia_nat_memmachine/src/nat/plugins/__init__.py new file mode 100644 index 0000000000..f05adda80b --- /dev/null +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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. + +# Namespace package - extend __path__ to include all nat.plugins locations +# This allows plugins installed in different locations (e.g., site-packages, editable installs) +# to be discovered as part of the same namespace +from pkgutil import extend_path +__path__ = extend_path(__path__, __name__) diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/__init__.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py new file mode 100644 index 0000000000..8fe73f7945 --- /dev/null +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py @@ -0,0 +1,490 @@ +import asyncio +from typing import Any + +import requests +from memmachine.common.api import MemoryType + +from nat.memory.interfaces import MemoryEditor +from nat.memory.models import MemoryItem + + +class MemMachineEditor(MemoryEditor): + """ + Wrapper class that implements NAT interfaces for MemMachine Integrations. + Uses the MemMachine Python SDK (MemMachineClient) as documented at: + https://github.com/MemMachine/MemMachine/blob/main/docs/examples/python.mdx + + Supports both episodic and semantic memory through the unified SDK interface. + + User needs to add MemMachine SDK ids as metadata to the MemoryItem: + - session_id + - agent_id + - group_id + - project_id + - org_id + + Group ID is optional. If not provided, the memory will be added to the 'default' group. + """ + + def __init__(self, memmachine_instance: Any): + """ + Initialize class with MemMachine instance. + + Args: + memmachine_instance: Preinstantiated MemMachineClient or Project object + from the MemMachine Python SDK. If a MemMachineClient is provided, + projects will be created/retrieved as needed. If a Project is provided, + it will be used directly. + """ + self._memmachine = memmachine_instance + # Check if it's a client or project + self._is_client = hasattr(memmachine_instance, 'create_project') + self._is_project = hasattr(memmachine_instance, 'memory') and not self._is_client + + def _get_memory_instance( + self, + user_id: str, + session_id: str, + agent_id: str, + group_id: str = "default", + project_id: str | None = None, + org_id: str | None = None + ) -> Any: + """ + Get or create a memory instance for the given context using the MemMachine SDK. + + Args: + user_id: User identifier + session_id: Session identifier + agent_id: Agent identifier + group_id: Group identifier (default: "default") + project_id: Optional project identifier (default: "default-project") + org_id: Optional organization identifier (default: "default-org") + + Returns: + Memory instance from MemMachine SDK + """ + # Use defaults if not provided + if not org_id: + org_id = "default-org" + if not project_id: + project_id = "default-project" + + # If we have a client, get or create the project first + if self._is_client: + # Use get_or_create_project which handles existing projects gracefully + # It will get the project if it exists, or create it if it doesn't + try: + project = self._memmachine.get_or_create_project( + org_id=org_id, + project_id=project_id, + description=f"Project for {user_id}" + ) + except requests.HTTPError as e: + # If get_or_create_project fails with 409 conflict, project already exists + # Get the existing project instead + if e.response.status_code == 409: + project = self._memmachine.get_project( + org_id=org_id, + project_id=project_id + ) + else: + # Re-raise other HTTP errors + raise + elif self._is_project: + # Use the project directly + project = self._memmachine + else: + # Fallback: assume it's already a memory instance or try to use it directly + return self._memmachine + + # Create memory instance from project + return project.memory( + user_id=user_id, + agent_id=agent_id, + session_id=session_id, + group_id=group_id + ) + + async def add_items(self, items: list[MemoryItem]) -> None: + """ + Insert Multiple MemoryItems into the memory using the MemMachine SDK. + Each MemoryItem is translated and uploaded through the MemMachine API. + + All memories are added to both episodic and semantic memory types. + """ + # Run synchronous operations in thread pool to make them async + tasks = [] + + for memory_item in items: + # Make a copy of metadata to avoid modifying the original + item_meta = memory_item.metadata.copy() if memory_item.metadata else {} + conversation = memory_item.conversation + user_id = memory_item.user_id + tags = memory_item.tags + memory_text = memory_item.memory + + # Extract session_id, agent_id, group_id, project_id, and org_id from metadata if present + session_id = item_meta.pop("session_id", "default_session") + agent_id = item_meta.pop("agent_id", "default_agent") + group_id = item_meta.pop("group_id", "default") + project_id = item_meta.pop("project_id", None) + org_id = item_meta.pop("org_id", None) + + # Get memory instance using MemMachine SDK + memory = self._get_memory_instance( + user_id, session_id, agent_id, group_id, project_id, org_id + ) + + # All memories are added to BOTH episodic and semantic memory types + memory_types = [MemoryType.Episodic, MemoryType.Semantic] + + # Prepare content for MemMachine + # If we have a conversation, add each message separately + # Otherwise, use memory_text or skip if no content + if conversation: + # Add each message in the conversation with its role + for msg in conversation: + msg_role = msg.get('role', 'user') + msg_content = msg.get('content', '') + + if not msg_content: + continue + + # Add tags to metadata if present + # MemMachine SDK expects tags as a string, not a list + metadata = item_meta.copy() if item_meta else {} + if tags: + # Convert list to comma-separated string + metadata["tags"] = ", ".join(tags) if isinstance(tags, list) else str(tags) + + # Capture variables in closure to avoid late binding issues + def add_memory(content=msg_content, role=msg_role, mem_types=memory_types, meta=metadata): + # Use MemMachine SDK add() method + # API: memory.add(content, role="user", metadata={}, memory_types=[...]) + # episode_type should be None (defaults to "message") or EpisodeType.MESSAGE + memory.add( + content=content, + role=role, + metadata=meta if meta else None, + memory_types=mem_types, + episode_type=None # Use default (MESSAGE) + ) + + task = asyncio.to_thread(add_memory) + tasks.append(task) + elif memory_text: + # Add as a single memory item (direct memory without conversation) + # Add tags to metadata if present + # MemMachine SDK expects tags as a string, not a list + metadata = item_meta.copy() if item_meta else {} + if tags: + # Convert list to comma-separated string + metadata["tags"] = ", ".join(tags) if isinstance(tags, list) else str(tags) + + def add_memory(): + # Use MemMachine SDK add() method + # API: memory.add(content, role="user", metadata={}, memory_types=[...]) + memory.add( + content=memory_text, + role="user", + metadata=metadata if metadata else None, + memory_types=memory_types, + episode_type=None # Use default (MESSAGE) + ) + + task = asyncio.to_thread(add_memory) + tasks.append(task) + + if tasks: + await asyncio.gather(*tasks) + + async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem]: + """ + Retrieve items relevant to the given query using the MemMachine SDK. + + Args: + query (str): The query string to match. + top_k (int): Maximum number of items to return. + kwargs: Other keyword arguments for search. + Must include 'user_id'. May include 'session_id', 'agent_id', 'project_id', 'org_id'. + + Returns: + list[MemoryItem]: The most relevant MemoryItems for the given query. + """ + user_id = kwargs.pop("user_id") # Ensure user ID is in keyword arguments + session_id = kwargs.pop("session_id", "default_session") + agent_id = kwargs.pop("agent_id", "default_agent") + group_id = kwargs.pop("group_id", "default") + project_id = kwargs.pop("project_id", None) + org_id = kwargs.pop("org_id", None) + + # Get memory instance using MemMachine SDK + memory = self._get_memory_instance( + user_id, session_id, agent_id, group_id, project_id, org_id + ) + + # Perform search using MemMachine SDK + def perform_search(): + # MemMachine SDK search() method signature: + # search(query, limit=None, filter_dict=None, timeout=None) + # Returns dict with 'episodic_memory', 'semantic_memory', 'episode_summary' + return memory.search(query=query, limit=top_k) + + search_results = await asyncio.to_thread(perform_search) + + # Construct MemoryItem instances from search results + memories = [] + + if not search_results: + return memories + + # MemMachine SDK returns a SearchResult Pydantic model with status and content fields + # The content field is a dict with episodic_memory and semantic_memory + # Extract the content dict from the SearchResult object + if hasattr(search_results, 'content'): + # SearchResult is a Pydantic model with content field + results_content = search_results.content + elif isinstance(search_results, dict): + # Fallback for dict response + results_content = search_results + else: + # Unknown format, return empty + return memories + + # episodic_memory is a dict with long_term_memory and short_term_memory + # Each contains an 'episodes' list with the actual memory episodes + # semantic_memory is a list of semantic features + episodic_memory_dict = results_content.get("episodic_memory", {}) + semantic_results = results_content.get("semantic_memory", []) + + # Extract episodes from the nested structure + # episodic_memory = { 'long_term_memory': { 'episodes': [...] }, 'short_term_memory': { 'episodes': [...] } } + episodic_results = [] + if isinstance(episodic_memory_dict, dict): + for memory_type in ['long_term_memory', 'short_term_memory']: + memory_data = episodic_memory_dict.get(memory_type, {}) + if isinstance(memory_data, dict): + episodes = memory_data.get('episodes', []) + if isinstance(episodes, list): + episodic_results.extend(episodes) + + # Process episodic memories - group by conversation if possible + # Episodes from the same conversation should be grouped together + episodic_by_conversation = {} # Key: conversation identifier, Value: list of episodes + standalone_episodic = [] # Episodes that don't belong to a conversation + + for episode in episodic_results: + if isinstance(episode, dict): + # Check if episode has role information (producer_role field) + episode_role = episode.get("producer_role") or episode.get("role") + episode_metadata = episode.get("metadata", {}) + + # Group episodes by test_id or similar identifier in metadata + # This groups episodes from the same conversation + conv_id = episode_metadata.get("test_id") or episode_metadata.get("conversation_id") + + if episode_role and conv_id: + # This is part of a conversation - group it + if conv_id not in episodic_by_conversation: + episodic_by_conversation[conv_id] = [] + episodic_by_conversation[conv_id].append(episode) + else: + # Standalone episode + standalone_episodic.append(episode) + else: + standalone_episodic.append(episode) + + # Reconstruct conversations from grouped episodes + for conv_key, conv_episodes in episodic_by_conversation.items(): + # Sort episodes by created_at timestamp if available + try: + conv_episodes.sort(key=lambda e: e.get("created_at") or e.get("timestamp") or "") + except: + pass + + # Extract conversation messages + conversation_messages = [] + memory_text = None + item_meta = {} + tags = [] + + for episode in conv_episodes: + # Get role from producer_role field + episode_role = episode.get("producer_role") or episode.get("role") or "user" + episode_content = episode.get("content") or episode.get("text") or "" + + if episode_content: + conversation_messages.append({ + "role": episode_role, + "content": episode_content + }) + + # Use first episode's metadata and tags + if not item_meta: + item_meta = episode.get("metadata", {}).copy() + + # Extract tags from metadata + if "tags" in item_meta: + tags_raw = item_meta.pop("tags", []) + if isinstance(tags_raw, str): + tags = [t.strip() for t in tags_raw.split(",") if t.strip()] + elif isinstance(tags_raw, list): + tags = tags_raw + else: + tags = [] + + # Create memory text from conversation (use first message or combine) + if conversation_messages: + memory_text = conversation_messages[0].get("content", "") + # Only set conversation if we have multiple messages + memories.append( + MemoryItem( + conversation=conversation_messages if len(conversation_messages) > 1 else None, + user_id=user_id, + memory=memory_text, + tags=tags, + metadata=item_meta + ) + ) + + # Process standalone episodic memories + for result in standalone_episodic: + memory_text = None + conversation = None + item_meta = {} + tags = [] + + if isinstance(result, dict): + memory_text = result.get("content") or result.get("text") + item_meta = result.get("metadata", {}) + + # Extract tags + if "tags" in item_meta: + tags_raw = item_meta.pop("tags", []) + if isinstance(tags_raw, str): + tags = [t.strip() for t in tags_raw.split(",") if t.strip()] + elif isinstance(tags_raw, list): + tags = tags_raw + else: + tags = [] + elif hasattr(result, 'content'): + memory_text = result.content + if hasattr(result, 'metadata'): + item_meta = result.metadata or {} + if hasattr(result, 'tags'): + tags = result.tags or [] + else: + memory_text = str(result) + + if memory_text: + memories.append( + MemoryItem( + conversation=conversation, + user_id=user_id, + memory=memory_text, + tags=tags, + metadata=item_meta + ) + ) + + # Process semantic memories + for result in semantic_results: + memory_text = None + item_meta = {} + tags = [] + + if isinstance(result, dict): + memory_text = result.get("feature") or result.get("content") or result.get("text") + item_meta = result.get("metadata", {}) + + # Extract tags + if "tags" in item_meta: + tags_raw = item_meta.pop("tags", []) + if isinstance(tags_raw, str): + tags = [t.strip() for t in tags_raw.split(",") if t.strip()] + elif isinstance(tags_raw, list): + tags = tags_raw + else: + tags = [] + elif hasattr(result, 'feature'): + memory_text = result.feature + if hasattr(result, 'metadata'): + item_meta = result.metadata or {} + else: + memory_text = str(result) + + if memory_text: + memories.append( + MemoryItem( + conversation=None, + user_id=user_id, + memory=memory_text, + tags=tags, + metadata=item_meta + ) + ) + + # Limit to top_k + return memories[:top_k] + + async def remove_items(self, **kwargs) -> None: + """ + Remove items using the MemMachine SDK. Additional parameters + needed for deletion can be specified in keyword arguments. + + Args: + kwargs (dict): Keyword arguments to pass to the remove-items method. + Should include either 'memory_id' (episodic_id or semantic_id) or 'user_id'. + May include 'session_id', 'agent_id', 'group_id', 'project_id', 'org_id'. + For memory_id deletion, may include 'memory_type' ('episodic' or 'semantic'). + """ + if "memory_id" in kwargs: + memory_id = kwargs.pop("memory_id") + memory_type = kwargs.pop("memory_type", "episodic") # Default to episodic + user_id = kwargs.pop("user_id", None) + session_id = kwargs.pop("session_id", "default_session") + agent_id = kwargs.pop("agent_id", "default_agent") + group_id = kwargs.pop("group_id", "default") + project_id = kwargs.pop("project_id", None) + org_id = kwargs.pop("org_id", None) + + if not user_id: + raise ValueError( + "user_id is required when deleting by memory_id. " + "A memory instance is needed to perform deletion, which requires user_id." + ) + + def delete_memory(): + memory = self._get_memory_instance( + user_id, session_id, agent_id, group_id, project_id, org_id + ) + # Use MemMachine SDK to delete specific memory + # API: memory.delete_episodic(episodic_id) or memory.delete_semantic(semantic_id) + if memory_type.lower() == "semantic": + memory.delete_semantic(semantic_id=memory_id) + else: + memory.delete_episodic(episodic_id=memory_id) + + await asyncio.to_thread(delete_memory) + + elif "user_id" in kwargs: + user_id = kwargs.pop("user_id") + session_id = kwargs.pop("session_id", "default_session") + agent_id = kwargs.pop("agent_id", "default_agent") + group_id = kwargs.pop("group_id", "default") + 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 method + # We would need to search for all memories and delete them individually + # For now, we'll raise a NotImplementedError with guidance + raise NotImplementedError( + "Bulk deletion by user_id is not directly supported by MemMachine SDK. " + "To delete all memories for a user, you would need to: " + "1. Search for all memories with that user_id " + "2. Extract memory IDs from results " + "3. Delete each memory individually using delete_episodic() or delete_semantic(). " + "Alternatively, delete specific memories using memory_id parameter." + ) \ No newline at end of file diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py new file mode 100644 index 0000000000..17e1db4a21 --- /dev/null +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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 +from nat.data_models.retry_mixin import RetryMixin +from nat.utils.exception_handlers.automatic_retries import patch_with_retry + + +class MemMachineMemoryClientConfig(MemoryBaseConfig, RetryMixin, name="memmachine_memory"): + """ + 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. + """ + base_url: str # Base URL of the MemMachine server (e.g., "http://localhost:8080") + org_id: str | None = None # Optional default organization ID + project_id: str | None = None # Optional default project ID + timeout: int = 30 # Request timeout in seconds + max_retries: int = 3 # Maximum number of retries for failed requests + + +@register_memory(config_type=MemMachineMemoryClientConfig) +async def memmachine_memory_client(config: MemMachineMemoryClientConfig, builder: Builder): + from .memmachine_editor import MemMachineEditor + # Import and initialize the MemMachine Python SDK + try: + from memmachine import MemMachineClient + except ImportError as e: + raise ImportError( + f"Could not import MemMachineClient from memmachine package. " + f"Error: {e}. " + "Please ensure memmachine package is installed: pip install memmachine. " + "See https://github.com/MemMachine/MemMachine/blob/main/docs/examples/python.mdx " + "for installation instructions." + ) from e + + # Initialize MemMachineClient with base_url + # This follows the documented SDK pattern for local instances: + # client = MemMachineClient(base_url="http://localhost:8080") + # Note: api_key is not needed for local/self-hosted MemMachine instances + try: + client = MemMachineClient( + base_url=config.base_url, + timeout=config.timeout, + max_retries=config.max_retries + ) + except Exception as e: + raise RuntimeError( + f"Failed to initialize MemMachineClient with base_url '{config.base_url}'. " + f"Error: {e}. " + "Please ensure the MemMachine server is running and the base_url is correct." + ) from e + + # If default org_id and project_id are provided, create/get the project + # Otherwise, the editor will create projects as needed + memmachine_instance = client + if config.org_id and config.project_id: + try: + # Use get_or_create_project to handle existing projects gracefully + project = client.get_or_create_project( + org_id=config.org_id, + project_id=config.project_id, + description=f"NeMo Agent toolkit project: {config.project_id}" + ) + memmachine_instance = project + except Exception as e: + # If project creation fails, fall back to using the client directly + # The editor will handle project creation on-demand + pass + + memory_editor = MemMachineEditor(memmachine_instance=memmachine_instance) + + 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 + ) + + yield memory_editor \ No newline at end of file diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py new file mode 100644 index 0000000000..9ce1b2c3e5 --- /dev/null +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py @@ -0,0 +1 @@ +from . import memory \ No newline at end of file diff --git a/packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md b/packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md new file mode 100644 index 0000000000..7cb0614bce --- /dev/null +++ b/packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md @@ -0,0 +1,129 @@ +# API Call Verification Tests + +## Overview + +The tests in `test_memmachine_api_calls.py` verify that the MemMachine integration makes **correct API calls** to the MemMachine SDK. These tests use spies to capture and validate: + +1. **Exact SDK methods called** - Verifies the right methods are invoked +2. **Parameter correctness** - Verifies parameters match SDK expectations +3. **Data transformations** - Verifies NAT format ā†’ MemMachine format conversion +4. **Memory type handling** - Verifies all memories are added to both episodic and semantic types + +## What Gets Tested + +### 1. Add Operations (`add_items`) + +#### Conversation Handling +- āœ… Each message in conversation calls `memory.add()` separately +- āœ… Messages preserve their roles (user/assistant/system) +- āœ… Messages are added in the correct order +- āœ… All memories are added to both episodic and semantic memory types + +#### Direct Memory (No Conversation) +- āœ… When `conversation=None`, uses `memory` field as content +- āœ… All memories are added to both episodic and semantic memory types +- āœ… Default `role` is "user" + +#### Metadata Handling +- āœ… Tags are included in metadata dict +- āœ… Custom metadata fields are preserved +- āœ… Special fields (session_id, agent_id, project_id, org_id) are extracted and NOT passed to `add()` +- āœ… Empty metadata becomes `None`, not empty dict + +#### Project/Org Handling +- āœ… Custom `project_id`/`org_id` in metadata triggers `create_project()` call +- āœ… `project.memory()` is called with correct user_id, session_id, agent_id, group_id + +### 2. Search Operations (`search`) + +#### Parameter Mapping +- āœ… NAT's `top_k` parameter is converted to SDK's `limit` parameter +- āœ… `query` parameter is passed correctly +- āœ… `project.memory()` is called with user_id, session_id, agent_id, group_id + +#### Custom Project/Org +- āœ… Custom `project_id`/`org_id` triggers `create_project()` call + +### 3. Remove Operations (`remove_items`) + +#### Episodic Deletion +- āœ… Calls `memory.delete_episodic(episodic_id=...)` with correct ID + +#### Semantic Deletion +- āœ… Calls `memory.delete_semantic(semantic_id=...)` with correct ID + +### 4. API Call Format + +#### Keyword Arguments +- āœ… All SDK methods are called with keyword arguments, not positional +- āœ… Parameter names match SDK exactly (`limit` not `top_k`, `episodic_id` not `memory_id`) + +## Test Structure + +Each test class focuses on a specific aspect: + +- **`TestAddItemsAPICalls`** - Verifies `add_items()` makes correct `memory.add()` calls +- **`TestSearchAPICalls`** - Verifies `search()` makes correct `memory.search()` calls +- **`TestRemoveItemsAPICalls`** - Verifies `remove_items()` makes correct delete calls +- **`TestAPICallParameterValidation`** - Verifies parameter names and formats +- **`TestDataTransformation`** - Verifies data transformations are correct + +## Running the Tests + +```bash +# Run all API call verification tests +pytest tests/test_memmachine_api_calls.py -v + +# Run a specific test class +pytest tests/test_memmachine_api_calls.py::TestAddItemsAPICalls -v + +# Run a specific test +pytest tests/test_memmachine_api_calls.py::TestAddItemsAPICalls::test_add_conversation_calls_add_with_correct_parameters -v +``` + +## Key Differences from Unit Tests + +| Aspect | Unit Tests (`test_memmachine_editor.py`) | API Call Tests (`test_memmachine_api_calls.py`) | +|--------|------------------------------------------|--------------------------------------------------| +| **Focus** | Integration logic, error handling | Exact API calls and parameters | +| **Verification** | Mocks return values | Spies capture actual calls | +| **What's Tested** | Code flow, edge cases | Parameter correctness, data transformation | +| **SDK Methods** | Mocked, not verified | Spied, parameters validated | + +## Example: What Gets Verified + +### Adding a Conversation + +**Input (NAT format):** +```python +MemoryItem( + conversation=[ + {"role": "user", "content": "I like pizza"}, + {"role": "assistant", "content": "Great!"} + ], + user_id="user123", + metadata={"session_id": "s1", "agent_id": "a1"}, + tags=["food"] +) +``` + +**Verified API Calls:** +1. āœ… `create_project(org_id="default-org", project_id="default-project", ...)` +2. āœ… `project.memory(user_id="user123", session_id="s1", agent_id="a1", group_id="default")` +3. āœ… `memory.add(content="I like pizza", role="user", memory_types=[Episodic, Semantic], metadata={"tags": ["food"]})` +4. āœ… `memory.add(content="Great!", role="assistant", memory_types=[Episodic, Semantic], metadata={"tags": ["food"]})` + +**What's Verified:** +- āœ… Two `add()` calls (one per message) +- āœ… Correct roles preserved +- āœ… All memories added to both episodic and semantic types +- āœ… Tags in metadata +- āœ… session_id/agent_id NOT in metadata (used for memory instance instead) +- āœ… All parameters are keyword arguments + +## Integration with Real Server + +For tests that verify actual API calls work with a real MemMachine server, see: +- `test_memmachine_integration.py` - Tests with real server +- `TESTING.md` - Guide for integration testing + diff --git a/packages/nvidia_nat_memmachine/tests/__init__.py b/packages/nvidia_nat_memmachine/tests/__init__.py new file mode 100644 index 0000000000..a1744724ec --- /dev/null +++ b/packages/nvidia_nat_memmachine/tests/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py new file mode 100644 index 0000000000..0662e3e40d --- /dev/null +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python3 +""" +Simple script to test adding memories and retrieving them. + +This script demonstrates the full integration: +1. Adds memories using the NAT integration +2. Retrieves them back +3. Verifies they match + +Usage: + python tests/test_add_and_retrieve.py + or + pytest tests/test_add_and_retrieve.py +""" + +import asyncio +import os +import uuid +from datetime import datetime + +from nat.builder.builder import Builder +from nat.memory.models import MemoryItem +from nat.plugins.memmachine.memory import MemMachineMemoryClientConfig + + +async def test_add_and_retrieve(): + """Test adding memories and retrieving them.""" + # Configuration + base_url = os.environ.get("MEMMACHINE_BASE_URL", "http://localhost:8080") + test_id = str(uuid.uuid4())[:8] + + config = MemMachineMemoryClientConfig( + base_url=base_url, + org_id=f"test_org_{test_id}", + project_id=f"test_project_{test_id}", + ) + + user_id = f"test_user_{uuid.uuid4().hex[:8]}" + session_id = "test_session" + agent_id = "test_agent" + + print("=" * 80) + print("MemMachine Integration Test: Add and Retrieve Memories") + print("=" * 80) + print(f"Base URL: {base_url}") + print(f"User ID: {user_id}") + print(f"Org ID: {config.org_id}") + print(f"Project ID: {config.project_id}") + print() + + builder = Builder() + + try: + async with builder: + async with builder.get_memory_client("memmachine_memory", config) as memory_client: + print("āœ“ Memory client initialized\n") + + # Test 1: Add conversation memory + print("Test 1: Adding conversation memory...") + conversation_memory = MemoryItem( + conversation=[ + {"role": "user", "content": "I love pizza and Italian food."}, + {"role": "assistant", "content": "I'll remember that you love pizza and Italian food."}, + ], + user_id=user_id, + memory="User loves pizza", + metadata={ + "session_id": session_id, + "agent_id": agent_id, + "test_timestamp": datetime.now().isoformat() + }, + tags=["food", "preference", "italian"] + ) + + await memory_client.add_items([conversation_memory]) + print("āœ“ Conversation memory added") + + # Wait a moment for indexing + await asyncio.sleep(2) + + # Retrieve it + print("\nRetrieving conversation memory...") + retrieved = await memory_client.search( + query="pizza Italian food", + top_k=10, + user_id=user_id, + session_id=session_id, + agent_id=agent_id + ) + + print(f"āœ“ Retrieved {len(retrieved)} memories") + if retrieved: + print(f" First memory: {retrieved[0].memory or str(retrieved[0].conversation)}") + print(f" Tags: {retrieved[0].tags}") + + # Test 2: Add direct memory (no conversation) + # All memories are added to both episodic and semantic memory types + print("\n" + "-" * 80) + print("Test 2: Adding direct memory...") + direct_memory = MemoryItem( + conversation=None, + user_id=user_id, + memory="User prefers working in the morning and is allergic to peanuts", + metadata={ + "session_id": session_id, + "agent_id": agent_id, + "test_timestamp": datetime.now().isoformat() + }, + tags=["preference", "allergy", "schedule"] + ) + + await memory_client.add_items([direct_memory]) + print("āœ“ Direct memory added") + + # Wait for indexing + await asyncio.sleep(2) + + # Retrieve it + print("\nRetrieving direct memory...") + retrieved = await memory_client.search( + query="morning work allergy peanuts", + top_k=10, + user_id=user_id, + session_id=session_id, + agent_id=agent_id + ) + + print(f"āœ“ Retrieved {len(retrieved)} memories") + if retrieved: + for i, mem in enumerate(retrieved[:3], 1): + print(f" Memory {i}: {mem.memory}") + print(f" Tags: {mem.tags}") + + # Test 3: Add multiple memories and retrieve all + print("\n" + "-" * 80) + print("Test 3: Adding multiple memories...") + multiple_memories = [ + MemoryItem( + conversation=[{"role": "user", "content": f"Fact {i}: I like item {i}"}], + user_id=user_id, + memory=f"Fact {i}", + metadata={ + "session_id": session_id, + "agent_id": agent_id, + "fact_number": i + }, + tags=[f"fact_{i}"] + ) + for i in range(1, 4) # Add 3 memories + ] + + await memory_client.add_items(multiple_memories) + print("āœ“ Added 3 memories") + + # Wait for indexing + await asyncio.sleep(2) + + # Retrieve all with broad query + print("\nRetrieving all memories (broad search)...") + all_memories = await memory_client.search( + query="*", # Broad query + top_k=20, + user_id=user_id, + session_id=session_id, + agent_id=agent_id + ) + + print(f"āœ“ Retrieved {len(all_memories)} total memories") + print("\nAll memories:") + for i, mem in enumerate(all_memories, 1): + content = mem.memory or (str(mem.conversation) if mem.conversation else "N/A") + print(f" {i}. {content[:60]}...") + print(f" Tags: {mem.tags}") + + print("\n" + "=" * 80) + print("āœ“ All tests completed successfully!") + print("=" * 80) + + except Exception as e: + print(f"\nāœ— Error: {e}") + import traceback + traceback.print_exc() + raise + + +if __name__ == "__main__": + asyncio.run(test_add_and_retrieve()) diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py new file mode 100644 index 0000000000..c8e1b56614 --- /dev/null +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py @@ -0,0 +1,619 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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. + +""" +Tests that verify actual MemMachine SDK API calls are made correctly. + +These tests use spies/wrappers to capture and verify: +1. The exact SDK methods called +2. The parameters passed to each method +3. The data transformations (NAT MemoryItem ā†’ MemMachine format) +4. That all memories are added to both episodic and semantic memory types +""" + +from unittest.mock import Mock, patch +from typing import Any + +import pytest + +from nat.memory.models import MemoryItem +from nat.plugins.memmachine.memmachine_editor import MemMachineEditor + + +class APICallSpy: + """Spy class to capture and verify actual SDK API calls.""" + + def __init__(self): + self.calls = [] + self.return_values = {} + + def record_call(self, method_name: str, args: tuple, kwargs: dict): + """Record an API call.""" + self.calls.append({ + 'method': method_name, + 'args': args, + 'kwargs': kwargs + }) + + def get_calls(self, method_name: str = None): + """Get all calls, optionally filtered by method name.""" + if method_name: + return [c for c in self.calls if c['method'] == method_name] + return self.calls + + def assert_called_with(self, method_name: str, **expected_kwargs): + """Assert a method was called with specific parameters.""" + calls = self.get_calls(method_name) + assert len(calls) > 0, f"Expected {method_name} to be called, but it wasn't" + + for call in calls: + call_kwargs = call['kwargs'] + # Check if all expected kwargs match + matches = all( + call_kwargs.get(key) == value + for key, value in expected_kwargs.items() + ) + if matches: + return call + + raise AssertionError( + f"Expected {method_name} to be called with {expected_kwargs}, " + f"but got calls: {[c['kwargs'] for c in calls]}" + ) + + +@pytest.fixture(name="api_spy") +def api_spy_fixture(): + """Fixture to provide an API call spy.""" + return APICallSpy() + + +@pytest.fixture(name="spied_memory_instance") +def spied_memory_instance_fixture(api_spy: APICallSpy): + """Create a memory instance with spied methods.""" + mock_memory = Mock() + + # Wrap the add method to spy on calls + original_add = Mock(return_value=True) + def spied_add(*args, **kwargs): + api_spy.record_call('add', args, kwargs) + return original_add(*args, **kwargs) + mock_memory.add = spied_add + + # Wrap the search method + original_search = Mock(return_value={ + "episodic_memory": [], + "semantic_memory": [], + "episode_summary": [] + }) + def spied_search(*args, **kwargs): + api_spy.record_call('search', args, kwargs) + return original_search(*args, **kwargs) + mock_memory.search = spied_search + + # Wrap delete methods + original_delete_episodic = Mock(return_value=True) + def spied_delete_episodic(*args, **kwargs): + api_spy.record_call('delete_episodic', args, kwargs) + return original_delete_episodic(*args, **kwargs) + mock_memory.delete_episodic = spied_delete_episodic + + original_delete_semantic = Mock(return_value=True) + def spied_delete_semantic(*args, **kwargs): + api_spy.record_call('delete_semantic', args, kwargs) + return original_delete_semantic(*args, **kwargs) + mock_memory.delete_semantic = spied_delete_semantic + + return mock_memory + + +@pytest.fixture(name="spied_project") +def spied_project_fixture(spied_memory_instance: Mock, api_spy: APICallSpy): + """Create a project instance with spied memory() method.""" + mock_project = Mock(spec=['memory', 'org_id', 'project_id']) + + def spied_memory(*args, **kwargs): + api_spy.record_call('project.memory', args, kwargs) + return spied_memory_instance + + mock_project.memory = spied_memory + mock_project.org_id = "test_org" + mock_project.project_id = "test_project" + return mock_project + + +@pytest.fixture(name="spied_client") +def spied_client_fixture(spied_project: Mock, api_spy: APICallSpy): + """Create a client instance with spied create_project and get_or_create_project methods.""" + mock_client = Mock(spec=['create_project', 'get_or_create_project', 'base_url']) + + def spied_create_project(*args, **kwargs): + api_spy.record_call('create_project', args, kwargs) + return spied_project + + def spied_get_or_create_project(*args, **kwargs): + api_spy.record_call('get_or_create_project', args, kwargs) + return spied_project + + mock_client.create_project = spied_create_project + mock_client.get_or_create_project = spied_get_or_create_project + mock_client.base_url = "http://localhost:8080" + return mock_client + + +@pytest.fixture(name="editor_with_spy") +def editor_with_spy_fixture(spied_client: Mock): + """Create an editor with spied SDK calls.""" + return MemMachineEditor(memmachine_instance=spied_client) + + +class TestAddItemsAPICalls: + """Test that add_items makes correct API calls to MemMachine SDK.""" + + async def test_add_conversation_calls_add_with_correct_parameters( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that adding a conversation calls memory.add() with correct parameters.""" + item = MemoryItem( + conversation=[ + {"role": "user", "content": "I like pizza"}, + {"role": "assistant", "content": "Great! What's your favorite topping?"} + ], + user_id="user123", + memory="User likes pizza", + metadata={"session_id": "session1", "agent_id": "agent1"}, + tags=["food", "preference"] + ) + + await editor_with_spy.add_items([item]) + + # Verify project.memory was called with correct parameters + api_spy.assert_called_with( + 'project.memory', + user_id="user123", + session_id="session1", + agent_id="agent1", + group_id="default" + ) + + # Verify add was called twice (once per message) + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 2, f"Expected 2 add calls, got {len(add_calls)}" + + # Verify first call (user message) - episodic by default + user_call = next( + (c for c in add_calls if c['kwargs'].get('role') == 'user'), + None + ) + assert user_call is not None, "Should have a call with role='user'" + assert user_call['kwargs']['content'] == "I like pizza" + assert user_call['kwargs']['role'] == "user" + # Now uses memory_types instead of episode_type + assert user_call['kwargs']['episode_type'] is None + assert 'memory_types' in user_call['kwargs'] + assert 'tags' in user_call['kwargs'].get('metadata', {}) + # MemMachine SDK expects tags as comma-separated string + assert user_call['kwargs']['metadata']['tags'] == "food, preference" + + # Verify second call (assistant message) + assistant_call = next( + (c for c in add_calls if c['kwargs'].get('role') == 'assistant'), + None + ) + assert assistant_call is not None, "Should have a call with role='assistant'" + assert assistant_call['kwargs']['content'] == "Great! What's your favorite topping?" + assert assistant_call['kwargs']['role'] == "assistant" + # Now uses memory_types instead of episode_type + assert assistant_call['kwargs']['episode_type'] is None + assert 'memory_types' in assistant_call['kwargs'] + + async def test_add_direct_memory_calls_add_with_both_types( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that direct memory (no conversation) calls add() with both memory types.""" + item = MemoryItem( + conversation=None, + user_id="user123", + memory="User prefers working in the morning", + metadata={ + "session_id": "session1", + "agent_id": "agent1" + }, + tags=["preference"] + ) + + await editor_with_spy.add_items([item]) + + # Verify add was called with both memory types + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 1 + assert add_calls[0]['kwargs']['content'] == "User prefers working in the morning" + assert add_calls[0]['kwargs']['role'] == "user" + assert add_calls[0]['kwargs']['episode_type'] is None + # Verify memory_types contains both Episodic and Semantic + memory_types = add_calls[0]['kwargs']['memory_types'] + assert len(memory_types) == 2, "Should have both episodic and semantic memory types" + + # Verify metadata includes tags (as comma-separated string) + assert add_calls[0]['kwargs']['metadata']['tags'] == "preference" + + async def test_add_conversation_memory_calls_add_with_both_types( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that conversation memory calls add() with both memory types.""" + item = MemoryItem( + conversation=[{"role": "user", "content": "Hello"}], + user_id="user123", + memory="Test", + metadata={ + "session_id": "session1", + "agent_id": "agent1" + }, + tags=[] + ) + + await editor_with_spy.add_items([item]) + + # Verify add was called with both memory types + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 1 + assert add_calls[0]['kwargs']['content'] == "Hello" + assert add_calls[0]['kwargs']['role'] == "user" + assert add_calls[0]['kwargs']['episode_type'] is None + # Verify memory_types contains both Episodic and Semantic + memory_types = add_calls[0]['kwargs']['memory_types'] + assert len(memory_types) == 2, "Should have both episodic and semantic memory types" + + async def test_add_with_custom_project_org_calls_get_or_create_project( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that custom project_id/org_id triggers get_or_create_project call.""" + item = MemoryItem( + conversation=[{"role": "user", "content": "Test"}], + user_id="user123", + memory="Test", + metadata={ + "session_id": "session1", + "agent_id": "agent1", + "project_id": "custom_project", + "org_id": "custom_org" + } + ) + + await editor_with_spy.add_items([item]) + + # Verify get_or_create_project was called with custom IDs + api_spy.assert_called_with( + 'get_or_create_project', + org_id="custom_org", + project_id="custom_project", + description="Project for user123" + ) + + async def test_add_preserves_metadata_except_special_fields( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that metadata is preserved except for special fields like session_id.""" + item = MemoryItem( + conversation=[{"role": "user", "content": "Test"}], + user_id="user123", + memory="Test", + metadata={ + "session_id": "session1", + "agent_id": "agent1", + "custom_field": "custom_value", + "another_field": 123 + }, + tags=["tag1"] + ) + + await editor_with_spy.add_items([item]) + + # Verify metadata in the API call + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 1 + + metadata = add_calls[0]['kwargs'].get('metadata', {}) + # Special fields should be removed (used for memory instance creation) + assert 'session_id' not in metadata, "session_id should be removed from metadata" + assert 'agent_id' not in metadata, "agent_id should be removed from metadata" + # Custom fields should be preserved + assert metadata['custom_field'] == "custom_value" + assert metadata['another_field'] == 123 + # MemMachine SDK expects tags as comma-separated string + assert metadata['tags'] == "tag1" + + +class TestSearchAPICalls: + """Test that search makes correct API calls to MemMachine SDK.""" + + async def test_search_calls_memory_search_with_correct_parameters( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy, + spied_memory_instance: Mock + ): + """Verify that search calls memory.search() with correct parameters.""" + # Set up search return value + spied_memory_instance.search.return_value = { + "episodic_memory": [{"content": "I like pizza", "metadata": {}}], + "semantic_memory": [], + "episode_summary": [] + } + + results = await editor_with_spy.search( + query="What do I like?", + top_k=10, + user_id="user123", + session_id="session1", + agent_id="agent1" + ) + + # Verify project.memory was called + api_spy.assert_called_with( + 'project.memory', + user_id="user123", + session_id="session1", + agent_id="agent1", + group_id="default" + ) + + # Verify search was called with correct parameters + api_spy.assert_called_with( + 'search', + query="What do I like?", + limit=10 + ) + + async def test_search_with_custom_project_org( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy, + spied_memory_instance: Mock + ): + """Verify search with custom project/org calls get_or_create_project.""" + spied_memory_instance.search.return_value = { + "episodic_memory": [], + "semantic_memory": [], + "episode_summary": [] + } + + await editor_with_spy.search( + query="test", + user_id="user123", + project_id="custom_project", + org_id="custom_org" + ) + + # Verify get_or_create_project was called + api_spy.assert_called_with( + 'get_or_create_project', + org_id="custom_org", + project_id="custom_project", + description="Project for user123" + ) + + +class TestRemoveItemsAPICalls: + """Test that remove_items makes correct API calls to MemMachine SDK.""" + + async def test_remove_episodic_calls_delete_episodic( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that removing episodic memory calls delete_episodic().""" + await editor_with_spy.remove_items( + memory_id="episodic_123", + memory_type="episodic", + user_id="user123", + session_id="session1", + agent_id="agent1" + ) + + # Verify delete_episodic was called with correct ID + api_spy.assert_called_with( + 'delete_episodic', + episodic_id="episodic_123" + ) + + async def test_remove_semantic_calls_delete_semantic( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that removing semantic memory calls delete_semantic().""" + await editor_with_spy.remove_items( + memory_id="semantic_456", + memory_type="semantic", + user_id="user123", + session_id="session1", + agent_id="agent1" + ) + + # Verify delete_semantic was called with correct ID + api_spy.assert_called_with( + 'delete_semantic', + semantic_id="semantic_456" + ) + + +class TestAPICallParameterValidation: + """Test that API calls use correct parameter names and formats.""" + + async def test_add_uses_keyword_arguments_not_positional( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that add() is called with keyword arguments, not positional.""" + item = MemoryItem( + conversation=[{"role": "user", "content": "Test"}], + user_id="user123", + memory="Test", + metadata={"session_id": "session1", "agent_id": "agent1"} + ) + + await editor_with_spy.add_items([item]) + + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 1 + + # Verify it was called with kwargs, not positional args + call = add_calls[0] + assert len(call['args']) == 0, "add() should be called with keyword arguments only" + assert 'content' in call['kwargs'] + assert 'role' in call['kwargs'] + assert 'episode_type' in call['kwargs'] + + async def test_search_uses_limit_not_top_k( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy, + spied_memory_instance: Mock + ): + """Verify that search() uses 'limit' parameter (SDK name), not 'top_k'.""" + spied_memory_instance.search.return_value = { + "episodic_memory": [], + "semantic_memory": [], + "episode_summary": [] + } + + await editor_with_spy.search( + query="test", + top_k=5, # NAT uses top_k + user_id="user123" + ) + + # Verify search was called with 'limit' (SDK parameter name) + search_calls = api_spy.get_calls('search') + assert len(search_calls) == 1 + assert 'limit' in search_calls[0]['kwargs'] + assert search_calls[0]['kwargs']['limit'] == 5 + assert 'top_k' not in search_calls[0]['kwargs'], "SDK uses 'limit', not 'top_k'" + + async def test_metadata_is_dict_or_none_not_empty_dict( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that metadata is passed as dict or None, never empty dict.""" + item = MemoryItem( + conversation=[{"role": "user", "content": "Test"}], + user_id="user123", + memory="Test", + metadata={"session_id": "session1", "agent_id": "agent1"}, + tags=[] # No tags + ) + + await editor_with_spy.add_items([item]) + + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 1 + + metadata = add_calls[0]['kwargs'].get('metadata') + # Should be None if empty, or a dict with content + assert metadata is None or isinstance(metadata, dict) + if metadata is not None: + assert len(metadata) > 0, "Metadata should not be empty dict, use None instead" + + +class TestDataTransformation: + """Test that data is correctly transformed between NAT and MemMachine formats.""" + + async def test_conversation_messages_preserved_in_order( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that conversation messages are added in the correct order.""" + item = MemoryItem( + conversation=[ + {"role": "user", "content": "First message"}, + {"role": "assistant", "content": "Second message"}, + {"role": "user", "content": "Third message"} + ], + user_id="user123", + memory="Test", + metadata={"session_id": "session1", "agent_id": "agent1"} + ) + + await editor_with_spy.add_items([item]) + + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 3 + + # Verify order and content + assert add_calls[0]['kwargs']['content'] == "First message" + assert add_calls[0]['kwargs']['role'] == "user" + assert add_calls[1]['kwargs']['content'] == "Second message" + assert add_calls[1]['kwargs']['role'] == "assistant" + assert add_calls[2]['kwargs']['content'] == "Third message" + assert add_calls[2]['kwargs']['role'] == "user" + + async def test_tags_included_in_metadata( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that tags are included in the metadata dict.""" + item = MemoryItem( + conversation=[{"role": "user", "content": "Test"}], + user_id="user123", + memory="Test", + metadata={"session_id": "session1", "agent_id": "agent1"}, + tags=["tag1", "tag2", "tag3"] + ) + + await editor_with_spy.add_items([item]) + + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 1 + + metadata = add_calls[0]['kwargs'].get('metadata', {}) + assert 'tags' in metadata + # MemMachine SDK expects tags as a comma-separated string + assert metadata['tags'] == "tag1, tag2, tag3" + + async def test_empty_conversation_uses_memory_text( + self, + editor_with_spy: MemMachineEditor, + api_spy: APICallSpy + ): + """Verify that when conversation is None, memory text is used.""" + item = MemoryItem( + conversation=None, + user_id="user123", + memory="This is the memory text", + metadata={"session_id": "session1", "agent_id": "agent1"} + ) + + await editor_with_spy.add_items([item]) + + add_calls = api_spy.get_calls('add') + assert len(add_calls) == 1 + assert add_calls[0]['kwargs']['content'] == "This is the memory text" + assert add_calls[0]['kwargs']['role'] == "user" # Default role + diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py new file mode 100644 index 0000000000..2ef531bad1 --- /dev/null +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py @@ -0,0 +1,511 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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, Mock + +import pytest + +from nat.memory.models import MemoryItem +from nat.plugins.memmachine.memmachine_editor import MemMachineEditor + + +@pytest.fixture(name="mock_memory_instance") +def mock_memory_instance_fixture(): + """Fixture to provide a mocked Memory instance from MemMachine SDK.""" + mock_memory = Mock() + mock_memory.add = Mock(return_value=True) + mock_memory.search = Mock(return_value={ + "episodic_memory": [], + "semantic_memory": [], + "episode_summary": [] + }) + mock_memory.delete_episodic = Mock(return_value=True) + mock_memory.delete_semantic = Mock(return_value=True) + return mock_memory + + +@pytest.fixture(name="mock_project") +def mock_project_fixture(mock_memory_instance): + """Fixture to provide a mocked Project instance from MemMachine SDK.""" + # Use spec to restrict attributes - Project should have 'memory' but NOT 'create_project' + # This ensures hasattr checks work correctly + mock_project = Mock(spec=['memory', 'org_id', 'project_id']) + mock_project.memory = Mock(return_value=mock_memory_instance) + mock_project.org_id = "test_org" + mock_project.project_id = "test_project" + # Explicitly ensure create_project doesn't exist (Mock with spec will raise AttributeError) + return mock_project + + +@pytest.fixture(name="mock_client") +def mock_client_fixture(mock_project): + """Fixture to provide a mocked MemMachineClient instance.""" + # Use spec to ensure create_project and get_or_create_project exist for hasattr checks + mock_client = Mock(spec=['create_project', 'get_or_create_project', 'base_url']) + mock_client.create_project = Mock(return_value=mock_project) + mock_client.get_or_create_project = Mock(return_value=mock_project) + mock_client.base_url = "http://localhost:8080" + return mock_client + + +@pytest.fixture(name="memmachine_editor_with_client") +def memmachine_editor_with_client_fixture(mock_client): + """Fixture to provide an instance of MemMachineEditor with a mocked client.""" + return MemMachineEditor(memmachine_instance=mock_client) + + +@pytest.fixture(name="memmachine_editor_with_project") +def memmachine_editor_with_project_fixture(mock_project): + """Fixture to provide an instance of MemMachineEditor with a mocked project.""" + return MemMachineEditor(memmachine_instance=mock_project) + + +@pytest.fixture(name="sample_memory_item") +def sample_memory_item_fixture(): + """Fixture to provide a sample MemoryItem.""" + conversation = [ + { + "role": "user", + "content": "Hi, I'm Alex. I'm a vegetarian and I'm allergic to nuts.", + }, + { + "role": "assistant", + "content": "Hello Alex! I've noted that you're a vegetarian and have a nut allergy.", + }, + ] + + return MemoryItem( + conversation=conversation, + user_id="user123", + memory="Sample memory", + metadata={"key1": "value1", "session_id": "session456", "agent_id": "agent789"}, + tags=["tag1", "tag2"] + ) + + +@pytest.fixture(name="sample_direct_memory_item") +def sample_direct_memory_item_fixture(): + """Fixture to provide a MemoryItem for direct memory (no conversation). + + Direct memories are added to both episodic and semantic memory types. + """ + return MemoryItem( + conversation=None, + user_id="user123", + memory="I prefer working in the morning", + metadata={"session_id": "session456", "agent_id": "agent789"}, + tags=["preference"] + ) + + +async def test_add_items_with_conversation( + memmachine_editor_with_client: MemMachineEditor, + mock_client: Mock, + mock_project: Mock, + mock_memory_instance: Mock, + sample_memory_item: MemoryItem +): + """Test adding MemoryItem objects with conversation successfully.""" + items = [sample_memory_item] + await memmachine_editor_with_client.add_items(items) + + # Verify project was created/retrieved + mock_client.get_or_create_project.assert_called_once() + + # Verify memory instance was created + mock_project.memory.assert_called_once_with( + user_id="user123", + agent_id="agent789", + session_id="session456", + group_id="default" + ) + + # Verify add was called for each message in conversation + # The await above should have completed all async tasks + assert mock_memory_instance.add.call_count == 2, f"Expected 2 calls, got {mock_memory_instance.add.call_count}. Calls: {mock_memory_instance.add.call_args_list}" + + # Get all calls + all_calls = mock_memory_instance.add.call_args_list + assert len(all_calls) == 2, f"Expected 2 calls in call_args_list, got {len(all_calls)}" + + # Extract roles and contents from all calls + calls_data = [] + for call in all_calls: + if call.kwargs: + role = call.kwargs.get("role") + content = call.kwargs.get("content") + episode_type = call.kwargs.get("episode_type") + metadata = call.kwargs.get("metadata", {}) + else: + # Handle positional args if needed + role = None + content = call.args[0] if call.args else None + episode_type = None + metadata = {} + if role and content: + calls_data.append({"role": role, "content": content, "episode_type": episode_type, "metadata": metadata}) + + # Verify we have both roles + roles = [c["role"] for c in calls_data] + assert "user" in roles, f"Expected 'user' role in calls, got: {roles}. Calls data: {calls_data}" + assert "assistant" in roles, f"Expected 'assistant' role in calls, got: {roles}. Calls data: {calls_data}" + + # Verify user message + user_call_data = next(c for c in calls_data if c["role"] == "user") + assert user_call_data["content"] == "Hi, I'm Alex. I'm a vegetarian and I'm allergic to nuts." + # Now uses memory_types instead of episode_type + assert user_call_data["episode_type"] is None + assert "tags" in user_call_data["metadata"] + + # Verify assistant message + assistant_call_data = next(c for c in calls_data if c["role"] == "assistant") + assert assistant_call_data["content"] == "Hello Alex! I've noted that you're a vegetarian and have a nut allergy." + # Now uses memory_types instead of episode_type + assert assistant_call_data["episode_type"] is None + + +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 +): + """Test adding MemoryItem for direct memory (no conversation). + + Direct memories are added to both episodic and semantic memory types. + """ + items = [sample_direct_memory_item] + await memmachine_editor_with_client.add_items(items) + + # Verify add was called + assert mock_memory_instance.add.call_count == 1 + + # Verify memory_types is used (episode_type is None) + call_kwargs = mock_memory_instance.add.call_args.kwargs + assert call_kwargs["content"] == "I prefer working in the morning" + assert call_kwargs["episode_type"] is None + assert "memory_types" in call_kwargs + # Verify both memory types are included + assert len(call_kwargs["memory_types"]) == 2, "Should use both episodic and semantic memory types" + assert call_kwargs["role"] == "user" + + +async def test_add_items_empty_list( + memmachine_editor_with_client: MemMachineEditor, + mock_memory_instance: Mock +): + """Test adding an empty list of MemoryItem objects.""" + await memmachine_editor_with_client.add_items([]) + + # Should not call add if list is empty + mock_memory_instance.add.assert_not_called() + + +async def test_add_items_with_memory_text_only( + memmachine_editor_with_client: MemMachineEditor, + mock_client: Mock, + mock_project: Mock, + mock_memory_instance: Mock +): + """Test adding MemoryItem with only memory text (no conversation).""" + item = MemoryItem( + conversation=None, + user_id="user123", + memory="This is a standalone memory", + metadata={"session_id": "session456", "agent_id": "agent789"}, + tags=[] + ) + + await memmachine_editor_with_client.add_items([item]) + + # Verify add was called once + assert mock_memory_instance.add.call_count == 1 + + # Verify memory_types is used (episode_type is None) + call_kwargs = mock_memory_instance.add.call_args.kwargs + assert call_kwargs["content"] == "This is a standalone memory" + assert call_kwargs["episode_type"] is None + assert "memory_types" in call_kwargs + + +async def test_search_success( + memmachine_editor_with_client: MemMachineEditor, + mock_client: Mock, + mock_project: Mock, + mock_memory_instance: Mock +): + """Test searching with a valid query and user ID.""" + # Mock search results with the new nested structure + # MemMachine SDK returns SearchResult with content containing nested episodic_memory + mock_search_result = Mock() + mock_search_result.content = { + "episodic_memory": { + "long_term_memory": { + "episodes": [ + { + "content": "I like pizza", + "metadata": {"key1": "value1", "tags": "food"} + } + ] + }, + "short_term_memory": { + "episodes": [] + } + }, + "semantic_memory": [ + { + "feature": "User prefers Italian food", + "metadata": {"key2": "value2"} + } + ] + } + mock_memory_instance.search.return_value = mock_search_result + + result = await memmachine_editor_with_client.search( + query="What do I like to eat?", + top_k=5, + user_id="user123", + session_id="session456", + agent_id="agent789" + ) + + # Verify search was called + mock_memory_instance.search.assert_called_once_with(query="What do I like to eat?", limit=5) + + # Verify results + assert len(result) == 2 # One episodic + one semantic + assert result[0].memory == "I like pizza" + assert result[0].tags == ["food"] + assert result[1].memory == "User prefers Italian food" + + +async def test_search_with_string_tags( + memmachine_editor_with_client: MemMachineEditor, + mock_memory_instance: Mock +): + """Test searching when tags come back as comma-separated string from SDK.""" + # Mock search results with the new nested structure + mock_search_result = Mock() + mock_search_result.content = { + "episodic_memory": { + "long_term_memory": { + "episodes": [ + { + "content": "I like pizza and pasta", + "metadata": {"tags": "food, preference, italian"} # String format + } + ] + }, + "short_term_memory": { + "episodes": [] + } + }, + "semantic_memory": [] + } + mock_memory_instance.search.return_value = mock_search_result + + result = await memmachine_editor_with_client.search( + query="What do I like?", + top_k=5, + user_id="user123" + ) + + assert len(result) == 1 + # Tags should be converted from string to list + assert result[0].tags == ["food", "preference", "italian"] + + +async def test_search_empty_results( + memmachine_editor_with_client: MemMachineEditor, + mock_memory_instance: Mock +): + """Test searching with empty results.""" + mock_search_result = Mock() + mock_search_result.content = { + "episodic_memory": { + "long_term_memory": {"episodes": []}, + "short_term_memory": {"episodes": []} + }, + "semantic_memory": [] + } + mock_memory_instance.search.return_value = mock_search_result + + result = await memmachine_editor_with_client.search( + query="test query", + top_k=5, + user_id="user123" + ) + + assert len(result) == 0 + + +async def test_search_missing_user_id(memmachine_editor_with_client: MemMachineEditor): + """Test searching without providing a user ID.""" + with pytest.raises(KeyError, match="user_id"): + await memmachine_editor_with_client.search(query="test query") + + +async def test_search_with_defaults( + memmachine_editor_with_client: MemMachineEditor, + mock_memory_instance: Mock +): + """Test searching with default session_id and agent_id.""" + mock_search_result = Mock() + mock_search_result.content = { + "episodic_memory": { + "long_term_memory": {"episodes": []}, + "short_term_memory": {"episodes": []} + }, + "semantic_memory": [] + } + mock_memory_instance.search.return_value = mock_search_result + + await memmachine_editor_with_client.search( + query="test query", + user_id="user123" + ) + + # Verify memory instance was created with defaults + # The editor should use default_session and default_agent + mock_memory_instance.search.assert_called_once() + + +async def test_remove_items_by_memory_id_episodic( + memmachine_editor_with_client: MemMachineEditor, + mock_client: Mock, + mock_project: Mock, + mock_memory_instance: Mock +): + """Test removing items by episodic memory ID.""" + await memmachine_editor_with_client.remove_items( + memory_id="episodic_123", + memory_type="episodic", + user_id="user123", + session_id="session456", + agent_id="agent789" + ) + + # Verify delete_episodic was called + mock_memory_instance.delete_episodic.assert_called_once_with(episodic_id="episodic_123") + + +async def test_remove_items_by_memory_id_semantic( + memmachine_editor_with_client: MemMachineEditor, + mock_client: Mock, + mock_project: Mock, + mock_memory_instance: Mock +): + """Test removing items by semantic memory ID.""" + await memmachine_editor_with_client.remove_items( + memory_id="semantic_123", + memory_type="semantic", + user_id="user123", + session_id="session456", + agent_id="agent789" + ) + + # Verify delete_semantic was called + mock_memory_instance.delete_semantic.assert_called_once_with(semantic_id="semantic_123") + + +async def test_remove_items_by_memory_id_without_user_id( + memmachine_editor_with_client: MemMachineEditor +): + """Test that removing items by memory_id without user_id raises ValueError.""" + with pytest.raises(ValueError, match="user_id is required"): + await memmachine_editor_with_client.remove_items(memory_id="episodic_123") + + +async def test_remove_items_by_user_id_not_implemented( + memmachine_editor_with_client: MemMachineEditor +): + """Test that removing all items by user_id raises NotImplementedError.""" + with pytest.raises(NotImplementedError, match="Bulk deletion by user_id"): + await memmachine_editor_with_client.remove_items(user_id="user123") + + +async def test_editor_with_project_instance( + memmachine_editor_with_project: MemMachineEditor, + mock_project: Mock, + mock_memory_instance: Mock, + sample_memory_item: MemoryItem +): + """Test that editor works correctly when initialized with a Project instance.""" + items = [sample_memory_item] + await memmachine_editor_with_project.add_items(items) + + # Verify project.memory was called directly (not create_project) + mock_project.memory.assert_called_once() + + # Verify add was called + assert mock_memory_instance.add.call_count == 2 + + +async def test_add_items_with_custom_project_and_org( + memmachine_editor_with_client: MemMachineEditor, + mock_client: Mock, + mock_project: Mock, + mock_memory_instance: Mock +): + """Test adding items with custom project_id and org_id in metadata.""" + item = MemoryItem( + conversation=[{"role": "user", "content": "Test"}], + user_id="user123", + memory="Test memory", + metadata={ + "session_id": "session456", + "agent_id": "agent789", + "project_id": "custom_project", + "org_id": "custom_org" + } + ) + + await memmachine_editor_with_client.add_items([item]) + + # Verify project was created/retrieved with custom org_id and project_id + mock_client.get_or_create_project.assert_called_once_with( + org_id="custom_org", + project_id="custom_project", + description="Project for user123" + ) + + +async def test_search_with_custom_project_and_org( + memmachine_editor_with_client: MemMachineEditor, + mock_client: Mock, + mock_project: Mock, + mock_memory_instance: Mock +): + """Test searching with custom project_id and org_id.""" + mock_memory_instance.search.return_value = { + "episodic_memory": [], + "semantic_memory": [], + "episode_summary": [] + } + + await memmachine_editor_with_client.search( + query="test", + user_id="user123", + project_id="custom_project", + org_id="custom_org" + ) + + # Verify project was created/retrieved with custom IDs + mock_client.get_or_create_project.assert_called_once_with( + org_id="custom_org", + project_id="custom_project", + description="Project for user123" + ) diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py new file mode 100644 index 0000000000..0ab35a0375 --- /dev/null +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py @@ -0,0 +1,431 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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. + +""" +Integration tests for MemMachine memory integration. + +These tests require a running MemMachine server. They test the full +integration by adding memories and then retrieving them. + +To run these tests: +1. Start MemMachine server (and databases) +2. Set MEMMACHINE_BASE_URL environment variable (defaults to http://localhost:8080) +3. Run: pytest tests/test_memmachine_integration.py -v +""" + +import os +import uuid + +import pytest + +from nat.builder.workflow_builder import WorkflowBuilder +from nat.data_models.config import GeneralConfig +from nat.memory.models import MemoryItem +from nat.plugins.memmachine.memory import MemMachineMemoryClientConfig + + +@pytest.fixture(name="memmachine_base_url") +def memmachine_base_url_fixture(): + """Get MemMachine base URL from environment or use default.""" + return os.environ.get("MEMMACHINE_BASE_URL", "http://localhost:8080") + + +@pytest.fixture(name="test_config") +def test_config_fixture(memmachine_base_url: str): + """Create a test configuration.""" + # Use unique org/project IDs for each test run to avoid conflicts + test_id = str(uuid.uuid4())[:8] + return MemMachineMemoryClientConfig( + base_url=memmachine_base_url, + org_id=f"test_org_{test_id}", + project_id=f"test_project_{test_id}", + timeout=30, + max_retries=3 + ) + + +@pytest.fixture(name="test_user_id") +def test_user_id_fixture(): + """Generate a unique user ID for testing.""" + return f"test_user_{uuid.uuid4().hex[:8]}" + + +@pytest.mark.integration +@pytest.mark.slow +async def test_add_and_retrieve_conversation_memory( + test_config: MemMachineMemoryClientConfig, + test_user_id: str +): + """Test adding a conversation memory and retrieving it.""" + general_config = GeneralConfig() + async with WorkflowBuilder(general_config=general_config) as builder: + await builder.add_memory_client("memmachine_memory", test_config) + memory_client = await builder.get_memory_client("memmachine_memory") + + # Create a test conversation memory + conversation = [ + {"role": "user", "content": "I love pizza and Italian food."}, + {"role": "assistant", "content": "I'll remember that you love pizza and Italian food."}, + ] + + memory_item = MemoryItem( + conversation=conversation, + user_id=test_user_id, + memory="User loves pizza", + metadata={ + "session_id": "test_session_1", + "agent_id": "test_agent_1", + "test_id": "conversation_test" + }, + tags=["food", "preference"] + ) + + # Add the memory + await memory_client.add_items([memory_item]) + + # Wait a moment for indexing (if needed) + import asyncio + await asyncio.sleep(1) + + # Retrieve the memory + retrieved_memories = await memory_client.search( + query="pizza Italian food", + top_k=10, + user_id=test_user_id, + session_id="test_session_1", + agent_id="test_agent_1" + ) + + # Verify we got results + assert len(retrieved_memories) > 0, "Should retrieve at least one memory" + + # Check that our memory is in the results + # Note: MemMachine may store conversation messages separately or process them, + # so we check for the content/keywords rather than exact conversation structure + found = False + for mem in retrieved_memories: + # Check if this is our memory by looking for the test_id in metadata + if mem.metadata.get("test_id") == "conversation_test": + found = True + # MemMachine may return individual messages, not full conversations + # So we check that the content is present (either in conversation or memory field) + content = mem.memory or (str(mem.conversation) if mem.conversation else "") + assert "pizza" in content.lower() or "italian" in content.lower(), \ + f"Should contain pizza/italian content. Got: {content}" + # Verify tags + assert "food" in mem.tags or "preference" in mem.tags, \ + f"Should have tags. Got: {mem.tags}" + break + + assert found, f"Should find the memory we just added. Found {len(retrieved_memories)} memories with metadata: {[m.metadata.get('test_id') for m in retrieved_memories]}" + + +@pytest.mark.integration +@pytest.mark.slow +async def test_add_and_retrieve_direct_memory( + test_config: MemMachineMemoryClientConfig, + test_user_id: str +): + """Test adding a direct memory (fact/preference without conversation) and retrieving it. + + All memories are now added to both episodic and semantic memory types. + """ + general_config = GeneralConfig() + async with WorkflowBuilder(general_config=general_config) as builder: + await builder.add_memory_client("memmachine_memory", test_config) + memory_client = await builder.get_memory_client("memmachine_memory") + + # Create a direct memory (no conversation) + direct_memory = MemoryItem( + conversation=None, + user_id=test_user_id, + memory="User prefers working in the morning and is allergic to peanuts", + metadata={ + "session_id": "test_session_2", + "agent_id": "test_agent_2", + "test_id": "direct_test" + }, + tags=["preference", "allergy"] + ) + + # Add the memory + await memory_client.add_items([direct_memory]) + + # Wait for memory ingestion + # Memories are processed asynchronously by MemMachine's background task + import asyncio + await asyncio.sleep(5) # Wait for background ingestion task + + # Try searching multiple times with retries (memory ingestion is async) + retrieved_memories = [] + for attempt in range(3): + retrieved_memories = await memory_client.search( + query="morning work allergy peanuts", + top_k=10, + user_id=test_user_id, + session_id="test_session_2", + agent_id="test_agent_2" + ) + if len(retrieved_memories) > 0: + break + await asyncio.sleep(2) # Wait another 2 seconds before retry + + # Verify we got results + if len(retrieved_memories) == 0: + # If no results, try a broader search + retrieved_memories = await memory_client.search( + query="preference allergy", # Broader query + top_k=20, + user_id=test_user_id, + session_id="test_session_2", + agent_id="test_agent_2" + ) + + # Check for related keywords + found = False + for mem in retrieved_memories: + # Check by test_id or by content keywords + if mem.metadata.get("test_id") == "direct_test": + found = True + break + content = mem.memory.lower() if mem.memory else "" + if any(keyword in content for keyword in ["morning", "peanut", "allergy", "prefer"]): + found = True + break + + # It's acceptable if we don't find exact match immediately due to async processing + if not found: + pytest.skip( + "Direct memory not found - this may be due to async processing delay. " + f"Found {len(retrieved_memories)} memories. " + "Memory ingestion can take several seconds." + ) + + +@pytest.mark.integration +@pytest.mark.slow +async def test_add_multiple_and_retrieve_all( + test_config: MemMachineMemoryClientConfig, + test_user_id: str +): + """Test adding multiple memories and retrieving them all.""" + general_config = GeneralConfig() + async with WorkflowBuilder(general_config=general_config) as builder: + await builder.add_memory_client("memmachine_memory", test_config) + memory_client = await builder.get_memory_client("memmachine_memory") + + # Create multiple test memories + memories = [ + MemoryItem( + conversation=[{"role": "user", "content": f"Memory {i}: I like item {i}"}], + user_id=test_user_id, + memory=f"Memory {i}", + metadata={ + "session_id": "test_session_3", + "agent_id": "test_agent_3", + "test_id": f"multi_test_{i}" + }, + tags=[f"item_{i}"] + ) + for i in range(1, 6) # Create 5 memories + ] + + # Add all memories + await memory_client.add_items(memories) + + # Wait for indexing + import asyncio + await asyncio.sleep(2) + + # Retrieve all memories with a broad query + retrieved_memories = await memory_client.search( + query="*", # Broad query to get all + top_k=20, + user_id=test_user_id, + session_id="test_session_3", + agent_id="test_agent_3" + ) + + # Verify we got results + assert len(retrieved_memories) >= 3, f"Should retrieve at least 3 memories, got {len(retrieved_memories)}" + + # Check that our test memories are in the results + found_ids = set() + for mem in retrieved_memories: + test_id = mem.metadata.get("test_id", "") + if test_id.startswith("multi_test_"): + found_ids.add(test_id) + + assert len(found_ids) >= 3, f"Should find at least 3 of our test memories, found: {found_ids}" + + +@pytest.mark.integration +@pytest.mark.slow +async def test_add_and_verify_conversation_content_match( + test_config: MemMachineMemoryClientConfig, + test_user_id: str +): + """Test that conversation memory content can be retrieved. + + All memories are added to both episodic and semantic memory types. + """ + general_config = GeneralConfig() + async with WorkflowBuilder(general_config=general_config) as builder: + await builder.add_memory_client("memmachine_memory", test_config) + memory_client = await builder.get_memory_client("memmachine_memory") + + # Create a conversation memory + original_content = "The user mentioned their favorite programming language is Python" + original_tags = ["programming", "preference"] + + memory_item = MemoryItem( + conversation=[{"role": "user", "content": original_content}], + user_id=test_user_id, + memory=original_content, + metadata={ + "session_id": "test_session_4", + "agent_id": "test_agent_4", + "test_id": "conversation_content_test" + }, + tags=original_tags + ) + + # Add the memory + await memory_client.add_items([memory_item]) + + # Wait for indexing + import asyncio + await asyncio.sleep(2) + + # Retrieve the memory + retrieved_memories = await memory_client.search( + query="Python programming language", + top_k=10, + user_id=test_user_id, + session_id="test_session_4", + agent_id="test_agent_4" + ) + + # Find our memory + found_memory = None + for mem in retrieved_memories: + if mem.metadata.get("test_id") == "conversation_content_test": + found_memory = mem + break + + assert found_memory is not None, f"Should find the conversation memory. Found {len(retrieved_memories)} memories" + + # Verify content + content = found_memory.memory.lower() if found_memory.memory else "" + assert "python" in content or "programming" in content, \ + f"Retrieved memory should contain 'Python' or 'programming'. Got: {found_memory.memory}" + + # Verify tags are preserved + assert len(found_memory.tags) > 0, "Should have tags" + assert any("programming" in tag.lower() or "preference" in tag.lower() for tag in found_memory.tags), \ + f"Should have relevant tags. Got: {found_memory.tags}" + + +@pytest.mark.integration +@pytest.mark.slow +async def test_conversation_and_direct_memory_both_retrievable( + test_config: MemMachineMemoryClientConfig, + test_user_id: str +): + """Test that both conversation and direct memories are stored and retrievable. + + All memories are now added to both episodic and semantic memory types. + """ + general_config = GeneralConfig() + async with WorkflowBuilder(general_config=general_config) as builder: + await builder.add_memory_client("memmachine_memory", test_config) + memory_client = await builder.get_memory_client("memmachine_memory") + + # Add conversation memory + conversation_memory = MemoryItem( + conversation=[ + {"role": "user", "content": "What's the weather today?"}, + {"role": "assistant", "content": "It's sunny and 75Ā°F."} + ], + user_id=test_user_id, + memory="Weather conversation", + metadata={ + "session_id": "test_session_5", + "agent_id": "test_agent_5", + "test_id": "conversation_type_test" + }, + tags=["weather"] + ) + + # Add direct memory (no conversation) + direct_memory = MemoryItem( + conversation=None, + user_id=test_user_id, + memory="User lives in San Francisco and works as a software engineer", + metadata={ + "session_id": "test_session_5", + "agent_id": "test_agent_5", + "test_id": "direct_type_test" + }, + tags=["location", "occupation"] + ) + + # Add both + await memory_client.add_items([conversation_memory, direct_memory]) + + # Wait for indexing + import asyncio + await asyncio.sleep(2) + + # Search for conversation memory + conversation_results = await memory_client.search( + query="weather sunny", + top_k=10, + user_id=test_user_id, + session_id="test_session_5", + agent_id="test_agent_5" + ) + + # Search for direct memory (with retries due to async processing) + direct_results = [] + for attempt in range(3): + direct_results = await memory_client.search( + query="San Francisco software engineer", + top_k=10, + user_id=test_user_id, + session_id="test_session_5", + agent_id="test_agent_5" + ) + if len(direct_results) > 0: + break + await asyncio.sleep(3) # Wait for memory ingestion + + # Verify conversation memory can be retrieved + conversation_found = any(m.metadata.get("test_id") == "conversation_type_test" for m in conversation_results) + assert conversation_found or len(conversation_results) > 0, "Should find conversation memory" + + # Check for direct memory + direct_found = any(m.metadata.get("test_id") == "direct_type_test" for m in direct_results) + direct_keywords_found = any( + any(keyword in (m.memory or "").lower() for keyword in ["san francisco", "software", "engineer"]) + for m in direct_results + ) + + # Direct memory may not be immediately available due to async processing + if not direct_found and not direct_keywords_found: + pytest.skip( + "Direct memory not found - may be due to async processing delay. " + "Memories are processed asynchronously." + ) diff --git a/packages/nvidia_nat_memmachine/tests/test_memory.py b/packages/nvidia_nat_memmachine/tests/test_memory.py new file mode 100644 index 0000000000..20d2a2668a --- /dev/null +++ b/packages/nvidia_nat_memmachine/tests/test_memory.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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 sys +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +import pytest + +from nat.builder.builder import Builder +from nat.plugins.memmachine.memory import MemMachineMemoryClientConfig, memmachine_memory_client + + +@pytest.fixture(name="mock_builder") +def mock_builder_fixture(): + """Fixture to provide a mocked Builder instance.""" + return Mock(spec=Builder) + + +@pytest.fixture(name="config") +def config_fixture(): + """Fixture to provide a MemMachineMemoryClientConfig instance.""" + return MemMachineMemoryClientConfig( + base_url="http://localhost:8080", + org_id="test_org", + project_id="test_project", + timeout=30, + max_retries=3 + ) + + +@pytest.fixture(name="config_minimal") +def config_minimal_fixture(): + """Fixture to provide a minimal MemMachineMemoryClientConfig instance.""" + return MemMachineMemoryClientConfig( + base_url="http://localhost:8080" + ) + + +@pytest.fixture(name="mock_memmachine_client") +def mock_memmachine_client_fixture(): + """Fixture to provide a mocked MemMachineClient.""" + mock_client = Mock() + mock_client.base_url = "http://localhost:8080" + return mock_client + + +@pytest.fixture(name="mock_project") +def mock_project_fixture(): + """Fixture to provide a mocked Project instance.""" + mock_project = Mock() + mock_project.org_id = "test_org" + mock_project.project_id = "test_project" + return mock_project + + +async def test_memmachine_memory_client_success( + config: MemMachineMemoryClientConfig, + mock_builder: Mock, + mock_memmachine_client: Mock, + mock_project: Mock +): + """Test successful initialization of memmachine memory client.""" + mock_memmachine_client.get_or_create_project.return_value = mock_project + + # Patch where the import happens - inside the function + with patch("memmachine.MemMachineClient", return_value=mock_memmachine_client): + # @register_memory wraps the function with asynccontextmanager, so use async with + async with memmachine_memory_client(config, mock_builder) as editor: + assert editor is not None + # Verify client was initialized correctly + mock_memmachine_client.get_or_create_project.assert_called_once_with( + org_id="test_org", + project_id="test_project", + description="NeMo Agent toolkit project: test_project" + ) + + +async def test_memmachine_memory_client_minimal_config( + config_minimal: MemMachineMemoryClientConfig, + mock_builder: Mock, + mock_memmachine_client: Mock +): + """Test initialization with minimal config (no org_id/project_id).""" + with patch("memmachine.MemMachineClient", return_value=mock_memmachine_client): + # @register_memory wraps the function with asynccontextmanager, so use async with + async with memmachine_memory_client(config_minimal, mock_builder) as editor: + assert editor is not None + # Should not create project if org_id/project_id not provided + mock_memmachine_client.get_or_create_project.assert_not_called() + + +async def test_memmachine_memory_client_import_error( + config: MemMachineMemoryClientConfig, + mock_builder: Mock +): + """Test that ImportError is raised when memmachine package is not installed.""" + # Mock the import to raise ImportError + # We need to patch the import inside the function, so patch where it's imported from + import builtins + original_import = builtins.__import__ + + def import_side_effect(name, *args, **kwargs): + if name == "memmachine": + raise ImportError("No module named 'memmachine'") + return original_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=import_side_effect): + with pytest.raises(ImportError, match="Could not import MemMachineClient"): + async with memmachine_memory_client(config, mock_builder): + pass + + +async def test_memmachine_memory_client_initialization_error( + config: MemMachineMemoryClientConfig, + mock_builder: Mock +): + """Test that RuntimeError is raised when client initialization fails.""" + with patch("memmachine.MemMachineClient", side_effect=ValueError("base_url is required")): + with pytest.raises(RuntimeError, match="Failed to initialize MemMachineClient"): + async with memmachine_memory_client(config, mock_builder): + pass + + +async def test_memmachine_memory_client_project_creation_failure( + config: MemMachineMemoryClientConfig, + mock_builder: Mock, + mock_memmachine_client: Mock +): + """Test that editor still works if project creation fails.""" + mock_memmachine_client.get_or_create_project.side_effect = Exception("Project creation failed") + + with patch("memmachine.MemMachineClient", return_value=mock_memmachine_client): + # Should not raise exception, should fall back to using client directly + # @register_memory wraps the function with asynccontextmanager, so use async with + async with memmachine_memory_client(config, mock_builder) as editor: + assert editor is not None + # Project creation should have been attempted + mock_memmachine_client.get_or_create_project.assert_called_once() + + + + +async def test_memmachine_memory_client_config_validation(): + """Test that MemMachineMemoryClientConfig validates required fields.""" + # base_url is required + with pytest.raises(Exception): # Pydantic validation error + MemMachineMemoryClientConfig() + + # Should work with base_url + config = MemMachineMemoryClientConfig(base_url="http://localhost:8080") + assert config.base_url == "http://localhost:8080" + assert config.timeout == 30 + assert config.max_retries == 3 + + +async def test_memmachine_memory_client_with_retry_mixin( + config: MemMachineMemoryClientConfig, + mock_builder: Mock, + mock_memmachine_client: Mock, + mock_project: Mock +): + """Test that retry mixin is applied when config has retry settings.""" + mock_memmachine_client.get_or_create_project.return_value = mock_project + + # Add retry configuration + config.num_retries = 5 + config.retry_on_status_codes = [500, 502, 503] + config.retry_on_errors = ["ConnectionError"] + + with patch("memmachine.MemMachineClient", return_value=mock_memmachine_client): + with patch("nat.plugins.memmachine.memory.patch_with_retry") as mock_patch: + mock_patch.return_value = Mock() + # @register_memory wraps the function with asynccontextmanager, so use async with + async with memmachine_memory_client(config, mock_builder) as editor: + assert editor is not None + # Verify patch_with_retry was called + mock_patch.assert_called_once() + call_kwargs = mock_patch.call_args.kwargs + assert call_kwargs["retries"] == 5 + assert call_kwargs["retry_codes"] == [500, 502, 503] + assert call_kwargs["retry_on_messages"] == ["ConnectionError"] From 07ff4dff170cff2c5a19c197b7b6a955720a5378 Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Fri, 23 Jan 2026 11:52:18 -0800 Subject: [PATCH 02/18] sS Address the PR review comments ... update nat version to 1.5, add missing license headers and change copyright years, add logging with proper exception handling, remove unnecessary files --- packages/nvidia_nat_all/pyproject.toml | 1 + packages/nvidia_nat_memmachine/pyproject.toml | 4 +- .../src/nat/plugins/__init__.py | 8 +- .../plugins/memmachine/memmachine_editor.py | 43 ++++-- .../src/nat/plugins/memmachine/memory.py | 36 +++-- .../src/nat/plugins/memmachine/register.py | 15 ++ .../tests/API_CALL_VERIFICATION.md | 129 ------------------ .../nvidia_nat_memmachine/tests/__init__.py | 14 -- .../tests/test_add_and_retrieve.py | 29 +++- .../tests/test_memmachine_api_calls.py | 6 +- .../tests/test_memmachine_integration.py | 45 ++++-- .../tests/test_memory.py | 6 +- 12 files changed, 147 insertions(+), 189 deletions(-) delete mode 100644 packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md delete mode 100644 packages/nvidia_nat_memmachine/tests/__init__.py diff --git a/packages/nvidia_nat_all/pyproject.toml b/packages/nvidia_nat_all/pyproject.toml index 13f78c3dc9..34d5791dee 100644 --- a/packages/nvidia_nat_all/pyproject.toml +++ b/packages/nvidia_nat_all/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "nvidia-nat-llama-index", "nvidia-nat-mcp", "nvidia-nat-mem0ai", + "nvidia-nat-memmachine", "nvidia-nat-mysql", "nvidia-nat-nemo-customizer", # nvidia-nat-openpipe-art cannot be part of all due to conflicts with nvidia-nat-crewai diff --git a/packages/nvidia_nat_memmachine/pyproject.toml b/packages/nvidia_nat_memmachine/pyproject.toml index 2e2b0c6ceb..e0fdd181da 100644 --- a/packages/nvidia_nat_memmachine/pyproject.toml +++ b/packages/nvidia_nat_memmachine/pyproject.toml @@ -20,8 +20,8 @@ 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~=1.4", - "memmachine-server~=0.2.2", + "nvidia-nat~=1.5", + "memmachine-server~=0.2.1", ] requires-python = ">=3.11,<3.14" description = "Subpackage for MemMachine integration in NeMo Agent toolkit. Requires a cfg.yml configuration file with database and AI model settings." diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/__init__.py b/packages/nvidia_nat_memmachine/src/nat/plugins/__init__.py index f05adda80b..d1d9d1e7ea 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/__init__.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -12,9 +12,3 @@ # 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. - -# Namespace package - extend __path__ to include all nat.plugins locations -# This allows plugins installed in different locations (e.g., site-packages, editable installs) -# to be discovered as part of the same namespace -from pkgutil import extend_path -__path__ = extend_path(__path__, __name__) diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py index 8fe73f7945..3b0b7fea93 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py @@ -1,4 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-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 +import logging from typing import Any import requests @@ -7,6 +23,8 @@ from nat.memory.interfaces import MemoryEditor from nat.memory.models import MemoryItem +logger = logging.getLogger(__name__) + class MemMachineEditor(MemoryEditor): """ @@ -182,14 +200,19 @@ def add_memory(content=msg_content, role=msg_role, mem_types=memory_types, meta= # Convert list to comma-separated string metadata["tags"] = ", ".join(tags) if isinstance(tags, list) else str(tags) - def add_memory(): + def add_memory( + content=memory_text, + mem=memory, + meta=metadata, + mem_types=memory_types + ): # Use MemMachine SDK add() method # API: memory.add(content, role="user", metadata={}, memory_types=[...]) - memory.add( - content=memory_text, + mem.add( + content=content, role="user", - metadata=metadata if metadata else None, - memory_types=memory_types, + metadata=meta if meta else None, + memory_types=mem_types, episode_type=None # Use default (MESSAGE) ) @@ -300,8 +323,12 @@ def perform_search(): # Sort episodes by created_at timestamp if available try: conv_episodes.sort(key=lambda e: e.get("created_at") or e.get("timestamp") or "") - except: - pass + except (TypeError, AttributeError, ValueError) as e: + # Skip sorting if timestamps are missing or incompatible + logger.warning( + f"Failed to sort episodes for conversation '{conv_key}': {e}. " + "Continuing without sorting." + ) # Extract conversation messages conversation_messages = [] @@ -475,7 +502,7 @@ def delete_memory(): group_id = kwargs.pop("group_id", "default") project_id = kwargs.pop("project_id", None) org_id = kwargs.pop("org_id", None) - delete_semantic = kwargs.pop("delete_semantic_memory", False) + # Note: delete_semantic_memory flag is not yet implemented for bulk deletion # Note: MemMachine SDK doesn't have a delete_all method # We would need to search for all memories and delete them individually diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py index 17e1db4a21..75f80ab52e 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,12 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging +from collections.abc import AsyncGenerator + from nat.builder.builder import Builder from nat.cli.register_workflow import register_memory from nat.data_models.memory import MemoryBaseConfig from nat.data_models.retry_mixin import RetryMixin +from nat.memory.interfaces import MemoryEditor from nat.utils.exception_handlers.automatic_retries import patch_with_retry +logger = logging.getLogger(__name__) + class MemMachineMemoryClientConfig(MemoryBaseConfig, RetryMixin, name="memmachine_memory"): """ @@ -39,7 +45,10 @@ class MemMachineMemoryClientConfig(MemoryBaseConfig, RetryMixin, name="memmachin @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[MemoryEditor, None]: from .memmachine_editor import MemMachineEditor # Import and initialize the MemMachine Python SDK try: @@ -82,19 +91,24 @@ async def memmachine_memory_client(config: MemMachineMemoryClientConfig, builder description=f"NeMo Agent toolkit project: {config.project_id}" ) memmachine_instance = project - except Exception as e: + except Exception: # If project creation fails, fall back to using the client directly # The editor will handle project creation on-demand - pass + logger.warning( + "Failed to create/get project '%s' in org '%s', falling back to client-level access", + config.project_id, + config.org_id, + exc_info=True, + ) memory_editor = MemMachineEditor(memmachine_instance=memmachine_instance) - 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 - ) + # Apply retry wrapper (config always inherits from 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 + ) yield memory_editor \ No newline at end of file diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py index 9ce1b2c3e5..e2e0382d6f 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/register.py @@ -1 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-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 . import memory \ No newline at end of file diff --git a/packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md b/packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md deleted file mode 100644 index 7cb0614bce..0000000000 --- a/packages/nvidia_nat_memmachine/tests/API_CALL_VERIFICATION.md +++ /dev/null @@ -1,129 +0,0 @@ -# API Call Verification Tests - -## Overview - -The tests in `test_memmachine_api_calls.py` verify that the MemMachine integration makes **correct API calls** to the MemMachine SDK. These tests use spies to capture and validate: - -1. **Exact SDK methods called** - Verifies the right methods are invoked -2. **Parameter correctness** - Verifies parameters match SDK expectations -3. **Data transformations** - Verifies NAT format ā†’ MemMachine format conversion -4. **Memory type handling** - Verifies all memories are added to both episodic and semantic types - -## What Gets Tested - -### 1. Add Operations (`add_items`) - -#### Conversation Handling -- āœ… Each message in conversation calls `memory.add()` separately -- āœ… Messages preserve their roles (user/assistant/system) -- āœ… Messages are added in the correct order -- āœ… All memories are added to both episodic and semantic memory types - -#### Direct Memory (No Conversation) -- āœ… When `conversation=None`, uses `memory` field as content -- āœ… All memories are added to both episodic and semantic memory types -- āœ… Default `role` is "user" - -#### Metadata Handling -- āœ… Tags are included in metadata dict -- āœ… Custom metadata fields are preserved -- āœ… Special fields (session_id, agent_id, project_id, org_id) are extracted and NOT passed to `add()` -- āœ… Empty metadata becomes `None`, not empty dict - -#### Project/Org Handling -- āœ… Custom `project_id`/`org_id` in metadata triggers `create_project()` call -- āœ… `project.memory()` is called with correct user_id, session_id, agent_id, group_id - -### 2. Search Operations (`search`) - -#### Parameter Mapping -- āœ… NAT's `top_k` parameter is converted to SDK's `limit` parameter -- āœ… `query` parameter is passed correctly -- āœ… `project.memory()` is called with user_id, session_id, agent_id, group_id - -#### Custom Project/Org -- āœ… Custom `project_id`/`org_id` triggers `create_project()` call - -### 3. Remove Operations (`remove_items`) - -#### Episodic Deletion -- āœ… Calls `memory.delete_episodic(episodic_id=...)` with correct ID - -#### Semantic Deletion -- āœ… Calls `memory.delete_semantic(semantic_id=...)` with correct ID - -### 4. API Call Format - -#### Keyword Arguments -- āœ… All SDK methods are called with keyword arguments, not positional -- āœ… Parameter names match SDK exactly (`limit` not `top_k`, `episodic_id` not `memory_id`) - -## Test Structure - -Each test class focuses on a specific aspect: - -- **`TestAddItemsAPICalls`** - Verifies `add_items()` makes correct `memory.add()` calls -- **`TestSearchAPICalls`** - Verifies `search()` makes correct `memory.search()` calls -- **`TestRemoveItemsAPICalls`** - Verifies `remove_items()` makes correct delete calls -- **`TestAPICallParameterValidation`** - Verifies parameter names and formats -- **`TestDataTransformation`** - Verifies data transformations are correct - -## Running the Tests - -```bash -# Run all API call verification tests -pytest tests/test_memmachine_api_calls.py -v - -# Run a specific test class -pytest tests/test_memmachine_api_calls.py::TestAddItemsAPICalls -v - -# Run a specific test -pytest tests/test_memmachine_api_calls.py::TestAddItemsAPICalls::test_add_conversation_calls_add_with_correct_parameters -v -``` - -## Key Differences from Unit Tests - -| Aspect | Unit Tests (`test_memmachine_editor.py`) | API Call Tests (`test_memmachine_api_calls.py`) | -|--------|------------------------------------------|--------------------------------------------------| -| **Focus** | Integration logic, error handling | Exact API calls and parameters | -| **Verification** | Mocks return values | Spies capture actual calls | -| **What's Tested** | Code flow, edge cases | Parameter correctness, data transformation | -| **SDK Methods** | Mocked, not verified | Spied, parameters validated | - -## Example: What Gets Verified - -### Adding a Conversation - -**Input (NAT format):** -```python -MemoryItem( - conversation=[ - {"role": "user", "content": "I like pizza"}, - {"role": "assistant", "content": "Great!"} - ], - user_id="user123", - metadata={"session_id": "s1", "agent_id": "a1"}, - tags=["food"] -) -``` - -**Verified API Calls:** -1. āœ… `create_project(org_id="default-org", project_id="default-project", ...)` -2. āœ… `project.memory(user_id="user123", session_id="s1", agent_id="a1", group_id="default")` -3. āœ… `memory.add(content="I like pizza", role="user", memory_types=[Episodic, Semantic], metadata={"tags": ["food"]})` -4. āœ… `memory.add(content="Great!", role="assistant", memory_types=[Episodic, Semantic], metadata={"tags": ["food"]})` - -**What's Verified:** -- āœ… Two `add()` calls (one per message) -- āœ… Correct roles preserved -- āœ… All memories added to both episodic and semantic types -- āœ… Tags in metadata -- āœ… session_id/agent_id NOT in metadata (used for memory instance instead) -- āœ… All parameters are keyword arguments - -## Integration with Real Server - -For tests that verify actual API calls work with a real MemMachine server, see: -- `test_memmachine_integration.py` - Tests with real server -- `TESTING.md` - Guide for integration testing - diff --git a/packages/nvidia_nat_memmachine/tests/__init__.py b/packages/nvidia_nat_memmachine/tests/__init__.py deleted file mode 100644 index a1744724ec..0000000000 --- a/packages/nvidia_nat_memmachine/tests/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, 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_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py index 0662e3e40d..a4f50cdb5f 100644 --- a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -1,4 +1,19 @@ #!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2024-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. + """ Simple script to test adding memories and retrieving them. @@ -14,14 +29,19 @@ """ import asyncio +import logging import os import uuid from datetime import datetime +import pytest + from nat.builder.builder import Builder from nat.memory.models import MemoryItem from nat.plugins.memmachine.memory import MemMachineMemoryClientConfig +logger = logging.getLogger(__name__) + async def test_add_and_retrieve(): """Test adding memories and retrieving them.""" @@ -178,10 +198,15 @@ async def test_add_and_retrieve(): except Exception as e: print(f"\nāœ— Error: {e}") - import traceback - traceback.print_exc() + logger.exception("Error during test execution") raise +@pytest.mark.integration +async def test_add_and_retrieve_integration(): + """Integration test for adding and retrieving memories.""" + await test_add_and_retrieve() + + if __name__ == "__main__": asyncio.run(test_add_and_retrieve()) diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py index c8e1b56614..ad46eb74c2 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -47,7 +47,7 @@ def record_call(self, method_name: str, args: tuple, kwargs: dict): 'kwargs': kwargs }) - def get_calls(self, method_name: str = None): + def get_calls(self, method_name: str | None = None): """Get all calls, optionally filtered by method name.""" if method_name: return [c for c in self.calls if c['method'] == method_name] @@ -363,7 +363,7 @@ async def test_search_calls_memory_search_with_correct_parameters( "episode_summary": [] } - results = await editor_with_spy.search( + await editor_with_spy.search( query="What do I like?", top_k=10, user_id="user123", diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py index 0ab35a0375..c3b7fc189f 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -19,10 +19,8 @@ These tests require a running MemMachine server. They test the full integration by adding memories and then retrieving them. -To run these tests: -1. Start MemMachine server (and databases) -2. Set MEMMACHINE_BASE_URL environment variable (defaults to http://localhost:8080) -3. Run: pytest tests/test_memmachine_integration.py -v +The tests will automatically skip if the MemMachine server is not available. +Set `MEMMACHINE_BASE_URL` environment variable to override default (http://localhost:8080). """ import os @@ -36,10 +34,35 @@ from nat.plugins.memmachine.memory import MemMachineMemoryClientConfig -@pytest.fixture(name="memmachine_base_url") -def memmachine_base_url_fixture(): - """Get MemMachine base URL from environment or use default.""" - return os.environ.get("MEMMACHINE_BASE_URL", "http://localhost:8080") +@pytest.fixture(name="memmachine_base_url", scope="session") +def memmachine_base_url_fixture(fail_missing: bool = False) -> str: + """ + Ensure MemMachine server is running and provide base URL. + + To run these tests, a MemMachine server must be running. + Set MEMMACHINE_BASE_URL environment variable to override default (http://localhost:8080). + """ + base_url = os.getenv("MEMMACHINE_BASE_URL", "http://localhost:8080") + if not base_url.startswith("http"): + base_url = f"http://{base_url}" + + try: + # Try to import and use MemMachineClient to check server availability + from memmachine import MemMachineClient + + client = MemMachineClient(base_url=base_url, timeout=5.0) + client.health_check(timeout=5.0) + return base_url + except ImportError: + reason = "memmachine package not installed. Install with: pip install memmachine" + if fail_missing: + raise RuntimeError(reason) from None + pytest.skip(reason=reason) + except Exception: + reason = f"Unable to connect to MemMachine server at {base_url}. Please ensure the server is running." + if fail_missing: + raise RuntimeError(reason) from None + pytest.skip(reason=reason) @pytest.fixture(name="test_config") @@ -170,7 +193,7 @@ async def test_add_and_retrieve_direct_memory( # Try searching multiple times with retries (memory ingestion is async) retrieved_memories = [] - for attempt in range(3): + for _attempt in range(3): retrieved_memories = await memory_client.search( query="morning work allergy peanuts", top_k=10, @@ -400,7 +423,7 @@ async def test_conversation_and_direct_memory_both_retrievable( # Search for direct memory (with retries due to async processing) direct_results = [] - for attempt in range(3): + for _attempt in range(3): direct_results = await memory_client.search( query="San Francisco software engineer", top_k=10, diff --git a/packages/nvidia_nat_memmachine/tests/test_memory.py b/packages/nvidia_nat_memmachine/tests/test_memory.py index 20d2a2668a..b46bc4e8cb 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memory.py +++ b/packages/nvidia_nat_memmachine/tests/test_memory.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -154,8 +154,10 @@ async def test_memmachine_memory_client_project_creation_failure( async def test_memmachine_memory_client_config_validation(): """Test that MemMachineMemoryClientConfig validates required fields.""" + from pydantic import ValidationError + # base_url is required - with pytest.raises(Exception): # Pydantic validation error + with pytest.raises(ValidationError): MemMachineMemoryClientConfig() # Should work with base_url From 3a7c40e766a043d414e08190fca823d8fd75a19b Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Fri, 23 Jan 2026 16:09:30 -0800 Subject: [PATCH 03/18] Remove group_id and other unused variable references and fix error logging --- .../plugins/memmachine/memmachine_editor.py | 36 +++++++++---------- .../src/nat/plugins/memmachine/memory.py | 6 ++-- .../tests/test_memmachine_api_calls.py | 17 ++------- .../tests/test_memmachine_editor.py | 3 +- .../tests/test_memmachine_integration.py | 14 +++----- 5 files changed, 26 insertions(+), 50 deletions(-) diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py index 3b0b7fea93..97702ab4ec 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py @@ -29,7 +29,7 @@ class MemMachineEditor(MemoryEditor): """ Wrapper class that implements NAT interfaces for MemMachine Integrations. - Uses the MemMachine Python SDK (MemMachineClient) as documented at: + Uses the MemMachine Python SDK as documented at: https://github.com/MemMachine/MemMachine/blob/main/docs/examples/python.mdx Supports both episodic and semantic memory through the unified SDK interface. @@ -37,11 +37,8 @@ class MemMachineEditor(MemoryEditor): User needs to add MemMachine SDK ids as metadata to the MemoryItem: - session_id - agent_id - - group_id - project_id - org_id - - Group ID is optional. If not provided, the memory will be added to the 'default' group. """ def __init__(self, memmachine_instance: Any): @@ -64,7 +61,6 @@ def _get_memory_instance( user_id: str, session_id: str, agent_id: str, - group_id: str = "default", project_id: str | None = None, org_id: str | None = None ) -> Any: @@ -75,7 +71,6 @@ def _get_memory_instance( user_id: User identifier session_id: Session identifier agent_id: Agent identifier - group_id: Group identifier (default: "default") project_id: Optional project identifier (default: "default-project") org_id: Optional organization identifier (default: "default-org") @@ -120,8 +115,7 @@ def _get_memory_instance( return project.memory( user_id=user_id, agent_id=agent_id, - session_id=session_id, - group_id=group_id + session_id=session_id ) async def add_items(self, items: list[MemoryItem]) -> None: @@ -142,16 +136,15 @@ async def add_items(self, items: list[MemoryItem]) -> None: tags = memory_item.tags memory_text = memory_item.memory - # Extract session_id, agent_id, group_id, project_id, and org_id from metadata if present + # Extract session_id, agent_id, project_id, and org_id from metadata if present session_id = item_meta.pop("session_id", "default_session") agent_id = item_meta.pop("agent_id", "default_agent") - group_id = item_meta.pop("group_id", "default") project_id = item_meta.pop("project_id", None) org_id = item_meta.pop("org_id", None) # Get memory instance using MemMachine SDK memory = self._get_memory_instance( - user_id, session_id, agent_id, group_id, project_id, org_id + user_id, session_id, agent_id, project_id, org_id ) # All memories are added to BOTH episodic and semantic memory types @@ -177,11 +170,17 @@ async def add_items(self, items: list[MemoryItem]) -> None: metadata["tags"] = ", ".join(tags) if isinstance(tags, list) else str(tags) # Capture variables in closure to avoid late binding issues - def add_memory(content=msg_content, role=msg_role, mem_types=memory_types, meta=metadata): + def add_memory( + content=msg_content, + role=msg_role, + mem=memory, + mem_types=memory_types, + meta=metadata, + ): # Use MemMachine SDK add() method # API: memory.add(content, role="user", metadata={}, memory_types=[...]) # episode_type should be None (defaults to "message") or EpisodeType.MESSAGE - memory.add( + mem.add( content=content, role=role, metadata=meta if meta else None, @@ -238,13 +237,12 @@ async def search(self, query: str, top_k: int = 5, **kwargs) -> list[MemoryItem] user_id = kwargs.pop("user_id") # Ensure user ID is in keyword arguments session_id = kwargs.pop("session_id", "default_session") agent_id = kwargs.pop("agent_id", "default_agent") - group_id = kwargs.pop("group_id", "default") project_id = kwargs.pop("project_id", None) org_id = kwargs.pop("org_id", None) # Get memory instance using MemMachine SDK memory = self._get_memory_instance( - user_id, session_id, agent_id, group_id, project_id, org_id + user_id, session_id, agent_id, project_id, org_id ) # Perform search using MemMachine SDK @@ -325,7 +323,7 @@ def perform_search(): conv_episodes.sort(key=lambda e: e.get("created_at") or e.get("timestamp") or "") except (TypeError, AttributeError, ValueError) as e: # Skip sorting if timestamps are missing or incompatible - logger.warning( + logger.exception( f"Failed to sort episodes for conversation '{conv_key}': {e}. " "Continuing without sorting." ) @@ -463,7 +461,7 @@ async def remove_items(self, **kwargs) -> None: Args: kwargs (dict): Keyword arguments to pass to the remove-items method. Should include either 'memory_id' (episodic_id or semantic_id) or 'user_id'. - May include 'session_id', 'agent_id', 'group_id', 'project_id', 'org_id'. + May include 'session_id', 'agent_id', 'project_id', 'org_id'. For memory_id deletion, may include 'memory_type' ('episodic' or 'semantic'). """ if "memory_id" in kwargs: @@ -472,7 +470,6 @@ async def remove_items(self, **kwargs) -> None: user_id = kwargs.pop("user_id", None) session_id = kwargs.pop("session_id", "default_session") agent_id = kwargs.pop("agent_id", "default_agent") - group_id = kwargs.pop("group_id", "default") project_id = kwargs.pop("project_id", None) org_id = kwargs.pop("org_id", None) @@ -484,7 +481,7 @@ async def remove_items(self, **kwargs) -> None: def delete_memory(): memory = self._get_memory_instance( - user_id, session_id, agent_id, group_id, project_id, org_id + user_id, session_id, agent_id, project_id, org_id ) # Use MemMachine SDK to delete specific memory # API: memory.delete_episodic(episodic_id) or memory.delete_semantic(semantic_id) @@ -499,7 +496,6 @@ def delete_memory(): user_id = kwargs.pop("user_id") session_id = kwargs.pop("session_id", "default_session") agent_id = kwargs.pop("agent_id", "default_agent") - group_id = kwargs.pop("group_id", "default") project_id = kwargs.pop("project_id", None) org_id = kwargs.pop("org_id", None) # Note: delete_semantic_memory flag is not yet implemented for bulk deletion diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py index 75f80ab52e..78573d2203 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py @@ -47,7 +47,7 @@ class MemMachineMemoryClientConfig(MemoryBaseConfig, RetryMixin, name="memmachin @register_memory(config_type=MemMachineMemoryClientConfig) async def memmachine_memory_client( config: MemMachineMemoryClientConfig, - builder: Builder, # Required by @register_memory contract, unused here + _builder: Builder, # Required by @register_memory contract ) -> AsyncGenerator[MemoryEditor, None]: from .memmachine_editor import MemMachineEditor # Import and initialize the MemMachine Python SDK @@ -55,9 +55,9 @@ async def memmachine_memory_client( from memmachine import MemMachineClient except ImportError as e: raise ImportError( - f"Could not import MemMachineClient from memmachine package. " + f"Could not import MemMachineClient from memmachine-server package. " f"Error: {e}. " - "Please ensure memmachine package is installed: pip install memmachine. " + "Please ensure memmachine-server package is installed: pip install memmachine-server. " "See https://github.com/MemMachine/MemMachine/blob/main/docs/examples/python.mdx " "for installation instructions." ) from e diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py index ad46eb74c2..6b6e5e15e9 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py @@ -181,13 +181,11 @@ async def test_add_conversation_calls_add_with_correct_parameters( await editor_with_spy.add_items([item]) - # Verify project.memory was called with correct parameters api_spy.assert_called_with( 'project.memory', user_id="user123", session_id="session1", - agent_id="agent1", - group_id="default" + agent_id="agent1" ) # Verify add was called twice (once per message) @@ -202,14 +200,11 @@ async def test_add_conversation_calls_add_with_correct_parameters( assert user_call is not None, "Should have a call with role='user'" assert user_call['kwargs']['content'] == "I like pizza" assert user_call['kwargs']['role'] == "user" - # Now uses memory_types instead of episode_type assert user_call['kwargs']['episode_type'] is None assert 'memory_types' in user_call['kwargs'] assert 'tags' in user_call['kwargs'].get('metadata', {}) - # MemMachine SDK expects tags as comma-separated string assert user_call['kwargs']['metadata']['tags'] == "food, preference" - # Verify second call (assistant message) assistant_call = next( (c for c in add_calls if c['kwargs'].get('role') == 'assistant'), None @@ -217,7 +212,6 @@ async def test_add_conversation_calls_add_with_correct_parameters( assert assistant_call is not None, "Should have a call with role='assistant'" assert assistant_call['kwargs']['content'] == "Great! What's your favorite topping?" assert assistant_call['kwargs']['role'] == "assistant" - # Now uses memory_types instead of episode_type assert assistant_call['kwargs']['episode_type'] is None assert 'memory_types' in assistant_call['kwargs'] @@ -240,17 +234,14 @@ async def test_add_direct_memory_calls_add_with_both_types( await editor_with_spy.add_items([item]) - # Verify add was called with both memory types add_calls = api_spy.get_calls('add') assert len(add_calls) == 1 assert add_calls[0]['kwargs']['content'] == "User prefers working in the morning" assert add_calls[0]['kwargs']['role'] == "user" assert add_calls[0]['kwargs']['episode_type'] is None - # Verify memory_types contains both Episodic and Semantic memory_types = add_calls[0]['kwargs']['memory_types'] assert len(memory_types) == 2, "Should have both episodic and semantic memory types" - # Verify metadata includes tags (as comma-separated string) assert add_calls[0]['kwargs']['metadata']['tags'] == "preference" async def test_add_conversation_memory_calls_add_with_both_types( @@ -272,13 +263,11 @@ async def test_add_conversation_memory_calls_add_with_both_types( await editor_with_spy.add_items([item]) - # Verify add was called with both memory types add_calls = api_spy.get_calls('add') assert len(add_calls) == 1 assert add_calls[0]['kwargs']['content'] == "Hello" assert add_calls[0]['kwargs']['role'] == "user" assert add_calls[0]['kwargs']['episode_type'] is None - # Verify memory_types contains both Episodic and Semantic memory_types = add_calls[0]['kwargs']['memory_types'] assert len(memory_types) == 2, "Should have both episodic and semantic memory types" @@ -302,7 +291,6 @@ async def test_add_with_custom_project_org_calls_get_or_create_project( await editor_with_spy.add_items([item]) - # Verify get_or_create_project was called with custom IDs api_spy.assert_called_with( 'get_or_create_project', org_id="custom_org", @@ -376,8 +364,7 @@ async def test_search_calls_memory_search_with_correct_parameters( 'project.memory', user_id="user123", session_id="session1", - agent_id="agent1", - group_id="default" + agent_id="agent1" ) # Verify search was called with correct parameters diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py index 2ef531bad1..3092404e6c 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py @@ -128,8 +128,7 @@ async def test_add_items_with_conversation( mock_project.memory.assert_called_once_with( user_id="user123", agent_id="agent789", - session_id="session456", - group_id="default" + session_id="session456" ) # Verify add was called for each message in conversation diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py index c3b7fc189f..f54261d0cd 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py @@ -27,6 +27,7 @@ import uuid import pytest +import requests from nat.builder.workflow_builder import WorkflowBuilder from nat.data_models.config import GeneralConfig @@ -47,17 +48,10 @@ def memmachine_base_url_fixture(fail_missing: bool = False) -> str: base_url = f"http://{base_url}" try: - # Try to import and use MemMachineClient to check server availability - from memmachine import MemMachineClient - - client = MemMachineClient(base_url=base_url, timeout=5.0) - client.health_check(timeout=5.0) + # Check if server is available via health endpoint + response = requests.get(f"{base_url}/api/v2/health", timeout=5) + response.raise_for_status() return base_url - except ImportError: - reason = "memmachine package not installed. Install with: pip install memmachine" - if fail_missing: - raise RuntimeError(reason) from None - pytest.skip(reason=reason) except Exception: reason = f"Unable to connect to MemMachine server at {base_url}. Please ensure the server is running." if fail_missing: From c0dfc7aca61bd72c7bfc70f034396e21bc349266 Mon Sep 17 00:00:00 2001 From: Charlie-Yi-2002 Date: Mon, 26 Jan 2026 11:05:26 -0800 Subject: [PATCH 04/18] Update packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Charlie-Yi-2002 --- packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py index a4f50cdb5f..ea09f55765 100644 --- a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -198,7 +198,7 @@ async def test_add_and_retrieve(): except Exception as e: print(f"\nāœ— Error: {e}") - logger.exception("Error during test execution") + logger.error("Error during test execution", exc_info=True) raise From 2cd9a7ed12756e089605fab6b7cbdb5a8296ba4c Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Mon, 26 Jan 2026 11:52:44 -0800 Subject: [PATCH 05/18] add return hints and fix exception handling Signed-off-by: Charlie Yi --- .../nat/plugins/memmachine/memmachine_editor.py | 14 +++++++------- .../tests/test_memmachine_integration.py | 9 +++++++-- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py index 97702ab4ec..23c7b6e0fd 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memmachine_editor.py @@ -28,17 +28,17 @@ class MemMachineEditor(MemoryEditor): """ - Wrapper class that implements NAT interfaces for MemMachine Integrations. - Uses the MemMachine Python SDK as documented at: + Wrapper class that implements `nat` interfaces for `MemMachine` integrations. + Uses the `MemMachine` Python SDK (`MemMachineClient`) as documented at: https://github.com/MemMachine/MemMachine/blob/main/docs/examples/python.mdx Supports both episodic and semantic memory through the unified SDK interface. - User needs to add MemMachine SDK ids as metadata to the MemoryItem: - - session_id - - agent_id - - project_id - - org_id + User needs to add `MemMachine` SDK ids as metadata to the MemoryItem: + - `session_id` + - `agent_id` + - `project_id` + - `org_id` """ def __init__(self, memmachine_instance: Any): diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py index f54261d0cd..cb944cd3a7 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py @@ -60,7 +60,7 @@ def memmachine_base_url_fixture(fail_missing: bool = False) -> str: @pytest.fixture(name="test_config") -def test_config_fixture(memmachine_base_url: str): +def test_config_fixture(memmachine_base_url: str) -> MemMachineMemoryClientConfig: """Create a test configuration.""" # Use unique org/project IDs for each test run to avoid conflicts test_id = str(uuid.uuid4())[:8] @@ -74,13 +74,14 @@ def test_config_fixture(memmachine_base_url: str): @pytest.fixture(name="test_user_id") -def test_user_id_fixture(): +def test_user_id_fixture() -> str: """Generate a unique user ID for testing.""" return f"test_user_{uuid.uuid4().hex[:8]}" @pytest.mark.integration @pytest.mark.slow +@pytest.mark.asyncio async def test_add_and_retrieve_conversation_memory( test_config: MemMachineMemoryClientConfig, test_user_id: str @@ -151,6 +152,7 @@ async def test_add_and_retrieve_conversation_memory( @pytest.mark.integration @pytest.mark.slow +@pytest.mark.asyncio async def test_add_and_retrieve_direct_memory( test_config: MemMachineMemoryClientConfig, test_user_id: str @@ -233,6 +235,7 @@ async def test_add_and_retrieve_direct_memory( @pytest.mark.integration @pytest.mark.slow +@pytest.mark.asyncio async def test_add_multiple_and_retrieve_all( test_config: MemMachineMemoryClientConfig, test_user_id: str @@ -290,6 +293,7 @@ async def test_add_multiple_and_retrieve_all( @pytest.mark.integration @pytest.mark.slow +@pytest.mark.asyncio async def test_add_and_verify_conversation_content_match( test_config: MemMachineMemoryClientConfig, test_user_id: str @@ -357,6 +361,7 @@ async def test_add_and_verify_conversation_content_match( @pytest.mark.integration @pytest.mark.slow +@pytest.mark.asyncio async def test_conversation_and_direct_memory_both_retrievable( test_config: MemMachineMemoryClientConfig, test_user_id: str From cd16e96342aacf52333ebd01de8d590bb4f1c6a3 Mon Sep 17 00:00:00 2001 From: Charlie-Yi-2002 Date: Mon, 26 Jan 2026 12:30:33 -0800 Subject: [PATCH 06/18] Update packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Charlie-Yi-2002 --- packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py index 3092404e6c..a43f3c3c3d 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); From ad52f840ddb760f2fb02ae2984277c0ef3cca132 Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Mon, 26 Jan 2026 12:36:55 -0800 Subject: [PATCH 07/18] add async decerator Signed-off-by: Charlie Yi --- packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py index ea09f55765..f631f24d85 100644 --- a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -203,6 +203,7 @@ async def test_add_and_retrieve(): @pytest.mark.integration +@pytest.mark.asyncio async def test_add_and_retrieve_integration(): """Integration test for adding and retrieving memories.""" await test_add_and_retrieve() From 4358bd87ddc38ed9e0fa82e3e854ae3086bfeeba Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Tue, 27 Jan 2026 12:06:44 -0800 Subject: [PATCH 08/18] add license, change default port to 8095, skip tests if memmachine not running Signed-off-by: Charlie Yi --- packages/nvidia_nat_memmachine/README.md | 10 ++++---- .../memmachine_memory_example.ipynb | 8 +++--- .../src/nat/plugins/memmachine/__init__.py | 14 +++++++++++ .../src/nat/plugins/memmachine/memory.py | 4 +-- .../tests/test_add_and_retrieve.py | 25 ++++++++++++++++--- .../tests/test_memmachine_api_calls.py | 2 +- .../tests/test_memmachine_editor.py | 2 +- .../tests/test_memmachine_integration.py | 6 ++--- .../tests/test_memory.py | 10 ++++---- 9 files changed, 57 insertions(+), 24 deletions(-) diff --git a/packages/nvidia_nat_memmachine/README.md b/packages/nvidia_nat_memmachine/README.md index d0d032450f..8c56e45288 100644 --- a/packages/nvidia_nat_memmachine/README.md +++ b/packages/nvidia_nat_memmachine/README.md @@ -43,11 +43,11 @@ memmachine-configure The wizard will guide you through setting up: -- **Neo4j Database**: Option to install Neo4j automatically or provide connection details for an existing instance +- **Neo4j Database**: Option to install Neo4j automatically or provide connection details for an existing instance. If you enter nothing, Neo4j is installed on your local disk by default. - **Large Language Model (LLM) Provider**: Choose from supported providers like OpenAI, AWS Bedrock, or Ollama -- **Model Selection**: Select specific LLM and embedding models +- **Model Selection**: Select specific LLM and embedding models. The default is OpenAI with `gpt-4o-mini` and `text-embedding-3-small`. - **API Keys and Credentials**: Input necessary API keys for your selected LLM provider -- **Server Settings**: Configure server host and port +- **Server Settings**: Configure server host and port. The default is `localhost:8095`. **Note**: - The wizard installs Neo4j and Java automatically when you choose to install Neo4j (platform-specific: Windows uses ZIP, macOS uses brew, Linux uses tar.gz) @@ -62,7 +62,7 @@ After completing the configuration wizard, start the server: memmachine-server ``` -The server will start on `http://localhost:8080` by default (or the port you configured). +The server will start on `http://localhost:8095` by default (or the port you configured). For more details, see the [MemMachine Configuration Wizard Documentation](https://docs.memmachine.ai/open_source/configuration-wizard). @@ -73,7 +73,7 @@ Add MemMachine memory to your workflow configuration: ```yaml memory: memmachine_memory: - base_url: "http://localhost:8080" # MemMachine server URL + base_url: "http://localhost:8095" # MemMachine server URL org_id: "my_org" # Optional: default organization ID project_id: "my_project" # Optional: default project ID ``` diff --git a/packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb b/packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb index d6b9f8128b..c8ab2346f9 100644 --- a/packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb +++ b/packages/nvidia_nat_memmachine/memmachine_memory_example.ipynb @@ -83,7 +83,7 @@ "- **LLM Provider**: Choose from OpenAI, AWS Bedrock, Ollama, etc\n", "- **Model Selection**: Select LLM and embedding models\n", "- **API Keys**: Input necessary credentials\n", - "- **Server Settings**: Configure host and port (default: localhost:8080)\n", + "- **Server Settings**: Configure host and port (default: localhost:8095)\n", "\n", "Configuration file will be saved at: `~/.config/memmachine/cfg.yml`\n", "\n", @@ -93,7 +93,7 @@ "memmachine-server\n", "```\n", "\n", - "The server will start on `http://localhost:8080` by default.\n", + "The server will start on `http://localhost:8095` by default.\n", "\n", "For more details, see the [MemMachine Documentation](https://docs.memmachine.ai/open_source/configuration-wizard)." ] @@ -387,7 +387,7 @@ "import requests\n", "\n", "# MemMachine server URL (default)\n", - "MEMMACHINE_BASE_URL = os.environ.get(\"MEMMACHINE_BASE_URL\", \"http://localhost:8080\")\n", + "MEMMACHINE_BASE_URL = os.environ.get(\"MEMMACHINE_BASE_URL\", \"http://localhost:8095\")\n", "\n", "try:\n", " response = requests.get(f\"{MEMMACHINE_BASE_URL}/api/v2/health\", timeout=5)\n", @@ -683,7 +683,7 @@ "memory:\n", " memmachine_memory:\n", " _type: memmachine_memory\n", - " base_url: \"http://localhost:8080\"\n", + " base_url: \"http://localhost:8095\"\n", " org_id: \"{memmachine_config.org_id}\"\n", " project_id: \"{memmachine_config.project_id}\"\n", "\n", diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/__init__.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/__init__.py index e69de29bb2..b16841385a 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/__init__.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-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. \ No newline at end of file diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py index 78573d2203..d8f7e7ce57 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py @@ -37,7 +37,7 @@ class MemMachineMemoryClientConfig(MemoryBaseConfig, RetryMixin, name="memmachin LLM API keys (e.g., OpenAI) are configured in the MemMachine cfg.yml file, not in this client configuration. """ - base_url: str # Base URL of the MemMachine server (e.g., "http://localhost:8080") + base_url: str # Base URL of the MemMachine server (e.g., "http://localhost:8095") org_id: str | None = None # Optional default organization ID project_id: str | None = None # Optional default project ID timeout: int = 30 # Request timeout in seconds @@ -64,7 +64,7 @@ async def memmachine_memory_client( # Initialize MemMachineClient with base_url # This follows the documented SDK pattern for local instances: - # client = MemMachineClient(base_url="http://localhost:8080") + # client = MemMachineClient(base_url="http://localhost:8095") # Note: api_key is not needed for local/self-hosted MemMachine instances try: client = MemMachineClient( diff --git a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py index f631f24d85..96c7f705cf 100644 --- a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -35,6 +35,7 @@ from datetime import datetime import pytest +import requests from nat.builder.builder import Builder from nat.memory.models import MemoryItem @@ -43,12 +44,30 @@ logger = logging.getLogger(__name__) +def _memmachine_available(base_url: str) -> bool: + """Return True if MemMachine server is reachable.""" + if not base_url.startswith("http"): + base_url = f"http://{base_url}" + try: + response = requests.get(f"{base_url}/api/v2/health", timeout=5) + response.raise_for_status() + return True + except Exception: + return False + + async def test_add_and_retrieve(): - """Test adding memories and retrieving them.""" + """Test adding memories and retrieving them. Skips if MemMachine server is not running.""" # Configuration - base_url = os.environ.get("MEMMACHINE_BASE_URL", "http://localhost:8080") + base_url = os.environ.get("MEMMACHINE_BASE_URL", "http://localhost:8095") + if not _memmachine_available(base_url): + pytest.skip( + f"MemMachine server not available at {base_url}. " + "Start the server or set MEMMACHINE_BASE_URL to run this test." + ) + test_id = str(uuid.uuid4())[:8] - + config = MemMachineMemoryClientConfig( base_url=base_url, org_id=f"test_org_{test_id}", diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py index 6b6e5e15e9..fefec2f3d6 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_api_calls.py @@ -149,7 +149,7 @@ def spied_get_or_create_project(*args, **kwargs): mock_client.create_project = spied_create_project mock_client.get_or_create_project = spied_get_or_create_project - mock_client.base_url = "http://localhost:8080" + mock_client.base_url = "http://localhost:8095" return mock_client diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py index a43f3c3c3d..fba4744069 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py @@ -56,7 +56,7 @@ def mock_client_fixture(mock_project): mock_client = Mock(spec=['create_project', 'get_or_create_project', 'base_url']) mock_client.create_project = Mock(return_value=mock_project) mock_client.get_or_create_project = Mock(return_value=mock_project) - mock_client.base_url = "http://localhost:8080" + mock_client.base_url = "http://localhost:8095" return mock_client diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py index cb944cd3a7..f5a4ff9f46 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_integration.py @@ -20,7 +20,7 @@ integration by adding memories and then retrieving them. The tests will automatically skip if the MemMachine server is not available. -Set `MEMMACHINE_BASE_URL` environment variable to override default (http://localhost:8080). +Set `MEMMACHINE_BASE_URL` environment variable to override default (http://localhost:8095). """ import os @@ -41,9 +41,9 @@ def memmachine_base_url_fixture(fail_missing: bool = False) -> str: Ensure MemMachine server is running and provide base URL. To run these tests, a MemMachine server must be running. - Set MEMMACHINE_BASE_URL environment variable to override default (http://localhost:8080). + Set MEMMACHINE_BASE_URL environment variable to override default (http://localhost:8095). """ - base_url = os.getenv("MEMMACHINE_BASE_URL", "http://localhost:8080") + base_url = os.getenv("MEMMACHINE_BASE_URL", "http://localhost:8095") if not base_url.startswith("http"): base_url = f"http://{base_url}" diff --git a/packages/nvidia_nat_memmachine/tests/test_memory.py b/packages/nvidia_nat_memmachine/tests/test_memory.py index b46bc4e8cb..38dbe75a20 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memory.py +++ b/packages/nvidia_nat_memmachine/tests/test_memory.py @@ -32,7 +32,7 @@ def mock_builder_fixture(): def config_fixture(): """Fixture to provide a MemMachineMemoryClientConfig instance.""" return MemMachineMemoryClientConfig( - base_url="http://localhost:8080", + base_url="http://localhost:8095", org_id="test_org", project_id="test_project", timeout=30, @@ -44,7 +44,7 @@ def config_fixture(): def config_minimal_fixture(): """Fixture to provide a minimal MemMachineMemoryClientConfig instance.""" return MemMachineMemoryClientConfig( - base_url="http://localhost:8080" + base_url="http://localhost:8095" ) @@ -52,7 +52,7 @@ def config_minimal_fixture(): def mock_memmachine_client_fixture(): """Fixture to provide a mocked MemMachineClient.""" mock_client = Mock() - mock_client.base_url = "http://localhost:8080" + mock_client.base_url = "http://localhost:8095" return mock_client @@ -161,8 +161,8 @@ async def test_memmachine_memory_client_config_validation(): MemMachineMemoryClientConfig() # Should work with base_url - config = MemMachineMemoryClientConfig(base_url="http://localhost:8080") - assert config.base_url == "http://localhost:8080" + config = MemMachineMemoryClientConfig(base_url="http://localhost:8095") + assert config.base_url == "http://localhost:8095" assert config.timeout == 30 assert config.max_retries == 3 From 4cf71dea328f0e4e6257bfd09ad08ed0bb3ae6b9 Mon Sep 17 00:00:00 2001 From: Charlie-Yi-2002 Date: Tue, 27 Jan 2026 12:51:57 -0800 Subject: [PATCH 09/18] Update packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Charlie-Yi-2002 --- .../nvidia_nat_memmachine/tests/test_memmachine_editor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py index fba4744069..d3e85d0284 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py +++ b/packages/nvidia_nat_memmachine/tests/test_memmachine_editor.py @@ -177,8 +177,8 @@ async def test_add_items_with_conversation( async def test_add_items_with_direct_memory( memmachine_editor_with_client: MemMachineEditor, - mock_client: Mock, - mock_project: Mock, + _mock_client: Mock, + _mock_project: Mock, mock_memory_instance: Mock, sample_direct_memory_item: MemoryItem ): From 86da880fbd560cdc5bdd3be4912b4267bea51257 Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Tue, 27 Jan 2026 13:08:55 -0800 Subject: [PATCH 10/18] mark tests as async Signed-off-by: Charlie Yi --- packages/nvidia_nat_memmachine/tests/test_memory.py | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/nvidia_nat_memmachine/tests/test_memory.py b/packages/nvidia_nat_memmachine/tests/test_memory.py index b46bc4e8cb..c76a9afdb4 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memory.py +++ b/packages/nvidia_nat_memmachine/tests/test_memory.py @@ -17,6 +17,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest +pytestmark = pytest.mark.asyncio from nat.builder.builder import Builder from nat.plugins.memmachine.memory import MemMachineMemoryClientConfig, memmachine_memory_client From 302354453057ebb31d681303775e70ea156a3c38 Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Tue, 27 Jan 2026 13:15:59 -0800 Subject: [PATCH 11/18] rename helper function to avoid double test collection Signed-off-by: Charlie Yi --- .../nvidia_nat_memmachine/tests/test_add_and_retrieve.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py index 96c7f705cf..9841b6db6e 100644 --- a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -56,8 +56,8 @@ def _memmachine_available(base_url: str) -> bool: return False -async def test_add_and_retrieve(): - """Test adding memories and retrieving them. Skips if MemMachine server is not running.""" +async def _run_add_and_retrieve(): + """Run add-and-retrieve flow. Skips if MemMachine server is not running. Not collected by pytest.""" # Configuration base_url = os.environ.get("MEMMACHINE_BASE_URL", "http://localhost:8095") if not _memmachine_available(base_url): @@ -225,8 +225,8 @@ async def test_add_and_retrieve(): @pytest.mark.asyncio 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()) + asyncio.run(_run_add_and_retrieve()) From 5351e16117aa8ffad2c0b52073dca6fe968a9dcb Mon Sep 17 00:00:00 2001 From: Charlie-Yi-2002 Date: Tue, 27 Jan 2026 13:22:41 -0800 Subject: [PATCH 12/18] Update packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Signed-off-by: Charlie-Yi-2002 --- packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py index 9841b6db6e..6e2d86c390 100644 --- a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -20,7 +20,7 @@ This script demonstrates the full integration: 1. Adds memories using the NAT integration 2. Retrieves them back -3. Verifies they match +3. Prints retrieved results to confirm API calls succeed Usage: python tests/test_add_and_retrieve.py From 6f406d05e2489d53f0dfad3199ba4b2e8a5ff163 Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Wed, 28 Jan 2026 11:32:44 -0800 Subject: [PATCH 13/18] add async slow decorators and fix error handling for ModuleNotFound Signed-off-by: Charlie Yi --- .../tests/test_add_and_retrieve.py | 16 +++++++++------- .../nvidia_nat_memmachine/tests/test_memory.py | 2 +- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py index 6e2d86c390..9692150ab9 100644 --- a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -35,7 +35,7 @@ from datetime import datetime import pytest -import requests +import httpx from nat.builder.builder import Builder from nat.memory.models import MemoryItem @@ -44,15 +44,16 @@ logger = logging.getLogger(__name__) -def _memmachine_available(base_url: str) -> bool: +async def _memmachine_available(base_url: str) -> bool: """Return True if MemMachine server is reachable.""" if not base_url.startswith("http"): base_url = f"http://{base_url}" try: - response = requests.get(f"{base_url}/api/v2/health", timeout=5) - response.raise_for_status() - return True - except Exception: + async with httpx.AsyncClient(timeout=5.0) as client: + response = await client.get(f"{base_url}/api/v2/health") + response.raise_for_status() + return True + except httpx.RequestError: return False @@ -60,7 +61,7 @@ async def _run_add_and_retrieve(): """Run add-and-retrieve flow. Skips if MemMachine server is not running. Not collected by pytest.""" # Configuration base_url = os.environ.get("MEMMACHINE_BASE_URL", "http://localhost:8095") - if not _memmachine_available(base_url): + if not await _memmachine_available(base_url): pytest.skip( f"MemMachine server not available at {base_url}. " "Start the server or set MEMMACHINE_BASE_URL to run this test." @@ -223,6 +224,7 @@ async def _run_add_and_retrieve(): @pytest.mark.integration @pytest.mark.asyncio +@pytest.mark.slow async def test_add_and_retrieve_integration(): """Integration test for adding and retrieving memories.""" await _run_add_and_retrieve() diff --git a/packages/nvidia_nat_memmachine/tests/test_memory.py b/packages/nvidia_nat_memmachine/tests/test_memory.py index 353ca33458..1d356266bb 100644 --- a/packages/nvidia_nat_memmachine/tests/test_memory.py +++ b/packages/nvidia_nat_memmachine/tests/test_memory.py @@ -114,7 +114,7 @@ async def test_memmachine_memory_client_import_error( def import_side_effect(name, *args, **kwargs): if name == "memmachine": - raise ImportError("No module named 'memmachine'") + raise ModuleNotFoundError return original_import(name, *args, **kwargs) with patch("builtins.__import__", side_effect=import_side_effect): From 504cf1c4de8490c2992c80d6762d308623a6c76b Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Wed, 28 Jan 2026 14:39:00 -0800 Subject: [PATCH 14/18] update memmachine-server to latest 0.2.5 Signed-off-by: Charlie Yi --- packages/nvidia_nat_memmachine/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/nvidia_nat_memmachine/pyproject.toml b/packages/nvidia_nat_memmachine/pyproject.toml index e0fdd181da..62c6d34eb5 100644 --- a/packages/nvidia_nat_memmachine/pyproject.toml +++ b/packages/nvidia_nat_memmachine/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ # 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.5", - "memmachine-server~=0.2.1", + "memmachine-server~=0.2.5", ] requires-python = ">=3.11,<3.14" description = "Subpackage for MemMachine integration in NeMo Agent toolkit. Requires a cfg.yml configuration file with database and AI model settings." From 09960eb13a0bbb5957242b222865593a6ce05b4d Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Wed, 28 Jan 2026 14:44:45 -0800 Subject: [PATCH 15/18] clarify default port in README Signed-off-by: Charlie Yi --- packages/nvidia_nat_memmachine/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/nvidia_nat_memmachine/README.md b/packages/nvidia_nat_memmachine/README.md index 8c56e45288..e624b46495 100644 --- a/packages/nvidia_nat_memmachine/README.md +++ b/packages/nvidia_nat_memmachine/README.md @@ -47,7 +47,7 @@ The wizard will guide you through setting up: - **Large Language Model (LLM) Provider**: Choose from supported providers like OpenAI, AWS Bedrock, or Ollama - **Model Selection**: Select specific LLM and embedding models. The default is OpenAI with `gpt-4o-mini` and `text-embedding-3-small`. - **API Keys and Credentials**: Input necessary API keys for your selected LLM provider -- **Server Settings**: Configure server host and port. The default is `localhost:8095`. +- **Server Settings**: Configure server host and port. The default is `localhost:8080` but it is recommended to bind to port `8095` as `8080` is a commonly used port. **Note**: - The wizard installs Neo4j and Java automatically when you choose to install Neo4j (platform-specific: Windows uses ZIP, macOS uses brew, Linux uses tar.gz) From 8529fc09807bc8ed51250d014d79c12ef3dd4f18 Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Wed, 28 Jan 2026 14:50:37 -0800 Subject: [PATCH 16/18] add license header Signed-off-by: Charlie Yi --- packages/nvidia_nat_memmachine/README.md | 15 +++++++++++++++ .../tests/test_add_and_retrieve.py | 1 - 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/packages/nvidia_nat_memmachine/README.md b/packages/nvidia_nat_memmachine/README.md index e624b46495..a877180628 100644 --- a/packages/nvidia_nat_memmachine/README.md +++ b/packages/nvidia_nat_memmachine/README.md @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-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. + # NVIDIA NeMo Agent toolkit - MemMachine Integration This package provides integration with MemMachine for memory management in NeMo Agent toolkit. diff --git a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py index 9692150ab9..9c64da3529 100644 --- a/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py +++ b/packages/nvidia_nat_memmachine/tests/test_add_and_retrieve.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # From 6692f7755e461f624b40288d33a9dce9c9dbe64b Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Wed, 28 Jan 2026 14:53:34 -0800 Subject: [PATCH 17/18] style fixes Signed-off-by: Charlie Yi --- packages/nvidia_nat_memmachine/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/nvidia_nat_memmachine/README.md b/packages/nvidia_nat_memmachine/README.md index a877180628..25d6a9e122 100644 --- a/packages/nvidia_nat_memmachine/README.md +++ b/packages/nvidia_nat_memmachine/README.md @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# NVIDIA NeMo Agent toolkit - MemMachine Integration +# NVIDIA NeMo Agent Toolkit - MemMachine Integration This package provides integration with MemMachine for memory management in NeMo Agent toolkit. @@ -77,7 +77,7 @@ After completing the configuration wizard, start the server: memmachine-server ``` -The server will start on `http://localhost:8095` by default (or the port you configured). +The server will start on `http://localhost:8080` by default (or the port you configured with the configuration wizard). For more details, see the [MemMachine Configuration Wizard Documentation](https://docs.memmachine.ai/open_source/configuration-wizard). From 3b2e00506927839414eec475d89f07de297c122c Mon Sep 17 00:00:00 2001 From: Charlie Yi Date: Thu, 29 Jan 2026 15:50:17 -0800 Subject: [PATCH 18/18] change dependency to memmachine-client at 0.2.5, update docs to list memmachine-server as optional to support a local deployment Signed-off-by: Charlie Yi --- packages/nvidia_nat_memmachine/README.md | 6 ++++-- packages/nvidia_nat_memmachine/pyproject.toml | 2 +- .../src/nat/plugins/memmachine/memory.py | 4 ++-- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/packages/nvidia_nat_memmachine/README.md b/packages/nvidia_nat_memmachine/README.md index 25d6a9e122..28dbba20d6 100644 --- a/packages/nvidia_nat_memmachine/README.md +++ b/packages/nvidia_nat_memmachine/README.md @@ -24,7 +24,7 @@ MemMachine is a unified memory management system that supports both episodic and ## Prerequisites - Python 3.11+ -- MemMachine server (install via `pip install memmachine-server`) +- **memmachine-client** is installed automatically with this package. You need a running MemMachine instance (local or hosted) to connect to. To run a local instance, install and configure **memmachine-server** separately (see [MemMachine Server Setup](#memmachine-server-setup) below). ## Installation @@ -42,7 +42,9 @@ uv pip install -e packages/nvidia_nat_memmachine ## MemMachine Server Setup -### Step 1: Install MemMachine Server +This section is optional. Only follow these steps if you want to run a **local** MemMachine instance. If you use a hosted MemMachine instance, configure `base_url` (and any auth) in your workflow config and skip this section. + +### Step 1: Install MemMachine Server (local instance only) ```bash pip install memmachine-server diff --git a/packages/nvidia_nat_memmachine/pyproject.toml b/packages/nvidia_nat_memmachine/pyproject.toml index 62c6d34eb5..ed2d891a5c 100644 --- a/packages/nvidia_nat_memmachine/pyproject.toml +++ b/packages/nvidia_nat_memmachine/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ # 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.5", - "memmachine-server~=0.2.5", + "memmachine-client~=0.2.5", ] requires-python = ">=3.11,<3.14" description = "Subpackage for MemMachine integration in NeMo Agent toolkit. Requires a cfg.yml configuration file with database and AI model settings." diff --git a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py index d8f7e7ce57..0cedde4ee8 100644 --- a/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py +++ b/packages/nvidia_nat_memmachine/src/nat/plugins/memmachine/memory.py @@ -55,9 +55,9 @@ async def memmachine_memory_client( from memmachine import MemMachineClient except ImportError as e: raise ImportError( - f"Could not import MemMachineClient from memmachine-server package. " + f"Could not import MemMachineClient from memmachine-client package. " f"Error: {e}. " - "Please ensure memmachine-server package is installed: pip install memmachine-server. " + "Please ensure memmachine-client package is installed: pip install memmachine-client. " "See https://github.com/MemMachine/MemMachine/blob/main/docs/examples/python.mdx " "for installation instructions." ) from e