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") + )