From a6e11c625c8677fb68b1b8faca4f312dfbe5814f Mon Sep 17 00:00:00 2001 From: Jacky020 Date: Wed, 23 Sep 2026 15:44:00 +0800 Subject: [PATCH] [feat-a-mem] Add a-mem policies with its test It is the reproduction according to the source repo https://github.com/WujiangXu/A-mem-sys . The retreival of neighbor note (through note links) is self implemented as not well described in the paper and the source repo. For the interpretation of the policy, please refer to the comments. --- README.md | 3 + examples/a_mem/e2e_demo.py | 917 ++++++++++++++++++++ src/agent_memory/__init__.py | 2 + src/agent_memory/memories/__init__.py | 2 + src/agent_memory/memories/a_mem/__init__.py | 15 + src/agent_memory/memories/a_mem/policy.py | 299 +++++++ src/agent_memory/memories/a_mem/prompts.py | 106 +++ src/agent_memory/memories/a_mem/storage.py | 127 +++ tests/test_a_mem_e2e_example.py | 505 +++++++++++ tests/test_a_mem_policy.py | 682 +++++++++++++++ 10 files changed, 2658 insertions(+) create mode 100644 examples/a_mem/e2e_demo.py create mode 100644 src/agent_memory/memories/a_mem/__init__.py create mode 100644 src/agent_memory/memories/a_mem/policy.py create mode 100644 src/agent_memory/memories/a_mem/prompts.py create mode 100644 src/agent_memory/memories/a_mem/storage.py create mode 100644 tests/test_a_mem_e2e_example.py create mode 100644 tests/test_a_mem_policy.py diff --git a/README.md b/README.md index c3062d2..de44f54 100644 --- a/README.md +++ b/README.md @@ -95,6 +95,9 @@ examples/claude/ examples/zep/ e2e_demo.py real three-row LOCOMO + Zep policy demo +examples/a_mem/ + e2e_demo.py real five-row LOCOMO + A-Mem note evolution demo + examples/helloworld/ helloworld_smoke.py real LOCOMO + LOTUS sem_filter/sem_map/top-k smoke ``` diff --git a/examples/a_mem/e2e_demo.py b/examples/a_mem/e2e_demo.py new file mode 100644 index 0000000..c88be66 --- /dev/null +++ b/examples/a_mem/e2e_demo.py @@ -0,0 +1,917 @@ +"""Run the isolated Agent A-Mem storage and recovery smoke.""" + +from __future__ import annotations + +import argparse +from collections.abc import Mapping, Sequence +from datetime import datetime +import hashlib +from importlib.metadata import PackageNotFoundError, version +import json +import math +import os +from pathlib import Path +import pickle +import subprocess +import sys +import time +from typing import Any, cast +from uuid import uuid4 + +from dotenv import load_dotenv +import pandas as pd + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(PROJECT_ROOT / "src")) + +import agent_memory as am # noqa: E402 +from agent_memory.adapters.lotus import LotusAdapter # noqa: E402 +from agent_memory.adapters.lotus.context import LotusExecutionConfig # noqa: E402 +from agent_memory.datasets.locomo import load_locomo_rows # noqa: E402 +from agent_memory.evaluation.locomo import ensure_locomo_dataset # noqa: E402 +from agent_memory.evaluation.pricing import PricingSnapshot # noqa: E402 +from agent_memory.evaluation.trace_metrics import ( # noqa: E402 + normalize_provider_calls, + summarize_provider_calls, +) +from agent_memory.memories.a_mem.prompts import ( # noqa: E402 + AMEM_SOURCE_COMMIT, + ANALYSE_CONTENT_PROMPT, + EMBEDDING_TEXT_PROMPT, + EVOLUTION_SYSTEM_PROMPT, + NOTE_CONSOLIDATION_INSTRUCTION, +) +from agent_memory.memories.a_mem.storage import ( # noqa: E402 + AMEM_BGE_M3, + AMEM_NEO4J_SCHEMA, + AMEM_NEO4J_STATEMENTS, +) +from agent_memory.policy.retrieval import RetrievalResult # noqa: E402 +from agent_memory.storage import StorageDeployment # noqa: E402 +from agent_memory.storage.neo4j import ( # noqa: E402 + Neo4jConnector, + SentenceTransformerEmbeddingProvider, +) +from agent_memory.tracing.semantic import semantic_trace_scope # noqa: E402 + + +LOCOMO_CACHE_PATH = PROJECT_ROOT / ".cache" / "agent-memory" / "locomo10.json" +MODEL = "deepseek/deepseek-v4-flash" +LOG_COLUMNS = ("content", "timestamp") +RETRIEVAL_CHANNELS = ("notes", "neighbors") +TRACE_RUN_KIND = "amem" +LOCOMO_TURN_LIMIT = 1000 + +# Free OpenRouter variants allow 20 requests per minute, so stay under it; LOTUS +# applies the limit as a per-request delay and caps the batch size with it. +LM_NUM_RETRIES = 12 +LM_TIMEOUT_SECONDS = 120 +LM_RATE_LIMIT_PER_MINUTE = 10 + +# One pinned LOCOMO exchange. Sample 0 is conv-26, a Caroline/Melanie dialogue, and +# the slice starts at D2:8, the first turn of the adoption-agency thread. All five +# turns advance one topic — researching agencies, then picking one, then the reason +# — which is what lets A-Mem link notes and rewrite an earlier note's context and +# tags; see EVOLUTION_SYSTEM_PROMPT in src/agent_memory/memories/a_mem/prompts.py. +SAMPLE_INDEX = 0 +START_TURN_ID = "D2:8" +ROW_LIMIT = 5 + +# Official LOCOMO question conv-26:q88. Its answer is stated in the last ingested +# turn (D2:12), so reaching it depends on the link chain and the evolved context +# A-Mem builds across the thread rather than on any single note. +QUESTION_ID = "conv-26:q88" +QUERY = "Why did Caroline choose the adoption agency?" +GOLD_ANSWER = "because of their inclusivity and support for LGBTQ+ individuals" + +RETRIEVAL_CONTRACT: dict[str, Any] = { + "profile": "amem-bge-m3", + "notes_method": "cosine_similarity", + "notes_limit": 5, + "neighbors_method": "bfs", + "neighbors_limit": 10, + "bfs_max_depth": 1, + "reranker": None, +} + + +def main() -> None: + """Run add, search, checkpoint restore and repeated search.""" + + args = parse_args() + output_dir = args.output_dir.resolve() + if output_dir.exists() and any(output_dir.iterdir()): + raise FileExistsError(f"output directory must be empty: {output_dir}") + output_dir.mkdir(parents=True, exist_ok=True) + provider = require_environment(args.model) + + trace_dir = output_dir / "trace" + dataset_path, raw_rows, rows = selected_rows(row_limit=args.row_limit) + write_json(output_dir / "input" / "locomo_rows.json", list(raw_rows)) + write_json(output_dir / "input" / "amem_rows.json", rows) + print( + f"Agent A-Mem storage and recovery smoke: {len(rows)} LOCOMO turns from " + f"{dataset_path}, namespace {args.namespace}" + ) + + # The deployment publishes a StorageConnector protocol, but this demo owns the + # Neo4j connections it opens and has to close them, so narrow to the concrete + # type at the point of construction. + connectors: list[Neo4jConnector] = [] + status = "failed" + timings: dict[str, float] = {} + memory: am.AMem | None = None + observed: dict[str, dict[str, Any]] = {} + try: + storage = create_storage_deployment(args.namespace) + connector = cast(Neo4jConnector, storage.connector) + connectors.append(connector) + adapter = LotusAdapter( + model=args.model, + config=LotusExecutionConfig( + lm_model_kwargs=model_kwargs(args.model, provider), + lm_num_retries=LM_NUM_RETRIES, + lm_timeout=LM_TIMEOUT_SECONDS, + lm_rate_limit=LM_RATE_LIMIT_PER_MINUTE, + lm_enable_cache=False, + semantic_trace_dir=trace_dir, + ), + ) + memory = am.AMem(adapter=adapter, storage=storage) + + insertion_started = time.perf_counter() + for index, row in enumerate(rows, start=1): + print(f"add[{index}]: {row['content'][:100]}") + with semantic_trace_scope( + run_kind=TRACE_RUN_KIND, + phase="add", + add_index=index, + source_description=row["source_description"], + ): + memory.add({name: row[name] for name in LOG_COLUMNS}) + record_new_notes(observed, memory, add_index=index) + timings["insertion_wall_seconds"] = round( + time.perf_counter() - insertion_started, 4 + ) + write_state_artifacts(memory, output_dir) + + llm_before_retrieval = provider_call_count(trace_dir) + retrieval_started = time.perf_counter() + with semantic_trace_scope(run_kind=TRACE_RUN_KIND, phase="retrieval"): + before = memory.query(QUERY) + timings["retrieval_before_wall_seconds"] = round( + time.perf_counter() - retrieval_started, 4 + ) + if not isinstance(before, RetrievalResult): + raise TypeError("Agent A-Mem query must return RetrievalResult") + if provider_call_count(trace_dir) != llm_before_retrieval: + raise AssertionError("Agent A-Mem retrieval unexpectedly called the LLM") + validate_retrieval_result(before) + write_retrieval(output_dir / "retrieval" / "before_restore.json", before) + + checkpoint_started = time.perf_counter() + snapshot = memory._runtime.snapshot_state() + checkpoint_path = output_dir / "checkpoint" / "state.pkl" + checkpoint_path.parent.mkdir(parents=True, exist_ok=True) + checkpoint_path.write_bytes(pickle.dumps(snapshot)) + timings["checkpoint_save_wall_seconds"] = round( + time.perf_counter() - checkpoint_started, 4 + ) + write_json( + output_dir / "checkpoint" / "metadata.json", + { + "schema_version": snapshot["schema_version"], + "policy_fingerprint": snapshot["plan_fingerprint"], + "storage_commit": snapshot.get("storage_commit"), + "storage_bound": True, + "bytes": checkpoint_path.stat().st_size, + }, + ) + + # Restore against a fresh connection, so the round trip proves the durable + # state replays into a new physical handle rather than reusing the old one. + connector.close() + connectors.remove(connector) + restored_storage = create_storage_deployment(args.namespace) + connectors.append(cast(Neo4jConnector, restored_storage.connector)) + restored = am.AMem(adapter=adapter, storage=restored_storage) + restore_started = time.perf_counter() + restored._runtime.restore_state(pickle.loads(checkpoint_path.read_bytes())) + timings["checkpoint_restore_wall_seconds"] = round( + time.perf_counter() - restore_started, 4 + ) + assert_same_state(memory, restored) + + llm_before_restored_retrieval = provider_call_count(trace_dir) + retrieval_started = time.perf_counter() + with semantic_trace_scope(run_kind=TRACE_RUN_KIND, phase="retrieval"): + after = restored.query(QUERY) + timings["retrieval_after_wall_seconds"] = round( + time.perf_counter() - retrieval_started, 4 + ) + if not isinstance(after, RetrievalResult): + raise TypeError("restored Agent A-Mem query must return RetrievalResult") + if provider_call_count(trace_dir) != llm_before_restored_retrieval: + raise AssertionError( + "restored Agent A-Mem retrieval unexpectedly called the LLM" + ) + validate_retrieval_result(after) + assert_same_retrieval(before, after) + write_retrieval(output_dir / "retrieval" / "after_restore.json", after) + + evolution = evolution_summary( + observed=observed, + memory=restored, + rows=rows, + neighbors_rows=len(after.channels["neighbors"]), + ) + write_json(output_dir / "evolution" / "summary.json", evolution) + write_json( + output_dir / "retrieval" / "summary.json", + { + "query": QUERY, + "namespace": args.namespace, + "channels": { + name: len(frame) for name, frame in after.channels.items() + }, + "checkpoint_order_verified": True, + }, + ) + write_state_artifacts(restored, output_dir) + print_evolution_report(evolution) + status = "completed" + except Exception as error: + write_json( + output_dir / "diagnostics" / "failure.json", + {"type": type(error).__name__, "message": str(error)}, + ) + raise + finally: + for connector in connectors: + connector.close() + write_metrics(output_dir, status=status, timings=timings) + write_json( + output_dir / "manifest.json", + manifest( + output_dir, + model=args.model, + namespace=args.namespace, + status=status, + rows=rows, + ), + ) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + """Parse the isolated Agent A-Mem smoke options.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--row-limit", + type=int, + default=ROW_LIMIT, + help=( + "Turns to ingest from START_TURN_ID. A-Mem issues four semantic calls " + "per note plus one evolution call per new/earlier note pair, so keep " + "this small." + ), + ) + parser.add_argument("--model", default=MODEL) + parser.add_argument("--output-dir", type=Path, default=default_output_dir()) + parser.add_argument("--namespace", default=f"amem-e2e-{uuid4()}") + return parser.parse_args(argv) + + +def default_output_dir() -> Path: + """Return a fresh timestamped output path for one run.""" + + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + return PROJECT_ROOT / ".memory-test" / "a_mem-e2e" / timestamp + + +def model_provider(model: str) -> str: + """Return the LiteLLM provider for ``model``, or fail when unrecognized.""" + + import litellm + + try: + _model, provider, _api_key, _api_base = litellm.get_llm_provider(model=model) + except Exception as error: + raise SystemExit( + f"unrecognized LiteLLM provider for model {model!r}: {error}" + ) from error + return provider + + +def model_kwargs(model: str, provider: str) -> dict[str, Any]: + """Return the provider-specific LOTUS model kwargs for one run. + + Retries, timeout and rate limiting are execution-owned options and must go on + ``LotusExecutionConfig`` instead; overriding them here raises. + """ + + # litellm defaults to "constant_retry", which fires every attempt back to back + # with no wait; a shared-pool 429 needs real backoff to clear. + kwargs: dict[str, Any] = {"retry_strategy": "exponential_backoff_retry"} + if provider == "deepseek": + kwargs["extra_body"] = {"thinking": {"type": "disabled"}} + elif provider == "openrouter": + kwargs["extra_body"] = {"reasoning": {"enabled": False}} + # OpenRouter ":free" variants cap output at 512 tokens. + kwargs["max_tokens"] = 512 if ":free" in model else 4096 + return kwargs + + +def require_environment(model: str = MODEL) -> str: + """Load the local credentials and fail before creating artifacts if absent. + + Returns the LiteLLM provider so callers can pick provider-specific settings. + """ + + load_dotenv(PROJECT_ROOT / ".env") + import litellm + + provider = model_provider(model) + validation = litellm.validate_environment(model=model) + if not validation["keys_in_environment"]: + raise SystemExit( + f"model {model!r} requires environment variable(s): " + + ", ".join(str(name) for name in validation["missing_keys"]) + ) + missing = [ + name + for name in ("AGENT_MEMORY_NEO4J_URI", "AGENT_MEMORY_NEO4J_PASSWORD") + if not os.getenv(name) + ] + if missing: + raise SystemExit( + "Agent A-Mem retrieval requires environment variables: " + + ", ".join(missing) + ) + return provider + + +# Input normalization. +def locomo_source_description(row: Mapping[str, Any]) -> str: + """Return a short human-readable provenance label for one LOCOMO turn.""" + + session_number = str(row.get("session_id", "")).removeprefix("session_") + return ( + f"LOCOMO sample {row.get('sample_index')} session {session_number} " + f"turn {row.get('turn_id', '')}" + ) + + +def amem_log_row(row: Mapping[str, Any]) -> dict[str, Any]: + """Map one normalized LOCOMO turn to the AMem log schema. + + Mirrors ``A-mem/test_advanced_robust.py:240-244``; the session ``date_time`` + is passed through unchanged (``A-mem/memory_layer.py:294-296``). + """ + + speaker = str(row.get("speaker", "")) + text = str(row.get("message", "")) + caption = str(row.get("blip_caption", "")).strip() + if caption: + # A-mem/load_dataset.py:63-70 inlines the caption as "[Image: ...]". + caption_text = f"[Image: {caption}]" + text = f"{caption_text} {text}" if text else caption_text + return { + "content": f"Speaker {speaker}says : {text}", + "timestamp": str(row.get("timestamp", "")), + } + + +def selected_rows( + *, + row_limit: int, +) -> tuple[Path, list[dict[str, Any]], list[dict[str, Any]]]: + """Load the pinned LOCOMO exchange and return raw and A-Mem-shaped rows.""" + + if row_limit < 1: + raise SystemExit("--row-limit must be at least 1") + dataset_path = ensure_locomo_dataset(LOCOMO_CACHE_PATH) + loaded = load_locomo_rows( + dataset_path, + sample_limit=SAMPLE_INDEX + 1, + turn_limit=LOCOMO_TURN_LIMIT, + ) + start = next( + ( + index + for index, row in enumerate(loaded) + if str(row.get("turn_id", "")) == START_TURN_ID + ), + None, + ) + if start is None: + raise SystemExit(f"LOCOMO turn {START_TURN_ID!r} is not in the dataset") + selected = loaded[start : start + row_limit] + if len(selected) != row_limit: + raise SystemExit( + f"Requested {row_limit} LOCOMO turns from {START_TURN_ID}, " + f"found {len(selected)}" + ) + rows: list[dict[str, Any]] = [] + for row in selected: + mapped = amem_log_row(row) + mapped["source_description"] = locomo_source_description(row) + rows.append(mapped) + return dataset_path, selected, rows + + +def policy_input_fingerprint(rows: Sequence[Mapping[str, Any]]) -> str: + """Hash the normalized source rows shared with the native baseline.""" + + payload = json.dumps( + rows, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +# Retrieval assertions. +def validate_retrieval_result(result: RetrievalResult) -> None: + """Require the A-Mem channels, a non-empty note channel and valid ranks.""" + + if tuple(result.channels) != RETRIEVAL_CHANNELS: + raise RuntimeError( + f"Agent A-Mem retrieval must return {RETRIEVAL_CHANNELS} channels, " + f"got {tuple(result.channels)}" + ) + for name, frame in result.channels.items(): + if frame.empty: + if name == "notes": + raise RuntimeError("Agent A-Mem retrieval channel 'notes' is empty") + continue + missing = [ + column + for column in ("record_id", "content", "rank", "score") + if column not in frame.columns + ] + if missing: + raise RuntimeError( + f"Agent A-Mem retrieval channel {name!r} is missing columns: {missing}" + ) + ranks = frame["rank"].tolist() + if ranks != list(range(1, len(frame) + 1)): + raise RuntimeError( + f"Agent A-Mem retrieval channel {name!r} has invalid ranks: {ranks}" + ) + + +def record_ids(frame: pd.DataFrame, *, channel: str) -> list[Any]: + """Return one channel's record identity column, or fail on a broken schema.""" + + if "record_id" not in frame.columns: + raise RuntimeError( + f"Agent A-Mem retrieval channel {channel!r} has no record_id column" + ) + return frame["record_id"].tolist() + + +def assert_same_retrieval( + before: RetrievalResult, + after: RetrievalResult, +) -> None: + """Verify checkpoint restore preserves result identities and ordering.""" + + if tuple(before.channels) != tuple(after.channels): + raise RuntimeError("retrieval channels changed after checkpoint restore") + for name in before.channels: + before_ids = record_ids(before.channels[name], channel=name) + after_ids = record_ids(after.channels[name], channel=name) + if before_ids != after_ids: + raise RuntimeError( + f"retrieval order changed after restore for {name!r}: " + f"{before_ids!r} != {after_ids!r}" + ) + + +def assert_same_state(memory: am.AMem, restored: am.AMem) -> None: + """Require the restored durable state to equal the saved state.""" + + before_state = memory._runtime._state + after_state = restored._runtime._state + if set(before_state) != set(after_state): + raise RuntimeError( + "checkpoint restore changed durable state keys: " + f"{sorted(before_state)} != {sorted(after_state)}" + ) + for name in sorted(before_state): + pd.testing.assert_frame_equal(before_state[name], after_state[name]) + + +# Evolution evidence (reported, never asserted). +def text_value(value: Any) -> str: + """Return a stripped string for one possibly-null frame cell.""" + + if value is None: + return "" + if isinstance(value, float) and math.isnan(value): + return "" + return str(value).strip() + + +def string_sequence(value: Any) -> tuple[str, ...]: + """Return a sorted string tuple for one list-valued frame cell.""" + + if value is None: + return () + if isinstance(value, (str, bytes)): + text = text_value(value) + return (text,) if text else () + try: + entries = list(value) + except TypeError: + text = text_value(value) + return (text,) if text else () + normalized = (text_value(entry) for entry in entries) + return tuple(sorted(item for item in normalized if item)) + + +def note_observations(memory: am.AMem) -> dict[str, tuple[str, tuple[str, ...], str]]: + """Return ``{note_id: (content, tags, context)}`` for the current note view.""" + + frame = memory._runtime._state.get("note", pd.DataFrame()) + observations: dict[str, tuple[str, tuple[str, ...], str]] = {} + for record in frame.to_dict("records"): + note_id = text_value(record.get("_row_id")) + if not note_id: + continue + observations[note_id] = ( + text_value(record.get("content")), + string_sequence(record.get("tags")), + text_value(record.get("context")), + ) + return observations + + +def record_new_notes( + observed: dict[str, dict[str, Any]], + memory: am.AMem, + *, + add_index: int, +) -> None: + """Remember each note's analysed values at the add where it first appears.""" + + for note_id, (content, tags, context) in note_observations(memory).items(): + if note_id in observed: + continue + observed[note_id] = { + "index": add_index, + "content": content, + "tags": tags, + "context": context, + } + + +def evolution_summary( + *, + observed: Mapping[str, Mapping[str, Any]], + memory: am.AMem, + rows: Sequence[Mapping[str, Any]], + neighbors_rows: int | None, +) -> dict[str, Any]: + """Report link creation and neighbour rewriting without asserting them. + + A note is judged only once a later note has been added, so a freshly created + note cannot count as preserved. + """ + + final = note_observations(memory) + last_add_index = len(rows) + rewritten: list[dict[str, Any]] = [] + preserved: list[str] = [] + missing: list[str] = [] + for note_id, first in observed.items(): + if first["index"] >= last_add_index: + continue + current = final.get(note_id) + if current is None: + missing.append(note_id) + continue + content, tags, context = current + if (tags, context) == (first["tags"], first["context"]): + preserved.append(note_id) + continue + rewritten.append( + { + "note_id": note_id, + "first_seen_add": first["index"], + "content": first["content"][:120], + "tags_changed": list(first["tags"]) != list(tags), + "context_changed": first["context"] != context, + "tags_before": list(first["tags"]), + "tags_after": list(tags), + "context_before": first["context"], + "context_after": context, + } + ) + return { + "note_rows": len(final), + "notes_observed": len(observed), + "rewritten_count": len(rewritten), + "preserved_count": len(preserved), + "missing_count": len(missing), + "missing_note_ids": missing, + "rewritten": rewritten, + "neighbors_rows": neighbors_rows, + "links_created": None if neighbors_rows is None else neighbors_rows > 0, + "asserted": False, + } + + +def print_evolution_report(evolution: Mapping[str, Any]) -> None: + """Print the reported, never asserted evolution evidence.""" + + print("\nevolution report (reported, never asserted):") + print( + f" notes: {evolution['notes_observed']} observed, " + f"{evolution['rewritten_count']} rewritten, " + f"{evolution['preserved_count']} preserved" + ) + print(f" neighbors channel rows: {evolution['neighbors_rows']}") + if not evolution["links_created"]: + print( + " warning: no RELATES_TO links were created, so the neighbors " + "channel is empty" + ) + if evolution["missing_count"]: + print( + f" warning: {evolution['missing_count']} observed note(s) are absent " + "from the final note view" + ) + + +# Physical deployment and artifacts. +def create_storage_deployment(namespace: str) -> StorageDeployment: + """Create the real CPU-only Neo4j deployment. + + Both A-Mem channels declare ``reranker=None``, so no cross-encoder provider. + """ + + connector = Neo4jConnector( + uri=os.environ["AGENT_MEMORY_NEO4J_URI"], + auth=( + os.getenv("AGENT_MEMORY_NEO4J_USER", "neo4j"), + os.environ["AGENT_MEMORY_NEO4J_PASSWORD"], + ), + database=os.getenv("AGENT_MEMORY_NEO4J_DATABASE", "neo4j"), + embedding_provider=SentenceTransformerEmbeddingProvider(AMEM_BGE_M3), + schema=AMEM_NEO4J_SCHEMA, + ) + try: + return StorageDeployment( + connector=connector, + statements=AMEM_NEO4J_STATEMENTS, + namespace=namespace, + ) + except Exception: + connector.close() + raise + + +def write_state_artifacts(memory: am.AMem, output_dir: Path) -> dict[str, Path]: + """Write the current source log and the public note view.""" + + state = memory._runtime._state + return { + "state/log": write_json( + output_dir / "state" / "log.json", + state.get("log", pd.DataFrame()).to_dict("records"), + ), + "views/note": write_json( + output_dir / "views" / "note.json", + state.get("note", pd.DataFrame()).to_dict("records"), + ), + } + + +def write_retrieval(path: Path, result: RetrievalResult) -> None: + """Write one JSON-safe retrieval artifact.""" + + write_json( + path, + { + "query": result.query, + "channels": { + name: frame.to_dict("records") + for name, frame in result.channels.items() + }, + "metrics": { + name: dict(metrics) for name, metrics in result.metrics.items() + }, + }, + ) + + +def write_metrics( + output_dir: Path, + *, + status: str, + timings: Mapping[str, float], +) -> None: + """Normalize provider attempts and write a reproducible metrics summary.""" + + events = read_jsonl(output_dir / "trace" / "events.jsonl") + provider_rows = normalize_provider_calls( + events, + output_dir=output_dir, + pricing=PricingSnapshot.deepseek_2026_07_17(), + ) + metrics_dir = output_dir / "metrics" + metrics_dir.mkdir(parents=True, exist_ok=True) + pd.DataFrame(provider_rows).to_csv( + metrics_dir / "provider_usage.csv", + index=False, + ) + write_json( + metrics_dir / "summary.json", + { + "status": status, + **dict(timings), + "provider_usage": summarize_provider_calls(provider_rows), + "llm_calls": provider_call_count(output_dir / "trace"), + }, + ) + + +# Provenance. +def manifest( + output_dir: Path, + *, + model: str, + namespace: str, + status: str, + rows: Sequence[Mapping[str, Any]], +) -> dict[str, Any]: + """Return the non-sensitive Agent A-Mem reproducibility manifest.""" + + return { + "run_mode": "integration-smoke", + "system": "agent-amem", + "status": status, + "source": source_state(PROJECT_ROOT), + "upstream_reference": "A-mem (read-only reference checkout at ../A-mem)", + "native_amem_reference_commit": AMEM_SOURCE_COMMIT, + "input_fingerprint": policy_input_fingerprint(rows), + "locomo_question": { + "question_id": QUESTION_ID, + "question": QUERY, + "gold_answer": GOLD_ANSWER, + "start_turn_id": START_TURN_ID, + "row_limit": len(rows), + }, + "memory_model": model, + "memory_thinking": "disabled", + "runtime_contract_fingerprint": runtime_contract_fingerprint(model), + "application_cache": "disabled", + "prompts": prompt_fingerprints(), + "embedding": {**AMEM_BGE_M3.to_dict(), "device": "cpu"}, + "retrieval": dict(RETRIEVAL_CONTRACT), + "storage": { + "kind": "neo4j", + "namespace": namespace, + "database": os.getenv("AGENT_MEMORY_NEO4J_DATABASE", "neo4j"), + "driver": package_version("neo4j"), + }, + "runtime": runtime_versions(), + "pricing": PricingSnapshot.deepseek_2026_07_17().to_dict(), + } + + +def prompt_fingerprints() -> dict[str, dict[str, str]]: + """Return the four A-Mem prompt paths and their content digests.""" + + prompts = ( + ("analyse_content", "ANALYSE_CONTENT_PROMPT", ANALYSE_CONTENT_PROMPT), + ("evolution_system", "EVOLUTION_SYSTEM_PROMPT", EVOLUTION_SYSTEM_PROMPT), + ( + "note_consolidation", + "NOTE_CONSOLIDATION_INSTRUCTION", + NOTE_CONSOLIDATION_INSTRUCTION, + ), + ("embedding_text", "EMBEDDING_TEXT_PROMPT", EMBEDDING_TEXT_PROMPT), + ) + return { + label: { + "path": f"src/agent_memory/memories/a_mem/prompts.py:{constant}", + "sha256": hashlib.sha256(text.encode("utf-8")).hexdigest(), + } + for label, constant, text in prompts + } + + +def runtime_contract_fingerprint(model: str) -> str: + """Hash the cross-system physical model and retrieval contract.""" + + payload = { + "memory_model": model.removeprefix("deepseek/"), + "thinking": "disabled", + "application_cache": "disabled", + "embedding_model": AMEM_BGE_M3.model, + "embedding_revision": AMEM_BGE_M3.revision, + "embedding_dimensions": AMEM_BGE_M3.dimensions, + "embedding_source_column": AMEM_BGE_M3.source_column, + "embedding_device": "cpu", + **RETRIEVAL_CONTRACT, + } + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + default=str, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def source_state(repo: Path) -> dict[str, Any]: + """Return commit, dirty state and uv lock digest.""" + + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repo, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + dirty = bool( + subprocess.run( + ["git", "status", "--porcelain"], + cwd=repo, + check=True, + capture_output=True, + text=True, + ).stdout + ) + return { + "commit": commit, + "dirty": dirty, + "lockfile": "uv.lock", + "lockfile_sha256": hashlib.sha256((repo / "uv.lock").read_bytes()).hexdigest(), + } + + +def runtime_versions() -> dict[str, Any]: + """Return the critical isolated runtime versions.""" + + return { + "python": sys.version.split()[0], + "agent-memory": package_version("agent-memory"), + "neo4j": package_version("neo4j"), + "sentence-transformers": package_version("sentence-transformers"), + "numpy": package_version("numpy"), + "pandas": package_version("pandas"), + "lotus-ai": package_version("lotus-ai"), + "torch": package_version("torch"), + "transformers": package_version("transformers"), + "litellm": package_version("litellm"), + } + + +def provider_call_count(trace_dir: Path) -> int: + """Count physical provider usage events in the semantic trace.""" + + return sum( + event.get("event_type") == "provider_usage" + for event in read_jsonl(trace_dir / "events.jsonl") + ) + + +def package_version(name: str) -> str | None: + """Return an installed package version, if present.""" + + try: + return version(name) + except PackageNotFoundError: + return None + + +def read_jsonl(path: Path) -> list[dict[str, Any]]: + """Read an optional JSONL artifact.""" + + if not path.exists(): + return [] + return [ + json.loads(line) + for line in path.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + + +def write_json(path: Path, value: Any) -> Path: + """Write one inspectable UTF-8 JSON artifact.""" + + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, default=str), + encoding="utf-8", + ) + return path + + +if __name__ == "__main__": + main() diff --git a/src/agent_memory/__init__.py b/src/agent_memory/__init__.py index fa2f566..bb74bf7 100644 --- a/src/agent_memory/__init__.py +++ b/src/agent_memory/__init__.py @@ -2,6 +2,7 @@ from .api import Memory, Message from .memories import ( + AMem, ClaudeMemory, Mem0Memory, Mem0MemoryEnhanced, @@ -29,6 +30,7 @@ __version__ = "0.0.1" __all__ = [ + "AMem", "ClaudeMemory", "DifferentiatedPolicy", "BFS", diff --git a/src/agent_memory/memories/__init__.py b/src/agent_memory/memories/__init__.py index c95e810..fd03274 100644 --- a/src/agent_memory/memories/__init__.py +++ b/src/agent_memory/memories/__init__.py @@ -1,10 +1,12 @@ """Built-in memory policies.""" +from .a_mem import AMem from .claude import ClaudeMemory from .mem0 import Mem0Memory, Mem0MemoryEnhanced from .zep import ZepMemory, ZepMemoryExtended __all__ = [ + "AMem", "ClaudeMemory", "Mem0Memory", "Mem0MemoryEnhanced", diff --git a/src/agent_memory/memories/a_mem/__init__.py b/src/agent_memory/memories/a_mem/__init__.py new file mode 100644 index 0000000..41f7413 --- /dev/null +++ b/src/agent_memory/memories/a_mem/__init__.py @@ -0,0 +1,15 @@ +"""A-Mem-style built-in memory policy.""" + +from .policy import AMem +from .storage import ( + AMEM_BGE_M3, + AMEM_NEO4J_SCHEMA, + AMEM_NEO4J_STATEMENTS, +) + +__all__ = [ + "AMem", + "AMEM_BGE_M3", + "AMEM_NEO4J_SCHEMA", + "AMEM_NEO4J_STATEMENTS", +] diff --git a/src/agent_memory/memories/a_mem/policy.py b/src/agent_memory/memories/a_mem/policy.py new file mode 100644 index 0000000..a2c525d --- /dev/null +++ b/src/agent_memory/memories/a_mem/policy.py @@ -0,0 +1,299 @@ +"""A-Mem style note evolution memory policy.""" + +from __future__ import annotations + +from agent_memory.policy.aggregates import min, sem_agg +from agent_memory.api import Log, Memory +from agent_memory.policy.logical import UserQuery +from agent_memory.policy.retrieval import ( + BFS, + CosineSimilarity, + RetrievalQuery, +) +from .prompts import ( + ANALYSE_CONTENT_PROMPT, + EMBEDDING_TEXT_PROMPT, + EVOLUTION_SYSTEM_PROMPT, + NOTE_CONSOLIDATION_INSTRUCTION, +) + + +_NOTE_COLUMNS = [ + "_row_id", + "content", + "keywords", + "context", + "tags", + "timestamp", + "embedding_text", +] + + +_NOTE_STATE_COLUMNS = [ + "_row_id", + "_add_seq", + "content", + "keywords", + "context", + "tags", + "timestamp", +] + + +_NOTE_LINK_COLUMNS = ["source_note_id", "target_note_id"] + +# Shared by the two retrieval channels (base similarity + BFS neighbour expansion): +# record_id/rank/score are engine-supplied, content carries the note payload. +_RETRIEVAL_COLUMNS = ["record_id", "content", "rank", "score"] + + +class AMem(Memory): + """ + A-Mem style note evolution memory policy. + """ + + # Input log + + log = Log( + { + "content": "The main text content of the memory.", + "timestamp": "Creation time in format YYYYMMDDHHMM.", + }, + system_columns=True, + ) + + # Analysis + + # Each log entry enriched with keywords, context, and tags through llm (§3.1 eq. 2). + _analysed_notes = log.sem_map( + input_cols=["content"], + output_cols={ + "analysis": ( + "The structured content analysis as one JSON object with keys " + "keywords, context, and tags." + ), + }, + instruction=ANALYSE_CONTENT_PROMPT, + ).unnest( + column="analysis", + fields={ + "keywords": "keywords", + "context": "context", + "tags": "tags", + }, + ) + + # Note pairs + + # Instead of retrieving the top-k most relevant memories for decision-making, (§3.2 eq. 4 and 5 originally) + # we consider all earlier notes as potential neighbors for evolution. + _earlier_notes = _analysed_notes.alias("earlier") + _later_notes = _analysed_notes.alias("later") + # Each row is one ordered (earlier, later) note pair, suffixed :earlier / :later. + _candidate_note_pairs = _earlier_notes.join( + _later_notes, + on=(_earlier_notes.col("_add_seq") < _later_notes.col("_add_seq")), + ) + + # Evolution decision (§3.1 eq. 6 and §3.2 eq. 7) + + # Each row is one evolution decision for a later note against its earlier neighbors. + # Two possible operations (can be none or either one or both): + # 1. "strengthen" -> LINK GENERATION: emit a RELATES_TO edge to each suggested neighbor. + # 2. "update_neighbor" -> MEMORY EVOLUTION: rewrite the neighbor's tags/context. + _evolution_output = _candidate_note_pairs.group_by( + ["_row_id:later", "_add_seq:later"] + ).agg( + sem_agg( + input_cols=[ + "_add_seq:earlier", + "content:earlier", + "context:earlier", + "keywords:earlier", + "tags:earlier", + "content:later", + "context:later", + "keywords:later", + "tags:later", + ], + output_cols={ + "evolution": ( + "The complete evolution decision as one JSON object with keys " + "should_evolve, actions, suggested_connections, tags_to_update, " + "neighbor_updates." + ), + }, + instruction=EVOLUTION_SYSTEM_PROMPT, + ), + ) + + # unnest: one column per decision key (note the list-valued payload keys below). + # should_evolve (bool) + # actions (list of "strengthen"/"update_neighbor") + # suggested_connections (list of [_add_seq:earlier] positions) + # tags_to_update (list) + # neighbor_updates (list of {neighbor_add_seq, new_context, new_tags}). + _evolution_decisions = _evolution_output.unnest( + column="evolution", + fields={ + "should_evolve": "should_evolve", + "actions": "actions", + "suggested_connections": "suggested_connections", + "tags_to_update": "tags_to_update", + "neighbor_updates": "neighbor_updates", + }, + ) + + # One row per action; the full decision payload repeats on every row. + _actions = _evolution_decisions.explode(column="actions", output_col="action") + + # Link generation (§3.2) + + # Each row is one note that gets link generation, plus its tags_to_update and suggested_connections. + _link_generation_writes = ( + _actions.filter( + _actions.col("should_evolve") & (_actions.col("action") == "strengthen") + ) + .assign(_row_id=_actions.col("_row_id:later")) + .select(["_row_id", "_row_id:later", "suggested_connections", "tags_to_update"]) + ) + + # Each row is one directed RELATES_TO edge (source -> target note). + # This will be stored in graphDb, and BFS neighbour expansion traverses them. + # The model names each neighbor by its [_add_seq:earlier] position, so resolve that + # position against the candidate pair relation to recover the earlier note id. + # A position naming no earlier neighbor of this later note produces no edge. + _note_links = ( + _link_generation_writes.select(["_row_id:later", "suggested_connections"]) + .explode(column="suggested_connections", output_col="_add_seq:earlier") + .join(_candidate_note_pairs, on=["_row_id:later", "_add_seq:earlier"]) + .assign( + source_note_id=_candidate_note_pairs.col("_row_id:later"), + target_note_id=_candidate_note_pairs.col("_row_id:earlier"), + ) + .select(_NOTE_LINK_COLUMNS) + .drop_duplicates(subset=_NOTE_LINK_COLUMNS) + ) + + # Note states + + # A note's state is assembled from one state row per state it passes through: + # - base state: before memory evolution + # - evolution state: one per later memory evolution + # Memory evolution rewriting an earlier note introduces a circular dependency on note. + # Note states mitigate this by turning mutation into append + # LLM then merges the states into one final state per note. + + # Notes with link-generation: update tags, keep keywords/context unchanged. + _link_generation_notes = ( + _analysed_notes.join(_link_generation_writes, on="_row_id", how="inner") + .assign(tags=_link_generation_writes.col("tags_to_update")) + .select(_NOTE_STATE_COLUMNS) + ) + + # Notes with no link-generation: keep keywords/context/tags unchanged. + _unchanged_notes = _analysed_notes.join( + _link_generation_writes, on="_row_id", how="left_anti" + ).select(_NOTE_STATE_COLUMNS) + + # This form the base state before memory evolution + _base_note_states = _link_generation_notes.concat(_unchanged_notes) + + # Memory evolution (§3.2) + + # Each row is one note that gets memory evolution, plus its new_context and new_tags. + # The rewritten note is named by its [_add_seq:earlier] position, so resolve that + # position against the candidate pair relation to recover the rewritten note id. + # _add_seq is added to guide consolidation to prefers the later rewrite. + _memory_evolution_writes = ( + _actions.filter( + _actions.col("should_evolve") + & (_actions.col("action") == "update_neighbor") + ) + .select(["_row_id:later", "neighbor_updates"]) + .explode(column="neighbor_updates", output_col="neighbor_update") + .unnest( + column="neighbor_update", + fields={ + "neighbor_add_seq": "_add_seq:earlier", + "new_context": "context", + "new_tags": "tags", + }, + ) + .join(_candidate_note_pairs, on=["_row_id:later", "_add_seq:earlier"]) + .assign( + _row_id=_candidate_note_pairs.col("_row_id:earlier"), + _add_seq=_candidate_note_pairs.col("_add_seq:later"), + ) + .select(["_row_id", "context", "tags", "_add_seq"]) + ) + + _rewrites = _memory_evolution_writes.alias("rewrite") + + # Applies memory evolution's change as a state of the rewritten note. + _memory_evolution_note_states = ( + _analysed_notes.join(_rewrites, on="_row_id") + .assign( + context=_rewrites.col("context"), + tags=_rewrites.col("tags"), + _add_seq=_rewrites.col("_add_seq"), + ) + .select(_NOTE_STATE_COLUMNS) + ) + + # Consolidation + + # An LLM consolidates a note's successive states (each representing a change) into one note. + note = ( + _base_note_states.concat(_memory_evolution_note_states) + .group_by("_row_id") + .agg( + # These columns require llm to consolidate the successive states into one final state. + sem_agg( + input_cols=["context", "tags", "_add_seq"], + output_cols={ + "context": "Accumulated note context.", + "tags": "Accumulated note tags.", + }, + instruction=NOTE_CONSOLIDATION_INSTRUCTION, + ), + # These columns keep constant across states. + min(column="content", output_col="content"), + min(column="keywords", output_col="keywords"), + min(column="timestamp", output_col="timestamp"), + ) + .sem_map( + # Produce embedding text by merging content, keywords, context, tags (§3.1 eq. 3) + input_cols=["content", "keywords", "context", "tags"], + output_cols={ + "embedding_text": ( + "The verbatim concatenation of content, keywords, context, and " + "tags, in that order, with no rephrasing. Output only the " + "concatenated text." + ), + }, + instruction=EMBEDDING_TEXT_PROMPT, + ) + .select(_NOTE_COLUMNS) + ) + + # Retrieval (§3.4) + + # Retrieve the top-5 notes from the note table. (§3.4 eq.8 & 9 & eq.10) + _retrieved_notes = note.search( + UserQuery(), + methods=[CosineSimilarity()], + reranker=None, + limit=5, + ).select(_RETRIEVAL_COLUMNS) + + # Neighbor expansion via BFS (Figure 2). + retrieval_query = RetrievalQuery( + notes=_retrieved_notes, + neighbors=note.search( + UserQuery(), # query text ignored by BFS + methods=[BFS(origins=_retrieved_notes, max_depth=1)], + reranker=None, + limit=10, # neighbor budget, independent of base k + ).select(_RETRIEVAL_COLUMNS), + ) diff --git a/src/agent_memory/memories/a_mem/prompts.py b/src/agent_memory/memories/a_mem/prompts.py new file mode 100644 index 0000000..a7a6545 --- /dev/null +++ b/src/agent_memory/memories/a_mem/prompts.py @@ -0,0 +1,106 @@ +"""Audited A-Mem prompt adaptation.""" + +from __future__ import annotations + + +AMEM_SOURCE_COMMIT = "0c8039f28fdcc08189a23c07a3437d9d2482f9c2" +AMEM_SOURCE_PROMPT_PATH = ( + "A-mem/llm_text_parsers.py:ANALYZE_CONTENT_PROMPT / EVOLUTION_DECISION_PROMPT" +) + +ANALYSE_CONTENT_PROMPT = r"""Generate a structured analysis of the following content by: + 1. Identifying the most salient keywords (focus on nouns, verbs, and key concepts) + 2. Extracting core themes and contextual elements + 3. Creating relevant categorical tags + + Format the response as a JSON object: + {{ + "keywords": [ + // several specific, distinct keywords that capture key concepts and terminology + // Order from most to least important + // Don't include keywords that are the name of the speaker or time + // At least three keywords, but don't be too redundant. + ], + "context": + // one sentence summarizing: + // - Main topic/domain + // - Key arguments/points + // - Intended audience/purpose + , + "tags": [ + // several broad categories/themes for classification + // Include domain, format, and type tags + // At least three tags, but don't be too redundant. + ] + }} + + Content for analysis: + {content} + + Return only a single JSON object. No markdown fences and no extra prose. + """ + + +EVOLUTION_SYSTEM_PROMPT = r""" +You are an AI memory evolution agent responsible for managing and evolving a knowledge base. +Each entry below pairs one new memory with one earlier neighbor memory. The neighbor's +position is given by its [_add_seq:earlier] field, which is a small integer. + +Refer to a neighbor only by copying that integer. Never write a field name, and never +write the words "earlier" or "later" in place of the number. + +Based on this information, determine: +1. Should this memory be evolved? Consider its relationships with the other memories. +2. What specific actions should be taken (strengthen, update_neighbor)? + 2.1 If strengthen, which neighbor should it be connected to? Give that neighbor's + [_add_seq:earlier] integer. Also give the updated tags for the new memory. + 2.2 If update_neighbor, give the updated context and tags for each neighbor. + If a neighbor's context or tags should not change, repeat the original values. + Emit one object per neighbor, each carrying its own neighbor_add_seq; + order does not matter. +Tags should reflect the content of these memories so they can be retrieved and categorized later. + +Return only a single JSON object. No markdown fences and no extra prose. +{{ + "should_evolve": true, + "actions": ["strengthen", "update_neighbor"], + "suggested_connections": [0, 1], + "tags_to_update": ["tag_1", "tag_n"], + "neighbor_updates": [ + {{"neighbor_add_seq": 0, "new_context": "new context", "new_tags": ["tag_1", "tag_n"]}} + ] +}} +""" + + +NOTE_CONSOLIDATION_INSTRUCTION = """ +Merge the grouped context and tags entries for one note into a single accumulated +{context} and {tags}. Use only information explicitly present in the grouped +entries; never infer beyond the evidence. + +Preserve every distinct fact, detail, and tag across all entries, including +changes over time. When entries conflict or supersede one another, prefer the +newer entry: each input row carries _add_seq, and a larger _add_seq means the +entry was produced later. An entry that simply omits a fact does not remove it; +only drop a fact when a newer entry explicitly contradicts or supersedes it. +""".strip() + + +EMBEDDING_TEXT_PROMPT = r"""Produce the exact text to be embedded for memory retrieval. + +Concatenate the following fields, IN ORDER, with no changes whatsoever: +1. content +2. keywords (join the list items with a single space; keep each keyword verbatim) +3. context +4. tags (join the list items with a single space; keep each tag verbatim) + +Rules: +- Output ONLY the concatenated text. No labels, no headings, no markdown fences, no extra prose. +- The original `content` must appear word-for-word at the start, unchanged. +- Do not translate, rewrite, reorder, summarize, or add any commentary to the source material. + +content: {content} +keywords: {keywords} +context: {context} +tags: {tags} +""" diff --git a/src/agent_memory/memories/a_mem/storage.py b/src/agent_memory/memories/a_mem/storage.py new file mode 100644 index 0000000..b199f38 --- /dev/null +++ b/src/agent_memory/memories/a_mem/storage.py @@ -0,0 +1,127 @@ +"""Neo4j storage profile for the A-Mem note-evolution policy. + +Implements equation 3 of the A-Mem design: the embedded text is the verbatim +concatenation of content + keywords + context + tags. The policy emits that +combined text in a single `embedding_text` column (see a_mem/prompts.py), and +this profile hands it to the engine's existing single-column EmbeddingSpec. +No engine code is modified; this module only consumes the public connector API. +""" + +from __future__ import annotations + +from agent_memory.storage import Schema, StatementSet, TableDescriptor +from agent_memory.storage.neo4j import ( + EmbeddingSpec, + Neo4jIdentity, + Neo4jNodeMapping, + Neo4jRelationshipMapping, + Neo4jSchema, +) + +from .policy import AMem + + +# Physical embedding and index contract for the A-Mem profile. Reuses the same +# BGE-M3 revision/dimensions as mem0/zep so vectors live in the same space. +AMEM_BGE_M3 = EmbeddingSpec( + source_column="embedding_text", + property_name="note_embedding", + model="BAAI/bge-m3", + revision="5617a9f61b028005a4858fdac845db406aefb181", + dimensions=1024, + normalize=True, +) + + +# Sink schema validates logical policy rows before any Neo4j lowering occurs. +_NOTE_SCHEMA = ( + Schema.new_builder() + .column("_row_id", "BIGINT") + .column("content", "STRING") + .column("keywords", "ARRAY") + .column("context", "STRING") + .column("tags", "ARRAY") + .column("timestamp", "STRING") + # Verbatim eq.3 concatenation; also stored as a node property so the + # cross-encoder reranker can read it back as passage text. + .column("embedding_text", "STRING") + .primary_key("_row_id") + .build() +) + + +_NOTE_TARGET = ( + TableDescriptor.for_connector("neo4j") + .schema(_NOTE_SCHEMA) + .mapping( + Neo4jNodeMapping( + label="Note", + identity=Neo4jIdentity(columns=("_row_id",), kind="note"), + properties={ + "content": "content", + "keywords": "keywords", + "context": "context", + "tags": "tags", + "timestamp": "timestamp", + "embedding_text": "embedding_text", + }, + embedding=AMEM_BGE_M3, + ) + ) + .build() +) + + +# Link edges live in their own relationship table, not as an array property on +# the Note node, so BFS neighbor expansion can traverse them. +_NOTE_LINK_SCHEMA = ( + Schema.new_builder() + .column("source_note_id", "STRING") + .column("target_note_id", "STRING") + .primary_key("source_note_id", "target_note_id") + .build() +) + + +_NOTE_LINK_TARGET = ( + TableDescriptor.for_connector("neo4j") + .schema(_NOTE_LINK_SCHEMA) + .mapping( + Neo4jRelationshipMapping( + relationship_type="RELATES_TO", + identity=Neo4jIdentity( + columns=("source_note_id", "target_note_id"), kind="note_link" + ), + source=Neo4jIdentity(columns=("source_note_id",), kind="note"), + target=Neo4jIdentity(columns=("target_note_id",), kind="note"), + ) + ) + .build() +) + + +AMEM_NEO4J_SCHEMA = Neo4jSchema( + queries=( + "CREATE CONSTRAINT agent_memory_amem_note_uuid IF NOT EXISTS " + "FOR (n:Note) REQUIRE n.uuid IS UNIQUE", + "CREATE INDEX agent_memory_amem_note_embedding IF NOT EXISTS " + "FOR (n:Note) ON (n.note_embedding)", + "CREATE INDEX agent_memory_amem_note_link_uuid IF NOT EXISTS " + "FOR ()-[e:RELATES_TO]-() ON (e.uuid)", + ), + traversal_relationship_types=("RELATES_TO",), +) + + +AMEM_NEO4J_STATEMENTS = ( + StatementSet() + .add_insert(_NOTE_TARGET, AMem.note) + .add_insert(_NOTE_LINK_TARGET, AMem._note_links) +) + + +__all__ = [ + "AMEM_BGE_M3", + "AMEM_NEO4J_SCHEMA", + "AMEM_NEO4J_STATEMENTS", +] diff --git a/tests/test_a_mem_e2e_example.py b/tests/test_a_mem_e2e_example.py new file mode 100644 index 0000000..e74a6ad --- /dev/null +++ b/tests/test_a_mem_e2e_example.py @@ -0,0 +1,505 @@ +"""Tests for the real A-Mem LOCOMO example input boundary.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pandas as pd +import pytest + +import agent_memory as am + + +EXAMPLE_PATH = ( + Path(__file__).resolve().parents[1] / "examples" / "a_mem" / "e2e_demo.py" +) + + +def _load_example() -> ModuleType: + assert EXAMPLE_PATH.exists(), "examples/a_mem/e2e_demo.py is required" + spec = importlib.util.spec_from_file_location("a_mem_e2e_demo", EXAMPLE_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _main_args(output_dir: Path) -> SimpleNamespace: + return SimpleNamespace( + row_limit=1, + model="test-model", + output_dir=output_dir, + namespace="amem-test", + ) + + +def _amem_rows() -> list[dict[str, object]]: + return [ + { + "content": "Speaker Carolinesays : Hello", + "timestamp": "1:14 pm on 25 May, 2023", + "source_description": "LOCOMO sample 0 session 2 turn D2:8", + } + ] + + +def _retrieval_result() -> am.RetrievalResult: + return am.RetrievalResult( + query="Caroline", + channels={ + "notes": pd.DataFrame( + [ + { + "record_id": "note-1", + "content": "Speaker Carolinesays : Hello", + "rank": 1, + "score": 0.9, + }, + { + "record_id": "note-3", + "content": "Speaker Bobsays : Later", + "rank": 2, + "score": 0.4, + }, + ] + ), + "neighbors": pd.DataFrame( + [ + { + "record_id": "note-2", + "content": "Speaker Bobsays : Hi", + "rank": 1, + "score": 0.5, + } + ] + ), + }, + metrics={"notes": {}, "neighbors": {"bfs_origins": ["note-1"]}}, + ) + + +def _memory(note_rows: list[dict[str, object]]) -> SimpleNamespace: + return SimpleNamespace( + add=Mock(), + query=Mock(return_value=_retrieval_result()), + _runtime=SimpleNamespace( + _state={ + "log": pd.DataFrame(), + "note": pd.DataFrame(note_rows), + }, + snapshot_state=Mock( + return_value={ + "schema_version": 2, + "plan_fingerprint": "test-plan", + "storage_commit": None, + } + ), + restore_state=Mock(), + ), + ) + + +def _patch_common( + example: ModuleType, + monkeypatch: pytest.MonkeyPatch, + *, + args: SimpleNamespace, + rows: list[dict[str, object]], +) -> None: + monkeypatch.setattr(example, "parse_args", lambda: args) + monkeypatch.setattr(example, "require_environment", lambda *a, **k: "test-provider") + monkeypatch.setattr( + example, + "selected_rows", + lambda **kwargs: (Path("locomo.json"), rows, rows), + ) + monkeypatch.setattr( + example, + "source_state", + lambda repo: { + "commit": "test-commit", + "dirty": False, + "lockfile": "uv.lock", + "lockfile_sha256": "0" * 64, + }, + ) + + +def test_amem_log_row_matches_upstream_speaker_and_caption_format() -> None: + example = _load_example() + + result = example.amem_log_row( + { + "message": "Caroline is researching adoption agencies.", + "speaker": "Caroline", + "session_id": "session_2", + "sample_index": 0, + "turn_id": "D2:8", + "timestamp": "1:14 pm on 25 May, 2023", + "blip_caption": "an adoption agency brochure", + } + ) + + assert result == { + "content": ( + "Speaker Carolinesays : [Image: an adoption agency brochure] " + "Caroline is researching adoption agencies." + ), + "timestamp": "1:14 pm on 25 May, 2023", + } + + +def test_amem_log_row_omits_an_empty_caption() -> None: + example = _load_example() + + result = example.amem_log_row( + { + "message": "Hello", + "speaker": "Bob", + "session_id": "session_1", + "sample_index": 0, + "turn_id": "D1:1", + "timestamp": "1:00 pm on 25 May, 2023", + "blip_caption": " ", + } + ) + + assert result["content"] == "Speaker Bobsays : Hello" + + +def test_locomo_source_description_keeps_turn_identity() -> None: + example = _load_example() + + assert ( + example.locomo_source_description( + {"sample_index": 0, "session_id": "session_2", "turn_id": "D2:8"} + ) + == "LOCOMO sample 0 session 2 turn D2:8" + ) + + +def test_policy_input_fingerprint_is_stable_and_content_sensitive() -> None: + example = _load_example() + rows = _amem_rows() + + digest = example.policy_input_fingerprint(rows) + + assert len(digest) == 64 + assert digest == example.policy_input_fingerprint(_amem_rows()) + changed = [dict(rows[0], content="Speaker Bobsays : Hello")] + assert example.policy_input_fingerprint(changed) != digest + + +def test_e2e_requires_neo4j_environment( + monkeypatch: pytest.MonkeyPatch, +) -> None: + example = _load_example() + monkeypatch.setattr(example, "load_dotenv", lambda *args, **kwargs: False) + monkeypatch.setenv("DEEPSEEK_API_KEY", "test-key") + monkeypatch.delenv("AGENT_MEMORY_NEO4J_URI", raising=False) + monkeypatch.delenv("AGENT_MEMORY_NEO4J_PASSWORD", raising=False) + + with pytest.raises(SystemExit, match="AGENT_MEMORY_NEO4J_URI"): + example.require_environment() + + +def test_parse_args_pins_the_exchange_and_generates_a_fresh_namespace() -> None: + example = _load_example() + + parsed = example.parse_args(["--namespace", "amem-test"]) + + assert parsed.row_limit == example.ROW_LIMIT + assert parsed.model == "deepseek/deepseek-v4-flash" + assert parsed.namespace == "amem-test" + + generated = example.parse_args([]) + assert generated.namespace.startswith("amem-e2e-") + + +def test_validate_retrieval_result_allows_an_empty_neighbors_channel() -> None: + example = _load_example() + result = am.RetrievalResult( + query="Caroline", + channels={ + "notes": pd.DataFrame( + [ + { + "record_id": "note-1", + "content": "Hello", + "rank": 1, + "score": 0.9, + } + ] + ), + "neighbors": pd.DataFrame(), + }, + ) + + example.validate_retrieval_result(result) + + +def test_validate_retrieval_result_rejects_empty_notes_and_broken_ranks() -> None: + example = _load_example() + empty_notes = am.RetrievalResult( + query="Caroline", + channels={"notes": pd.DataFrame(), "neighbors": pd.DataFrame()}, + ) + broken_ranks = am.RetrievalResult( + query="Caroline", + channels={ + "notes": pd.DataFrame( + [ + { + "record_id": "note-1", + "content": "Hello", + "rank": 2, + "score": 0.9, + } + ] + ), + "neighbors": pd.DataFrame(), + }, + ) + missing_columns = am.RetrievalResult( + query="Caroline", + channels={ + "notes": pd.DataFrame([{"rank": 1}]), + "neighbors": pd.DataFrame(), + }, + ) + + with pytest.raises(RuntimeError, match="channel 'notes' is empty"): + example.validate_retrieval_result(empty_notes) + with pytest.raises(RuntimeError, match="invalid ranks"): + example.validate_retrieval_result(broken_ranks) + with pytest.raises(RuntimeError, match="missing columns"): + example.validate_retrieval_result(missing_columns) + + +def test_assert_same_retrieval_detects_reordering() -> None: + example = _load_example() + before = _retrieval_result() + reordered = am.RetrievalResult( + query="Caroline", + channels={ + "notes": before.channels["notes"].iloc[::-1].reset_index(drop=True), + "neighbors": before.channels["neighbors"], + }, + ) + + example.assert_same_retrieval(before, before) + with pytest.raises(RuntimeError, match="retrieval order changed"): + example.assert_same_retrieval(before, reordered) + + +def test_evolution_summary_reports_rewrites_without_asserting() -> None: + example = _load_example() + memory = _memory( + [ + { + "_row_id": "note-1", + "content": "Speaker Carolinesays : Hello", + "tags": ["adoption", "research"], + "context": "Caroline researches adoption agencies", + } + ] + ) + observed = { + "note-1": { + "index": 1, + "content": "Speaker Carolinesays : Hello", + "tags": ("research",), + "context": "Caroline research", + } + } + + summary = example.evolution_summary( + observed=observed, + memory=memory, + rows=[{}, {}], + neighbors_rows=2, + ) + + assert summary["asserted"] is False + assert summary["rewritten_count"] == 1 + assert summary["preserved_count"] == 0 + assert summary["missing_count"] == 0 + assert summary["links_created"] is True + assert summary["rewritten"][0]["tags_changed"] is True + + +def test_evolution_summary_skips_notes_first_seen_at_the_last_add() -> None: + example = _load_example() + memory = _memory( + [ + { + "_row_id": "note-1", + "content": "Speaker Carolinesays : Hello", + "tags": ["research"], + "context": "Caroline research", + } + ] + ) + observed = { + "note-1": { + "index": 2, + "content": "Speaker Carolinesays : Hello", + "tags": ("research",), + "context": "Caroline research", + } + } + + summary = example.evolution_summary( + observed=observed, + memory=memory, + rows=[{}, {}], + neighbors_rows=0, + ) + + assert summary["notes_observed"] == 1 + assert summary["rewritten_count"] == 0 + assert summary["preserved_count"] == 0 + assert summary["links_created"] is False + + +def test_write_state_artifacts_accepts_an_empty_state(tmp_path: Path) -> None: + example = _load_example() + memory = SimpleNamespace(_runtime=SimpleNamespace(_state={})) + + written = example.write_state_artifacts(memory, tmp_path) + + assert set(written) == {"state/log", "views/note"} + assert all(path.exists() for path in written.values()) + + +def test_storage_deployment_closes_connector_when_binding_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + example = _load_example() + connector = SimpleNamespace(close=Mock()) + monkeypatch.setenv("AGENT_MEMORY_NEO4J_URI", "bolt://localhost:7687") + monkeypatch.setenv("AGENT_MEMORY_NEO4J_PASSWORD", "test-password") + monkeypatch.setattr(example, "Neo4jConnector", lambda **kwargs: connector) + monkeypatch.setattr( + example, + "SentenceTransformerEmbeddingProvider", + lambda spec: object(), + ) + + def fail_deployment(**kwargs: object) -> object: + raise RuntimeError("deployment binding failed") + + monkeypatch.setattr(example, "StorageDeployment", fail_deployment) + + with pytest.raises(RuntimeError, match="deployment binding failed"): + example.create_storage_deployment("amem-test") + + connector.close.assert_called_once_with() + + +@pytest.mark.parametrize( + ("failure_phase", "expected_close_count"), + [ + ("setup_storage", 0), + ("setup_adapter", 1), + ("setup_memory", 1), + ], +) +def test_main_records_setup_failure_and_closes_owned_storage( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + failure_phase: str, + expected_close_count: int, +) -> None: + example = _load_example() + output_dir = tmp_path / failure_phase + connector = SimpleNamespace(close=Mock()) + storage = SimpleNamespace(connector=connector) + args = _main_args(output_dir) + rows = _amem_rows() + _patch_common(example, monkeypatch, args=args, rows=rows) + + if failure_phase == "setup_storage": + monkeypatch.setattr( + example, + "create_storage_deployment", + lambda namespace: (_ for _ in ()).throw( + RuntimeError("storage setup failed") + ), + ) + else: + monkeypatch.setattr( + example, + "create_storage_deployment", + lambda namespace: storage, + ) + + if failure_phase == "setup_adapter": + monkeypatch.setattr( + example, + "LotusAdapter", + lambda **kwargs: (_ for _ in ()).throw( + RuntimeError("adapter setup failed") + ), + ) + else: + monkeypatch.setattr(example, "LotusAdapter", lambda **kwargs: object()) + + if failure_phase == "setup_memory": + monkeypatch.setattr( + example.am, + "AMem", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("memory setup failed")), + ) + + with pytest.raises(RuntimeError, match="setup failed"): + example.main() + + failure = (output_dir / "diagnostics" / "failure.json").read_text(encoding="utf-8") + assert '"type": "RuntimeError"' in failure + assert connector.close.call_count == expected_close_count + + +def test_main_closes_storage_once_after_retrieval_round_trip( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + example = _load_example() + output_dir = tmp_path / "success" + connector = SimpleNamespace(close=Mock()) + storage = SimpleNamespace(connector=connector) + args = _main_args(output_dir) + rows = _amem_rows() + memory = _memory( + [ + { + "_row_id": "note-1", + "content": "Speaker Carolinesays : Hello", + "tags": ["research"], + "context": "Caroline research", + } + ] + ) + _patch_common(example, monkeypatch, args=args, rows=rows) + monkeypatch.setattr(example, "create_storage_deployment", lambda namespace: storage) + monkeypatch.setattr(example, "LotusAdapter", lambda **kwargs: object()) + monkeypatch.setattr(example.am, "AMem", lambda **kwargs: memory) + + example.main() + + memory.add.assert_called_once_with( + {"content": rows[0]["content"], "timestamp": rows[0]["timestamp"]} + ) + assert [item.args for item in memory.query.call_args_list] == [ + (example.QUERY,), + (example.QUERY,), + ] + # The restore reopens the namespace on a second connection, so the run closes + # both the original and the restored connector. + assert connector.close.call_count == 2 + assert (output_dir / "retrieval" / "summary.json").exists() + assert (output_dir / "evolution" / "summary.json").exists() + assert (output_dir / "manifest.json").exists() diff --git a/tests/test_a_mem_policy.py b/tests/test_a_mem_policy.py new file mode 100644 index 0000000..7551e14 --- /dev/null +++ b/tests/test_a_mem_policy.py @@ -0,0 +1,682 @@ +"""Tests for the A-Mem note-evolution memory policy.""" + +from __future__ import annotations + +import json + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters.lotus import LotusAdapter +from agent_memory.adapters.lotus.relational import _minimum_value +from agent_memory.adapters.lotus.sem_agg import aggregate_groups_with_keys +from agent_memory.memories.a_mem import AMem, AMEM_NEO4J_STATEMENTS +from agent_memory.memories.a_mem.prompts import ( + ANALYSE_CONTENT_PROMPT, + EMBEDDING_TEXT_PROMPT, + EVOLUTION_SYSTEM_PROMPT, + NOTE_CONSOLIDATION_INSTRUCTION, +) +from agent_memory.planner import PolicyDifferentiator, RetrievalPlan +from agent_memory.planner.rules import DifferentialRules +from agent_memory.policy.aggregates import MinAggregateSpec, SemanticAggregateSpec +from agent_memory.policy.schema import output_columns +from agent_memory.runtime import MemoryRuntime + + +# Helpers and fake adapter + + +def _missing(value: object) -> bool: + """True for None and float NaN (pandas NULL representation).""" + return value is None or (isinstance(value, float) and pd.isna(value)) + + +def _as_list(value: object) -> list[object]: + """Coerce a scalar/list value into a list, treating NULL as empty.""" + if _missing(value): + return [] + return list(value) if isinstance(value, (list, tuple)) else [value] + + +def _ordered_union(*lists: list[object]) -> list[object]: + """Order-preserving union of several lists.""" + result: list[object] = [] + for items in lists: + for item in items: + if item not in result: + result.append(item) + return result + + +# The deterministic analyses the fake returns for the coffee/tea corpus. +COFFEE_TEA_ANALYSES: dict[str, dict[str, object]] = { + "I love coffee": { + "keywords": ["coffee"], + "context": "likes coffee", + "tags": ["coffee"], + }, + "I switched to tea": { + "keywords": ["tea"], + "context": "drinks tea", + "tags": ["tea"], + }, + "I only drink coffee on weekends": { + "keywords": ["coffee", "weekend"], + "context": "coffee on weekends", + "tags": ["coffee", "weekend"], + }, +} + + +class _FakeAMemAdapter(LotusAdapter): + """Deterministic stand-in for AMem's semantic operators. + + Deterministic operators run on the real adapter; ``sem_map`` and an ``agg`` + holding a semantic aggregate are replaced wholesale, because the engine calls + the model inside them and exposes no narrower hook. No prompt is read. + """ + + def __init__( + self, + *, + analyses: dict[str, dict[str, object]] | None = None, + actions_for: dict[str, tuple[str, ...]] | None = None, + ) -> None: + """Configure the fake model answers. + + ``actions_for`` maps a later note's content to the actions the model + returns; repeating ``"strengthen"`` n times emits n identical + ``suggested_connections`` entries, which the link relation is expected to + collapse. + """ + + super().__init__() + self.analyses = dict(analyses or {}) + self.actions_for = dict(actions_for or {}) + self.evolution_inputs: list[dict[str, object]] = [] + + def execute(self, query, inputs): + if query.op == "sem_map": + return self._sem_map(query, inputs) + if query.op == "agg" and self._has_semantic_aggregate(query): + return self._grouped_agg(query, inputs) + return super().execute(query, inputs) + + @staticmethod + def _has_semantic_aggregate(query) -> bool: + return any( + isinstance(spec, SemanticAggregateSpec) + for spec in query.params.get("aggregates", ()) + ) + + def _sem_map(self, query, inputs): + source = self.execute(query.inputs[0], inputs) + raw_outputs = query.params["output_cols"] + output_names = ( + [str(name) for name in raw_outputs] + if isinstance(raw_outputs, dict) + else [column.name for column in raw_outputs] + ) + result = source.copy() + if output_names == ["analysis"]: + # Fixture: the analysis the model would have returned for this content. + result["analysis"] = [ + json.dumps(self.analyses[str(row["content"])]) + for _, row in source.iterrows() + ] + elif output_names == ["embedding_text"]: + result["embedding_text"] = [ + self._embedding(row) for _, row in source.iterrows() + ] + elif output_names == ["evolution"]: + result["evolution"] = [ + self._latest_lane(row, "evolution") for _, row in source.iterrows() + ] + elif set(output_names) <= {"context", "tags"}: + for name in output_names: + result[name] = [ + self._merged_lane(row, name) for _, row in source.iterrows() + ] + else: + raise NotImplementedError( + f"unsupported sem_map output columns {output_names!r}" + ) + return result + + @classmethod + def _latest_lane(cls, row, column: str) -> object: + """Keep the stored value: the right side of a join-map join is the view.""" + + for name in (f"{column}:right", f"{column}:left"): + if name in row and not _missing(row[name]): + return row[name] + return None + + @classmethod + def _merged_lane(cls, row, column: str) -> object: + """Combine one column across the joined sides, as the group merge does. + + Must stay idempotent: incremental maintenance re-aggregates the + already-accumulated row. + """ + + values = [ + value + for name in (f"{column}:right", f"{column}:left") + if name in row + for value in _as_list(row[name]) + if not _missing(value) + ] + if column == "tags": + return _ordered_union(values) + return "; ".join(str(value) for value in values) if values else None + + @staticmethod + def _embedding(row) -> str: + """Concatenate the eq.3 fields in order.""" + + parts = [ + str(row.get("content", "")), + " ".join(str(item) for item in _as_list(row.get("keywords"))), + str(row.get("context", "")), + " ".join(str(item) for item in _as_list(row.get("tags"))), + ] + return " ".join(part for part in parts if part) + + def _grouped_agg(self, query, inputs): + """Mirror execute_agg, replacing only the per-group model call.""" + + source = self.execute(query.inputs[0], inputs) + aggregates = tuple(query.params.get("aggregates", ())) + names = list(output_columns(query)) + if source.empty: + return pd.DataFrame(columns=names) + rows: list[dict[str, object]] = [] + for key_values, group in aggregate_groups_with_keys(source): + row: dict[str, object] = dict(key_values) + for spec in aggregates: + if isinstance(spec, SemanticAggregateSpec): + for column, value in self._semantic(spec, group).items(): + row[column] = value + elif isinstance(spec, MinAggregateSpec): + row[spec.output_col] = _minimum_value(group, spec.columns) + else: + raise TypeError(f"unsupported aggregate spec {type(spec).__name__}") + rows.append(row) + return pd.DataFrame(rows, columns=names) + + def _semantic( + self, spec: SemanticAggregateSpec, group: pd.DataFrame + ) -> dict[str, object]: + """Return the per-group answer for one semantic aggregate.""" + + names = [column.name for column in spec.output_cols] + if "evolution" in names: + return {"evolution": self._evolution(group)} + if "context" in names and "tags" in names: + return { + "context": self._merge_context(group), + "tags": self._merge_tags(group), + } + raise NotImplementedError(f"unsupported semantic aggregate outputs {names!r}") + + def _evolution(self, group: pd.DataFrame) -> str: + """Build one decision JSON, defaulting to a single neighbour update.""" + + # Maintenance remerge: the group already holds computed `evolution` JSON, + # so pass the latest through instead of recomputing from candidate pairs. + if "evolution" in group.columns: + values = [value for value in group["evolution"] if not _missing(value)] + if values: + return values[-1] + if "_add_seq:earlier" in group.columns: + earlier = group.sort_values("_add_seq:earlier").iloc[0] + else: + earlier = group.iloc[0] + neighbor_add_seq = int(earlier["_add_seq:earlier"]) + later_content = str(group["content:later"].iloc[0]) + earliest_earlier_tags = _as_list(earlier.get("tags:earlier")) + self.evolution_inputs.append( + {"later": later_content, "earliest_earlier_tags": earliest_earlier_tags} + ) + new_tags = _ordered_union( + earliest_earlier_tags, + _as_list(group["tags:later"].iloc[0]), + ) + actions = self.actions_for.get(later_content, ("update_neighbor",)) + has_strengthen = "strengthen" in actions + has_update = "update_neighbor" in actions + return json.dumps( + { + "should_evolve": True, + "actions": list(actions), + # Repeating the action repeats the suggestion: the link relation is + # expected to collapse the duplicates. + "suggested_connections": [neighbor_add_seq] + * actions.count("strengthen"), + "tags_to_update": new_tags if has_strengthen else [], + "neighbor_updates": ( + [ + { + "neighbor_add_seq": neighbor_add_seq, + "new_context": "", + "new_tags": new_tags, + } + ] + if has_update + else [] + ), + } + ) + + @staticmethod + def _merge_tags(group: pd.DataFrame) -> list[object]: + return _ordered_union(*(_as_list(value) for value in group["tags"])) + + @staticmethod + def _merge_context(group: pd.DataFrame) -> object: + parts = [str(value) for value in group["context"] if not _missing(value)] + return "; ".join(parts) if parts else None + + +# Spec topology + + +def test_a_mem_exports_public_api() -> None: + """AMem is part of the public built-in policy surface.""" + + assert am.AMem is AMem + + +def test_a_mem_spec_exposes_only_note_view() -> None: + """Only the consolidated note is public; all stages stay private.""" + + spec = AMem.spec() + + assert tuple(spec.views) == ("note",) + assert tuple(spec.private_relations) == ( + "_analysed_notes", + "_earlier_notes", + "_later_notes", + "_candidate_note_pairs", + "_evolution_output", + "_evolution_decisions", + "_actions", + "_link_generation_writes", + "_note_links", + "_link_generation_notes", + "_unchanged_notes", + "_base_note_states", + "_memory_evolution_writes", + "_rewrites", + "_memory_evolution_note_states", + ) + assert tuple(spec.retrieval_queries) == ("default",) + assert "retrieval_query" not in spec.views + assert "_retrieved_notes" not in spec.private_relations + + +# Prompt pinning + + +def test_a_mem_analysis_and_embedding_prompts_pin_contract() -> None: + """Analysis and embedding prompts keep their structured contracts.""" + + assert "Generate a structured analysis" in ANALYSE_CONTENT_PROMPT + assert "keywords" in ANALYSE_CONTENT_PROMPT + assert "context" in ANALYSE_CONTENT_PROMPT + assert "tags" in ANALYSE_CONTENT_PROMPT + + # Content analysis is not evolution: the decision keys belong to the evolution + # prompt only (they were copy-pasted here once). + assert "should_evolve" not in ANALYSE_CONTENT_PROMPT + assert "actions" not in ANALYSE_CONTENT_PROMPT + assert "suggested_connections" not in ANALYSE_CONTENT_PROMPT + assert "tags_to_update" not in ANALYSE_CONTENT_PROMPT + assert "neighbor_updates" not in ANALYSE_CONTENT_PROMPT + + assert "Concatenate the following fields" in EMBEDDING_TEXT_PROMPT + assert "word-for-word" in EMBEDDING_TEXT_PROMPT + assert "content" in EMBEDDING_TEXT_PROMPT + assert "keywords" in EMBEDDING_TEXT_PROMPT + assert "context" in EMBEDDING_TEXT_PROMPT + assert "tags" in EMBEDDING_TEXT_PROMPT + + +def test_a_mem_evolution_and_consolidation_prompts_pin_contract() -> None: + """Evolution and consolidation prompts keep their action/merge contracts.""" + + assert "strengthen" in EVOLUTION_SYSTEM_PROMPT + assert "update_neighbor" in EVOLUTION_SYSTEM_PROMPT + assert "should_evolve" in EVOLUTION_SYSTEM_PROMPT + assert "neighbor_updates" in EVOLUTION_SYSTEM_PROMPT + assert "new_context" in EVOLUTION_SYSTEM_PROMPT + assert "new_tags" in EVOLUTION_SYSTEM_PROMPT + + consolidation = " ".join(NOTE_CONSOLIDATION_INSTRUCTION.split()) + assert "Merge the grouped context and tags" in consolidation + assert "Preserve every distinct fact" in consolidation + assert "prefer the newer entry" in consolidation + assert "larger _add_seq" in consolidation + + +# Retrieval shape + + +def test_a_mem_policy_differentiates() -> None: + """The policy compiles through the differential planner.""" + + policy = PolicyDifferentiator().differentiate( + AMem.spec(), + statements=AMEM_NEO4J_STATEMENTS, + ) + assert isinstance(policy.retrieval_queries["default"], RetrievalPlan) + + +def test_a_mem_retrieval_is_two_channel_storage_backed_dag() -> None: + """The two retrieval channels compile to two storage search nodes.""" + + spec = AMem.spec() + policy = PolicyDifferentiator().differentiate( + spec, statements=AMEM_NEO4J_STATEMENTS + ) + plan = policy.retrieval_queries["default"] + assert isinstance(plan, RetrievalPlan) + assert ( + len([node for node in plan.nodes.values() if node.execution_kind == "search"]) + == 2 + ) + assert ( + policy.fingerprint + == PolicyDifferentiator() + .differentiate(spec, statements=AMEM_NEO4J_STATEMENTS) + .fingerprint + ) + + +def test_a_mem_query_requires_a_search_capable_storage_backend() -> None: + """The default retrieval needs a storage backend, not an in-memory scan.""" + + memory = AMem(adapter=_FakeAMemAdapter()) + with pytest.raises(NotImplementedError, match="requires a storage backend"): + memory.query("What does the user remember?") + + +# Behavioural runtime + + +def _note_by_content(memory: AMem, content: str) -> dict[str, object]: + """Return the consolidated note with this exact content.""" + + for record in memory._runtime._state["note"].to_dict("records"): + if record["content"] == content: + return record + raise AssertionError(f"no consolidated note with content {content!r}") + + +def _note_link_rows( + memory: AMem, + adapter: _FakeAMemAdapter, +) -> list[tuple[object, object]]: + """Recompute the RELATES_TO edges from the current log. + + ``_note_links`` is a storage sink: it becomes a node only when the policy is + compiled with the statements, so the edges are recomputed from the log instead of + read out of a node state. + """ + + links = adapter.execute( + AMem.spec().private_relations["_note_links"], + {"log": memory._runtime._state["log"]}, + ) + return [ + (row["source_note_id"], row["target_note_id"]) + for row in links.to_dict("records") + ] + + +GROUPED_AGG_RULE_FAMILIES = ( + "rule-all-group", + "rule-all-group-optimized", + "rule-join-map", + "rule-re-group", +) + + +def _maintained_memory( + grouped_agg_rule: str, + adapter: _FakeAMemAdapter, + contents: tuple[str, ...], +) -> AMem: + """Return one AMem whose adds ran under the requested grouped-aggregate rule.""" + + policy = PolicyDifferentiator( + rules=DifferentialRules(grouped_agg_rule=grouped_agg_rule) + ).differentiate(AMem.spec()) + memory = AMem(adapter=adapter) + memory._runtime = MemoryRuntime(policy, adapter=adapter) + for index, content in enumerate(contents, start=1): + memory.add({"content": content, "timestamp": f"2026010{index}0000"}) + return memory + + +def _canonical_state_records( + records: list[dict[str, object]], +) -> list[dict[str, object]]: + """Return order-insensitive, NULL-safe records for state comparison.""" + + rows: list[dict[str, object]] = [] + for record in records: + row: dict[str, object] = {} + for key, value in record.items(): + if _missing(value): + row[key] = "" + elif isinstance(value, (list, tuple)): + row[key] = sorted(str(item) for item in value) + else: + row[key] = str(value) + rows.append(row) + return sorted(rows, key=repr) + + +def test_a_mem_accumulates_neighbor_updates() -> None: + """The engine accumulates note states across adds.""" + + adapter = _FakeAMemAdapter(analyses=COFFEE_TEA_ANALYSES) + memory = AMem(adapter=adapter) + + memory.add({"content": "I love coffee", "timestamp": "202601010000"}) + assert _note_by_content(memory, "I love coffee")["tags"] == ["coffee"] + + memory.add({"content": "I switched to tea", "timestamp": "202601020000"}) + assert _note_by_content(memory, "I love coffee")["tags"] == ["coffee", "tea"] + + memory.add( + {"content": "I only drink coffee on weekends", "timestamp": "202601030000"} + ) + assert _note_by_content(memory, "I love coffee")["tags"] == [ + "coffee", + "tea", + "weekend", + ] + + assert _note_by_content(memory, "I switched to tea")["tags"] == ["tea"] + assert _note_by_content(memory, "I only drink coffee on weekends")["tags"] == [ + "coffee", + "weekend", + ] + + # Frozen-analysis pin: note3's decision saw note1's *frozen* tags, not the + # accumulated [coffee, tea] — accumulation happens in consolidation. + weekend_input = next( + entry + for entry in adapter.evolution_inputs + if entry["later"] == "I only drink coffee on weekends" + ) + assert weekend_input["earliest_earlier_tags"] == ["coffee"] + + # min-recovery: the identity columns are identical on every state row of a note, + # so min() carries them through unchanged. + coffee = _note_by_content(memory, "I love coffee") + assert coffee["keywords"] == ["coffee"] + assert coffee["timestamp"] == "202601010000" + + +@pytest.mark.parametrize("grouped_agg_rule", GROUPED_AGG_RULE_FAMILIES) +def test_a_mem_incremental_note_matches_full_recompute(grouped_agg_rule: str) -> None: + """Incremental maintenance reproduces a from-scratch recompute of note.""" + + adapter = _FakeAMemAdapter(analyses=COFFEE_TEA_ANALYSES) + try: + memory = _maintained_memory( + grouped_agg_rule, + adapter, + ("I love coffee", "I switched to tea", "I only drink coffee on weekends"), + ) + except NotImplementedError as error: + pytest.skip(f"{grouped_agg_rule} is not expressible here: {error}") + + def canonical(records: list[dict[str, object]]) -> list[dict[str, object]]: + rows = [] + for record in records: + record = dict(record) + record["tags"] = sorted(record["tags"]) + record.pop("embedding_text", None) + rows.append(record) + return sorted(rows, key=lambda row: str(row["content"])) + + full = adapter.execute( + AMem.spec().views["note"].query, + {"log": memory._runtime._state["log"]}, + ) + incremental = memory._runtime._state["note"] + assert canonical(full.to_dict("records")) == canonical( + incremental.to_dict("records") + ) + + +@pytest.mark.parametrize("grouped_agg_rule", GROUPED_AGG_RULE_FAMILIES) +def test_a_mem_maintenance_query_ignores_an_empty_change( + grouped_agg_rule: str, +) -> None: + """Every generated maintenance query is a fixed point with no new rows.""" + + adapter = _FakeAMemAdapter( + analyses=COFFEE_TEA_ANALYSES, + actions_for={"I switched to tea": ("strengthen", "update_neighbor")}, + ) + try: + memory = _maintained_memory( + grouped_agg_rule, adapter, ("I love coffee", "I switched to tea") + ) + except NotImplementedError as error: + pytest.skip(f"{grouped_agg_rule} is not expressible here: {error}") + + policy = memory._runtime.policy + node_state = memory._runtime._engine.node_state + checked = 0 + for node_id in policy.execution_order: + node = policy.nodes[node_id] + current = node_state[node_id] + if node.maintenance_query is None or current.empty: + continue + inputs = dict(node_state) + inputs[node_id] = current + for parent_id in node.input_node_ids: + inputs[f"{parent_id}__inserted"] = node_state[parent_id].iloc[0:0] + maintained = adapter.execute(node.maintenance_query, inputs) + assert _canonical_state_records( + maintained.to_dict("records") + ) == _canonical_state_records(current.to_dict("records")), ( + f"{node_id} ({node.execution_kind}) is not a fixed point" + ) + checked += 1 + # Only populated nodes can be checked; an empty node has nothing to maintain. + expected = sum( + 1 + for node_id, node in policy.nodes.items() + if node.maintenance_query is not None and not node_state[node_id].empty + ) + assert checked == expected + + +def test_a_mem_single_add_handles_strengthen_and_update_neighbor() -> None: + """One decision fans out to both strengthen and update_neighbor.""" + + adapter = _FakeAMemAdapter( + analyses=COFFEE_TEA_ANALYSES, + actions_for={"I switched to tea": ("strengthen", "update_neighbor")}, + ) + memory = AMem(adapter=adapter) + + memory.add({"content": "I love coffee", "timestamp": "202601010000"}) + memory.add({"content": "I switched to tea", "timestamp": "202601020000"}) + + # update_neighbor: note1 was rewritten by note2. + assert _note_by_content(memory, "I love coffee")["tags"] == ["coffee", "tea"] + # strengthen: note2's own tags were overridden (otherwise they'd be ["tea"]). + assert _note_by_content(memory, "I switched to tea")["tags"] == ["coffee", "tea"] + + +def test_a_mem_strengthen_alone_overrides_tags_and_leaves_neighbor() -> None: + """strengthen (LINK) alone overrides the new note's tags, not the neighbor's.""" + + adapter = _FakeAMemAdapter( + analyses=COFFEE_TEA_ANALYSES, + actions_for={"I switched to tea": ("strengthen",)}, + ) + memory = AMem(adapter=adapter) + + memory.add({"content": "I love coffee", "timestamp": "202601010000"}) + memory.add({"content": "I switched to tea", "timestamp": "202601020000"}) + + # strengthen overrides note2's tags (would otherwise be ["tea"]). + assert _note_by_content(memory, "I switched to tea")["tags"] == ["coffee", "tea"] + # no update_neighbor: note1 is untouched. + assert _note_by_content(memory, "I love coffee")["tags"] == ["coffee"] + + +def test_a_mem_strengthen_emits_one_deduplicated_note_link() -> None: + """Repeated strengthen suggestions collapse to one RELATES_TO edge.""" + + adapter = _FakeAMemAdapter( + analyses=COFFEE_TEA_ANALYSES, + actions_for={"I switched to tea": ("strengthen", "strengthen")}, + ) + memory = AMem(adapter=adapter) + + memory.add({"content": "I love coffee", "timestamp": "202601010000"}) + memory.add({"content": "I switched to tea", "timestamp": "202601020000"}) + + coffee_id = _note_by_content(memory, "I love coffee")["_row_id"] + tea_id = _note_by_content(memory, "I switched to tea")["_row_id"] + # The pair is the link's identity, so a duplicate row would be rejected by the + # storage sink's inserted-key validation. + assert _note_link_rows(memory, adapter) == [(tea_id, coffee_id)] + + +def test_a_mem_evolution_state_carries_identity_and_the_writer_order() -> None: + """A rewrite's state row is the rewritten note, stamped with the writer's _add_seq.""" + + adapter = _FakeAMemAdapter(analyses=COFFEE_TEA_ANALYSES) + memory = AMem(adapter=adapter) + + memory.add({"content": "I love coffee", "timestamp": "202601010000"}) + memory.add({"content": "I switched to tea", "timestamp": "202601020000"}) + + states = adapter.execute( + AMem.spec().private_relations["_memory_evolution_note_states"], + {"log": memory._runtime._state["log"]}, + ).to_dict("records") + + # note2 rewrote note1: identity comes from the rewritten note, _add_seq from the + # rewriting one, and no identity column is left NULL. + assert [row["content"] for row in states] == ["I love coffee"] + assert [row["_add_seq"] for row in states] == [1] + assert not any( + _missing(row[column]) + for row in states + for column in ("content", "keywords", "timestamp") + )