diff --git a/src/agent_memory/__init__.py b/src/agent_memory/__init__.py index 31954a0..ad49680 100644 --- a/src/agent_memory/__init__.py +++ b/src/agent_memory/__init__.py @@ -22,6 +22,7 @@ Relation, UserQuery, array_agg, + arg_min, avg, case_when, collect_list, @@ -59,6 +60,7 @@ "ZepMemoryExtended", "__version__", "array_agg", + "arg_min", "avg", "case_when", "collect_list", diff --git a/src/agent_memory/adapters/lotus/adapter.py b/src/agent_memory/adapters/lotus/adapter.py index 030b3d8..8517108 100644 --- a/src/agent_memory/adapters/lotus/adapter.py +++ b/src/agent_memory/adapters/lotus/adapter.py @@ -2,9 +2,10 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sized from dataclasses import dataclass, field from hashlib import sha256 +import re from typing import Any from uuid import uuid4 @@ -53,6 +54,7 @@ from agent_memory.adapters.lotus.sem_map import execute_sem_map from agent_memory.adapters.lotus.sem_topk import execute_sem_topk from agent_memory.tracing.semantic import ( + measure_semantic_trace_io, query_digest, semantic_trace_scope, write_compact_operator_trace, @@ -66,6 +68,12 @@ ) from agent_memory.policy.logical import QueryExpr from agent_memory.storage.embedding import EmbeddingProvider +from agent_memory.planner.differential_policy import DifferentiatedPolicy +from agent_memory.planner.physical import COMBINED_VERSION, FUSION_VERSION, SUMMARY_VERSION, REPRESENTATIVE_VERSION, optimize_policy, walk +from agent_memory.adapters.lotus.fusion import execute_target_state + +_FUSION_VERSIONS = {"zep-target-state": FUSION_VERSION, "zep-combined": COMBINED_VERSION, + "zep-fact-summary": SUMMARY_VERSION, "zep-representative": REPRESENTATIVE_VERSION} DEFAULT_LOTUS_MODEL = "deepseek/deepseek-v4-pro" @@ -78,6 +86,7 @@ class LotusAdapter: config: LotusExecutionConfig = field(default_factory=LotusExecutionConfig) pair_embedding_provider: EmbeddingProvider | None = None _context: LotusExecutionContext = field(init=False, repr=False) + _independent_context: LotusExecutionContext | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: """Initialize shared LOTUS execution context.""" @@ -88,10 +97,245 @@ def __post_init__(self) -> None: pair_embedding_provider=self.pair_embedding_provider, ) + @property + def supports_legacy_restore(self) -> bool: + """Physical rewrites require the node-state executor and its identity.""" + return self.config.physical_fusion == "disabled" and not self.config.pair_filter_batching and not self.config.listwise_join_batching and not self.config.groupby_prompt_batching and self.config.sem_agg_prompt_batching is None and not self.config.predicate_reuse_sites and not any(p.top_k is not None and p.mode != "oracle-only" for p in self.config.semantic_pair_profiles.values()) + + def bounded_predicate_scope(self, query: QueryExpr) -> tuple[tuple[str, ...], tuple[str, ...]] | None: + """Declare non-row-local filter dependencies and complete Top-k buckets.""" + profile = self.config.semantic_pair_profiles.get(query_digest(query)) + if query.op != "sem_filter" or profile is None or profile.mode == "oracle-only" or profile.top_k is None: + return None + from agent_memory.policy.schema import output_columns + + columns = output_columns(query.inputs[0]) + referenced = tuple(dict.fromkeys(re.findall(r"\{([^{}]+)\}", str(query.params["instruction"])))) + # Unknown dependencies disable the metadata-only shortcut, not screening. + semantic = referenced if referenced and set(referenced) <= set(columns) else tuple(columns) + dependencies = tuple(dict.fromkeys((*semantic, *profile.left_id_columns, + *profile.right_id_columns, *profile.left_text_columns, *profile.right_text_columns))) + buckets = (profile.right_id_columns if profile.direction == "right-to-left" + else profile.left_id_columns if profile.direction == "left-to-right" else ()) + return dependencies, buckets + @property def maintenance_execution_fingerprint(self) -> str: """Identify physical settings that can change maintained state.""" + base = self._maintenance_execution_fingerprint() + if self.config.parallel_fact_extraction: + base = sha256(f"{base}\nindependent-fact-extraction-v2:cache={self.config.lm_enable_cache}".encode()).hexdigest() + if any(p.top_k is not None and p.mode != "oracle-only" for p in self.config.semantic_pair_profiles.values()): + base = sha256(f"{base}\ncomplete-candidate-topk-v2".encode()).hexdigest() + if self.config.reuse_unchanged_entity_name: + base = sha256(f"{base}\nunchanged-entity-name-v1".encode()).hexdigest() + if self.config.groupby_prompt_batching: + from agent_memory.planner.serialization import stable_json + base = sha256(stable_json({"base": base, "groupby-prompt-batching-v1": { + site: setting.to_dict() for site, setting in sorted(self.config.groupby_prompt_batching.items()) + }}).encode()).hexdigest() + if self.config.listwise_join_batching: + from agent_memory.planner.serialization import stable_json + base = sha256(stable_json({"base": base, "listwise-join-batching-v1": { + site: setting.to_dict() for site, setting in sorted(self.config.listwise_join_batching.items()) + }}).encode()).hexdigest() + if self.config.predicate_reuse_sites: + base = sha256(f"{base}\npredicate-reuse-v1:{sorted(self.config.predicate_reuse_sites)}".encode()).hexdigest() + if self.config.sem_agg_prompt_batching is not None: + base = sha256(f"{base}\nsem-agg-prompt-batching-v1:{self.config.sem_agg_prompt_batching.fingerprint}".encode()).hexdigest() + if self.config.pair_filter_batching: + from agent_memory.planner.serialization import stable_json + base = sha256(stable_json({"base": base, "pair_filter_batching": { + site: config.to_dict() for site, config in sorted(self.config.pair_filter_batching.items()) + }}).encode()).hexdigest() + if self.config.physical_fusion == "disabled": + return base + version = _FUSION_VERSIONS[self.config.physical_fusion] + if self.config.physical_fusion in {"zep-combined", "zep-fact-summary", "zep-representative"}: + from agent_memory.planner.serialization import stable_json + base = sha256(stable_json({"base": base, "model": self.model, + "lm_model_kwargs": dict(self.config.lm_model_kwargs), + "structured_max_tokens": self.config.structured_max_tokens, + "structured_parse_retries": self.config.structured_parse_retries, + "sem_groupby_default": self.config.sem_groupby_default, + "sem_groupby_pair_batch_size": self.config.sem_groupby_pair_batch_size, + }).encode()).hexdigest() + return sha256(f"{base}\n{version}".encode()).hexdigest() + + def prepare_policy(self, policy: DifferentiatedPolicy) -> DifferentiatedPolicy: + """Validate and lower explicitly enabled physical execution strategies.""" + prepared = optimize_policy(policy, strategy=self.config.physical_fusion) + if self.config.groupby_prompt_batching: + from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id + from agent_memory.adapters.lotus.sem_groupby import validate_site_groupby_batching + found = set() + roots = [node.query for node in prepared.nodes.values()] + roots.extend(node.maintenance_query for node in prepared.nodes.values() if node.maintenance_query is not None) + for root in roots: + for query in walk(root): + if query.op != "sem_groupby": + continue + site = semantic_pair_site_id(query) + if site in self.config.groupby_prompt_batching: + validate_site_groupby_batching(query, self.config) + found.add(site) + if found != set(self.config.groupby_prompt_batching): + raise ValueError("configured groupby batching site not found in executable plan") + if self.config.listwise_join_batching: + from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id + found = set() + pending = [node.query for node in prepared.nodes.values()] + pending.extend(node.maintenance_query for node in prepared.nodes.values() if node.maintenance_query is not None) + while pending: + query = pending.pop() + if query.op == "fused_target_state": + continue + pending.extend(query.inputs) + if query.op != "sem_join": + continue + try: + site = semantic_pair_site_id(query) + except ValueError: + continue + if site not in self.config.listwise_join_batching: + continue + profile = self.config.semantic_pair_profiles.get(query_digest(query)) + if not query.params.get("k") or (profile is not None and profile.mode == "proxy-only"): + raise ValueError("site listwise batching requires a top-k oracle join") + found.add(site) + if found != set(self.config.listwise_join_batching): + raise ValueError("configured listwise batching site not found in executable plan") + if self.config.predicate_reuse_sites: + from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id + from agent_memory.planner.physical import enable_predicate_reuse + sites: dict[str, set[str]] = {} + for node_id, node in prepared.nodes.items(): + if node.query.op != "sem_filter": + continue + try: + site = semantic_pair_site_id(node.query) + except ValueError: + continue + sites.setdefault(site, set()).add(node_id) + if not set(self.config.predicate_reuse_sites) <= sites.keys(): + raise ValueError("configured predicate reuse site not found in plan") + prepared = enable_predicate_reuse(prepared, { + node_id for site in self.config.predicate_reuse_sites for node_id in sites[site] + }) + if self.config.pair_filter_batching: + from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id, semantic_pair_site_physical_contract + from agent_memory.policy.schema import output_columns + found = set() + queries = [n.query for n in prepared.nodes.values()] + queries.extend(n.maintenance_query for n in prepared.nodes.values() if n.maintenance_query is not None) + for root in queries: + for query in walk(root): + if query.op != "sem_filter": + continue + try: + site = semantic_pair_site_id(query) + except ValueError: + continue + setting = self.config.pair_filter_batching.get(site) + if setting is None: + continue + semantic_pair_site_physical_contract(query) + columns = output_columns(query.inputs[0]) + if not set(setting.group_by) <= set(columns) or any( + not col.split(":", 1)[0].endswith("_id") or ":later" not in col for col in setting.group_by + ): + raise ValueError("site batching requires stable later endpoint ID columns") + profile = self.config.semantic_pair_profiles.get(query_digest(query)) + if profile is not None and profile.mode == "proxy-only": + raise ValueError("site batching cannot replace proxy-only decisions") + found.add(site) + if found != set(self.config.pair_filter_batching): + raise ValueError("configured pair-filter batching site not found in plan") + for node in prepared.nodes.values(): + if node.maintenance_query is None: + continue + for query in walk(node.maintenance_query): + if query.op != "fused_target_state": + continue + if self.config.physical_fusion not in _FUSION_VERSIONS: + raise ValueError("fused plan requires matching physical_fusion configuration") + profile = self.config.semantic_pair_profiles.get(query.params.get("join_profile_digest", query_digest(query.inputs[0]))) + if profile is not None and profile.mode not in {"oracle-only", "search-filter"}: + raise ValueError("target-state fusion cannot replace proxy-only decisions") + self.independent_node_id(prepared) + return prepared + + def independent_node_id(self, policy: DifferentiatedPolicy) -> str | None: + """Select only the registered Zep extraction with raw entity inputs.""" + if not self.config.parallel_fact_extraction: + return None + from agent_memory.memories.zep.representative import ZepRepresentativeMemory + + template = next(q for q in walk(ZepRepresentativeMemory.facts.expr) + if q.op == "sem_flat_map" and "entities" in q.params["input_cols"]) + candidates = [n for n in policy.nodes.values() + if n.execution_kind == "semantic_row" and n.query.op == template.op + and n.query.params == template.params] + if len(candidates) != 1: + raise ValueError("parallel fact extraction requires exactly one registered extraction node") + node = candidates[0] + # Check the compiled upstream graph, not just prompt text or output columns. + reference = ZepRepresentativeMemory.differentiate_policy() + pending = [node.node_id] + seen: set[str] = set() + while pending: + current = pending.pop() + if current in seen: + continue + seen.add(current) + actual, expected = policy.nodes[current], reference.nodes.get(current) + if expected is None or actual.query != expected.query or actual.input_node_ids != expected.input_node_ids: + raise ValueError("parallel fact extraction input contract differs from registered Zep view") + pending.extend(actual.input_node_ids) + return node.node_id + + def execute_independent(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: + """Execute one extraction using private request metadata and usage counters.""" + if not self.config.parallel_fact_extraction or query.op != "sem_flat_map": + raise ValueError("independent execution is not enabled for this operator") + if self._independent_context is None: + self._independent_context = self._context.fork() + context = self._independent_context + if context._scoped_lm is None: + raise RuntimeError("independent LM context was not configured") + with context._scoped_lm.scope(context._lm), measure_semantic_trace_io() as trace_io, semantic_trace_scope( + semantic_operator=query.op, operator_call_id=f"{query.op}-{query_digest(query)}-{uuid4().hex[:8]}", + query_digest=query_digest(query), execution_lane="independent-fact-extraction-v2", + semantic_trace_snapshot_mode=self.config.semantic_trace_snapshot_mode, + ): + try: + result = execute_sem_flat_map( + query, inputs, self.execute, context, model=context._lm, + ) + except Exception as exc: + if self.config.lm_enable_cache is True: + try: + self._write_framework_cache_usage(query, + usage=context.consume_cache_usage_delta(), status="error", error=exc) + except Exception as trace_error: + exc.add_note(f"LOTUS cache usage trace failed: {trace_error}") + raise + if self.config.lm_enable_cache is True: + self._write_framework_cache_usage(query, + usage=context.consume_cache_usage_delta(), status="success", output=result) + # Worker I/O can overlap main-thread work; do not subtract it again + # from event wall time or mutate the caller's measurement object. + write_trace_event(self.config.trace_dir(), operator=query.op, + event_type="independent_execution_result", payload={ + "trace_io_latency_ms": trace_io.latency_ms, + "trace_bytes_written": trace_io.bytes_written, + }) + return result + + def _maintenance_execution_fingerprint(self) -> str: + """Keep the pre-fusion default fingerprint byte-for-byte stable.""" + pair_fingerprint = semantic_pair_profiles_fingerprint( self.config.semantic_pair_profiles ) @@ -140,8 +384,28 @@ def _base_maintenance_execution_fingerprint(self, pair_fingerprint: str) -> str: def execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: """Execute a logical query expression through LOTUS.""" + from agent_memory.adapters.lotus.identity_reuse import IDENTITY_DECISIONS_INPUT, identity_scope + if IDENTITY_DECISIONS_INPUT in inputs: + scoped = dict(inputs) + values = scoped.pop(IDENTITY_DECISIONS_INPUT) + with identity_scope(values): + return self._execute(query, scoped) + return self._execute(query, inputs) + + def _execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: + """Dispatch within the optional executor-owned identity scope.""" + if self.config.physical_fusion in {"zep-fact-summary", "zep-representative"} and query.op in {"agg", "sem_map"}: + from agent_memory.planner.physical import lower_fact_summary_node + query = lower_fact_summary_node(query) match query.op: + case "fused_target_state": + if self.config.physical_fusion not in _FUSION_VERSIONS: + raise ValueError("fused plan requires matching physical_fusion configuration") + version = _FUSION_VERSIONS[self.config.physical_fusion] + if query.params.get("version") != version: + raise ValueError("unsupported target-state fusion version") + return self._execute_traced_semantic(query, inputs, execute_target_state) case "log": return execute_log(inputs) case "materialized_view": @@ -219,6 +483,9 @@ def execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: case "sem_agg": return self._execute_traced_semantic(query, inputs, execute_sem_agg) case "sem_map": + if query.params.get("fact_summary_map"): + from agent_memory.adapters.lotus.relational import execute_fact_summary_map + return self._execute_traced_semantic(query, inputs, execute_fact_summary_map) return self._execute_traced_semantic(query, inputs, execute_sem_map) case "sem_topk": return self._execute_traced_semantic(query, inputs, execute_sem_topk) @@ -315,7 +582,7 @@ def _write_framework_cache_usage( **usage, } output_frame = output[0] if isinstance(output, tuple) and output else output - if hasattr(output_frame, "__len__"): + if isinstance(output_frame, Sized): payload["output_row_count"] = len(output_frame) if error is not None: payload.update( diff --git a/src/agent_memory/adapters/lotus/context.py b/src/agent_memory/adapters/lotus/context.py index 97171e9..3c388f9 100644 --- a/src/agent_memory/adapters/lotus/context.py +++ b/src/agent_memory/adapters/lotus/context.py @@ -5,6 +5,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from pathlib import Path +from threading import Lock from typing import Any from agent_memory.adapters.lotus.pair_execution import ( @@ -16,6 +17,8 @@ ) from agent_memory.adapters.lotus.provider_usage_lm import provider_usage_tracing_lm_class from agent_memory.adapters.lotus.traced_lm import TracedLM +from agent_memory.adapters.lotus.scoped_lm import ScopedLM +from agent_memory.adapters.lotus.site_batching import PairFilterBatching from agent_memory.storage.embedding import EmbeddingProvider DEFAULT_STRUCTURED_MAX_TOKENS = 8192 @@ -68,7 +71,15 @@ class LotusExecutionConfig: default_factory=dict ) prompt_batching: PromptBatching | None = None + sem_agg_prompt_batching: PromptBatching | None = None structured_output_transport: str = "chat-json-object" + physical_fusion: str = "disabled" + pair_filter_batching: Mapping[str, PairFilterBatching] = field(default_factory=dict) + listwise_join_batching: Mapping[str, PromptBatching] = field(default_factory=dict) + groupby_prompt_batching: Mapping[str, PromptBatching] = field(default_factory=dict) + reuse_unchanged_entity_name: bool = False + parallel_fact_extraction: bool = False + predicate_reuse_sites: tuple[str, ...] = () sem_filter_examples: Sequence[Mapping[str, Any]] | None = None sem_filter_helper_examples: Sequence[Mapping[str, Any]] | None = None @@ -114,6 +125,73 @@ def __post_init__(self) -> None: """Validate bounded semantic execution settings.""" validate_lm_max_ctx_len(self.lm_model_kwargs.get("max_ctx_len")) + if not isinstance(self.parallel_fact_extraction, bool): + raise TypeError("parallel_fact_extraction must be boolean") + if self.parallel_fact_extraction and ( + self.physical_fusion != "zep-representative" or not isinstance(self.lm_enable_cache, bool) + or self.prompt_batching is not None + ): + raise ValueError("parallel fact extraction requires zep-representative, explicit cache mode, and no global batching") + if not isinstance(self.reuse_unchanged_entity_name, bool): + raise TypeError("reuse_unchanged_entity_name must be boolean") + if self.reuse_unchanged_entity_name and self.physical_fusion != "zep-representative": + raise ValueError("unchanged entity names require registered zep-representative fusion") + if self.groupby_prompt_batching: + if self.prompt_batching is not None or self.sem_groupby_default: + raise ValueError("site groupby batching requires no global batching and default=False") + for site, setting in self.groupby_prompt_batching.items(): + if (not isinstance(site, str) or not site.startswith("sem_groupby:") + or not isinstance(setting, PromptBatching) or setting.max_tasks is None): + raise ValueError("groupby batching requires sem_groupby sites and bounded PromptBatching") + if self.listwise_join_batching: + if (self.prompt_batching is not None or self.sem_join_topk_method != "listwise" + or self.sem_join_examples is not None or self.sem_join_strategy is not None + or self.sem_join_cascade_args is not None or self.sem_join_safe_mode): + raise ValueError("site listwise batching requires plain listwise execution without global batching") + for site, setting in self.listwise_join_batching.items(): + if (not isinstance(site, str) or not site.startswith("sem_join:") + or not isinstance(setting, PromptBatching) or setting.max_tasks is None): + raise ValueError("listwise batching requires sem_join sites and bounded PromptBatching") + if self.sem_agg_prompt_batching is not None: + if not isinstance(self.sem_agg_prompt_batching, PromptBatching): + raise TypeError("sem_agg_prompt_batching must be PromptBatching or None") + if self.prompt_batching is not None or self.sem_agg_dispatch != "sequential": + raise ValueError("aggregate-only batching requires no global batching and sequential dispatch") + if self.sem_agg_safe_mode or self.sem_agg_model_kwargs: + raise ValueError("aggregate-only batching does not support safe mode or custom aggregate model kwargs") + if self.predicate_reuse_sites: + if not isinstance(self.predicate_reuse_sites, tuple) or any( + not isinstance(site, str) or not site.startswith("sem_filter:") + for site in self.predicate_reuse_sites + ) or len(set(self.predicate_reuse_sites)) != len(self.predicate_reuse_sites): + raise ValueError("predicate_reuse_sites must contain distinct sem_filter sites") + from agent_memory.adapters.lotus.sem_filter_batch_prompting import validate_batch_prompting_sem_filter_config + validate_batch_prompting_sem_filter_config(self) + if self.pair_filter_batching: + if self.prompt_batching is not None or self.structured_output_transport != "chat-json-object": + raise ValueError("site batching requires no global batching and chat-json-object") + for site, setting in self.pair_filter_batching.items(): + if not site.startswith("sem_filter:") or not isinstance(setting, PairFilterBatching): + raise ValueError("site batching supports registered pair-shaped sem_filter sites only") + from agent_memory.adapters.lotus.sem_filter_batch_prompting import validate_batch_prompting_sem_filter_config + validate_batch_prompting_sem_filter_config(self) + if self.physical_fusion not in {"disabled", "zep-target-state", "zep-combined", "zep-fact-summary", "zep-representative"}: + raise ValueError(f"unknown physical fusion strategy: {self.physical_fusion!r}") + if self.physical_fusion != "disabled": + unsupported = ( + self.prompt_batching is not None + or self.structured_output_transport != "chat-json-object" + or self.sem_join_topk_method != "listwise" + or self.sem_join_examples is not None + or self.sem_join_strategy is not None + or self.sem_join_cascade_args is not None + or self.sem_join_safe_mode + or self.sem_agg_safe_mode + or bool(self.sem_agg_model_kwargs) + or self.sem_agg_dispatch != "sequential" + ) + if unsupported: + raise ValueError("target-state fusion requires plain listwise/chat-json-object execution") validate_structured_output_transport(self.structured_output_transport) if self.prompt_batching is not None and not isinstance( self.prompt_batching, PromptBatching @@ -179,7 +257,9 @@ class LotusExecutionContext: config: LotusExecutionConfig = field(default_factory=LotusExecutionConfig) pair_embedding_provider: EmbeddingProvider | None = None _configured: bool = field(default=False, init=False, repr=False) + _configure_lock: Lock = field(default_factory=Lock, init=False, repr=False, compare=False) _lm: Any | None = field(default=None, init=False, repr=False) + _scoped_lm: ScopedLM | None = field(default=None, init=False, repr=False) _reported_cache_usage: dict[str, int] = field( default_factory=dict, init=False, @@ -196,10 +276,36 @@ def __post_init__(self) -> None: def configure(self) -> None: """Configure LOTUS before invoking semantic dataframe operators.""" - if self._configured: - return + with self._configure_lock: + if self._configured: + return + + import lotus + + lm = self.new_lm() + scoped_lm = ScopedLM(lm) if self.config.parallel_fact_extraction else None + settings_kwargs: dict[str, Any] = {"lm": scoped_lm if scoped_lm is not None else lm} + if self.config.lm_enable_cache is not None: + settings_kwargs["enable_cache"] = self.config.lm_enable_cache + lotus.settings.configure(**settings_kwargs) + self._lm = lm + self._scoped_lm = scoped_lm + self._configured = True + + def fork(self) -> LotusExecutionContext: + """Create private LM/cache counters without reconfiguring global settings.""" + self.configure() + if self._scoped_lm is None: + raise RuntimeError("independent LM context requires parallel execution") + context = LotusExecutionContext(self.model, self.config, self.pair_embedding_provider) + context._lm = self.new_lm() + context._scoped_lm = self._scoped_lm + context._configured = True + return context + + def new_lm(self) -> Any: + """Build a private LM with the same contract, without changing LOTUS settings.""" - import lotus from lotus.cache import InMemoryCache from lotus.models import LM @@ -235,16 +341,10 @@ def configure(self) -> None: lm_class = deepseek_responses_lm_class(LM) base_lm = ( provider_usage_tracing_lm_class(lm_class)(**lm_kwargs, trace_dir=trace_dir) - if trace_dir is not None or self.config.prompt_batching is not None + if trace_dir is not None or self.config.prompt_batching is not None or self.config.sem_agg_prompt_batching is not None or self.config.pair_filter_batching or self.config.listwise_join_batching or self.config.groupby_prompt_batching else lm_class(**lm_kwargs) ) - lm = TracedLM(base_lm, trace_dir) if trace_dir is not None else base_lm - self._lm = lm - settings_kwargs: dict[str, Any] = {"lm": lm} - if self.config.lm_enable_cache is not None: - settings_kwargs["enable_cache"] = self.config.lm_enable_cache - lotus.settings.configure(**settings_kwargs) - self._configured = True + return TracedLM(base_lm, trace_dir) if trace_dir is not None else base_lm def cache_usage_snapshot(self) -> dict[str, int]: """Return physical, virtual, and framework-cache counters for this LM.""" diff --git a/src/agent_memory/adapters/lotus/fusion.py b/src/agent_memory/adapters/lotus/fusion.py new file mode 100644 index 0000000..0f3c8ad --- /dev/null +++ b/src/agent_memory/adapters/lotus/fusion.py @@ -0,0 +1,363 @@ +"""Joint target selection and state synthesis, with deterministic writeback.""" + +from collections.abc import Callable, Mapping +from dataclasses import replace +import json +from typing import Any + +import pandas as pd + +from agent_memory.adapters.lotus.context import LotusExecutionContext +from agent_memory.adapters.lotus.pair_execution import ( + PAIR_LEFT_ID_COLUMN, + PAIR_RIGHT_ID_COLUMN, +) +from agent_memory.adapters.lotus.sem_join import ( + assemble_join_frame, + join_series, + semantic_join_pair_candidates, +) +from agent_memory.adapters.lotus.sem_topk_join import _apply_pair_profile +from agent_memory.adapters.lotus.structured import execute_structured_lm_retry_result +from agent_memory.planner.physical import replace_query +from agent_memory.planner.rules import JOIN_MAP_BINDING_NAME, JOIN_MAP_RIGHT_ID_COLUMN +from agent_memory.policy.aggregates import SemanticAggregateSpec +from agent_memory.policy.logical import QueryExpr +from agent_memory.policy.schema import output_columns +from agent_memory.tracing.semantic import query_digest, write_trace_event + +_STATE_BINDING = "__fused_target_states" +_SYSTEM = ( + "Resolve the supplied semantic join and state-consolidation tasks jointly. " + "For every incoming ID choose one eligible target ID or null. Copy the target " + "ID exactly from that incoming item's eligible_targets list; never use a target " + "ID from another incoming item. Select only true " + "matches, not merely related targets. If several match, choose the strongest. " + "For each selected target, consolidate its old state and ALL incoming states " + "assigned to it, exactly once, following the consolidation instruction. " + "Do not generate states for unselected targets. Do not edit IDs. " + "Return only the requested JSON object." +) + + +def parse_resolution(raw: str, eligible: Mapping[str, set[str]], *, + state_columns: tuple[str, ...] = ("relation_type", "fact")) -> dict[str, Any]: + """Validate complete assignments, candidate membership, and canonical states.""" + try: + value = json.loads(raw) + except json.JSONDecodeError as error: + raise ValueError("fusion output is not valid JSON") from error + if not isinstance(value, dict) or set(value) != {"matches", "states"}: + raise ValueError("fusion output requires only matches and states") + if not isinstance(value["matches"], list) or not isinstance(value["states"], list): + raise ValueError("fusion matches and states must be arrays") + assigned: dict[str, str | None] = {} + for item in value["matches"]: + if not isinstance(item, dict) or set(item) != {"left_id", "right_id"}: + raise ValueError("invalid fusion assignment schema") + left, right = item["left_id"], item["right_id"] + if not isinstance(left, str) or left not in eligible or left in assigned: + raise ValueError("unknown or duplicate incoming ID") + if right is not None and ( + not isinstance(right, str) or right not in eligible[left] + ): + raise ValueError( + f"target ID {right!r} is not eligible for incoming row {left!r}; " + f"choose one of {sorted(eligible[left])!r} or null" + ) + assigned[left] = right + if set(assigned) != set(eligible): + raise ValueError("fusion output omitted incoming IDs") + states: dict[str, dict[str, str]] = {} + for item in value["states"]: + if not isinstance(item, dict) or set(item) != {"right_id", *state_columns}: + raise ValueError("invalid canonical-state schema") + if any(not isinstance(v, str) for v in item.values()): + raise ValueError("canonical-state fields must be strings") + identifier = item["right_id"] + if identifier in states: + raise ValueError("duplicate canonical target state") + states[identifier] = {key: item[key] for key in state_columns} + if set(states) != {target for target in assigned.values() if target is not None}: + raise ValueError("canonical states must cover exactly the selected targets") + return {"matches": assigned, "states": states} + + +def _unchanged_entity_states( + eligible: Mapping[str, set[str]], + fixed: Mapping[str, str | None], + left: pd.DataFrame, + right: pd.DataFrame, + left_ids: Mapping[str, Any], + right_ids: Mapping[str, Any], +) -> dict[str, dict[str, str]] | None: + """Copy names only after every screened identity decision is established.""" + if not eligible or set(fixed) != set(eligible): + return None + states: dict[str, dict[str, str]] = {} + for incoming, target in fixed.items(): + if target is None: + continue + old = right.loc[right_ids[target], "name"] + new = left.loc[left_ids[incoming], "name"] + if not isinstance(old, str) or not isinstance(new, str) or old != new: + return None + states[target] = {"name": old} + return states + + +def execute_target_state( + query: QueryExpr, + inputs: Mapping[str, Any], + execute: Callable[[QueryExpr, Mapping[str, Any]], Any], + context: LotusExecutionContext, +) -> pd.DataFrame: + """Execute one fused request, retaining relational provenance/temporal logic.""" + import lotus + + joined, body = query.inputs + aggregate = query.params["aggregate"] + semantic = next( + s + for s in aggregate.params["aggregates"] + if isinstance(s, SemanticAggregateSpec) + ) + state_columns = tuple(c.name for c in semantic.output_cols) + left = execute(joined.inputs[0], inputs) + right = execute(joined.inputs[1], inputs) + if not left.index.is_unique or not right.index.is_unique: + raise ValueError("fusion requires unique occurrence indices") + left_ids = {f"l{i}": idx for i, idx in enumerate(left.index)} + right_ids = {f"r{i}": idx for i, idx in enumerate(right.index)} + left_names = {idx: name for name, idx in left_ids.items()} + right_names = {idx: name for name, idx in right_ids.items()} + eligible: dict[str, set[str]] = {} + profile_digest = query.params.get("join_profile_digest", query_digest(joined)) + profile = context.config.semantic_pair_profiles.get(profile_digest) + if profile is not None and profile.mode not in {"oracle-only", "search-filter"}: + raise ValueError("target-state fusion cannot replace proxy-only decisions") + ls, rs = pd.Series(dtype=object), pd.Series(dtype=object) + if not left.empty and not right.empty: + ls, rs, _, _, instruction = join_series( + left, right, str(joined.params["instruction"]) + ) + pairs = semantic_join_pair_candidates( + ls, + rs, + left_frame=left, + right_frame=right, + on=tuple(joined.params.get("on", ())), + query=joined, + ) + candidates, _ = _apply_pair_profile(joined, pairs, context, profile=profile) + for _, row in candidates.iterrows(): + eligible.setdefault(left_names[row[PAIR_LEFT_ID_COLUMN]], set()).add( + right_names[row[PAIR_RIGHT_ID_COLUMN]] + ) + else: + instruction = str(joined.params["instruction"]) + from agent_memory.adapters.lotus.identity_reuse import current_identity_decisions + decisions = current_identity_decisions.get() if query.params.get("identity_reuse") else None + ordered_targets = {name: sorted(targets) for name, targets in eligible.items()} + decision_keys = {name: decisions.key("fused-entity-match-v1", instruction, + str(ls.loc[left_ids[name]]), + [(target, str(rs.loc[right_ids[target]])) for target in targets]) + for name, targets in ordered_targets.items()} if decisions is not None else {} + fixed: dict[str, str | None] = {} + if decisions is not None: + for name, targets in ordered_targets.items(): + selected = decisions.lookup(decision_keys[name], len(targets)) + if selected is not None: + if len(selected) > 1: + raise ValueError("exclusive fusion cached multiple targets") + fixed[name] = targets[selected[0]] if selected else None + resolution: dict[str, Any] = {"matches": {}, "states": {}} + unchanged_states = None + if context.config.reuse_unchanged_entity_name and query.params.get("identity_reuse"): + from agent_memory.memories.zep.fact_summary import IDENTITY_SPEC + from agent_memory.planner.physical import REPRESENTATIVE_VERSION + if query.params.get("version") == REPRESENTATIVE_VERSION and semantic == IDENTITY_SPEC: + unchanged_states = _unchanged_entity_states(eligible, fixed, left, right, left_ids, right_ids) + state_regenerated = False + if eligible and len(fixed) == len(eligible) and all(v is None for v in fixed.values()): + resolution = {"matches": fixed, "states": {}} + elif unchanged_states is not None: + resolution = {"matches": fixed, "states": unchanged_states} + write_trace_event(context.config.trace_dir(), operator="fused_target_state", + event_type="unchanged_entity_name", payload={ + "version": "unchanged-entity-name-v1", + "reused_matches": len(fixed), "unchanged_targets": len(unchanged_states), + "provider_calls": 0, + }, parsed_output=resolution) + elif eligible: + context.configure() + # Fixed decisions need only their chosen state. Unknown assignments + # still need every screened candidate; never prune their alternatives. + prompt_targets: dict[str, set[str]] = {} + for name, targets in eligible.items(): + if name not in fixed: + prompt_targets[name] = targets + else: + target = fixed[name] + prompt_targets[name] = {target} if target is not None else set() + # Match context is exactly the ordinary join's text projection. State + # synthesis receives only the semantic aggregate's declared inputs. + payload = { + "version": query.params["version"], + "join_instruction": instruction, + "consolidation_instruction": semantic.instruction, + "incoming": [ + { + "id": name, + "join_text": str(ls.loc[left_ids[name]]), + "state": { + c: left.loc[left_ids[name], c] + for c in semantic.input_cols or () + }, + "eligible_targets": sorted(targets), + } + for name, targets in prompt_targets.items() + ], + "targets": [ + { + "id": name, + "join_text": str(rs.loc[right_ids[name]]), + "state": { + c: right.loc[right_ids[name], c] + for c in semantic.input_cols or () + }, + } + for name in sorted(set().union(*prompt_targets.values())) + ], + "output_schema": { + "matches": [ + { + "left_id": "one supplied incoming ID", + "right_id": "one ID from that incoming item's eligible_targets, or null", + } + ], + "states": [ + {"right_id": "r0", **{c: "string" for c in state_columns}} + ], + }, + } + if fixed: + payload["fixed_matches"] = fixed + payload["fixed_match_contract"] = ( + "These identity decisions are already established for identical join inputs. " + "Copy them unchanged into matches. Still consolidate CURRENT states for each " + "selected target, including all current incoming contributions." + ) + prompts = [ + [ + {"role": "system", "content": _SYSTEM}, + { + "role": "user", + "content": json.dumps(payload, ensure_ascii=False, default=str), + }, + ] + ] + + def validate(raw: str) -> None: + parsed = parse_resolution(raw, eligible, state_columns=state_columns) + if any(parsed["matches"][name] != target for name, target in fixed.items()): + raise ValueError("fusion changed an established identity decision") + + def build_retry_prompt(prompt: Any, _raw: str, error: str) -> Any: + if not isinstance(prompt, list): + raise TypeError("fusion retry requires a message-list prompt") + return [ + *prompt, + { + "role": "user", + "content": ( + f"The previous response was rejected: {error}. " + "Re-read each incoming item's own eligible_targets list. " + "Return the complete corrected JSON object and nothing else." + ), + }, + ] + + # A malformed fusion assignment is unsafe to apply, but it is also a + # recoverable provider-contract failure. Give it one bounded retry even + # when the generic semantic retry budget is zero; never relax validation. + result = execute_structured_lm_retry_result( + lotus.settings.lm, + prompts, + lm_kwargs={ + "response_format": {"type": "json_object"}, + "max_tokens": context.config.structured_max_tokens, + "progress_bar_desc": "Fused target-state", + }, + output_cols=(), + shape="object", + require_explanation=False, + operator="fused_target_state", + max_retries=max(1, context.config.structured_parse_retries), + output_validator=validate, + retry_prompt_builder=build_retry_prompt, + ) + if not result.invalid_indices: + resolution = parse_resolution(result.raw_outputs[0], eligible, state_columns=state_columns) + state_regenerated = bool(resolution["states"]) + write_trace_event( + context.config.trace_dir(), + operator="fused_target_state", + event_type="fusion_resolution", + payload={ + "version": query.params["version"], + "join_query_digest": query_digest(joined), + "left_occurrences": left_ids, + "right_occurrences": right_ids, + "valid": not result.invalid_indices, + "eligible_pair_count": sum(map(len, eligible.values())), + }, + raw_output=result.raw_output_attempts, + parsed_output=resolution if not result.invalid_indices else None, + ) + if result.invalid_indices: + raise ValueError( + f"invalid target-state output; artifacts: {result.failure_artifact_paths}" + ) + if decisions is not None: + for name, target in resolution["matches"].items(): + candidates_for_name = ordered_targets[name] + decisions.store(decision_keys[name], len(candidates_for_name), + () if target is None else (candidates_for_name.index(target),)) + write_trace_event(context.config.trace_dir(), operator="fused_target_state", + event_type="identity_reuse", payload={ + "reused": len(fixed), "total": len(eligible), + "state_regenerated": state_regenerated, + }) + matches = [ + (left_ids[left_id], right_ids[right_id], None) + for left_id, right_id in resolution["matches"].items() + if right_id is not None + ] + frame = assemble_join_frame( + left, right, matches, how="outer", id_columns=tuple(joined.params["id_columns"]) + ) + bound = {**inputs, JOIN_MAP_BINDING_NAME: frame} + deterministic = replace( + aggregate, + params={ + "aggregates": tuple( + s + for s in aggregate.params["aggregates"] + if not isinstance(s, SemanticAggregateSpec) + ) + }, + ) + merged = execute(deterministic, bound) + for column in state_columns: + merged[column] = [ + resolution["states"][right_names[idx]][column] + for idx in merged[JOIN_MAP_RIGHT_ID_COLUMN] + ] + reference = QueryExpr( + op="materialized_view", + params={"name": _STATE_BINDING, "columns": output_columns(aggregate)}, + ) + return execute( + replace_query(body, aggregate, reference), {**bound, _STATE_BINDING: merged} + ) diff --git a/src/agent_memory/adapters/lotus/identity_reuse.py b/src/agent_memory/adapters/lotus/identity_reuse.py new file mode 100644 index 0000000..4a851ea --- /dev/null +++ b/src/agent_memory/adapters/lotus/identity_reuse.py @@ -0,0 +1,49 @@ +"""Executor-owned identity decisions, scoped to registered entity maintenance.""" + +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from hashlib import sha256 +import json +from typing import Any + +IDENTITY_DECISIONS_INPUT = "__physical_identity_decisions" + + +@dataclass +class IdentityDecisions: + """Store a complete selection mask, never a prior row or occurrence ID.""" + + values: dict[str, bool] + + @staticmethod + def key(*semantic_input: Any) -> str: + return sha256(json.dumps(semantic_input, ensure_ascii=False, sort_keys=True, + default=str).encode()).hexdigest() + + def lookup(self, key: str, size: int) -> tuple[int, ...] | None: + keys = [f"identity-v1:{key}:{size}:{i}" for i in range(size)] + if not keys or any(k not in self.values for k in keys): + return None + return tuple(i for i, k in enumerate(keys) if self.values[k]) + + def store(self, key: str, size: int, selected: Sequence[int]) -> None: + if len(set(selected)) != len(selected) or any(type(i) is not int or not 0 <= i < size for i in selected): + raise ValueError("identity selection must contain distinct candidate positions") + self.values.update({f"identity-v1:{key}:{size}:{i}": i in selected for i in range(size)}) + + +current_identity_decisions: ContextVar[IdentityDecisions | None] = ContextVar( + "entity_identity_decisions", default=None, +) + + +@contextmanager +def identity_scope(values: dict[str, bool]) -> Iterator[None]: + """Keep nested operators on the current executor transaction's decision map.""" + token = current_identity_decisions.set(IdentityDecisions(values)) + try: + yield + finally: + current_identity_decisions.reset(token) diff --git a/src/agent_memory/adapters/lotus/pair_execution.py b/src/agent_memory/adapters/lotus/pair_execution.py index f40ed95..67e25c0 100644 --- a/src/agent_memory/adapters/lotus/pair_execution.py +++ b/src/agent_memory/adapters/lotus/pair_execution.py @@ -9,6 +9,7 @@ import json import math import re +from pathlib import Path from time import perf_counter from typing import Any @@ -354,10 +355,7 @@ def select_semantic_pair_candidates( vectors = embedding_provider.embed(profile.embedding, unique_texts) embedding_latency_ms = (perf_counter() - started) * 1000 vectors_by_text = _validated_vectors(unique_texts, vectors, profile.embedding) - scores = tuple( - _cosine(vectors_by_text[left], vectors_by_text[right]) - for left, right in pair_texts - ) + scores = _memoized_cosines(vectors_by_text, pair_texts) selected = _select_positions( scores, left_ids=left_ids, @@ -398,7 +396,7 @@ def semantic_pair_profiles_fingerprint( def write_semantic_pair_execution_trace( - trace_dir: object, + trace_dir: Path | str | None, *, operator: str, query_digest_value: str, @@ -493,11 +491,18 @@ def _validated_vectors( return result -def _cosine(left: Sequence[float], right: Sequence[float]) -> float: - dot = sum(a * b for a, b in zip(left, right, strict=True)) - left_norm = math.sqrt(sum(value * value for value in left)) - right_norm = math.sqrt(sum(value * value for value in right)) - return dot / (left_norm * right_norm) +def _memoized_cosines( + vectors: Mapping[str, Sequence[float]], pairs: Sequence[tuple[str, str]], +) -> tuple[float, ...]: + """Reuse work within one selection, preserving the original scalar arithmetic.""" + norms = {text: math.sqrt(sum(value * value for value in vector)) + for text, vector in vectors.items()} + scores: dict[tuple[str, str], float] = {} + for left, right in pairs: + if (left, right) not in scores: + dot = sum(a * b for a, b in zip(vectors[left], vectors[right], strict=True)) + scores[left, right] = dot / (norms[left] * norms[right]) + return tuple(scores[pair] for pair in pairs) def _select_positions( diff --git a/src/agent_memory/adapters/lotus/predicate_reuse.py b/src/agent_memory/adapters/lotus/predicate_reuse.py new file mode 100644 index 0000000..e96bb52 --- /dev/null +++ b/src/agent_memory/adapters/lotus/predicate_reuse.py @@ -0,0 +1,43 @@ +"""Reuse fixed text predicates independently of non-semantic row metadata.""" + +from collections.abc import Callable +import json + +import pandas as pd + + +def reuse_predicate_decisions( + source: pd.DataFrame, instruction: str, decisions: dict[str, bool], + evaluate: Callable[[pd.DataFrame, list[int]], pd.DataFrame], +) -> tuple[pd.DataFrame, set[int]]: + """Screening precedes this function; results always carry current row values.""" + from lotus.nl_expression import parse_cols + from lotus.templates.task_instructions import df2multimodal_info + + columns = tuple(parse_cols(instruction)) + if not columns or any(c not in source for c in columns): + return evaluate(source, list(range(len(source)))), set() + # Do not guess dependencies for multimodal or nested input values. + if any(not isinstance(v, str) and v is not None and v is not pd.NA + and not isinstance(v, (int, float, bool)) + for row in source.loc[:, list(columns)].itertuples(index=False, name=None) for v in row): + return evaluate(source, list(range(len(source)))), set() + docs = df2multimodal_info(source, list(columns)) + if any(doc.get("image") for doc in docs): + return evaluate(source, list(range(len(source)))), set() + keys = [json.dumps([instruction, columns, doc.get("text", "")], ensure_ascii=False) for doc in docs] + missing: dict[str, int] = {} + reused = {i for i, key in enumerate(keys) if key in decisions} + for i, key in enumerate(keys): + if key not in decisions: + missing.setdefault(key, i) + positions = list(missing.values()) + if positions: + pending = source.iloc[positions].copy() + pending.index = pd.Index(positions) + result = evaluate(pending, positions) + if not result.index.is_unique or not set(result.index) <= set(positions): + raise ValueError("predicate output must preserve unique input positions") + selected = set(result.index) + decisions.update({key: position in selected for key, position in missing.items()}) + return source.iloc[[i for i, key in enumerate(keys) if decisions[key]]].copy(), reused diff --git a/src/agent_memory/adapters/lotus/prompt_batching.py b/src/agent_memory/adapters/lotus/prompt_batching.py index 61451fc..37ad8e9 100644 --- a/src/agent_memory/adapters/lotus/prompt_batching.py +++ b/src/agent_memory/adapters/lotus/prompt_batching.py @@ -142,6 +142,7 @@ def run_prompt_batches( operator: str, trace_dir: Any = None, model_kwargs: Mapping[str, Any] | None = None, + task_partitions: Sequence[Sequence[int]] | None = None, ) -> PromptBatchRunResult[ResultT]: """Run ready semantic tasks in deterministic, context-bounded prompts.""" @@ -177,17 +178,18 @@ def run_prompt_batches( return result try: - requests = _build_requests( - task_sequence, - task_ids=task_ids, - build_request=build_request, - model=model, - config=config, - output_schema=( - output_schema - if structured_output_transport == "responses-json-schema" - else None - ), + partitions = (tuple(range(len(task_sequence))),) if task_partitions is None else task_partitions + if sorted(position for group in partitions for position in group) != list(range(len(task_sequence))): + raise ValueError("task partitions must cover each task exactly once") + requests = tuple( + request + for group in partitions + for request in _build_requests( + tuple(task_sequence[i] for i in group), + task_ids=tuple(task_ids[i] for i in group), + build_request=build_request, model=model, config=config, + output_schema=output_schema if structured_output_transport == "responses-json-schema" else None, + ) ) except ValueError as error: _write_batch_failure( diff --git a/src/agent_memory/adapters/lotus/relational.py b/src/agent_memory/adapters/lotus/relational.py index 56db059..a8b8a05 100644 --- a/src/agent_memory/adapters/lotus/relational.py +++ b/src/agent_memory/adapters/lotus/relational.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence +from dataclasses import replace import json from typing import Any @@ -10,6 +11,7 @@ import pandas as pd from agent_memory.policy.aggregates import ( + ArgMinAggregateSpec, ArrayAggregateSpec, CollectListAggregateSpec, MinAggregateSpec, @@ -462,9 +464,25 @@ def execute_agg( raise ValueError("agg expects group_by or sem_groupby input") if source.empty: return pd.DataFrame(columns=list(output_columns(query))) + if query.params.get("fact_summary") or query.params.get("identity_singleton"): + from agent_memory.memories.zep.fact_summary import IDENTITY_SPEC + from agent_memory.planner.physical import summary_specs + semantic_specs = tuple(s for s in aggregates if isinstance(s, SemanticAggregateSpec)) + if query.params.get("fact_summary") and semantic_specs not in tuple((s,) for s in summary_specs()): + raise ValueError("fact summary shortcut requires the registered summary contract") + if query.params.get("identity_singleton") and semantic_specs != (IDENTITY_SPEC,): + raise ValueError("identity singleton requires the registered entity contract") + return _execute_agg_with_group_batching(query, source, aggregates, context=context) + if query.params.get("singleton_identity"): + from agent_memory.memories.zep.policy import ZepMemory + template = ZepMemory._deduplicated_facts.expr.inputs[0] + if aggregates != tuple(template.params["aggregates"]) or group_query.params != template.inputs[0].params: + raise ValueError("singleton identity is restricted to registered new Zep facts") + return _execute_agg_with_group_batching(query, source, aggregates, context=context) if ( context.config.sem_agg_dispatch != "sequential" or context.config.prompt_batching is not None + or context.config.sem_agg_prompt_batching is not None ): return _execute_agg_with_group_batching( query, @@ -479,6 +497,9 @@ def execute_agg( ): row: dict[str, Any] = dict(key_values) for aggregate in aggregates: + if isinstance(aggregate, ArgMinAggregateSpec): + _apply_deterministic_aggregate(aggregate, group, row) + continue if isinstance(aggregate, ArrayAggregateSpec): if aggregate.output_col in row: continue @@ -542,11 +563,32 @@ def _execute_agg_with_group_batching( rows = [dict(key_values) for key_values, _group in grouped] for aggregate in aggregates: if isinstance(aggregate, SemanticAggregateSpec): - semantic_values = _execute_grouped_semantic_aggregate_spec_many( + passthrough = { + i: {c.name: group.iloc[0][c.name] for c in aggregate.output_cols} + for i, group in enumerate(groups) + if (query.params.get("singleton_identity") or query.params.get("identity_singleton")) and len(group) == 1 + and all(isinstance(group.iloc[0][c.name], str) for c in aggregate.output_cols) + } + pending = [i for i in range(len(groups)) if i not in passthrough] + if query.params.get("fact_summary"): + for row, values in zip(rows, _fact_summary_values(groups, context=context), strict=True): + row.update(values) + continue + if query.params.get("singleton_identity"): + from agent_memory.tracing.semantic import write_trace_event + write_trace_event(context.config.trace_dir(), operator="agg", + event_type="singleton_fact_identity", payload={ + "passthrough_groups": len(passthrough), + "provider_groups": len(pending), + "approximate": True, + }) + generated = _execute_grouped_semantic_aggregate_spec_many( aggregate, - groups, + [groups[i] for i in pending], context=context, ) + values_by_group = {**passthrough, **dict(zip(pending, generated, strict=True))} + semantic_values = [values_by_group[i] for i in range(len(groups))] for row, values in zip(rows, semantic_values, strict=True): for column, value in values.items(): if column not in row: @@ -565,6 +607,77 @@ def _execute_agg_with_group_batching( return result +def _truncate_summary(text: str, max_chars: int = 1000) -> str: + """Apply Graphiti's sentence-boundary length policy (Apache-2.0).""" + import re + if len(text) <= max_chars: + return text + prefix = text[:max_chars] + boundaries = list(re.finditer(r"[.!?](?:\s|$)", prefix)) + return (text[:boundaries[-1].end()] if boundaries else prefix).rstrip() + + +def _fact_summary_values(groups: Sequence[pd.DataFrame], *, context: Any) -> list[dict[str, str]]: + """Append short evidence verbatim; compress only over-threshold groups.""" + from agent_memory.memories.zep.fact_summary import SUMMARY_SPEC + from agent_memory.tracing.semantic import write_trace_event + direct: dict[int, dict[str, str]] = {} + for i, group in enumerate(groups): + texts = group["summary"].tolist() + if any(not isinstance(text, str) for text in texts): + raise TypeError("fact summaries require string evidence") + joined = "\n".join(texts) + if len(joined) <= 2000: + direct[i] = {"summary": joined} + pending = [i for i in range(len(groups)) if i not in direct] + compression_spec = replace(SUMMARY_SPEC, instruction=( + "Compress the supplied {summary} evidence into one concise factual summary of at most " + "1000 characters including spaces and punctuation. Preserve supported names, relationships, " + "dates and qualifiers. Remove repetition first. Do not invent information." + )) + generated = _execute_grouped_semantic_aggregate_spec_many( + compression_spec, [groups[i] for i in pending], context=context) if pending else [] + if len(generated) != len(pending): + raise ValueError("summary compression returned an incorrect number of groups") + if any(not isinstance(value.get("summary"), str) for value in generated): + raise TypeError("summary compression must return a string") + for i, value in enumerate(generated): + raw = value["summary"] + if not raw.strip(): + raise ValueError("summary compression returned blank text") + clipped = _truncate_summary(raw) + if not clipped.strip(): + raise ValueError("summary clipping returned blank text") + if clipped != raw: + write_trace_event(context.config.trace_dir(), operator="agg", event_type="fact_summary_truncation", + payload={"version": "summary-sentence-boundary-v1", "group_index": pending[i], + "raw_summary": raw, "summary": clipped, + "before_chars": len(raw), "after_chars": len(clipped)}) + generated[i] = {"summary": clipped} + values = {**direct, **dict(zip(pending, generated, strict=True))} + write_trace_event(context.config.trace_dir(), operator="agg", event_type="fact_summary_threshold", + payload={"direct_groups": len(direct), "compression_groups": len(pending), + "append_limit_chars": 2000, "compression_limit_chars": 1000}) + return [values[i] for i in range(len(groups))] + + +def execute_fact_summary_map(query: QueryExpr, inputs: Mapping[str, Any], execute: Callable, context: Any) -> pd.DataFrame: + """Merge old summaries and new evidence without regenerating short text.""" + from agent_memory.memories.zep.fact_summary import SUMMARY_SPEC + from agent_memory.planner.rules import DifferentialInstructionRewriter + expected = DifferentialInstructionRewriter().agg_to_map( + SUMMARY_SPEC.instruction, input_cols=("summary",), output_cols=SUMMARY_SPEC.output_cols) + if query.params.get("instruction") != expected or query.params.get("output_cols") != SUMMARY_SPEC.output_cols: + raise ValueError("summary map requires the registered summary contract") + frame = execute(query.inputs[0], inputs).copy() + groups = [] + for _, row in frame.iterrows(): + texts = [row[c] for c in ("summary:right", "summary:left") if pd.notna(row[c])] + groups.append(pd.DataFrame({"summary": texts})) + frame["summary"] = [v["summary"] for v in _fact_summary_values(groups, context=context)] + return frame + + def _apply_deterministic_aggregate( aggregate: object, group: pd.DataFrame, @@ -572,6 +685,21 @@ def _apply_deterministic_aggregate( ) -> None: """Apply one non-semantic aggregate spec to one grouped frame.""" + if isinstance(aggregate, ArgMinAggregateSpec): + columns = (*aggregate.order_by, *aggregate.columns) + records = group.loc[:, list(columns)].to_dict("records") + if not records: + raise ValueError("arg_min requires a nonempty group") + if group.loc[:, list(aggregate.order_by)].isna().any().any(): + raise ValueError("arg_min ordering keys cannot be null") + def key(record: dict[str, Any]) -> tuple[Any, ...]: + return tuple(record[c] for c in aggregate.order_by) + selected = min(records, key=key) + tied = group.loc[[key(record) == key(selected) for record in records], list(aggregate.columns)] + if len(tied.drop_duplicates()) != 1: + raise ValueError("arg_min tied ordering keys have conflicting payloads") + row.update(selected) + return if isinstance(aggregate, ArrayAggregateSpec): if aggregate.output_col not in row: row[aggregate.output_col] = _array_records_json(group, aggregate.columns) diff --git a/src/agent_memory/adapters/lotus/scoped_lm.py b/src/agent_memory/adapters/lotus/scoped_lm.py new file mode 100644 index 0000000..6f5bef9 --- /dev/null +++ b/src/agent_memory/adapters/lotus/scoped_lm.py @@ -0,0 +1,28 @@ +"""Context-local model routing for LOTUS's global operator-cache lookup.""" + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Any + + +class ScopedLM: + """Keep LOTUS settings stable while independent lanes own separate LM state.""" + + def __init__(self, model: Any) -> None: + self._current: ContextVar[Any] = ContextVar("lotus_execution_lm", default=model) + + def __getattr__(self, name: str) -> Any: + return getattr(self._current.get(), name) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + return self._current.get()(*args, **kwargs) + + @contextmanager + def scope(self, model: Any) -> Iterator[None]: + """Route model calls and native cache/stat accesses within this context.""" + token = self._current.set(model) + try: + yield + finally: + self._current.reset(token) diff --git a/src/agent_memory/adapters/lotus/sem_agg.py b/src/agent_memory/adapters/lotus/sem_agg.py index e2af101..fad35f5 100644 --- a/src/agent_memory/adapters/lotus/sem_agg.py +++ b/src/agent_memory/adapters/lotus/sem_agg.py @@ -199,7 +199,7 @@ def execute_native_sem_agg( def _prompt_batching(config: LotusExecutionConfig) -> PromptBatching | None: """Return the shared prompt batching contract.""" - return config.prompt_batching + return config.sem_agg_prompt_batching or config.prompt_batching def execute_native_sem_agg_groups( @@ -663,15 +663,15 @@ def write_sem_agg_audit( } if ( config.sem_agg_dispatch != "sequential" - or config.prompt_batching is not None + or _prompt_batching(config) is not None ): audit_row.update( { "dispatch": config.sem_agg_dispatch, "prompt_batching": ( None - if config.prompt_batching is None - else config.prompt_batching.to_dict() + if (batching := _prompt_batching(config)) is None + else batching.to_dict() ), "structured_output_repaired": syntax_repair_method is not None, "structured_output_repair_method": syntax_repair_method or "", diff --git a/src/agent_memory/adapters/lotus/sem_filter.py b/src/agent_memory/adapters/lotus/sem_filter.py index cb4c5d7..0a045f4 100644 --- a/src/agent_memory/adapters/lotus/sem_filter.py +++ b/src/agent_memory/adapters/lotus/sem_filter.py @@ -6,6 +6,7 @@ from collections.abc import Callable, Mapping, Sequence import json import re +from pathlib import Path from typing import Any import pandas as pd @@ -20,7 +21,9 @@ PairCandidateSelection, select_semantic_pair_candidates, write_semantic_pair_execution_trace, + semantic_pair_site_id, ) +from agent_memory.adapters.lotus.site_batching import execute_site_batching from agent_memory.adapters.lotus.sem_filter_batch_prompting import ( execute_batch_prompted_sem_filter, validate_batch_prompting_sem_filter_config, @@ -33,6 +36,8 @@ ) from agent_memory.adapters.lotus.structured import examples_dataframe, normalize_strategy from agent_memory.policy.logical import QueryExpr +from agent_memory.planner.physical import PREDICATE_DECISIONS_INPUT +from agent_memory.adapters.lotus.predicate_reuse import reuse_predicate_decisions QUALIFIED_PLACEHOLDER_PATTERN = re.compile( r"(? pd.DataFrame: + nonlocal prompt_count, prompt_retry_count, verified_tuple_count + verified_tuple_count = len(frame) + if frame.empty or profile is not None and profile.mode == "proxy-only": + return frame.copy() + if single_schema_predicate(context.config): + decision = execute_schema_predicate( + frame, instruction=instruction, config=context.config, operator="sem_filter", + ) + return frame.iloc[[i for i, keep in enumerate(decision.decisions) if keep]].copy() + if site_batching is not None: + prompted = execute_site_batching( + frame, identities=oracle_source.iloc[positions], instruction=instruction, + context=context, config=site_batching, + ) + elif prompt_batching is not None: + prompted = execute_batch_prompted_sem_filter( + frame, instruction=instruction, context=context, prompt_batching=prompt_batching, + ) + else: + return frame.sem_filter(instruction, **native_sem_filter_kwargs(context.config)) prompt_count = prompted.prompt_count prompt_retry_count = prompted.retry_count - else: - result = lotus_source.sem_filter( - instruction, - **native_sem_filter_kwargs(context.config), + return prompted.frame + + reused_positions: set[int] = set() + decisions = inputs.get(PREDICATE_DECISIONS_INPUT) + if decisions is not None and (profile is None or profile.mode != "proxy-only"): + if not isinstance(decisions, dict): + raise TypeError("predicate decisions must be executor-owned state") + verified_tuple_count = 0 + result, reused_positions = reuse_predicate_decisions( + lotus_source, instruction, decisions, evaluate, ) + elif not batch_prompting and not single_schema_predicate(context.config) and selection is None: + result = lotus_source.sem_filter(instruction, **native_sem_filter_kwargs(context.config)) + else: + result = evaluate(lotus_source, list(range(len(lotus_source)))) if restore_columns: result = result.rename(columns=restore_columns) with semantic_trace_scope( @@ -137,6 +158,7 @@ def execute_sem_filter( result=result, instruction=str(query.params["instruction"]), profile=profile, + reused_positions=reused_positions, ) trace_payload: dict[str, Any] = { "instruction": str(query.params["instruction"]), @@ -161,6 +183,13 @@ def execute_sem_filter( "verified_tuple_count": verified_tuple_count, } ) + if site_batching is not None: + trace_payload["pair_filter_batching"] = site_batching.to_dict() + trace_payload["semantic_pair_site_id"] = semantic_pair_site_id(query) + if decisions is not None: + trace_payload["predicate_reused_tuple_count"] = len(reused_positions) + trace_payload["predicate_evaluated_tuple_count"] = verified_tuple_count + trace_payload["predicate_coalesced_tuple_count"] = len(lotus_source) - len(reused_positions) - int(verified_tuple_count or 0) write_compact_operator_trace( context.config.trace_dir(), operator="sem_filter", @@ -173,12 +202,13 @@ def execute_sem_filter( def write_pairwise_sem_filter_trace( - trace_dir: object, + trace_dir: Path | str | None, *, source: pd.DataFrame, result: pd.DataFrame, instruction: str, profile: Any, + reused_positions: set[int] | None = None, ) -> None: """Record pair-shaped filter decisions without full relation snapshots.""" @@ -193,7 +223,7 @@ def write_pairwise_sem_filter_trace( "proxy" if profile is not None and profile.mode == "proxy-only" else "oracle" ) rows: list[dict[str, Any]] = [] - for _, row in source.iterrows(): + for position, (_, row) in enumerate(source.iterrows()): signature = _row_signature(row) matched = result_counts[signature] > 0 if matched: @@ -202,7 +232,7 @@ def write_pairwise_sem_filter_trace( { "instruction": instruction, "direction": direction, - "decision_source": decision_source, + "decision_source": "reused" if position in (reused_positions or ()) else decision_source, "left_id": _endpoint_id(row, left_id_columns), "right_id": _endpoint_id(row, right_id_columns), "left": _endpoint_text(row, left_text_columns), diff --git a/src/agent_memory/adapters/lotus/sem_filter_batch_prompting.py b/src/agent_memory/adapters/lotus/sem_filter_batch_prompting.py index b74a0eb..7e60ac6 100644 --- a/src/agent_memory/adapters/lotus/sem_filter_batch_prompting.py +++ b/src/agent_memory/adapters/lotus/sem_filter_batch_prompting.py @@ -53,6 +53,7 @@ class _BatchPromptingTask: position: int row_id: str context: str + shared_context: str = "" @dataclass @@ -73,10 +74,13 @@ def __call__( trace_dir: Any = None, operator: str = "sem_filter", structured_output_transport: str = "chat-json-object", + task_partitions: Sequence[Sequence[int]] | None = None, + shared_columns: tuple[str, ...] = (), ) -> BatchPromptingResult: """Evaluate independent tuples in bounded structured prompts.""" import lotus + from lotus.nl_expression import parse_cols from lotus.templates import task_instructions lm = lotus.settings.lm @@ -84,7 +88,7 @@ def __call__( raise ValueError( "batch-prompting sem_filter requires a configured language model" ) - columns = tuple(lotus.nl_expression.parse_cols(instruction)) + columns = tuple(parse_cols(instruction)) if not columns: raise ValueError( "batch-prompting sem_filter requires instruction column placeholders" @@ -101,6 +105,15 @@ def __call__( ) claim = _render_claim(instruction, columns) tasks = _build_tasks(docs) + if shared_columns: + if not set(shared_columns) <= set(columns): + raise ValueError("shared context columns must be predicate dependencies") + remaining = [c for c in columns if c not in shared_columns] + contexts = task_instructions.df2multimodal_info(self._obj, remaining) + shared = task_instructions.df2multimodal_info(self._obj, list(shared_columns)) + tasks = tuple(_BatchPromptingTask(i, f"row_{i}", str(doc.get("text", "")), + str(common.get("text", ""))) + for i, (doc, common) in enumerate(zip(contexts, shared, strict=True))) if not tasks: return BatchPromptingResult( frame=self._obj.copy(), @@ -131,6 +144,7 @@ def __call__( model_kwargs={ "show_progress_bar": True, }, + task_partitions=task_partitions, ) selected_positions = [ task.position @@ -154,6 +168,8 @@ def execute_batch_prompted_sem_filter( instruction: str, context: LotusExecutionContext, prompt_batching: PromptBatching, + task_partitions: Sequence[Sequence[int]] | None = None, + shared_columns: tuple[str, ...] = (), ) -> BatchPromptingResult: """Execute one explicitly configured batch-prompting semantic filter.""" @@ -166,6 +182,8 @@ def execute_batch_prompted_sem_filter( progress_bar_desc=context.config.sem_filter_progress_bar_desc, trace_dir=context.config.trace_dir(), operator="sem_filter", + task_partitions=task_partitions, + shared_columns=shared_columns, structured_output_transport=context.config.structured_output_transport, ) @@ -181,6 +199,8 @@ def execute_batch_prompted_predicate( trace_dir: Any = None, operator: str = "sem_filter", structured_output_transport: str = "chat-json-object", + task_partitions: Sequence[Sequence[int]] | None = None, + shared_columns: tuple[str, ...] = (), ) -> BatchPromptingResult: """Evaluate independent boolean predicate tasks in shared prompts.""" @@ -193,6 +213,8 @@ def execute_batch_prompted_predicate( trace_dir=trace_dir, operator=operator, structured_output_transport=structured_output_transport, + task_partitions=task_partitions, + shared_columns=shared_columns, ) @@ -256,6 +278,22 @@ def _build_request( ], "output_schema": {"decisions": [{"row_id": "row_id", "keep": True}]}, } + if any(task.shared_context for task in tasks): + shared = {task.shared_context for task in tasks} + if len(shared) == 1: + payload["shared_context"] = tasks[0].shared_context + payload["context_contract"] = "Each row's context consists of shared_context plus that row's context." + else: + contexts = {text: f"context_{i}" for i, text in enumerate(dict.fromkeys(task.shared_context for task in tasks))} + payload["contexts"] = {key: text for text, key in contexts.items()} + payload["rows"] = [ + {"row_id": task.row_id, "context_id": contexts[task.shared_context], "context": task.context} + for task in tasks + ] + payload["context_contract"] = ( + "Evaluate each row using ONLY contexts[row.context_id] plus that row's context. " + "Other contexts and rows are unrelated tasks, not evidence for this row." + ) return PromptBatchRequest( task_ids=tuple(task.row_id for task in tasks), prompt=[ diff --git a/src/agent_memory/adapters/lotus/sem_flat_map.py b/src/agent_memory/adapters/lotus/sem_flat_map.py index 614a84d..d28cb86 100644 --- a/src/agent_memory/adapters/lotus/sem_flat_map.py +++ b/src/agent_memory/adapters/lotus/sem_flat_map.py @@ -22,14 +22,16 @@ def execute_sem_flat_map( inputs: Mapping[str, Any], execute: Callable[[QueryExpr, Mapping[str, Any]], Any], context: LotusExecutionContext, + *, model: Any = None, ) -> Any: """Execute sem_flat_map as structured JSON rows-wrapper generation.""" - context.configure() + if model is None: + context.configure() source = execute(query.inputs[0], inputs) cols = output_columns(query, operator="sem_flat_map") input_cols = resolve_input_cols(source, query, operator="sem_flat_map") - executor = StructuredLMExecutor(source) + executor = StructuredLMExecutor(source, model=model) generation = executor( input_cols=input_cols, output_cols=cols, @@ -104,6 +106,6 @@ def apply_flat_map_outputs( columns.append(column.name) if ordinal_col is not None: columns.append(ordinal_col) - result = pd.DataFrame(rows, columns=columns) + result = pd.DataFrame(rows, columns=pd.Index(columns)) result.index = pd.Index(row_indexes, dtype="object") return result diff --git a/src/agent_memory/adapters/lotus/sem_groupby.py b/src/agent_memory/adapters/lotus/sem_groupby.py index ca5b3c6..1c2938b 100644 --- a/src/agent_memory/adapters/lotus/sem_groupby.py +++ b/src/agent_memory/adapters/lotus/sem_groupby.py @@ -8,13 +8,14 @@ import pandas as pd -from agent_memory.adapters.lotus.context import LotusExecutionContext +from agent_memory.adapters.lotus.context import LotusExecutionConfig, LotusExecutionContext from agent_memory.adapters.lotus.pair_execution import ( PAIR_LEFT_ID_COLUMN, PAIR_LEFT_TEXT_COLUMN, PAIR_RIGHT_ID_COLUMN, PAIR_RIGHT_TEXT_COLUMN, SemanticPairExecutionProfile, + semantic_pair_site_id, select_semantic_pair_candidates, write_semantic_pair_execution_trace, ) @@ -35,6 +36,14 @@ ) +def validate_site_groupby_batching(query: QueryExpr, config: LotusExecutionConfig) -> None: + """Reject sites that do not execute independent pairwise group decisions.""" + profile = config.semantic_pair_profiles.get(query_digest(query)) + if (query.params.get("labels") or query.params.get("membership") not in {None, "exclusive"} + or (profile is not None and profile.mode == "proxy-only")): + raise ValueError("site groupby batching requires open-ended exclusive oracle grouping") + + def execute_sem_groupby( query: QueryExpr, inputs: Mapping[str, Any], @@ -54,6 +63,14 @@ def execute_sem_groupby( partition_by = tuple(str(column) for column in query.params.get("partition_by", ())) labels = tuple(query.params.get("labels") or ()) digest = query_digest(query) + prompt_batching = context.config.prompt_batching + pair_batch_size = context.config.sem_groupby_pair_batch_size + if context.config.groupby_prompt_batching: + setting = context.config.groupby_prompt_batching.get(semantic_pair_site_id(query)) + if setting is not None: + validate_site_groupby_batching(query, context.config) + prompt_batching = setting + pair_batch_size = None profile = context.config.semantic_pair_profiles.get(digest) if labels and profile is not None: raise ValueError( @@ -80,13 +97,13 @@ def execute_sem_groupby( label_col=str(query.params.get("label_col", "_label")), instruction=str(query.params["instruction"]), default=context.config.sem_groupby_default, - pair_batch_size=context.config.sem_groupby_pair_batch_size, + pair_batch_size=pair_batch_size, pair_batch_retries=context.config.sem_groupby_pair_batch_retries, trace_dir=context.config.trace_dir(), query_digest_value=digest, profile=profile, embedding_provider=context.pair_embedding_provider, - prompt_batching=context.config.prompt_batching, + prompt_batching=prompt_batching, structured_output_transport=context.config.structured_output_transport, structured_parse_retries=context.config.structured_parse_retries, structured_max_tokens=context.config.structured_max_tokens, @@ -98,7 +115,7 @@ def execute_sem_groupby( labels=labels, label_col=str(query.params.get("label_col", "_label")), instruction=str(query.params["instruction"]), - prompt_batching=context.config.prompt_batching, + prompt_batching=prompt_batching, structured_output_transport=context.config.structured_output_transport, ) @@ -108,13 +125,13 @@ def execute_sem_groupby( input_cols=input_cols, instruction=str(query.params["instruction"]), default=context.config.sem_groupby_default, - pair_batch_size=context.config.sem_groupby_pair_batch_size, + pair_batch_size=pair_batch_size, pair_batch_retries=context.config.sem_groupby_pair_batch_retries, trace_dir=context.config.trace_dir(), query_digest_value=digest, profile=profile, embedding_provider=context.pair_embedding_provider, - prompt_batching=context.config.prompt_batching, + prompt_batching=prompt_batching, structured_output_transport=context.config.structured_output_transport, structured_parse_retries=context.config.structured_parse_retries, structured_max_tokens=context.config.structured_max_tokens, @@ -471,6 +488,25 @@ def evaluate_group_matches( "{left} and {right} satisfy this semantic grouping condition: " f"{lowered_instruction}" ) + from agent_memory.adapters.lotus.identity_reuse import current_identity_decisions + decisions = current_identity_decisions.get() + original_pairs = pairs + keys = [decisions.key("group", user_instruction, default, row["left"], row["right"]) + for _, row in pairs.iterrows()] if decisions is not None else [] + known = {i: decisions.lookup(key, 1) for i, key in enumerate(keys)} if decisions is not None else {} + missing = [i for i in range(len(pairs)) if known.get(i) is None] + if decisions is not None: + from agent_memory.tracing.semantic import write_trace_event + write_trace_event(trace_dir, operator="sem_groupby", event_type="identity_reuse", + payload={"reused": len(pairs) - len(missing), "total": len(pairs)}) + pairs = pairs.iloc[missing].reset_index(drop=True) + if pairs.empty: + write_groupby_pair_trace(trace_dir, original_pairs, source_instruction=instruction, + instruction=user_instruction, outputs=[bool(known[i]) for i in range(len(original_pairs))], + raw_outputs=[""] * len(original_pairs), + explanations=["reused identity decision"] * len(original_pairs), default=default) + return [(int(row["_left_unique_id"]), int(row["_right_unique_id"])) + for i, (_, row) in enumerate(original_pairs.iterrows()) if known[i]] if prompt_batching is not None: if default: raise ValueError( @@ -528,6 +564,16 @@ def evaluate_group_matches( aligned_batch_values(output, "explanations", len(pair_batch), "") ) + if decisions is not None: + for position, keep in zip(missing, parsed_outputs, strict=True): + decisions.store(keys[position], 1, (0,) if keep else ()) + generated = dict(zip(missing, zip(parsed_outputs, raw_outputs, explanations, strict=True), strict=True)) + aligned = [generated[i] if i in generated else (bool(known[i]), "", "reused identity decision") + for i in range(len(original_pairs))] + parsed_outputs = [v[0] for v in aligned] + raw_outputs = [v[1] for v in aligned] + explanations = [v[2] for v in aligned] + pairs = original_pairs write_groupby_pair_trace( trace_dir, pairs, diff --git a/src/agent_memory/adapters/lotus/sem_topk_join.py b/src/agent_memory/adapters/lotus/sem_topk_join.py index be6fc7a..339559d 100644 --- a/src/agent_memory/adapters/lotus/sem_topk_join.py +++ b/src/agent_memory/adapters/lotus/sem_topk_join.py @@ -12,6 +12,7 @@ from agent_memory.adapters.lotus.context import LotusExecutionContext from agent_memory.adapters.lotus.prompt_batching import ( ParsedPromptBatch, + PromptBatching, PromptBatchItem, PromptBatchRequest, run_prompt_batches, @@ -22,6 +23,7 @@ PAIR_RIGHT_ID_COLUMN, PAIR_RIGHT_TEXT_COLUMN, SemanticPairExecutionProfile, + semantic_pair_site_id, select_semantic_pair_candidates, write_semantic_pair_execution_trace, ) @@ -94,6 +96,15 @@ def evaluate_sem_topk_join( ) method = context.config.sem_join_topk_method + site_batching = None + if context.config.listwise_join_batching: + try: + site = semantic_pair_site_id(query) + except ValueError: + site = "" + site_batching = context.config.listwise_join_batching.get(site) + if site_batching is not None and profile is not None and profile.mode == "proxy-only": + raise ValueError("site listwise batching cannot replace proxy-only decisions") if profile is not None and profile.mode == "proxy-only": join_results = _proxy_topk(candidates, scores=scores, k=k) method = "proxy-only" @@ -106,6 +117,7 @@ def evaluate_sem_topk_join( right_label=right_label, k=k, context=context, + prompt_batching=site_batching, ) else: join_results = _pairwise_topk( @@ -127,6 +139,7 @@ def evaluate_sem_topk_join( "oracle_candidate_count": len(candidates), "selected_pair_count": len(join_results), "structured_retry_count": retry_count, + **({"listwise_join_batching": site_batching.to_dict()} if site_batching is not None else {}), } @@ -172,6 +185,7 @@ def _listwise_topk( right_label: str, k: int, context: LotusExecutionContext, + prompt_batching: PromptBatching | None = None, ) -> tuple[list[tuple[Any, Any, str | None]], int]: """Resolve zero-to-k right matches for each left row in listwise batches.""" @@ -190,19 +204,29 @@ def _listwise_topk( right_label=right_label, k=k, ) - if context.config.prompt_batching is not None: + from agent_memory.adapters.lotus.identity_reuse import current_identity_decisions + decisions = current_identity_decisions.get() + all_tasks = tasks + keys = {task.task_id: decisions.key("listwise-v1", task.prompt, k) for task in tasks} if decisions else {} + cached = {task.left_id: tuple(list(task.candidate_positions.values())[i] for i in selected) + for task in tasks + if decisions is not None + and (selected := decisions.lookup(keys[task.task_id], len(task.candidate_positions))) is not None} + tasks = [task for task in tasks if task.left_id not in cached] + batching = prompt_batching or context.config.prompt_batching + if batching is not None: task_by_id = {task.task_id: task for task in tasks} execution = run_prompt_batches( tasks, task_id=lambda task: task.task_id, - build_request=_build_listwise_batch_request, + build_request=_build_compact_listwise_request if prompt_batching is not None else _build_listwise_batch_request, parse_results=lambda raw_output: _parse_listwise_prompt_batch( raw_output, tasks=task_by_id, k=k, ), model=lm, - config=context.config.prompt_batching, + config=batching, output_schema=_listwise_batch_schema(), structured_output_transport=context.config.structured_output_transport, max_retries=context.config.structured_parse_retries, @@ -226,6 +250,13 @@ def _listwise_topk( max_retries=context.config.structured_parse_retries, ) + if decisions is not None: + for task in tasks: + positions = list(task.candidate_positions.values()) + decisions.store(keys[task.task_id], len(positions), + [positions.index(p) for p in selected_by_left[task.left_id]]) + selected_by_left.update(cached) + tasks = all_tasks selected_positions = { position for positions in selected_by_left.values() @@ -239,7 +270,7 @@ def _listwise_topk( { "operator": "sem_join", "instruction": instruction, - "decision_source": "listwise", + "decision_source": "identity-reuse" if row[PAIR_LEFT_ID_COLUMN] in cached else "listwise", "left_id": row[PAIR_LEFT_ID_COLUMN], "right_id": row[PAIR_RIGHT_ID_COLUMN], "left": row[PAIR_LEFT_TEXT_COLUMN], @@ -388,6 +419,27 @@ def _listwise_batch_schema() -> dict[str, Any]: } +def _build_compact_listwise_request(tasks: tuple[_ListwiseTask, ...]) -> PromptBatchRequest: + """Share only the instruction; each anchor retains its own candidate domain.""" + payloads = [json.loads(task.prompt[1]["content"]) for task in tasks] + instruction = payloads[0]["join_condition"] + if any(payload["join_condition"] != instruction for payload in payloads): + raise ValueError("listwise tasks must share the same instruction") + payload = { + "join_condition": instruction, + "tasks": [{"task_id": task.task_id, **{key: item[key] for key in + ("max_matches", "left", "right_candidates")}} + for task, item in zip(tasks, payloads, strict=True)], + "output_schema": {"results": [{"task_id": "task_id", "selected_ids": ["candidate_id"]}]}, + } + return PromptBatchRequest( + task_ids=tuple(task.task_id for task in tasks), + prompt=[{"role": "system", "content": LISTWISE_JOIN_SYSTEM_PROMPT + " " + LISTWISE_JOIN_BATCH_SYSTEM_PROMPT}, + {"role": "user", "content": json.dumps(payload, ensure_ascii=False, default=str)}], + max_tokens=LISTWISE_JOIN_MAX_TOKENS, + ) + + def _build_listwise_batch_request( tasks: tuple[_ListwiseTask, ...], ) -> PromptBatchRequest: diff --git a/src/agent_memory/adapters/lotus/site_batching.py b/src/agent_memory/adapters/lotus/site_batching.py new file mode 100644 index 0000000..4460b66 --- /dev/null +++ b/src/agent_memory/adapters/lotus/site_batching.py @@ -0,0 +1,94 @@ +"""Opt-in, site-scoped batching of pair predicates after candidate screening.""" +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import pandas as pd + +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.sem_filter_batch_prompting import ( + BatchPromptingResult, + execute_batch_prompted_sem_filter, +) + +if TYPE_CHECKING: + from agent_memory.adapters.lotus.context import LotusExecutionContext + + +@dataclass(frozen=True) +class PairFilterBatching: + """Batch one pair-filter site, partitioned by stable endpoint columns.""" + + group_by: tuple[str, ...] + prompt_batching: PromptBatching + shared_columns: tuple[str, ...] = () + pack_small_groups: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.pack_small_groups, bool): + raise TypeError("pack_small_groups must be boolean") + if not isinstance(self.group_by, tuple) or not self.group_by or any( + not isinstance(c, str) or not c for c in self.group_by + ) or len(set(self.group_by)) != len(self.group_by): + raise ValueError("group_by must contain distinct column names") + if not isinstance(self.prompt_batching, PromptBatching) or self.prompt_batching.max_tasks is None: + raise ValueError("site batching requires an explicit positive max_tasks") + if not isinstance(self.shared_columns, tuple) or any(not isinstance(c, str) or not c for c in self.shared_columns): + raise ValueError("shared_columns must be a tuple of column names") + + def to_dict(self) -> dict[str, Any]: + result = {"version": "pair-filter-batching-v2", "group_by": list(self.group_by), + "prompt_batching": self.prompt_batching.to_dict()} + if self.shared_columns: + result.update(version="pair-filter-shared-context-v1", shared_columns=list(self.shared_columns)) + if self.pack_small_groups: + result.update(version="pair-filter-packed-groups-v1", pack_small_groups=True) + return result + + +def execute_site_batching( + source: pd.DataFrame, *, identities: pd.DataFrame, instruction: str, + context: LotusExecutionContext, config: PairFilterBatching, +) -> BatchPromptingResult: + """Evaluate groups independently and restore original occurrence positions.""" + if len(source) != len(identities): + raise ValueError("pair identities and predicate rows differ") + groups: dict[tuple[Any, ...], list[int]] = defaultdict(list) + for position, row in enumerate(identities.loc[:, list(config.group_by)].itertuples(index=False, name=None)): + if any(value is None or value is pd.NA or isinstance(value, float) and pd.isna(value) for value in row): + raise ValueError("site batching requires non-null endpoint identities") + try: + groups[row].append(position) + except TypeError as error: + raise ValueError("site batching requires hashable endpoint identities") from error + partitions = tuple(tuple(positions) for positions in groups.values()) + if config.pack_small_groups: + assert config.prompt_batching.max_tasks is not None + partitions = _pack_partitions(partitions, config.prompt_batching.max_tasks) + return execute_batch_prompted_sem_filter( + source, instruction=instruction, context=context, + prompt_batching=config.prompt_batching, + task_partitions=partitions, + shared_columns=config.shared_columns, + ) + + +def _pack_partitions(groups: tuple[tuple[int, ...], ...], limit: int) -> tuple[tuple[int, ...], ...]: + """Pack consecutive small groups without dropping identities or splitting them.""" + packed: list[tuple[int, ...]] = [] + pending: tuple[int, ...] = () + for group in groups: + if len(pending) + len(group) > limit: + if pending: + packed.append(pending) + pending = () + if len(group) > limit: + # The existing executor splits oversized groups using the same limit. + packed.append(group) + else: + pending += group + if pending: + packed.append(pending) + return tuple(packed) diff --git a/src/agent_memory/adapters/lotus/structured.py b/src/agent_memory/adapters/lotus/structured.py index 8dd9880..95aad69 100644 --- a/src/agent_memory/adapters/lotus/structured.py +++ b/src/agent_memory/adapters/lotus/structured.py @@ -168,12 +168,12 @@ def resolve_input_cols(source: Any, query: QueryExpr, *, operator: str) -> tuple if explicit is not None: columns = tuple(str(column) for column in explicit) else: - import lotus + from lotus.nl_expression import parse_cols try: parsed = tuple( column - for column in lotus.nl_expression.parse_cols( + for column in parse_cols( str(query.params["instruction"]) ) if column in source.columns @@ -463,6 +463,7 @@ class StructuredLMExecutor: """Internal executor that reuses LOTUS operator caching without pandas accessors.""" _obj: Any + model: Any = None @operator_cache def __call__( @@ -492,7 +493,8 @@ def __call__( from lotus.templates import task_instructions from lotus.utils import show_safe_mode - if lotus.settings.lm is None: + model = self.model if self.model is not None else lotus.settings.lm + if model is None: raise ValueError( "The language model must be an instance of LM. Please configure " "a valid language model using lotus.settings.configure()" @@ -517,7 +519,9 @@ def __call__( input_cols=input_cols, output_cols=output_cols, ) - formatted_instruction = lotus.nl_expression.nle2str( + from lotus.nl_expression import nle2str + + formatted_instruction = nle2str( formatter_instruction, list(input_cols), ) @@ -542,7 +546,7 @@ def __call__( prompts = [ task_instructions.map_formatter( - lotus.settings.lm, + model, doc, user_instruction, examples_multimodal_data=examples_multimodal_data, @@ -555,15 +559,15 @@ def __call__( ] if safe_mode: - estimated_cost = sum(lotus.settings.lm.count_tokens(prompt) for prompt in prompts) + estimated_cost = sum(model.count_tokens(prompt) for prompt in prompts) show_safe_mode(estimated_cost, len(prompts)) - current_max_tokens = int(getattr(lotus.settings.lm, "max_tokens", 512) or 512) + current_max_tokens = int(getattr(model, "max_tokens", 512) or 512) max_tokens = max(current_max_tokens, structured_max_tokens) if prompt_batching is not None: execution = _execute_structured_prompt_batches( prompts, - model=lotus.settings.lm, + model=model, prompt_batching=prompt_batching, structured_output_transport=structured_output_transport, output_cols=output_cols, @@ -603,7 +607,7 @@ def __call__( snapshots={"input": self._obj.loc[:, list(input_cols)].copy()}, ) if safe_mode: - lotus.settings.lm.print_total_usage() + model.print_total_usage() return StructuredGenerationResult( parsed_outputs=tuple(parsed_outputs), raw_outputs=tuple(raw_outputs), @@ -618,7 +622,7 @@ def __call__( "response_format": {"type": "json_object"}, } retry_result = execute_structured_lm_retry_result( - lotus.settings.lm, + model, prompts, lm_kwargs=lm_kwargs, output_cols=output_cols, @@ -667,7 +671,7 @@ def __call__( ) if safe_mode: - lotus.settings.lm.print_total_usage() + model.print_total_usage() return StructuredGenerationResult( parsed_outputs=parsed_outputs, @@ -996,6 +1000,7 @@ def execute_structured_lm_retry_result( max_retries: int, failure_extra_by_index: Mapping[int, Mapping[str, Any]] | None = None, output_validator: Callable[[str], None] | None = None, + retry_prompt_builder: Callable[[Any, str, str], Any] | None = None, ) -> StructuredLMRetryResult: """Call the LM and return retry metadata without hiding parse failures.""" @@ -1017,7 +1022,20 @@ def execute_structured_lm_retry_result( retries_left = max_retries while invalid and retries_left > 0: - retry_prompts = [prompts[index] for index in invalid] + retry_prompts = [] + for index in invalid: + prompt = prompts[index] + if retry_prompt_builder is not None: + error = structured_parse_error( + raw_outputs[index], + output_cols=output_cols, + shape=shape, + require_explanation=require_explanation, + operator=operator, + output_validator=output_validator, + ) + prompt = retry_prompt_builder(prompt, raw_outputs[index], error) + retry_prompts.append(prompt) _STRUCTURED_RETRY_STATS.retry_batches += 1 _STRUCTURED_RETRY_STATS.retry_rows += len(retry_prompts) retry_kwargs = dict(lm_kwargs) diff --git a/src/agent_memory/evaluation/agent_memory_drivers.py b/src/agent_memory/evaluation/agent_memory_drivers.py index 3c1ce53..ab935da 100644 --- a/src/agent_memory/evaluation/agent_memory_drivers.py +++ b/src/agent_memory/evaluation/agent_memory_drivers.py @@ -9,6 +9,7 @@ from pathlib import Path import pickle from typing import Any +from agent_memory.memories.zep.fact_summary import zep_memory_type, zep_storage_statements import pandas as pd @@ -21,6 +22,7 @@ semantic_pair_site_physical_contract, ) from agent_memory.adapters.lotus.context import validate_lm_max_ctx_len +from agent_memory.adapters.lotus.site_batching import PairFilterBatching from agent_memory.adapters.lotus.prompt_batching import ( PromptBatching, validate_structured_output_transport, @@ -403,8 +405,12 @@ def event_to_zep_log_row(event: BenchmarkEvent) -> dict[str, str]: if not event.timestamp: raise ValueError("Zep benchmark events require an explicit timestamp") + content = f"{event.speaker}: {event.text}" + caption = event.metadata.get("blip_caption") + if isinstance(caption, str) and caption.strip(): + content += f"\n(description of attached image: {caption.strip()})" return { - "content": f"{event.speaker}: {event.text}", + "content": content, "role": event.speaker, "speaker": event.speaker, "reference_time": event.timestamp, @@ -805,6 +811,16 @@ def __init__( prompt_batching: PromptBatching | None = None, sem_agg_dispatch: str = "sequential", structured_output_transport: str = "chat-json-object", + physical_fusion: str = "disabled", + pair_filter_batching: dict[str, PairFilterBatching] | None = None, + listwise_join_batching: dict[str, PromptBatching] | None = None, + groupby_prompt_batching: dict[str, PromptBatching] | None = None, + reuse_unchanged_entity_name: bool = False, + parallel_fact_extraction: bool = False, + sem_agg_prompt_batching: PromptBatching | None = None, + predicate_reuse_sites: tuple[str, ...] = (), + memory_num_retries: int = BENCHMARK_LM_NUM_RETRIES, + structured_parse_retries: int = 3, lm_max_ctx_len: int | None = None, embedding_device: str = "cpu", semantic_trace_snapshot_mode: str = "compact", @@ -837,6 +853,16 @@ def __init__( self.prompt_batching = prompt_batching self.sem_agg_dispatch = sem_agg_dispatch self.structured_output_transport = structured_output_transport + self.physical_fusion = physical_fusion + self.pair_filter_batching = dict(pair_filter_batching or {}) + self.listwise_join_batching = dict(listwise_join_batching or {}) + self.groupby_prompt_batching = dict(groupby_prompt_batching or {}) + self.reuse_unchanged_entity_name = reuse_unchanged_entity_name + self.parallel_fact_extraction = parallel_fact_extraction + self.sem_agg_prompt_batching = sem_agg_prompt_batching + self.predicate_reuse_sites = predicate_reuse_sites + self.memory_num_retries = memory_num_retries + self.structured_parse_retries = structured_parse_retries self.lm_max_ctx_len = lm_max_ctx_len embedding_contract = _semantic_pair_embedding_contract( self.semantic_pair_profiles @@ -880,6 +906,16 @@ def from_environment( prompt_batching: PromptBatching | None = None, sem_agg_dispatch: str = "sequential", structured_output_transport: str = "chat-json-object", + physical_fusion: str = "disabled", + pair_filter_batching: dict[str, PairFilterBatching] | None = None, + listwise_join_batching: dict[str, PromptBatching] | None = None, + groupby_prompt_batching: dict[str, PromptBatching] | None = None, + reuse_unchanged_entity_name: bool = False, + parallel_fact_extraction: bool = False, + sem_agg_prompt_batching: PromptBatching | None = None, + predicate_reuse_sites: tuple[str, ...] = (), + memory_num_retries: int = BENCHMARK_LM_NUM_RETRIES, + structured_parse_retries: int = 3, lm_max_ctx_len: int | None = None, embedding_device: str = "cpu", semantic_trace_snapshot_mode: str = "compact", @@ -937,6 +973,16 @@ def from_environment( prompt_batching=prompt_batching, sem_agg_dispatch=sem_agg_dispatch, structured_output_transport=structured_output_transport, + physical_fusion=physical_fusion, + pair_filter_batching=pair_filter_batching, + listwise_join_batching=listwise_join_batching, + groupby_prompt_batching=groupby_prompt_batching, + reuse_unchanged_entity_name=reuse_unchanged_entity_name, + parallel_fact_extraction=parallel_fact_extraction, + sem_agg_prompt_batching=sem_agg_prompt_batching, + predicate_reuse_sites=predicate_reuse_sites, + memory_num_retries=memory_num_retries, + structured_parse_retries=structured_parse_retries, lm_max_ctx_len=lm_max_ctx_len, embedding_device=embedding_device, semantic_trace_snapshot_mode=semantic_trace_snapshot_mode, @@ -971,7 +1017,6 @@ def __call__( import agent_memory as am from agent_memory.adapters.lotus import LotusAdapter from agent_memory.adapters.lotus.context import LotusExecutionConfig - from agent_memory.memories.zep.storage import GRAPHITI_NEO4J_STATEMENTS from agent_memory.planner import DifferentialRules, PolicyDifferentiator from agent_memory.runtime import MemoryRuntime from agent_memory.storage import StorageDeployment @@ -987,14 +1032,22 @@ def __call__( self.connector.embedding_provider = traced_embedding_provider storage = StorageDeployment( connector=self.connector, - statements=GRAPHITI_NEO4J_STATEMENTS, + statements=zep_storage_statements(self.physical_fusion), namespace=namespace, ) adapter = LotusAdapter( model=self.model_id, config=LotusExecutionConfig( semantic_trace_dir=trace_dir, - lm_num_retries=BENCHMARK_LM_NUM_RETRIES, + lm_num_retries=self.memory_num_retries, + structured_parse_retries=self.structured_parse_retries, + pair_filter_batching=self.pair_filter_batching, + listwise_join_batching=self.listwise_join_batching, + groupby_prompt_batching=self.groupby_prompt_batching, + reuse_unchanged_entity_name=self.reuse_unchanged_entity_name, + parallel_fact_extraction=self.parallel_fact_extraction, + sem_agg_prompt_batching=self.sem_agg_prompt_batching, + predicate_reuse_sites=self.predicate_reuse_sites, lm_model_kwargs={ **( {"max_ctx_len": self.lm_max_ctx_len} @@ -1015,6 +1068,7 @@ def __call__( prompt_batching=self.prompt_batching, sem_agg_dispatch=self.sem_agg_dispatch, structured_output_transport=self.structured_output_transport, + physical_fusion=self.physical_fusion, semantic_trace_snapshot_mode=self.semantic_trace_snapshot_mode, ), pair_embedding_provider=( @@ -1024,10 +1078,10 @@ def __call__( policy = PolicyDifferentiator( rules=DifferentialRules(grouped_agg_rule=self.grouped_agg_rule) ).differentiate( - am.ZepMemory.spec(), + zep_memory_type(self.physical_fusion).spec(), statements=storage.statements, ) - memory = am.ZepMemory(adapter=adapter) + memory = zep_memory_type(self.physical_fusion)(adapter=adapter) if self.refresh_every > 1: memory._runtime = MemoryRuntime( policy, diff --git a/src/agent_memory/evaluation/artifacts.py b/src/agent_memory/evaluation/artifacts.py index 5c877ca..1891cf8 100644 --- a/src/agent_memory/evaluation/artifacts.py +++ b/src/agent_memory/evaluation/artifacts.py @@ -346,6 +346,7 @@ def initialize( expected["answer_parser_contracts"] = dict( sorted(answer_parser_contracts.items()) ) + expected = _json_safe(expected) manifest_path = self.output_dir / "manifest.json" if manifest_path.exists(): actual = json.loads(manifest_path.read_text(encoding="utf-8")) diff --git a/src/agent_memory/evaluation/locomo_contracts.py b/src/agent_memory/evaluation/locomo_contracts.py index a365c3a..80e9065 100644 --- a/src/agent_memory/evaluation/locomo_contracts.py +++ b/src/agent_memory/evaluation/locomo_contracts.py @@ -28,7 +28,7 @@ LOCOMO_OFFICIAL_SCORER_DIGEST = sha256( b"locomo_answer_score:category-1-5:v1" ).hexdigest() -LOCOMO_ANSWER_PARSER_ID = "locomo-structured-or-raw-text:v1" +LOCOMO_ANSWER_PARSER_ID = "locomo-structured-or-raw-text:v2" def _parse_json_object(response: str) -> Mapping[str, Any]: @@ -47,12 +47,15 @@ def _parse_answer(response: str) -> str: if not raw: raise ValueError("LOCOMO answer must be non-empty") try: - answer = _parse_json_object(raw).get("answer") + parsed = _parse_json_object(raw) except (json.JSONDecodeError, ValueError): + if raw.startswith(("{", "[", "```")): + raise ValueError("LOCOMO structured answer must be a valid JSON object") from None return raw + answer = parsed.get("answer") if isinstance(answer, str) and answer.strip(): return answer.strip() - return raw + raise ValueError("LOCOMO structured answer requires a non-empty answer string") def _parse_zep_judge(response: str) -> Mapping[str, str]: diff --git a/src/agent_memory/evaluation/run.py b/src/agent_memory/evaluation/run.py index a95ef2c..c8f4b93 100644 --- a/src/agent_memory/evaluation/run.py +++ b/src/agent_memory/evaluation/run.py @@ -9,6 +9,7 @@ from pathlib import Path import re from typing import Any +from agent_memory.memories.zep.fact_summary import zep_memory_type, zep_storage_statements from .agent_memory_drivers import ( ClaudeMemoryDriverFactory, @@ -33,6 +34,8 @@ SEM_JOIN_TOPK_METHODS, ) from agent_memory.adapters.lotus.json_output import JSON_REPAIR_VERSION +from agent_memory.adapters.lotus.site_batching import PairFilterBatching +from agent_memory.planner.serialization import stable_json from agent_memory.adapters.lotus.prompt_batching import ( PromptBatching, validate_structured_output_transport, @@ -62,7 +65,7 @@ "claude-memory-declared-sem-topk:v1", ), "zep-memory": ( - "benchmark-event-to-zep-log:v1", + "benchmark-event-to-zep-log:v2", "zep-memory-entity-rrf-fact-bfs-cross-encoder:v1", ), "mem0-memory": ( @@ -128,6 +131,17 @@ def run_agent_memory_bundle( sem_agg_dispatch: str = "sequential", prompt_batching: PromptBatching | None = None, structured_output_transport: str = "chat-json-object", + physical_fusion: str = "disabled", + pair_filter_batching: Mapping[str, PairFilterBatching] | None = None, + listwise_join_batching: Mapping[str, PromptBatching] | None = None, + groupby_prompt_batching: Mapping[str, PromptBatching] | None = None, + reuse_unchanged_entity_name: bool = False, + parallel_fact_extraction: bool = False, + sem_agg_prompt_batching: PromptBatching | None = None, + predicate_reuse_sites: tuple[str, ...] = (), + memory_num_retries: int | None = None, + structured_parse_retries: int | None = None, + parse_attempts: int = 4, lm_max_ctx_len: int | None = None, semantic_pair_profile: str = "oracle-only", semantic_pair_top_k: int | None = None, @@ -152,6 +166,65 @@ def run_agent_memory_bundle( structured_output_transport, model=memory_provider_model_id ) grouped_agg_rule = resolve_grouped_agg_rule(system_id, grouped_agg_rule) + site_execution_options: dict[str, Any] = {} + if parallel_fact_extraction: + from agent_memory.adapters.lotus.context import LotusExecutionConfig + LotusExecutionConfig(parallel_fact_extraction=True, physical_fusion=physical_fusion, + lm_enable_cache=lotus_cache_mode != "disabled", prompt_batching=prompt_batching) + site_execution_options["parallel_fact_extraction"] = True + if groupby_prompt_batching or reuse_unchanged_entity_name: + from agent_memory.adapters.lotus.context import LotusExecutionConfig + LotusExecutionConfig(groupby_prompt_batching=dict(groupby_prompt_batching or {}), + reuse_unchanged_entity_name=reuse_unchanged_entity_name, physical_fusion=physical_fusion, + prompt_batching=prompt_batching, structured_output_transport=structured_output_transport) + if groupby_prompt_batching: + site_execution_options["groupby_prompt_batching"] = dict(groupby_prompt_batching) + if reuse_unchanged_entity_name: + site_execution_options["reuse_unchanged_entity_name"] = True + if listwise_join_batching: + from agent_memory.adapters.lotus.context import LotusExecutionConfig + LotusExecutionConfig(listwise_join_batching=listwise_join_batching, + prompt_batching=prompt_batching, sem_join_topk_method=sem_join_topk_method or "listwise", + structured_output_transport=structured_output_transport) + site_execution_options["listwise_join_batching"] = dict(listwise_join_batching) + if sem_agg_prompt_batching is not None: + from agent_memory.adapters.lotus.context import LotusExecutionConfig + LotusExecutionConfig( + sem_agg_prompt_batching=sem_agg_prompt_batching, + prompt_batching=prompt_batching, + sem_agg_dispatch=sem_agg_dispatch, + structured_output_transport=structured_output_transport, + ) + site_execution_options["sem_agg_prompt_batching"] = sem_agg_prompt_batching + if pair_filter_batching: + site_execution_options["pair_filter_batching"] = dict(pair_filter_batching) + if predicate_reuse_sites: + from agent_memory.adapters.lotus.context import LotusExecutionConfig + LotusExecutionConfig(predicate_reuse_sites=predicate_reuse_sites) + site_execution_options["predicate_reuse_sites"] = predicate_reuse_sites + for name, value in (("memory_num_retries", memory_num_retries), ("structured_parse_retries", structured_parse_retries)): + if value is not None: + if type(value) is not int or value < 0: + raise ValueError(f"{name} must be a nonnegative integer") + site_execution_options[name] = value + if site_execution_options and system_id != "zep-memory": + raise ValueError("site execution options currently support Zep only") + if physical_fusion != "disabled": + from agent_memory.adapters.lotus.context import LotusExecutionConfig + if system_id != "zep-memory": + raise ValueError("target-state fusion is registered only for Zep") + LotusExecutionConfig( + physical_fusion=physical_fusion, + structured_output_transport=structured_output_transport, + prompt_batching=prompt_batching, + sem_agg_dispatch=sem_agg_dispatch, + sem_join_topk_method=sem_join_topk_method or "listwise", + pair_filter_batching=dict(pair_filter_batching or {}), + listwise_join_batching=dict(listwise_join_batching or {}), + groupby_prompt_batching=dict(groupby_prompt_batching or {}), + reuse_unchanged_entity_name=reuse_unchanged_entity_name, + sem_agg_prompt_batching=sem_agg_prompt_batching, + ) if ( isinstance(refresh_every, bool) or not isinstance(refresh_every, int) @@ -280,14 +353,13 @@ def run_agent_memory_bundle( else: from agent_memory.memories.zep.storage import ( GRAPHITI_BGE_M3, - GRAPHITI_NEO4J_STATEMENTS, ) policy = PolicyDifferentiator( rules=DifferentialRules(grouped_agg_rule=grouped_agg_rule) ).differentiate( - am.ZepMemory.spec(), - statements=GRAPHITI_NEO4J_STATEMENTS, + zep_memory_type(physical_fusion).spec(), + statements=zep_storage_statements(physical_fusion), ) operators = ("sem_filter", "sem_join", "sem_groupby") embedding = GRAPHITI_BGE_M3 @@ -326,15 +398,14 @@ def run_agent_memory_bundle( import agent_memory as am from agent_memory.memories.zep.storage import ( GRAPHITI_BGE_M3, - GRAPHITI_NEO4J_STATEMENTS, ) from agent_memory.planner import DifferentialRules, PolicyDifferentiator policy = PolicyDifferentiator( rules=DifferentialRules(grouped_agg_rule=grouped_agg_rule) ).differentiate( - am.ZepMemory.spec(), - statements=GRAPHITI_NEO4J_STATEMENTS, + zep_memory_type(physical_fusion).spec(), + statements=zep_storage_statements(physical_fusion), ) semantic_pair_profiles = build_operator_semantic_pair_profiles( policy, @@ -405,6 +476,23 @@ def run_agent_memory_bundle( if part ] maintenance_execution_id = "|".join(maintenance_execution_parts) + from agent_memory.planner.physical import FUSION_VERSION, COMBINED_VERSION, SUMMARY_VERSION, REPRESENTATIVE_VERSION + fusion_version: str | None = None + if physical_fusion != "disabled": + fusion_version = {"zep-target-state": FUSION_VERSION, "zep-combined": COMBINED_VERSION, + "zep-fact-summary": SUMMARY_VERSION, "zep-representative": REPRESENTATIVE_VERSION}[physical_fusion] + maintenance_execution_id += f"|physical-fusion:{fusion_version}" + if site_execution_options: + identity = {**site_execution_options, "pair_filter_batching": { + site: config.to_dict() for site, config in sorted((pair_filter_batching or {}).items()) + }} + if sem_agg_prompt_batching is not None: + identity["sem_agg_prompt_batching"] = sem_agg_prompt_batching.to_dict() + if listwise_join_batching: + identity["listwise_join_batching"] = {site: setting.to_dict() for site, setting in sorted(listwise_join_batching.items())} + if groupby_prompt_batching: + identity["groupby_prompt_batching"] = {site: setting.to_dict() for site, setting in sorted(groupby_prompt_batching.items())} + maintenance_execution_id += "|site-execution:" + sha256(stable_json(identity).encode()).hexdigest() framework_cache_mode = ( LOTUS_MEMORY_CACHE_ID if lotus_cache_mode == "memory" else "disabled" ) @@ -470,6 +558,17 @@ def run_agent_memory_bundle( "sem_agg_dispatch": sem_agg_dispatch, } ) + if fusion_version is not None: + lotus_execution_provenance["physical_fusion"] = physical_fusion + lotus_execution_provenance["fusion_version"] = fusion_version + if site_execution_options: + lotus_execution_provenance["site_execution"] = { + **{k: v for k, v in site_execution_options.items() if k not in {"pair_filter_batching", "sem_agg_prompt_batching", "listwise_join_batching", "groupby_prompt_batching"}}, + **({"groupby_prompt_batching": {site: setting.to_dict() for site, setting in sorted(groupby_prompt_batching.items())}} if groupby_prompt_batching else {}), + **({"listwise_join_batching": {site: setting.to_dict() for site, setting in sorted(listwise_join_batching.items())}} if listwise_join_batching else {}), + **({"sem_agg_prompt_batching": sem_agg_prompt_batching.to_dict()} if sem_agg_prompt_batching is not None else {}), + "pair_filter_batching": {site: config.to_dict() for site, config in sorted((pair_filter_batching or {}).items())}, + } if prompt_batching is not None: lotus_execution_provenance["prompt_batching"] = ( prompt_batching.to_dict() @@ -595,6 +694,8 @@ def run_agent_memory_bundle( ) elif system_id == "zep-memory": driver_factory = ZepMemoryDriverFactory.from_environment( + **site_execution_options, + **({"physical_fusion": physical_fusion} if physical_fusion != "disabled" else {}), base_namespace=base_namespace or _namespace(bundle.benchmark_id, output_dir), model_id=memory_provider_model_id, @@ -664,6 +765,7 @@ def run_agent_memory_bundle( else None ) BenchmarkRunner( + parse_attempts=parse_attempts, system_contract=system_contract, contracts=contracts, driver_factory=driver_factory, diff --git a/src/agent_memory/evaluation/trace_metrics.py b/src/agent_memory/evaluation/trace_metrics.py index ea0a877..fc941c6 100644 --- a/src/agent_memory/evaluation/trace_metrics.py +++ b/src/agent_memory/evaluation/trace_metrics.py @@ -44,7 +44,8 @@ def normalize_provider_calls( tuple[str, int], deque[Mapping[str, Any]] ] = defaultdict(deque) for event in events: - if event.get("event_type") == "llm_call" and "status" not in event: + if (event.get("event_type") == "llm_call" and "status" not in event + and not _lotus_cache_only_call(event)): key = _agent_call_key(event, item_key="llm_item_index") if key is not None: pending_semantic_calls[key].append(event) @@ -76,6 +77,8 @@ def normalize_provider_calls( elif event_type == "llm_call" and "status" in event: rows.append(_native_graphiti_row(event, pricing)) elif event_type in {"llm_call", "llm_batch_error"}: + if _lotus_cache_only_call(event): + continue if int(event.get("llm_item_index") or 0) != 0: continue if event_type == "llm_call" and id(event) in paired_semantic_calls: @@ -84,6 +87,20 @@ def normalize_provider_calls( return rows +def _lotus_cache_only_call(event: Mapping[str, Any]) -> bool: + """Exclude only proven all-cache successes, never unknown or failed calls.""" + size = _optional_int(event.get("llm_batch_size")) + hits = _optional_int(event.get("usage_cache_hits")) + return ( + event.get("event_type") == "llm_call" + and "status" not in event + and size is not None and size > 0 + and hits is not None and hits >= size + and all(_optional_int(event.get(f"usage_physical_{kind}_tokens")) == 0 + for kind in ("prompt", "completion", "total")) + ) + + def summarize_provider_calls(rows: Sequence[Mapping[str, Any]]) -> dict[str, Any]: """Summarize normalized calls without replacing unknown usage with guesses.""" @@ -229,23 +246,23 @@ def _agent_runner_row( pricing: PricingSnapshot | None, ) -> dict[str, Any]: error = event.get("event_type") == "llm_batch_error" + prompt = _optional_int(event.get("usage_prompt_tokens")) + completion = _optional_int(event.get("usage_completion_tokens")) return _row( event, source="agent-runner", status="error" if error else "success", latency_ms=_seconds_to_milliseconds(event.get("latency_sec")), - prompt_tokens=_optional_int(event.get("usage_prompt_tokens")), + prompt_tokens=prompt, cache_hit_tokens=_optional_int( event.get("usage_prompt_cache_hit_tokens") ), cache_miss_tokens=_optional_int( event.get("usage_prompt_cache_miss_tokens") ), - completion_tokens=_optional_int(event.get("usage_completion_tokens")), + completion_tokens=completion, reasoning_tokens=_optional_int(event.get("usage_reasoning_tokens")), - usage_available=not error and any( - key.startswith("usage_") for key in event - ), + usage_available=not error and prompt is not None and completion is not None, pricing=pricing, ) diff --git a/src/agent_memory/evaluation/zep/answering.py b/src/agent_memory/evaluation/zep/answering.py index caab42a..012a3c3 100644 --- a/src/agent_memory/evaluation/zep/answering.py +++ b/src/agent_memory/evaluation/zep/answering.py @@ -82,6 +82,11 @@ {context} Question: {question} + +Return a JSON object with an "answer" field containing your actual answer as a +non-empty string. The schema below describes the response, not the answer: +do not copy the schema or return its properties, required, title, or type fields. +If evidence is missing, put "No information available" in the "answer" field. """.strip() ZEP_JUDGE_SYSTEM_PROMPT = ( diff --git a/src/agent_memory/memories/zep/fact_summary.py b/src/agent_memory/memories/zep/fact_summary.py new file mode 100644 index 0000000..544858c --- /dev/null +++ b/src/agent_memory/memories/zep/fact_summary.py @@ -0,0 +1,95 @@ +"""Entity identities and summaries derived from incident facts.""" + +from dataclasses import replace + +from agent_memory.api import Memory +from agent_memory.memories.zep.policy import ZepMemory +from agent_memory.policy.aggregates import SemanticAggregateSpec, sem_agg +from agent_memory.policy.expressions import case_when +from agent_memory.policy.logical import UserQuery +from agent_memory.policy.relation import Relation +from agent_memory.policy.retrieval import BFS, BM25, RRF, CosineSimilarity, CrossEncoder, RetrievalQuery +from agent_memory.planner.physical import replace_query, walk +from agent_memory.storage import StatementSet + +IDENTITY_SPEC = sem_agg( + input_cols=["name"], output_cols={"name": "Canonical entity name."}, + instruction="Choose one canonical {name} from the supplied names. Prefer the most complete name for the same entity. Never invent a name.", +) +SUMMARY_SPEC = sem_agg( + input_cols=["summary"], output_cols={"summary": "Fact-derived entity summary."}, + instruction=( + "Join the supplied {summary} texts with newlines, preserving their order and content. " + "If the result is at most 2000 characters, return it unchanged. Otherwise compress it " + "to at most 1000 characters, preserving supported names, relationships, dates and " + "qualifiers. State facts directly. Do not invent information." + ), +) + +_old_aggregate = next(q for q in walk(ZepMemory.entities.expr) if q.op == "agg") +_identity_aggregate = replace(_old_aggregate, params={"aggregates": tuple( + IDENTITY_SPEC if isinstance(s, SemanticAggregateSpec) else s + for s in _old_aggregate.params["aggregates"] +)}) +_identity_expr = replace_query(ZepMemory.entities.expr, _old_aggregate, _identity_aggregate) +_identity_expr = replace(_identity_expr, params={"columns": ( + "entity_id", "name", "entity_type", "mentions", +)}) + + +class ZepFactSummaryMemory(Memory): + """Maintain entity identity first, then summarize incident facts separately.""" + + log = ZepMemory.log + episodes = ZepMemory.episodes + _identities = Relation(_identity_expr) + _episode_entities = Relation(replace_query(ZepMemory._episode_entities.expr, ZepMemory.entities.expr, _identity_expr)) + facts = Relation(replace_query(ZepMemory.facts.expr, ZepMemory.entities.expr, _identity_expr)) + # Each fact contributes to its endpoints, not every entity in its message. + _source_facts = facts.assign(entity_id=facts.col("source_entity_id"), summary=facts.col("fact")).select(["entity_id", "summary"]) + _target_facts = facts.filter(facts.col("source_entity_id") != facts.col("target_entity_id")).assign( + entity_id=facts.col("target_entity_id"), summary=facts.col("fact"), + ).select(["entity_id", "summary"]) + _incident_facts = _source_facts.union_by_name(_target_facts).join(_identities.select(["entity_id"]), on="entity_id") + _summaries = _incident_facts.group_by("entity_id").agg(SUMMARY_SPEC) + _joined = _identities.join(_summaries, on="entity_id", how="left") + entities = _joined.assign(summary=case_when(_joined.col("summary").is_null(), "", _joined.col("summary"))).select( + ["entity_id", "name", "entity_type", "summary", "mentions"] + ) + _retrieved_entities = entities.search( + UserQuery(), methods=[BM25(), CosineSimilarity()], reranker=RRF(), limit=20, + ).select(["record_id", "name", "summary", "rank", "score"]) + retrieval_query = RetrievalQuery( + entities=_retrieved_entities, + facts=facts.search(UserQuery(), methods=[BM25(), CosineSimilarity(), BFS(origins=_retrieved_entities, max_depth=3)], + reranker=CrossEncoder(model="BAAI/bge-reranker-v2-m3"), limit=20).select( + ["record_id", "fact", "valid_at", "invalid_at", "expired_at", "rank", "score"]), + ) + + +def zep_memory_type(strategy: str) -> type[Memory]: + """Select the explicit logical variant for benchmark configuration.""" + import agent_memory as am + if strategy == "zep-representative": + from agent_memory.memories.zep.representative import ZepRepresentativeMemory + return ZepRepresentativeMemory + return ZepFactSummaryMemory if strategy == "zep-fact-summary" else am.ZepMemory + + +def zep_storage_statements(strategy: str) -> StatementSet: + """Bind the selected logical views to the unchanged Graphiti targets.""" + from agent_memory.memories.zep.storage import GRAPHITI_NEO4J_STATEMENTS + if strategy not in {"zep-fact-summary", "zep-representative"}: + return GRAPHITI_NEO4J_STATEMENTS + from agent_memory.memories.zep.representative import ZepRepresentativeMemory + memory = ZepRepresentativeMemory if strategy == "zep-representative" else ZepFactSummaryMemory + replacements = ( + (ZepMemory.entities, memory.entities), + (ZepMemory.facts, memory.facts), + (ZepMemory._episode_entities, memory._episode_entities), + ) + statements = StatementSet() + for statement in GRAPHITI_NEO4J_STATEMENTS.statements: + relation = next((new for old, new in replacements if statement.query == old.expr), Relation(statement.query)) + statements = statements.add_insert(statement.target, relation) + return statements diff --git a/src/agent_memory/memories/zep/representative.py b/src/agent_memory/memories/zep/representative.py new file mode 100644 index 0000000..27a9c90 --- /dev/null +++ b/src/agent_memory/memories/zep/representative.py @@ -0,0 +1,62 @@ +"""Fact groups represented by their earliest occurrence. + +Provenance and temporal fields retain their aggregation rules. +""" + +from dataclasses import replace + +from agent_memory.api import Memory +from agent_memory.memories.zep.fact_summary import ZepFactSummaryMemory +from agent_memory.memories.zep.policy import ZepMemory +from agent_memory.planner.physical import replace_query +from agent_memory.policy.aggregates import MinAggregateSpec, SemanticAggregateSpec, arg_min +from agent_memory.policy.relation import Relation, SearchRelation +from agent_memory.policy.retrieval import RetrievalQuery +from agent_memory.policy.logical import QueryExpr + + +_old = ZepMemory._deduplicated_facts.expr.inputs[0] +_representative = replace(_old, params={"aggregates": tuple( + arg_min(order_by=("add_seq", "fact_ordinal"), columns=("relation_type", "fact")) + if isinstance(spec, SemanticAggregateSpec) else spec + for spec in _old.params["aggregates"] + if not (isinstance(spec, MinAggregateSpec) and spec.output_col == "add_seq") +)}) + +_TEMPORAL_GROUP_INSTRUCTION = str(_old.inputs[0].params["instruction"]) + """ +For time-specific events, use the resolved event time {valid_at} to interpret +relative dates in the fact text. Separate occurrences of an event are not +duplicates just because their wording is similar. An unknown event time is +not evidence that two occurrences are the same. For ongoing or timeless +claims, a difference in timestamps alone does not establish a different fact. +""" + + +def representative_query(query: QueryExpr) -> QueryExpr: + """Replace fact synthesis, preserving provenance and temporal aggregation.""" + # Entity identity has already been replaced in the fact-summary variant. + from agent_memory.planner.physical import walk + candidates = [q for q in walk(query) if q.op == "agg" and q.params == _old.params] + for candidate in candidates: + group = candidate.inputs[0] + temporal_group = replace(group, params={**group.params, + "input_cols": (*group.params["input_cols"], "valid_at"), + "instruction": _TEMPORAL_GROUP_INSTRUCTION.strip()}) + query = replace_query(query, candidate, replace(candidate, params=_representative.params, + inputs=(temporal_group,))) + return query + + +class ZepRepresentativeMemory(Memory): + """Use stable fact representatives with the existing fact-summary contract.""" + + log = ZepFactSummaryMemory.log + episodes = ZepFactSummaryMemory.episodes + _identities = ZepFactSummaryMemory._identities + _episode_entities = ZepFactSummaryMemory._episode_entities + facts = Relation(representative_query(ZepFactSummaryMemory.facts.expr)) + entities = Relation(representative_query(ZepFactSummaryMemory.entities.expr)) + retrieval_query = RetrievalQuery(**{ + name: SearchRelation(representative_query(query)) + for name, query in ZepFactSummaryMemory.retrieval_query.channels.items() + }) diff --git a/src/agent_memory/planner/differential_policy.py b/src/agent_memory/planner/differential_policy.py index 6175b67..c1a7972 100644 --- a/src/agent_memory/planner/differential_policy.py +++ b/src/agent_memory/planner/differential_policy.py @@ -7,7 +7,7 @@ from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType -from agent_memory.policy.aggregates import SemanticAggregateSpec +from agent_memory.policy.aggregates import ArgMinAggregateSpec, SemanticAggregateSpec from agent_memory.policy.logical import MemorySpec, MemoryView, QueryExpr from agent_memory.policy.retrieval import RetrievalQuery from agent_memory.policy.schema import output_columns @@ -370,7 +370,8 @@ def _compile_grouped_aggregate(self, query: QueryExpr, *, spec: MemorySpec) -> s semantic = query.op == "sem_agg" or ( query.op == "agg" and any( - isinstance(specification, SemanticAggregateSpec) + (isinstance(specification, SemanticAggregateSpec) + or (grouped.op == "sem_groupby" and isinstance(specification, ArgMinAggregateSpec))) for specification in query.params.get("aggregates", ()) ) ) diff --git a/src/agent_memory/planner/physical.py b/src/agent_memory/planner/physical.py new file mode 100644 index 0000000..1318438 --- /dev/null +++ b/src/agent_memory/planner/physical.py @@ -0,0 +1,264 @@ +"""Explicit physical rewrites over an already differentiated policy.""" + +from dataclasses import replace +from hashlib import sha256 +from collections.abc import Callable, Iterator + +from agent_memory.planner.differential_policy import DifferentiatedPolicy +from agent_memory.planner.rules import JOIN_MAP_BINDING_NAME +from agent_memory.planner.serialization import stable_json +from agent_memory.policy.aggregates import SemanticAggregateSpec +from agent_memory.policy.logical import QueryExpr + +FUSION_VERSION = "zep-target-state-v1" +COMBINED_VERSION = "zep-combined-v2" +SUMMARY_VERSION = "zep-fact-summary-v2" +REPRESENTATIVE_VERSION = "zep-representative-v2" +PREDICATE_DECISIONS_INPUT = "__physical_predicate_decisions" + + +def enable_predicate_reuse( + policy: DifferentiatedPolicy, node_ids: set[str], +) -> DifferentiatedPolicy: + """Register fixed textual predicate decisions as physical executor state.""" + if not node_ids: + return policy + nodes = dict(policy.nodes) + for node_id in node_ids: + node = nodes[node_id] + if node.query.op != "sem_filter" or node.execution_kind != "semantic_row": + raise ValueError("predicate reuse requires a row-local sem_filter node") + nodes[node_id] = replace(node, execution_kind="semantic_predicate") + fingerprint = sha256(stable_json({ + "base": policy.fingerprint, "predicate-reuse-v1": sorted(node_ids), + }).encode()).hexdigest() + return replace(policy, nodes=nodes, fingerprint=fingerprint) + + +def walk(query: QueryExpr) -> Iterator[QueryExpr]: + """Visit a query tree, including repeated references.""" + yield query + for child in query.inputs: + yield from walk(child) + + +def replace_query(query: QueryExpr, old: QueryExpr, new: QueryExpr) -> QueryExpr: + """Substitute a selected subtree without modifying the original plan.""" + if query == old: + return new + inputs = tuple(replace_query(q, old, new) for q in query.inputs) + if all(a is b for a, b in zip(inputs, query.inputs, strict=True)): + return query + return replace(query, inputs=inputs) + + +def _zep_target_state(policy: DifferentiatedPolicy) -> DifferentiatedPolicy: + from agent_memory.memories.zep.policy import ZepMemory + + template = ZepMemory._deduplicated_facts.expr.inputs[0] + return _fuse_registered_state(policy, template, FUSION_VERSION, ("relation_type", "fact")) + + +def _fuse_registered_state( + policy: DifferentiatedPolicy, template: QueryExpr, version: str, + state_columns: tuple[str, ...], +) -> DifferentiatedPolicy: + matches = [ + node + for node in policy.nodes.values() + if node.query.op == "agg" + and node.query.params == template.params + and node.query.inputs[0].op == "sem_groupby" + and node.query.inputs[0].params == template.inputs[0].params + ] + if policy.grouped_agg_rule != "rule-join-map" or len(matches) != 1: + raise ValueError( + "zep-target-state requires exactly one registered Zep fact join-map site" + ) + node = matches[0] + query = node.maintenance_query + if query is not None and query.op == "fused_target_state": + if query.params.get("version") != version: + raise ValueError("unsupported target-state fusion version") + return policy + if ( + query is None + or query.op != "let" + or query.params.get("name") != JOIN_MAP_BINDING_NAME + ): + raise ValueError("unsupported Zep join-map maintenance shape") + joined, body = query.inputs + if ( + joined.op != "sem_join" + or joined.params.get("k") != 1 + or joined.params.get("how") != "outer" + ): + raise ValueError("target-state fusion requires an exclusive outer join") + aggregates = [q for q in walk(body) if q.op == "agg"] + if len(aggregates) != 1: + raise ValueError("target-state fusion requires one state aggregate") + aggregate = aggregates[0] + semantic = [ + s + for s in aggregate.params["aggregates"] + if isinstance(s, SemanticAggregateSpec) + ] + if len(semantic) != 1 or tuple(c.name for c in semantic[0].output_cols) != state_columns: + raise ValueError("unsupported target-state semantic output contract") + fused = QueryExpr( + op="fused_target_state", + inputs=query.inputs, + params={"version": version, "aggregate": aggregate}, + ) + nodes = dict(policy.nodes) + nodes[node.node_id] = replace(node, maintenance_query=fused) + fingerprint = sha256( + stable_json( + { + "base": policy.fingerprint, + "fusion": version, + "site": node.node_id, + "query": fused, + } + ).encode() + ).hexdigest() + return replace(policy, nodes=nodes, fingerprint=fingerprint) + + +def _zep_combined(policy: DifferentiatedPolicy, *, fact_summary: bool = False) -> DifferentiatedPolicy: + from agent_memory.memories.zep.policy import ZepMemory + + version = SUMMARY_VERSION if fact_summary else COMBINED_VERSION + entity_template = next(q for q in walk(ZepMemory.entities.expr) if q.op == "agg") + entity_columns = ("name", "summary") + if fact_summary: + from agent_memory.memories.zep.fact_summary import ZepFactSummaryMemory + entity_template = next(q for q in walk(ZepFactSummaryMemory._identities.expr) if q.op == "agg") + entity_columns = ("name",) + + existing = [n.maintenance_query for n in policy.nodes.values() + if n.maintenance_query is not None and n.maintenance_query.op == "fused_target_state"] + if len(existing) == 2 and all(q.params.get("version") == version for q in existing): + if (sum(bool(q.params.get("identity_reuse")) for q in existing) == 1 + and sum(bool(child.params.get("singleton_identity")) for q in existing for child in walk(q)) == 1): + return policy + prepared = _fuse_registered_state( + policy, ZepMemory._deduplicated_facts.expr.inputs[0], version, + ("relation_type", "fact"), + ) + prepared = _fuse_registered_state( + prepared, entity_template, version, entity_columns, + ) + nodes = dict(prepared.nodes) + for node_id, node in nodes.items(): + query = node.maintenance_query + if query is None or query.op != "fused_target_state": + continue + if tuple(c.name for s in query.params["aggregate"].params["aggregates"] + if isinstance(s, SemanticAggregateSpec) for c in s.output_cols) != ("relation_type", "fact"): + if fact_summary: + from agent_memory.tracing.semantic import query_digest + joined = query.inputs[0] + delta = next(q for q in walk(joined.inputs[0]) if q.op == "agg") + query = replace_query(query, delta, replace(delta, params={**delta.params, "identity_singleton": True})) + query = replace(query, params={**query.params, "join_profile_digest": query_digest(joined)}) + nodes[node_id] = replace(node, maintenance_query=replace( + query, params={**query.params, "identity_reuse": True})) + continue + joined = query.inputs[0] + delta_aggs = [q for q in walk(joined.inputs[0]) if q.op == "agg"] + if len(delta_aggs) != 1: + raise ValueError("singleton fact strategy requires one new-fact aggregate") + delta = delta_aggs[0] + marked = replace(delta, params={**delta.params, "singleton_identity": True}) + from agent_memory.tracing.semantic import query_digest + rewritten = replace_query(query, delta, marked) + nodes[node_id] = replace(node, maintenance_query=replace( + rewritten, params={**rewritten.params, "join_profile_digest": query_digest(joined)})) + fingerprint = sha256(stable_json({"base": prepared.fingerprint, + "version": version, "singleton_fact_identity": True}).encode()).hexdigest() + return replace(prepared, nodes=nodes, fingerprint=fingerprint) + + +def summary_specs() -> tuple[SemanticAggregateSpec, SemanticAggregateSpec]: + """Exact registered raw and state-merge contracts, never prompt substring matching.""" + from agent_memory.memories.zep.fact_summary import SUMMARY_SPEC + from agent_memory.planner.rules import DifferentialInstructionRewriter + return SUMMARY_SPEC, replace(SUMMARY_SPEC, instruction=DifferentialInstructionRewriter().state_reaggregation( + SUMMARY_SPEC.instruction, state_cols=("summary",), raw_input_cols=("summary",))) + + +def lower_fact_summary_node(query: QueryExpr) -> QueryExpr: + """Use the same registered shortcuts for direct and incremental execution.""" + from agent_memory.memories.zep.fact_summary import SUMMARY_SPEC, IDENTITY_SPEC + from agent_memory.planner.rules import DifferentialInstructionRewriter + map_instruction = DifferentialInstructionRewriter().agg_to_map( + SUMMARY_SPEC.instruction, input_cols=("summary",), output_cols=SUMMARY_SPEC.output_cols) + if query.op == "agg": + specs = tuple(query.params.get("aggregates", ())) + if specs in tuple((s,) for s in summary_specs()): + return replace(query, params={**query.params, "fact_summary": True}) + if query.inputs[0].op == "sem_groupby" and tuple(s for s in specs if isinstance(s, SemanticAggregateSpec)) == (IDENTITY_SPEC,): + return replace(query, params={**query.params, "identity_singleton": True}) + if query.op == "sem_map" and query.params.get("instruction") == map_instruction and query.params.get("output_cols") == SUMMARY_SPEC.output_cols: + return replace(query, params={**query.params, "fact_summary_map": True}) + return query + + +def _zep_fact_summary(policy: DifferentiatedPolicy) -> DifferentiatedPolicy: + prepared = _zep_combined(policy, fact_summary=True) + def mark(query: QueryExpr) -> QueryExpr: + rewritten = replace(query, inputs=tuple(mark(q) for q in query.inputs)) + return lower_fact_summary_node(rewritten) + nodes = {key: replace(node, query=mark(node.query), maintenance_query=( + mark(node.maintenance_query) if node.maintenance_query is not None else None)) + for key, node in prepared.nodes.items()} + return replace(prepared, nodes=nodes) + + +def _zep_representative(policy: DifferentiatedPolicy) -> DifferentiatedPolicy: + """Fuse entity identity only; fact representatives need no synthesis call.""" + from agent_memory.memories.zep.representative import ZepRepresentativeMemory + from agent_memory.policy.aggregates import ArgMinAggregateSpec + if not any(any(isinstance(s, ArgMinAggregateSpec) for s in node.query.params.get("aggregates", ())) + for node in policy.nodes.values()): + raise ValueError("zep-representative requires the representative logical view") + fused = [n.maintenance_query for n in policy.nodes.values() + if n.maintenance_query is not None and n.maintenance_query.op == "fused_target_state"] + if len(fused) == 1 and fused[0].params.get("version") == REPRESENTATIVE_VERSION: + return policy + template = next(q for q in walk(ZepRepresentativeMemory._identities.expr) if q.op == "agg") + prepared = _fuse_registered_state(policy, template, REPRESENTATIVE_VERSION, ("name",)) + def mark(query: QueryExpr) -> QueryExpr: + return lower_fact_summary_node(replace(query, inputs=tuple(mark(q) for q in query.inputs))) + nodes = {} + for key, node in prepared.nodes.items(): + maintenance = node.maintenance_query + if maintenance is not None: + from agent_memory.tracing.semantic import query_digest + profile_digest = query_digest(maintenance.inputs[0]) if maintenance.op == "fused_target_state" else None + maintenance = mark(maintenance) + if maintenance.op == "fused_target_state": + maintenance = replace(maintenance, params={**maintenance.params, "identity_reuse": True, + "join_profile_digest": profile_digest}) + nodes[key] = replace(node, query=mark(node.query), maintenance_query=maintenance) + return replace(prepared, nodes=nodes) + + +_REWRITES: dict[str, Callable[[DifferentiatedPolicy], DifferentiatedPolicy]] = { + "zep-target-state": _zep_target_state, + "zep-combined": _zep_combined, + "zep-fact-summary": _zep_fact_summary, + "zep-representative": _zep_representative, +} + + +def optimize_policy( + policy: DifferentiatedPolicy, *, strategy: str = "disabled" +) -> DifferentiatedPolicy: + """Apply one registered rewrite; disabled returns the original object.""" + if strategy == "disabled": + return policy + if strategy not in _REWRITES: + raise ValueError(f"unknown physical fusion strategy: {strategy!r}") + return _REWRITES[strategy](policy) diff --git a/src/agent_memory/planner/rules.py b/src/agent_memory/planner/rules.py index 2e65dbb..3836702 100644 --- a/src/agent_memory/planner/rules.py +++ b/src/agent_memory/planner/rules.py @@ -6,6 +6,7 @@ import re from agent_memory.policy.aggregates import ( + ArgMinAggregateSpec, ArrayAggregateSpec, CollectListAggregateSpec, MinAggregateSpec, @@ -702,6 +703,9 @@ def _remerge_aggregate_specs( specs: list[object] = [] for spec in aggregate.params.get("aggregates", ()): + if isinstance(spec, ArgMinAggregateSpec): + specs.append(spec) + continue if isinstance(spec, SemanticAggregateSpec): state_cols = aggregate_output_names(spec) specs.append( diff --git a/src/agent_memory/policy/__init__.py b/src/agent_memory/policy/__init__.py index 3c8d588..88943b1 100644 --- a/src/agent_memory/policy/__init__.py +++ b/src/agent_memory/policy/__init__.py @@ -1,6 +1,6 @@ """Logical policy authoring objects and immutable query descriptors.""" -from .aggregates import array_agg, avg, collect_list, count, min, sem_agg, sum +from .aggregates import arg_min, array_agg, avg, collect_list, count, min, sem_agg, sum from .expressions import case_when, least, try_cast from .logical import ColumnSpec, MemorySpec, MemoryView, QueryExpr, UserQuery from .relation import Log, Relation, SearchRelation @@ -26,6 +26,7 @@ "SearchRelation", "UserQuery", "array_agg", + "arg_min", "avg", "case_when", "collect_list", diff --git a/src/agent_memory/policy/aggregates.py b/src/agent_memory/policy/aggregates.py index 357dc85..9d175df 100644 --- a/src/agent_memory/policy/aggregates.py +++ b/src/agent_memory/policy/aggregates.py @@ -18,6 +18,28 @@ class SemanticAggregateSpec: instruction: str +@dataclass(frozen=True) +class ArgMinAggregateSpec: + """Copy one row's payload and ordering key, retaining sufficient merge state.""" + + order_by: tuple[str, ...] + columns: tuple[str, ...] + + +def arg_min(*, order_by: Sequence[str], columns: Sequence[str]) -> ArgMinAggregateSpec: + """Select the lexicographically earliest row; retain keys for incremental merge. + + Keys must be non-null and comparable. Tied keys must have identical payloads; + ambiguous ties fail rather than silently depending on DataFrame row order. + """ + keys, values = tuple(order_by), tuple(columns) + if not keys or not values or any(not isinstance(c, str) or not c for c in (*keys, *values)): + raise ValueError("arg_min requires nonempty ordering and payload column names") + if len(set((*keys, *values))) != len(keys) + len(values): + raise ValueError("arg_min ordering and payload columns must be distinct") + return ArgMinAggregateSpec(keys, values) + + @dataclass(frozen=True) class ArrayAggregateSpec: """Serializable descriptor for one grouped array aggregate.""" @@ -70,6 +92,7 @@ class AvgAggregateSpec: ) AggregateSpec: TypeAlias = ( SemanticAggregateSpec + | ArgMinAggregateSpec | ArrayAggregateSpec | CollectListAggregateSpec | MinAggregateSpec @@ -185,6 +208,8 @@ def normalize_min_columns( def aggregate_output_names(spec: AggregateSpec) -> tuple[str, ...]: """Return output column names produced by one aggregate descriptor.""" + if isinstance(spec, ArgMinAggregateSpec): + return (*spec.order_by, *spec.columns) if isinstance(spec, SemanticAggregateSpec): return tuple(column.name for column in spec.output_cols) if isinstance( @@ -215,6 +240,7 @@ def normalize_aggregate_specs(specs: Sequence[AggregateSpec]) -> tuple[Aggregate spec, ( SemanticAggregateSpec, + ArgMinAggregateSpec, ArrayAggregateSpec, CollectListAggregateSpec, MinAggregateSpec, @@ -226,7 +252,8 @@ def normalize_aggregate_specs(specs: Sequence[AggregateSpec]) -> tuple[Aggregate ] if invalid: raise TypeError( - "agg accepts only agent_memory.sem_agg(...), agent_memory.array_agg(...), " + "agg accepts only agent_memory.sem_agg(...), agent_memory.arg_min(...), " + "agent_memory.array_agg(...), " "agent_memory.collect_list(...), agent_memory.min(...), " "agent_memory.count(...), agent_memory.sum(...), or " "agent_memory.avg(...) " diff --git a/src/agent_memory/policy/schema.py b/src/agent_memory/policy/schema.py index bd6dc47..9146043 100644 --- a/src/agent_memory/policy/schema.py +++ b/src/agent_memory/policy/schema.py @@ -5,6 +5,7 @@ from collections.abc import Sequence from .aggregates import ( + ArgMinAggregateSpec, ArrayAggregateSpec, AvgAggregateSpec, CollectListAggregateSpec, @@ -328,7 +329,9 @@ def _aggregate_output_columns(aggregates: object) -> tuple[str, ...]: outputs: list[str] = [] for aggregate in aggregates: # type: ignore[union-attr] - if isinstance(aggregate, SemanticAggregateSpec): + if isinstance(aggregate, ArgMinAggregateSpec): + names = (*aggregate.order_by, *aggregate.columns) + elif isinstance(aggregate, SemanticAggregateSpec): names = tuple(column.name for column in aggregate.output_cols) elif isinstance( aggregate, diff --git a/src/agent_memory/runtime/executor.py b/src/agent_memory/runtime/executor.py index 1bf7023..2c8a48e 100644 --- a/src/agent_memory/runtime/executor.py +++ b/src/agent_memory/runtime/executor.py @@ -3,10 +3,13 @@ from __future__ import annotations from collections import Counter, defaultdict, deque -from collections.abc import Callable, Hashable, Mapping +from collections.abc import Callable, Hashable, Iterator, Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from contextvars import copy_context from copy import deepcopy from dataclasses import dataclass from datetime import date, datetime, timezone +from heapq import merge from typing import Any from uuid import uuid4 @@ -15,6 +18,7 @@ from agent_memory.api import Message, MessageInput from agent_memory.planner.differential_policy import DifferentiatedPolicy from agent_memory.planner.retrieval import RetrievalPlan +from agent_memory.planner.physical import PREDICATE_DECISIONS_INPUT from agent_memory.policy.logical import MemoryView, QueryExpr, UserQuery from agent_memory.policy.retrieval import RetrievalResult from agent_memory.policy.relation import ( @@ -25,6 +29,7 @@ ) from agent_memory.policy.schema import output_columns from agent_memory.runtime.window import WINDOW_SOURCE_INPUT, completed_count_windows +from agent_memory.runtime.row_outputs import RowOutputCache from agent_memory.storage.connector import ( StorageCommit, StorageConflictError, @@ -59,6 +64,81 @@ def between(cls, old: pd.DataFrame, new: pd.DataFrame) -> "NodeOutputUpdate": ) +@dataclass +class _CandidateBucketIndex: + """Derived membership for joins that retain survivors and append new IDs.""" + + columns: tuple[str, ...] + buckets: dict[tuple[Any, ...], dict[Hashable, int]] + next_order: int + row_count: int + + @staticmethod + def _keys(frame: pd.DataFrame, columns: tuple[str, ...]) -> Iterator[tuple[Any, ...]]: + # Input is read-only during iteration. Retain the values as well as their + # IDs to prevent Python ID reuse; bound this per-call memo to 64 entries. + recent: dict[tuple[int, ...], tuple[tuple[Any, ...], tuple[Any, ...]]] = {} + for row in frame.loc[:, list(columns)].itertuples(index=False, name=None): + identity = tuple(id(value) for value in row) + entry = recent.get(identity) + if entry is None: + key = _row_key(row) + if len(recent) == 64: + recent.clear() + recent[identity] = (row, key) + else: + key = entry[1] + yield key + + @classmethod + def build(cls, frame: pd.DataFrame, columns: tuple[str, ...]) -> "_CandidateBucketIndex": + buckets: dict[tuple[Any, ...], dict[Hashable, int]] = {} + for order, (occurrence, key) in enumerate(zip(frame.index, cls._keys(frame, columns), strict=True)): + buckets.setdefault(key, {})[occurrence] = order + return cls(columns, buckets, len(frame), len(frame)) + + def updated(self, removed: pd.DataFrame, added: pd.DataFrame) -> "_CandidateBucketIndex": + if len(removed) == self.row_count: + return type(self).build(added, self.columns) + buckets = dict(self.buckets) + touched: set[tuple[Any, ...]] = set() + next_order = self.next_order + for frame, inserting in ((removed, False), (added, True)): + previous_key = None + members: dict[Hashable, int] = {} + for occurrence, key in zip(frame.index, self._keys(frame, self.columns), strict=True): + # Consecutive members of one bucket need only one lookup of the + # potentially nested key. The per-call key memo retains it. + if key is not previous_key: + if key not in touched: + buckets[key] = dict(buckets.get(key, {})) + touched.add(key) + members = buckets[key] + previous_key = key + if inserting: + members[occurrence] = next_order + next_order += 1 + else: + del members[occurrence] + for key in touched: + if not buckets[key]: + del buckets[key] + return type(self)(self.columns, buckets, next_order, self.row_count - len(removed) + len(added)) + + def select(self, frame: pd.DataFrame, keys: set[tuple[Any, ...]]) -> pd.DataFrame: + if len(keys) >= len(self.buckets) and self.buckets.keys() <= keys: + return frame.iloc[range(len(frame))] + # Each bucket is already ordered. Merge only the affected memberships. + streams = [self.buckets[key].items() for key in keys if key in self.buckets] + occurrences = [occurrence for occurrence, _ in merge(*streams, key=lambda item: item[1])] + return frame.loc[occurrences] if occurrences else frame.iloc[:0] + + +_BucketUpdates = dict[ + tuple[str, tuple[str, ...]], tuple[_CandidateBucketIndex, _CandidateBucketIndex] +] + + class PolicyExecutor: """Execute one differentiated policy step and commit all states atomically.""" @@ -69,18 +149,28 @@ def __init__( adapter: Any | None = None, storage: StorageDeployment | None = None, ) -> None: - self.policy = policy - self.spec = policy.spec self.adapter = adapter if adapter is not None else self._default_adapter() + prepare = getattr(self.adapter, "prepare_policy", None) + self.policy = policy if prepare is None else prepare(policy) + policy = self.policy + self.spec = policy.spec + self._independent_node_id = getattr( + self.adapter, "independent_node_id", lambda policy: None, + )(policy) + if self._independent_node_id is not None: + selected = policy.nodes[self._independent_node_id] + if selected.execution_kind != "semantic_row" or selected.query.op != "sem_flat_map": + raise ValueError("independent execution requires a row-local sem_flat_map") self.storage = storage self._validate_storage_plan() self._state: dict[str, pd.DataFrame] = {} self._node_state: dict[str, pd.DataFrame] = {} - self._semantic_output_cache: dict[ - str, - dict[Hashable, pd.DataFrame], - ] = {} + self._semantic_output_cache: dict[str, RowOutputCache] = {} self._window_next_start: dict[str, int] = {} + self._predicate_decisions: dict[str, dict[str, bool]] = {} + # Derived from committed rows; rebuild lazily after restore, not in snapshots. + self._join_row_indexes: dict[str, dict[tuple[Any, ...], tuple[Hashable, ...]]] = {} + self._candidate_bucket_indexes: dict[tuple[str, tuple[str, ...]], _CandidateBucketIndex] = {} self._next_occurrence = 0 self._storage_commit = ( None @@ -129,11 +219,23 @@ def add(self, message: MessageInput) -> None: """Append one source row and propagate its change through the shared DAG.""" row = self.normalize_source_row(message) - self.apply_delta(pd.DataFrame([row], columns=self.source_columns)) + self.apply_delta(pd.DataFrame([row], columns=pd.Index(self.source_columns))) def apply_delta(self, changed_rows: pd.DataFrame) -> None: """Propagate one relation-valued source delta and publish it atomically.""" + if self._independent_node_id is None: + self._apply_delta(changed_rows) + else: + # Always drain the worker before returning, including on main-branch failure. + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="memory-extraction") as pool: + self._apply_delta(changed_rows, pool=pool) + + def _apply_delta( + self, changed_rows: pd.DataFrame, *, pool: ThreadPoolExecutor | None = None, + ) -> None: + """Keep state mutation and publication on the calling thread.""" + if not isinstance(changed_rows, pd.DataFrame): raise TypeError("changed_rows must be a pandas DataFrame") if changed_rows.empty: @@ -171,14 +273,37 @@ def allocate(node_id: str) -> str: staged_node_state = dict(self._node_state) staged_cache = { - node_id: dict(lineage) + node_id: lineage.copy() for node_id, lineage in self._semantic_output_cache.items() } staged_window_start = dict(self._window_next_start) + staged_predicates = {node: dict(decisions) for node, decisions in self._predicate_decisions.items()} + staged_join_indexes = dict(self._join_row_indexes) + staged_bucket_updates: _BucketUpdates = {} updates: dict[str, NodeOutputUpdate] = {} - for node_id in self.policy.execution_order: + remaining = list(self.policy.execution_order) + pending: Future[Any] | None = None + while remaining: + if pending is not None and pending.done(): + pending.result() # Surface extraction failure before starting more work. + node_id = next((candidate for candidate in remaining + if not (pending is not None and candidate == self._independent_node_id) + and all(parent in updates for parent in self.policy.nodes[candidate].input_node_ids)), None) + if node_id is None: + if pending is None or self._independent_node_id not in remaining: + raise RuntimeError("maintenance plan has unresolved dependencies") + node_id = self._independent_node_id node = self.policy.nodes[node_id] + if node_id == self._independent_node_id and pending is None and pool is not None: + inserted = updates[node.input_node_ids[0]].inserted_rows + if not inserted.empty: + pending = pool.submit( + copy_context().run, self.adapter.execute_independent, node.query, + {node.input_node_ids[0]: _copy_frame_with_objects(inserted)}, + ) + continue + remaining.remove(node_id) if node.execution_kind == "source": staged_node_state[node_id] = next_log updates[node_id] = NodeOutputUpdate( @@ -200,7 +325,22 @@ def allocate(node_id: str) -> str: and not needs_global_aggregate_identity ): staged_node_state.setdefault(node_id, old_state) - updates[node_id] = NodeOutputUpdate.between(old_state, old_state.copy()) + updates[node_id] = NodeOutputUpdate( + old_state, old_state.iloc[:0], old_state.iloc[:0], + ) + continue + + if node.execution_kind == "relational_state": + row_index = self._join_row_indexes.get(node_id) + row_index = _index_join_rows(old_state) if row_index is None else dict(row_index) + update = self._execute_relational_join_update( + node_id, old_state, parent_updates, + allocate=lambda node_id=node_id: allocate(node_id), + row_index=row_index, + ) + staged_join_indexes[node_id] = row_index + staged_node_state[node_id] = update.output_rows + updates[node_id] = update continue if node.execution_kind == "deterministic": @@ -211,22 +351,25 @@ def allocate(node_id: str) -> str: old_state=old_state, parent_update=parent_updates[0], ) - elif node.execution_kind in { - "relational_state", - "semantic_binary_state", - }: + elif node.execution_kind == "semantic_binary_state": next_state = self._execute_binary_state( node_id, old_state=old_state, parent_updates=parent_updates, staged_node_state=staged_node_state, ) - elif node.execution_kind == "semantic_row": + elif node.execution_kind in {"semantic_row", "semantic_predicate"}: + completed = pending.result() if pending is not None and node_id == self._independent_node_id else None next_state = self._execute_semantic_row( node_id, old_state=old_state, parent_update=parent_updates[0], staged_cache=staged_cache, + staged_bucket_updates=staged_bucket_updates, + predicate_decisions=(staged_predicates.setdefault(node_id, {}) + if node.execution_kind == "semantic_predicate" else None), + execute=(lambda query, inputs: completed) + if pending is not None and node_id == self._independent_node_id else None, ) elif node.execution_kind == "semantic_state": next_state = self._execute_semantic_state( @@ -234,6 +377,9 @@ def allocate(node_id: str) -> str: old_state=old_state, parent_update=parent_updates[0], staged_node_state=staged_node_state, + identity_decisions=(staged_predicates.setdefault(node_id, {}) + if node.maintenance_query is not None + and node.maintenance_query.params.get("identity_reuse") else None), ) elif node.execution_kind == "process_window": next_state = self._execute_process_window( @@ -256,14 +402,16 @@ def allocate(node_id: str) -> str: ) next_state = self._align_output(node_id, next_state) - if node.execution_kind != "semantic_row": - next_state = preserve_occurrence_index( + if node.execution_kind not in {"semantic_row", "semantic_predicate"}: + update = _preserve_occurrences_and_diff( old_state, next_state, allocate=lambda node_id=node_id: allocate(node_id), ) - staged_node_state[node_id] = next_state - updates[node_id] = NodeOutputUpdate.between(old_state, next_state) + else: + update = NodeOutputUpdate.between(old_state, next_state) + staged_node_state[node_id] = update.output_rows + updates[node_id] = update next_public_state: dict[str, pd.DataFrame] = { "log": next_log.reset_index(drop=True).copy() @@ -282,6 +430,9 @@ def allocate(node_id: str) -> str: self._state = next_public_state self._node_state = staged_node_state self._semantic_output_cache = staged_cache + self._predicate_decisions = staged_predicates + self._join_row_indexes = staged_join_indexes + self._candidate_bucket_indexes.update({key: pair[1] for key, pair in staged_bucket_updates.items()}) self._window_next_start = staged_window_start self._next_occurrence = staged_next_occurrence self._storage_commit = next_storage_commit @@ -306,6 +457,9 @@ def snapshot_state(self) -> dict[str, Any]: snapshot["adapter_execution_fingerprint"] = ( adapter_execution_fingerprint ) + if any(n.execution_kind == "semantic_predicate" or n.maintenance_query is not None + and n.maintenance_query.params.get("identity_reuse") for n in self.policy.nodes.values()): + snapshot["predicate_decisions"] = {node: dict(values) for node, values in self._predicate_decisions.items()} if self.storage is not None: commit = self._require_storage_commit() if commit.is_initial: @@ -371,12 +525,23 @@ def restore_state(self, snapshot: Mapping[str, Any]) -> None: staged_state = dict(state) staged_node_state = dict(node_state) staged_semantic_output_cache = { - str(node_id): dict(_require_nested_mapping(cache, "semantic output cache")) + str(node_id): RowOutputCache(_require_nested_mapping(cache, "semantic output cache")) for node_id, cache in semantic_output_cache.items() } staged_window_next_start = { str(node_id): int(value) for node_id, value in window_next_start.items() } + reuse_nodes = {key for key, node in self.policy.nodes.items() + if node.execution_kind == "semantic_predicate" or node.maintenance_query is not None + and node.maintenance_query.params.get("identity_reuse")} + raw_predicates = _require_mapping(snapshot, "predicate_decisions") if reuse_nodes else {} + staged_predicates: dict[str, dict[str, bool]] = {} + for node_id, values in raw_predicates.items(): + if node_id not in reuse_nodes or not isinstance(values, Mapping) or any( + not isinstance(key, str) or type(value) is not bool for key, value in values.items() + ): + raise ValueError("invalid predicate decision state") + staged_predicates[str(node_id)] = dict(values) staged_storage_commit: StorageCommit | None = None if self.storage is not None: @@ -416,6 +581,9 @@ def restore_state(self, snapshot: Mapping[str, Any]) -> None: self._state = staged_state self._node_state = staged_node_state self._semantic_output_cache = staged_semantic_output_cache + self._predicate_decisions = staged_predicates + self._join_row_indexes = {} + self._candidate_bucket_indexes = {} self._window_next_start = staged_window_next_start self._next_occurrence = next_occurrence self._storage_commit = staged_storage_commit @@ -681,11 +849,9 @@ def _execute_binary_state( if len(node.input_node_ids) != 2 or len(parent_updates) != 2: raise RuntimeError("Binary state nodes require exactly two inputs") if any(not update.retracted_rows.empty for update in parent_updates): - if node.execution_kind == "semantic_binary_state": - raise NotImplementedError( - "Semantic binary state maintenance requires append-only input changes" - ) - return self._execute_deterministic(node_id, staged_node_state) + raise NotImplementedError( + "Semantic binary state maintenance requires append-only input changes" + ) if node.maintenance_query is None: raise RuntimeError(f"Binary state node {node_id} has no maintenance query") @@ -702,6 +868,68 @@ def _execute_binary_state( inputs[f"{parent_id}__inserted"] = update.inserted_rows return self.adapter.execute(node.maintenance_query, inputs) + def _execute_relational_join_update( + self, node_id: str, old_state: pd.DataFrame, + updates: tuple[NodeOutputUpdate, ...], + *, allocate: Callable[[], Hashable], + row_index: dict[tuple[Any, ...], tuple[Hashable, ...]] | None = None, + ) -> NodeOutputUpdate: + """Propagate exact inner-join changes without re-diffing its full output.""" + node = self.policy.nodes[node_id] + left, right = updates + old_left, old_right = ( + self._node_state.get(p, self._empty_node_frame(p)) + for p in node.input_node_ids + ) + surviving_left = ( + _multiset_difference(left.retracted_rows, old_left)[0] + if not left.retracted_rows.empty else old_left + ) + + def bind(template: QueryExpr, name: str, frame: pd.DataFrame) -> QueryExpr: + if template.op == "alias": + return QueryExpr(op="alias", params=template.params, + inputs=(bind(template.inputs[0], name, frame),)) + return QueryExpr(op="materialized_view", params={ + "name": name, "columns": tuple(frame.columns), + }) + + def join(a: pd.DataFrame, b: pd.DataFrame) -> pd.DataFrame: + if a.empty or b.empty: + return old_state.iloc[:0].copy() + query = QueryExpr(op="join", params=node.query.params, inputs=( + bind(node.query.inputs[0], "__join_left", a), + bind(node.query.inputs[1], "__join_right", b), + )) + return self._align_output( + node_id, + self.adapter.execute(query, {"__join_left": a, "__join_right": b}), + ) + + # Remove each old match once, then add each new match once. Using the + # surviving left side avoids double-counting simultaneous side changes. + def concat(parts: list[pd.DataFrame]) -> pd.DataFrame: + nonempty = [part for part in parts if not part.empty] + return pd.concat(nonempty, ignore_index=True) if nonempty else old_state.iloc[:0] + + removed = concat([ + join(left.retracted_rows, old_right), + join(surviving_left, right.retracted_rows), + ]) + if left.retracted_rows.empty and right.retracted_rows.empty: + # Preserve the compiled append plan's task order as well as its bag. + added = concat([ + join(left.inserted_rows, old_right), + join(old_left, right.inserted_rows), + join(left.inserted_rows, right.inserted_rows), + ]) + else: + added = concat([ + join(left.inserted_rows, right.output_rows), + join(surviving_left, right.inserted_rows), + ]) + return _apply_join_delta(old_state, removed, added, allocate=allocate, row_index=row_index) + def _execute_algebraic_state( self, node_id: str, @@ -734,20 +962,64 @@ def _execute_semantic_row( *, old_state: pd.DataFrame, parent_update: NodeOutputUpdate, - staged_cache: dict[str, dict[Hashable, pd.DataFrame]], + staged_cache: dict[str, RowOutputCache], + staged_bucket_updates: _BucketUpdates | None = None, + predicate_decisions: dict[str, bool] | None = None, + execute: Callable[[QueryExpr, Mapping[str, Any]], Any] | None = None, ) -> pd.DataFrame: node = self.policy.nodes[node_id] if len(node.input_node_ids) != 1: raise NotImplementedError("Row-local semantic nodes require one input") - cache = staged_cache.setdefault(node_id, {}) - for occurrence in set(parent_update.retracted_rows.index): + cache = staged_cache.get(node_id) + if not isinstance(cache, RowOutputCache): + cache = RowOutputCache(cache) + staged_cache[node_id] = cache + scope = getattr(self.adapter, "bounded_predicate_scope", lambda query: None)(node.query) + retracted = parent_update.retracted_rows + inserted = parent_update.inserted_rows + if scope is not None: + dependencies, bucket_columns = scope + parent_id = node.input_node_ids[0] + previous = self._node_state.get(parent_id) + if previous is None: + previous = parent_update.output_rows.iloc[:0] + bucket_pair = None + if (bucket_columns and staged_bucket_updates is not None + and self.policy.nodes[parent_id].execution_kind == "relational_state"): + key = (parent_id, tuple(bucket_columns)) + bucket_pair = staged_bucket_updates.get(key) + if bucket_pair is None: + before = self._candidate_bucket_indexes.get(key) + if before is None: + before = _CandidateBucketIndex.build(previous, key[1]) + bucket_pair = (before, before.updated(retracted, inserted)) + staged_bucket_updates[key] = bucket_pair + if _carry_predicate_replacements(cache, retracted, inserted, dependencies): + frames = cache.nonempty_frames() + return pd.concat(frames) if frames else old_state.iloc[:0].copy() + if bucket_columns: + retracted, inserted = _carry_unchanged_predicate_buckets( + cache, retracted, inserted, dependencies, bucket_columns, + ) + if bucket_pair is None: + retracted, inserted = _affected_candidate_buckets( + previous, parent_update.output_rows, retracted, inserted, bucket_columns, + ) + else: + keys = {_row_key(row) for frame in (retracted, inserted) + for row in frame.loc[:, list(bucket_columns)].itertuples(index=False, name=None)} + retracted = bucket_pair[0].select(previous, keys) + inserted = bucket_pair[1].select(parent_update.output_rows, keys) + for occurrence in set(retracted.index): cache.pop(occurrence, None) - inserted = parent_update.inserted_rows if not inserted.empty: parent_id = node.input_node_ids[0] - outputs = self.adapter.execute(node.query, {parent_id: inserted}) + inputs: dict[str, Any] = {parent_id: inserted} + if predicate_decisions is not None: + inputs[PREDICATE_DECISIONS_INPUT] = predicate_decisions + outputs = (execute or self.adapter.execute)(node.query, inputs) outputs = self._align_output(node_id, outputs) unknown_indexes = set(outputs.index).difference(inserted.index) if unknown_indexes: @@ -755,14 +1027,14 @@ def _execute_semantic_row( f"{node.query.op} did not preserve semantic input indexes: " f"{sorted(unknown_indexes, key=repr)}" ) - empty_output = outputs.iloc[0:0].copy() + empty_output = cache.empty_like(outputs) for occurrence in dict.fromkeys(inserted.index): mask = outputs.index == occurrence cache[occurrence] = ( - outputs.loc[mask].copy() if mask.any() else empty_output.copy() + outputs.loc[mask].copy() if mask.any() else empty_output ) - frames = [frame for frame in cache.values() if not frame.empty] + frames = cache.nonempty_frames() if not frames: return old_state.iloc[0:0].copy() return pd.concat(frames, axis=0) @@ -774,6 +1046,7 @@ def _execute_semantic_state( old_state: pd.DataFrame, parent_update: NodeOutputUpdate, staged_node_state: Mapping[str, pd.DataFrame], + identity_decisions: dict[str, bool] | None = None, ) -> pd.DataFrame: node = self.policy.nodes[node_id] parent_id = node.input_node_ids[0] @@ -784,9 +1057,11 @@ def _execute_semantic_state( ) if node.maintenance_query is None: raise RuntimeError(f"Semantic state node {node_id} has no maintenance query") - inputs = dict(staged_node_state) + inputs: dict[str, Any] = dict(staged_node_state) inputs[node_id] = old_state inputs[f"{parent_id}__inserted"] = parent_update.inserted_rows + if identity_decisions is not None: + inputs["__physical_identity_decisions"] = identity_decisions return self.adapter.execute(node.maintenance_query, inputs) def _execute_process_window( @@ -882,10 +1157,10 @@ def _align_output(self, node_id: str, frame: pd.DataFrame) -> pd.DataFrame: return frame.loc[:, expected].copy() def _empty_node_frame(self, node_id: str) -> pd.DataFrame: - return pd.DataFrame(columns=self.policy.nodes[node_id].output_columns) + return pd.DataFrame(columns=pd.Index(self.policy.nodes[node_id].output_columns)) def _empty_view_frame(self, view: MemoryView) -> pd.DataFrame: - return pd.DataFrame(columns=self.query_output_columns(view.query)) + return pd.DataFrame(columns=pd.Index(self.query_output_columns(view.query))) def normalize_source_row( self, @@ -984,6 +1259,148 @@ def _default_adapter(self) -> Any: return LotusAdapter() +def _carry_predicate_replacements( + cache: RowOutputCache, removed: pd.DataFrame, + added: pd.DataFrame, dependencies: tuple[str, ...], +) -> bool: + """Carry membership, not old metadata, when semantic input bags are unchanged.""" + if len(removed) != len(added) or any(i not in cache for i in removed.index): + return False + old_keys = [_row_key(row) for row in removed.loc[:, list(dependencies)].itertuples(index=False, name=None)] + new_keys = [_row_key(row) for row in added.loc[:, list(dependencies)].itertuples(index=False, name=None)] + if Counter(old_keys) != Counter(new_keys): + return False + kept: dict[tuple[Any, ...], deque[bool]] = defaultdict(deque) + for key, occurrence in zip(old_keys, removed.index, strict=True): + kept[key].append(not cache[occurrence].empty) + for occurrence in removed.index: + cache.pop(occurrence) + empty = cache.empty_like(added) + for position, (key, occurrence) in enumerate(zip(new_keys, added.index, strict=True)): + cache[occurrence] = added.iloc[[position]].copy() if kept[key].popleft() else empty + return True + + +def _carry_unchanged_predicate_buckets( + cache: RowOutputCache, removed: pd.DataFrame, + added: pd.DataFrame, dependencies: tuple[str, ...], columns: tuple[str, ...], +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Carry metadata-only buckets even when other buckets have semantic changes.""" + def buckets(frame: pd.DataFrame) -> dict[tuple[Any, ...], list[int]]: + result: dict[tuple[Any, ...], list[int]] = defaultdict(list) + for position, row in enumerate(frame.loc[:, list(columns)].itertuples(index=False, name=None)): + result[_row_key(row)].append(position) + return result + + old_buckets, new_buckets = buckets(removed), buckets(added) + carried_old: set[int] = set() + carried_new: set[int] = set() + for key, old_positions in old_buckets.items(): + new_positions = new_buckets.get(key, []) + if _carry_predicate_replacements( + cache, removed.iloc[old_positions], added.iloc[new_positions], dependencies, + ): + carried_old.update(old_positions) + carried_new.update(new_positions) + return ( + removed.iloc[[i for i in range(len(removed)) if i not in carried_old]], + added.iloc[[i for i in range(len(added)) if i not in carried_new]], + ) + + +def _affected_candidate_buckets( + previous: pd.DataFrame, current: pd.DataFrame, removed: pd.DataFrame, + added: pd.DataFrame, columns: tuple[str, ...], +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Refresh complete changed Top-k buckets, including displaced surviving rows.""" + if not columns: + # Symmetric Top-k is the union of both endpoints' choices. A local subset + # cannot safely determine membership, so use the full candidate domain. + return previous, current + keys = {_row_key(row) for frame in (removed, added) + for row in frame.loc[:, list(columns)].itertuples(index=False, name=None)} + + def select(frame: pd.DataFrame) -> pd.DataFrame: + positions = [i for i, row in enumerate(frame.loc[:, list(columns)].itertuples(index=False, name=None)) + if _row_key(row) in keys] + return frame.iloc[positions] + + return select(previous), select(current) + + +def _index_join_rows(frame: pd.DataFrame) -> dict[tuple[Any, ...], tuple[Hashable, ...]]: + """Index bag occurrences once; subsequent commits update only changed keys.""" + if not frame.index.is_unique: + raise ValueError("join state requires unique occurrence IDs") + rows: dict[tuple[Any, ...], list[Hashable]] = defaultdict(list) + for occurrence, row in zip(frame.index, frame.itertuples(index=False, name=None), strict=True): + rows[_row_key(row)].append(occurrence) + return {key: tuple(ids) for key, ids in rows.items()} + + +def _apply_join_delta( + old: pd.DataFrame, + removed: pd.DataFrame, + added: pd.DataFrame, + *, + allocate: Callable[[], Hashable], + row_index: dict[tuple[Any, ...], tuple[Hashable, ...]] | None = None, +) -> NodeOutputUpdate: + """Cancel changed bags, preserving old occurrence IDs and indexing new rows.""" + + if any(list(frame.columns) != list(old.columns) for frame in (removed, added)): + raise ValueError("Join delta columns must match the materialized output") + # A replacement can remove and recreate the same output. Only its net + # change may propagate to downstream semantic operators. + if not removed.empty: + added, removed = _multiset_difference(removed, added) + positions: list[int] = [] + if not removed.empty: + remaining = Counter( + _row_key(row) for row in removed.itertuples(index=False, name=None) + ) + if row_index is not None: + occurrences: list[Hashable] = [] + for key, number in remaining.items(): + matches = row_index.get(key, ()) + if len(matches) < number: + raise RuntimeError("join retractions exceed materialized output multiplicity") + occurrences.extend(matches[-number:]) + if len(matches) == number: + del row_index[key] + else: + row_index[key] = matches[:-number] + positions = old.index.get_indexer(occurrences).tolist() + if -1 in positions: + raise RuntimeError("join occurrence index does not match materialized state") + remaining.clear() + else: + # Keep the standalone helper's scan path; committed executors supply + # an index. Both paths retain the earliest equal bag occurrences. + for reverse_position, row in enumerate(old.iloc[::-1].itertuples(index=False, name=None)): + key = _row_key(row) + if remaining.get(key, 0): + positions.append(len(old) - 1 - reverse_position) + remaining[key] -= 1 + if remaining[key] == 0: + del remaining[key] + if not remaining: + break + if remaining: + raise RuntimeError("join retractions exceed materialized output multiplicity") + retracted = old.iloc[sorted(positions)] + retained = old.drop(index=retracted.index) if positions else old + inserted = added.copy() + inserted.index = pd.Index([allocate() for _ in range(len(inserted))], dtype="object") + if row_index is not None: + for key, new_occurrences in _index_join_rows(inserted).items(): + row_index[key] = row_index.get(key, ()) + new_occurrences + output = retained if inserted.empty else ( + inserted if retained.empty else pd.concat([retained, inserted]) + ) + return NodeOutputUpdate(output, inserted, retracted) + + def _copy_frame_with_objects(frame: pd.DataFrame) -> pd.DataFrame: """Detach caller-owned nested values at public input/output boundaries.""" @@ -1025,6 +1442,38 @@ def _multiset_difference( ) +def _preserve_occurrences_and_diff( + old: pd.DataFrame, + new: pd.DataFrame, + *, + allocate: Callable[[], Hashable], +) -> NodeOutputUpdate: + """Pair equal occurrences once, retaining the existing bag and ID order.""" + + if list(old.columns) != list(new.columns): + raise ValueError("Relation states must have identical ordered columns") + available: dict[tuple[Any, ...], deque[int]] = defaultdict(deque) + old_indexes: list[Hashable] = list(old.index) + for position, row in enumerate(old.itertuples(index=False, name=None)): + available[_row_key(row)].append(position) + retained = bytearray(len(old)) + indexes: list[Hashable] = [] + inserted: list[int] = [] + for position, row in enumerate(new.itertuples(index=False, name=None)): + matches = available.get(_row_key(row)) + if matches: + old_position = matches.popleft() + retained[old_position] = 1 + indexes.append(old_indexes[old_position]) + else: + indexes.append(allocate()) + inserted.append(position) + result = new.copy() + result.index = pd.Index(indexes, dtype="object") + retracted = [position for position, keep in enumerate(retained) if not keep] + return NodeOutputUpdate(result, result.iloc[inserted].copy(), old.iloc[retracted].copy()) + + def preserve_occurrence_index( old: pd.DataFrame, new: pd.DataFrame, diff --git a/src/agent_memory/runtime/row_outputs.py b/src/agent_memory/runtime/row_outputs.py new file mode 100644 index 0000000..8293c02 --- /dev/null +++ b/src/agent_memory/runtime/row_outputs.py @@ -0,0 +1,89 @@ +"""Private, transaction-local indexing of semantic occurrence outputs.""" + +from collections.abc import Hashable, Iterator, Mapping, MutableMapping + +import pandas as pd + + +class RowOutputCache(MutableMapping[Hashable, pd.DataFrame]): + """Retain empty membership cheaply and enumerate only nonempty outputs. + + Stored frames are read-only by convention, just like executor node state. + Updates replace frames; forks share frames but never mutable bookkeeping. + Snapshots serialize the mapping, not this derived index. + """ + + def __init__(self, frames: Mapping[Hashable, pd.DataFrame] | None = None) -> None: + self._frames: dict[Hashable, pd.DataFrame] = {} + self._positions: dict[Hashable, int] = {} + self._nonempty: dict[Hashable, pd.DataFrame] = {} + self._next_position = 0 + self._empty_templates: list[pd.DataFrame] = [] + if frames is not None: + for key, frame in frames.items(): + self[key] = frame + + def __getitem__(self, key: Hashable) -> pd.DataFrame: + return self._frames[key] + + def __setitem__(self, key: Hashable, frame: pd.DataFrame) -> None: + empty = frame.empty + if empty: + frame = self._intern_empty(frame) + if key not in self._frames: + self._positions[key] = self._next_position + self._next_position += 1 + self._frames[key] = frame + if empty: + self._nonempty.pop(key, None) + else: + self._nonempty[key] = frame + + def __delitem__(self, key: Hashable) -> None: + del self._frames[key] + del self._positions[key] + self._nonempty.pop(key, None) + + def __iter__(self) -> Iterator[Hashable]: + return iter(self._frames) + + def __len__(self) -> int: + return len(self._frames) + + def copy(self) -> "RowOutputCache": + """Fork bookkeeping without copying shared, unmodified output frames.""" + result = RowOutputCache() + result._frames = self._frames.copy() + result._positions = self._positions.copy() + result._nonempty = self._nonempty.copy() + result._next_position = self._next_position + result._empty_templates = self._empty_templates.copy() + return result + + def empty_like(self, frame: pd.DataFrame) -> pd.DataFrame: + """Get a read-only template once before recording a batch of empties.""" + return self._intern_empty(frame.iloc[:0]) + + def nonempty_frames(self) -> list[pd.DataFrame]: + """Return outputs in mapping insertion order, without visiting empties.""" + return [self._nonempty[key] for key in sorted(self._nonempty, key=self._positions.__getitem__)] + + def _intern_empty(self, frame: pd.DataFrame) -> pd.DataFrame: + for template in self._empty_templates: + if frame is template: + return template + # Arbitrary attrs may contain mutable/non-comparable values. Keep those + # frames independent rather than silently dropping metadata. + if frame.attrs: + return frame.copy(deep=True) + for template in self._empty_templates: + if ( + frame.columns.identical(template.columns) + and frame.index.identical(template.index) + and frame.dtypes.equals(template.dtypes) + and frame.flags.allows_duplicate_labels == template.flags.allows_duplicate_labels + ): + return template + template = frame.copy(deep=True) + self._empty_templates.append(template) + return template diff --git a/src/agent_memory/runtime/runtime.py b/src/agent_memory/runtime/runtime.py index 76ee256..3925fe9 100644 --- a/src/agent_memory/runtime/runtime.py +++ b/src/agent_memory/runtime/runtime.py @@ -37,6 +37,7 @@ def __init__( adapter=adapter, storage=storage, ) + self.policy = self._engine.policy @property def _state(self) -> dict[str, Any]: @@ -124,6 +125,8 @@ def restore_state(self, snapshot: Mapping[str, Any]) -> None: "count refresh contract requires a schema-v3 runtime checkpoint" ) if schema_version == 1: + if not getattr(self._engine.adapter, "supports_legacy_restore", True): + raise ValueError("physical fusion cannot restore a legacy schema-v1 checkpoint") if self.storage is not None: raise NotImplementedError( "schema-v1 checkpoints cannot be restored with storage" diff --git a/tests/fixtures/zep_contradiction_contrasts.json b/tests/fixtures/zep_contradiction_contrasts.json new file mode 100644 index 0000000..b7fd4ff --- /dev/null +++ b/tests/fixtures/zep_contradiction_contrasts.json @@ -0,0 +1,35 @@ +{ + "kind": "synthetic-controlled-contrasts", + "groups": [ + {"topic": "residence", "new_fact": "As of today, Alice has moved her sole current primary residence from London to Paris; she no longer lives in London.", "candidates": [ + {"fact": "Alice currently has her primary residence in London.", "expected": true, "reason": "The new fact explicitly supersedes the same person's current residence."}, + {"fact": "Bob currently lives in London.", "expected": false, "reason": "Different person."}, + {"fact": "Alice lived in London during 2019.", "expected": false, "reason": "A past residence remains true after a move."}, + {"fact": "Alice enjoys visiting museums.", "expected": false, "reason": "Independent preference."}]}, + {"topic": "occupation", "new_fact": "Mira has left her job as a teacher at Elm School. Her only current job is now a nurse at City Hospital.", "candidates": [ + {"fact": "Mira worked as a teacher at Elm School in 2018.", "expected": false, "reason": "Historical employment is preserved."}, + {"fact": "Mira currently works as a teacher at Elm School.", "expected": true, "reason": "Explicitly left that current job."}, + {"fact": "Leon currently teaches at Elm School.", "expected": false, "reason": "Different employee."}, + {"fact": "Mira holds a teaching qualification.", "expected": false, "reason": "Changing jobs does not revoke qualifications."}]}, + {"topic": "preference", "new_fact": "Nora says her current favorite beverage has changed: it is tea, not coffee; she no longer considers coffee her favorite.", "candidates": [ + {"fact": "Nora's favorite dessert is cake.", "expected": false, "reason": "Different preference attribute."}, + {"fact": "Nora preferred coffee to tea in 2020.", "expected": false, "reason": "Past preference is not contradicted."}, + {"fact": "Nora's current favorite beverage is coffee.", "expected": true, "reason": "Explicit replacement of current favorite."}, + {"fact": "Omar's favorite beverage is coffee.", "expected": false, "reason": "Different person."}]}, + {"topic": "plan", "new_fact": "Iris has canceled her planned Kyoto trip for October 2027 and will not take that trip; she will instead stay home that month.", "candidates": [ + {"fact": "Iris visited Kyoto in April 2024.", "expected": false, "reason": "Past completed trip."}, + {"fact": "Iris plans to visit Rome in December 2027.", "expected": false, "reason": "Different trip and time."}, + {"fact": "Kai plans to visit Kyoto in October 2027.", "expected": false, "reason": "Different traveler."}, + {"fact": "Iris currently plans to visit Kyoto in October 2027.", "expected": true, "reason": "Same planned trip explicitly canceled."}]}, + {"topic": "count correction", "new_fact": "The final attendance count for the 2025 Harbor Workshop is corrected to 18 people. The earlier total of 12 was a counting error, not a different session.", "candidates": [ + {"fact": "The final attendance at the 2025 Harbor Workshop was 12 people.", "expected": true, "reason": "Explicit correction for the identical event and measure."}, + {"fact": "The 2024 Harbor Workshop had 12 attendees.", "expected": false, "reason": "Different year's event."}, + {"fact": "The 2025 Harbor Workshop offered 12 presentations.", "expected": false, "reason": "Different counted attribute."}, + {"fact": "The 2025 Hill Workshop had 12 attendees.", "expected": false, "reason": "Different event."}]}, + {"topic": "relationship status", "new_fact": "Lena and Sam have divorced. They are no longer married to each other.", "candidates": [ + {"fact": "Lena and Sam married each other in 2016.", "expected": false, "reason": "Divorce does not invalidate the historical marriage event."}, + {"fact": "Lena and Sam are currently married to each other.", "expected": true, "reason": "Current marital relationship explicitly ended."}, + {"fact": "Lena and Sam have a daughter named Ava.", "expected": false, "reason": "Parenthood can persist after divorce."}, + {"fact": "Ruth and Joel are currently married to each other.", "expected": false, "reason": "Different couple."}]} + ] +} diff --git a/tests/test_agent_memory_benchmark_drivers.py b/tests/test_agent_memory_benchmark_drivers.py index 71a11e6..2928d19 100644 --- a/tests/test_agent_memory_benchmark_drivers.py +++ b/tests/test_agent_memory_benchmark_drivers.py @@ -140,6 +140,46 @@ def test_zep_event_mapping_is_benchmark_neutral() -> None: } +def test_zep_driver_keeps_caption_outside_scoring_metadata(tmp_path: Path) -> None: + event = replace( + _event(), + metadata={ + "blip_caption": " a cup with a dog face ", + "answer": "scoring-only answer", + "evidence": ["scoring-only evidence"], + }, + ) + memory = _Memory(None) + ZepMemoryDriver(memory, trace_dir=tmp_path).add(event) + + assert memory.rows == [{ + **event_to_zep_log_row(_event()), + "content": "user: I moved to Paris.\n" + "(description of attached image: a cup with a dog face)", + }] + + +@pytest.mark.parametrize("caption", [None, "", " ", 123, ["not text"]]) +def test_zep_empty_or_invalid_caption_preserves_plain_input(caption: object) -> None: + event = replace(_event(), metadata={"blip_caption": caption}) + assert event_to_zep_log_row(event) == event_to_zep_log_row(_event()) + + +def test_zep_caption_contract_cannot_reuse_old_input_identity() -> None: + from agent_memory.evaluation.harness import MemorySystemContract + from agent_memory.evaluation.run import _BUILT_IN_CONTRACTS + + input_id, retrieval_id = _BUILT_IN_CONTRACTS["zep-memory"] + current = MemorySystemContract( + system_id="zep-memory", memory_model_id="fake", + memory_provider_model_id="fake", + input_adapter_id=input_id, retrieval_recipe_id=retrieval_id, + ) + old = replace(current, input_adapter_id="benchmark-event-to-zep-log:v1") + assert current.input_adapter_digest != old.input_adapter_digest + assert current.maintenance_fingerprint != old.maintenance_fingerprint + + def test_zep_driver_preserves_native_entity_and_fact_channels(tmp_path) -> None: retrieval = RetrievalResult( query="Where do I live?", @@ -658,6 +698,89 @@ def run(self, bundle) -> None: assert captured["closed"] is True +@pytest.mark.parametrize( + ("strategy", "expected_version"), + [ + ("disabled", None), + ("zep-target-state", "zep-target-state-v1"), + ("zep-combined", "zep-combined-v2"), + ("zep-fact-summary", "zep-fact-summary-v2"), + ("zep-representative", "zep-representative-v2"), + ], +) +@pytest.mark.parametrize("site_options", [False, True]) +def test_run_records_actual_fusion_version( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + strategy: str, + expected_version: str | None, + site_options: bool, +) -> None: + import agent_memory.evaluation.run as run_module + + captured: dict[str, Any] = {} + + class FakeFactory: + def __init__(self, **kwargs: Any) -> None: + pass + + @classmethod + def from_environment(cls, **kwargs: Any) -> "FakeFactory": + return cls(**kwargs) + + def close(self) -> None: + pass + + def runtime_provenance(self) -> dict[str, Any]: + return {} + + class FakeRunner: + def __init__(self, **kwargs: Any) -> None: + captured.update(kwargs) + + def run(self, bundle: BenchmarkBundle) -> None: + pass + + monkeypatch.setattr(run_module, "_require_environment", lambda _system: None) + monkeypatch.setattr(run_module, "ZepMemoryDriverFactory", FakeFactory) + monkeypatch.setattr(run_module, "BenchmarkRunner", FakeRunner) + monkeypatch.setattr( + run_module, "collect_runtime_provenance", + lambda *args, **kwargs: {"source": {}, "runtime": {}}, + ) + monkeypatch.setattr(run_module, "validate_run_provenance", lambda *args, **kwargs: None) + bundle = BenchmarkBundle( + "locomo", + "revision", + "sha256", + ( + BenchmarkCase( + case_id="case-1", + task_id="locomo", + events=(_event(),), + questions=(BenchmarkQuestion("q1", "case-1", "?", "a", ()),), + ), + ), + {"run_mode": "integration-smoke"}, + ) + run_agent_memory_bundle( + bundle=bundle, + contracts={}, + system_id="zep-memory", + output_dir=tmp_path / "run", + memory_provider_model_id="provider-model", + grouped_agg_rule="rule-join-map", + physical_fusion=strategy, + structured_parse_retries=0 if site_options else None, + ) + execution = captured["runtime_provenance"]["runtime"]["lotus_execution"] + if expected_version is None: + assert "fusion_version" not in execution + else: + assert execution["fusion_version"] == expected_version + assert expected_version in captured["system_contract"].maintenance_execution_id + + def test_run_records_and_propagates_count_refresh_contract( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/tests/test_aggregate_only_batching.py b/tests/test_aggregate_only_batching.py new file mode 100644 index 0000000..984e7e4 --- /dev/null +++ b/tests/test_aggregate_only_batching.py @@ -0,0 +1,99 @@ +"""Aggregate-only packing composes with fusion without global batching.""" +from dataclasses import replace +from collections.abc import Mapping +import json +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.sem_agg import execute_structured_sem_agg_groups +from agent_memory.policy.logical import ColumnSpec, QueryExpr +from agent_memory.policy.aggregates import ArrayAggregateSpec, SemanticAggregateSpec +from agent_memory.adapters.lotus.relational import execute_agg + + +class Model: + max_tokens = 512 + max_ctx_len = 32768 + + def __init__(self, invalid: bool = False) -> None: + self.calls: list[Any] = [] + self.invalid = invalid + + def count_tokens(self, value: Any) -> int: + return 1 + + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + outputs = [] + for prompt in prompts: + data = json.loads(prompt[1]["content"]) + self.calls.append(data) + outputs.append(json.dumps({"results": [ + {"group_id": g["group_id"], "output": {"summary": str(g["documents"])}} + for g in reversed(data["groups"][1:] if self.invalid else data["groups"]) + ]})) + return SimpleNamespace(outputs=outputs) + + +def config(size: int = 4) -> LotusExecutionConfig: + return LotusExecutionConfig(physical_fusion="zep-target-state", + sem_agg_prompt_batching=PromptBatching(max_tasks=size), structured_parse_retries=0) + + +@pytest.mark.parametrize("invalid", [False, True]) +def test_groups_are_packed_and_restored_without_global_batching(monkeypatch: pytest.MonkeyPatch, invalid: bool) -> None: + import lotus + model = Model(invalid) + monkeypatch.setattr(lotus.settings, "lm", model) + q = QueryExpr(op="sem_agg", params={"instruction": "Summarize {body}.", + "input_cols": ("body",), "output_cols": (ColumnSpec("summary"),)}) + groups = [pd.DataFrame({"body": [f"value-{i}"], "provenance": ["MUST_NOT_SEND"]}) for i in range(5)] + def execute(frames: list[pd.DataFrame]) -> list[Mapping[str, Any]]: + return execute_structured_sem_agg_groups(q, frames, ("body",), (ColumnSpec("summary"),), config()) + if invalid: + with pytest.raises(ValueError): + execute(groups) + else: + result = execute(groups) + assert all(f"value-{i}" in row["summary"] for i, row in enumerate(result)) + assert [len(c["groups"]) for c in model.calls] == [4, 1] + assert "MUST_NOT_SEND" not in str(model.calls) + assert execute([]) == [] + assert len(model.calls) == 2 + assert config().prompt_batching is None + + +def test_configuration_conflicts_and_execution_identity() -> None: + c = config() + for overrides in ({"prompt_batching": PromptBatching(max_tasks=4)}, + {"sem_agg_dispatch": "provider-batched"}, {"sem_agg_safe_mode": True}): + with pytest.raises(ValueError): + replace(c, **overrides) + identities = {LotusAdapter(config=x).maintenance_execution_fingerprint + for x in (LotusExecutionConfig(), replace(c, sem_agg_prompt_batching=None), c, config(2))} + assert len(identities) == 4 + assert not LotusAdapter(config=c).supports_legacy_restore + + +def test_mixed_aggregate_dispatches_packing_and_preserves_provenance(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + model = Model() + monkeypatch.setattr(lotus.settings, "lm", model) + source = pd.DataFrame({"key": ["a", "a", "b"], "body": ["one", "two", "three"], "id": [1, 2, 3]}) + source.attrs["agent_memory_groupby_keys"] = ("key",) + q = QueryExpr(op="agg", inputs=(QueryExpr(op="group_by", params={"keys": ("key",)}),), + params={"aggregates": ( + SemanticAggregateSpec(input_cols=("body",), output_cols=(ColumnSpec("summary"),), instruction="Summarize {body}."), + ArrayAggregateSpec(columns=("id",), output_col="provenance"), + )}) + result = execute_agg(q, {}, lambda _q, _inputs: source, SimpleNamespace(config=config())) + assert [len(c["groups"]) for c in model.calls] == [2] + assert result.key.tolist() == ["a", "b"] + assert [json.loads(v) for v in result.provenance] == [[{"id": 1}, {"id": 2}], [{"id": 3}]] + assert "one" in result.summary.iloc[0] and "two" in result.summary.iloc[0] + assert "three" in result.summary.iloc[1] diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py index 67dd1fc..dda27ac 100644 --- a/tests/test_benchmark_harness.py +++ b/tests/test_benchmark_harness.py @@ -141,7 +141,6 @@ def test_checkpoint_restore_rejects_changed_framework_cache_mode( ('{"answer":"Speyer"}', "Speyer"), ('```json\n{"answer":"Speyer"}\n```', "Speyer"), ("Speyer", "Speyer"), - ('{"wrong_field":"schema echo"}', '{"wrong_field":"schema echo"}'), ], ) def test_locomo_answer_parser_accepts_structured_or_raw_text( @@ -165,22 +164,20 @@ def test_locomo_answer_parser_rejects_empty_and_fingerprints_contract() -> None: assert contract.fingerprint != previous.fingerprint -def test_locomo_schema_echo_is_scored_as_zero_without_special_case() -> None: +@pytest.mark.parametrize("response", [ + '{"type":"object","properties":{"answer":{"type":"string"}}}', + '{"wrong_field":"schema echo"}', + '{"answer":null}', + '{"answer":42}', + '{"answer":" "}', + '{"answer":', + '[{"answer":"Speyer"}]', + '```json\n{"wrong_field":"echo"}\n```', +]) +def test_locomo_invalid_structured_answer_is_not_scored(response: str) -> None: contract = locomo_task_contract() - question = BenchmarkQuestion( - question_id="q167", - sample_id="case", - question="What is not in memory?", - gold_answer="No information available", - evidence_event_ids=(), - category="5", - ) - schema_echo = '{"type":"object","properties":{"answer":{"type":"string"}}}' - answer = contract.answer_parser(schema_echo) - assert contract.deterministic_scorer is not None - - assert answer == schema_echo - assert contract.deterministic_scorer(question, answer).score == 0.0 + with pytest.raises(ValueError, match="LOCOMO structured answer"): + contract.answer_parser(response) @dataclass @@ -916,6 +913,27 @@ def provenance(mode: str) -> dict[str, object]: changed.run(_bundle()) +def test_resume_accepts_json_safe_contract_values(tmp_path: Path) -> None: + provenance = { + "source": {"commit": "source", "dirty": False}, + "runtime": { + "python": "3.12", + "lotus_execution": {"predicate_reuse_sites": ("site-a",)}, + }, + } + + for _ in range(2): + BenchmarkRunner( + system_contract=_system_contract(), + contracts={"task-1": _contract()}, + driver_factory=lambda case_id, state_dir, trace_dir: _Driver(state_dir), + answer_model=_Model(), + judge_model=_Model(), + artifacts=BenchmarkArtifactStore(tmp_path), + runtime_provenance=provenance, + ).run(_bundle()) + + def test_retrieval_system_error_scores_zero_and_completed_case_is_skipped( tmp_path: Path, ) -> None: @@ -1290,7 +1308,8 @@ def add(self, event: BenchmarkEvent) -> dict[str, int]: runner.run(_bundle()) -def test_invalid_answer_retries_only_current_model_step(tmp_path) -> None: +@pytest.mark.parametrize("invalid_response", ["", '{"properties":{"answer":{"type":"string"}}}']) +def test_invalid_answer_retries_only_current_model_step(tmp_path, invalid_response) -> None: drivers: list[_Driver] = [] def factory(case_id, state_dir, trace_dir): @@ -1300,10 +1319,10 @@ def factory(case_id, state_dir, trace_dir): drivers.append(driver) return driver - model = _Model(["", "one", "two"]) + model = _Model([invalid_response, "one", "two"]) runner = BenchmarkRunner( system_contract=_system_contract(), - contracts={"task-1": _contract()}, + contracts={"task-1": replace(_contract(), answer_parser=locomo_task_contract().answer_parser)}, driver_factory=factory, answer_model=model, judge_model=model, diff --git a/tests/test_benchmark_trace_metrics.py b/tests/test_benchmark_trace_metrics.py index a9a120b..003b941 100644 --- a/tests/test_benchmark_trace_metrics.py +++ b/tests/test_benchmark_trace_metrics.py @@ -3,6 +3,8 @@ import json from pathlib import Path +import pytest + from agent_memory.evaluation.trace_metrics import ( normalize_framework_cache_usage, normalize_provider_calls, @@ -11,6 +13,80 @@ ) +def _cached_lotus_call(**changes: object) -> dict: + return { + "event_type": "llm_call", "operator_call_id": "op-1", + "llm_item_index": 0, "llm_batch_size": 1, + "latency_sec": 0.001, "usage_scope": "batch", + "usage_physical_prompt_tokens": 0, + "usage_physical_completion_tokens": 0, + "usage_physical_total_tokens": 0, "usage_cache_hits": 1, + "usage_virtual_prompt_tokens": 564, + "usage_virtual_completion_tokens": 86, + **changes, + } + + +@pytest.mark.parametrize("size", [1, 3]) +def test_cache_only_lotus_batch_is_not_a_provider_call(tmp_path: Path, size: int) -> None: + events = [_cached_lotus_call(llm_batch_size=size, llm_item_index=i, + usage_cache_hits=size) for i in range(size)] + assert normalize_provider_calls(events, output_dir=tmp_path) == [] + + +def test_cache_hit_does_not_take_later_provider_calls_latency(tmp_path: Path) -> None: + # One operator may first hit cache and then issue another request. + events = [_cached_lotus_call(), { + "event_type": "llm_call", "operator_call_id": "op-1", + "llm_item_index": 0, "latency_sec": 0.5, + }, { + "event_type": "provider_usage", "operator_call_id": "op-1", + "provider_item_index": 0, "provider_usage_available": True, + "provider_prompt_tokens": 10, "provider_completion_tokens": 2, + "provider_prompt_cache_hit_tokens": 0, "provider_prompt_cache_miss_tokens": 10, + }] + rows = normalize_provider_calls(events, output_dir=tmp_path) + assert len(rows) == 1 + assert rows[0]["latency_ms"] == 500 + assert summarize_provider_calls(rows)["estimated_cost_usd"] is not None + + +@pytest.mark.parametrize("changes", [ + {"usage_cache_hits": 0}, + {"usage_physical_total_tokens": None}, + {"llm_batch_size": 2}, # Only one of two items is known to be cached. + {"event_type": "llm_batch_error"}, +]) +def test_cache_counters_do_not_hide_unknown_or_failed_calls(tmp_path: Path, changes: dict) -> None: + rows = normalize_provider_calls([_cached_lotus_call(**changes)], output_dir=tmp_path) + assert len(rows) == 1 + summary = summarize_provider_calls(rows) + assert summary["usage_complete"] is False + assert summary["estimated_cost_usd"] is None + + +def test_partial_cache_batch_keeps_real_provider_response(tmp_path: Path) -> None: + events = [_cached_lotus_call(llm_batch_size=2, usage_physical_prompt_tokens=10, + usage_physical_completion_tokens=2, usage_physical_total_tokens=12), { + "event_type": "provider_usage", "operator_call_id": "op-1", + "provider_item_index": 0, "provider_usage_available": True, + "provider_prompt_tokens": 10, "provider_completion_tokens": 2, + "provider_prompt_cache_hit_tokens": 10, "provider_prompt_cache_miss_tokens": 0, + }] + rows = normalize_provider_calls(events, output_dir=tmp_path) + assert len(rows) == 1 and rows[0]["source"] == "agent-provider" + assert rows[0]["total_tokens"] == 12 + + +def test_explicit_provider_event_is_never_removed_by_cache_metadata(tmp_path: Path) -> None: + events = [_cached_lotus_call(), {"event_type": "provider_usage", + "operator_call_id": "op-1", "provider_item_index": 0, + "provider_usage_available": False}] + rows = normalize_provider_calls(events, output_dir=tmp_path) + assert len(rows) == 1 and rows[0]["source"] == "agent-provider" + assert summarize_provider_calls(rows)["usage_complete"] is False + + def test_normalize_provider_calls_can_omit_pricing(tmp_path: Path) -> None: events = [{ "event_type": "provider_usage", "provider_usage_available": True, diff --git a/tests/test_execution_review_regressions.py b/tests/test_execution_review_regressions.py new file mode 100644 index 0000000..c617b24 --- /dev/null +++ b/tests/test_execution_review_regressions.py @@ -0,0 +1,191 @@ +"""Offline regressions for execution setup and experiment failure handling.""" + +from concurrent.futures import ThreadPoolExecutor +import json +from pathlib import Path +import runpy +import sys +from threading import Event, Lock +from typing import Any + +import pandas as pd +import pytest + +from agent_memory.adapters.lotus.context import LotusExecutionConfig, LotusExecutionContext +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.sem_filter_batch_prompting import execute_batch_prompted_predicate + + +def test_concurrent_configuration_installs_one_shared_wrapper(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + + context = LotusExecutionContext("deepseek/deepseek-flash", LotusExecutionConfig( + physical_fusion="zep-representative", parallel_fact_extraction=True, lm_enable_cache=False, + )) + first_started, release_first, second_started, duplicate_build = (Event() for _ in range(4)) + counter_lock = Lock() + models: list[object] = [] + + def new_lm(self: LotusExecutionContext) -> object: + model = object() + with counter_lock: + models.append(model) + first = len(models) == 1 + if first: + first_started.set() + assert release_first.wait(5) + else: + duplicate_build.set() + return model + + def second_configure() -> None: + second_started.set() + context.configure() + + monkeypatch.setattr(LotusExecutionContext, "new_lm", new_lm) + monkeypatch.setattr(lotus.settings, "lm", None) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(context.configure) + assert first_started.wait(5) + second = pool.submit(second_configure) + assert second_started.wait(5) + duplicate_build.wait(0.2) + release_first.set() + first.result(timeout=5) + second.result(timeout=5) + assert len(models) == 1 + assert lotus.settings.lm is context._scoped_lm + child = context.fork() + assert child._scoped_lm is lotus.settings.lm + assert child._lm is not context._lm + + +def test_configuration_failure_does_not_publish_partial_state(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + + context = LotusExecutionContext("deepseek/deepseek-flash", LotusExecutionConfig( + physical_fusion="zep-representative", parallel_fact_extraction=True, lm_enable_cache=False, + )) + monkeypatch.setattr(LotusExecutionContext, "new_lm", lambda self: object()) + attempts: list[dict[str, Any]] = [] + + def configure(**kwargs: Any) -> None: + attempts.append(kwargs) + if len(attempts) == 1: + raise RuntimeError("configuration failed") + + monkeypatch.setattr(lotus.settings, "configure", configure) + with pytest.raises(RuntimeError, match="configuration failed"): + context.configure() + assert context._lm is None and context._scoped_lm is None + assert not context._configured + context.configure() + assert context._configured and context._scoped_lm is attempts[-1]["lm"] + + +@pytest.mark.parametrize("finish_reason", ["stop", "length"]) +def test_groupby_only_batching_preserves_completion_metadata( + monkeypatch: pytest.MonkeyPatch, finish_reason: str, +) -> None: + import lotus + import lotus.models + from litellm import ModelResponse + from lotus.models import LM + + class FakeLM(LM): + def _process_uncached_messages(self, uncached_data: Any, all_kwargs: Any, + show_progress_bar: bool, progress_bar_desc: str) -> Any: + return [ModelResponse(model="deepseek/deepseek-flash", choices=[{ + "index": 0, "message": {"role": "assistant", "content": json.dumps({ + "decisions": [{"row_id": "row_0", "keep": True}], + })}, "finish_reason": finish_reason, + }], usage={"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}) + for _ in uncached_data] + + monkeypatch.setattr(lotus.models, "LM", FakeLM) + monkeypatch.setattr(lotus.settings, "lm", None) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + batching = PromptBatching(max_tasks=16) + context = LotusExecutionContext("deepseek/deepseek-flash", LotusExecutionConfig( + groupby_prompt_batching={"sem_groupby:test": batching}, + structured_parse_retries=0, lm_enable_cache=False, + )) + context.configure() + + def execute() -> Any: + return execute_batch_prompted_predicate( + pd.DataFrame({"left": ["a"], "right": ["a"]}), + instruction="Do {left} and {right} match?", prompt_batching=batching, + structured_parse_retries=0, structured_max_tokens=128, + progress_bar_desc="Offline grouping", operator="sem_groupby", + ) + + if finish_reason == "length": + with pytest.raises(ValueError, match="incomplete structured response"): + execute() + else: + assert execute().decisions == (True,) + + +@pytest.mark.parametrize("state", ["missing", "empty", "preflighted"]) +def test_followup_cli_precondition_failure_preserves_status( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, state: str, +) -> None: + import tools.zep_frozen_followup as probe + + output = tmp_path / "run" + if state != "missing": + output.mkdir() + status = output / "status.json" + if state == "preflighted": + status.write_text(json.dumps({"stage": "preflight_passed"})) + monkeypatch.setattr(sys, "argv", [probe.__file__, "run", "--output", str(output)]) + with pytest.raises(FileNotFoundError): + runpy.run_path(str(probe.__file__), run_name="__main__") + if state == "preflighted": + assert json.loads(status.read_text()) == {"stage": "preflight_passed"} + else: + assert not status.exists() + assert output.exists() == (state == "empty") + + +@pytest.mark.parametrize("stage", ["completed", "answer_replay", "contradictions", "failed"]) +def test_followup_cli_rejected_replay_preserves_status( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, stage: str, +) -> None: + import tools.zep_frozen_followup as probe + + status = tmp_path / "status.json" + original = json.dumps({"stage": stage, "details": "original execution"}, indent=2) + status.write_text(original) + monkeypatch.setattr(sys, "argv", [probe.__file__, "run", "--output", str(tmp_path)]) + with pytest.raises(ValueError, match="refusing replay"): + runpy.run_path(str(probe.__file__), run_name="__main__") + assert status.read_text() == original + + +def test_followup_execution_failure_records_status( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + import tools.zep_frozen_followup as probe + + probe.save(tmp_path / "status.json", {"stage": "preflight_passed"}) + probe.save(tmp_path / "frozen.json", []) + probe.save(tmp_path / "answer-inputs.json", {}) + probe.save(tmp_path / "manifest.json", { + "input_hashes": {}, "script_sha256": probe.digest(Path(probe.__file__)), + "frozen_sha256": probe.digest(tmp_path / "frozen.json"), + "answer-inputs_sha256": probe.digest(tmp_path / "answer-inputs.json"), + }) + + def fail_replay(output: Path) -> None: + assert json.loads((output / "status.json").read_text())["stage"] == "answer_replay" + raise RuntimeError("offline injected failure") + + monkeypatch.setattr(probe, "replay_answers", fail_replay) + with pytest.raises(RuntimeError, match="offline injected failure"): + probe.run(tmp_path) + assert json.loads((tmp_path / "status.json").read_text()) == { + "stage": "failed", "error": "offline injected failure", "type": "RuntimeError", + } diff --git a/tests/test_independent_execution.py b/tests/test_independent_execution.py new file mode 100644 index 0000000..b9a1851 --- /dev/null +++ b/tests/test_independent_execution.py @@ -0,0 +1,249 @@ +"""Offline checks for one independent extraction branch and atomic publication.""" + +import json +from concurrent.futures import ThreadPoolExecutor +from contextvars import copy_context +from dataclasses import replace +from pathlib import Path +from threading import Barrier, Event, current_thread +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.memories.zep.representative import ZepRepresentativeMemory +from agent_memory.runtime.executor import PolicyExecutor + + +def test_registered_branch_has_no_entity_resolution_dependency() -> None: + adapter = LotusAdapter(config=LotusExecutionConfig( + physical_fusion="zep-representative", parallel_fact_extraction=True, + lm_enable_cache=False, + )) + policy = adapter.prepare_policy(ZepRepresentativeMemory.differentiate_policy()) + selected = adapter.independent_node_id(policy) + assert selected is not None + assert policy.nodes[selected].query.op == "sem_flat_map" + assert "entities" in policy.nodes[selected].query.params["input_cols"] + + +class ForkMemory(am.Memory): + log = am.Log({"text": "Input"}) + extracted = log.sem_flat_map( + input_cols=["text"], output_cols={"fact": "Fact"}, instruction="Extract {text} into {fact}", + ) + resolved = log.sem_map( + input_cols=["text"], output_cols={"entity": "Entity"}, instruction="Resolve {text} into {entity}", + ) + result = extracted.join(resolved.select(["text", "entity"]), on="text") + + +class ForkAdapter(LotusAdapter): + """Require both branches to be in flight without relying on sleep timing.""" + + def __init__(self, *, fail: str | None = None, parallel: bool = True) -> None: + super().__init__() + self.fail = fail + self.parallel = parallel + self.started = Event() + self.released = Event() + self.finished = Event() + self.calls: list[str] = [] + + def independent_node_id(self, policy: Any) -> str | None: + return next(n.node_id for n in policy.nodes.values() if n.query.op == "sem_flat_map") if self.parallel else None + + def execute_independent(self, query: Any, inputs: Any) -> Any: + self.calls.append("fact") + self.started.set() + try: + assert self.released.wait(5), "independent branch was not scheduled" + if self.fail == "fact": + raise ValueError("fact failed") + source = super().execute(query.inputs[0], inputs) + return source.assign(fact=source.text) + finally: + self.finished.set() + + def execute(self, query: Any, inputs: Any) -> Any: + if query.op == "sem_flat_map": + assert not self.parallel, "deferred extraction was executed twice" + self.calls.append("fact") + source = super().execute(query.inputs[0], inputs) + return source.assign(fact=source.text) + if query.op == "sem_map": + assert current_thread().name == "MainThread" + self.calls.append("entity") + if self.parallel: + assert self.started.wait(5) + self.released.set() + if self.fail == "entity": + raise ValueError("entity failed") + source = super().execute(query.inputs[0], inputs) + return source.assign(entity=source.text) + return super().execute(query, inputs) + + +def test_parallel_branch_matches_serial_and_restores() -> None: + adapter = ForkAdapter() + executor = PolicyExecutor(ForkMemory.differentiate_policy(), adapter=adapter) + serial = PolicyExecutor(ForkMemory.differentiate_policy(), adapter=ForkAdapter(parallel=False)) + for _ in range(2): + executor.add({"text": "same"}) + serial.add({"text": "same"}) + assert adapter.calls.count("fact") == 2 + assert adapter.calls.count("entity") == 2 + assert adapter.finished.is_set() + # Four join rows, not two: duplicate occurrences must survive. + result = executor.read_view("result") + assert len(result) == 4 + assert result[["text", "fact", "entity"]].equals(serial.read_view("result")[["text", "fact", "entity"]]) + restored = PolicyExecutor(ForkMemory.differentiate_policy(), adapter=ForkAdapter()) + restored.restore_state(executor.snapshot_state()) + pd.testing.assert_frame_equal(restored.read_view("result"), result) + restored.add({"text": "new"}) + assert len(restored.read_view("result")) == 5 + + +@pytest.mark.parametrize("branch", ["fact", "entity"]) +def test_either_failure_drains_worker_and_does_not_commit(branch: str) -> None: + adapter = ForkAdapter(fail=branch) + executor = PolicyExecutor(ForkMemory.differentiate_policy(), adapter=adapter) + before = executor.snapshot_state() + with pytest.raises(ValueError, match=f"{branch} failed"): + executor.add({"text": "x"}) + assert adapter.finished.is_set() + assert executor.source_row_count == 0 + assert executor.snapshot_state() == before + + +def test_config_and_restore_boundaries() -> None: + base = LotusExecutionConfig(physical_fusion="zep-representative", lm_enable_cache=False) + policy = ZepRepresentativeMemory.differentiate_policy() + serial = PolicyExecutor(policy, adapter=LotusAdapter(config=base)) + parallel = PolicyExecutor(policy, adapter=LotusAdapter(config=replace(base, parallel_fact_extraction=True))) + with pytest.raises(ValueError, match="fingerprint"): + parallel.restore_state(serial.snapshot_state()) + restored = PolicyExecutor(policy, adapter=LotusAdapter(config=replace(base, parallel_fact_extraction=True))) + restored.restore_state(parallel.snapshot_state()) + cached = PolicyExecutor(policy, adapter=LotusAdapter(config=replace(base, + parallel_fact_extraction=True, lm_enable_cache=True))) + with pytest.raises(ValueError, match="fingerprint"): + cached.restore_state(parallel.snapshot_state()) + for kwargs in ({"lm_enable_cache": None}, {"physical_fusion": "disabled"}): + with pytest.raises(ValueError, match="requires"): + replace(base, parallel_fact_extraction=True, **kwargs) + assert LotusAdapter(config=base).independent_node_id(policy) is None + with pytest.raises(ValueError, match="registered"): + parallel.adapter.independent_node_id(ForkMemory.differentiate_policy()) + selected = parallel._independent_node_id + assert selected is not None + parent_id = policy.nodes[selected].input_node_ids[0] + changed = dict(policy.nodes) + changed[parent_id] = replace(changed[parent_id], query=replace(changed[parent_id].query, params={"changed": True})) + with pytest.raises(ValueError, match="input contract"): + parallel.adapter.independent_node_id(replace(policy, nodes=changed)) + + +@pytest.mark.parametrize("cache_enabled", [False, True]) +def test_private_lm_keeps_global_model_usage_metadata_and_trace_isolated(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, cache_enabled: bool) -> None: + import lotus + import lotus.models + from litellm import ModelResponse + from lotus.models import LM + from agent_memory.tracing.semantic import semantic_trace_scope + + barrier = Barrier(2, timeout=5) + + class FakeLM(LM): + def _process_uncached_messages(self, uncached_data: Any, all_kwargs: Any, + show_progress_bar: bool, progress_bar_desc: str) -> Any: + if not uncached_data: + return [] + barrier.wait() + responses = [] + for prompt, _ in uncached_data: + word = "worker" if "worker-input" in str(prompt) else "main" + tokens = 17 if word == "worker" else 11 + responses.append(ModelResponse(model="deepseek/deepseek-flash", choices=[{ + "index": 0, "message": {"role": "assistant", "content": json.dumps({"rows": [{"fact": word}]})}, + "finish_reason": "stop", + }], usage={"prompt_tokens": tokens, "completion_tokens": 3, "total_tokens": tokens + 3})) + return responses + + monkeypatch.setattr(lotus.models, "LM", FakeLM) + # Restore global settings after the test, including the prior model. + monkeypatch.setattr(lotus.settings, "lm", None) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + adapter = LotusAdapter(model="deepseek/deepseek-flash", config=LotusExecutionConfig( + physical_fusion="zep-representative", parallel_fact_extraction=True, + lm_enable_cache=cache_enabled, semantic_trace_dir=tmp_path, structured_parse_retries=0, + )) + adapter._context.configure() + global_lm = lotus.settings.lm + assert global_lm is not None + query = am.Source({"text": "Input"}).sem_flat_map(input_cols=["text"], + output_cols={"fact": "Fact"}, instruction="Extract {text} into {fact}").expr + with semantic_trace_scope(event_id="event-1", phase="insertion"): + with ThreadPoolExecutor(max_workers=1) as pool: + future = pool.submit(copy_context().run, adapter.execute_independent, query, + {"log": pd.DataFrame({"text": ["worker-input"]})}) + main = adapter.execute(query, {"log": pd.DataFrame({"text": ["main-input"]})}) + worker = future.result() + assert main.fact.tolist() == ["main"] and worker.fact.tolist() == ["worker"] + assert lotus.settings.lm is global_lm + assert adapter._independent_context is not None + worker_lm = adapter._independent_context._lm + assert worker_lm is not None + assert worker_lm is not global_lm + assert global_lm.stats.physical_usage.total_tokens == 14 + assert worker_lm.stats.physical_usage.total_tokens == 20 + if cache_enabled: + # Exact operator reuse bypasses the LM. Metadata-only differences miss + # the operator cache but preserve the prompt and hit the native LM cache. + for frame in (pd.DataFrame({"text": ["worker-input"]}), + pd.DataFrame({"text": ["worker-input"], "metadata": [1]})): + adapter.execute_independent(query, {"log": frame}) + assert worker_lm.stats.operator_cache_hits == 1 + assert worker_lm.stats.cache_hits == 1 + assert worker_lm.stats.physical_usage.total_tokens == 20 + assert worker_lm.stats.virtual_usage.total_tokens == 60 + assert global_lm.stats.operator_cache_hits == 0 + assert global_lm.stats.cache_hits == 0 + assert global_lm.stats.virtual_usage.total_tokens == 14 + events = [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()] + calls = [e for e in events if e["event_type"] == "llm_call"] + real_calls = [c for c in calls if c["usage_physical_total_tokens"]] + assert len(real_calls) == 2 and {c["event_id"] for c in real_calls} == {"event-1"} + assert sorted(c["usage_physical_total_tokens"] for c in real_calls) == [14, 20] + assert len({c["operator_call_id"] for c in real_calls}) == 2 + assert sum("execution_lane" in c for c in real_calls) == 1 + usage = [e for e in events if e["event_type"] == "provider_usage"] + assert sorted(e["provider_total_tokens"] for e in usage) == [14, 20] + assert all(e["provider_finish_reason"] == "stop" for e in usage) + independent = [e for e in events if e["event_type"] == "independent_execution_result"] + assert len(independent) == (3 if cache_enabled else 1) and independent[0]["trace_bytes_written"] > 0 + if cache_enabled: + cached = [e for e in events if e["event_type"] == "framework_cache_usage"] + assert sum(e["physical_total_tokens"] for e in cached) == 34 + assert sum(e["virtual_total_tokens"] for e in cached) == 74 + assert sum(e["lm_cache_hits"] for e in cached) == 1 + assert sum(e["operator_cache_hits"] for e in cached) == 1 + assert all("execution_lane" in e for e in cached if e["lm_cache_hits"] or e["operator_cache_hits"]) + + +def test_scoped_model_resets_after_failure() -> None: + from types import SimpleNamespace + from agent_memory.adapters.lotus.scoped_lm import ScopedLM + + main = SimpleNamespace(cache="main") + worker = SimpleNamespace(cache="worker") + model = ScopedLM(main) + with pytest.raises(ValueError, match="failed"): + with model.scope(worker): + assert model.cache == "worker" + raise ValueError("failed") + assert model.cache == "main" diff --git a/tests/test_maintenance_work.py b/tests/test_maintenance_work.py new file mode 100644 index 0000000..0ba712a --- /dev/null +++ b/tests/test_maintenance_work.py @@ -0,0 +1,227 @@ +"""Regression checks for physical maintenance work, without a provider.""" + +from types import SimpleNamespace +from typing import Any +import json + +import pandas as pd +import pytest + +from agent_memory.adapters import LotusAdapter +from agent_memory.policy.logical import QueryExpr +from agent_memory.runtime.executor import NodeOutputUpdate, PolicyExecutor + + +@pytest.mark.parametrize("same_side", [False, True]) +def test_inner_join_replacements_never_recompute_old_pairs(same_side: bool) -> None: + left = pd.DataFrame({"k": [1, 1, 2], "v": ["a", "a", "b"]}) + right = pd.DataFrame({"k": [1, 2], "w": ["x", "y"]}) + new_left = pd.DataFrame({"k": [1, 2, 3], "v": ["a", "B", "c"]}) + new_right = right if same_side else pd.DataFrame({"k": [1, 1, 3], "w": ["x", "z", "q"]}) + leaves = tuple(QueryExpr(op="materialized_view", params={"name": n, "columns": tuple(f.columns)}) + for n, f in [("l", left), ("r", right)]) + query = QueryExpr(op="join", inputs=leaves, params={"on": ["k"], "how": "inner"}) + adapter = LotusAdapter() + old = adapter.execute(query, {"l": left, "r": right}) + expected = adapter.execute(query, {"l": new_left, "r": new_right}) + node = SimpleNamespace(input_node_ids=("l", "r"), execution_kind="relational_state", query=query) + executor: Any = SimpleNamespace(policy=SimpleNamespace(nodes={"j": node}), adapter=adapter, + _node_state={"l": left, "r": right}) + executor._execute_deterministic = lambda *args: pytest.fail("full join fallback") + executor._empty_node_frame = lambda name: executor._node_state[name].iloc[:0] + executor._align_output = lambda name, frame: frame.loc[:, old.columns] + from itertools import count + ids = count() + update = PolicyExecutor._execute_relational_join_update(executor, "j", old, + (NodeOutputUpdate.between(left, new_left), NodeOutputUpdate.between(right, new_right)), + allocate=lambda: f"new:{next(ids)}") + delta = NodeOutputUpdate.between(update.output_rows, expected) + assert delta.is_empty + + +def test_site_groups_use_one_provider_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: + from test_site_batching import adapter, query, rows + a, lm = adapter(monkeypatch) + batches: list[int] = [] + original = type(lm).__call__ + + def record(self: Any, prompts: Any, **kwargs: Any) -> Any: + batches.append(len(prompts)) + return original(self, prompts, **kwargs) + + monkeypatch.setattr(type(lm), "__call__", record) + result = a.execute(query(), {"pairs": rows()}) + pd.testing.assert_frame_equal(result, rows()) + assert batches == [3] + + +def test_predicate_reuse_snapshot_and_current_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + from dataclasses import replace + import agent_memory as am + from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id + from test_site_batching import adapter, query, rows + + a, lm = adapter(monkeypatch) + a.config = replace(a.config, predicate_reuse_sites=(semantic_pair_site_id(query()),)) + a._context.config = a.config + source = am.Source({c: c for c in rows().columns}) + view = source.sem_filter(instruction=str(query().params["instruction"])) + flow = am.SemanticDataflow(source=source, views={"result": view}, adapter=a) + first = rows().iloc[:1].copy() + flow.apply(first) + calls = len(lm.calls) + snapshot = flow.snapshot_state() + restored = am.SemanticDataflow(source=source, views={"result": view}, adapter=a) + restored.restore_state(snapshot) + changed = first.copy() + changed["occurrence"] = 99 + restored.apply(changed) + assert len(lm.calls) == calls + assert restored.view("result").occurrence.tolist() == [0, 99] + changed["fact:earlier_added"] = "different text" + restored.apply(changed) + assert len(lm.calls) == calls + 1 + disabled, _ = adapter(monkeypatch) + other = am.SemanticDataflow(source=source, views={"result": view}, adapter=disabled) + with pytest.raises(ValueError, match="fingerprint"): + other.restore_state(snapshot) + + +def test_metadata_replacement_in_real_runtime(monkeypatch: pytest.MonkeyPatch) -> None: + from dataclasses import replace + import agent_memory as am + from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id + from test_site_batching import adapter, query, rows + + a, lm = adapter(monkeypatch) + a.config = replace(a.config, predicate_reuse_sites=(semantic_pair_site_id(query()),)) + a._context.config = a.config + source = am.Source({c: c for c in rows().columns}) + grouped = source.group_by(["fact:earlier_added", "fact:later_added", "fact_id:later_added"]).array_agg( + columns=["occurrence"], output_col="provenance") + view = grouped.sem_filter(instruction=str(query().params["instruction"])) + flow = am.SemanticDataflow(source=source, views={"result": view}, adapter=a) + flow.apply(rows().iloc[:1]) + calls = len(lm.calls) + extra = rows().iloc[:1].copy() + extra["occurrence"] = 99 + flow.apply(extra) + assert len(lm.calls) == calls + assert json.loads(flow.view("result").iloc[0].provenance) == [{"occurrence": 0}, {"occurrence": 99}] + + +def test_inequality_self_join_replacement_scales_with_changes() -> None: + import agent_memory as am + + class Tracking(LotusAdapter): + def __init__(self) -> None: + super().__init__() + self.join_sizes: list[int] = [] + + def execute(self, query: QueryExpr, inputs: Any) -> Any: + result = super().execute(query, inputs) + if query.op == "join": + self.join_sizes.append(len(result)) + return result + + source = am.Source({"id": "id", "value": "number"}) + state = source.group_by("id").min(column="value", output_col="value") + left, right = state.alias("l"), state.alias("r") + view = left.join(right, on=left.col("id") < right.col("id")) + adapter = Tracking() + flow = am.SemanticDataflow(source=source, views={"result": view}, adapter=adapter) + data = pd.DataFrame({"id": list(range(40)), "value": [10] * 40}) + flow.apply(data) + adapter.join_sizes.clear() + delta = pd.DataFrame({"id": [0, 0, 39], "value": [5, 5, 3]}) + flow.apply(delta) + expected = LotusAdapter().execute(view.expr, {"log": pd.concat([data, delta], ignore_index=True)}) + assert NodeOutputUpdate.between(flow.view("result"), expected).is_empty + assert max(adapter.join_sizes) <= 2 * 40 + snapshot = flow.snapshot_state() + restored = am.SemanticDataflow(source=source, views={"result": view}, adapter=LotusAdapter()) + restored.restore_state(snapshot) + for value in [2, 1]: + event = pd.DataFrame({"id": [20], "value": [value]}) + restored.apply(event) + flow.apply(event) + assert NodeOutputUpdate.between(flow.view("result"), restored.view("result")).is_empty + + +def test_reuse_boolean_null_direction_duplicates_and_failure() -> None: + from agent_memory.adapters.lotus.predicate_reuse import reuse_predicate_decisions + + cache: dict[str, bool] = {} + calls: list[int] = [] + + def evaluate(frame: pd.DataFrame, positions: list[int]) -> pd.DataFrame: + calls.append(len(frame)) + return frame.loc[frame["a"] == "yes"].copy() + + data = pd.DataFrame({"a": ["yes", "yes", None, "no"], "b": ["no", "no", "x", "yes"], "meta": range(4)}) + result, _ = reuse_predicate_decisions(data, "Compare {a} to {b}", cache, evaluate) + assert calls == [3] and result.meta.tolist() == [0, 1] + data["meta"] += 10 + result, reused = reuse_predicate_decisions(data, "Compare {a} to {b}", cache, evaluate) + assert calls == [3] and reused == {0, 1, 2, 3} and result.meta.tolist() == [10, 11] + before = dict(cache) + + def fail(frame: pd.DataFrame, positions: list[int]) -> pd.DataFrame: + raise ValueError("provider failed") + + with pytest.raises(ValueError, match="provider failed"): + reuse_predicate_decisions(data, "Changed instruction {a} {b}", cache, fail) + assert cache == before + empty, _ = reuse_predicate_decisions(data.iloc[:0], "Compare {a} to {b}", cache, fail) + assert empty.empty + + +def test_combined_groupby32_aggregate4_reuse_prepares_without_model() -> None: + from agent_memory.adapters.lotus.context import LotusExecutionConfig + from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id + from agent_memory.adapters.lotus.prompt_batching import PromptBatching + from agent_memory.memories.zep.policy import ZepMemory + from test_site_batching import setting + + site = semantic_pair_site_id(ZepMemory._contradictory_fact_pairs.expr) + config = LotusExecutionConfig(physical_fusion="zep-target-state", + pair_filter_batching={site: setting()}, predicate_reuse_sites=(site,), + sem_groupby_pair_batch_size=32, sem_agg_prompt_batching=PromptBatching(max_tasks=4)) + plan = LotusAdapter(config=config).prepare_policy(ZepMemory.differentiate_policy()) + assert any(n.execution_kind == "semantic_predicate" for n in plan.nodes.values()) + assert any(n.maintenance_query is not None and n.maintenance_query.op == "fused_target_state" + for n in plan.nodes.values()) + + +def test_warm_join_retraction_does_not_key_unchanged_pair_rows(monkeypatch: pytest.MonkeyPatch) -> None: + import sys + import agent_memory as am + import agent_memory.runtime.executor as module + + source = am.Source({"id": "number", "value": "number"}) + state = source.group_by("id").min(column="value", output_col="value") + left, right = state.alias("l"), state.alias("r") + view = left.join(right, on=left.col("id") < right.col("id")) + flow = am.SemanticDataflow(source=source, views={"result": view}, adapter=LotusAdapter()) + initial = pd.DataFrame({"id": range(80), "value": [10] * 80}) + flow.apply(initial) + original = module._row_key + scanned = 0 + + def count_scan(row: Any) -> Any: + nonlocal scanned + if sys._getframe(1).f_code.co_name == "_apply_join_delta": + scanned += 1 + return original(row) + + monkeypatch.setattr(module, "_row_key", count_scan) + delta = pd.DataFrame({"id": [0], "value": [5]}) + flow.apply(delta) + assert scanned == 0 + expected = LotusAdapter().execute(view.expr, {"log": pd.concat([initial, delta], ignore_index=True)}) + assert NodeOutputUpdate.between(flow.view("result"), expected).is_empty + restored = am.SemanticDataflow(source=source, views={"result": view}, adapter=LotusAdapter()) + restored.restore_state(flow.snapshot_state()) + restored.apply(delta.assign(value=4)) + flow.apply(delta.assign(value=4)) + assert NodeOutputUpdate.between(flow.view("result"), restored.view("result")).is_empty diff --git a/tests/test_packed_pair_groups.py b/tests/test_packed_pair_groups.py new file mode 100644 index 0000000..f54531a --- /dev/null +++ b/tests/test_packed_pair_groups.py @@ -0,0 +1,66 @@ +"""Packed groups retain each pair's complete input and occurrence identity.""" +from dataclasses import replace +import json +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.site_batching import PairFilterBatching, _pack_partitions +from test_site_batching import adapter, query, rows + + +def test_pack_preserves_groups_and_oversized_partition() -> None: + assert _pack_partitions(((0,), (1, 2), (3, 4), (5,), tuple(range(6, 12))), 4) == ( + (0, 1, 2), (3, 4, 5), tuple(range(6, 12))) + assert _pack_partitions((), 4) == () + + +def test_packed_contexts_preserve_pair_decisions(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + a, _ = adapter(monkeypatch) + site = next(iter(a.config.pair_filter_batching)) + config = PairFilterBatching(("fact_id:later_added",), PromptBatching(max_tasks=4), + shared_columns=("fact_later_added",), pack_small_groups=True) + a._context.config = replace(a.config, pair_filter_batching={site: config}) + seen: list[dict[str, Any]] = [] + + class Oracle: + max_tokens = 512 + cache = None + + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + outputs = [] + for prompt in prompts: + payload = json.loads(prompt[1]["content"]) + seen.append(payload) + assert len(payload["rows"]) <= 4 + decisions = [] + for row in payload["rows"]: + common = payload.get("shared_context") or payload["contexts"][row["context_id"]] + assert "same old" in row["context"] + decisions.append({"row_id": row["row_id"], "keep": "new A" in common}) + outputs.append(json.dumps({"decisions": list(reversed(decisions))})) + return SimpleNamespace(outputs=outputs) + + monkeypatch.setattr(lotus.settings, "lm", Oracle()) + data = rows().iloc[:3] + result = a.execute(query(), {"pairs": data}) + pd.testing.assert_frame_equal(result, data.iloc[[0, 2]]) + assert len(seen) == 1 and len(seen[0]["contexts"]) == 2 + assert len(a.execute(query(), {"pairs": data.iloc[:0]})) == 0 + assert len(seen) == 1 + + +def test_packing_changes_execution_identity_and_default_stays_unchanged() -> None: + from agent_memory.adapters import LotusAdapter + from agent_memory.adapters.lotus.context import LotusExecutionConfig + original = PairFilterBatching(("id",), PromptBatching(max_tasks=16), shared_columns=("fact",)) + packed = replace(original, pack_small_groups=True) + assert original.to_dict()["version"] == "pair-filter-shared-context-v1" + assert "pack_small_groups" not in original.to_dict() + assert LotusAdapter(config=LotusExecutionConfig(pair_filter_batching={"sem_filter:test": original})).maintenance_execution_fingerprint != LotusAdapter(config=LotusExecutionConfig(pair_filter_batching={"sem_filter:test": packed})).maintenance_execution_fingerprint + with pytest.raises(TypeError): + replace(original, pack_small_groups=1) diff --git a/tests/test_physical_fusion.py b/tests/test_physical_fusion.py new file mode 100644 index 0000000..c29cf81 --- /dev/null +++ b/tests/test_physical_fusion.py @@ -0,0 +1,565 @@ +"""Offline target-state fusion tests using the real Zep maintenance contract.""" + +from dataclasses import replace +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +from agent_memory.adapters.lotus.adapter import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.fusion import parse_resolution +from agent_memory.adapters.lotus.pair_execution import ( + SemanticPairExecutionProfile, + PAIR_LEFT_ID_COLUMN, + PAIR_RIGHT_ID_COLUMN, + PAIR_LEFT_TEXT_COLUMN, + PAIR_RIGHT_TEXT_COLUMN, +) +from agent_memory.memories.zep.policy import ZepMemory +from agent_memory.planner.physical import optimize_policy +from agent_memory.policy.logical import QueryExpr +from agent_memory.tracing.semantic import query_digest + + +class FakeLM: + def __init__(self, answers: list[str]) -> None: + self.answers = list(answers) + self.calls: list[Any] = [] + + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + self.calls.append((prompts, kwargs)) + return SimpleNamespace(outputs=[self.answers.pop(0)]) + + +def fused_query() -> QueryExpr: + policy = optimize_policy( + ZepMemory.differentiate_policy(), strategy="zep-target-state" + ) + return next( + n.maintenance_query + for n in policy.nodes.values() + if n.maintenance_query is not None + and n.maintenance_query.op == "fused_target_state" + ) + + +def state(fact: str, seq: int, *, entity: str = "a") -> dict[str, Any]: + return { + "fact_id": (seq, 0), + "source_entity_id": entity, + "target_entity_id": "b", + "relation_type": "LIKES", + "fact": fact, + "valid_at": None, + "invalid_at": None, + "expired_at": None, + "provenance": json.dumps( + [ + { + "episode_id": str(seq), + "fact_ordinal": 0, + "created_at": "2026-01-01", + "content": fact, + } + ] + ), + "created_at": "2026-01-01", + "add_seq": seq, + } + + +def response( + matches: list[tuple[str, str | None]], states: list[tuple[str, str]] +) -> str: + return json.dumps( + { + "matches": [ + {"left_id": left, "right_id": right} for left, right in matches + ], + "states": [ + {"right_id": r, "relation_type": "LIKES", "fact": f} for r, f in states + ], + } + ) + + +def execute_fused( + monkeypatch: pytest.MonkeyPatch, + left: pd.DataFrame, + right: pd.DataFrame, + answers: list[str], + *, + pair_profile: SemanticPairExecutionProfile | None = None, + provider: Any = None, + **config: Any, +) -> tuple[pd.DataFrame, FakeLM]: + import lotus + + query = fused_query() + model = FakeLM(answers) + monkeypatch.setattr(lotus.settings, "lm", model) + joined = query.inputs[0] + materialized_left = QueryExpr( + op="materialized_view", + params={"name": "incoming", "columns": tuple(left.columns)}, + ) + materialized_right = QueryExpr( + op="materialized_view", params={"name": "old", "columns": tuple(right.columns)} + ) + query = replace( + query, + inputs=( + replace(joined, inputs=(materialized_left, materialized_right)), + query.inputs[1], + ), + ) + if pair_profile is not None: + config["semantic_pair_profiles"] = {query_digest(query.inputs[0]): pair_profile} + adapter = LotusAdapter( + config=LotusExecutionConfig(physical_fusion="zep-target-state", **config), + pair_embedding_provider=provider, + ) + monkeypatch.setattr(adapter._context, "configure", lambda: None) + return adapter.execute(query, {"incoming": left, "old": right}), model + + +def test_default_and_registered_rewrite() -> None: + original = ZepMemory.differentiate_policy() + assert optimize_policy(original) is original + fused = optimize_policy(original, strategy="zep-target-state") + assert ( + fused.fingerprint + == optimize_policy(original, strategy="zep-target-state").fingerprint + ) + assert fused.fingerprint != original.fingerprint + assert fused.spec is original.spec + assert optimize_policy(fused, strategy="zep-target-state") is fused + changed = [key for key in original.nodes if original.nodes[key] != fused.nodes[key]] + assert len(changed) == 1 + for key in original.nodes: + assert original.nodes[key].query == fused.nodes[key].query + + +def test_matched_new_untouched_and_many_to_one(monkeypatch: pytest.MonkeyPatch) -> None: + left = pd.DataFrame( + [state("new one", 2), state("new two", 3), state("new only", 4, entity="x")] + ) + right = pd.DataFrame([state("old", 0), state("untouched", 1, entity="z")]) + result, model = execute_fused( + monkeypatch, + left, + right, + [response([("l0", "r0"), ("l1", "r0")], [("r0", "merged")])], + ) + assert result["fact"].tolist() == ["untouched", "merged", "new only"] + merged = result[result.fact == "merged"].iloc[0] + assert merged.fact_id == (0, 0) + provenance = json.loads(merged.provenance) + assert len(provenance) == 3 + assert {p["episode_id"] for p in provenance} == {"0", "2", "3"} + assert len(model.calls) == 1 + prompt = json.loads(model.calls[0][0][0][1]["content"]) + assert len(prompt["incoming"]) == 2 + assert len(prompt["targets"]) == 1 + assert "join_instruction" in prompt and "consolidation_instruction" in prompt + + +def test_no_selected_target_keeps_both(monkeypatch: pytest.MonkeyPatch) -> None: + result, model = execute_fused( + monkeypatch, + pd.DataFrame([state("new", 1)]), + pd.DataFrame([state("old", 0)]), + [response([("l0", None)], [])], + ) + assert result.fact.tolist() == ["old", "new"] + assert len(model.calls) == 1 + + +@pytest.mark.parametrize("side", ["left", "right", "both"]) +def test_empty_sides_do_not_call_provider( + monkeypatch: pytest.MonkeyPatch, side: str +) -> None: + frame = pd.DataFrame([state("old", 0)]) + left = frame.iloc[:0] if side in {"left", "both"} else frame + right = frame.iloc[:0] if side in {"right", "both"} else frame + result, model = execute_fused(monkeypatch, left, right, []) + assert len(result) == len(left) + len(right) + assert not model.calls + + +@pytest.mark.parametrize( + "raw", + [ + "{}", + response([("l0", "r9")], [("r9", "bad")]), + response([], []), + response([("l0", "r0"), ("l0", "r0")], [("r0", "bad")]), + response([("l0", "r0")], []), + response([("l0", None)], [("r0", "extra")]), + '{"matches":[{"left_id":"l0","right_id":"r0"}],"states":[{"right_id":"r0","relation_type":"x","fact":false}]}', + ], +) +def test_invalid_output_rejected(raw: str) -> None: + with pytest.raises(ValueError): + parse_resolution(raw, {"l0": {"r0"}}) + + +def test_cross_row_target_assignment_is_rejected() -> None: + raw = response( + [("l0", "r1"), ("l1", "r1")], + [("r1", "merged")], + ) + with pytest.raises( + ValueError, + match=r"target ID 'r1' is not eligible for incoming row 'l0'", + ): + parse_resolution(raw, {"l0": {"r0"}, "l1": {"r1"}}) + + +def test_shared_retry_contract(monkeypatch: pytest.MonkeyPatch) -> None: + result, model = execute_fused( + monkeypatch, + pd.DataFrame([state("new", 1)]), + pd.DataFrame([state("old", 0)]), + ["{}", response([("l0", "r0")], [("r0", "merged")])], + structured_parse_retries=1, + ) + assert result.fact.tolist() == ["merged"] + assert len(model.calls) == 2 + assert model.calls[0][0] != model.calls[1][0] + retry_feedback = model.calls[1][0][0][-1] + assert retry_feedback["role"] == "user" + assert "fusion output requires only matches and states" in retry_feedback["content"] + assert "eligible_targets" in retry_feedback["content"] + + +def test_fusion_keeps_one_safety_retry_when_generic_retries_are_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + result, model = execute_fused( + monkeypatch, + pd.DataFrame([state("new", 1)]), + pd.DataFrame([state("old", 0)]), + ["{}", response([("l0", "r0")], [("r0", "merged")])], + structured_parse_retries=0, + ) + assert result.fact.tolist() == ["merged"] + assert len(model.calls) == 2 + + +def test_unsupported_plan_and_config() -> None: + from agent_memory.memories.mem0.policy import Mem0Memory + + with pytest.raises(ValueError, match="registered"): + optimize_policy(Mem0Memory.differentiate_policy(), strategy="zep-target-state") + with pytest.raises(ValueError, match="unknown"): + optimize_policy(ZepMemory.differentiate_policy(), strategy="invented") + with pytest.raises(ValueError, match="plain"): + LotusExecutionConfig( + physical_fusion="zep-target-state", sem_join_topk_method="pairwise-naive" + ) + + +def fact_flow( + monkeypatch: pytest.MonkeyPatch, answers: list[str], *, enabled: bool = True +) -> tuple[Any, FakeLM]: + """Use the real fact plan; only initial extraction/grouping is an oracle.""" + import agent_memory as am + import lotus + from agent_memory.policy.aggregates import SemanticAggregateSpec + from agent_memory.policy.relation import Relation + from agent_memory.policy.schema import output_columns + + class FactOracleAdapter(LotusAdapter): + def execute(self, query: QueryExpr, inputs: Any) -> Any: + if query.op == "agg" and query.inputs[0].op == "sem_groupby": + rows = self.execute(query.inputs[0].inputs[0], inputs) + # These fixtures have one distinct fact per incoming message. + assert len(rows) <= 1 + return pd.DataFrame( + [ + state( + row["content"], + row["add_seq"], + entity=row["source_entity_id"], + ) + for _, row in rows.iterrows() + ], + columns=pd.Index(output_columns(query)), + ) + if query.op == "agg": + assert not any( + isinstance(s, SemanticAggregateSpec) + for s in query.params["aggregates"] + ), "fused merge executed twice" + if query.op == "sem_join": + raise AssertionError("fused matching executed twice") + return super().execute(query, inputs) + + raw = {**state("old", 0), "episode_id": "0", "fact_ordinal": 0, "content": "old"} + source = am.Source({name: name for name in raw}) + original = ZepMemory._deduplicated_facts.expr + aggregate = original.inputs[0] + grouped = replace(aggregate.inputs[0], inputs=(source.expr,)) + view = Relation(replace(original, inputs=(replace(aggregate, inputs=(grouped,)),))) + model = FakeLM(answers) + monkeypatch.setattr(lotus.settings, "lm", model) + adapter = FactOracleAdapter( + config=LotusExecutionConfig( + physical_fusion="zep-target-state" if enabled else "disabled", + structured_parse_retries=0, + ) + ) + monkeypatch.setattr(adapter._context, "configure", lambda: None) + return am.SemanticDataflow( + source=source, views={"facts": view}, adapter=adapter + ), model + + +def fact_input(fact: str, seq: int, *, entity: str = "a") -> pd.DataFrame: + return pd.DataFrame( + [ + { + **state(fact, seq, entity=entity), + "episode_id": str(seq), + "fact_ordinal": 0, + "content": fact, + } + ] + ) + + +def test_real_runtime_restore_and_single_execution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + flow, model = fact_flow(monkeypatch, [response([("l0", "r0")], [("r0", "merged")])]) + flow.apply(fact_input("old", 0)) + assert not model.calls + flow.apply(fact_input("new", 1)) + assert flow.view("facts").fact.tolist() == ["merged"] + assert len(model.calls) == 1 + saved = flow.snapshot_state() + restored, second_model = fact_flow(monkeypatch, []) + restored.restore_state(saved) + restored.apply(fact_input("unrelated", 2, entity="other")) + assert restored.view("facts").fact.tolist() == ["merged", "unrelated"] + assert not second_model.calls + plain, _ = fact_flow(monkeypatch, [], enabled=False) + with pytest.raises(ValueError, match="fingerprint"): + plain.restore_state(saved) + + +def test_failed_fused_update_does_not_commit( + monkeypatch: pytest.MonkeyPatch, tmp_path: Any +) -> None: + from agent_memory.adapters.lotus import structured + + monkeypatch.setattr(structured, "STRUCTURED_FAILURE_DIR", tmp_path) + flow, model = fact_flow(monkeypatch, ["{}", "{}"]) + flow.apply(fact_input("old", 0)) + before = flow.view("facts") + with pytest.raises(ValueError, match="invalid target-state"): + flow.apply(fact_input("new", 1)) + pd.testing.assert_frame_equal(flow.view("facts"), before) + assert len(flow.snapshot_state()["state"]["log"]) == 1 + assert len(model.calls) == 2 + + +def test_screening_and_partition_contract(monkeypatch: pytest.MonkeyPatch) -> None: + from agent_memory.storage import EmbeddingSpec + + class Embeddings: + def embed(self, spec: Any, texts: Any) -> list[list[float]]: + return [[0.0, 1.0] if "bad" in text else [1.0, 0.0] for text in texts] + + profile = SemanticPairExecutionProfile( + mode="search-filter", + direction="left-to-right", + left_id_columns=(PAIR_LEFT_ID_COLUMN,), + right_id_columns=(PAIR_RIGHT_ID_COLUMN,), + left_text_columns=(PAIR_LEFT_TEXT_COLUMN,), + right_text_columns=(PAIR_RIGHT_TEXT_COLUMN,), + embedding=EmbeddingSpec( + source_column="text", + property_name="embedding", + model="fake", + revision="1", + dimensions=2, + normalize=True, + ), + top_k=1, + min_similarity=0.5, + ) + result, model = execute_fused( + monkeypatch, + pd.DataFrame([state("new", 3)]), + pd.DataFrame( + [state("bad", 0), state("good", 1), state("other", 2, entity="other")] + ), + [response([("l0", "r1")], [("r1", "merged")])], + pair_profile=profile, + provider=Embeddings(), + ) + payload = json.loads(model.calls[0][0][0][1]["content"]) + assert payload["incoming"][0]["eligible_targets"] == ["r1"] + assert result.fact.tolist() == ["bad", "other", "merged"] + + +def test_legacy_restore_rejected_before_mutation() -> None: + from agent_memory.runtime.runtime import MemoryRuntime + + runtime = MemoryRuntime( + ZepMemory.differentiate_policy(), + adapter=LotusAdapter( + config=LotusExecutionConfig(physical_fusion="zep-target-state") + ), + ) + with pytest.raises(ValueError, match="legacy"): + runtime.restore_state({"schema_version": 1}) + + +def test_multiple_match_membership_not_silently_narrowed() -> None: + original = ZepMemory.differentiate_policy() + nodes = dict(original.nodes) + key = next( + key + for key, node in nodes.items() + if node.query.op == "agg" + and node.query.params == ZepMemory._deduplicated_facts.expr.inputs[0].params + ) + node = nodes[key] + grouped = replace( + node.query.inputs[0], + params={**node.query.inputs[0].params, "membership": "overlapping"}, + ) + nodes[key] = replace(node, query=replace(node.query, inputs=(grouped,))) + with pytest.raises(ValueError, match="registered"): + optimize_policy(replace(original, nodes=nodes), strategy="zep-target-state") + + +def test_trace_preserves_both_logical_results( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + raw = response([("l0", "r0")], [("r0", "merged")]) + _, model = execute_fused( + monkeypatch, + pd.DataFrame([state("new", 1)]), + pd.DataFrame([state("old", 0)]), + [raw], + semantic_trace_dir=tmp_path, + ) + events = [ + json.loads(line) + for line in (tmp_path / "events.jsonl").read_text().splitlines() + ] + (event,) = [item for item in events if item["event_type"] == "fusion_resolution"] + parsed = json.loads(Path(event["parsed_output_path"]).read_text()) + assert parsed["matches"] == {"l0": "r0"} + assert parsed["states"]["r0"]["fact"] == "merged" + assert json.loads(Path(event["raw_output_path"]).read_text()) == [[raw]] + assert event["eligible_pair_count"] == 1 + assert len(model.calls) == 1 + + +def test_prepared_plan_cannot_run_with_disabled_adapter() -> None: + from agent_memory.runtime.executor import PolicyExecutor + + fused = optimize_policy( + ZepMemory.differentiate_policy(), strategy="zep-target-state" + ) + with pytest.raises(ValueError, match="configuration"): + PolicyExecutor(fused, adapter=LotusAdapter()) + + +def test_default_execution_identity_unchanged() -> None: + from agent_memory.runtime.executor import PolicyExecutor + + original = ZepMemory.differentiate_policy() + adapter = LotusAdapter() + assert PolicyExecutor(original, adapter=adapter).policy is original + assert adapter.maintenance_execution_fingerprint == "" + fused_adapter = LotusAdapter( + config=LotusExecutionConfig(physical_fusion="zep-target-state") + ) + assert ( + fused_adapter.maintenance_execution_fingerprint + != adapter.maintenance_execution_fingerprint + ) + + +def test_one_provider_call_not_two_logical_charges( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + from agent_memory.evaluation.trace_metrics import normalize_provider_calls + from agent_memory.tracing.semantic import ( + write_llm_call_trace, + write_provider_usage_trace, + ) + + original_call = FakeLM.__call__ + + def traced_call(model: FakeLM, prompts: Any, **kwargs: Any) -> Any: + result = original_call(model, prompts, **kwargs) + write_llm_call_trace( + tmp_path, + model="fake", + messages=prompts, + kwargs=kwargs, + outputs=result.outputs, + latency_sec=0.01, + ) + write_provider_usage_trace( + tmp_path, + model="fake", + responses=[ + SimpleNamespace( + usage={ + "prompt_tokens": 100, + "completion_tokens": 20, + "total_tokens": 120, + } + ) + ], + ) + return result + + monkeypatch.setattr(FakeLM, "__call__", traced_call) + execute_fused( + monkeypatch, + pd.DataFrame([state("new", 1)]), + pd.DataFrame([state("old", 0)]), + [response([("l0", "r0")], [("r0", "merged")])], + semantic_trace_dir=tmp_path, + ) + events = [ + json.loads(line) + for line in (tmp_path / "events.jsonl").read_text().splitlines() + ] + calls = normalize_provider_calls(events, output_dir=tmp_path, include_cost=False) + assert len(calls) == 1 + assert calls[0]["operator"] == "fused_target_state" + assert calls[0]["total_tokens"] == 120 + + +def test_benchmark_factory_passes_fusion(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + from agent_memory.adapters import lotus + from agent_memory.evaluation.agent_memory_drivers import ZepMemoryDriverFactory + + class AdapterReached(Exception): + pass + + def inspect_adapter(**kwargs: Any) -> Any: + assert kwargs["config"].physical_fusion == "zep-target-state" + raise AdapterReached + + factory = ZepMemoryDriverFactory(connector=SimpleNamespace(embedding_provider=None), + base_namespace="test", physical_fusion="zep-target-state", + neo4j_image="fake", neo4j_image_digest="fake") + monkeypatch.setattr(lotus, "LotusAdapter", inspect_adapter) + with pytest.raises(AdapterReached): + factory("case", tmp_path / "state", tmp_path / "trace") diff --git a/tests/test_relational_join_ivm.py b/tests/test_relational_join_ivm.py index c43c00c..6320aec 100644 --- a/tests/test_relational_join_ivm.py +++ b/tests/test_relational_join_ivm.py @@ -10,6 +10,7 @@ import pytest import agent_memory as am +from agent_memory.runtime import executor as executor_module from agent_memory.adapters import LotusAdapter from agent_memory.planner.differential_policy import DifferentialNode from agent_memory.policy.logical import QueryExpr @@ -226,7 +227,7 @@ def test_downstream_sem_filter_receives_only_new_join_pairs() -> None: assert len(memory._runtime._state["links"]) == 3 -def test_parent_retraction_uses_full_join_recompute() -> None: +def test_parent_retraction_uses_affected_join_parts() -> None: adapter = _TrackingAdapter() memory = _RetractingParentMemory(adapter=adapter) node = _only_join_node(_RetractingParentMemory) @@ -235,7 +236,8 @@ def test_parent_retraction_uses_full_join_recompute() -> None: adapter.root_queries.clear() memory.add({"key": "k", "value": 3}) - assert node.query in adapter.root_queries + assert node.query not in adapter.root_queries + assert any(query.op == "join" for query in adapter.root_queries) assert node.maintenance_query not in adapter.root_queries assert memory._runtime._state["pairs"].to_dict(orient="records") == [ {"key": "k", "value:left": 3, "value:right": 3} @@ -272,3 +274,128 @@ def test_inner_join_snapshot_round_trip_and_fingerprint_validation() -> None: incompatible["plan_fingerprint"] = "pre-relational-state-plan" with pytest.raises(ValueError, match="plan fingerprint"): _TemporalSelfJoinMemory()._runtime.restore_state(incompatible) + + +@pytest.mark.parametrize("value", ["next", "unchanged"]) +def test_join_does_not_rekey_old_pairs_on_append_or_no_change( + monkeypatch: pytest.MonkeyPatch, value: str, +) -> None: + class FilteredMemory(am.Memory): + log = am.Log({"value": "Value."}, system_columns=True) + _kept = log.filter(log.col("value") != "unchanged") + _left = _kept.alias("left") + _right = _kept.alias("right") + pairs = _left.join( + _right, on=_left.col(LOG_ADD_SEQ_COLUMN) < _right.col(LOG_ADD_SEQ_COLUMN), + ) + + memory = FilteredMemory() + for v in ("a", "b", "c"): + memory.add({"value": v}) + node = _only_join_node(FilteredMemory) + old = memory._runtime.snapshot_state()["node_state"][node.node_id].copy() + original = executor_module._row_key + old_keys = {original(row) for row in old.itertuples(index=False, name=None)} + inspected = 0 + + def count_pair_keys(row: tuple[Any, ...]) -> tuple[Any, ...]: + nonlocal inspected + if len(row) == len(old.columns) and original(row) in old_keys: + inspected += 1 + return original(row) + + monkeypatch.setattr(executor_module, "_row_key", count_pair_keys) + memory.add({"value": value}) + current = memory._runtime.snapshot_state()["node_state"][node.node_id] + pd.testing.assert_frame_equal(current.loc[old.index], old) + assert inspected == 0 + assert len(current) == (6 if value == "next" else 3) + + +@pytest.mark.parametrize("indexed", [False, True]) +def test_join_delta_cancels_replacements_and_preserves_nested_null_rows(indexed: bool) -> None: + old = pd.DataFrame({"value": [None, None, {"a": [1, None]}]}, index=pd.Index(["a", "b", "c"])) + removed = old.iloc[[0, 2]] + added = pd.DataFrame({"value": [{"a": [1, None]}, "new"]}) + index = executor_module._index_join_rows(old) if indexed else None + update = executor_module._apply_join_delta(old, removed, added, allocate=lambda: "d", row_index=index) + assert update.retracted_rows.index.tolist() == ["b"] + assert update.inserted_rows.index.tolist() == ["d"] + pd.testing.assert_frame_equal(update.output_rows.loc[["a", "c"]], old.loc[["a", "c"]]) + assert update.output_rows.loc["d", "value"] == "new" + if indexed: + assert index == executor_module._index_join_rows(update.output_rows) + with pytest.raises(RuntimeError, match="retractions exceed"): + executor_module._apply_join_delta(old, pd.concat([old, old]), old.iloc[:0], allocate=lambda: "unused") + + +def test_replacing_unprojected_join_input_does_not_repeat_semantic_work() -> None: + class ProjectedMemory(am.Memory): + log = am.Log({"key": "key", "value": "value"}) + _state = log.group_by("key").min(column="value", output_col="value") + _left, _right = _state.alias("l"), _state.alias("r") + pairs = _left.join(_right, on="key").select(["key"]) + result = pairs.sem_filter(instruction="Keep all rows.") + + adapter = _TrackingAdapter() + memory = ProjectedMemory(adapter=adapter) + memory.add({"key": "a", "value": 9}) + memory.add({"key": "b", "value": 5}) + before = memory._runtime.snapshot_state() + calls = list(adapter.sem_filter_input_sizes) + memory.add({"key": "a", "value": 3}) + assert adapter.sem_filter_input_sizes == calls + assert _bag(memory._runtime._state["pairs"]) == Counter({("a",): 1, ("b",): 1}) + node = _only_join_node(ProjectedMemory) + old = before["node_state"][node.node_id] + unchanged = old.loc[old["key"] == "b"] + current = memory._runtime.snapshot_state()["node_state"][node.node_id] + pd.testing.assert_frame_equal(current.loc[unchanged.index], unchanged) + + +def test_join_append_retains_compiled_task_order() -> None: + memory = _ExplodedSelfJoinMemory() + memory.add({"items": '["a", "b"]'}) + before = memory._runtime.snapshot_state()["node_state"] + memory.add({"items": '["c", "d"]'}) + after = memory._runtime.snapshot_state()["node_state"] + node = _only_join_node(_ExplodedSelfJoinMemory) + inputs = {node.node_id: before[node.node_id]} + for parent in node.input_node_ids: + inputs[parent] = before[parent] + inputs[f"{parent}__inserted"] = executor_module.NodeOutputUpdate.between( + before[parent], after[parent], + ).inserted_rows + assert node.maintenance_query is not None + expected = LotusAdapter().execute(node.maintenance_query, inputs) + pd.testing.assert_frame_equal( + after[node.node_id].reset_index(drop=True), expected.reset_index(drop=True), + ) + + +@pytest.mark.parametrize("indexed", [False, True]) +def test_join_delta_random_bags_match_net_changes(indexed: bool) -> None: + import random + from itertools import count + + rng = random.Random(41) + for _ in range(50): + values = [rng.randrange(4) for _ in range(rng.randrange(20))] + old = pd.DataFrame({"value": values}, index=pd.Index([f"old:{i}" for i in range(len(values))])) + positions = rng.sample(range(len(values)), rng.randrange(len(values) + 1)) + removed = old.iloc[positions] + new_values = [rng.randrange(4) for _ in range(rng.randrange(10))] + added = pd.DataFrame({"value": new_values}) + ids = count() + index = executor_module._index_join_rows(old) if indexed else None + update = executor_module._apply_join_delta( + old, removed, added, allocate=lambda: f"new:{next(ids)}", row_index=index, + ) + expected = Counter(v for i, v in enumerate(values) if i not in positions) + Counter(new_values) + actual = Counter(update.output_rows["value"]) + assert actual == expected + assert Counter(update.inserted_rows["value"]) == expected - Counter(values) + assert Counter(update.retracted_rows["value"]) == Counter(values) - expected + assert update.output_rows.index.is_unique + if indexed: + assert index == executor_module._index_join_rows(update.output_rows) diff --git a/tests/test_replacement_screening.py b/tests/test_replacement_screening.py new file mode 100644 index 0000000..f368e84 --- /dev/null +++ b/tests/test_replacement_screening.py @@ -0,0 +1,235 @@ +"""Fixed-oracle regressions for Top-k membership across source replacements.""" + +from math import cos, sin +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.pair_execution import SemanticPairExecutionProfile, semantic_pair_site_id +from agent_memory.planner.physical import walk +from agent_memory.runtime.executor import NodeOutputUpdate +from agent_memory.storage.embedding import EmbeddingSpec +from agent_memory.tracing.semantic import query_digest + + +class FixedVectors: + calls: int = 0 + + def embed(self, spec: EmbeddingSpec, texts: list[str]) -> list[list[float]]: + self.calls += 1 + angles = {"A": 1.2, "B": .1, "C": .2, "D": .3, "E": .4, "F": .5, "G": 0, "H": .05} + return [[cos(angles[t.rsplit(': ', 1)[-1]]), sin(angles[t.rsplit(': ', 1)[-1]])] for t in texts] + + +def build(monkeypatch: pytest.MonkeyPatch, direction: str = "right-to-left") -> tuple[Any, ...]: + source = am.Source({"id": "number", "fact": "text", "episode": "text"}) + state = source.group_by(["id", "fact"]).array_agg(columns=["episode"], output_col="provenance") + left, right = state.alias("earlier"), state.alias("later") + view = left.join(right, on=left.col("id") < right.col("id")).sem_filter( + instruction="Compare {fact:earlier} and {fact:later}.") + policy = am.SemanticDataflow(source=source, views={"result": view}, adapter=LotusAdapter())._executor.policy + digests = {query_digest(q) for n in policy.nodes.values() + for root in (n.query, n.maintenance_query) if root is not None + for q in walk(root) if q.op == "sem_filter"} + profile = SemanticPairExecutionProfile(mode="search-filter", direction=direction, + left_id_columns=("id:earlier",), right_id_columns=("id:later",), + left_text_columns=("fact:earlier",), right_text_columns=("fact:later",), + embedding=EmbeddingSpec("fact", "vector", "fake", "v1", 2, True), top_k=5) + vectors = FixedVectors() + adapter = LotusAdapter(config=LotusExecutionConfig( + semantic_pair_profiles={d: profile for d in digests}, + predicate_reuse_sites=(semantic_pair_site_id(view.expr),)), pair_embedding_provider=vectors) + calls: list[tuple[str, str]] = [] + + def oracle(frame: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame: + calls.extend(zip(frame["fact_earlier"], frame["fact_later"])) + return frame.copy() + + monkeypatch.setattr(adapter._context, "configure", lambda: None) + monkeypatch.setattr(pd.DataFrame, "sem_filter", oracle) + full = LotusAdapter(config=LotusExecutionConfig( + semantic_pair_profiles={query_digest(view.expr): profile}), pair_embedding_provider=FixedVectors()) + monkeypatch.setattr(full._context, "configure", lambda: None) + flow = am.SemanticDataflow(source=source, views={"result": view}, adapter=adapter) + return source, view, flow, adapter, full, vectors, calls + + +@pytest.mark.parametrize("direction", ["right-to-left", "left-to-right", "symmetric"]) +def test_metadata_only_replacement_keeps_full_topk_and_makes_no_calls(monkeypatch: pytest.MonkeyPatch, direction: str) -> None: + source, view, flow, adapter, full, vectors, calls = build(monkeypatch, direction) + initial = pd.DataFrame({"id": range(7), "fact": list("ABCDEFG"), "episode": ["original"] * 7}) + flow.apply(initial) + count, embedded = len(calls), vectors.calls + extra = initial.iloc[[0]].assign(episode="another source") + flow.apply(extra) + assert len(calls) == count + assert vectors.calls == embedded + expected = full.execute(view.expr, {"log": pd.concat([initial, extra], ignore_index=True)}) + assert NodeOutputUpdate.between(flow.view("result"), expected).is_empty + restored = am.SemanticDataflow(source=source, views={"result": view}, adapter=adapter) + restored.restore_state(flow.snapshot_state()) + count = len(calls) + restored.apply(extra.assign(episode="third source")) + assert len(calls) == count + expected = full.execute(view.expr, {"log": pd.concat([initial, extra, extra.assign(episode="third source")], ignore_index=True)}) + assert NodeOutputUpdate.between(restored.view("result"), expected).is_empty + + +def test_new_stronger_candidate_retracts_displaced_old_pair(monkeypatch: pytest.MonkeyPatch) -> None: + _, view, flow, _, full, _, _ = build(monkeypatch) + # H arrives later but has an earlier logical ID; it displaces a candidate of G. + initial = pd.DataFrame({"id": [0, 2, 3, 4, 5, 6, 7], "fact": list("ABCDEFG"), "episode": ["original"] * 7}) + flow.apply(initial) + extra = pd.DataFrame({"id": [1], "fact": ["H"], "episode": ["new"]}) + flow.apply(extra) + expected = full.execute(view.expr, {"log": pd.concat([initial, extra], ignore_index=True)}) + assert NodeOutputUpdate.between(flow.view("result"), expected).is_empty + + +def test_mixed_update_does_not_rescreen_metadata_only_buckets(monkeypatch: pytest.MonkeyPatch) -> None: + from agent_memory.adapters.lotus import sem_filter as filter_module + + _, view, flow, _, full, _, _ = build(monkeypatch) + initial = pd.DataFrame({"id": range(7), "fact": list("ABCDEFG"), "episode": ["original"] * 7}) + flow.apply(initial) + screened: list[tuple[str, str]] = [] + select = filter_module.select_semantic_pair_candidates + + def record(source: pd.DataFrame, **kwargs: Any) -> Any: + screened.extend(zip(source["fact:earlier"], source["fact:later"])) + return select(source, **kwargs) + + monkeypatch.setattr(filter_module, "select_semantic_pair_candidates", record) + extra = pd.DataFrame({"id": [0, 7], "fact": ["A", "H"], "episode": ["new source", "new fact"]}) + flow.apply(extra) + assert screened == [(fact, "H") for fact in "ABCDEFG"] + expected = full.execute(view.expr, {"log": pd.concat([initial, extra], ignore_index=True)}) + assert NodeOutputUpdate.between(flow.view("result"), expected).is_empty + + +@pytest.mark.parametrize("direction", ["right-to-left", "left-to-right", "symmetric"]) +def test_mixed_candidate_changes_match_full_after_restore(monkeypatch: pytest.MonkeyPatch, direction: str) -> None: + source, view, flow, adapter, full, _, _ = build(monkeypatch, direction) + initial = pd.DataFrame({"id": [0, 2, 3, 4, 5, 6, 7], "fact": list("ABCDEFG"), "episode": ["original"] * 7}) + flow.apply(initial) + restored = am.SemanticDataflow(source=source, views={"result": view}, adapter=adapter) + restored.restore_state(flow.snapshot_state()) + # A's metadata changes while H competes in existing candidate buckets. + extra = pd.DataFrame({"id": [0, 1], "fact": ["A", "H"], "episode": ["new source", "new fact"]}) + restored.apply(extra) + expected = full.execute(view.expr, {"log": pd.concat([initial, extra], ignore_index=True)}) + assert NodeOutputUpdate.between(restored.view("result"), expected).is_empty + + +def test_representative_matching_declares_resolved_time_on_both_paths() -> None: + from agent_memory.memories.zep.representative import ZepRepresentativeMemory + from agent_memory.policy.aggregates import ArgMinAggregateSpec + + agg = next(q for q in walk(ZepRepresentativeMemory.facts.expr) + if any(isinstance(s, ArgMinAggregateSpec) for s in q.params.get("aggregates", ()))) + assert "valid_at" in agg.inputs[0].params["input_cols"] + assert "{valid_at}" in agg.inputs[0].params["instruction"] + policy = ZepRepresentativeMemory.differentiate_policy() + joins = [q for node in policy.nodes.values() if node.maintenance_query is not None + for q in walk(node.maintenance_query) if q.op == "sem_join" + and "relation_type:left" in q.params["instruction"]] + assert joins + assert all("{valid_at:left}" in q.params["instruction"] and "{valid_at:right}" in q.params["instruction"] for q in joins) + + +def test_representative_full_grouping_sends_resolved_dates(monkeypatch: pytest.MonkeyPatch) -> None: + import json + import lotus + from agent_memory.adapters.lotus.prompt_batching import PromptBatching + from agent_memory.adapters.lotus.sem_groupby import execute_sem_groupby, GROUP_ID_COLUMN + from agent_memory.memories.zep.representative import ZepRepresentativeMemory + from agent_memory.policy.aggregates import ArgMinAggregateSpec + from test_zep_clean_maintenance import Oracle, offline_context + + group = next(q.inputs[0] for q in walk(ZepRepresentativeMemory.facts.expr) + if any(isinstance(s, ArgMinAggregateSpec) for s in q.params.get("aggregates", ()))) + rows = pd.DataFrame({"source_entity_id": ["a", "a"], "target_entity_id": ["b", "b"], + "relation_type": ["ATTENDED"] * 2, + "fact": ["Alice attended a parade last Friday.", "Alice attended a parade last week."], + "valid_at": ["2023-06-02", "2023-08-04"]}) + config = LotusExecutionConfig(groupby_prompt_batching={semantic_pair_site_id(group): PromptBatching(max_tasks=16)}) + oracle = Oracle() + monkeypatch.setattr(lotus.settings, "lm", oracle) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + result = execute_sem_groupby(group, {}, lambda *_: rows, offline_context(config, monkeypatch)) + assert result[GROUP_ID_COLUMN].nunique() == 2 # Fake oracle is False; not a model quality assertion. + assert len(oracle.requests) == 1 + prompt = json.dumps(oracle.requests[0]) + assert "2023-06-02" in prompt and "2023-08-04" in prompt and "valid_at" in prompt + + +@pytest.mark.parametrize("direction", ["right-to-left", "left-to-right", "symmetric"]) +def test_removal_promotion_text_change_duplicates_and_empty(monkeypatch: pytest.MonkeyPatch, direction: str) -> None: + _, _, flow, adapter, _, _, _ = build(monkeypatch, direction) + executor = flow._executor + node_id, node = next((i, n) for i, n in executor.policy.nodes.items() if n.execution_kind == "semantic_predicate") + parent = node.input_node_ids[0] + base = pd.DataFrame({ + "id:earlier": range(6), "fact:earlier": list("ABCDEF"), + "provenance:earlier": ["original"] * 6, + "id:later": [6] * 6, "fact:later": ["G"] * 6, + "provenance:later": ["original"] * 6, + }, index=pd.Index([f"p{i}" for i in range(6)])) + duplicate = base.iloc[[0]].copy() + duplicate.index = pd.Index(["duplicate"]) + frames = [base, base.drop(index="p1"), pd.concat([base.drop(index="p1"), duplicate])] + modified = frames[-1].copy() + modified.loc[modified["id:earlier"] == 0, "fact:earlier"] = "H" + frames.extend([modified, modified.iloc[:0], base]) + cache: dict = {} + decisions: dict = {} + previous = base.iloc[:0] + result = previous + for current in frames: + executor._node_state[parent] = previous + result = executor._execute_semantic_row(node_id, old_state=result, + parent_update=NodeOutputUpdate.between(previous, current), staged_cache=cache, + predicate_decisions=decisions) + expected = adapter.execute(node.query, {parent: current}) + assert NodeOutputUpdate.between(result, expected).is_empty + previous = current + + +@pytest.mark.parametrize("mixed", [False, True]) +def test_failure_leaves_join_index_membership_and_predicates_uncommitted(monkeypatch: pytest.MonkeyPatch, mixed: bool) -> None: + from copy import deepcopy + + source, view, flow, adapter, _, _, _ = build(monkeypatch) + initial = pd.DataFrame({"id": [0, 2, 3, 4, 5, 6, 7], "fact": list("ABCDEFG"), "episode": ["original"] * 7}) + flow.apply(initial) + snapshot = flow.snapshot_state() + indexes = deepcopy(flow._executor._join_row_indexes) + cache = deepcopy(flow._executor._predicate_decisions) + + def fail(*args: Any, **kwargs: Any) -> Any: + raise ValueError("synthetic predicate failure") + + monkeypatch.setattr(pd.DataFrame, "sem_filter", fail) + with pytest.raises(ValueError, match="synthetic predicate failure"): + extra = pd.DataFrame({"id": [1], "fact": ["H"], "episode": ["new"]}) + if mixed: + extra = pd.DataFrame({"id": [0, 8], "fact": ["A", "H"], "episode": ["new source", "new fact"]}) + flow.apply(extra) + assert flow._executor._join_row_indexes == indexes + assert flow._executor._predicate_decisions == cache + restored = am.SemanticDataflow(source=source, views={"result": view}, adapter=adapter) + restored.restore_state(snapshot) + pd.testing.assert_frame_equal(flow.view("result"), restored.view("result")) + assert restored._executor._join_row_indexes == {} + for node_id, cached in restored._executor._semantic_output_cache.items(): + current = flow._executor._semantic_output_cache[node_id] + assert current.keys() == cached.keys() + for occurrence, frame in cached.items(): + pd.testing.assert_frame_equal(current[occurrence], frame) + snapshot["adapter_execution_fingerprint"] = "old-invocation-local-top5" + with pytest.raises(ValueError, match="fingerprint"): + restored.restore_state(snapshot) diff --git a/tests/test_row_output_cache.py b/tests/test_row_output_cache.py new file mode 100644 index 0000000..272c437 --- /dev/null +++ b/tests/test_row_output_cache.py @@ -0,0 +1,195 @@ +"""Offline regressions for compact semantic output bookkeeping.""" + +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.runtime.row_outputs import RowOutputCache + + +class FilterOracle: + """Keep positive inputs; relational work still uses the real adapter.""" + + def __init__(self) -> None: + self.inputs: list[pd.DataFrame] = [] + self.delegate = LotusAdapter() + + def execute(self, query: Any, inputs: Any) -> pd.DataFrame: + if query.op == "sem_filter": + frame = inputs[str(query.inputs[0].params["name"])] + self.inputs.append(frame.copy(deep=True)) + return frame.loc[frame["value"] > 0].copy() + return self.delegate.execute(query, inputs) + + +def build_flow(adapter: FilterOracle) -> Any: + source = am.Source({"value": "number"}) + return am.SemanticDataflow( + source=source, + views={"result": source.sem_filter(instruction="Keep positive {value}.")}, + adapter=adapter, + ) + + +def test_rejected_occurrences_share_empty_output_across_events_and_restore() -> None: + adapter = FilterOracle() + flow = build_flow(adapter) + for _ in range(3): + flow.apply(pd.DataFrame({"value": [0] * 100})) + snapshot = flow.snapshot_state() + cache = next(iter(snapshot["semantic_output_cache"].values())) + assert len(cache) == 300 + assert all(frame.empty for frame in cache.values()) + assert len({id(frame) for frame in cache.values()}) == 1 + assert snapshot["schema_version"] == 2 + assert type(cache) is dict + restored = build_flow(adapter) + restored.restore_state(snapshot) + restored.apply(pd.DataFrame({"value": [1, 0, 1]})) + assert restored.view("result")["value"].tolist() == [1, 1] + assert [len(frame) for frame in adapter.inputs] == [100, 100, 100, 3] + + +def test_restore_compacts_legacy_independent_empties_without_changing_snapshot() -> None: + adapter = FilterOracle() + flow = build_flow(adapter) + flow.apply(pd.DataFrame({"value": [0] * 30 + [1]})) + snapshot = flow.snapshot_state() + node, cache = next(iter(snapshot["semantic_output_cache"].items())) + legacy = {key: frame.copy(deep=True) for key, frame in cache.items()} + snapshot["semantic_output_cache"][node] = legacy + restored = build_flow(adapter) + restored.restore_state(snapshot) + current = restored.snapshot_state()["semantic_output_cache"][node] + assert len({id(v) for v in legacy.values() if v.empty}) == 30 + assert len({id(v) for v in current.values() if v.empty}) == 1 + for key in legacy: + pd.testing.assert_frame_equal(legacy[key], current[key]) + assert len(adapter.inputs) == 1 + wrong = {**snapshot, "plan_fingerprint": "different"} + with pytest.raises(ValueError, match="fingerprint"): + restored.restore_state(wrong) + + +def test_membership_transitions_reproduce_dictionary_order_and_multiplicity() -> None: + empty = pd.DataFrame({"x": pd.Series(dtype="Int64")}) + first = pd.DataFrame({"x": pd.array([1, None, 1], dtype="Int64")}, index=pd.Index(["a"] * 3)) + second = first.assign(x=pd.array([2, None, 2], dtype="Int64")) + cache = RowOutputCache() + reference: dict[Any, pd.DataFrame] = {} + changes = [("a", empty), ("b", first), ("a", second), ("b", empty), + ("c", first), ("b", second), ("a", None), ("a", first)] + for key, frame in changes: + if frame is None: + del cache[key] + del reference[key] + else: + cache[key] = frame + reference[key] = frame + assert list(cache) == list(reference) + actual = cache.nonempty_frames() + expected = [v for v in reference.values() if not v.empty] + assert len(actual) == len(expected) + for left, right in zip(actual, expected, strict=True): + pd.testing.assert_frame_equal(left, right) + + +def test_forked_bookkeeping_and_output_frames_do_not_mutate_committed_state() -> None: + empty = pd.DataFrame({"x": pd.Series(dtype="int64")}) + old = pd.DataFrame({"x": [1], "metadata": [["original"]]}) + cache = RowOutputCache({"empty": empty, "kept": old}) + fork = cache.copy() + fork.pop("kept") + fork["empty"] = pd.DataFrame({"x": [2], "metadata": [["new"]]}) + fork["new"] = empty.rename(columns={"x": "other"}) + assert list(cache) == ["empty", "kept"] + assert cache["empty"].empty + pd.testing.assert_frame_equal(cache["kept"], old) + assert list(fork) == ["empty", "new"] + assert fork.nonempty_frames()[0]["metadata"].tolist() == [["new"]] + + +def test_nonempty_enumeration_does_not_inspect_empty_frames(monkeypatch: pytest.MonkeyPatch) -> None: + empty = pd.DataFrame({"x": pd.Series(dtype="int64")}) + cache = RowOutputCache({i: empty for i in range(10000)}) + kept = pd.DataFrame({"x": [1]}) + cache["kept"] = kept + + def forbidden(_: pd.DataFrame) -> bool: + raise AssertionError("enumeration inspected a frame's emptiness") + + monkeypatch.setattr(pd.DataFrame, "empty", property(forbidden)) + result = cache.nonempty_frames() + assert len(result) == 1 and result[0] is kept + + +@pytest.mark.parametrize("kind", ["dtype", "column", "index_name", "index_dtype", "category", "flags", "attrs"]) +def test_empty_schema_and_metadata_are_not_conflated(kind: str) -> None: + base = pd.DataFrame({"x": pd.Series(dtype="int64")}) + other = base.copy(deep=True) + if kind == "dtype": + other = other.astype({"x": "Int64"}) + elif kind == "column": + other.columns = pd.Index(["other"]) + elif kind == "index_name": + other.index.name = "occurrence" + elif kind == "index_dtype": + other.index = pd.Index([], dtype="object") + elif kind == "category": + other["x"] = pd.Categorical([], categories=["one", "two"], ordered=True) + elif kind == "flags": + other.flags.allows_duplicate_labels = False + else: + other.attrs["source"] = {"nested": ["value"]} + cache = RowOutputCache({"a": base, "b": other}) + assert cache["a"] is not cache["b"] + pd.testing.assert_frame_equal(cache["a"], base) + pd.testing.assert_frame_equal(cache["b"], other) + assert cache["b"].attrs == other.attrs + + +def test_empty_template_does_not_alias_mutable_caller_frame() -> None: + empty = pd.DataFrame({"x": pd.Series(dtype="int64")}) + cache = RowOutputCache({"a": empty}) + empty.columns = pd.Index(["changed"]) + assert cache["a"].columns.tolist() == ["x"] + cache["b"] = empty + assert cache["b"].columns.tolist() == ["changed"] + + +def test_rejected_predicate_membership_survives_metadata_replacement_and_restore(monkeypatch: pytest.MonkeyPatch) -> None: + from test_replacement_screening import build + + _, view, flow, adapter, full, _, _ = build(monkeypatch) + calls: list[int] = [] + + def predicate(frame: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame: + calls.append(len(frame)) + return frame.loc[~frame["fact_earlier"].isin(["A", "C"])].copy() + + monkeypatch.setattr(pd.DataFrame, "sem_filter", predicate) + initial = pd.DataFrame({"id": range(7), "fact": list("ABCDEFG"), "episode": ["old"] * 7}) + flow.apply(initial) + snapshot = flow.snapshot_state() + flow.restore_state(snapshot) + call_count = len(calls) + extra = initial.iloc[[0, 2]].assign(episode="new source") + flow.apply(extra) + assert len(calls) == call_count + expected = full.execute(view.expr, {"log": pd.concat([initial, extra], ignore_index=True)}) + actual = flow.view("result") + assert len(actual) == len(expected) + keys = ["id:earlier", "id:later"] + pd.testing.assert_frame_equal( + actual.sort_values(keys).reset_index(drop=True), + expected.sort_values(keys).reset_index(drop=True), + ) + # Full has its own row order. Maintenance must preserve the legacy cache's + # insertion order, including rows re-appended after metadata replacement. + cache = next(iter(flow._executor._semantic_output_cache.values())) + legacy_order = pd.concat([frame for frame in cache.values() if not frame.empty]) + pd.testing.assert_frame_equal(actual.reset_index(drop=True), legacy_order.reset_index(drop=True)) + assert adapter is flow._executor.adapter diff --git a/tests/test_runtime_state_indexing.py b/tests/test_runtime_state_indexing.py new file mode 100644 index 0000000..6ec78e1 --- /dev/null +++ b/tests/test_runtime_state_indexing.py @@ -0,0 +1,324 @@ +"""Exact occurrence and candidate-order contracts for derived runtime indexes.""" + +from collections import deque +from copy import deepcopy +from typing import Any + +import pandas as pd +import pytest + +import agent_memory.runtime.executor as runtime +import agent_memory as am +from agent_memory.planner.physical import walk +from agent_memory.tracing.semantic import query_digest +from test_replacement_screening import build + + +def test_row_changes_keep_duplicate_identity_and_allocator_order(monkeypatch: pytest.MonkeyPatch) -> None: + old = pd.DataFrame({"v": ["a", "a", "b", "c", "a"]}, index=pd.Index(["a0", "a1", "b", "c", "a2"])) + new = pd.DataFrame({"v": ["b", "a", "d", "a", "e"]}) + allocated: list[str] = [] + keys: list[tuple[Any, ...]] = [] + original = runtime._row_key + + def key(row: tuple[Any, ...]) -> tuple[Any, ...]: + keys.append(row) + return original(row) + + def allocate() -> str: + value = f"new{len(allocated)}" + allocated.append(value) + return value + + monkeypatch.setattr(runtime, "_row_key", key) + update = runtime._preserve_occurrences_and_diff(old, new, allocate=allocate) + assert list(update.output_rows.index) == ["b", "a0", "new0", "a1", "new1"] + assert list(update.inserted_rows.index) == ["new0", "new1"] + assert list(update.retracted_rows.index) == ["c", "a2"] + assert allocated == ["new0", "new1"] + assert len(keys) == len(old) + len(new) + + +@pytest.mark.parametrize("empty", [None, "old", "new", "both"]) +def test_row_changes_match_existing_complex_cell_contract(empty: str | None) -> None: + old = pd.DataFrame({"v": [None, {"a": [1, None]}, (2, {3, 4}), pd.NA, float("nan")], + "n": pd.array([1, 2, 3, 1, 1], dtype="Int64")}) + old.index = pd.Index(["o0", "o1", "o2", "o3", "o4"], name="old_name") + new = old.iloc[[2, 0, 1, 1]].copy() + if empty in {"old", "both"}: + old = old.iloc[:0] + if empty in {"new", "both"}: + new = new.iloc[:0] + ids = deque(f"n{i}" for i in range(20)) + expected = runtime.preserve_occurrence_index(old, new, allocate=ids.popleft) + expected_update = runtime.NodeOutputUpdate.between(old, expected) + ids = deque(f"n{i}" for i in range(20)) + actual = runtime._preserve_occurrences_and_diff(old, new, allocate=ids.popleft) + for field in ("output_rows", "inserted_rows", "retracted_rows"): + pd.testing.assert_frame_equal(getattr(actual, field), getattr(expected_update, field)) + with pytest.raises(ValueError, match="ordered columns"): + runtime._preserve_occurrences_and_diff(old, new[["n", "v"]], allocate=ids.popleft) + + +def test_row_changes_preserve_occurrence_scalar_types() -> None: + old = pd.DataFrame({"value": [1, 2]}, index=pd.Index([11, 12])) + expected = runtime.preserve_occurrence_index(old, old, allocate=lambda: 99) + actual = runtime._preserve_occurrences_and_diff(old, old, allocate=lambda: 99).output_rows + assert [(type(i), i) for i in actual.index] == [(type(i), i) for i in expected.index] + + +def test_bucket_index_copies_only_touched_members_and_preserves_order() -> None: + old = pd.DataFrame({"bucket": [1, 2, 1, 3], "entity": ["x", "x", "y", "x"], + "metadata": [0, 0, 0, 0]}, index=pd.Index(["a", "b", "c", "d"])) + columns = ("bucket", "entity") + index = runtime._CandidateBucketIndex.build(old, columns) + removed = old.loc[["a"]] + added = pd.DataFrame({"bucket": [1, 2], "entity": ["x", "x"], "metadata": [1, 1]}, index=pd.Index(["e", "f"])) + current = pd.concat([old.drop(index="a"), added]) + updated = index.updated(removed, added) + assert updated.buckets[runtime._row_key((3, "x"))] is index.buckets[runtime._row_key((3, "x"))] + keys = {runtime._row_key((1, "x")), runtime._row_key((2, "x"))} + assert list(index.select(old, keys).index) == ["a", "b"] + assert list(updated.select(current, keys).index) == ["b", "e", "f"] + assert list(updated.select(current, keys)["metadata"]) == [0, 1, 1] + empty = updated.updated(current, current.iloc[:0]) + assert empty.buckets == {} + reborn = empty.updated(current.iloc[:0], current) + pd.testing.assert_frame_equal(reborn.select(current, set(reborn.buckets)), current) + + +@pytest.mark.parametrize("values", [("a", "b"), (1, 2), (True, 1), ("a", 1), (("fact", 1), ("fact", 2))]) +def test_bucket_updates_do_not_refreeze_each_repeated_identity(monkeypatch: pytest.MonkeyPatch, values: tuple[Any, Any]) -> None: + old = pd.DataFrame({"bucket": list(values) * 100}, index=pd.Index([f"o{i}" for i in range(200)])) + index = runtime._CandidateBucketIndex.build(old, ("bucket",)) + removed = old.iloc[:100] + added = removed.copy() + added.index = pd.Index([f"n{i}" for i in range(100)]) + calls: list[tuple[Any, ...]] = [] + original = runtime._row_key + + def key(row: tuple[Any, ...]) -> tuple[Any, ...]: + calls.append(row) + return original(row) + + monkeypatch.setattr(runtime, "_row_key", key) + updated = index.updated(removed, added) + assert len(calls) <= 4 + assert updated.row_count == 200 + + +def test_complete_withdrawal_builds_only_new_members(monkeypatch: pytest.MonkeyPatch) -> None: + old = pd.DataFrame({"bucket": [1, 2, 3]}, index=pd.Index(["a", "b", "c"])) + index = runtime._CandidateBucketIndex.build(old, ("bucket",)) + added = pd.DataFrame({"bucket": [4, 5]}, index=pd.Index(["d", "e"])) + original = runtime._row_key + rows: list[tuple[Any, ...]] = [] + + def key(row: tuple[Any, ...]) -> tuple[Any, ...]: + rows.append(row) + return original(row) + + monkeypatch.setattr(runtime, "_row_key", key) + updated = index.updated(old, added) + assert rows == [(4,), (5,)] + assert updated.row_count == 2 + assert index.row_count == 3 + + +def test_scalar_key_reuse_preserves_type_distinctions_and_bounded_work() -> None: + class Text(str): + pass + + values = ["1", Text("1"), 1, True, 1.0, None, pd.NA, {"a": [1]}] + [str(i) for i in range(200)] + frame = pd.DataFrame({"key": values}, index=pd.Index([f"p{i}" for i in range(len(values))])) + index = runtime._CandidateBucketIndex.build(frame, ("key",)) + expected: dict = {} + for occurrence, value in zip(frame.index, values, strict=True): + expected.setdefault(runtime._row_key((value,)), []).append(occurrence) + assert {k: list(v) for k, v in index.buckets.items()} == expected + + +def test_bucket_memo_does_not_alias_temporary_timestamp_objects() -> None: + frame = pd.DataFrame({"key": pd.date_range("2020-01-01", periods=400, tz="UTC")}, + index=pd.Index([f"p{i}" for i in range(400)])) + index = runtime._CandidateBucketIndex.build(frame, ("key",)) + assert len(index.buckets) == 400 + for occurrence, row in zip(frame.index, frame.itertuples(index=False, name=None), strict=True): + assert list(index.buckets[runtime._row_key(row)]) == [occurrence] + + +def test_warm_lookup_visits_no_unaffected_bucket_or_frame_rows() -> None: + frame = pd.DataFrame({"bucket": range(100), "value": range(100)}, index=pd.Index([f"p{i}" for i in range(100)])) + index = runtime._CandidateBucketIndex.build(frame, ("bucket",)) + + class NoScan(dict): + def __iter__(self) -> Any: + raise AssertionError("scanned all buckets") + + def items(self) -> Any: + raise AssertionError("scanned all buckets") + + def values(self) -> Any: + raise AssertionError("scanned all buckets") + + index.buckets = NoScan(index.buckets) + pd.testing.assert_frame_equal(index.select(frame, {runtime._row_key((5,))}), frame.iloc[[5]]) + + +def test_full_bucket_coverage_needs_no_member_merge(monkeypatch: pytest.MonkeyPatch) -> None: + frame = pd.DataFrame({"bucket": [2, 1, 2], "value": [0, 1, 2]}, index=pd.Index(["a", "b", "c"])) + index = runtime._CandidateBucketIndex.build(frame, ("bucket",)) + + def unnecessary(*args: Any, **kwargs: Any) -> Any: + raise AssertionError("full coverage must not merge all occurrence IDs") + + monkeypatch.setattr(runtime, "merge", unnecessary) + selected = index.select(frame, set(index.buckets)) + pd.testing.assert_frame_equal(selected, frame) + selected.loc["a", "value"] = 99 + assert frame.loc["a", "value"] == 0 + + +def test_bucket_keys_preserve_nulls_nested_values_and_duplicate_occurrences() -> None: + frame = pd.DataFrame({"key": [None, {"nested": [1, None]}, pd.NA, {"nested": [1, None]}], + "value": ["first", "second", "third", "fourth"]}, index=pd.Index(["a", "b", "c", "d"])) + index = runtime._CandidateBucketIndex.build(frame, ("key",)) + pd.testing.assert_frame_equal(index.select(frame, {runtime._row_key((float("nan"),))}), frame.iloc[[0, 2]]) + removed = frame.iloc[[1]] + added = frame.iloc[[1]].copy() + added.index = pd.Index(["e"]) + added["value"] = "current metadata" + current = pd.concat([frame.drop(index="b"), added]) + updated = index.updated(removed, added) + chosen = updated.select(current, {runtime._row_key(({"nested": [1, None]},))}) + assert list(chosen.index) == ["d", "e"] + assert list(chosen["value"]) == ["fourth", "current metadata"] + + +@pytest.mark.parametrize("direction", ["right-to-left", "left-to-right", "symmetric"]) +def test_runtime_index_syncs_metadata_and_rebuilds_after_restore(monkeypatch: pytest.MonkeyPatch, direction: str) -> None: + _, _, flow, _, _, _, calls = build(monkeypatch, direction) + initial = pd.DataFrame({"id": range(7), "fact": list("ABCDEFG"), "episode": ["original"] * 7}) + flow.apply(initial) + executor = flow._executor + count = len(calls) + flow.apply(initial.iloc[[0]].assign(episode="new metadata")) + assert len(calls) == count + indexes = executor._candidate_bucket_indexes + assert bool(indexes) == (direction != "symmetric") + for (parent, columns), index in indexes.items(): + frame = executor.node_state[parent] + rebuilt = runtime._CandidateBucketIndex.build(frame, columns) + assert set(index.buckets) == set(rebuilt.buckets) + pd.testing.assert_frame_equal(index.select(frame, set(index.buckets)), frame) + assert all(list(index.buckets[k]) == list(rebuilt.buckets[k]) for k in index.buckets) + snapshot = flow.snapshot_state() + assert "candidate_bucket_indexes" not in snapshot + flow.restore_state(snapshot) + assert executor._candidate_bucket_indexes == {} + flow.apply(initial.iloc[[0]].assign(episode="after restore")) + assert len(calls) == count + assert bool(executor._candidate_bucket_indexes) == (direction != "symmetric") + + +@pytest.mark.parametrize("failure", ["model", "storage"]) +def test_failed_step_never_commits_bucket_changes(monkeypatch: pytest.MonkeyPatch, failure: str) -> None: + _, _, flow, _, _, _, _ = build(monkeypatch) + initial = pd.DataFrame({"id": [0, 2, 3, 4, 5, 6, 7], "fact": list("ABCDEFG"), "episode": ["original"] * 7}) + flow.apply(initial) + executor = flow._executor + before = {key: deepcopy(index.buckets) for key, index in executor._candidate_bucket_indexes.items()} + original = getattr(pd.DataFrame, "sem_filter") if failure == "model" else executor._write_storage_updates + + def fail(*args: Any, **kwargs: Any) -> Any: + raise RuntimeError("injected failure") + + target, name = (pd.DataFrame, "sem_filter") if failure == "model" else (executor, "_write_storage_updates") + monkeypatch.setattr(target, name, fail) + extra = pd.DataFrame({"id": [1], "fact": ["H"], "episode": ["new"]}) + with pytest.raises(RuntimeError, match="injected failure"): + flow.apply(extra) + assert {k: i.buckets for k, i in executor._candidate_bucket_indexes.items()} == before + monkeypatch.setattr(target, name, original) + flow.apply(extra) + for (parent, _), index in executor._candidate_bucket_indexes.items(): + frame = executor.node_state[parent] + pd.testing.assert_frame_equal(index.select(frame, set(index.buckets)), frame) + + +def test_full_runtime_matches_scan_path_including_calls_and_private_ids(monkeypatch: pytest.MonkeyPatch) -> None: + _, _, optimized, _, _, _, _ = build(monkeypatch) + _, _, reference, _, _, _, calls = build(monkeypatch) + prompts: list[pd.DataFrame] = [] + oracle = getattr(pd.DataFrame, "sem_filter") + + def observe(frame: pd.DataFrame, *args: Any, **kwargs: Any) -> pd.DataFrame: + prompts.append(frame.copy(deep=True)) + return oracle(frame, *args, **kwargs) + + monkeypatch.setattr(pd.DataFrame, "sem_filter", observe) + updates = [ + pd.DataFrame({"id": [0, 2, 3, 4, 5, 6, 7], "fact": list("ABCDEFG"), "episode": ["original"] * 7}), + pd.DataFrame({"id": [0], "fact": ["A"], "episode": ["metadata"]}), + pd.DataFrame({"id": [1], "fact": ["H"], "episode": ["stronger"]}), + pd.DataFrame({"id": [2, 9], "fact": ["B", "H"], "episode": ["again", "new"]}), + ] + semantic = runtime.PolicyExecutor._execute_semantic_row + + def scan(self: Any, *args: Any, **kwargs: Any) -> pd.DataFrame: + kwargs["staged_bucket_updates"] = None + return semantic(self, *args, **kwargs) + + def separate(old: pd.DataFrame, new: pd.DataFrame, *, allocate: Any) -> runtime.NodeOutputUpdate: + return runtime.NodeOutputUpdate.between(old, runtime.preserve_occurrence_index(old, new, allocate=allocate)) + + for rows in updates: + calls.clear() + prompts.clear() + optimized.apply(rows) + expected_calls = list(calls) + expected_prompts = list(prompts) + calls.clear() + prompts.clear() + with monkeypatch.context() as patch: + patch.setattr(runtime.PolicyExecutor, "_execute_semantic_row", scan) + patch.setattr(runtime, "_preserve_occurrences_and_diff", separate) + reference.apply(rows) + assert calls == expected_calls + assert len(prompts) == len(expected_prompts) + for actual, expected in zip(prompts, expected_prompts, strict=True): + pd.testing.assert_frame_equal(actual, expected) + assert optimized._executor._next_occurrence == reference._executor._next_occurrence + for node_id, frame in optimized._executor.node_state.items(): + pd.testing.assert_frame_equal(frame, reference._executor.node_state[node_id]) + + +def test_two_predicates_share_one_parent_index_update(monkeypatch: pytest.MonkeyPatch) -> None: + source, view, _, adapter, _, _, _ = build(monkeypatch) + state = source.group_by(["id", "fact"]).array_agg(columns=["episode"], output_col="provenance") + left, right = state.alias("earlier"), state.alias("later") + other = left.join(right, on=left.col("id") < right.col("id")).sem_filter( + instruction="Again compare {fact:earlier} and {fact:later}.") + profile = next(iter(adapter.config.semantic_pair_profiles.values())) + unprepared = am.SemanticDataflow(source=source, views={"first": view, "second": other})._executor.policy + for node in unprepared.nodes.values(): + for root in (node.query, node.maintenance_query): + if root is not None: + for query in walk(root): + if query.op == "sem_filter": + adapter.config.semantic_pair_profiles[query_digest(query)] = profile + flow = am.SemanticDataflow(source=source, views={"first": view, "second": other}, adapter=adapter) + updated = runtime._CandidateBucketIndex.updated + counts: list[int] = [] + + def record(self: Any, removed: pd.DataFrame, added: pd.DataFrame) -> Any: + counts.append(len(added)) + return updated(self, removed, added) + + monkeypatch.setattr(runtime._CandidateBucketIndex, "updated", record) + initial = pd.DataFrame({"id": range(7), "fact": list("ABCDEFG"), "episode": ["original"] * 7}) + flow.apply(initial) + flow.apply(initial.iloc[[0]].assign(episode="metadata")) + assert len(counts) == 2 + assert len(flow._executor._candidate_bucket_indexes) == 1 + pd.testing.assert_frame_equal(flow.view("first"), flow.view("second")) diff --git a/tests/test_site_batching.py b/tests/test_site_batching.py new file mode 100644 index 0000000..3db02fe --- /dev/null +++ b/tests/test_site_batching.py @@ -0,0 +1,186 @@ +"""Site-only batching preserves screening boundaries and pair occurrences.""" +from dataclasses import replace +import json +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.site_batching import PairFilterBatching +from agent_memory.memories.zep.policy import ZepMemory, _CONTRADICTORY_FACT_INSTRUCTION +from agent_memory.policy.logical import QueryExpr + + +def setting(size: int = 4) -> PairFilterBatching: + return PairFilterBatching(("fact_id:later_added",), PromptBatching(max_tasks=size)) + + +def query() -> QueryExpr: + return QueryExpr(op="sem_filter", inputs=(QueryExpr(op="materialized_view", params={ + "name": "pairs", "columns": ("fact:earlier_added", "fact:later_added", "fact_id:later_added", "occurrence")}),), + params={"instruction": _CONTRADICTORY_FACT_INSTRUCTION}) + + +class Provider: + max_tokens = 512 + cache = None + + def __init__(self, invalid: bool = False) -> None: + self.calls: list[Any] = [] + self.invalid = invalid + + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + outputs = [] + for prompt in prompts: + data = json.loads(prompt[1]["content"]) + self.calls.append(data) + rows = data["rows"] + assert len(rows) <= 4 + outputs.append(json.dumps({"decisions": [{"row_id": row["row_id"], "keep": True} + for row in reversed(rows[1:] if self.invalid else rows)]})) + return SimpleNamespace(outputs=outputs) + + +def adapter(monkeypatch: pytest.MonkeyPatch, *, invalid: bool = False) -> tuple[LotusAdapter, Provider]: + import lotus + lm = Provider(invalid) + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + a = LotusAdapter(config=LotusExecutionConfig( + pair_filter_batching={semantic_pair_site_id(query()): setting()}, structured_parse_retries=0)) + monkeypatch.setattr(a._context, "configure", lambda: None) + return a, lm + + +def rows() -> pd.DataFrame: + return pd.DataFrame({"fact:earlier_added": ["same old"] * 7, + "fact:later_added": ["new A", "new B", "new A", "new A", "new A", "new A", "new B"], + "fact_id:later_added": [(1, 0), (2, 0), (1, 0), (1, 0), (1, 0), (1, 0), (2, 0)], + "occurrence": list(range(7))}, index=pd.Index([0] * 7)) + + +def test_group_boundary_split_order_and_duplicate_occurrences(monkeypatch: pytest.MonkeyPatch) -> None: + a, lm = adapter(monkeypatch) + data = rows() + result = a.execute(query(), {"pairs": data}) + pd.testing.assert_frame_equal(result, data) + assert [len(call["rows"]) for call in lm.calls] == [4, 1, 2] + assert all(not ("new A" in str(call) and "new B" in str(call)) for call in lm.calls) + assert len(a.execute(query(), {"pairs": data.iloc[:0]})) == 0 + assert len(lm.calls) == 3 + + +def test_missing_decision_fails_without_retry(monkeypatch: pytest.MonkeyPatch) -> None: + a, lm = adapter(monkeypatch, invalid=True) + with pytest.raises(ValueError): + a.execute(query(), {"pairs": rows().iloc[:1]}) + assert len(lm.calls) == 1 + + +def test_unconfigured_predicate_remains_native(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + a, lm = adapter(monkeypatch) + assert lotus.settings.lm is lm + changed = replace(query(), params={"instruction": "Is {fact:earlier_added} true?"}) + calls = [] + monkeypatch.setattr(pd.DataFrame, "sem_filter", lambda frame, *args, **kw: calls.append(args) or frame.copy()) + a.execute(changed, {"pairs": rows()}) + assert len(calls) == 1 and not lm.calls + + +def test_zep_plan_binding_and_identity() -> None: + config = LotusExecutionConfig(physical_fusion="zep-target-state", + pair_filter_batching={semantic_pair_site_id(query()): setting()}) + a = LotusAdapter(config=config) + plan = a.prepare_policy(ZepMemory.differentiate_policy()) + assert sum(n.maintenance_query is not None and n.maintenance_query.op == "fused_target_state" + for n in plan.nodes.values()) == 1 + plain = LotusAdapter(config=LotusExecutionConfig(physical_fusion="zep-target-state")) + assert a.maintenance_execution_fingerprint != plain.maintenance_execution_fingerprint + assert a.maintenance_execution_fingerprint != LotusAdapter(config=replace(config, + pair_filter_batching={semantic_pair_site_id(query()): setting(2)})).maintenance_execution_fingerprint + with pytest.raises(ValueError, match="not found"): + LotusAdapter(config=replace(config, pair_filter_batching={"sem_filter:missing": setting()})).prepare_policy(ZepMemory.differentiate_policy()) + with pytest.raises(ValueError): + replace(config, prompt_batching=PromptBatching(max_tasks=4)) + + +def test_real_runtime_maintenance_and_restore(monkeypatch: pytest.MonkeyPatch) -> None: + a, lm = adapter(monkeypatch) + source = am.Source({c: c for c in rows().columns}) + view = source.sem_filter(instruction=_CONTRADICTORY_FACT_INSTRUCTION) + flow = am.SemanticDataflow(source=source, views={"result": view}, adapter=a) + flow.apply(rows().iloc[:2]) + flow.apply(rows().iloc[2:]) + assert flow.view("result").occurrence.tolist() == list(range(7)) + assert len(lm.calls) == 4 + snapshot = flow.snapshot_state() + plain = am.SemanticDataflow(source=source, views={"result": view}, adapter=LotusAdapter()) + with pytest.raises(ValueError, match="fingerprint"): + plain.restore_state(snapshot) + + +@pytest.mark.parametrize("aggregate_batching", [None, PromptBatching(max_tasks=4)]) +def test_fusion_and_site_batching_in_one_real_dataflow(monkeypatch: pytest.MonkeyPatch, aggregate_batching: PromptBatching | None) -> None: + import lotus + from agent_memory.policy.relation import Relation + from agent_memory.policy.schema import output_columns + from agent_memory.policy.aggregates import SemanticAggregateSpec + from test_physical_fusion import state, fact_input + + calls = {"fusion": 0, "predicate": 0} + class Model: + max_tokens = 512 + cache = None + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + outputs = [] + for prompt in prompts: + payload = json.loads(prompt[1]["content"]) + if "incoming" in payload: + calls["fusion"] += 1 + outputs.append(json.dumps({"matches": [{"left_id": r["id"], "right_id": None} + for r in payload["incoming"]], "states": []})) + else: + calls["predicate"] += 1 + outputs.append(json.dumps({"decisions": [{"row_id": r["row_id"], "keep": True} + for r in payload["rows"]]})) + return SimpleNamespace(outputs=outputs) + + class OracleAdapter(LotusAdapter): + def execute(self, query: QueryExpr, inputs: Any) -> Any: + q = query + if q.op == "agg" and q.inputs[0].op == "sem_groupby": + frame = self.execute(q.inputs[0].inputs[0], inputs) + return pd.DataFrame([state(r["content"], r["add_seq"]) for _, r in frame.iterrows()], + columns=pd.Index(output_columns(q))) + if q.op == "agg": + assert not any(isinstance(s, SemanticAggregateSpec) for s in q.params["aggregates"]) + assert q.op != "sem_join", "fused matching was repeated" + return super().execute(q, inputs) + + source = am.Source({c: c for c in fact_input("old", 0).columns}) + original = ZepMemory._deduplicated_facts.expr + agg = original.inputs[0] + facts = Relation(replace(original, inputs=(replace(agg, + inputs=(replace(agg.inputs[0], inputs=(source.expr,)),)),))) + earlier, later = facts.alias("earlier_added"), facts.alias("later_added") + pairs = earlier.join(later, on=earlier.col("fact_id") < later.col("fact_id")) + view = pairs.sem_filter(instruction=_CONTRADICTORY_FACT_INSTRUCTION) + a = OracleAdapter(config=LotusExecutionConfig(physical_fusion="zep-target-state", structured_parse_retries=0, + sem_agg_prompt_batching=aggregate_batching, + pair_filter_batching={semantic_pair_site_id(view.expr): setting()})) + monkeypatch.setattr(a._context, "configure", lambda: None) + monkeypatch.setattr(lotus.settings, "lm", Model()) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + flow = am.SemanticDataflow(source=source, views={"facts": facts, "pairs": view}, adapter=a) + for i in range(3): + flow.apply(fact_input(f"fact-{i}", i)) + assert len(flow.view("facts")) == 3 and len(flow.view("pairs")) == 3 + assert calls["fusion"] == 2 + assert calls["predicate"] >= 2 diff --git a/tests/test_site_listwise_batching.py b/tests/test_site_listwise_batching.py new file mode 100644 index 0000000..8421e8a --- /dev/null +++ b/tests/test_site_listwise_batching.py @@ -0,0 +1,201 @@ +"""Site-scoped listwise packing preserves each anchor's candidate domain.""" + +import json +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.sem_join import execute_sem_join +from test_sem_topk_join import _query +from test_zep_clean_maintenance import offline_context +from agent_memory.policy.logical import QueryExpr + + +def site_query() -> QueryExpr: + query = _query(k=1, on=("tenant",)) + return QueryExpr(query.op, query.inputs, {**query.params, "instruction": "{name:left} and {name:right} describe the same entity."}) + + +class PackedLM: + max_ctx_len = 100_000 + max_tokens = 1024 + cache = None + + def __init__(self, invalid: str = "") -> None: + self.requests: list[dict[str, Any]] = [] + self.invalid = invalid + + def count_tokens(self, messages: Any) -> int: + return len(str(messages)) + + def __call__(self, messages: Any, **kwargs: Any) -> Any: + outputs = [] + for prompt in messages: + payload = json.loads(prompt[1]["content"]) + self.requests.append(payload) + tasks = payload["tasks"] + results = [{"task_id": t["task_id"], "selected_ids": [t["right_candidates"][0]["id"]]} + for t in reversed(tasks)] + if self.invalid == "missing": + results.pop() + elif self.invalid == "duplicate": + results.append(results[0]) + elif self.invalid == "cross-task": + results[0]["selected_ids"] = [tasks[0]["right_candidates"][0]["id"]] + outputs.append(json.dumps({"results": results})) + return SimpleNamespace(outputs=outputs) + + +def test_site_packs_anchors_without_changing_candidates(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + + query = site_query() + site = semantic_pair_site_id(query) + config = LotusExecutionConfig(listwise_join_batching={site: PromptBatching(max_tasks=16)}) + left = pd.DataFrame({"tenant": list(range(17)), "name": ["same"] * 17}, index=range(100, 117)) + right = pd.DataFrame({"tenant": list(range(17)), "name": ["same"] * 17}, index=range(200, 217)) + lm = PackedLM() + monkeypatch.setattr(lotus.settings, "lm", lm) + output = execute_sem_join(query, {}, lambda q, _: left if q.params["name"] == "left" else right, + offline_context(config, monkeypatch)) + assert len(output) == 17 # Identical text never deduplicates occurrences. + assert [len(p["tasks"]) for p in lm.requests] == [16, 1] + for request in lm.requests: + assert request["join_condition"] == query.params["instruction"] + assert all(len(t["right_candidates"]) == 1 for t in request["tasks"]) + assert all("messages" not in t and "join_condition" not in t for t in request["tasks"]) + assert config.prompt_batching is None + # A different semantic site still gets the original unbatched path. + assert site not in LotusExecutionConfig().listwise_join_batching + + +@pytest.mark.parametrize("invalid", ["missing", "duplicate", "cross-task"]) +def test_site_rejects_invalid_or_cross_task_results(monkeypatch: pytest.MonkeyPatch, invalid: str) -> None: + import lotus + + query = site_query() + config = LotusExecutionConfig(listwise_join_batching={semantic_pair_site_id(query): PromptBatching(max_tasks=16)}, + structured_parse_retries=0) + rows = pd.DataFrame({"tenant": [1, 2], "name": ["same", "same"]}) + monkeypatch.setattr(lotus.settings, "lm", PackedLM(invalid)) + with pytest.raises(ValueError): + execute_sem_join(query, {}, lambda *_: rows, offline_context(config, monkeypatch)) + + +def test_site_execution_identity_and_plan_validation() -> None: + from agent_memory.memories.zep.representative import ZepRepresentativeMemory + from agent_memory.planner.physical import walk + from agent_memory.runtime.executor import PolicyExecutor + + policy = ZepRepresentativeMemory.differentiate_policy() + fact_join = next(q for n in policy.nodes.values() if n.maintenance_query is not None + for q in walk(n.maintenance_query) if q.op == "sem_join" and q.params.get("on")) + site = semantic_pair_site_id(fact_join) + old = LotusAdapter(config=LotusExecutionConfig(physical_fusion="zep-representative")) + packed = LotusAdapter(config=LotusExecutionConfig(physical_fusion="zep-representative", + listwise_join_batching={site: PromptBatching(max_tasks=16)})) + packed.prepare_policy(policy) + assert old.maintenance_execution_fingerprint != packed.maintenance_execution_fingerprint + with pytest.raises(ValueError, match="fingerprint"): + PolicyExecutor(policy, adapter=packed).restore_state(PolicyExecutor(policy, adapter=old).snapshot_state()) + invalid = LotusAdapter(config=LotusExecutionConfig(listwise_join_batching={"sem_join:missing": PromptBatching(max_tasks=16)})) + with pytest.raises(ValueError, match="not found"): + invalid.prepare_policy(policy) + + +def test_site_rejects_global_or_pairwise_configuration() -> None: + setting = {"sem_join:test": PromptBatching(max_tasks=16)} + with pytest.raises(ValueError): + LotusExecutionConfig(listwise_join_batching=setting, prompt_batching=PromptBatching()) + with pytest.raises(ValueError): + LotusExecutionConfig(listwise_join_batching=setting, sem_join_topk_method="pairwise-naive") + + +def test_other_sites_and_empty_candidates_do_not_pack(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + from test_sem_topk_join import _BatchListwiseLM + + config = LotusExecutionConfig(listwise_join_batching={"sem_join:other": PromptBatching(max_tasks=16)}) + rows = pd.DataFrame({"name": ["x", "y"]}) + lm = _BatchListwiseLM(['{"selected_ids": []}', '{"selected_ids": []}']) + monkeypatch.setattr(lotus.settings, "lm", lm) + result = execute_sem_join(_query(k=1), {}, lambda *_: rows, offline_context(config, monkeypatch)) + assert result.empty + assert len(lm.calls[0][0]) == 2 + assert all("tasks" not in json.loads(prompt[1]["content"]) for prompt in lm.calls[0][0]) + config = LotusExecutionConfig(listwise_join_batching={semantic_pair_site_id(site_query()): PromptBatching(max_tasks=16)}) + empty_lm = PackedLM() + monkeypatch.setattr(lotus.settings, "lm", empty_lm) + for empty in (False, True): + left = pd.DataFrame({"tenant": [1], "name": ["x"]}) + right = pd.DataFrame({"tenant": [2], "name": ["x"]}) + if empty: + left = left.iloc[:0] + result = execute_sem_join(site_query(), {}, lambda q, _: left if q.params["name"] == "left" else right, + offline_context(config, monkeypatch)) + assert result.empty + assert not empty_lm.requests + + +def test_real_runtime_uses_packed_join_and_restores(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + import agent_memory as am + from agent_memory.adapters.lotus import adapter as adapter_module + from agent_memory.adapters.lotus.sem_groupby import GROUP_ID_COLUMN + from agent_memory.memories.zep.representative import ZepRepresentativeMemory + from agent_memory.planner.physical import replace_query, walk + from agent_memory.policy.aggregates import ArgMinAggregateSpec + from agent_memory.policy.relation import Relation + from agent_memory.runtime.executor import NodeOutputUpdate + + def group(query: Any, inputs: Any, execute: Any, context: Any) -> pd.DataFrame: + frame = execute(query.inputs[0], inputs).copy() + frame[GROUP_ID_COLUMN] = pd.factorize(frame.fact)[0] + frame.attrs["agent_memory_groupby_input_cols"] = tuple(query.params["input_cols"]) + frame.attrs["agent_memory_sem_groupby_partition_by"] = tuple(query.params.get("partition_by", ())) + return frame + + monkeypatch.setattr(adapter_module, "execute_sem_groupby", group) + aggregate = next(q for q in walk(ZepRepresentativeMemory.facts.expr) + if any(isinstance(s, ArgMinAggregateSpec) for s in q.params.get("aggregates", ()))) + raw = dict(source_entity_id="a", target_entity_id="b", relation_type="LIKES", fact="Alice likes tea.", + episode_id="0", fact_ordinal=0, add_seq=0, created_at="2026-01-01", + valid_at=None, invalid_at=None, expired_at=None, content="original") + source = am.Source({c: c for c in raw}) + view = Relation(replace_query(aggregate, aggregate.inputs[0].inputs[0], source.expr)) + # Obtain the actual generated site, rather than guessing a digest. + full_policy = ZepRepresentativeMemory.differentiate_policy() + site = next(semantic_pair_site_id(q) for n in full_policy.nodes.values() if n.maintenance_query is not None + for q in walk(n.maintenance_query) if q.op == "sem_join" and q.params.get("on")) + adapter = LotusAdapter(config=LotusExecutionConfig(listwise_join_batching={site: PromptBatching(max_tasks=16)})) + flow = am.SemanticDataflow(source=source, views={"facts": view}, adapter=adapter) + monkeypatch.setattr(adapter._context, "configure", lambda: None) + lm = PackedLM() + monkeypatch.setattr(lotus.settings, "lm", lm) + rows: list[dict[str, Any]] = [] + for seq in range(2): + batch = [{**raw, "episode_id": str(seq), "add_seq": seq, "valid_at": f"2026-01-0{seq + 1}"}, + {**raw, "source_entity_id": "c", "target_entity_id": "d", "fact": "Carol likes coffee.", + "episode_id": str(seq), "add_seq": seq, "fact_ordinal": 1, "valid_at": f"2026-01-0{seq + 1}"}] + rows.extend(batch) + flow.apply(pd.DataFrame(batch)) + actual = flow.view("facts") + expected = adapter.execute(view.expr, {"log": pd.DataFrame(rows)}) + for frame in (actual, expected): + frame["provenance"] = frame.provenance.map(lambda s: json.dumps(sorted(json.loads(s), key=lambda r: r["episode_id"]), sort_keys=True)) + assert NodeOutputUpdate.between(actual, expected).is_empty + restored = am.SemanticDataflow(source=source, views={"facts": view}, adapter=adapter) + restored.restore_state(flow.snapshot_state()) + flow = restored + assert [len(request["tasks"]) for request in lm.requests] == [2] + for task in lm.requests[0]["tasks"]: + incoming = json.dumps(task["left"]) + previous = json.dumps(task["right_candidates"]) + assert "valid_at" in incoming and "2026-01-02" in incoming + assert "valid_at" in previous and "2026-01-01" in previous diff --git a/tests/test_zep_clean_maintenance.py b/tests/test_zep_clean_maintenance.py new file mode 100644 index 0000000..4a6985f --- /dev/null +++ b/tests/test_zep_clean_maintenance.py @@ -0,0 +1,366 @@ +"""Offline checks for bounded maintenance shortcuts and real prompt packing.""" + +import json +import math +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig, LotusExecutionContext +from agent_memory.adapters.lotus.fusion import _unchanged_entity_states +from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.sem_groupby import GROUP_ID_COLUMN, execute_sem_groupby +from agent_memory.memories.zep.representative import ZepRepresentativeMemory +from agent_memory.planner.physical import replace_query, walk +from agent_memory.policy.aggregates import ArgMinAggregateSpec +from agent_memory.policy.logical import QueryExpr +from agent_memory.policy.relation import Relation + + +class Oracle: + """Independent false predicates, exact entity assignments, and no fact matches.""" + + max_tokens = 1024 + max_ctx_len = 100_000 + cache = None + + def __init__(self, invalid: str = "") -> None: + self.requests: list[dict[str, Any]] = [] + self.invalid = invalid + + def count_tokens(self, messages: Any) -> int: + return len(str(messages)) + + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + outputs = [] + for prompt in prompts: + data = json.loads(prompt[1]["content"]) + self.requests.append(data) + if "rows" in data: + decisions = [{"row_id": row["row_id"], "keep": False} for row in reversed(data["rows"])] + if self.invalid == "missing": + decisions.pop() + elif self.invalid == "duplicate": + decisions.append(decisions[0]) + elif self.invalid == "boolean": + decisions[0]["keep"] = "false" + outputs.append(json.dumps({"decisions": decisions})) + elif "incoming" in data: + matches = [{"left_id": row["id"], "right_id": row["eligible_targets"][0]} + for row in data["incoming"]] + targets = {m["right_id"] for m in matches} + states = [{"right_id": row["id"], "name": row["state"]["name"]} + for row in data["targets"] if row["id"] in targets] + outputs.append(json.dumps({"matches": matches, "states": states})) + else: + outputs.append('{"selected_ids": []}') + return SimpleNamespace(outputs=outputs) + + +def grouping(partitioned: bool = False) -> QueryExpr: + return QueryExpr("sem_groupby", (QueryExpr("materialized_view", params={"name": "rows"}),), + {"input_cols": ("name",), "instruction": "The same entity: {name}.", + "partition_by": ("tenant",) if partitioned else ()}) + + +def offline_context(config: LotusExecutionConfig, monkeypatch: pytest.MonkeyPatch) -> LotusExecutionContext: + context = LotusExecutionContext(model="fake", config=config) + monkeypatch.setattr(context, "configure", lambda: None) + return context + + +def test_groupby_prompt16_preserves_partitions_duplicates_and_order(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + import lotus + from agent_memory.adapters.lotus.identity_reuse import identity_scope + + query = grouping(True) + config = LotusExecutionConfig(groupby_prompt_batching={semantic_pair_site_id(query): PromptBatching(max_tasks=16)}, + sem_groupby_pair_batch_size=32, semantic_trace_dir=tmp_path) + rows = pd.DataFrame({"tenant": [0] * 8 + [1] * 2, + "name": [f"first-{i}" for i in range(7)] + ["first-0", "second-a", "second-b"], + "metadata": list(range(10))}) + oracle = Oracle() + monkeypatch.setattr(lotus.settings, "lm", oracle) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + context = offline_context(config, monkeypatch) + cache: dict[str, bool] = {} + with identity_scope(cache): + result = execute_sem_groupby(query, {}, lambda *_: rows, context) + assert [len(r["rows"]) for r in oracle.requests] == [16, 5, 1] + assert all(not ("first-" in str(r) and "second-" in str(r)) for r in oracle.requests) + assert result.metadata.tolist() == list(range(10)) + assert result[GROUP_ID_COLUMN].nunique() == 9 + assert result.iloc[0][GROUP_ID_COLUMN] == result.iloc[7][GROUP_ID_COLUMN] + assert config.prompt_batching is None and config.sem_groupby_pair_batch_size == 32 + rows.metadata += 100 + reused = execute_sem_groupby(query, {}, lambda *_: rows, context) + assert len(oracle.requests) == 3 + assert reused.metadata.tolist() == list(range(100, 110)) + events = [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()] + batching = [event for event in events if event["event_type"] == "prompt_batching"] + assert sum(event["prompt_count"] for event in batching) == 3 + assert sum(event["task_count"] for event in batching) == 22 + assert [size for event in batching for size in event["chunk_sizes"]] == [16, 5, 1] + + +@pytest.mark.parametrize("invalid", ["missing", "duplicate", "boolean"]) +def test_groupby_invalid_results_fail(monkeypatch: pytest.MonkeyPatch, invalid: str) -> None: + import lotus + query = grouping() + config = LotusExecutionConfig(groupby_prompt_batching={semantic_pair_site_id(query): PromptBatching(max_tasks=16)}, + structured_parse_retries=0) + oracle = Oracle(invalid) + monkeypatch.setattr(lotus.settings, "lm", oracle) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + with pytest.raises(ValueError): + execute_sem_groupby(query, {}, lambda *_: pd.DataFrame({"name": ["a", "b", "c"]}), offline_context(config, monkeypatch)) + assert len(oracle.requests) == 1 + + +def test_empty_and_unselected_grouping_keep_original_behavior(monkeypatch: pytest.MonkeyPatch) -> None: + from agent_memory.adapters.lotus import sem_groupby + query = grouping() + seen: list[int] = [] + def evaluate(*args: Any, **kwargs: Any) -> Any: + seen.append(len(kwargs["docs"])) + return SimpleNamespace(outputs=[False] * len(kwargs["docs"])) + monkeypatch.setattr(sem_groupby, "evaluate_group_match_batch", evaluate) + config = LotusExecutionConfig(groupby_prompt_batching={"sem_groupby:other": PromptBatching(max_tasks=16)}, + sem_groupby_pair_batch_size=2) + rows = pd.DataFrame({"name": ["a", "b", "c"]}) + execute_sem_groupby(query, {}, lambda *_: rows, offline_context(config, monkeypatch)) + assert seen == [2, 1] + empty_config = replace(config, groupby_prompt_batching={semantic_pair_site_id(query): PromptBatching(max_tasks=16)}) + result = execute_sem_groupby(query, {}, lambda *_: rows.iloc[:0], offline_context(empty_config, monkeypatch)) + assert result.empty and seen == [2, 1] + + +@pytest.mark.parametrize("names,fixed,expected", [ + (["Alice", "Alice", "new"], {"l0": "r0", "l1": "r0", "l2": None}, {"r0": {"name": "Alice"}}), + (["Alice", "Alice Smith", "new"], {"l0": "r0", "l1": "r0", "l2": None}, None), + (["Alice", "Alice", "new"], {"l0": "r0", "l2": None}, None), + ([None, "Alice", "new"], {"l0": "r0", "l1": "r0", "l2": None}, None), +]) +def test_name_shortcut_requires_every_decision_and_identical_names(names: list[Any], fixed: dict[str, Any], expected: Any) -> None: + left = pd.DataFrame({"name": names}, index=pd.Index([10, 11, 12])) + right = pd.DataFrame({"name": ["Alice", "unrelated"]}, index=pd.Index([20, 21])) + assert _unchanged_entity_states({"l0": {"r0"}, "l1": {"r0"}, "l2": {"r1"}}, fixed, + left, right, {"l0": 10, "l1": 11, "l2": 12}, {"r0": 20, "r1": 21}) == expected + + +def make_identity_flow(monkeypatch: pytest.MonkeyPatch, oracle: Oracle, enabled: bool, + trace_dir: Path, packed: bool = False) -> tuple[Any, dict[str, Any]]: + import lotus + from test_physical_fusion import state + raw = {**state("fact", 0), "name": "Alice", "entity_type": "Entity", "episode_id": "0", + "entity_ordinal": 0, "fact_ordinal": 0, "content": "original"} + source = am.Source({c: c for c in raw}) + identities = ZepRepresentativeMemory._identities.expr + aggregate = next(q for q in walk(identities) if q.op == "agg") + identities = replace_query(identities, aggregate.inputs[0].inputs[0], source.expr) + # Keep the actual representative fact aggregate in the registered plan. + facts = next(q for q in walk(ZepRepresentativeMemory.facts.expr) + if any(isinstance(s, ArgMinAggregateSpec) for s in q.params.get("aggregates", ()))) + facts = replace_query(facts, facts.inputs[0].inputs[0], source.expr) + sites = {semantic_pair_site_id(q): PromptBatching(max_tasks=16) + for view in (identities, facts) for q in walk(view) if q.op == "sem_groupby"} if packed else {} + adapter = LotusAdapter(config=LotusExecutionConfig(physical_fusion="zep-representative", + reuse_unchanged_entity_name=enabled, semantic_trace_dir=trace_dir, + groupby_prompt_batching=sites, sem_groupby_pair_batch_size=32)) + monkeypatch.setattr(adapter._context, "configure", lambda: None) + monkeypatch.setattr(lotus.settings, "lm", oracle) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + return am.SemanticDataflow(source=source, views={"entities": Relation(identities), "facts": Relation(facts)}, adapter=adapter), raw + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_real_identity_runtime_reuses_after_restore_and_propagates_metadata(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, enabled: bool) -> None: + oracle = Oracle() + flow, raw = make_identity_flow(monkeypatch, oracle, enabled, tmp_path) + for seq in range(3): + flow.apply(pd.DataFrame([{**raw, "add_seq": seq, "episode_id": str(seq), "content": f"source-{seq}"}])) + saved = flow.snapshot_state() + flow, _ = make_identity_flow(monkeypatch, oracle, enabled, tmp_path) + flow.restore_state(saved) + requests = [r for r in oracle.requests if "incoming" in r] + assert len(requests) == (1 if enabled else 2) + row = flow.view("entities").iloc[0] + assert row["name"] == "Alice" and row["entity_id"] == (0, 0) + assert len(json.loads(row["mentions"])) == 3 + assert "source-2" in row["mentions"] + if enabled: + assert "unchanged_entity_name" in (tmp_path / "events.jsonl").read_text() + # A different name must return to the model even after a match is cached. + for seq in (3, 4): + flow.apply(pd.DataFrame([{**raw, "name": "Alice Smith", "add_seq": seq, "episode_id": str(seq)}])) + assert len([r for r in oracle.requests if "incoming" in r]) == len(requests) + 2 + + +def test_new_options_reject_incompatible_restores_and_unbound_sites() -> None: + from agent_memory.runtime.executor import PolicyExecutor + policy = ZepRepresentativeMemory.differentiate_policy() + sites = {semantic_pair_site_id(q) for n in policy.nodes.values() for q in walk(n.query) if q.op == "sem_groupby"} + assert len(sites) == 2 + base = LotusExecutionConfig(physical_fusion="zep-representative", sem_groupby_pair_batch_size=32) + old = PolicyExecutor(policy, adapter=LotusAdapter(config=base)) + for config in (replace(base, reuse_unchanged_entity_name=True), + replace(base, groupby_prompt_batching={s: PromptBatching(max_tasks=16) for s in sites})): + current = PolicyExecutor(policy, adapter=LotusAdapter(config=config)) + with pytest.raises(ValueError, match="fingerprint"): + current.restore_state(old.snapshot_state()) + restored = PolicyExecutor(policy, adapter=LotusAdapter(config=config)) + restored.restore_state(current.snapshot_state()) + bad = LotusAdapter(config=replace(base, groupby_prompt_batching={"sem_groupby:missing": PromptBatching(max_tasks=16)})) + with pytest.raises(ValueError, match="not found"): + bad.prepare_policy(policy) + invalid_configs: tuple[dict[str, Any], ...] = ({"physical_fusion": "disabled", "reuse_unchanged_entity_name": True}, + {"reuse_unchanged_entity_name": 1}, + {"groupby_prompt_batching": {"sem_join:bad": PromptBatching(max_tasks=16)}}, + {"groupby_prompt_batching": {"sem_groupby:bad": PromptBatching()}}, + {"groupby_prompt_batching": {"sem_groupby:x": PromptBatching(max_tasks=16)}, "sem_groupby_default": True}, + {"groupby_prompt_batching": {"sem_groupby:x": PromptBatching(max_tasks=16)}, "prompt_batching": PromptBatching()}) + for kwargs in invalid_configs: + with pytest.raises((ValueError, TypeError)): + LotusExecutionConfig(**kwargs) + + +def test_real_runtime_packs_both_grouping_sites(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + oracle = Oracle() + flow, raw = make_identity_flow(monkeypatch, oracle, True, tmp_path, packed=True) + rows = pd.DataFrame([{**raw, "name": f"person-{i}", "fact": f"fact-{i}", + "entity_ordinal": i, "fact_ordinal": i} for i in range(7)]) + flow.apply(rows) + assert [len(r["rows"]) for r in oracle.requests] == [16, 5, 16, 5] + assert len(flow.view("entities")) == len(flow.view("facts")) == 7 + saved = flow.snapshot_state() + restored, _ = make_identity_flow(monkeypatch, oracle, True, tmp_path, packed=True) + restored.restore_state(saved) + pd.testing.assert_frame_equal(restored.view("entities"), flow.view("entities")) + rows["add_seq"] = 1 + rows["episode_id"] = "1" + rows["fact"] = [f"new-fact-{i}" for i in range(7)] + restored.apply(rows) + # Identity grouping is reused; new fact grouping still uses prompt16. + grouped = [r for r in oracle.requests if "rows" in r] + assert [len(r["rows"]) for r in grouped] == [16, 5, 16, 5, 16, 5] + assert len(restored.view("facts")) == 14 + + +def test_groupby_rejects_unsupported_sites_before_input() -> None: + from agent_memory.adapters.lotus.sem_groupby import validate_site_groupby_batching + from agent_memory.adapters.lotus.pair_execution import SemanticPairExecutionProfile + from agent_memory.tracing.semantic import query_digest + from test_semantic_pair_execution import PAIR_EMBEDDING + query = grouping() + for params in ({"labels": (("x", "label"),)}, {"membership": "overlapping"}): + with pytest.raises(ValueError, match="exclusive oracle"): + validate_site_groupby_batching(replace(query, params={**query.params, **params}), LotusExecutionConfig()) + profile = SemanticPairExecutionProfile(mode="proxy-only", embedding=PAIR_EMBEDDING, + min_similarity=0.5, direction="symmetric", left_id_columns=("l",), right_id_columns=("r",), + left_text_columns=("left",), right_text_columns=("right",)) + with pytest.raises(ValueError, match="exclusive oracle"): + validate_site_groupby_batching(query, LotusExecutionConfig(semantic_pair_profiles={query_digest(query): profile})) + + +@pytest.mark.parametrize("direction", ["symmetric", "left-to-right", "right-to-left"]) +def test_memoized_scores_are_bitwise_equal_and_keep_selection(direction: str, monkeypatch: pytest.MonkeyPatch) -> None: + from agent_memory.adapters.lotus import pair_execution as pairs + vectors = {"a": [0.1, 0.2, 0.3], "b": [1.0, 2.0, 3.0], "c": [-0.2, 0.4, 0.7], "d": [0.0, 1.0, 0.0]} + text_pairs = [(a, b) for a in vectors for b in vectors] * 3 + expected = tuple(sum(x * y for x, y in zip(vectors[a], vectors[b], strict=True)) / + (math.sqrt(sum(x*x for x in vectors[a])) * math.sqrt(sum(x*x for x in vectors[b]))) + for a, b in text_pairs) + sqrt = math.sqrt + norms: list[float] = [] + def counted(value: float) -> float: + norms.append(value) + return sqrt(value) + monkeypatch.setattr(pairs.math, "sqrt", counted) + actual = pairs._memoized_cosines(vectors, text_pairs) + assert [v.hex() for v in actual] == [v.hex() for v in expected] + assert len(norms) == len(vectors) + for threshold in (None, expected[1], math.nextafter(expected[1], math.inf)): + options: dict[str, Any] = dict(left_ids=[a for a, _ in text_pairs], right_ids=[b for _, b in text_pairs], + direction=direction, top_k=2, min_similarity=threshold) + assert pairs._select_positions(actual, **options) == pairs._select_positions(expected, **options) + assert pairs._memoized_cosines({}, []) == () + + +def test_factory_propagates_new_options(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + from agent_memory.adapters import lotus + from agent_memory.evaluation.agent_memory_drivers import ZepMemoryDriverFactory + class Reached(Exception): + pass + config = {"sem_groupby:test": PromptBatching(max_tasks=16)} + def inspect(**kwargs: Any) -> Any: + assert kwargs["config"].groupby_prompt_batching == config + assert kwargs["config"].reuse_unchanged_entity_name is True + raise Reached + factory = ZepMemoryDriverFactory(connector=SimpleNamespace(embedding_provider=None), + base_namespace="test", physical_fusion="zep-representative", groupby_prompt_batching=config, + reuse_unchanged_entity_name=True, neo4j_image="fake", neo4j_image_digest="fake") + monkeypatch.setattr(lotus, "LotusAdapter", inspect) + with pytest.raises(Reached): + factory("case", tmp_path / "state", tmp_path / "trace") + + +@pytest.mark.parametrize("parallel", [False, True]) +@pytest.mark.parametrize("cache_mode", ["disabled", "memory"]) +def test_benchmark_records_options_and_changes_resume_identity(monkeypatch: pytest.MonkeyPatch, tmp_path: Path, parallel: bool, cache_mode: str) -> None: + import agent_memory.evaluation.run as run_module + from test_benchmark_harness import _bundle + captured: dict[str, Any] = {} + + class Factory: + @classmethod + def from_environment(cls, **kwargs: Any) -> Any: + captured["factory"] = kwargs + return cls() + + def close(self) -> None: + pass + + def runtime_provenance(self) -> dict[str, str]: + return {"connector": "neo4j", "image": "fake", "image_digest": "sha256:fake", + "server_version": "fake", "driver_version": "fake"} + + class Runner: + def __init__(self, **kwargs: Any) -> None: + captured["runner"] = kwargs + + def run(self, bundle: Any) -> None: + pass + + monkeypatch.setattr(run_module, "_require_environment", lambda _: None) + monkeypatch.setattr(run_module, "collect_runtime_provenance", lambda *a, **k: {"source": {}, "runtime": {}}) + monkeypatch.setattr(run_module, "validate_run_provenance", lambda *a, **k: None) + monkeypatch.setattr(run_module, "ZepMemoryDriverFactory", Factory) + monkeypatch.setattr(run_module, "BenchmarkRunner", Runner) + sites = {semantic_pair_site_id(q): PromptBatching(max_tasks=16) + for n in ZepRepresentativeMemory.differentiate_policy().nodes.values() + for q in walk(n.query) if q.op == "sem_groupby"} + identities = [] + for enabled in (False, True): + run_module.run_agent_memory_bundle(bundle=_bundle(), contracts={}, system_id="zep-memory", + output_dir=tmp_path / str(enabled), physical_fusion="zep-representative", + parallel_fact_extraction=parallel and enabled, + lotus_cache_mode=cache_mode, + reuse_unchanged_entity_name=enabled, groupby_prompt_batching=sites if enabled else None) + identities.append(captured["runner"]["system_contract"].maintenance_execution_id) + assert identities[0] != identities[1] + assert captured["factory"]["groupby_prompt_batching"] == sites + assert captured["factory"]["reuse_unchanged_entity_name"] is True + assert captured["factory"].get("parallel_fact_extraction", False) is parallel + assert captured["factory"]["lotus_cache_mode"] == cache_mode + provenance = captured["runner"]["runtime_provenance"]["runtime"]["lotus_execution"]["site_execution"] + assert provenance["groupby_prompt_batching"] == {s: {"max_tasks": 16} for s in sites} + assert provenance["reuse_unchanged_entity_name"] is True + assert provenance.get("parallel_fact_extraction", False) is parallel + json.dumps(provenance) diff --git a/tests/test_zep_combined_physical.py b/tests/test_zep_combined_physical.py new file mode 100644 index 0000000..27dbb09 --- /dev/null +++ b/tests/test_zep_combined_physical.py @@ -0,0 +1,264 @@ +"""Opt-in combined physical lowering, without provider calls.""" + +import json + +import pytest + +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.fusion import parse_resolution +from agent_memory.memories.zep.policy import ZepMemory +from agent_memory.planner.physical import optimize_policy, walk + + +def test_shared_context_only_once() -> None: + from agent_memory.adapters.lotus.sem_filter_batch_prompting import ( + _BatchPromptingTask, _build_request, + ) + tasks = tuple(_BatchPromptingTask(i, f"r{i}", f"old-{i}", "new-fact") for i in range(4)) + request = _build_request(tasks, claim="opposite", max_tokens=100) + body = request.prompt[1]["content"] + assert body.count("new-fact") == 1 + payload = json.loads(body) + assert len(payload["rows"]) == 4 + assert "shared_context" in payload + + +def test_singleton_shortcut_does_not_touch_multiline_groups(monkeypatch: pytest.MonkeyPatch) -> None: + import pandas as pd + from agent_memory.adapters.lotus import relational + template = ZepMemory._deduplicated_facts.expr.inputs[0] + from dataclasses import replace + query = replace(template, params={**template.params, "singleton_identity": True}) + # Use the real deterministic specs, with complete rows and explicit group IDs. + from test_physical_fusion import state + rows = pd.DataFrame([state("single", 0), state("first", 1), state("second", 2)]) + rows["episode_id"] = ["0", "1", "2"] + rows["fact_ordinal"] = 0 + rows["content"] = rows.fact + from agent_memory.adapters.lotus.sem_groupby import GROUP_ID_COLUMN + rows[GROUP_ID_COLUMN] = [0, 1, 1] + rows.attrs["agent_memory_groupby_input_cols"] = ("relation_type", "fact") + seen = [] + def aggregate(spec, groups, *, context): + seen.extend(len(g) for g in groups) + return [{"relation_type": "LIKES", "fact": "merged"} for _ in groups] + monkeypatch.setattr(relational, "_execute_grouped_semantic_aggregate_spec_many", aggregate) + from types import SimpleNamespace + context = SimpleNamespace(config=LotusExecutionConfig()) + result = relational.execute_agg(query, {}, lambda *_: rows, context) + assert result.fact.tolist() == ["single", "merged"] + assert seen == [2] + + +def test_identity_decisions_require_complete_ordered_candidates() -> None: + from agent_memory.adapters.lotus.identity_reuse import IdentityDecisions + cache: dict[str, bool] = {} + decisions = IdentityDecisions(cache) + key = decisions.key("match", "Alice", ["Alice", "Bob"]) + assert decisions.lookup(key, 2) is None + decisions.store(key, 2, [0]) + assert decisions.lookup(key, 2) == (0,) + assert decisions.lookup(decisions.key("match", "Alice", ["Bob", "Alice"]), 2) is None + assert decisions.lookup(decisions.key("changed", "Alice", ["Alice", "Bob"]), 2) is None + assert decisions.lookup(decisions.key("match", "Alice", ["Alice"]), 1) is None + with pytest.raises(ValueError): + decisions.store(key, 2, [2]) + + +def test_group_identity_reuses_only_same_semantic_input(monkeypatch: pytest.MonkeyPatch) -> None: + import pandas as pd + from types import SimpleNamespace + from agent_memory.adapters.lotus import sem_groupby + from agent_memory.adapters.lotus.identity_reuse import identity_scope + calls = [] + def evaluate(*args, **kwargs): + calls.append(kwargs["docs"]) + return SimpleNamespace(outputs=[True] * len(kwargs["docs"])) + monkeypatch.setattr(sem_groupby, "evaluate_group_match_batch", evaluate) + rows = pd.DataFrame({"name": ["Alice", "Alicia"], "metadata": [1, 2]}) + cache: dict[str, bool] = {} + with identity_scope(cache): + assert sem_groupby.evaluate_group_matches(rows, input_cols=("name",), instruction="same {name}") == [(0, 1)] + rows["metadata"] = [3, 4] + assert sem_groupby.evaluate_group_matches(rows, input_cols=("name",), instruction="same {name}") == [(0, 1)] + assert len(calls) == 1 + sem_groupby.evaluate_group_matches(rows.iloc[::-1].reset_index(drop=True), input_cols=("name",), instruction="same {name}") + assert len(calls) == 2 + rows.loc[1, "name"] = "Bob" + sem_groupby.evaluate_group_matches(rows, input_cols=("name",), instruction="same {name}") + assert len(calls) == 3 + + +def test_combined_identity_rejects_changed_model_on_restore() -> None: + from agent_memory.adapters import LotusAdapter + from agent_memory.runtime.executor import PolicyExecutor + first = PolicyExecutor(ZepMemory.differentiate_policy(), adapter=LotusAdapter(model="model-a", config=LotusExecutionConfig(physical_fusion="zep-combined"))) + second = PolicyExecutor(ZepMemory.differentiate_policy(), adapter=LotusAdapter(model="model-b", config=LotusExecutionConfig(physical_fusion="zep-combined"))) + with pytest.raises(ValueError, match="fingerprint"): + second.restore_state(first.snapshot_state()) + + +def test_singleton_rewrite_preserves_screening_identity() -> None: + from agent_memory.tracing.semantic import query_digest + original = optimize_policy(ZepMemory.differentiate_policy(), strategy="zep-target-state") + combined = optimize_policy(ZepMemory.differentiate_policy(), strategy="zep-combined") + fact = next(n for n in original.nodes.values() if n.maintenance_query is not None and n.maintenance_query.op == "fused_target_state") + rewritten = combined.nodes[fact.node_id].maintenance_query + assert rewritten.params["join_profile_digest"] == query_digest(fact.maintenance_query.inputs[0]) + + +def test_shared_site_preserves_order_duplicates_and_current_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + from dataclasses import replace + import pandas as pd + from test_site_batching import adapter, query, rows + a, lm = adapter(monkeypatch) + site, setting = next(iter(a.config.pair_filter_batching.items())) + a.config = replace(a.config, pair_filter_batching={site: replace(setting, shared_columns=("fact_later_added",))}) + a._context.config = a.config + result = a.execute(query(), {"pairs": rows()}) + pd.testing.assert_frame_equal(result, rows()) + assert [len(c["rows"]) for c in lm.calls] == [4, 1, 2] + assert all("shared_context" in c for c in lm.calls) + assert all("new A" not in str(c["rows"]) and "new B" not in str(c["rows"]) for c in lm.calls) + + +def test_combined_real_runtime_snapshot_and_current_state(monkeypatch: pytest.MonkeyPatch) -> None: + from dataclasses import replace + from types import SimpleNamespace + from typing import Any + import pandas as pd + import lotus + import agent_memory as am + from agent_memory.adapters import LotusAdapter + from agent_memory.policy.aggregates import SemanticAggregateSpec + from agent_memory.policy.relation import Relation + from agent_memory.policy.logical import QueryExpr + from test_physical_fusion import state + + calls: list[dict[str, Any]] = [] + class Provider: + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + outputs = [] + for prompt in prompts: + data = json.loads(prompt[1]["content"]) + calls.append(data) + matches = [{"left_id": item["id"], "right_id": item["eligible_targets"][0]} for item in data["incoming"]] + states = [] + for target in data["targets"]: + assigned = [item for item, match in zip(data["incoming"], matches) if match["right_id"] == target["id"]] + if not assigned: + continue + if "summary" in data["output_schema"]["states"][0]: + states.append({"right_id": target["id"], "name": "Alice", "summary": "current merged summary"}) + else: + states.append({"right_id": target["id"], "relation_type": "LIKES", "fact": "current merged fact"}) + outputs.append(json.dumps({"matches": matches, "states": states})) + return SimpleNamespace(outputs=outputs) + + class Oracle(LotusAdapter): + def execute(self, query: QueryExpr, inputs: Any) -> Any: + if query.op == "agg" and query.inputs[0].op == "sem_groupby": + specs = [s for s in query.params["aggregates"] if isinstance(s, SemanticAggregateSpec)] + if specs and tuple(c.name for c in specs[0].output_cols) == ("name", "summary"): + rows = super().execute(query.inputs[0], inputs) + deterministic = replace(query, params={"aggregates": tuple(s for s in query.params["aggregates"] if not isinstance(s, SemanticAggregateSpec))}) + result = super().execute(deterministic, inputs) + result["name"] = [rows.iloc[0]["name"]] + result["summary"] = [rows.iloc[0]["content"]] + return result + return super().execute(query, inputs) + + raw = {**state("fact", 0), "name": "Alice", "entity_type": "Entity", "episode_id": "0", "entity_ordinal": 0, + "fact_ordinal": 0, "content": "old"} + source = am.Source({c: c for c in raw}) + entity = ZepMemory.entities.expr + entity_agg = next(q for q in walk(entity) if q.op == "agg") + from agent_memory.planner.physical import replace_query + entity = replace_query(entity, entity_agg.inputs[0].inputs[0], source.expr) + fact = ZepMemory._deduplicated_facts.expr + fact_agg = fact.inputs[0] + fact = replace_query(fact, fact_agg.inputs[0].inputs[0], source.expr) + adapter = Oracle(config=LotusExecutionConfig(physical_fusion="zep-combined")) + monkeypatch.setattr(adapter._context, "configure", lambda: None) + monkeypatch.setattr(lotus.settings, "lm", Provider()) + def make_flow(): + return am.SemanticDataflow(source=source, views={"entities": Relation(entity), "facts": Relation(fact)}, adapter=adapter) + flow = make_flow() + flow.apply(pd.DataFrame([raw])) + assert not calls + flow.apply(pd.DataFrame([{**raw, "add_seq": 1, "episode_id": "1", "content": "new"}])) + assert len(calls) == 2 + assert flow.view("entities").iloc[0]["entity_id"] == (0, 0) + assert len(json.loads(flow.view("entities").iloc[0]["mentions"])) == 2 + snapshot = flow.snapshot_state() + assert snapshot["predicate_decisions"] + restored = make_flow() + restored.restore_state(snapshot) + restored.apply(pd.DataFrame([{**raw, "add_seq": 2, "episode_id": "2", "content": "latest"}])) + assert len(calls) == 4 + entity_calls = [c for c in calls if "summary" in c["output_schema"]["states"][0]] + assert entity_calls[-1]["fixed_matches"] == {"l0": "r0"} + assert "latest" in str(entity_calls[-1]["incoming"]) + assert len(json.loads(restored.view("entities").iloc[0]["mentions"])) == 3 + + +def test_combined_rewrites_entities_and_facts_without_changing_view() -> None: + original = ZepMemory.differentiate_policy() + optimized = optimize_policy(original, strategy="zep-combined") + assert optimize_policy(optimized, strategy="zep-combined") is optimized + fused = [n for n in optimized.nodes.values() + if n.maintenance_query is not None + and n.maintenance_query.op == "fused_target_state"] + assert len(fused) == 2 + assert original.spec is optimized.spec + assert all(n.query == original.nodes[k].query for k, n in optimized.nodes.items()) + assert optimized.fingerprint != optimize_policy(original, strategy="zep-target-state").fingerprint + assert LotusExecutionConfig(physical_fusion="zep-combined").physical_fusion == "zep-combined" + marked = [q for n in fused for q in walk(n.maintenance_query) + if q.params.get("singleton_identity")] + assert len(marked) == 1 + assert marked[0].op == "agg" + + +def test_entity_resolution_strict_fields_and_many_to_one() -> None: + raw = json.dumps({"matches": [ + {"left_id": "l0", "right_id": "r0"}, + {"left_id": "l1", "right_id": "r0"}], + "states": [{"right_id": "r0", "name": "Alice", "summary": "Both facts."}]}) + result = parse_resolution(raw, {"l0": {"r0"}, "l1": {"r0"}}, + state_columns=("name", "summary")) + assert result["states"]["r0"] == {"name": "Alice", "summary": "Both facts."} + with pytest.raises(ValueError): + parse_resolution(raw.replace('"name": "Alice"', '"name": 3'), + {"l0": {"r0"}, "l1": {"r0"}}, state_columns=("name", "summary")) + + +@pytest.mark.parametrize("unknown_second", [False, True]) +def test_cached_targets_prune_only_unneeded_state( + monkeypatch: pytest.MonkeyPatch, unknown_second: bool, +) -> None: + from dataclasses import replace + import pandas as pd + import test_physical_fusion as fixture + from agent_memory.adapters.lotus.identity_reuse import identity_scope + + original = fixture.fused_query() + monkeypatch.setattr(fixture, "fused_query", lambda: replace( + original, params={**original.params, "identity_reuse": True})) + old = pd.DataFrame([fixture.state("selected", 0), fixture.state("unselected", 1)]) + incoming = pd.DataFrame([fixture.state("new", 2)]) + cache: dict[str, bool] = {} + with identity_scope(cache): + fixture.execute_fused(monkeypatch, incoming, old, [fixture.response( + [("l0", "r0")], [("r0", "merged")])]) + if unknown_second: + incoming = pd.concat([incoming, pd.DataFrame([fixture.state("unknown", 3)])], ignore_index=True) + assignments = [("l0", "r0")] + ([("l1", None)] if unknown_second else []) + result, model = fixture.execute_fused(monkeypatch, incoming, old, [fixture.response( + assignments, [("r0", "current merged")])]) + payload = json.loads(model.calls[0][0][0][1]["content"]) + assert payload["incoming"][0]["eligible_targets"] == ["r0"] + assert payload["fixed_matches"] == {"l0": "r0"} + assert [t["id"] for t in payload["targets"]] == (["r0", "r1"] if unknown_second else ["r0"]) + assert "unselected" in result.fact.tolist() + assert "current merged" in result.fact.tolist() diff --git a/tests/test_zep_combined_smoke_scope.py b/tests/test_zep_combined_smoke_scope.py new file mode 100644 index 0000000..3497f7b --- /dev/null +++ b/tests/test_zep_combined_smoke_scope.py @@ -0,0 +1,44 @@ +"""Workload scope checks stop before source validation or any provider setup.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from tools import zep_combined_smoke as smoke + + +@pytest.mark.parametrize("counts, error", [ + ((419, 198), "wrong AM source imported"), + ((128, 198), "expected frozen bundle"), + ((419, 69), "expected frozen bundle"), +]) +def test_full_combined_scope_keeps_exact_sample_size( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, counts: tuple[int, int], error: str, +) -> None: + case = SimpleNamespace( + events=tuple(SimpleNamespace(event_id=f"e{i}") for i in range(counts[0])), + questions=tuple(SimpleNamespace(evidence_event_ids=()) for _ in range(counts[1])), + ) + monkeypatch.setattr(smoke, "read_bundle", lambda _: SimpleNamespace(cases=(case,))) + with pytest.raises(ValueError, match=error): + smoke.execute(tmp_path, "preflight", full_sample0=True, maintenance_work=True, + combined_physical=True, fact_summary=True, representative=True, + node_batch_size=16, pack_small_groups=True, parallel_fact_extraction=True, + lotus_cache_mode="memory") + + +def test_prefix_scope_still_requires_complete_evidence( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + case = SimpleNamespace(events=(SimpleNamespace(event_id="e1"),), + questions=(SimpleNamespace(evidence_event_ids=()),)) + monkeypatch.setattr(smoke, "read_bundle", lambda _: SimpleNamespace(cases=(case,))) + with pytest.raises(ValueError, match="complete evidence"): + smoke.execute(tmp_path, "preflight", prefix128=True, maintenance_work=True, + combined_physical=True, fact_summary=True, representative=True) + + +def test_full_and_prefix_are_mutually_exclusive(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="choose full Sample 0"): + smoke.execute(tmp_path, "preflight", full_sample0=True, prefix128=True) diff --git a/tests/test_zep_fact_summary.py b/tests/test_zep_fact_summary.py new file mode 100644 index 0000000..f5c381f --- /dev/null +++ b/tests/test_zep_fact_summary.py @@ -0,0 +1,227 @@ +"""Offline checks of the distinct fact-derived summary condition.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pandas as pd +import pytest + +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus import relational +from agent_memory.memories.zep.fact_summary import SUMMARY_SPEC, ZepFactSummaryMemory +from agent_memory.memories.zep.policy import ZepMemory +from agent_memory.planner.physical import optimize_policy, walk +from agent_memory.policy.logical import QueryExpr + + +def test_query_and_execution_identity_are_distinct() -> None: + original = ZepMemory.differentiate_policy() + policy = ZepFactSummaryMemory.differentiate_policy() + adapter = LotusAdapter(config=LotusExecutionConfig(physical_fusion="zep-fact-summary")) + prepared = adapter.prepare_policy(policy) + assert policy.fingerprint != original.fingerprint + assert optimize_policy(prepared, strategy="zep-fact-summary").fingerprint == prepared.fingerprint + assert any(n.query.params.get("fact_summary") for n in prepared.nodes.values()) + from agent_memory.policy.schema import output_columns + assert "summary" not in output_columns(ZepFactSummaryMemory._identities.expr) + assert all("summary" not in q.params.get("columns", ()) for q in walk(ZepFactSummaryMemory._identities.expr) if q.op == "select") + + +@pytest.mark.parametrize("size", [0, 1999, 2000, 2001]) +def test_summary_threshold_uses_only_required_model_work(monkeypatch: pytest.MonkeyPatch, size: int) -> None: + query = next(q for q in walk(ZepFactSummaryMemory._summaries.expr) if q.op == "agg") + query = replace(query, params={**query.params, "fact_summary": True}) + rows = pd.DataFrame({"entity_id": ["a"], "summary": ["x" * size]}) + rows.attrs["agent_memory_groupby_keys"] = ("entity_id",) + calls = [] + def generate(spec, groups, *, context): + calls.extend(groups) + return [{"summary": "compressed"} for _ in groups] + monkeypatch.setattr(relational, "_execute_grouped_semantic_aggregate_spec_many", generate) + result = relational.execute_agg(query, {}, lambda *_: rows, SimpleNamespace(config=LotusExecutionConfig())) + assert len(calls) == int(size > 2000) + assert result.iloc[0]["summary"] == ("compressed" if size > 2000 else "x" * size) + + +def test_summary_clips_overlong_model_output(monkeypatch: pytest.MonkeyPatch) -> None: + query = next(q for q in walk(ZepFactSummaryMemory._summaries.expr) if q.op == "agg") + query = replace(query, params={**query.params, "fact_summary": True}) + rows = pd.DataFrame({"entity_id": ["a"], "summary": ["x" * 2001]}) + rows.attrs["agent_memory_groupby_keys"] = ("entity_id",) + monkeypatch.setattr(relational, "_execute_grouped_semantic_aggregate_spec_many", + lambda *a, **kw: [{"summary": "x" * 1001}]) + result = relational.execute_agg(query, {}, lambda *_: rows, SimpleNamespace(config=LotusExecutionConfig())) + assert result.iloc[0]["summary"] == "x" * 1000 + + +def test_unregistered_aggregate_cannot_use_shortcut() -> None: + query = QueryExpr(op="agg", inputs=(QueryExpr(op="group_by", params={"keys": ("entity_id",)}),), + params={"aggregates": (replace(SUMMARY_SPEC, instruction="different"),), "fact_summary": True}) + with pytest.raises(ValueError, match="registered"): + relational.execute_agg(query, {}, lambda *_: pd.DataFrame({"entity_id": ["a"], "summary": ["fact"]}), + SimpleNamespace(config=LotusExecutionConfig())) + + +def test_summary_compresses_once_and_records_truncation(monkeypatch: pytest.MonkeyPatch) -> None: + groups = [pd.DataFrame({"summary": [text]}) for text in ("short", "a" * 2100, "b" * 2100)] + calls = [] + events = [] + def generate(spec, selected, *, context): + calls.append((spec, selected)) + return [{"summary": "valid"}, {"summary": "x" * 2034}] if len(calls) == 1 else [{"summary": "repaired"}] + monkeypatch.setattr(relational, "_execute_grouped_semantic_aggregate_spec_many", generate) + monkeypatch.setattr("agent_memory.tracing.semantic.write_trace_event", lambda *args, **kwargs: events.append(kwargs)) + result = relational._fact_summary_values(groups, context=SimpleNamespace(config=LotusExecutionConfig())) + assert result == [{"summary": "short"}, {"summary": "valid"}, {"summary": "x" * 1000}] + assert len(calls) == 1 + assert calls[0][1][1] is groups[2] + assert "1000" in calls[0][0].instruction + assert "2000" not in calls[0][0].instruction + event = next(e for e in events if e["event_type"] == "fact_summary_truncation") + assert event["payload"]["raw_summary"] == "x" * 2034 + assert event["payload"]["before_chars"] == 2034 + assert event["payload"]["after_chars"] == 1000 + + +@pytest.mark.parametrize("bad", [None, 42, "", " "]) +def test_summary_invalid_output_fails_without_length_retry(monkeypatch: pytest.MonkeyPatch, bad: object) -> None: + calls = [] + def generate(*args, **kwargs): + calls.append(1) + return [{"summary": bad}] + monkeypatch.setattr(relational, "_execute_grouped_semantic_aggregate_spec_many", generate) + with pytest.raises((TypeError, ValueError)): + relational._fact_summary_values([pd.DataFrame({"summary": ["x" * 2100]})], + context=SimpleNamespace(config=LotusExecutionConfig())) + assert len(calls) == 1 + + +@pytest.mark.parametrize("raw,expected", [ + ("x" * 1000, "x" * 1000), + ("x" * 1076, "x" * 1000), + ("First sentence. " + "x" * 1074, "First sentence."), + ("x" * 999 + ". tail", "x" * 999 + "."), + ("中文" * 600, "中文" * 500), +]) +def test_native_sentence_boundary_policy(raw: str, expected: str) -> None: + assert relational._truncate_summary(raw) == expected + + +def test_summary_non_string_and_wrong_count_fail_without_retry(monkeypatch: pytest.MonkeyPatch) -> None: + for invalid in ([{"summary": None}], []): + calls = [] + def generate(*args, **kwargs): + calls.append(1) + return invalid + monkeypatch.setattr(relational, "_execute_grouped_semantic_aggregate_spec_many", generate) + with pytest.raises((TypeError, ValueError)): + relational._fact_summary_values([pd.DataFrame({"summary": ["x" * 2100]})], + context=SimpleNamespace(config=LotusExecutionConfig())) + assert len(calls) == 1 + + +def test_real_dataflow_fact_summary_and_restore(monkeypatch: pytest.MonkeyPatch) -> None: + import json + import lotus + import agent_memory as am + from agent_memory.planner.physical import replace_query + from agent_memory.policy.relation import Relation + from test_physical_fusion import state + + raw = {**state("Alice likes tea.", 0), "name": "Alice", "entity_type": "Entity", + "episode_id": "0", "entity_ordinal": 0, "fact_ordinal": 0, + "content": "Not a summary: Bob dislikes rain."} + source = am.Source({c: c for c in raw}) + identity = ZepFactSummaryMemory._identities.expr + identity_agg = next(q for q in walk(identity) if q.op == "agg") + identity = replace_query(identity, identity_agg.inputs[0].inputs[0], source.expr) + facts = ZepFactSummaryMemory.facts.expr + fact_agg = next(q for q in walk(facts) if q.op == "agg" and q.inputs[0].op == "sem_groupby") + facts = replace_query(fact_agg, fact_agg.inputs[0].inputs[0], source.expr) + entities = replace_query(ZepFactSummaryMemory.entities.expr, ZepFactSummaryMemory.facts.expr, facts) + entities = replace_query(entities, ZepFactSummaryMemory._identities.expr, identity) + calls = [] + class Provider: + def __call__(self, prompts, **kwargs): + answers = [] + for prompt in prompts: + payload = json.loads(prompt[1]["content"]) + calls.append(payload) + is_identity = set(payload["output_schema"]["states"][0]) == {"right_id", "name"} + matches = [{"left_id": item["id"], "right_id": item["eligible_targets"][0] if is_identity else None} + for item in payload["incoming"]] + states = [{"right_id": target, "name": "Alice"} for target in sorted({m["right_id"] for m in matches if m["right_id"] is not None})] + answers.append(json.dumps({"matches": matches, "states": states})) + return SimpleNamespace(outputs=answers) + adapter = LotusAdapter(config=LotusExecutionConfig(physical_fusion="zep-fact-summary")) + monkeypatch.setattr(adapter._context, "configure", lambda: None) + monkeypatch.setattr(lotus.settings, "lm", Provider()) + # This fixture exercises deduplicated facts and summaries, not contradictions. + def make_flow(): + return am.SemanticDataflow(source=source, views={"entities": Relation(entities), "facts": Relation(facts)}, adapter=adapter) + flow = make_flow() + flow.apply(pd.DataFrame([raw])) + assert flow.view("entities").iloc[0]["summary"] == "" # entity_id is an occurrence ID, not fixture endpoint "a" + assert not calls + snapshot = flow.snapshot_state() + restored = make_flow() + restored.restore_state(snapshot) + # Match the real occurrence-based ID to connect this fact to Alice. + restored.apply(pd.DataFrame([{**raw, "add_seq": 1, "episode_id": "1", "source_entity_id": (0, 0), + "fact": "Alice likes coffee."}])) + summary = restored.view("entities").iloc[0]["summary"] + assert summary == "Alice likes coffee." + assert "Bob" not in summary + assert all("summary" not in str(c["output_schema"]) for c in calls) + compression_groups = [] + def compress(spec, groups, *, context): + compression_groups.extend(groups) + return [{"summary": "Alice: compact retained facts."} for _ in groups] + monkeypatch.setattr(relational, "_execute_grouped_semantic_aggregate_spec_many", compress) + restored.apply(pd.DataFrame([{**raw, "add_seq": 2, "episode_id": "2", "source_entity_id": (0, 0), + "fact": "x" * 1990}])) + assert len(compression_groups) == 1 + assert restored.view("entities").iloc[0]["summary"] == "Alice: compact retained facts." + restarted = make_flow() + restarted.restore_state(restored.snapshot_state()) + restarted.apply(pd.DataFrame([{**raw, "add_seq": 3, "episode_id": "3", "source_entity_id": (0, 0), + "fact": "Alice likes cocoa."}])) + assert len(compression_groups) == 1 + assert restarted.view("entities").iloc[0]["summary"] == "Alice: compact retained facts.\nAlice likes cocoa." + + +def test_restore_rejects_previous_query_and_strategy() -> None: + from agent_memory.runtime.executor import PolicyExecutor + previous = PolicyExecutor(ZepMemory.differentiate_policy(), adapter=LotusAdapter( + config=LotusExecutionConfig(physical_fusion="zep-combined"))) + current = PolicyExecutor(ZepFactSummaryMemory.differentiate_policy(), adapter=LotusAdapter( + config=LotusExecutionConfig(physical_fusion="zep-fact-summary"))) + with pytest.raises(ValueError, match="fingerprint"): + current.restore_state(previous.snapshot_state()) + + +def test_direct_summary_execution_preserves_duplicates_without_model(monkeypatch: pytest.MonkeyPatch) -> None: + import agent_memory as am + source = am.Source({"entity_id": "id", "summary": "fact text"}) + query = source.group_by("entity_id").agg(SUMMARY_SPEC) + adapter = LotusAdapter(config=LotusExecutionConfig(physical_fusion="zep-fact-summary")) + monkeypatch.setattr(adapter._context, "configure", lambda: pytest.fail("unexpected model setup")) + result = adapter.execute(query.expr, {"log": pd.DataFrame({"entity_id": ["a", "a"], "summary": ["same", "same"]})}) + assert result.iloc[0]["summary"] == "same\nsame" + empty = adapter.execute(query.expr, {"log": pd.DataFrame(columns=["entity_id", "summary"])}) + assert empty.empty + with pytest.raises(TypeError, match="string evidence"): + adapter.execute(query.expr, {"log": pd.DataFrame({"entity_id": ["a"], "summary": [None]})}) +def test_fact_summary_storage_and_retrieval_plan() -> None: + from agent_memory.memories.zep.fact_summary import zep_storage_statements + from agent_memory.memories.zep.storage import GRAPHITI_NEO4J_STATEMENTS + from agent_memory.planner import PolicyDifferentiator + + assert zep_storage_statements("none") is GRAPHITI_NEO4J_STATEMENTS + statements = zep_storage_statements("zep-fact-summary") + assert len(statements.statements) == 4 + assert statements.statements[1].query == ZepFactSummaryMemory.entities.expr + assert statements.statements[2].query == ZepFactSummaryMemory.facts.expr + assert statements.statements[3].query == ZepFactSummaryMemory._episode_entities.expr + PolicyDifferentiator().differentiate(ZepFactSummaryMemory.spec(), statements=statements) diff --git a/tests/test_zep_frozen_followup.py b/tests/test_zep_frozen_followup.py new file mode 100644 index 0000000..9858e09 --- /dev/null +++ b/tests/test_zep_frozen_followup.py @@ -0,0 +1,187 @@ +"""Offline boundary checks for the fixed experiment.""" +from types import SimpleNamespace +from typing import Any, cast +import json +from pathlib import Path + +import pandas as pd +import pytest + +from tools.zep_frozen_followup import evaluate_group, freeze, validate_native_outputs +from tools.zep_frozen_followup import synthetic_preflight, run +from agent_memory.adapters.lotus.context import LotusExecutionConfig, LotusExecutionContext +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.sem_filter_batch_prompting import execute_batch_prompted_sem_filter + + +def events() -> list[dict[str, Any]]: + result = [] + for g in range(3): + call = f"call-{g}" + result.append({"operator": "sem_filter", "event_type": "operator_result", "operator_call_id": call, + "lowered_instruction": "{fact_later_added} contradicts {fact_earlier_added}"}) + for p in range(12): + result.append({"operator": "sem_filter", "event_type": "pair_decision", "trace_id": f"{g}-{p}", + "operator_call_id": call, "event_id": str(g), "right_id": "new", "left_id": "old", + "direction": "right-to-left", "decision_source": "oracle", "left": "fact: identical", + "right": "fact: repeated", "pair_index": p, "decision": False}) + return result + + +def test_freeze_preserves_duplicate_occurrences_order_and_caps() -> None: + groups = freeze(events()) + assert [len(g["rows"]) for g in groups] == [10, 10, 10] + assert [r["occurrence"] for r in groups[0]["rows"]] == [f"0-{i}" for i in range(10)] + changed = events() + for e in changed: + if "decision" in e: + e["decision"] = True + assert [[r["occurrence"] for r in g["rows"]] for g in freeze(changed)] == [ + [r["occurrence"] for r in g["rows"]] for g in groups] + + +@pytest.mark.parametrize("damage", ["identity", "instruction", "direction"]) +def test_freeze_fails_closed(damage: str) -> None: + data = events() + if damage == "identity": + data[2]["trace_id"] = data[1]["trace_id"] + elif damage == "instruction": + data[0]["lowered_instruction"] = "{missing}" + else: + data[1]["direction"] = "left-to-right" + with pytest.raises(ValueError): + freeze(data) + + +def test_empty_group_makes_no_call() -> None: + assert evaluate_group(None, {"rows": []}) == [] + + +@pytest.mark.parametrize("output", ["Maybe", "", "True and False", "{}"]) +def test_native_defaults_are_rejected(tmp_path: Path, output: str) -> None: + (tmp_path / "raw.json").write_text(json.dumps({"output": output})) + with pytest.raises(ValueError, match="unambiguous"): + validate_native_outputs([{"event_type": "llm_call", "raw_output_path": "raw.json"}], tmp_path, [False]) + + +def test_existing_batch_path_preserves_occurrences(monkeypatch: pytest.MonkeyPatch) -> None: + import lotus + + class Model: + max_tokens = 512 + cache = None + + def __call__(self, messages: Any, **kwargs: Any) -> Any: + assert len(messages) == 1 + return SimpleNamespace(outputs=[json.dumps({"decisions": [ + {"row_id": f"row_{i}", "keep": i != 1} for i in range(3)]})]) + + monkeypatch.setattr(lotus.settings, "lm", Model()) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + frame = pd.DataFrame({"fact_earlier_added": ["a", "b", "a"], + "fact_later_added": ["c"] * 3, "occurrence": ["0", "1", "2"]}) + result = execute_batch_prompted_sem_filter(frame, instruction="{fact_later_added} contradicts {fact_earlier_added}", + context=cast(LotusExecutionContext, SimpleNamespace(config=LotusExecutionConfig(structured_parse_retries=0))), + prompt_batching=PromptBatching(max_tasks=10)) + assert result.decisions == (True, False, True) + assert result.frame["occurrence"].tolist() == ["0", "2"] + + +def test_synthetic_freeze_and_label_isolation(tmp_path: Path) -> None: + fixture = Path(__file__).parent / "fixtures/zep_contradiction_contrasts.json" + output = tmp_path / "run" + synthetic_preflight(fixture, output) + groups = json.loads((output / "frozen.json").read_text()) + labels = json.loads((output / "assessment.json").read_text()) + assert len(groups) == 6 and len(labels) == 24 + assert sum(row["expected"] for row in labels) == 6 + class Adapter: + def execute(self, query: Any, views: Any) -> Any: + frame = views["frozen"] + assert list(frame.columns) == ["occurrence", "fact_earlier_added", "fact_later_added"] + assert "expected" not in query.params["instruction"] + assert frame["fact_later_added"].nunique() == 1 + return frame.iloc[[0, 2]] + for group in groups: + assert evaluate_group(Adapter(), group) == [True, False, True, False] + assert all(set(row) == {"occurrence", "fact_earlier_added", "fact_later_added"} for row in group["rows"]) + with pytest.raises(FileExistsError): + synthetic_preflight(fixture, output) + (output / "frozen.json").write_text("[]") + with pytest.raises(ValueError, match="evidence changed"): + run(output) + + +def test_synthetic_failure_stops_without_answer_replay(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + import tools.zep_frozen_followup as probe + fixture = Path(__file__).parent / "fixtures/zep_contradiction_contrasts.json" + output = tmp_path / "run" + synthetic_preflight(fixture, output) + calls = [] + def fail(**kwargs: Any) -> Any: + calls.append(kwargs) + raise RuntimeError("provider unavailable") + monkeypatch.setattr(probe, "LotusAdapter", fail) + monkeypatch.setattr(probe, "replay_answers", lambda _: pytest.fail("unexpected answer replay")) + with pytest.raises(RuntimeError, match="unavailable"): + run(output) + assert len(calls) == 1 + with pytest.raises(ValueError, match="refusing replay"): + run(output) + + +@pytest.mark.parametrize("artifact", ["frozen", "answer-inputs"]) +@pytest.mark.parametrize("change", ["none", "content", "missing_digest"]) +def test_historical_frozen_artifacts_checked_before_execution( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, artifact: str, change: str, +) -> None: + import tools.zep_frozen_followup as probe + import agent_memory.evaluation.provenance as provenance + + monkeypatch.setattr(provenance, "build_source_evidence", lambda _: {}) + old = tmp_path / "old" + trace_events = events() + for event in list(trace_events): + if event["event_type"] != "pair_decision": + continue + prompt = f"prompts/{event['trace_id']}.json" + probe.save(old / "fused" / prompt, [{"content": "identical repeated"}]) + trace_events.append({"event_type": "llm_call", "operator": "sem_filter", + "operator_call_id": event["operator_call_id"], + "llm_item_index": event["pair_index"], "prompt_path": prompt}) + trace = old / "fused/trace/events.jsonl" + trace.parent.mkdir(parents=True) + trace.write_text("".join(json.dumps(event) + "\n" for event in trace_events)) + for mode in ("unfused", "fused"): + for relative, row in ( + ("input/questions.jsonl", {"question_id": "q1"}), + ("cases/conv-26-2aac22fc/retrieval.jsonl", {"question_id": "q1", "context": "original"}), + ): + path = old / mode / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(row) + "\n") + output = tmp_path / "run" + probe.preflight(old, output) + manifest = json.loads((output / "manifest.json").read_text()) + for name in ("frozen", "answer-inputs"): + assert manifest[f"{name}_sha256"] == probe.digest(output / f"{name}.json") + if change == "content": + with (output / f"{artifact}.json").open("a") as stream: + stream.write("\n") + elif change == "missing_digest": + manifest.pop(f"{artifact}_sha256") + probe.save(output / "manifest.json", manifest) + assert all(probe.digest(Path(path)) == value for path, value in manifest["input_hashes"].items()) + + def stop_before_provider(_: Path) -> None: + raise RuntimeError("offline execution reached") + + monkeypatch.setattr(probe, "replay_answers", stop_before_provider) + monkeypatch.setattr(probe, "LotusAdapter", lambda **_: pytest.fail("unexpected provider setup")) + if change == "none": + with pytest.raises(RuntimeError, match="offline execution reached"): + run(output) + else: + with pytest.raises(ValueError, match=f"evidence changed or digest missing: {artifact}"): + run(output) + assert json.loads((output / "status.json").read_text()) == {"stage": "preflight_passed"} diff --git a/tests/test_zep_locomo_benchmark.py b/tests/test_zep_locomo_benchmark.py index a097a59..2ab2e16 100644 --- a/tests/test_zep_locomo_benchmark.py +++ b/tests/test_zep_locomo_benchmark.py @@ -278,9 +278,9 @@ def test_pinned_dataset_validation_rejects_wrong_bytes(tmp_path: Path) -> None: raise AssertionError("wrong LOCOMO bytes should be rejected") -def test_answer_and_judge_contracts_match_locked_native_digests() -> None: +def test_answer_format_fix_is_versioned_and_judge_digest_is_unchanged() -> None: assert ANSWER_PROMPT_DIGEST == ( - "dbd2f002e7cfafb4552a18d5c17e591cb6dd7d5eb44c762afd34094736e7fb6b" + "1fdf5c73bc3c5b36d6bd0dce56a596a5b2fbe657d949cf65e6b6823ec5d896f5" ) assert ZEP_JUDGE_PROMPT_DIGEST == ( "e74e7fbe755a002e5a8a05b3c661ed159a13a15c75c1953841019fc72d3cf83a" diff --git a/tests/test_zep_representative.py b/tests/test_zep_representative.py new file mode 100644 index 0000000..1afd5fa --- /dev/null +++ b/tests/test_zep_representative.py @@ -0,0 +1,220 @@ +"""Offline stable-representative checks, including real incremental execution.""" + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.memories.zep.representative import ZepRepresentativeMemory +from agent_memory.memories.zep.fact_summary import ZepFactSummaryMemory + + +@pytest.mark.parametrize("batch", [None, PromptBatching(max_tasks=16)]) +def test_arg_min_direct_and_incremental_restore(batch): + source = am.Source({"group": "key", "seq": "arrival", "ordinal": "ordinal", "fact": "text"}) + view = source.group_by("group").agg(am.arg_min(order_by=["seq", "ordinal"], columns=["fact"])) + adapter = LotusAdapter(config=LotusExecutionConfig(prompt_batching=batch)) + rows = pd.DataFrame([ + {"group": "a", "seq": 2, "ordinal": 0, "fact": "later"}, + {"group": "a", "seq": 1, "ordinal": 1, "fact": "first"}, + {"group": "a", "seq": 1, "ordinal": 1, "fact": "first"}, + {"group": "b", "seq": 3, "ordinal": 0, "fact": "independent"}, + ]) + flow = am.SemanticDataflow(source=source, views={"facts": view}, adapter=adapter) + for i in range(len(rows)): + flow.apply(rows.iloc[[i]]) + expected = adapter.execute(view.expr, {"log": rows.iloc[:i+1]}) + pd.testing.assert_frame_equal(flow.view("facts").sort_values("group").reset_index(drop=True), + expected.sort_values("group").reset_index(drop=True)) + restored = am.SemanticDataflow(source=source, views={"facts": view}, adapter=adapter) + restored.restore_state(flow.snapshot_state()) + flow = restored + + +def test_representative_rebinds_profiles_and_retains_listwise() -> None: + from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id + from agent_memory.evaluation.agent_memory_drivers import build_site_semantic_pair_profiles + from agent_memory.evaluation.semantic_pair_config import SemanticPairSiteBinding + from agent_memory.memories.zep.policy import ZepMemory + from agent_memory.memories.zep.storage import GRAPHITI_BGE_M3 + from agent_memory.planner.physical import walk + from agent_memory.tracing.semantic import query_digest + + old = ZepFactSummaryMemory.differentiate_policy() + new = ZepRepresentativeMemory.differentiate_policy() + old_joins = [q for n in old.nodes.values() if n.maintenance_query is not None + for q in walk(n.maintenance_query) if q.op == "sem_join"] + with pytest.raises(ValueError, match="not found"): + build_site_semantic_pair_profiles(new, bindings=tuple( + SemanticPairSiteBinding(semantic_pair_site_id(q), "search-filter", 5, None) + for q in old_joins), operators=("sem_filter", "sem_join", "sem_groupby"), + embedding=GRAPHITI_BGE_M3) + joins = [q for n in new.nodes.values() if n.maintenance_query is not None + for q in walk(n.maintenance_query) if q.op == "sem_join"] + bindings = tuple(SemanticPairSiteBinding(semantic_pair_site_id(q), "search-filter", 5, None) + for q in joins) + bindings += (SemanticPairSiteBinding(semantic_pair_site_id(ZepMemory._contradictory_fact_pairs.expr), "search-filter", 5, None),) + profiles, _ = build_site_semantic_pair_profiles(new, bindings=bindings, + operators=("sem_filter", "sem_join", "sem_groupby"), embedding=GRAPHITI_BGE_M3) + adapter = LotusAdapter(config=LotusExecutionConfig( + physical_fusion="zep-representative", semantic_pair_profiles=profiles, + sem_join_topk_method="listwise", sem_groupby_pair_batch_size=32, + sem_agg_prompt_batching=PromptBatching(max_tasks=16))) + prepared = adapter.prepare_policy(new) + fused = [n.maintenance_query for n in prepared.nodes.values() + if n.maintenance_query is not None and n.maintenance_query.op == "fused_target_state"] + assert len(fused) == 1 + assert profiles[fused[0].params["join_profile_digest"]].top_k == 5 + fact_joins = [q for n in prepared.nodes.values() if n.maintenance_query is not None and n.maintenance_query.op != "fused_target_state" + for q in walk(n.maintenance_query) if q.op == "sem_join"] + assert fact_joins + for query in fact_joins: + assert query.params["k"] == 1 + assert profiles[query_digest(query)].top_k == 5 + bad = SemanticPairSiteBinding("sem_filter:missing", "search-filter", 5, None) + with pytest.raises(ValueError, match="not found"): + build_site_semantic_pair_profiles(new, bindings=(bad,), + operators=("sem_filter", "sem_join", "sem_groupby"), embedding=GRAPHITI_BGE_M3) + + +def test_representative_keeps_distinct_details_and_multiple_targets(monkeypatch: pytest.MonkeyPatch) -> None: + import json + from agent_memory.adapters.lotus import adapter as adapter_module + from agent_memory.adapters.lotus.sem_groupby import GROUP_ID_COLUMN + from agent_memory.adapters.lotus.sem_join import assemble_join_frame + from agent_memory.planner.physical import replace_query, walk + from agent_memory.policy.aggregates import ArgMinAggregateSpec + from agent_memory.policy.relation import Relation + from agent_memory.runtime.executor import NodeOutputUpdate + + equivalent = {"Alice enjoys tea.": "Alice likes tea."} + + def key(text: str) -> str: + return equivalent.get(text, text) + + def group(query, inputs, execute, context): + frame = execute(query.inputs[0], inputs).copy() + frame[GROUP_ID_COLUMN] = pd.factorize(frame.fact.map(key))[0] + frame.attrs["agent_memory_groupby_input_cols"] = tuple(query.params["input_cols"]) + frame.attrs["agent_memory_sem_groupby_partition_by"] = tuple(query.params.get("partition_by", ())) + return frame + + def join(query, inputs, execute, context): + left, right = [execute(q, inputs) for q in query.inputs] + assert query.params["k"] == 1 + matches = [(i, j, None) for i in left.index for j in right.index + if key(left.loc[i, "fact"]) == key(right.loc[j, "fact"])] + return assemble_join_frame(left, right, matches, how="outer", id_columns=tuple(query.params["id_columns"])) + + monkeypatch.setattr(adapter_module, "execute_sem_groupby", group) + monkeypatch.setattr(adapter_module, "execute_sem_join", join) + aggregate = next(q for q in walk(ZepRepresentativeMemory.facts.expr) + if any(isinstance(s, ArgMinAggregateSpec) for s in q.params.get("aggregates", ()))) + raw = dict(source_entity_id="a", target_entity_id="b", relation_type="LIKES", fact="Alice likes tea.", + episode_id="0", fact_ordinal=0, add_seq=0, created_at="2026-01-01", + valid_at=None, invalid_at=None, expired_at=None, content="original") + source = am.Source({c: c for c in raw}) + view = Relation(replace_query(aggregate, aggregate.inputs[0].inputs[0], source.expr)) + adapter = LotusAdapter() + monkeypatch.setattr(adapter._context, "configure", lambda: pytest.fail("unexpected model request")) + flow = am.SemanticDataflow(source=source, views={"facts": view}, adapter=adapter) + rows = [] + for i, text in enumerate(["Alice likes tea.", "Alice likes coffee.", "Alice enjoys tea.", "Alice likes green tea.", "Alice likes coffee."]): + rows.append({**raw, "fact": text, "episode_id": str(i), "add_seq": i, "content": f"source-{i}"}) + flow.apply(pd.DataFrame([rows[-1]])) + actual = flow.view("facts") + expected = adapter.execute(view.expr, {"log": pd.DataFrame(rows)}) + for frame in (actual, expected): + frame["provenance"] = frame.provenance.map(lambda s: json.dumps( + sorted(json.loads(s), key=lambda r: r["episode_id"]), sort_keys=True)) + assert NodeOutputUpdate.between(actual, expected).is_empty + assert sum(len(json.loads(s)) for s in actual.provenance) == i + 1 + restored = am.SemanticDataflow(source=source, views={"facts": view}, adapter=adapter) + restored.restore_state(flow.snapshot_state()) + flow = restored + assert set(flow.view("facts").fact) == {"Alice likes tea.", "Alice likes coffee.", "Alice likes green tea."} + + +def test_arg_min_rejects_ambiguous_and_null_keys(): + source = am.Source({"g": "group", "seq": "order", "fact": "text"}) + query = source.group_by("g").agg(am.arg_min(order_by=["seq"], columns=["fact"])) + for rows, error in [([("a", 1, "x"), ("a", 1, "y")], "conflicting"), + ([("a", None, "x")], "null")]: + with pytest.raises(ValueError, match=error): + LotusAdapter().execute(query.expr, {"log": pd.DataFrame(rows, columns=pd.Index(["g", "seq", "fact"]))}) + + +def test_representative_view_has_new_identity_and_no_fact_synthesis(): + from agent_memory.planner.physical import walk + from agent_memory.policy.aggregates import ArgMinAggregateSpec + query = ZepRepresentativeMemory.facts.expr + assert any(isinstance(s, ArgMinAggregateSpec) for q in walk(query) for s in q.params.get("aggregates", ())) + policy = ZepRepresentativeMemory.differentiate_policy() + assert policy.fingerprint != ZepFactSummaryMemory.differentiate_policy().fingerprint + node = next(n for n in policy.nodes.values() if any(isinstance(s, ArgMinAggregateSpec) for s in n.query.params.get("aggregates", ()))) + assert node.execution_kind == "semantic_state" + assert node.maintenance_query is not None + from agent_memory.memories.zep.fact_summary import zep_memory_type, zep_storage_statements + assert zep_memory_type('zep-representative') is ZepRepresentativeMemory + assert zep_storage_statements('zep-representative').statements[2].query == ZepRepresentativeMemory.facts.expr + adapter = LotusAdapter(config=LotusExecutionConfig(physical_fusion='zep-representative')) + prepared = adapter.prepare_policy(policy) + assert adapter.prepare_policy(prepared).fingerprint == prepared.fingerprint + fused = [n for n in prepared.nodes.values() if n.maintenance_query is not None and n.maintenance_query.op == 'fused_target_state'] + assert len(fused) == 1 # Entity identity only, no model-generated fact state. + with pytest.raises(ValueError, match='representative logical view'): + adapter.prepare_policy(ZepFactSummaryMemory.differentiate_policy()) + + +def test_semantic_join_map_preserves_text_and_combines_provenance(monkeypatch: pytest.MonkeyPatch) -> None: + import json + from agent_memory.adapters.lotus import adapter as adapter_module + from agent_memory.adapters.lotus.sem_groupby import GROUP_ID_COLUMN + from agent_memory.adapters.lotus.sem_join import assemble_join_frame + from agent_memory.planner.physical import replace_query, walk + from agent_memory.policy.aggregates import ArgMinAggregateSpec + from agent_memory.policy.relation import Relation + from agent_memory.runtime.executor import NodeOutputUpdate + + def group(query, inputs, execute, context): + frame = execute(query.inputs[0], inputs).copy() + frame[GROUP_ID_COLUMN] = 0 + frame.attrs['agent_memory_groupby_input_cols'] = tuple(query.params['input_cols']) + frame.attrs['agent_memory_sem_groupby_partition_by'] = tuple(query.params.get('partition_by', ())) + return frame + + def join(query, inputs, execute, context): + left, right = [execute(q, inputs) for q in query.inputs] + matches = [(i, j, None) for i in left.index for j in right.index] + return assemble_join_frame(left, right, matches, how='outer', id_columns=tuple(query.params['id_columns'])) + + monkeypatch.setattr(adapter_module, 'execute_sem_groupby', group) + monkeypatch.setattr(adapter_module, 'execute_sem_join', join) + agg = next(q for q in walk(ZepRepresentativeMemory.facts.expr) + if any(isinstance(s, ArgMinAggregateSpec) for s in q.params.get('aggregates', ()))) + raw = dict(source_entity_id='a', target_entity_id='b', relation_type='LIKES', fact='Alice likes tea.', + episode_id='0', fact_ordinal=0, add_seq=0, created_at='2026-01-01', + valid_at=None, invalid_at=None, expired_at=None, content='original') + source = am.Source({c: c for c in raw}) + view = Relation(replace_query(agg, agg.inputs[0].inputs[0], source.expr)) + adapter = LotusAdapter() + monkeypatch.setattr(adapter._context, 'configure', lambda: pytest.fail('unexpected model request')) + flow = am.SemanticDataflow(source=source, views={'facts': view}, adapter=adapter) + rows = [] + for i, fact in enumerate(['Alice likes tea.', 'Alice enjoys tea.', 'Alice likes drinking tea.']): + rows.append({**raw, 'add_seq': i, 'episode_id': str(i), 'fact': fact}) + flow.apply(pd.DataFrame([rows[-1]])) + current = flow.view('facts') + assert current.iloc[0].fact == rows[0]['fact'] + assert len(json.loads(current.iloc[0].provenance)) == i + 1 + expected = adapter.execute(view.expr, {'log': pd.DataFrame(rows)}) + # Existing provenance aggregation retains occurrences, not arrival order. + for frame in (current, expected): + frame['provenance'] = frame['provenance'].map(lambda value: json.dumps( + sorted(json.loads(value), key=lambda record: record['episode_id']), sort_keys=True)) + assert NodeOutputUpdate.between(current, expected).is_empty + restored = am.SemanticDataflow(source=source, views={'facts': view}, adapter=adapter) + restored.restore_state(flow.snapshot_state()) + flow = restored diff --git a/tools/zep_combined_smoke.py b/tools/zep_combined_smoke.py new file mode 100644 index 0000000..b925b2c --- /dev/null +++ b/tools/zep_combined_smoke.py @@ -0,0 +1,290 @@ +"""One frozen combined condition using the existing benchmark runner.""" +from __future__ import annotations + +import argparse +from hashlib import sha256 +from importlib.metadata import version +import json +from pathlib import Path +import time +from typing import Any + +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.pair_execution import semantic_pair_site_id +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.site_batching import PairFilterBatching +from agent_memory.evaluation.bundle import read_bundle +from agent_memory.evaluation.locomo_contracts import locomo_task_contract +from agent_memory.evaluation.provenance import build_source_evidence +from agent_memory.evaluation.run import run_agent_memory_bundle +from agent_memory.memories.zep.policy import ZepMemory +from agent_memory.planner import PolicyDifferentiator + +MODEL = "deepseek/deepseek-flash" + + +def save(path: Path, data: Any) -> None: + """Atomically record progress in this run only.""" + temp = path.with_suffix(".tmp") + temp.write_text(json.dumps(data, indent=2, default=str)) + temp.replace(path) + + +def hashes(root: Path) -> dict[str, str]: + """Bind the experiment to its copied source and fixed inputs.""" + paths = [root / "site-profiles.json", Path(__file__)] + paths.extend(p for p in (root / "bundle").rglob("*") if p.is_file()) + paths.extend((root / "source/src").rglob("*.py")) + return {str(p): sha256(p.read_bytes()).hexdigest() for p in sorted(set(paths))} + + +def execute(root: Path, mode: str, *, full_sample0: bool = False, maintenance_work: bool = False, prefix128: bool = False, combined_physical: bool = False, fact_summary: bool = False, representative: bool = False, node_batch_size: int = 4, pack_small_groups: bool = False, listwise_batch_size: int | None = None, groupby_batch_size: int | None = None, reuse_unchanged_entity_name: bool = False, parallel_fact_extraction: bool = False, lotus_cache_mode: str = "disabled") -> None: + """Preflight, start, or resume one condition without replaying its prefix.""" + import agent_memory + from agent_memory.memories.zep.fact_summary import zep_memory_type, zep_storage_statements + if lotus_cache_mode not in {"disabled", "memory"}: + raise ValueError("lotus_cache_mode must be disabled or memory") + if isinstance(node_batch_size, bool) or node_batch_size not in (4, 16): + raise ValueError("node_batch_size must be 4 or 16") + if representative and not fact_summary: + raise ValueError("representative requires --fact-summary") + if listwise_batch_size is not None and (not representative or type(listwise_batch_size) is not int or listwise_batch_size != 16): + raise ValueError("listwise batching requires representative and the fixed batch16 condition") + if groupby_batch_size is not None and (not representative or type(groupby_batch_size) is not int or groupby_batch_size != 16): + raise ValueError("groupby batching requires representative and the fixed batch16 condition") + if reuse_unchanged_entity_name and not representative: + raise ValueError("unchanged names require representative") + if parallel_fact_extraction and not representative: + raise ValueError("parallel fact extraction requires representative") + if (representative or node_batch_size != 4 or pack_small_groups) and not (fact_summary and combined_physical and maintenance_work and (prefix128 or full_sample0)): + raise ValueError("representative/batch16/packing requires a frozen fact-summary condition") + if fact_summary and not combined_physical: + raise ValueError("fact-summary requires --combined-physical") + if full_sample0 and prefix128: + raise ValueError("choose full Sample 0 or the frozen 128-event prefix") + if combined_physical and (not (prefix128 or full_sample0) or not maintenance_work): + raise ValueError("combined physical experiment requires a frozen maintenance condition") + bundle = read_bundle(root / "bundle") + config = {semantic_pair_site_id(ZepMemory._contradictory_fact_pairs.expr): + PairFilterBatching(("fact_id:later_added",), PromptBatching(max_tasks=node_batch_size), + shared_columns=("fact_later_added",) if combined_physical else (), + pack_small_groups=pack_small_groups)} + fusion = "zep-combined" if combined_physical else "zep-target-state" + if fact_summary: + fusion = "zep-fact-summary" + if representative: + fusion = "zep-representative" + memory_type = zep_memory_type(fusion) + groupby_config: dict[str, PromptBatching] = {} + if groupby_batch_size is not None: + from agent_memory.planner.physical import optimize_policy, walk + logical = PolicyDifferentiator().differentiate(memory_type.spec(), statements=zep_storage_statements(fusion)) + physical = optimize_policy(logical, strategy=fusion) + sites = {semantic_pair_site_id(q) for node in physical.nodes.values() + for root_query in (node.query, node.maintenance_query) if root_query is not None + for q in walk(root_query) if q.op == "sem_groupby"} + if len(sites) != 2: + raise ValueError("representative must have exactly two entity/fact grouping sites") + groupby_config = {site: PromptBatching(max_tasks=groupby_batch_size) for site in sites} + listwise_config: dict[str, PromptBatching] = {} + if listwise_batch_size is not None: + from agent_memory.planner.physical import optimize_policy, walk + logical = PolicyDifferentiator().differentiate(memory_type.spec(), statements=zep_storage_statements(fusion)) + physical = optimize_policy(logical, strategy=fusion) + sites = {semantic_pair_site_id(q) for node in physical.nodes.values() + if node.maintenance_query is not None and node.maintenance_query.op != "fused_target_state" + for q in walk(node.maintenance_query) if q.op == "sem_join"} + if len(sites) != 1: + raise ValueError("representative must have exactly one executable fact matching site") + listwise_config = {site: PromptBatching(max_tasks=listwise_batch_size) for site in sites} + aggregate_batching = PromptBatching(max_tasks=node_batch_size) if full_sample0 or maintenance_work else None + reuse_sites = tuple(config) if maintenance_work else () + groupby_size = 32 if maintenance_work else None + expected = (419, 198) if full_sample0 else ((128, 1) if prefix128 else (24, 1)) + if combined_physical and not full_sample0: + case = bundle.cases[0] + event_ids = {event.event_id for event in case.events} + if not case.questions or any(not q.evidence_event_ids or not set(q.evidence_event_ids) <= event_ids for q in case.questions): + raise ValueError("quality questions require complete evidence in the 128-event prefix") + expected = (128, len(case.questions)) + condition = "zep-combined-aggregate4-sample0" if full_sample0 else "zep-combined-24e" + if full_sample0 and combined_physical: + condition = "zep-combined-sample0" + if prefix128: + condition = "zep-combined-128e" + if maintenance_work: + condition += "-maintenance-work-v1" + if combined_physical: + condition += "-combined-physical-v1" + if fact_summary: + condition += "-fact-summary-v1" + if representative: + condition += "-representative-v1" + if node_batch_size != 4: + condition += f"-node-batch{node_batch_size}" + if pack_small_groups: + condition += "-packed-groups-v1" + if listwise_batch_size is not None: + condition += f"-listwise-batch{listwise_batch_size}-v1" + if groupby_batch_size is not None: + condition += f"-groupby-prompt{groupby_batch_size}-v1" + if reuse_unchanged_entity_name: + condition += "-unchanged-entity-name-v1" + if parallel_fact_extraction: + condition += "-parallel-fact-extraction-v2" + if lotus_cache_mode == "memory": + condition += "-lotus-memory-cache" + if mode == "preflight": + if (root / "status.json").exists(): + raise FileExistsError("preflight already exists") + if len(bundle.cases) != 1 or (len(bundle.cases[0].events), len(bundle.cases[0].questions)) != expected: + raise ValueError(f"expected frozen bundle with events/questions={expected}") + if not Path(agent_memory.__file__).resolve().is_relative_to(root / "source"): + raise ValueError("wrong AM source imported") + logical = PolicyDifferentiator().differentiate(memory_type.spec(), statements=zep_storage_statements(fusion)) + extra_config: dict[str, Any] = {"lm_enable_cache": lotus_cache_mode == "memory"} + site_digests: dict[str, list[str]] = {} + profile_count = 0 + if combined_physical: + from agent_memory.evaluation.agent_memory_drivers import build_site_semantic_pair_profiles, BENCHMARK_STRUCTURED_MAX_TOKENS + from agent_memory.evaluation.semantic_pair_config import load_semantic_pair_profile_config + from agent_memory.memories.zep.storage import GRAPHITI_BGE_M3 + profiles, sites = build_site_semantic_pair_profiles(logical, + bindings=load_semantic_pair_profile_config(root / "site-profiles.json").bindings, + operators=("sem_filter", "sem_join", "sem_groupby"), embedding=GRAPHITI_BGE_M3, + embedding_device="cuda") + missing = set(config).union(reuse_sites) - set(sites) + if missing: + raise ValueError(f"batching/reuse sites not found in policy: {sorted(missing)}") + site_digests = {s: list(site.query_digests) for s, site in sites.items()} + profile_count = len(profiles) + extra_config = {"semantic_pair_profiles": profiles, "lm_num_retries": 0, + "structured_parse_retries": 0, "structured_max_tokens": BENCHMARK_STRUCTURED_MAX_TOKENS, + "lm_model_kwargs": {"extra_body": {"thinking": {"type": "disabled"}}}, + "lm_enable_cache": lotus_cache_mode == "memory"} + adapter = LotusAdapter(model=MODEL, config=LotusExecutionConfig(physical_fusion=fusion, pair_filter_batching=config, + listwise_join_batching=listwise_config, + groupby_prompt_batching=groupby_config, reuse_unchanged_entity_name=reuse_unchanged_entity_name, + parallel_fact_extraction=parallel_fact_extraction, + sem_agg_prompt_batching=aggregate_batching, predicate_reuse_sites=reuse_sites, + sem_groupby_pair_batch_size=groupby_size, sem_join_topk_method="listwise", **extra_config)) + policy = adapter.prepare_policy(logical) + from agent_memory.planner.physical import walk + from agent_memory.policy.aggregates import ArgMinAggregateSpec + joins = [q for n in policy.nodes.values() if n.maintenance_query is not None + for q in walk(n.maintenance_query) if q.op == "sem_join"] + if representative: + fused = [n for n in policy.nodes.values() if n.maintenance_query is not None + and n.maintenance_query.op == "fused_target_state"] + if len(fused) != 1 or not any(any(isinstance(s, ArgMinAggregateSpec) for s in n.query.params.get("aggregates", ())) for n in policy.nodes.values()): + raise ValueError("representative requires entity-only fusion and deterministic fact state") + if not joins or any(q.params.get("k") != 1 for q in joins): + raise ValueError("representative fact matching must retain top-1 listwise execution") + from agent_memory.tracing.semantic import query_digest + bound_profiles = extra_config["semantic_pair_profiles"] + fused_query = fused[0].maintenance_query + assert fused_query is not None + required_digests = {fused_query.params["join_profile_digest"]} + required_digests.update(query_digest(q) for n in policy.nodes.values() + if n.maintenance_query is not None and n.maintenance_query.op != "fused_target_state" + for q in walk(n.maintenance_query) if q.op == "sem_join") + if not required_digests <= set(bound_profiles): + raise ValueError("physical join rewrite lost its candidate screening profile") + save(root / "source-evidence.json", build_source_evidence(root / "source")) + save(root / "manifest.json", {"bundle_fingerprint": bundle.fingerprint, "hashes": hashes(root), + "source_path": agent_memory.__file__, "model": MODEL, "thinking": False, + "cache": lotus_cache_mode == "memory", "lotus_cache_mode": lotus_cache_mode, + "provider_retries": 0, "structured_parse_retries": 0, "answer_parse_attempts": 1, + "physical_fusion": fusion, "prompt_batching": None, + "node_batch_size": node_batch_size, "pack_small_groups": pack_small_groups, + **({"listwise_join_batching": {site: setting.to_dict() for site, setting in listwise_config.items()}} if listwise_config else {}), + **({"groupby_prompt_batching": {site: setting.to_dict() for site, setting in groupby_config.items()}} if groupby_config else {}), + **({"reuse_unchanged_entity_name": True} if reuse_unchanged_entity_name else {}), + **({"parallel_fact_extraction": True} if parallel_fact_extraction else {}), + "lm_max_batch_size": adapter.config.lm_max_batch_size, + "sem_join_topk_method": "listwise", + "site_query_digests": site_digests, + "matched_profile_query_count": profile_count, + "fact_state": "earliest representative" if representative else "semantic synthesis", + "logical_variant": memory_type.__name__, + "summary_input": "incident facts" if fact_summary else "entity mentions", + "fusion_minimum_parse_retries": 1, + "approximate_singleton_fact_identity": combined_physical and not representative, + "quality_scope": "full Sample 0" if full_sample0 else ("evidence-contained prefix questions; not full Sample 0 accuracy" if combined_physical else "historical fixed questions"), + "condition": condition, "event_count": expected[0], "question_count": expected[1], + "sem_agg_prompt_batching": aggregate_batching.to_dict() if aggregate_batching else None, + "predicate_reuse_sites": reuse_sites, "sem_groupby_pair_batch_size": groupby_size, + "pair_filter_batching": {s: c.to_dict() for s, c in config.items()}, + "execution_fingerprint": adapter.maintenance_execution_fingerprint, + "plan_fingerprint": policy.fingerprint, + "dependencies": {n: version(n) for n in ("lotus-ai", "litellm", "pandas", "neo4j")}, + "note": "one new combined condition; historical comparison, not paired ablation"}) + save(root / "status.json", {"stage": "preflight_passed"}) + return + status = json.loads((root / "status.json").read_text()) + if mode == "run" and status["stage"] != "preflight_passed": + raise ValueError("refusing replay of a started experiment") + if mode == "resume": + if status["stage"] not in {"failed", "interrupted"}: + raise ValueError("resume requires a failed or interrupted run") + checkpoints = list((root / "combined" / "cases").glob("*/checkpoints/current.json")) + if not checkpoints: + raise ValueError("resume requires a durable checkpoint") + manifest = json.loads((root / "manifest.json").read_text()) + if hashes(root) != manifest["hashes"]: + raise ValueError("source or input changed since preflight") + if manifest.get("condition") != condition: + raise ValueError("execution condition changed since preflight") + save( + root / "status.json", + { + "stage": "running", + "mode": mode, + "started_at": time.time(), + "resuming": mode == "resume", + }, + ) + try: + run_agent_memory_bundle( + bundle=bundle, contracts={"locomo": locomo_task_contract(judge_model_id=MODEL)}, + system_id="zep-memory", output_dir=root / "combined", base_namespace=root.name, + memory_model_id=MODEL, memory_provider_model_id=MODEL, answer_model_id=MODEL, judge_model_id=MODEL, + grouped_agg_rule="rule-join-map", sem_join_topk_method="listwise", physical_fusion=fusion, + pair_filter_batching=config, memory_num_retries=0, structured_parse_retries=0, parse_attempts=1, + listwise_join_batching=listwise_config, + groupby_prompt_batching=groupby_config, reuse_unchanged_entity_name=reuse_unchanged_entity_name, + parallel_fact_extraction=parallel_fact_extraction, + sem_agg_prompt_batching=aggregate_batching, + predicate_reuse_sites=reuse_sites, sem_groupby_pair_batch_size=groupby_size, + semantic_pair_profile_config=root / "site-profiles.json", embedding_device="cuda", + lotus_cache_mode=lotus_cache_mode, memory_thinking_enabled=False, condition_id=condition, + ) + except KeyboardInterrupt: + save(root / "status.json", {"stage": "interrupted", "error": "execution interrupted; not automatically resumed"}) + raise + except Exception as error: + save(root / "status.json", {"stage": "failed", "error": str(error), "type": type(error).__name__}) + raise + save(root / "status.json", {"stage": "completed", "finished_at": time.time()}) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("mode", choices=("preflight", "run", "resume")) + parser.add_argument("--root", type=Path, required=True) + parser.add_argument("--full-sample0", action="store_true") + parser.add_argument("--maintenance-work", action="store_true") + parser.add_argument("--prefix128", action="store_true") + parser.add_argument("--combined-physical", action="store_true") + parser.add_argument("--fact-summary", action="store_true") + parser.add_argument("--representative", action="store_true") + parser.add_argument("--node-batch-size", type=int, default=4) + parser.add_argument("--pack-small-groups", action="store_true") + parser.add_argument("--listwise-batch-size", type=int) + parser.add_argument("--groupby-batch-size", type=int) + parser.add_argument("--reuse-unchanged-entity-name", action="store_true") + parser.add_argument("--parallel-fact-extraction", action="store_true") + parser.add_argument("--lotus-cache-mode", choices=("disabled", "memory"), default="disabled") + args = parser.parse_args() + execute(args.root.resolve(), args.mode, full_sample0=args.full_sample0, maintenance_work=args.maintenance_work, prefix128=args.prefix128, combined_physical=args.combined_physical, fact_summary=args.fact_summary, representative=args.representative, node_batch_size=args.node_batch_size, pack_small_groups=args.pack_small_groups, listwise_batch_size=args.listwise_batch_size, groupby_batch_size=args.groupby_batch_size, reuse_unchanged_entity_name=args.reuse_unchanged_entity_name, parallel_fact_extraction=args.parallel_fact_extraction, lotus_cache_mode=args.lotus_cache_mode) diff --git a/tools/zep_frozen_followup.py b/tools/zep_frozen_followup.py new file mode 100644 index 0000000..380ac9c --- /dev/null +++ b/tools/zep_frozen_followup.py @@ -0,0 +1,302 @@ +"""Bounded answer replay and frozen contradiction experiment; no memory updates.""" +from __future__ import annotations + +import argparse +from dataclasses import asdict +from hashlib import sha256 +from importlib.metadata import version +import json +from pathlib import Path +import re +import time +from typing import Any + +import pandas as pd + +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.evaluation.locomo_contracts import locomo_task_contract +from agent_memory.evaluation.models import LiteLLMBenchmarkModel +from agent_memory.evaluation.trace_metrics import normalize_provider_calls, summarize_provider_calls +from agent_memory.evaluation.types import BenchmarkQuestion +from agent_memory.policy.logical import QueryExpr + +MODEL = "deepseek/deepseek-flash" + + +def read_lines(path: Path) -> list[dict[str, Any]]: + """Read complete JSONL evidence, failing on partial records.""" + return [json.loads(line) for line in path.read_text().splitlines() if line.strip()] + + +def save(path: Path, value: Any) -> None: + """Atomically write a new experiment artifact.""" + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2, default=str)) + temporary.replace(path) + + +def digest(path: Path) -> str: + return sha256(path.read_bytes()).hexdigest() + + +def freeze(events: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Select the first three multi-candidate groups without reading decisions.""" + results = {e["operator_call_id"]: e for e in events + if e.get("operator") == "sem_filter" and e.get("event_type") == "operator_result"} + groups: dict[tuple[str, str], list[dict[str, Any]]] = {} + identities: set[str] = set() + for event in events: + if event.get("operator") != "sem_filter" or event.get("event_type") != "pair_decision": + continue + if event["trace_id"] in identities: + raise ValueError("duplicate occurrence trace identity") + identities.add(event["trace_id"]) + if event["decision_source"] != "oracle" or event["direction"] != "right-to-left": + raise ValueError("expected post-screening, directed oracle pairs") + groups.setdefault((event["event_id"], event["right_id"]), []).append(event) + selected = [] + for (event_id, right_id), rows in groups.items(): + if len(rows) < 2: + continue + instruction = results[rows[0]["operator_call_id"]]["lowered_instruction"] + if set(re.findall(r"\{([^{}]+)\}", instruction)) != {"fact_earlier_added", "fact_later_added"}: + raise ValueError("cannot restore every predicate input") + frozen = [] + for row in rows[:10]: + if results[row["operator_call_id"]]["lowered_instruction"] != instruction: + raise ValueError("mixed predicate instructions") + if not row["left"].startswith("fact: ") or not row["right"].startswith("fact: "): + raise ValueError("unrecoverable fact text") + frozen.append({"occurrence": row["trace_id"], "operator_call_id": row["operator_call_id"], + "pair_index": row["pair_index"], "left_id": row["left_id"], + "right_id": row["right_id"], "direction": row["direction"], + "fact_earlier_added": row["left"][6:], "fact_later_added": row["right"][6:], + "historical_decision": row["decision"]}) + selected.append({"event_id": event_id, "right_id": right_id, + "instruction": instruction, "rows": frozen}) + if len(selected) == 3: + break + if len(selected) != 3: + raise ValueError("fewer than three multi-candidate groups") + return selected + + +def preflight(old: Path, output: Path) -> None: + """Freeze exact evidence and verify its original provider prompts.""" + output.mkdir(parents=True, exist_ok=False) + import agent_memory + from agent_memory.evaluation.provenance import build_source_evidence + + source_root = Path(agent_memory.__file__).resolve().parents[2] + save(output / "source-evidence.json", build_source_evidence(source_root)) + trace = old / "fused/trace/events.jsonl" + events = read_lines(trace) + groups = freeze(events) + calls = {(e["operator_call_id"], e["llm_item_index"]): e for e in events + if e.get("event_type") == "llm_call" and e.get("operator") == "sem_filter"} + hashes = {str(trace): digest(trace)} + for group in groups: + for row in group["rows"]: + call = calls[(row["operator_call_id"], row["pair_index"])] + path = old / "fused" / call["prompt_path"] + messages = json.loads(path.read_text()) + text = "\n".join(m["content"] for m in messages) + if not all(row[k] in text for k in ("fact_earlier_added", "fact_later_added")): + raise ValueError("frozen text not found in original provider context") + hashes[str(path)] = digest(path) + row["original_messages"] = messages + save(output / "frozen.json", groups) + saved_answers = {} + for mode in ("unfused", "fused"): + questions_path = old / mode / "input/questions.jsonl" + retrieval_path = old / mode / "cases/conv-26-2aac22fc/retrieval.jsonl" + q, r = read_lines(questions_path), read_lines(retrieval_path) + if len(q) != 1 or len(r) != 1 or q[0]["question_id"] != r[0]["question_id"]: + raise ValueError("ambiguous question/retrieval identity") + saved_answers[mode] = {"question": q[0], "context": r[0]["context"]} + hashes.update({str(p): digest(p) for p in (questions_path, retrieval_path)}) + save(output / "answer-inputs.json", saved_answers) + contract = locomo_task_contract(judge_model_id=MODEL) + save(output / "manifest.json", {"model": MODEL, "input_hashes": hashes, + "frozen_sha256": digest(output / "frozen.json"), + "answer-inputs_sha256": digest(output / "answer-inputs.json"), + "script_sha256": digest(Path(__file__)), "source_path": agent_memory.__file__, + "dependencies": {name: version(name) for name in ("lotus-ai", "litellm", "pandas")}, + "answer_prompt_digest": contract.answer_prompt_digest, "answer_parser_id": contract.answer_parser_id, + "order": ["answers", "pointwise", "joint"], "cache": False, "provider_retries": 0, + "pair_parse_retries": 0, "groups": len(groups), "pairs": sum(len(g["rows"]) for g in groups), + "note": "frozen negatives only; not full Sample 0 or contradiction recall"}) + save(output / "status.json", {"stage": "preflight_passed"}) + + +def replay_answers(output: Path) -> None: + contract = locomo_task_contract(judge_model_id=MODEL) + model = LiteLLMBenchmarkModel(model_id=MODEL) + for mode, data in json.loads((output / "answer-inputs.json").read_text()).items(): + q = {k: v for k, v in data["question"].items() if k != "case_id"} + q["evidence_event_ids"] = tuple(q["evidence_event_ids"]) + question = BenchmarkQuestion(**q) + prompt = contract.answer_prompt(question, data["context"]) + directory = output / "answers" / mode + save(directory / "prompt.json", asdict(prompt)) + answer = None + for attempt in (1, 2): + response = model.complete(prompt, attempt=attempt) + save(directory / f"answer-attempt-{attempt}.json", asdict(response)) + try: + answer = contract.answer_parser(response.text) + except ValueError as error: + save(directory / f"parse-error-{attempt}.json", {"error": str(error)}) + else: + break + if answer is None: + save(directory / "result.json", {"status": "parse_failed"}) + continue + assert contract.deterministic_scorer is not None + official = contract.deterministic_scorer(question, answer) + grader = contract.additional_graders[0] + assert grader.judge_plan is not None and grader.judge_reducer is not None + steps = grader.judge_plan(question, answer) + if len(steps) != 1: + raise ValueError("expected one judge call") + response = model.complete(steps[0].prompt, attempt=1) + save(directory / "judge.json", {"prompt": asdict(steps[0].prompt), "response": asdict(response)}) + judged = grader.judge_reducer(question, answer, (steps[0].parse(response.text),)) + save(directory / "result.json", {"status": "completed", "answer": answer, + "gold": question.gold_answer, "official": asdict(official), "judge": asdict(judged)}) + + +def synthetic_preflight(fixture: Path, output: Path) -> None: + """Freeze controlled contrasts; never pass assessment labels to the model.""" + from agent_memory.memories.zep.policy import _CONTRADICTORY_FACT_INSTRUCTION + + data = json.loads(fixture.read_text()) + if data.get("kind") != "synthetic-controlled-contrasts" or len(data["groups"]) != 6: + raise ValueError("expected six explicitly synthetic groups") + instruction = _CONTRADICTORY_FACT_INSTRUCTION.replace("{fact:later_added}", "{fact_later_added}").replace( + "{fact:earlier_added}", "{fact_earlier_added}") + groups, assessment = [], [] + for index, group in enumerate(data["groups"]): + candidates = group["candidates"] + if len(candidates) != 4 or any(type(c["expected"]) is not bool for c in candidates): + raise ValueError("expected four strictly boolean-labeled candidates") + if sum(c["expected"] for c in candidates) != 1: + raise ValueError("expected exactly one contradiction per group") + rows = [] + for pair, candidate in enumerate(candidates): + if not all(isinstance(value, str) and value.strip() for value in + (group["new_fact"], candidate["fact"], candidate["reason"])): + raise ValueError("empty fact or assessment reason") + occurrence = f"synthetic-{index}-{pair}" + rows.append({"occurrence": occurrence, "fact_earlier_added": candidate["fact"], + "fact_later_added": group["new_fact"]}) + assessment.append({"occurrence": occurrence, "expected": candidate["expected"], + "reason": candidate["reason"]}) + groups.append({"topic": group["topic"], "instruction": instruction, "rows": rows}) + output.mkdir(parents=True, exist_ok=False) + save(output / "frozen.json", groups) + save(output / "assessment.json", assessment) + import agent_memory + save(output / "manifest.json", {"model": MODEL, "input_hashes": {str(fixture): digest(fixture)}, + "script_sha256": digest(Path(__file__)), "frozen_sha256": digest(output / "frozen.json"), + "assessment_sha256": digest(output / "assessment.json"), "source_path": agent_memory.__file__, + "dependencies": {name: version(name) for name in ("lotus-ai", "litellm", "pandas")}, + "kind": data["kind"], "order": ["pointwise", "joint"], "pairs": 24, "groups": 6, + "maximum_requests": 30, "cache": False, "provider_retries": 0, "pair_parse_retries": 0}) + save(output / "status.json", {"stage": "preflight_passed"}) + + +def evaluate_group(adapter: Any, group: dict[str, Any]) -> list[bool]: + """Execute the existing filter on frozen rows without candidate regeneration.""" + frame = pd.DataFrame([{k: row[k] for k in ("occurrence", "fact_earlier_added", "fact_later_added")} + for row in group["rows"]], columns=pd.Index(["occurrence", "fact_earlier_added", "fact_later_added"])) + if frame.empty: + return [] + source = QueryExpr(op="materialized_view", params={"name": "frozen", "columns": tuple(frame.columns)}) + query = QueryExpr(op="sem_filter", inputs=(source,), params={"instruction": group["instruction"]}) + result = adapter.execute(query, {"frozen": frame}) + kept = result["occurrence"].tolist() + if len(kept) != len(set(kept)) or not set(kept) <= set(frame["occurrence"]): + raise ValueError("invalid output occurrence identities") + return [identity in kept for identity in frame["occurrence"]] + + +def validate_native_outputs(events: list[dict[str, Any]], directory: Path, decisions: list[bool]) -> None: + """Reject native parser defaults rather than silently accepting malformed text.""" + raw = [json.loads((directory / e["raw_output_path"]).read_text())["output"] + for e in events if e.get("event_type") == "llm_call" and "raw_output_path" in e] + if len(raw) != len(decisions): + raise ValueError("native output count mismatch") + for output, decision in zip(raw, decisions, strict=True): + match = re.fullmatch(r"\s*(?:Answer:\s*)?(True|False)\s*", output, flags=re.IGNORECASE) + if match is None or (match[1].lower() == "true") != decision: + raise ValueError("native output is not an unambiguous boolean decision") + + +def run(output: Path) -> None: + if json.loads((output / "status.json").read_text())["stage"] != "preflight_passed": + raise ValueError("refusing replay of a started experiment") + manifest = json.loads((output / "manifest.json").read_text()) + if any(digest(Path(p)) != expected for p, expected in manifest["input_hashes"].items()): + raise ValueError("historical evidence changed") + if digest(Path(__file__)) != manifest["script_sha256"]: + raise ValueError("script changed after preflight") + names = (("frozen", "assessment") if manifest.get("kind") == "synthetic-controlled-contrasts" + else ("frozen", "answer-inputs")) + for name in names: + if digest(output / f"{name}.json") != manifest.get(f"{name}_sha256"): + raise ValueError(f"frozen evidence changed or digest missing: {name}") + try: + if manifest.get("kind") != "synthetic-controlled-contrasts": + save(output / "status.json", {"stage": "answer_replay"}) + replay_answers(output) + groups = json.loads((output / "frozen.json").read_text()) + summaries = {} + for mode in ("pointwise", "joint"): + directory = output / mode + save(output / "status.json", {"stage": "contradictions", "mode": mode}) + adapter = LotusAdapter(model=MODEL, config=LotusExecutionConfig( + semantic_trace_dir=directory / "trace", lm_enable_cache=False, lm_num_retries=0, + lm_max_batch_size=1, structured_parse_retries=0, structured_max_tokens=8192, + prompt_batching=PromptBatching(max_tasks=10) if mode == "joint" else None, + lm_model_kwargs={"extra_body": {"thinking": {"type": "disabled"}}})) + results = [] + for index, group in enumerate(groups): + trace = directory / "trace/events.jsonl" + before_count = len(read_lines(trace)) if trace.exists() else 0 + started = time.perf_counter() + decisions = evaluate_group(adapter, group) + if mode == "pointwise": + validate_native_outputs(read_lines(trace)[before_count:], directory, decisions) + results.append({"group": index, "decisions": decisions, "wall_seconds": time.perf_counter() - started}) + save(directory / "results.json", results) + calls = normalize_provider_calls(read_lines(directory / "trace/events.jsonl"), output_dir=directory, include_cost=False) + summaries[mode] = {"results": results, "usage": summarize_provider_calls(calls)} + save(directory / "calls.json", calls) + save(output / "summary.json", summaries) + save(output / "status.json", {"stage": "completed"}) + except Exception as error: + save(output / "status.json", {"stage": "failed", "error": str(error), "type": type(error).__name__}) + raise + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("mode", choices=("preflight", "synthetic-preflight", "run")) + parser.add_argument("--old", type=Path) + parser.add_argument("--fixture", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.mode == "preflight": + if args.old is None: + parser.error("preflight requires --old") + preflight(args.old, args.output) + elif args.mode == "synthetic-preflight": + if args.fixture is None: + parser.error("synthetic-preflight requires --fixture") + synthetic_preflight(args.fixture, args.output) + else: + run(args.output)