From 484b7c884e64b0dc895877bf3455af9730c374af Mon Sep 17 00:00:00 2001 From: AijooseFactory Date: Sun, 1 Mar 2026 16:36:30 -0600 Subject: [PATCH 1/4] Add memory_saved_after extension hook after memory insert --- plugins/memory/helpers/memory.py | 25 ++++++++- plugins/tests/test_memory_hook.py | 87 +++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 1 deletion(-) create mode 100644 plugins/tests/test_memory_hook.py diff --git a/plugins/memory/helpers/memory.py b/plugins/memory/helpers/memory.py index c3edba92d..49afea51e 100644 --- a/plugins/memory/helpers/memory.py +++ b/plugins/memory/helpers/memory.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from datetime import datetime from typing import Any, List, Sequence from langchain.storage import InMemoryByteStore, LocalFileStore @@ -387,9 +389,29 @@ class Memory: self._save_db() # persist return rem_docs - async def insert_text(self, text, metadata: dict = {}): + async def insert_text(self, text, metadata: dict | None = None): + metadata = metadata or {} doc = Document(text, metadata=metadata) ids = await self.insert_documents([doc]) + + # Fire post-save memory hook extensions (best-effort, never breaks memory save) + try: + from python.helpers.extension import call_extensions + await call_extensions( + "memory_saved_after", + agent=getattr(self, "agent", None), + text=text, + metadata=metadata, + doc_id=ids[0], + memory_subdir=self.memory_subdir, + ) + except Exception as e: + try: + from python.helpers.print_style import PrintStyle + PrintStyle.warning(f"memory_saved_after hook failed: {type(e).__name__}") + except Exception: + pass + return ids[0] async def insert_documents(self, docs: list[Document]): @@ -570,3 +592,4 @@ def get_knowledge_subdirs_by_memory_subdir( default.append(get_project_meta(memory_subdir[9:], "knowledge")) return default + diff --git a/plugins/tests/test_memory_hook.py b/plugins/tests/test_memory_hook.py new file mode 100644 index 000000000..e65786ae3 --- /dev/null +++ b/plugins/tests/test_memory_hook.py @@ -0,0 +1,87 @@ +import sys +import pytest +from unittest.mock import patch, AsyncMock, MagicMock +from types import SimpleNamespace + +# Catch-all mock importer to avoid heavy Agent Zero dependencies locally +from importlib.machinery import ModuleSpec + +class MockLoader: + def create_module(self, spec): + if spec.name in sys.modules: + return sys.modules[spec.name] + mock = MagicMock() + # Ensure submodules can be accessed via attributes + mock.__path__ = [] + sys.modules[spec.name] = mock + return mock + + def exec_module(self, module): + pass + +class MockImporter: + def find_spec(self, fullname, path, target=None): + catch_prefixes = [ + 'langchain', 'faiss', 'simpleeval', 'webcolors', 'litellm', + 'openai', 'cryptography', 'nest_asyncio', 'whisper', 'git', + 'tiktoken', 'browser_use', 'docker', 'duckduckgo_search', 'bs4', + 'html2text', 'yaml', 'aiohttp', 'jinja2', 'markdown', 'requests', + 'sentence_transformers', 'regex', 'pydantic', 'rich', 'pymupdf', + 'playwright', 'pathspec', 'tenacity', 'dotenv' + ] + if any(fullname.startswith(p) for p in catch_prefixes): + return ModuleSpec(fullname, MockLoader(), is_package=True) + return None + +sys.meta_path.insert(0, MockImporter()) + +@pytest.fixture +def mock_memory(): + from plugins.memory.helpers.memory import Memory + # Create a dummy Memory object bypassing init to avoid Faiss overhead + mem = Memory.__new__(Memory) + mem.memory_subdir = "test_subdir" + mem.agent = SimpleNamespace(name="TestAgent") + # Mock insert_documents since we only test the post-save hook behavior + mem.insert_documents = AsyncMock(return_value=["doc-123"]) + return mem + +@pytest.mark.asyncio +async def test_memory_saved_after_hook_called(mock_memory): + text = "Hello world" + metadata = {"source": "test"} + + with patch("python.helpers.extension.call_extensions", new_callable=AsyncMock) as mock_call_ext: + doc_id = await mock_memory.insert_text(text, metadata=metadata) + + assert doc_id == "doc-123" + mock_call_ext.assert_called_once_with( + "memory_saved_after", + agent=mock_memory.agent, + text=text, + metadata=metadata, + doc_id="doc-123", + memory_subdir="test_subdir" + ) + +@pytest.mark.asyncio +async def test_memory_saved_after_hook_failure_isolation(mock_memory): + text = "Hello world" + metadata = {"source": "test"} + + with patch("python.helpers.extension.call_extensions", new_callable=AsyncMock) as mock_call_ext: + # Simulate a crash in an extension + mock_call_ext.side_effect = Exception("Extension crashed") + + # Patch PrintStyle.warning to ensure it logs the error without raising + with patch("python.helpers.print_style.PrintStyle.warning") as mock_print: + doc_id = await mock_memory.insert_text(text, metadata=metadata) + + # Memory save succeeds despite extension crash + assert doc_id == "doc-123" + + # Warning was emitted + mock_print.assert_called_once() + args, _ = mock_print.call_args + assert "memory_saved_after hook failed: Exception" in args[0] + From 33d220921bb161879835815c14eab0495ac6a4f3 Mon Sep 17 00:00:00 2001 From: AijooseFactory Date: Fri, 13 Mar 2026 14:47:50 -0500 Subject: [PATCH 2/4] feat: complete memory_save_before/after hooks per collaborator feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add memory_save_before hook with mutable {object} dict - Extensions can modify text/metadata before save - If object['text'] is None, save is skipped (returns None) - Rename memory_saved_after -> memory_save_after - Remove try/catch — extensions handle their own errors - Update tests: before/after hooks, skip-on-None, modification, propagation Addresses feedback from frdel on PR #1176 --- plugins/memory/helpers/memory.py | 50 ++++++------- plugins/tests/test_memory_hook.py | 112 +++++++++++++++++++++++------- 2 files changed, 113 insertions(+), 49 deletions(-) diff --git a/plugins/memory/helpers/memory.py b/plugins/memory/helpers/memory.py index 49afea51e..0a2345ac0 100644 --- a/plugins/memory/helpers/memory.py +++ b/plugins/memory/helpers/memory.py @@ -9,8 +9,6 @@ from python.helpers import guids # from langchain_chroma import Chroma from langchain_community.vectorstores import FAISS -# faiss needs to be patched for python 3.12 on arm #TODO remove once not needed -from python.helpers import faiss_monkey_patch import faiss @@ -18,9 +16,8 @@ from langchain_community.docstore.in_memory import InMemoryDocstore from langchain_community.vectorstores.utils import ( DistanceStrategy, ) -from langchain_core.embeddings import Embeddings - -import os, json +import os +import json import numpy as np @@ -28,7 +25,7 @@ from python.helpers.print_style import PrintStyle from python.helpers import files, plugins, projects from langchain_core.documents import Document from . import knowledge_import -from python.helpers.log import Log, LogItem +from python.helpers.log import LogItem from enum import Enum from agent import Agent, AgentContext import models @@ -390,27 +387,32 @@ class Memory: return rem_docs async def insert_text(self, text, metadata: dict | None = None): + from python.helpers.extension import call_extensions + metadata = metadata or {} - doc = Document(text, metadata=metadata) + + # memory_save_before: pass mutable object so extensions can edit or skip + obj = {"text": text, "metadata": metadata, "memory_subdir": self.memory_subdir} + await call_extensions( + "memory_save_before", + agent=getattr(self, "agent", None), + object=obj, + ) + + # If an extension set text to None, skip the save + if obj["text"] is None: + return None + + doc = Document(obj["text"], metadata=obj["metadata"]) ids = await self.insert_documents([doc]) - # Fire post-save memory hook extensions (best-effort, never breaks memory save) - try: - from python.helpers.extension import call_extensions - await call_extensions( - "memory_saved_after", - agent=getattr(self, "agent", None), - text=text, - metadata=metadata, - doc_id=ids[0], - memory_subdir=self.memory_subdir, - ) - except Exception as e: - try: - from python.helpers.print_style import PrintStyle - PrintStyle.warning(f"memory_saved_after hook failed: {type(e).__name__}") - except Exception: - pass + # memory_save_after: notify extensions after successful persist + obj["doc_id"] = ids[0] + await call_extensions( + "memory_save_after", + agent=getattr(self, "agent", None), + object=obj, + ) return ids[0] diff --git a/plugins/tests/test_memory_hook.py b/plugins/tests/test_memory_hook.py index e65786ae3..4b0bf21f8 100644 --- a/plugins/tests/test_memory_hook.py +++ b/plugins/tests/test_memory_hook.py @@ -1,6 +1,6 @@ import sys import pytest -from unittest.mock import patch, AsyncMock, MagicMock +from unittest.mock import patch, AsyncMock, MagicMock, call from types import SimpleNamespace # Catch-all mock importer to avoid heavy Agent Zero dependencies locally @@ -42,46 +42,108 @@ def mock_memory(): mem = Memory.__new__(Memory) mem.memory_subdir = "test_subdir" mem.agent = SimpleNamespace(name="TestAgent") - # Mock insert_documents since we only test the post-save hook behavior + # Mock insert_documents since we only test the hook behavior mem.insert_documents = AsyncMock(return_value=["doc-123"]) return mem + @pytest.mark.asyncio -async def test_memory_saved_after_hook_called(mock_memory): +async def test_memory_save_before_called_with_object(mock_memory): + """memory_save_before receives a mutable {object} dict.""" text = "Hello world" metadata = {"source": "test"} with patch("python.helpers.extension.call_extensions", new_callable=AsyncMock) as mock_call_ext: doc_id = await mock_memory.insert_text(text, metadata=metadata) - + assert doc_id == "doc-123" - mock_call_ext.assert_called_once_with( - "memory_saved_after", + # memory_save_before is the first call + before_call = mock_call_ext.call_args_list[0] + assert before_call == call( + "memory_save_before", agent=mock_memory.agent, - text=text, - metadata=metadata, - doc_id="doc-123", - memory_subdir="test_subdir" + object={"text": text, "metadata": metadata, "memory_subdir": "test_subdir"}, ) + @pytest.mark.asyncio -async def test_memory_saved_after_hook_failure_isolation(mock_memory): +async def test_memory_save_after_called_with_doc_id(mock_memory): + """memory_save_after receives the object with doc_id after persist.""" text = "Hello world" metadata = {"source": "test"} with patch("python.helpers.extension.call_extensions", new_callable=AsyncMock) as mock_call_ext: - # Simulate a crash in an extension - mock_call_ext.side_effect = Exception("Extension crashed") - - # Patch PrintStyle.warning to ensure it logs the error without raising - with patch("python.helpers.print_style.PrintStyle.warning") as mock_print: - doc_id = await mock_memory.insert_text(text, metadata=metadata) - - # Memory save succeeds despite extension crash - assert doc_id == "doc-123" - - # Warning was emitted - mock_print.assert_called_once() - args, _ = mock_print.call_args - assert "memory_saved_after hook failed: Exception" in args[0] + doc_id = await mock_memory.insert_text(text, metadata=metadata) + assert doc_id == "doc-123" + # memory_save_after is the second call + after_call = mock_call_ext.call_args_list[1] + assert after_call == call( + "memory_save_after", + agent=mock_memory.agent, + object={ + "text": text, + "metadata": metadata, + "memory_subdir": "test_subdir", + "doc_id": "doc-123", + }, + ) + + +@pytest.mark.asyncio +async def test_memory_save_skipped_when_text_none(mock_memory): + """Save is skipped when memory_save_before sets object['text'] to None.""" + text = "Hello world" + metadata = {"source": "test"} + + async def nullify_text(*args, **kwargs): + # Simulate an extension setting text to None + obj = kwargs.get("object") + if obj is not None: + obj["text"] = None + + with patch("python.helpers.extension.call_extensions", new_callable=AsyncMock) as mock_call_ext: + mock_call_ext.side_effect = nullify_text + doc_id = await mock_memory.insert_text(text, metadata=metadata) + + # Save was skipped + assert doc_id is None + # insert_documents was never called + mock_memory.insert_documents.assert_not_called() + # Only memory_save_before was called (no after) + assert mock_call_ext.call_count == 1 + + +@pytest.mark.asyncio +async def test_memory_save_before_can_modify_text(mock_memory): + """Extensions can modify the text via memory_save_before.""" + text = "Original" + metadata = {"source": "test"} + + async def modify_text(*args, **kwargs): + obj = kwargs.get("object") + if obj is not None and obj.get("text") == "Original": + obj["text"] = "Modified by extension" + + with patch("python.helpers.extension.call_extensions", new_callable=AsyncMock) as mock_call_ext: + mock_call_ext.side_effect = modify_text + doc_id = await mock_memory.insert_text(text, metadata=metadata) + + assert doc_id == "doc-123" + # Verify insert_documents was called with the modified text + call_args = mock_memory.insert_documents.call_args + doc = call_args[0][0][0] + assert doc.page_content == "Modified by extension" + + +@pytest.mark.asyncio +async def test_extension_exceptions_propagate(mock_memory): + """No try/catch — extension errors propagate to the caller.""" + text = "Hello world" + metadata = {"source": "test"} + + with patch("python.helpers.extension.call_extensions", new_callable=AsyncMock) as mock_call_ext: + mock_call_ext.side_effect = RuntimeError("Extension crashed") + + with pytest.raises(RuntimeError, match="Extension crashed"): + await mock_memory.insert_text(text, metadata=metadata) From 59d75f4f722df79dc4634c156339cd591c5dc228 Mon Sep 17 00:00:00 2001 From: AijooseFactory Date: Fri, 13 Mar 2026 17:10:16 -0500 Subject: [PATCH 3/4] fix: rename call_extensions to call_extensions_async for API alignment --- plugins/_memory/helpers/memory.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/plugins/_memory/helpers/memory.py b/plugins/_memory/helpers/memory.py index f4df462f3..332979af6 100644 --- a/plugins/_memory/helpers/memory.py +++ b/plugins/_memory/helpers/memory.py @@ -389,13 +389,13 @@ class Memory: return rem_docs async def insert_text(self, text, metadata: dict | None = None): - from helpers.extension import call_extensions + from helpers.extension import call_extensions_async metadata = metadata or {} # memory_save_before: pass mutable object so extensions can edit or skip obj = {"text": text, "metadata": metadata, "memory_subdir": self.memory_subdir} - await call_extensions( + await call_extensions_async( "memory_save_before", agent=getattr(self, "agent", None), object=obj, @@ -410,7 +410,7 @@ class Memory: # memory_save_after: notify extensions after successful persist obj["doc_id"] = ids[0] - await call_extensions( + await call_extensions_async( "memory_save_after", agent=getattr(self, "agent", None), object=obj, From f98ace6d28c3a0e0d9c50337ab8fa6fd5368cf29 Mon Sep 17 00:00:00 2001 From: AijooseFactory Date: Fri, 13 Mar 2026 17:17:23 -0500 Subject: [PATCH 4/4] ci: add Pytest CI workflow to verify memory hooks --- .github/workflows/pytest.yml | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 .github/workflows/pytest.yml diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml new file mode 100644 index 000000000..732411249 --- /dev/null +++ b/.github/workflows/pytest.yml @@ -0,0 +1,28 @@ +name: Pytest CI + +on: + push: + branches: [ feature/memory-saved-after-hook ] + pull_request: + branches: [ development ] + +jobs: + test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install pytest pytest-asyncio pyyaml + + - name: Run memory hook tests + run: | + pytest tests/test_memory_hook.py