From e87f25799c9e10826e78213faa4a4b160c906cce Mon Sep 17 00:00:00 2001 From: Von Date: Mon, 31 Aug 2026 00:01:28 +0800 Subject: [PATCH 1/5] [feat-semantic] Add cardinality-bounded semantic joins - extend semantic joins with exact-key and per-left top-k contracts - add listwise and pairwise top-k execution with optional candidate filtering - lower exclusive semantic grouping through partition-aware join-map maintenance - document and test the operator, membership, and physical execution contracts Signed-off-by: Von --- docs/design/groupby_agg.md | 127 ++++- docs/design/operator-api.md | 142 +++++- src/agent_memory/adapters/lotus/adapter.py | 13 +- src/agent_memory/adapters/lotus/context.py | 7 + .../adapters/lotus/pair_execution.py | 92 +++- src/agent_memory/adapters/lotus/relational.py | 20 + .../adapters/lotus/sem_groupby.py | 7 + src/agent_memory/adapters/lotus/sem_join.py | 194 +++++++- .../adapters/lotus/sem_topk_join.py | 459 ++++++++++++++++++ src/agent_memory/planner/rules.py | 299 ++++++++++-- src/agent_memory/policy/relation.py | 51 +- src/agent_memory/policy/schema.py | 4 + tests/test_differential_runtime.py | 12 - tests/test_interface_scaffold.py | 322 ++++++++++-- tests/test_sem_topk_join.py | 313 ++++++++++++ tests/test_semantic_pair_execution.py | 33 ++ 16 files changed, 1950 insertions(+), 145 deletions(-) create mode 100644 src/agent_memory/adapters/lotus/sem_topk_join.py create mode 100644 tests/test_sem_topk_join.py diff --git a/docs/design/groupby_agg.md b/docs/design/groupby_agg.md index c1c2eed..fff756e 100644 --- a/docs/design/groupby_agg.md +++ b/docs/design/groupby_agg.md @@ -423,10 +423,35 @@ entities = ( group_id, name, entity_type, summary ``` +### 1.10 `sem_groupby` 的 `membership` + +`membership` 说明一条 row 最终允许属于几个 semantic groups。它不决定模型调用 +是 pairwise 还是 listwise，也不决定是否启用 embedding Search-Filter。 + +```text +membership=None + 保留旧合同。当前静态执行仍给每条row一个group ID；join-map不额外限制一个 + changed group可以匹配多少个current groups。 + +membership="exclusive" + 每条row只属于一个semantic group。join-map把这个合同lower为每个changed group + 最多选择一个current group，也就是sem_join(k=1)。没有合适target时仍可新建group。 + +membership="overlapping" + 一条row可以属于多个semantic groups。API保留了这个名字，但当前静态执行和 + differential lowering尚未实现，调用时明确失败。 +``` + +这里最重要的边界是：`exclusive`是logical membership contract；`k=1`是compiler +为join-map选择的关系表达；pairwise/listwise和Search-Filter则是更下面的physical +execution。三层不能混在一起。 + ## 2. Paper-wise differential rules -这一节只写 paper-wise rules。公式尽量一行写完，不引入中间变量。这里默认 -`V` row 是可继续 merge 的 aggregate state。 +这一节只写paper-wise rules，不展开implementation-only的列投影。简单rule保持一行； +join-map显式写出changed groups和join result，避免把确定性keys、semantic predicate +和target merge压在一条难读的公式里。这里默认`V` row是可继续merge的aggregate +state。 记号： @@ -438,6 +463,8 @@ Ks = semantic keys Kr = partition_by relational keys θg = sem_groupby instruction θa = sem_agg / sem_map instruction +M = sem_groupby membership contract +κ(M)= 1 when M is exclusive; omitted when M is unspecified O = sem_agg output columns C = array_agg input columns o = array_agg output column @@ -446,24 +473,52 @@ m = min output column A* = A1, A2, ..., An ``` +本文用下面的缩写表示changed aggregate groups和current view之间的semantic +join-map： + +```text +JoinGroups(GΔ, V, Kr, θg, M) = + GΔ.sem_join( + V, + instruction=group_match(θg), + how="outer", + on=Kr when Kr is non-empty, + k=κ(M), + ) +``` + +`group_match(θg)`是compiler从原`sem_groupby` instruction改写出的“changed group +和current group是否属于同一组”谓词。`Ks`只是这个semantic predicate读取的内容， +绝不能写成`on=Ks`：`on`只接受必须精确相等的`partition_by` keys `Kr`。 + +`MergeByTarget`表示：matched rows按选中的current target合并；没有target的changed +group形成新row；未被触及的current row原样保留。如果多个changed groups命中同一个 +target，只生成一次合并后的target row。 + ### 2.1 `sem_groupby(...).sem_agg(...)` ```text -V = D.sem_groupby(Ks, θg).sem_agg(O, θa) +V = D.sem_groupby(Ks, partition_by=Kr, membership=M, θg).sem_agg(O, θa) ``` rule-re-group: ```text -V' = ΔD.sem_groupby(Ks, θg).sem_agg(O, θa).union(V).sem_groupby(Ks, θg).sem_agg(O, θa) +V' = ΔD.sem_groupby(Ks, partition_by=Kr, membership=M, θg).sem_agg(O, θa).union(V).sem_groupby(Ks, partition_by=Kr, membership=M, θg).sem_agg(O, θa) ``` rule-join-map: ```text -V' = ΔD.sem_groupby(Ks, θg).sem_agg(O, θa).sem_outer_join(V, on=Ks).sem_map(output_cols=O, instruction=θa) +GΔ = ΔD.sem_groupby(Ks, partition_by=Kr, membership=M, θg).sem_agg(O, θa) +J = JoinGroups(GΔ, V, Kr, θg, M) +V' = MergeByTarget(J, sem_agg(O, θa)).select(Kr + O) ``` +说人话：delta先在自己的确定性partition内形成changed groups；changed groups只和 +同partition的current groups做semantic匹配；最后只更新命中的旧group，并把未命中的 +delta作为新group加入。旧groups彼此不会重新比较。 + ### 2.2 `group_by(...).sem_agg(...)` ```text @@ -543,19 +598,21 @@ V' = ΔD.group_by(K).agg(A*).full_outer_join(V, on=K).merge(A*).select(K + outpu ### 2.6 `sem_groupby(...).agg(A1, ..., An)` ```text -V = D.sem_groupby(Ks, θg).agg(A1, A2, ..., An) +V = D.sem_groupby(Ks, partition_by=Kr, membership=M, θg).agg(A1, A2, ..., An) ``` rule-re-group: ```text -V' = ΔD.sem_groupby(Ks, θg).agg(A*).union(V).sem_groupby(Ks, θg).agg_merge(A*) +V' = ΔD.sem_groupby(Ks, partition_by=Kr, membership=M, θg).agg(A*).union(V).sem_groupby(Ks, partition_by=Kr, membership=M, θg).agg_merge(A*) ``` rule-join-map: ```text -V' = ΔD.sem_groupby(Ks, θg).agg(A*).sem_outer_join(V, on=Ks).merge(A*).select(outputs(A*)) +GΔ = ΔD.sem_groupby(Ks, partition_by=Kr, membership=M, θg).agg(A*) +J = JoinGroups(GΔ, V, Kr, θg, M) +V' = MergeByTarget(J, A*).select(Kr + outputs(A*)) ``` ### 2.7 `merge` and `agg_merge` @@ -583,7 +640,8 @@ agg_merge(min(c, m)) = min(m, m) ### 2.8 `partition_by` -`partition_by` 不改变 rule 形状。它只是把： +`partition_by` 不创造新的semantic rule family。它把输入先按普通关系key分区， +然后在每个分区内应用同一条semantic rule。也就是说，它把： ```text sem_groupby(Ks, θg) @@ -595,8 +653,13 @@ sem_groupby(Ks, θg) sem_groupby(Ks, partition_by=Kr, θg) ``` -输出和 join key 中额外带上 deterministic partition keys `Kr`。semantic keys `Ks` -仍然按 `sem_groupby` 的规则处理。 +输出中额外带上deterministic partition keys `Kr`。在join-map中，`Kr`还会下推成 +`sem_join(on=Kr)`，从而保证不同partition的rows根本不会成为semantic candidates。 +semantic keys `Ks`仍然只参与`θg`，不要求精确相等。 + +普通deterministic `group_by(Kr)`的增量维护不是semantic join。delta row可以直接按 +`Kr`找到受影响的partition；只有这些partition内部的changed semantic groups需要继续 +执行join-map或re-group。未被delta触及的partition原样保留。 ## 3. Implementation-wise differential rules @@ -633,19 +696,21 @@ aggregate state 的情况。 ### 3.2 `sem_groupby(...).sem_agg(...)` ```text -V = D.sem_groupby(Ks, θg).sem_agg(input_cols=I, output_cols=O, instruction=θa) +V = D.sem_groupby(Ks, partition_by=Kr, membership=M, θg).sem_agg(input_cols=I, output_cols=O, instruction=θa) ``` rule-join-map: ```text -V' = ΔD.sem_groupby(Ks, θg).sem_agg(input_cols=I, output_cols=O, instruction=θa).sem_outer_join(V, on=Ks).sem_map(output_cols=O, instruction=θa) +GΔ = ΔD.sem_groupby(Ks, partition_by=Kr, membership=M, θg).sem_agg(input_cols=I, output_cols=O, instruction=θa) +J = JoinGroups(GΔ, V, Kr, θg, M) +V' = MergeByTarget(J, sem_agg(input_cols=None, output_cols=O, instruction=θa)).select(Kr + O) ``` rule-re-group: ```text -V' = ΔD.sem_groupby(Ks, θg).sem_agg(input_cols=I, output_cols=O, instruction=θa).union(V).sem_groupby(Ks, θg).sem_agg(input_cols=None, output_cols=O, instruction=θa) +V' = ΔD.sem_groupby(Ks, partition_by=Kr, membership=M, θg).sem_agg(input_cols=I, output_cols=O, instruction=θa).union(V).sem_groupby(Ks, partition_by=Kr, membership=M, θg).sem_agg(input_cols=None, output_cols=O, instruction=θa) ``` rule-all-group: @@ -769,19 +834,23 @@ state，而不是把 old aggregate-state row 当作 raw evidence row 再做普 ### 3.7 `sem_groupby(...).agg(A1, ..., An)` ```text -V = D.sem_groupby(Ks, θg).agg(A1, A2, ..., An) +V = D.sem_groupby(Ks, partition_by=Kr, membership=M, θg).agg(A1, A2, ..., An) ``` rule-join-map: ```text -V' = ΔD.sem_groupby(Ks, θg).agg(A*).sem_outer_join(V, on=Ks).merge(A*).select(outputs(A*)) +GΔ = ΔD.sem_groupby(Ks, partition_by=Kr, membership=M, θg).agg(A*) +J = JoinGroups(GΔ, V, Kr, θg, M) +V' = MergeByTarget(J, A*).select(Kr + outputs(A*)) ``` rule-re-group: ```text -V' = ΔD.sem_groupby(Ks, θg).agg(A*).union(V).sem_groupby(Ks, θg).agg_merge(A*) +V' = ΔD.sem_groupby(Ks, partition_by=Kr, membership=M, θg).agg(A*) + .union(V) + .sem_groupby(Ks, partition_by=Kr, membership=M, θg).agg_merge(A*) ``` 和 deterministic `group_by(...).agg(A*)` 一样，mixed semantic `.agg(A*)` 暂不定义 @@ -812,14 +881,24 @@ agg_merge(min(c, m)) = min(m, m) ### 3.9 `partition_by` -`partition_by` 不改变 implementation rule shape。它只是在 -`sem_groupby(..., partition_by=Kr, ...)` 的 output、join key 和 touched-group key -里额外带上 deterministic partition keys `Kr`。 +`partition_by` 不改变implementation rule family，但会改变候选范围和output +identity： + +```text +static execution: + 先按Kr拆分DataFrame，再在每个partition内执行sem_groupby + +rule-re-group: + Kr随aggregate state保留，重新分组仍然只发生在各partition内部 + +rule-join-map: + compiler把Kr下推为sem_join(on=Kr) + 先做exact-key candidate restriction，再做semantic predicate +``` -当前 `rule-re-group` 支持 `partition_by`。当前 `rule-join-map` 明确拒绝 -`sem_groupby(..., partition_by=...)`，因为 semantic join 还不能同时强制 -deterministic partition equality；不能只在 join 后 filter，否则 outer-join 语义 -会改变。 +当前`rule-re-group`和`rule-join-map`都支持`partition_by`。这里不允许“先做全局 +semantic outer join，再在结果上filter Kr”，因为那会丢失本应保留的unmatched rows， +改变outer-join语义。 ## 4. 本文档不定义什么 diff --git a/docs/design/operator-api.md b/docs/design/operator-api.md index 3f7c0e2..6411174 100644 --- a/docs/design/operator-api.md +++ b/docs/design/operator-api.md @@ -733,6 +733,10 @@ Semantic grouping. In v0, grouped aggregation is expressed by chaining df.sem_groupby( input_cols=[...], instruction="...", + partition_by=None, + labels=None, + label_col="_label", + membership=None, ).sem_agg(...) ``` @@ -769,6 +773,35 @@ input and output. By default, `sem_groupby` is open-world grouping: groups are discovered from the rows using the membership condition. +`membership` states how many semantic groups one row is allowed to belong to. +It is a logical contract, not a choice between pairwise and listwise execution: + +```text +membership=None + Preserve the legacy contract. The current static executor still emits one + group id per row, but differential join-map does not assume that each + changed group has only one matching current group. + +membership="exclusive" + Each row belongs to exactly one semantic group. Under rule-join-map, each + changed group may therefore select at most one current group; the compiler + lowers this contract to sem_join(k=1). + +membership="overlapping" + One row may belong to multiple semantic groups. The name is reserved by the + API, but the static executor and differential planner do not implement this + output contract yet and fail explicitly. +``` + +In plain language, `exclusive` means “put this row in one group.” It does not +mean “the first semantically similar row wins,” and it does not prescribe how +the backend finds or ranks candidates. Those are physical execution choices. + +The current open-world static implementation uses pairwise semantic judgments +and forms one group from connected positive judgments. A future listwise +implementation may produce the same logical output through a different model +call shape. That physical choice is intentionally not part of this API. + For closed-world grouping, policy authors may provide explicit `labels`. The model must assign each row to exactly one declared label; if an `other` bucket is desired, it must be declared explicitly. The assigned label is written @@ -788,6 +821,8 @@ papers = rows.sem_groupby( ``` Multi-label and hierarchical labels are not part of the current contract. +Closed-world `labels` already define an exclusive classification task; do not +use `membership="overlapping"` with that form. `partition_by` optionally adds deterministic partition keys to semantic grouping. The semantic grouping rule is applied only within each partition, and @@ -801,6 +836,12 @@ entities = extracted_entities.sem_groupby( ) ``` +This is equivalent in meaning to “first split rows by exact `group_id`, then +run semantic grouping independently inside each exact group.” Rows from +different partitions never reach the embedding model or LLM as semantic pairs. +`partition_by` is therefore not a fuzzy grouping hint and is not another +Search-Filter threshold. + ### `group_by(...).agg(...)` and `sem_groupby(...).agg(...)` Grouped `agg(...)` accepts explicit aggregate specs: @@ -927,6 +968,8 @@ left.sem_join( right, instruction=instruction, how="inner", + on=None, + k=None, ) ``` @@ -948,13 +991,56 @@ left.sem_right_join(right, instruction=instruction) left.sem_outer_join(right, instruction=instruction) ``` -The canonical documentation form is `sem_join(..., instruction=..., how=...)`, -because it follows DataFrame style. +The canonical documentation form is +`sem_join(..., instruction=..., how=..., on=..., k=...)`, because it keeps the +logical join condition, deterministic keys, and match cardinality in one +DataFrame operator. + +`on` is an optional exact equality condition evaluated before the semantic +predicate: + +```python +changed_facts.sem_join( + current_facts, + on=["source_entity_id", "target_entity_id"], + instruction="The two rows express the same factual relationship.", + how="outer", +) +``` + +Only rows with equal `source_entity_id` and `target_entity_id` values are +eligible for semantic comparison. The columns named by `on` are ordinary +relational keys. They are not embedded and are not interpreted by the LLM. + +`k` optionally limits the number of matched right rows for each left row: + +```text +k=None Keep every right row that satisfies the semantic predicate. +k=1 Keep zero or one matching right row per left row. +k=K Keep zero to K matching right rows per left row. +``` + +The limit is left-to-right. It does not force a match: when no right row +satisfies the instruction, the left row remains unmatched and `how` determines +whether it appears in the assembled result. `k` must be a positive integer +when supplied. + +For `k` joins, the backend may use a listwise resolver or a pairwise resolver. +The listwise resolver gives one left row and its eligible right rows to the LM +and asks for at most `k` matching IDs. The pairwise resolver first verifies the +predicate pair by pair, then ranks verified matches. These are physical access +paths for the same `sem_join(..., k=K)` policy expression; policy code does not +choose them. Backend execution choices such as cascade, helper models, examples, and explanation tracing belong to the runtime/adapter layer. They should not appear in the policy author's logical join definition. +Search-Filter is also a physical optimization. When enabled for this join, it +uses embedding similarity to remove unlikely pairs before the listwise or +pairwise semantic resolver runs. It does not change `on`, `k`, `how`, or the +natural-language predicate, and it is not configured through this API. + Example: ```python @@ -966,6 +1052,7 @@ joined = topic_candidates.sem_join( contradictions, supersession, or forget/delete targets. """, how="left", + k=1, ) ``` @@ -1148,42 +1235,51 @@ V_prime = V.union( ) ``` -Semantic group-by with aggregation can use a coarse full-next-view maintenance -rule: +Semantic group-by with aggregation can use either re-group or join-map +maintenance. These are differential-rule choices, not alternative meanings of +the static `sem_groupby` API. + +For join-map, the compiler first aggregates the delta into changed groups, then +matches those changed groups against the current view: ```python V = ( D - .sem_groupby(input_cols=[...], instruction=group_instruction) + .sem_groupby( + input_cols=[...], + partition_by=partition_keys, + membership="exclusive", + instruction=group_instruction, + ) .sem_agg(...) ) delta_groups = ( delta_D - .sem_groupby(input_cols=[...], instruction=group_instruction) + .sem_groupby( + input_cols=[...], + partition_by=partition_keys, + membership="exclusive", + instruction=group_instruction, + ) .sem_agg(...) ) -V_prime = ( - delta_groups - .sem_join(V, instruction=join_instruction_prime, how="outer") - .sem_map( - output_cols=V.columns, - instruction=""" - Produce one next-view row: - merge matched delta group and existing view row; - keep unmatched existing view row; - add unmatched delta group row. - """, - ) - .select(V.columns) +joined = delta_groups.sem_join( + V, + instruction=join_instruction_prime, + how="outer", + on=partition_keys, + k=1, ) ``` -This is a coarse full-view rule. It does not split `delta_minus` and -`delta_plus`, and it does not introduce a separate delta-application operator. -The final `sem_map` is a column-level semantic merge over joined rows, not a -`sem_agg`. +The compiler then merges rows by the selected current target, emits unmatched +delta groups as new targets, and carries untouched current targets forward. +When `partition_keys` is empty, `on` is omitted. When membership is not +declared, `k` is omitted and the legacy all-matches join-map contract is +preserved. See `groupby_agg.md` for the complete paper-wise and +implementation-wise formulas. Join follows the usual relational delta shape: diff --git a/src/agent_memory/adapters/lotus/adapter.py b/src/agent_memory/adapters/lotus/adapter.py index 1bd5c0c..3e2b28c 100644 --- a/src/agent_memory/adapters/lotus/adapter.py +++ b/src/agent_memory/adapters/lotus/adapter.py @@ -4,6 +4,7 @@ from collections.abc import Mapping from dataclasses import dataclass, field +from hashlib import sha256 from typing import Any from uuid import uuid4 @@ -28,6 +29,7 @@ execute_flatten, execute_group_by, execute_join, + execute_let, execute_min, execute_select, execute_subtract, @@ -82,9 +84,16 @@ def __post_init__(self) -> None: def maintenance_execution_fingerprint(self) -> str: """Identify physical settings that can change maintained state.""" - return semantic_pair_profiles_fingerprint( + pair_fingerprint = semantic_pair_profiles_fingerprint( self.config.semantic_pair_profiles ) + if self.config.sem_join_topk_method == "listwise": + return pair_fingerprint + payload = ( + f"semantic_pair_profiles={pair_fingerprint}\n" + f"sem_join_topk_method={self.config.sem_join_topk_method}" + ) + return sha256(payload.encode("utf-8")).hexdigest() def execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: """Execute a logical query expression through LOTUS.""" @@ -106,6 +115,8 @@ def execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: return self._execute_traced_relational(query, inputs, execute_filter) case "group_by": return self._execute_traced_relational(query, inputs, execute_group_by) + case "let": + return self._execute_traced_relational(query, inputs, execute_let) case "array_agg": return self._execute_traced_relational(query, inputs, execute_array_agg) case "array_cat": diff --git a/src/agent_memory/adapters/lotus/context.py b/src/agent_memory/adapters/lotus/context.py index 796362c..2c5fb1d 100644 --- a/src/agent_memory/adapters/lotus/context.py +++ b/src/agent_memory/adapters/lotus/context.py @@ -24,6 +24,7 @@ "pairwise-heap", "listwise", ) +SEM_JOIN_TOPK_METHODS = SEM_TOPK_METHODS _LM_OWNED_KWARGS = { "cache", "max_batch_size", @@ -81,6 +82,7 @@ class LotusExecutionConfig: sem_join_cascade_args: Mapping[str, Any] | None = None sem_join_safe_mode: bool = False sem_join_progress_bar_desc: str = "Join comparisons" + sem_join_topk_method: str = "listwise" sem_groupby_default: bool = False sem_groupby_pair_batch_size: int | None = None @@ -104,6 +106,11 @@ def __post_init__(self) -> None: raise ValueError( "semantic_trace_snapshot_mode must be 'compact' or 'full'" ) + if self.sem_join_topk_method not in SEM_JOIN_TOPK_METHODS: + raise ValueError( + "sem_join_topk_method must be one of: " + + ", ".join(SEM_JOIN_TOPK_METHODS) + ) for query_digest, profile in self.semantic_pair_profiles.items(): if not isinstance(query_digest, str) or not query_digest: raise ValueError("semantic pair profile keys must be query digests") diff --git a/src/agent_memory/adapters/lotus/pair_execution.py b/src/agent_memory/adapters/lotus/pair_execution.py index 079fd49..f40ed95 100644 --- a/src/agent_memory/adapters/lotus/pair_execution.py +++ b/src/agent_memory/adapters/lotus/pair_execution.py @@ -146,6 +146,11 @@ class SemanticPairSite: instruction: str semantic_columns: tuple[str, ...] partition_by: tuple[str, ...] + direction: str + left_id_columns: tuple[str, ...] + right_id_columns: tuple[str, ...] + left_text_columns: tuple[str, ...] + right_text_columns: tuple[str, ...] query_digests: tuple[str, ...] def to_dict(self) -> dict[str, object]: @@ -158,6 +163,11 @@ def to_dict(self) -> dict[str, object]: "instruction": self.instruction, "semantic_columns": list(self.semantic_columns), "partition_by": list(self.partition_by), + "direction": self.direction, + "left_id_columns": list(self.left_id_columns), + "right_id_columns": list(self.right_id_columns), + "left_text_columns": list(self.left_text_columns), + "right_text_columns": list(self.right_text_columns), "query_digests": list(self.query_digests), } @@ -177,9 +187,15 @@ def semantic_pair_site_contract(query: Any) -> dict[str, object]: instruction = str(params.get("instruction") or "") semantic_columns = tuple(sorted(set(re.findall(r"\{([^{}]+)\}", instruction)))) partition_by = () + elif operator == "sem_filter": + instruction = str(params.get("instruction") or "") + semantic_columns = tuple( + sorted(set(re.findall(r"\{([^{}]+:[^{}]+)\}", instruction))) + ) + partition_by = () else: raise ValueError( - "semantic pair sites support sem_join and pairwise sem_groupby" + "semantic pair sites support sem_filter, sem_join, and pairwise sem_groupby" ) instruction = str(params.get("instruction") or "") if not instruction: @@ -194,6 +210,79 @@ def semantic_pair_site_contract(query: Any) -> dict[str, object]: } +def semantic_pair_site_physical_contract(query: Any) -> dict[str, object]: + """Return endpoint columns for one supported semantic pair site.""" + + if query.op == "sem_join": + return { + "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], + } + if query.op == "sem_groupby": + return { + "direction": "symmetric", + "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], + } + if query.op != "sem_filter": + raise ValueError(f"unsupported semantic pair site operator {query.op!r}") + + instruction = str(query.params.get("instruction") or "") + placeholders = tuple( + dict.fromkeys(re.findall(r"\{([^{}]+):([^{}]+)\}", instruction)) + ) + qualifiers = {qualifier for _column, qualifier in placeholders} + earlier = [qualifier for qualifier in qualifiers if qualifier.startswith("earlier")] + later = [qualifier for qualifier in qualifiers if qualifier.startswith("later")] + if len(earlier) != 1 or len(later) != 1 or len(qualifiers) != 2: + raise ValueError( + "pair-shaped sem_filter sites require one earlier and one later qualifier" + ) + earlier_qualifier = earlier[0] + later_qualifier = later[0] + left_text = tuple( + f"{column}:{qualifier}" + for column, qualifier in placeholders + if qualifier == earlier_qualifier + ) + right_text = tuple( + f"{column}:{qualifier}" + for column, qualifier in placeholders + if qualifier == later_qualifier + ) + input_columns = tuple( + str(column) for column in query.inputs[0].params.get("columns", ()) + ) + left_ids = tuple( + column + for column in input_columns + if column.endswith(f":{earlier_qualifier}") + and column.split(":", 1)[0].endswith("_id") + ) + right_ids = tuple( + column + for column in input_columns + if column.endswith(f":{later_qualifier}") + and column.split(":", 1)[0].endswith("_id") + ) + if not left_text or not right_text: + raise ValueError( + "pair-shaped sem_filter sites require semantic text on both sides" + ) + return { + "direction": "right-to-left", + "left_id_columns": list(left_ids or left_text), + "right_id_columns": list(right_ids or right_text), + "left_text_columns": list(left_text), + "right_text_columns": list(right_text), + } + + def semantic_pair_site_id(query: Any) -> str: """Return a stable identity for one natural-language pair predicate.""" @@ -486,6 +575,7 @@ def _endpoint_pair_id(left_id: str, right_id: str) -> str: "select_semantic_pair_candidates", "semantic_pair_site_contract", "semantic_pair_site_id", + "semantic_pair_site_physical_contract", "semantic_pair_profiles_fingerprint", "write_semantic_pair_execution_trace", ] diff --git a/src/agent_memory/adapters/lotus/relational.py b/src/agent_memory/adapters/lotus/relational.py index e8db7a7..ce03ab0 100644 --- a/src/agent_memory/adapters/lotus/relational.py +++ b/src/agent_memory/adapters/lotus/relational.py @@ -74,6 +74,26 @@ def execute_group_by( return source +def execute_let( + query: QueryExpr, + inputs: Mapping[str, Any], + execute: Callable[[QueryExpr, Mapping[str, Any]], Any], +) -> Any: + """Evaluate one relation once and bind it while executing a local query body.""" + + if len(query.inputs) != 2: + raise ValueError("let expects one bound relation and one query body") + name = str(query.params.get("name", "")) + if not name: + raise ValueError("let requires a non-empty binding name") + if name in inputs: + raise ValueError(f"let binding conflicts with an existing input: {name!r}") + value = execute(query.inputs[0], inputs) + scoped_inputs = dict(inputs) + scoped_inputs[name] = value + return execute(query.inputs[1], scoped_inputs) + + def execute_concat( query: QueryExpr, inputs: Mapping[str, Any], diff --git a/src/agent_memory/adapters/lotus/sem_groupby.py b/src/agent_memory/adapters/lotus/sem_groupby.py index bc58639..2d618b6 100644 --- a/src/agent_memory/adapters/lotus/sem_groupby.py +++ b/src/agent_memory/adapters/lotus/sem_groupby.py @@ -42,6 +42,13 @@ def execute_sem_groupby( ) -> Any: """Assign deterministic semantic group ids to rows.""" + membership = query.params.get("membership") + if membership == "overlapping": + raise NotImplementedError( + "overlapping sem_groupby membership is not implemented" + ) + if membership not in {None, "exclusive"}: + raise ValueError(f"Unsupported sem_groupby membership: {membership!r}") input_cols = tuple(str(column) for column in query.params["input_cols"]) partition_by = tuple(str(column) for column in query.params.get("partition_by", ())) labels = tuple(query.params.get("labels") or ()) diff --git a/src/agent_memory/adapters/lotus/sem_join.py b/src/agent_memory/adapters/lotus/sem_join.py index 30c051a..c281a91 100644 --- a/src/agent_memory/adapters/lotus/sem_join.py +++ b/src/agent_memory/adapters/lotus/sem_join.py @@ -2,7 +2,9 @@ from __future__ import annotations +from collections import defaultdict from collections.abc import Callable, Mapping, Sequence +import json from typing import Any import pandas as pd @@ -37,6 +39,16 @@ def execute_sem_join( digest = query_digest(query) profile = context.config.semantic_pair_profiles.get(digest) + id_columns = _join_id_columns(query.params.get("id_columns")) + if query.params.get("k") is not None and context.config.sem_join_cascade_args is not None: + raise ValueError("top-k sem_join cannot be combined with LOTUS sem_join cascade") + if query.params.get("on") and context.config.sem_join_cascade_args is not None: + # LOTUS 1.1.4 cascades rebuild the full Cartesian product and cannot + # consume exact-key candidate pairs. Keep that backend boundary explicit + # rather than silently dropping `on` semantics or forking LOTUS here. + raise ValueError( + "exact-key sem_join cannot be combined with LOTUS sem_join cascade" + ) if profile is not None and profile.mode in {"search-filter", "proxy-only"}: if context.config.sem_join_cascade_args is not None: raise ValueError( @@ -55,6 +67,7 @@ def execute_sem_join( right, (), how=str(query.params.get("how", "inner")), + id_columns=id_columns, ) write_sem_join_result_trace( context.config.trace_dir(), @@ -64,10 +77,61 @@ def execute_sem_join( result=result, payload={ "skipped_pairwise": True, + **( + { + "sem_join_topk_method": context.config.sem_join_topk_method, + "k": int(query.params["k"]), + } + if query.params.get("k") is not None + else {} + ), }, ) return result - if profile is not None and profile.mode in {"search-filter", "proxy-only"}: + topk_payload: Mapping[str, Any] = {} + if query.params.get("k") is not None: + from agent_memory.adapters.lotus.sem_topk_join import evaluate_sem_topk_join + + join_results, topk_payload = evaluate_sem_topk_join( + query, + left, + right, + context, + profile=profile, + ) + elif query.params.get("on"): + left_series, right_series, left_label, right_label, instruction = join_series( + left, + right, + str(query.params["instruction"]), + ) + pairs = semantic_join_pair_candidates( + left_series, + right_series, + left_frame=left, + right_frame=right, + on=tuple(str(column) for column in query.params["on"]), + ) + if profile is not None and profile.mode in {"search-filter", "proxy-only"}: + assert context.pair_embedding_provider is not None + join_results = evaluate_profiled_semantic_join( + query, + left, + right, + context.config, + profile=profile, + embedding_provider=context.pair_embedding_provider, + pairs=pairs, + ) + else: + join_results = verify_semantic_join_candidates( + pairs, + left_label=left_label, + right_label=right_label, + instruction=instruction, + config=context.config, + ) + elif profile is not None and profile.mode in {"search-filter", "proxy-only"}: assert context.pair_embedding_provider is not None join_results = evaluate_profiled_semantic_join( query, @@ -84,6 +148,7 @@ def execute_sem_join( right, join_results, how=str(query.params.get("how", "inner")), + id_columns=id_columns, ) write_sem_join_result_trace( context.config.trace_dir(), @@ -102,6 +167,7 @@ def execute_sem_join( and profile.mode in {"search-filter", "proxy-only"} else {} ), + **topk_payload, }, ) return result @@ -220,6 +286,7 @@ def evaluate_profiled_semantic_join( *, profile: SemanticPairExecutionProfile, embedding_provider: EmbeddingProvider, + pairs: pd.DataFrame | None = None, ) -> list[tuple[Any, Any, str | None]]: """Execute one profiled semantic join over embedding-selected pairs.""" @@ -228,7 +295,8 @@ def evaluate_profiled_semantic_join( right, str(query.params["instruction"]), ) - pairs = semantic_join_pair_candidates(left_series, right_series) + if pairs is None: + pairs = semantic_join_pair_candidates(left_series, right_series) selection = select_semantic_pair_candidates( pairs, profile=profile, @@ -254,6 +322,28 @@ def evaluate_profiled_semantic_join( for _row_index, row in candidates.iterrows() ] + return verify_semantic_join_candidates( + candidates, + left_label=left_label, + right_label=right_label, + instruction=instruction, + config=config, + ) + + +def verify_semantic_join_candidates( + candidates: pd.DataFrame, + *, + left_label: str, + right_label: str, + instruction: str, + config: LotusExecutionConfig, +) -> list[tuple[Any, Any, str | None]]: + """Verify canonical semantic-join candidate pairs with the LOTUS oracle.""" + + if candidates.empty: + return [] + import lotus from lotus.sem_ops.sem_filter import sem_filter from lotus.templates import task_instructions @@ -320,20 +410,73 @@ def evaluate_profiled_semantic_join( def semantic_join_pair_candidates( left: pd.Series, right: pd.Series, + *, + left_frame: pd.DataFrame | None = None, + right_frame: pd.DataFrame | None = None, + on: Sequence[str] = (), ) -> pd.DataFrame: - """Return the semantic join cross product in a canonical pair schema.""" + """Return semantic join pairs after optional exact-key restriction.""" + + if on and (left_frame is None or right_frame is None): + raise ValueError("exact-key semantic join candidates require both input frames") + missing = ( + sorted( + set(on).difference(left_frame.columns).union( + set(on).difference(right_frame.columns) + ) + ) + if left_frame is not None and right_frame is not None + else [] + ) + if missing: + raise ValueError(f"sem_join on columns not found on both sides: {missing}") - return pd.DataFrame( - [ + right_by_key: dict[str, list[Any]] = defaultdict(list) + if on: + assert right_frame is not None + for right_id in right.index: + right_by_key[_semantic_join_key(right_frame.loc[right_id], on)].append( + right_id + ) + + rows: list[dict[str, Any]] = [] + for left_id, left_value in left.items(): + if on: + assert left_frame is not None + right_ids: Sequence[Any] = right_by_key.get( + _semantic_join_key(left_frame.loc[left_id], on), + (), + ) + else: + right_ids = right.index + rows.extend( { PAIR_LEFT_ID_COLUMN: left_id, PAIR_RIGHT_ID_COLUMN: right_id, PAIR_LEFT_TEXT_COLUMN: left_value, - PAIR_RIGHT_TEXT_COLUMN: right_value, + PAIR_RIGHT_TEXT_COLUMN: right.loc[right_id], } - for left_id, left_value in left.items() - for right_id, right_value in right.items() - ] + for right_id in right_ids + ) + return pd.DataFrame( + rows, + columns=pd.Index( + ( + PAIR_LEFT_ID_COLUMN, + PAIR_RIGHT_ID_COLUMN, + PAIR_LEFT_TEXT_COLUMN, + PAIR_RIGHT_TEXT_COLUMN, + ) + ), + ) + + +def _semantic_join_key(row: pd.Series, columns: Sequence[str]) -> str: + return json.dumps( + [row[column] for column in columns], + ensure_ascii=False, + default=str, + separators=(",", ":"), ) @@ -521,6 +664,7 @@ def assemble_join_frame( how: str, return_explanations: bool = False, explanation_column: str = "explanation_join", + id_columns: tuple[str, str] | None = None, ) -> pd.DataFrame: """Assemble join output and deterministic unmatched rows.""" @@ -530,6 +674,11 @@ def assemble_join_frame( left_columns, right_columns = renamed_columns(left, right) output_columns = list(left_columns.values()) + list(right_columns.values()) + if id_columns is not None: + conflicts = sorted(set(id_columns).intersection(output_columns)) + if conflicts: + raise ValueError(f"sem_join id columns conflict with output columns: {conflicts}") + output_columns.extend(id_columns) if return_explanations: output_columns.append(explanation_column) @@ -540,6 +689,9 @@ def assemble_join_frame( matched_left.add(left_id) matched_right.add(right_id) row = join_row(left.loc[left_id], right.loc[right_id], left_columns, right_columns) + if id_columns is not None: + row[id_columns[0]] = left_id + row[id_columns[1]] = right_id if return_explanations: row[explanation_column] = explanation rows.append(row) @@ -548,6 +700,9 @@ def assemble_join_frame( for left_id in left.index: if left_id not in matched_left: row = left_only_row(left.loc[left_id], left_columns, right_columns) + if id_columns is not None: + row[id_columns[0]] = left_id + row[id_columns[1]] = pd.NA if return_explanations: row[explanation_column] = pd.NA rows.append(row) @@ -556,6 +711,9 @@ def assemble_join_frame( for right_id in right.index: if right_id not in matched_right: row = right_only_row(right.loc[right_id], left_columns, right_columns) + if id_columns is not None: + row[id_columns[0]] = pd.NA + row[id_columns[1]] = right_id if return_explanations: row[explanation_column] = pd.NA rows.append(row) @@ -563,6 +721,24 @@ def assemble_join_frame( return pd.DataFrame(rows, columns=output_columns) +def _join_id_columns(value: object) -> tuple[str, str] | None: + """Validate optional internal left/right row-ID output columns.""" + + if value is None: + return None + if ( + not isinstance(value, Sequence) + or isinstance(value, (str, bytes)) + or len(value) != 2 + or any(not isinstance(column, str) or not column for column in value) + ): + raise ValueError("sem_join id_columns must contain two non-empty column names") + left_id_column, right_id_column = (str(column) for column in value) + if left_id_column == right_id_column: + raise ValueError("sem_join id_columns must be distinct") + return left_id_column, right_id_column + + def cascade_args_from_mapping(value: Any) -> Any: """Convert row-like cascade args into LOTUS CascadeArgs.""" diff --git a/src/agent_memory/adapters/lotus/sem_topk_join.py b/src/agent_memory/adapters/lotus/sem_topk_join.py new file mode 100644 index 0000000..95122ee --- /dev/null +++ b/src/agent_memory/adapters/lotus/sem_topk_join.py @@ -0,0 +1,459 @@ +"""Cardinality-bounded semantic join execution.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import json +from typing import Any, cast + +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_LEFT_TEXT_COLUMN, + PAIR_RIGHT_ID_COLUMN, + PAIR_RIGHT_TEXT_COLUMN, + SemanticPairExecutionProfile, + select_semantic_pair_candidates, + write_semantic_pair_execution_trace, +) +from agent_memory.adapters.lotus.sem_topk import LOTUS_PAIRWISE_METHODS +from agent_memory.adapters.lotus.structured import normalize_strategy +from agent_memory.policy.logical import QueryExpr +from agent_memory.tracing.semantic import query_digest, write_pair_trace + +LISTWISE_JOIN_MAX_TOKENS = 1024 +LISTWISE_JOIN_SYSTEM_PROMPT = ( + "The user will provide one left row, a semantic join condition, and candidate " + "right rows. Your job is to select the right rows that satisfy the join " + "condition for that left row. Select at most max_matches candidates. Do not " + "select merely related candidates or fill the limit when fewer candidates " + "match. If too many candidates match, select the strongest matches. Return an " + "empty selected_ids array when none match. Return only a JSON object." +) + + +@dataclass(frozen=True) +class _ListwiseTask: + left_id: Any + prompt: list[dict[str, Any]] + candidate_positions: Mapping[str, int] + + +def evaluate_sem_topk_join( + query: QueryExpr, + left: pd.DataFrame, + right: pd.DataFrame, + context: LotusExecutionContext, + *, + profile: SemanticPairExecutionProfile | None, +) -> tuple[list[tuple[Any, Any, str | None]], Mapping[str, Any]]: + """Evaluate a per-left top-k semantic join with one configured access path.""" + + from agent_memory.adapters.lotus.sem_join import join_series + from agent_memory.adapters.lotus.sem_join import semantic_join_pair_candidates + + k = _positive_k(query.params.get("k")) + left_series, right_series, left_label, right_label, instruction = join_series( + left, + right, + str(query.params["instruction"]), + ) + pairs = semantic_join_pair_candidates( + left_series, + right_series, + left_frame=left, + right_frame=right, + on=tuple(str(column) for column in query.params.get("on", ())), + ) + candidates, scores = _apply_pair_profile( + query, + pairs, + context, + profile=profile, + ) + + method = context.config.sem_join_topk_method + if profile is not None and profile.mode == "proxy-only": + join_results = _proxy_topk(candidates, scores=scores, k=k) + method = "proxy-only" + retry_count = 0 + elif method == "listwise": + join_results, retry_count = _listwise_topk( + candidates, + instruction=instruction, + left_label=left_label, + right_label=right_label, + k=k, + context=context, + ) + else: + join_results = _pairwise_topk( + candidates, + instruction=instruction, + left_label=left_label, + right_label=right_label, + k=k, + method=method, + context=context, + ) + retry_count = 0 + + return join_results, { + "sem_join_topk_method": method, + "k": k, + "on": list(query.params.get("on", ())), + "eligible_pair_count": len(pairs), + "oracle_candidate_count": len(candidates), + "selected_pair_count": len(join_results), + "structured_retry_count": retry_count, + } + + +def _apply_pair_profile( + query: QueryExpr, + pairs: pd.DataFrame, + context: LotusExecutionContext, + *, + profile: SemanticPairExecutionProfile | None, +) -> tuple[pd.DataFrame, tuple[float, ...]]: + """Apply an optional candidate profile without changing join semantics by default.""" + + if profile is None or profile.mode == "oracle-only": + return pairs.reset_index(drop=True), () + if profile.direction != "left-to-right": + raise ValueError("top-k sem_join pair profiles must use left-to-right direction") + if context.pair_embedding_provider is None: + raise ValueError(f"{profile.mode} requires a pair embedding provider") + + selection = select_semantic_pair_candidates( + pairs, + profile=profile, + embedding_provider=context.pair_embedding_provider, + ) + write_semantic_pair_execution_trace( + context.config.trace_dir(), + operator="sem_join", + query_digest_value=query_digest(query), + profile=profile, + selection=selection, + ) + positions = list(selection.selected_positions) + selected = pairs.iloc[positions].reset_index(drop=True) + selected_scores = tuple(selection.similarity_scores[position] for position in positions) + return selected, selected_scores + + +def _listwise_topk( + candidates: pd.DataFrame, + *, + instruction: str, + left_label: str, + right_label: str, + k: int, + context: LotusExecutionContext, +) -> tuple[list[tuple[Any, Any, str | None]], int]: + """Resolve zero-to-k right matches for each left row in listwise batches.""" + + if candidates.empty: + return [], 0 + + import lotus + + lm = lotus.settings.lm + if lm is None: + raise ValueError("listwise top-k sem_join requires a configured language model") + tasks = _listwise_tasks( + candidates, + instruction=instruction, + left_label=left_label, + right_label=right_label, + k=k, + ) + pending = list(tasks) + selected_by_left: dict[Any, tuple[int, ...]] = {} + retry_count = 0 + last_error: ValueError | None = None + for attempt in range(context.config.structured_parse_retries + 1): + if not pending: + break + lm_call: Any = lm + output = lm_call( + [task.prompt for task in pending], + progress_bar_desc="Listwise join resolution", + max_tokens=LISTWISE_JOIN_MAX_TOKENS, + response_format={"type": "json_object"}, + ) + raw_outputs = list(getattr(output, "outputs", ())) + if len(raw_outputs) != len(pending): + raise ValueError( + "listwise top-k sem_join returned an unexpected number of outputs: " + f"expected {len(pending)}, got {len(raw_outputs)}" + ) + retry: list[_ListwiseTask] = [] + for task, raw_output in zip(pending, raw_outputs, strict=True): + raw = str(raw_output) + try: + selected_ids = _parse_listwise_ids( + raw, + valid_ids=set(task.candidate_positions), + k=k, + ) + except ValueError as error: + last_error = error + retry.append(task) + continue + selected_by_left[task.left_id] = tuple( + task.candidate_positions[candidate_id] for candidate_id in selected_ids + ) + if retry: + retry_count += len(retry) + pending = retry + + if pending: + assert last_error is not None + raise ValueError( + "listwise top-k sem_join returned invalid structured output after " + f"{context.config.structured_parse_retries + 1} attempt(s): {last_error}" + ) from last_error + + selected_positions = { + position + for positions in selected_by_left.values() + for position in positions + } + rows = [] + results: list[tuple[Any, Any, str | None]] = [] + for position, row in candidates.iterrows(): + matched = position in selected_positions + rows.append( + { + "operator": "sem_join", + "instruction": instruction, + "decision_source": "listwise", + "left_id": row[PAIR_LEFT_ID_COLUMN], + "right_id": row[PAIR_RIGHT_ID_COLUMN], + "left": row[PAIR_LEFT_TEXT_COLUMN], + "right": row[PAIR_RIGHT_TEXT_COLUMN], + "parsed_output": matched, + } + ) + if matched: + results.append( + ( + row[PAIR_LEFT_ID_COLUMN], + row[PAIR_RIGHT_ID_COLUMN], + None, + ) + ) + write_pair_trace( + context.config.trace_dir(), + operator="sem_join", + rows=rows, + ) + return results, retry_count + + +def _listwise_tasks( + candidates: pd.DataFrame, + *, + instruction: str, + left_label: str, + right_label: str, + k: int, +) -> list[_ListwiseTask]: + tasks: list[_ListwiseTask] = [] + for left_id, group in candidates.groupby(PAIR_LEFT_ID_COLUMN, sort=False): + candidate_positions = { + f"candidate_{position}": int(position) for position in group.index + } + right_candidates = [ + { + "id": candidate_id, + right_label: candidates.loc[position, PAIR_RIGHT_TEXT_COLUMN], + } + for candidate_id, position in candidate_positions.items() + ] + payload = { + "join_condition": instruction, + "max_matches": min(k, len(group)), + "left": {left_label: group.iloc[0][PAIR_LEFT_TEXT_COLUMN]}, + "right_candidates": right_candidates, + "output_schema": {"selected_ids": ["candidate_id"]}, + } + tasks.append( + _ListwiseTask( + left_id=left_id, + prompt=[ + {"role": "system", "content": LISTWISE_JOIN_SYSTEM_PROMPT}, + { + "role": "user", + "content": json.dumps(payload, ensure_ascii=False, default=str), + }, + ], + candidate_positions=candidate_positions, + ) + ) + return tasks + + +def _parse_listwise_ids( + raw_output: str, + *, + valid_ids: set[str], + k: int, +) -> tuple[str, ...]: + """Parse a strict zero-to-k listwise join response.""" + + if not raw_output.strip(): + raise ValueError("listwise top-k sem_join output is empty") + try: + payload = json.loads(raw_output) + except json.JSONDecodeError as error: + raise ValueError("listwise top-k sem_join output is not valid JSON") from error + if not isinstance(payload, dict): + raise ValueError("listwise top-k sem_join output must be a JSON object") + selected_ids = payload.get("selected_ids") + if not isinstance(selected_ids, list) or any( + not isinstance(candidate_id, str) for candidate_id in selected_ids + ): + raise ValueError("listwise top-k sem_join selected_ids must be a list of strings") + if len(selected_ids) > k: + raise ValueError(f"listwise top-k sem_join selected more than k={k} candidates") + if len(set(selected_ids)) != len(selected_ids): + raise ValueError("listwise top-k sem_join selected_ids must be unique") + unknown = sorted(set(selected_ids).difference(valid_ids)) + if unknown: + raise ValueError( + f"listwise top-k sem_join selected_ids contain unknown IDs: {unknown}" + ) + return tuple(selected_ids) + + +def _pairwise_topk( + candidates: pd.DataFrame, + *, + instruction: str, + left_label: str, + right_label: str, + k: int, + method: str, + context: LotusExecutionContext, +) -> list[tuple[Any, Any, str | None]]: + """Verify candidate pairs, then rank excess positive matches with LOTUS top-k.""" + + if candidates.empty: + return [] + + from agent_memory.adapters.lotus.sem_join import verify_semantic_join_candidates + + lotus_method = LOTUS_PAIRWISE_METHODS.get(method) + if lotus_method is None: + raise ValueError(f"Unsupported top-k sem_join method {method!r}") + verified = verify_semantic_join_candidates( + candidates, + left_label=left_label, + right_label=right_label, + instruction=instruction, + config=context.config, + ) + explanations = {(left_id, right_id): explanation for left_id, right_id, explanation in verified} + positives = candidates[ + [ + (row[PAIR_LEFT_ID_COLUMN], row[PAIR_RIGHT_ID_COLUMN]) in explanations + for _, row in candidates.iterrows() + ] + ] + results: list[tuple[Any, Any, str | None]] = [] + for left_id, grouped in positives.groupby(PAIR_LEFT_ID_COLUMN, sort=False): + group = cast(pd.DataFrame, grouped) + selected = group + if len(group) > k: + selected = _rank_pairwise_group( + group, + instruction=instruction, + left_label=left_label, + right_label=right_label, + k=k, + lotus_method=lotus_method, + context=context, + ) + for _, row in selected.iterrows(): + right_id = row[PAIR_RIGHT_ID_COLUMN] + results.append((left_id, right_id, explanations[(left_id, right_id)])) + return results + + +def _rank_pairwise_group( + group: pd.DataFrame, + *, + instruction: str, + left_label: str, + right_label: str, + k: int, + lotus_method: str, + context: LotusExecutionContext, +) -> pd.DataFrame: + left_text = str(group.iloc[0][PAIR_LEFT_TEXT_COLUMN]) + source = pd.DataFrame( + { + left_label: [left_text] * len(group), + right_label: group[PAIR_RIGHT_TEXT_COLUMN].tolist(), + "_pair_position": group.index.tolist(), + } + ) + ranking_instruction = ( + f"Rank {{{right_label}}} by how strongly it satisfies this semantic join " + f"condition with the fixed {{{left_label}}}: {instruction}" + ) + result = source.sem_topk( + ranking_instruction, + K=k, + method=lotus_method, + strategy=normalize_strategy(context.config.sem_topk_strategy), + cascade_threshold=context.config.sem_topk_cascade_threshold, + return_stats=context.config.sem_topk_return_stats, + safe_mode=context.config.sem_topk_safe_mode, + return_explanations=False, + ) + result_frame = result[0] if isinstance(result, tuple) else result + positions = [int(position) for position in result_frame["_pair_position"]] + return group.loc[positions] + + +def _proxy_topk( + candidates: pd.DataFrame, + *, + scores: Sequence[float], + k: int, +) -> list[tuple[Any, Any, str | None]]: + """Use embedding-selected pairs directly and cap each left side by score.""" + + if len(candidates) != len(scores): + raise ValueError("proxy-only top-k sem_join requires one score per candidate") + ranked = candidates.copy() + ranked["_similarity_score"] = list(scores) + results: list[tuple[Any, Any, str | None]] = [] + for _left_id, group in ranked.groupby(PAIR_LEFT_ID_COLUMN, sort=False): + ordered = group.sort_values( + by=["_similarity_score", PAIR_RIGHT_ID_COLUMN], + ascending=[False, True], + kind="stable", + ).head(k) + results.extend( + ( + row[PAIR_LEFT_ID_COLUMN], + row[PAIR_RIGHT_ID_COLUMN], + None, + ) + for _, row in ordered.iterrows() + ) + return results + + +def _positive_k(value: object) -> int: + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise ValueError("top-k sem_join requires a positive integer k") + return value diff --git a/src/agent_memory/planner/rules.py b/src/agent_memory/planner/rules.py index bbe2e61..3c5178d 100644 --- a/src/agent_memory/planner/rules.py +++ b/src/agent_memory/planner/rules.py @@ -26,6 +26,9 @@ CHANGED_FRAGMENT_BINARY_OPS = {"concat", "union"} GROUP_ID_COLUMN = "_agent_memory_group_id" CHANGED_MARKER_COLUMN = "_changed" +JOIN_MAP_LEFT_ID_COLUMN = "_agent_memory_join_map_left_id" +JOIN_MAP_RIGHT_ID_COLUMN = "_agent_memory_join_map_right_id" +JOIN_MAP_BINDING_NAME = "__agent_memory_join_map_pairs" GROUPED_AGG_RULE_ALIASES = { "compressed": "rule-all-group", "changed-aware": "rule-all-group-optimized", @@ -37,7 +40,6 @@ "join-map", "rule-join-map", "rule-re-group", - "prefer-join-map", "rule-all-group", "rule-all-group-optimized", ) @@ -157,15 +159,6 @@ def grouped_agg_rule(self) -> str: return self._grouped_agg_rule - def _grouped_agg_rule_for(self, groupby: QueryExpr) -> str: - """Select a supported grouped rule without executing a speculative rewrite.""" - - if self._grouped_agg_rule != "prefer-join-map": - return self._grouped_agg_rule - if groupby.op == "sem_groupby" and groupby.params.get("partition_by"): - return "rule-re-group" - return "rule-join-map" - def differentiate( self, query: QueryExpr, @@ -445,7 +438,7 @@ def _differentiate_min_view( ) keys = tuple(str(key) for key in source.params["keys"]) output_col = str(aggregate.params["output_col"]) - if self._grouped_agg_rule_for(source) == "rule-join-map": + if self._grouped_agg_rule == "rule-join-map": joined = QueryExpr( op="join", inputs=(changed_aggregate, current_view), @@ -620,7 +613,7 @@ def _differentiate_grouped_agg_view( source_query=source_query, instruction_rewriter=instruction_rewriter, ) - if self._grouped_agg_rule_for(groupby) == "rule-join-map": + if self._grouped_agg_rule == "rule-join-map": return self._build_grouped_agg_join_map_candidate( aggregate=aggregate, groupby=groupby, @@ -777,12 +770,6 @@ def _build_grouped_agg_join_map_candidate( ) -> QueryExpr: """Build rule-join-map for grouped agg(A*) patterns without array specs.""" - if groupby.op == "sem_groupby" and groupby.params.get("partition_by"): - raise NotImplementedError( - "sem_groupby partition_by is not supported by rule-join-map; " - "the semantic join cannot yet enforce deterministic partition keys." - ) - changed_groupby = QueryExpr( op=groupby.op, inputs=(changed_group_input,), @@ -793,6 +780,15 @@ def _build_grouped_agg_join_map_candidate( inputs=(changed_groupby,), params=aggregate.params, ) + if groupby.op == "sem_groupby": + return self._build_semantic_join_map( + aggregate=aggregate, + groupby=groupby, + changed_aggregate=changed_aggregate, + current_view=current_view, + final_columns=final_columns, + instruction_rewriter=instruction_rewriter, + ) if groupby.op == "group_by": joined: QueryExpr = QueryExpr( op="join", @@ -800,17 +796,7 @@ def _build_grouped_agg_join_map_candidate( params={"on": tuple(str(key) for key in groupby.params["keys"]), "how": "outer"}, ) else: - joined = QueryExpr( - op="sem_join", - inputs=(changed_aggregate, current_view), - params={ - "instruction": instruction_rewriter.groupby_to_join( - str(groupby.params["instruction"]), - input_cols=tuple(str(column) for column in groupby.params["input_cols"]), - ), - "how": "outer", - }, - ) + raise TypeError(f"Unsupported grouped aggregate input: {groupby.op}") preserved_columns = self._grouped_preserved_columns(groupby) mapped = joined array_outputs: list[str] = [] @@ -869,6 +855,245 @@ def _build_grouped_agg_join_map_candidate( params={"columns": tuple(final_columns)}, ) + def _build_semantic_join_map( + self, + *, + aggregate: QueryExpr, + groupby: QueryExpr, + changed_aggregate: QueryExpr, + current_view: QueryExpr, + final_columns: Sequence[str], + instruction_rewriter: DifferentialInstructionRewriter, + ) -> QueryExpr: + """Join changed groups to current targets and merge each target once.""" + + reserved = { + JOIN_MAP_LEFT_ID_COLUMN, + JOIN_MAP_RIGHT_ID_COLUMN, + }.intersection(final_columns) + if reserved: + raise ValueError( + "join-map internal ID columns conflict with view columns: " + f"{sorted(reserved)}" + ) + partition_keys = tuple( + str(key) for key in groupby.params.get("partition_by", ()) + ) + join_params: dict[str, object] = { + "instruction": instruction_rewriter.groupby_to_join( + str(groupby.params["instruction"]), + input_cols=tuple( + str(column) for column in groupby.params["input_cols"] + ), + ), + "how": "outer", + "id_columns": ( + JOIN_MAP_LEFT_ID_COLUMN, + JOIN_MAP_RIGHT_ID_COLUMN, + ), + } + membership = groupby.params.get("membership") + if membership == "exclusive": + join_params["k"] = 1 + elif membership == "overlapping": + raise NotImplementedError( + "overlapping sem_groupby membership is not implemented" + ) + elif membership is not None: + raise ValueError(f"Unsupported sem_groupby membership: {membership!r}") + if partition_keys: + join_params["on"] = partition_keys + joined = QueryExpr( + op="sem_join", + inputs=(changed_aggregate, current_view), + params=join_params, + ) + joined_ref = QueryExpr( + op="materialized_view", + params={ + "name": JOIN_MAP_BINDING_NAME, + "columns": output_columns(joined), + }, + ) + + matched = self._filter_join_map_rows( + joined_ref, + left_present=True, + right_present=True, + ) + left_only = self._filter_join_map_rows( + joined_ref, + left_present=True, + right_present=False, + ) + right_only = self._filter_join_map_rows( + joined_ref, + left_present=False, + right_present=True, + ) + + matched_changed = self._project_join_map_side( + matched, + side="left", + final_columns=final_columns, + include_target_id=True, + ) + matched_current = QueryExpr( + op="drop_duplicates", + inputs=( + self._project_join_map_side( + matched, + side="right", + final_columns=final_columns, + include_target_id=True, + ), + ), + params={"subset": (JOIN_MAP_RIGHT_ID_COLUMN,)}, + ) + state_rows = QueryExpr( + op="concat", + inputs=(matched_changed, matched_current), + ) + grouped_states = QueryExpr( + op="group_by", + inputs=(state_rows,), + params={ + "keys": ( + JOIN_MAP_RIGHT_ID_COLUMN, + *partition_keys, + ) + }, + ) + updates = self._reaggregate_join_map_states( + grouped_states, + aggregate=aggregate, + final_columns=final_columns, + instruction_rewriter=instruction_rewriter, + ) + new_groups = self._project_join_map_side( + left_only, + side="left", + final_columns=final_columns, + ) + untouched = self._project_join_map_side( + right_only, + side="right", + final_columns=final_columns, + ) + kept_and_updated = QueryExpr( + op="union_by_name", + inputs=(untouched, updates), + params={"allow_missing_columns": False}, + ) + body = QueryExpr( + op="union_by_name", + inputs=(kept_and_updated, new_groups), + params={"allow_missing_columns": False}, + ) + return QueryExpr( + op="let", + inputs=(joined, body), + params={"name": JOIN_MAP_BINDING_NAME}, + ) + + def _filter_join_map_rows( + self, + joined: QueryExpr, + *, + left_present: bool, + right_present: bool, + ) -> QueryExpr: + """Select one matched or unmatched side of an outer semantic join.""" + + left_id = ColumnExpr(JOIN_MAP_LEFT_ID_COLUMN) + right_id = ColumnExpr(JOIN_MAP_RIGHT_ID_COLUMN) + predicate = ( + left_id.is_not_null() if left_present else left_id.is_null() + ) & ( + right_id.is_not_null() if right_present else right_id.is_null() + ) + return QueryExpr( + op="filter", + inputs=(joined,), + params={"predicate": predicate.to_param()}, + ) + + def _project_join_map_side( + self, + source: QueryExpr, + *, + side: str, + final_columns: Sequence[str], + include_target_id: bool = False, + ) -> QueryExpr: + """Project one outer-join side back to the materialized view schema.""" + + assignments = { + str(column): ColumnExpr(str(column), qualifier=side).to_param() + for column in final_columns + } + columns: tuple[str, ...] = tuple(str(column) for column in final_columns) + if include_target_id: + assignments[JOIN_MAP_RIGHT_ID_COLUMN] = ColumnExpr( + JOIN_MAP_RIGHT_ID_COLUMN + ).to_param() + columns = (JOIN_MAP_RIGHT_ID_COLUMN, *columns) + assigned = QueryExpr( + op="assign", + inputs=(source,), + params={"assignments": assignments}, + ) + return QueryExpr( + op="select", + inputs=(assigned,), + params={"columns": columns}, + ) + + def _reaggregate_join_map_states( + self, + grouped_states: QueryExpr, + *, + aggregate: QueryExpr, + final_columns: Sequence[str], + instruction_rewriter: DifferentialInstructionRewriter, + ) -> QueryExpr: + """Merge current and changed aggregate state once for each selected target.""" + + if aggregate.op == "agg": + merged = self._build_aggregate_remerge( + grouped=grouped_states, + aggregate=aggregate, + state_cols=final_columns, + instruction_rewriter=instruction_rewriter, + ) + return self._select_after_array_remerge_flatten( + merged, + aggregate, + final_columns, + ) + if aggregate.op != "sem_agg": + raise TypeError(f"Unsupported join-map aggregate: {aggregate.op}") + + reaggregate_params = dict(aggregate.params) + reaggregate_params["input_cols"] = None + reaggregate_params["instruction"] = instruction_rewriter.state_reaggregation( + str(aggregate.params["instruction"]), + state_cols=final_columns, + raw_input_cols=tuple( + str(column) for column in (aggregate.params.get("input_cols") or ()) + ), + ) + reaggregated = QueryExpr( + op="sem_agg", + inputs=(grouped_states,), + params=reaggregate_params, + ) + return QueryExpr( + op="select", + inputs=(reaggregated,), + params={"columns": tuple(final_columns)}, + ) + def _differentiate_sem_groupby_agg_view( self, query: QueryExpr, @@ -893,7 +1118,7 @@ def _differentiate_sem_groupby_agg_view( instruction_rewriter=instruction_rewriter, ) - selected_rule = self._grouped_agg_rule_for(groupby) + selected_rule = self._grouped_agg_rule if selected_rule == "rule-re-group": return self._build_grouped_sem_agg_re_group_candidate( aggregate=aggregate, @@ -1121,11 +1346,6 @@ def _build_sem_groupby_agg_join_map_candidate( return None groupby = aggregate.inputs[0] - if groupby.op == "sem_groupby" and groupby.params.get("partition_by"): - raise NotImplementedError( - "sem_groupby partition_by is not supported by rule-join-map; " - "the semantic join cannot yet enforce deterministic partition keys." - ) changed_groupby = QueryExpr( op="sem_groupby", inputs=(changed_group_input,), @@ -1136,6 +1356,15 @@ def _build_sem_groupby_agg_join_map_candidate( inputs=(changed_groupby,), params=aggregate.params, ) + if groupby.op == "sem_groupby": + return self._build_semantic_join_map( + aggregate=aggregate, + groupby=groupby, + changed_aggregate=changed_aggregate, + current_view=current_view, + final_columns=final_columns, + instruction_rewriter=instruction_rewriter, + ) preserved_columns = self._grouped_preserved_columns(groupby) if groupby.op == "group_by": joined = QueryExpr( diff --git a/src/agent_memory/policy/relation.py b/src/agent_memory/policy/relation.py index a0c9fcb..c1e5147 100644 --- a/src/agent_memory/policy/relation.py +++ b/src/agent_memory/policy/relation.py @@ -41,6 +41,7 @@ ColumnInput = Sequence[str] | None ColumnOutput = Sequence[str] | Mapping[str, str] JoinOn = str | Sequence[str] | object +SEM_GROUPBY_MEMBERSHIPS = ("exclusive", "overlapping") def _normalize_input_cols(input_cols: ColumnInput) -> tuple[str, ...] | None: @@ -75,6 +76,19 @@ def _normalize_group_labels( return tuple(ColumnSpec(name=name, description=desc) for name, desc in labels.items()) +def _normalize_sem_groupby_membership(membership: str | None) -> str | None: + """Normalize an optional semantic-group membership contract.""" + + if membership is None: + return None + if membership not in SEM_GROUPBY_MEMBERSHIPS: + raise ValueError( + "sem_groupby membership must be one of: " + + ", ".join(SEM_GROUPBY_MEMBERSHIPS) + ) + return membership + + def _require_relation(value: Any, *, argument: str) -> "Relation": """Validate inputs to binary relation operators.""" @@ -120,6 +134,20 @@ def _normalize_group_keys(keys: str | Sequence[str]) -> tuple[str, ...]: return normalized +def _normalize_sem_join_keys(keys: str | Sequence[str]) -> tuple[str, ...]: + """Normalize exact semantic-join key columns.""" + + if isinstance(keys, str): + normalized = (keys,) + else: + normalized = tuple(str(key) for key in keys) + if not normalized: + raise ValueError("sem_join on requires at least one key column") + if any(not key for key in normalized): + raise ValueError("sem_join on key columns cannot be empty") + return normalized + + def _normalize_drop_duplicates_subset( subset: str | Sequence[str], ) -> tuple[str, ...]: @@ -578,10 +606,12 @@ def sem_groupby( partition_by: str | Sequence[str] | None = None, labels: Mapping[str, str] | None = None, label_col: str = "_label", + membership: str | None = None, ) -> "GroupedRelation": """Create a grouped semantic relation expression.""" normalized_partition_by = _normalize_partition_by(partition_by) + normalized_membership = _normalize_sem_groupby_membership(membership) params: dict[str, Any] = { "input_cols": _normalize_input_cols(input_cols), "instruction": instruction, @@ -590,6 +620,8 @@ def sem_groupby( } if normalized_partition_by is not None: params["partition_by"] = normalized_partition_by + if normalized_membership is not None: + params["membership"] = normalized_membership expr = QueryExpr( op="sem_groupby", inputs=(self.expr,), @@ -603,16 +635,23 @@ def sem_join( *, instruction: str, how: str = "inner", + on: str | Sequence[str] | None = None, + k: int | None = None, ) -> "Relation": """Add a semantic join expression.""" other = _require_relation(other, argument="other") - return self._derive( - "sem_join", - inputs=(self.expr, other.expr), - instruction=instruction, - how=how, - ) + params: dict[str, Any] = { + "instruction": instruction, + "how": how, + } + if on is not None: + params["on"] = _normalize_sem_join_keys(on) + if k is not None: + if not isinstance(k, int) or isinstance(k, bool) or k < 1: + raise ValueError("sem_join k must be a positive integer") + params["k"] = k + return self._derive("sem_join", inputs=(self.expr, other.expr), **params) def sem_topk( self, diff --git a/src/agent_memory/policy/schema.py b/src/agent_memory/policy/schema.py index f57081a..ea7f2e7 100644 --- a/src/agent_memory/policy/schema.py +++ b/src/agent_memory/policy/schema.py @@ -46,6 +46,8 @@ def output_columns( return source_columns + metadata_columns if query.op in {"alias", "group_by"}: return output_columns(query.inputs[0], window_source_columns=window_source_columns) + if query.op == "let": + return output_columns(query.inputs[1], window_source_columns=window_source_columns) if query.op == "over": return output_columns(query.inputs[0], window_source_columns=window_source_columns) if query.op == "array_agg": @@ -287,6 +289,8 @@ def semantic_join_columns( columns.append(f"{column}:left" if column in overlapping else column) for column in right_columns: columns.append(f"{column}:right" if column in overlapping else column) + for column in query.params.get("id_columns", ()): + columns.append(str(column)) return tuple(columns) diff --git a/tests/test_differential_runtime.py b/tests/test_differential_runtime.py index d7c09d6..ea2eed1 100644 --- a/tests/test_differential_runtime.py +++ b/tests/test_differential_runtime.py @@ -219,18 +219,6 @@ class GroupedMemory(am.Memory): assert join_map.grouped_agg_rule == "rule-join-map" -def test_claude_policy_prefer_join_map_matches_strict_join_map_plan() -> None: - strict = planner.PolicyDifferentiator( - rules=DifferentialRules(grouped_agg_rule="rule-join-map") - ).differentiate(am.ClaudeMemory.spec()) - preferred = planner.PolicyDifferentiator( - rules=DifferentialRules(grouped_agg_rule="prefer-join-map") - ).differentiate(am.ClaudeMemory.spec()) - - assert preferred.grouped_agg_rule == "prefer-join-map" - assert preferred.fingerprint == strict.fingerprint - - def test_policy_compiler_rejects_process_window_external_relation_capture() -> None: source_log = am.Log({"message": "Message body."}) external = source_log.select(["message"]) diff --git a/tests/test_interface_scaffold.py b/tests/test_interface_scaffold.py index a381003..010297b 100644 --- a/tests/test_interface_scaffold.py +++ b/tests/test_interface_scaffold.py @@ -36,6 +36,7 @@ execute_filter, execute_flatten, execute_join, + execute_let, execute_min, execute_subtract, execute_unnest, @@ -562,6 +563,42 @@ def test_grouped_relation_sem_agg_returns_normal_relation() -> None: assert aggregated.expr.inputs[0].op == "sem_groupby" +def test_sem_groupby_membership_contract_is_optional_and_validated() -> None: + log = am.Log({"topic_name": "Topic name."}) + unspecified = log.sem_groupby( + input_cols=["topic_name"], + instruction="Group matching topics.", + ) + exclusive = log.sem_groupby( + input_cols=["topic_name"], + instruction="Group matching topics.", + membership="exclusive", + ) + overlapping = log.sem_groupby( + input_cols=["topic_name"], + instruction="Group matching topics.", + membership="overlapping", + ) + + assert "membership" not in unspecified.expr.params + assert exclusive.expr.params["membership"] == "exclusive" + assert overlapping.expr.params["membership"] == "overlapping" + with pytest.raises(ValueError, match="membership must be one of"): + log.sem_groupby( + input_cols=["topic_name"], + instruction="Group matching topics.", + membership="unknown", + ) + with pytest.raises( + NotImplementedError, + match="overlapping sem_groupby membership", + ): + LotusAdapter().execute( + overlapping.expr, + {"log": pd.DataFrame({"topic_name": ["a", "b"]})}, + ) + + def test_grouped_relation_agg_accepts_public_aggregate_specs() -> None: log = am.Log( { @@ -1157,6 +1194,40 @@ def test_sem_join_query_expr_keeps_only_logical_params() -> None: ) +def test_sem_join_accepts_exact_keys_and_per_left_result_limit() -> None: + log = am.Log({"tenant_id": "Exact tenant key.", "message": "Raw message."}) + + joined = log.sem_join( + log, + instruction="{message:left} and {message:right} describe the same entity.", + how="outer", + on="tenant_id", + k=3, + ) + + assert joined.expr.params == { + "instruction": "{message:left} and {message:right} describe the same entity.", + "how": "outer", + "on": ("tenant_id",), + "k": 3, + } + + +@pytest.mark.parametrize("k", [0, -1, True, 1.5]) +def test_sem_join_rejects_invalid_result_limit(k: Any) -> None: + log = am.Log({"message": "Raw message."}) + + with pytest.raises(ValueError, match="positive integer"): + log.sem_join(log, instruction="Rows match.", k=k) + + +def test_sem_join_rejects_empty_exact_key_domain() -> None: + log = am.Log({"message": "Raw message."}) + + with pytest.raises(ValueError, match="at least one key"): + log.sem_join(log, instruction="Rows match.", on=[]) + + def test_sem_filter_query_expr_keeps_only_logical_params() -> None: log = am.Log({"message": "Raw message."}) @@ -1859,7 +1930,7 @@ class TopicEvidenceMemory(am.Memory): assert assignments["earliest"]["kind"] == "least" -def test_join_map_rejects_partitioned_semantic_grouping() -> None: +def test_join_map_composes_partitioned_semantic_grouping_with_exact_keys() -> None: log = am.Log( {"group_id": "Partition.", "name": "Name.", "body": "Evidence."} ) @@ -1876,6 +1947,7 @@ def test_join_map_rejects_partitioned_semantic_grouping() -> None: input_cols=["name"], partition_by="group_id", instruction="Rows with {name} refer to the same entity.", + membership="exclusive", ).agg( am.sem_agg( input_cols=["name", "body"], @@ -1885,20 +1957,59 @@ def test_join_map_rejects_partitioned_semantic_grouping() -> None: am.min(column="body", output_col="first_body"), ).expr - for query in (direct, mixed): - with pytest.raises( - NotImplementedError, - match="partition_by.*rule-join-map", - ): - DifferentialRules(grouped_agg_rule="rule-join-map").differentiate( - query, - source_input=QueryExpr(op="log"), - current_view=QueryExpr( - op="materialized_view", - params={"name": "view", "columns": output_columns(query)}, - ), - is_view_boundary=True, - ) + for query, expected_k in ((direct, None), (mixed, 1)): + differentiated = DifferentialRules( + grouped_agg_rule="rule-join-map" + ).differentiate( + query, + source_input=QueryExpr(op="log"), + current_view=QueryExpr( + op="materialized_view", + params={"name": "view", "columns": output_columns(query)}, + ), + is_view_boundary=True, + ) + sem_join_expr = next( + node for node in _query_nodes(differentiated) if node.op == "sem_join" + ) + assert sem_join_expr.params["on"] == ("group_id",) + if expected_k is None: + assert "k" not in sem_join_expr.params + else: + assert sem_join_expr.params["k"] == expected_k + assert sem_join_expr.params["how"] == "outer" + assert sem_join_expr.params["id_columns"] == ( + "_agent_memory_join_map_left_id", + "_agent_memory_join_map_right_id", + ) + assert output_columns(differentiated) == output_columns(query) + + +def test_join_map_rejects_unimplemented_overlapping_membership() -> None: + log = am.Log({"name": "Name.", "body": "Evidence."}) + query = log.sem_groupby( + input_cols=["name"], + instruction="Rows with {name} refer to the same entity.", + membership="overlapping", + ).sem_agg( + input_cols=["name", "body"], + output_cols={"name": "Canonical name.", "summary": "Summary."}, + instruction="Return {name} and summarize {body} as {summary}.", + ).expr + + with pytest.raises( + NotImplementedError, + match="overlapping sem_groupby membership", + ): + DifferentialRules(grouped_agg_rule="rule-join-map").differentiate( + query, + source_input=QueryExpr(op="log"), + current_view=QueryExpr( + op="materialized_view", + params={"name": "view", "columns": output_columns(query)}, + ), + is_view_boundary=True, + ) def test_differential_rules_reject_array_agg_outside_view_boundary() -> None: @@ -2096,17 +2207,17 @@ def test_join_map_grouped_agg_rule_emits_join_map_shape() -> None: ).differentiate(view) ops = _query_ops(differentiated) - assert differentiated.op == "select" - assert differentiated.params["columns"] == ("name", "description", "type", "body") - sem_map_expr = differentiated.inputs[0] - assert sem_map_expr.op == "sem_map" - sem_join_expr = sem_map_expr.inputs[0] + assert differentiated.op == "let" + sem_join_expr = next(node for node in _query_nodes(differentiated) if node.op == "sem_join") assert sem_join_expr.op == "sem_join" assert sem_join_expr.params["how"] == "outer" + assert "k" not in sem_join_expr.params assert "sem_join" in ops - assert "sem_map" in ops + assert "sem_map" not in ops + assert "sem_agg" in ops + assert ops.count("filter") >= 3 assert "{name:left} and {name:right}" in sem_join_expr.params["instruction"] - assert "{body:left} and {body:right}" in sem_map_expr.params["instruction"] + assert output_columns(differentiated) == ("name", "description", "type", "body") def test_changed_aware_grouped_agg_execution_keeps_old_only_groups() -> None: @@ -2133,6 +2244,134 @@ def test_changed_aware_grouped_agg_execution_keeps_old_only_groups() -> None: assert GROUP_ID_COLUMN not in result.columns +def test_join_map_grouped_agg_execution_keeps_untouched_groups_without_lm() -> None: + view = am.ClaudeMemory.spec().views["topics"] + query = QueryDifferentiator( + rules=DifferentialRules(grouped_agg_rule="rule-join-map"), + ).differentiate(view) + current = pd.DataFrame( + [ + { + "name": "caroline_adoption_goal", + "description": "Caroline's adoption plan.", + "type": "profile", + "body": "Caroline is researching adoption agencies.", + } + ] + ) + changed_log = pd.DataFrame(columns=["role", "message", "timestamp", "session_id"]) + + result = LotusAdapter().execute(query, {"log": changed_log, "topics": current}) + + assert result.to_dict(orient="records") == current.to_dict(orient="records") + assert "_agent_memory_join_map_left_id" not in result.columns + assert "_agent_memory_join_map_right_id" not in result.columns + + +def test_join_map_merges_multiple_changed_groups_into_one_selected_target( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import agent_memory.adapters.lotus.adapter as lotus_adapter_module + + log = am.Log({"name": "Name.", "body": "Body."}) + view = log.sem_groupby( + input_cols=["name"], + instruction="Rows with {name} refer to the same entity.", + ).sem_agg( + input_cols=["name", "body"], + output_cols={"name": "Canonical name.", "summary": "Summary."}, + instruction="Return {name} and summarize {body} as {summary}.", + ) + query = DifferentialRules(grouped_agg_rule="rule-join-map").differentiate( + view.expr, + source_input=QueryExpr(op="log"), + current_view=QueryExpr( + op="materialized_view", + params={"name": "entities", "columns": ("name", "summary")}, + ), + is_view_boundary=True, + ) + changed = pd.DataFrame( + [ + {"name": "Melanie", "summary": "First update."}, + {"name": "Mel", "summary": "Second update."}, + ] + ) + current = pd.DataFrame( + [ + {"name": "Melanie", "summary": "Existing summary."}, + {"name": "Max", "summary": "Untouched summary."}, + ] + ) + calls = {"join": 0, "reaggregate": 0} + + def execute_sem_join_for_test( + expr: QueryExpr, + inputs: dict[str, Any], + execute: Any, + _context: Any, + ) -> pd.DataFrame: + calls["join"] += 1 + left = execute(expr.inputs[0], inputs) + right = execute(expr.inputs[1], inputs) + return assemble_join_frame( + left, + right, + [(0, 0, None), (1, 0, None)], + how="outer", + id_columns=tuple(expr.params["id_columns"]), + ) + + def execute_sem_agg_for_test( + expr: QueryExpr, + inputs: dict[str, Any], + execute: Any, + _context: Any, + ) -> pd.DataFrame: + if expr.inputs[0].op == "sem_groupby": + return changed.copy() + calls["reaggregate"] += 1 + state_rows = execute(expr.inputs[0], inputs) + assert len(state_rows) == 3 + return pd.DataFrame( + [ + { + "_agent_memory_join_map_right_id": 0, + "name": "Melanie", + "summary": "Existing summary. First update. Second update.", + } + ] + ) + + monkeypatch.setattr( + lotus_adapter_module, + "execute_sem_join", + execute_sem_join_for_test, + ) + monkeypatch.setattr( + lotus_adapter_module, + "execute_sem_agg", + execute_sem_agg_for_test, + ) + + result = LotusAdapter().execute( + query, + { + "log": pd.DataFrame(columns=["name", "body"]), + "entities": current, + }, + ) + + assert calls == {"join": 1, "reaggregate": 1} + assert result.to_dict(orient="records") == [ + {"name": "Max", "summary": "Untouched summary."}, + { + "name": "Melanie", + "summary": "Existing summary. First update. Second update.", + }, + ] + + def test_differential_rules_reject_unknown_grouped_agg_strategy() -> None: with pytest.raises(ValueError, match="grouped_agg_rule"): DifferentialRules(grouped_agg_rule="unknown") @@ -2153,23 +2392,16 @@ def test_claude_topics_join_map_candidate_builder_rewrites_placeholders() -> Non ) assert candidate is not None - assert candidate.op == "select" - assert candidate.params["columns"] == ("name", "description", "type", "body") - sem_map_expr = candidate.inputs[0] - assert sem_map_expr.op == "sem_map" - sem_join_expr = sem_map_expr.inputs[0] + assert candidate.op == "let" + sem_join_expr = next(node for node in _query_nodes(candidate) if node.op == "sem_join") assert sem_join_expr.op == "sem_join" assert sem_join_expr.params["how"] == "outer" + assert "k" not in sem_join_expr.params assert "{name:left} and {name:right}" in sem_join_expr.params["instruction"] assert "{description:left} and {description:right}" in sem_join_expr.params["instruction"] assert "{type:left} and {type:right}" in sem_join_expr.params["instruction"] - assert "{body:left} and {body:right}" in sem_map_expr.params["instruction"] - assert tuple(col.name for col in sem_map_expr.params["output_cols"]) == ( - "name", - "description", - "type", - "body", - ) + assert "sem_map" not in _query_ops(candidate) + assert output_columns(candidate) == ("name", "description", "type", "body") def test_differential_query_planner_recomputes_views_from_materialized_dependencies() -> None: @@ -4264,6 +4496,28 @@ def test_relational_execution_ops_follow_dataframe_semantics() -> None: ] +def test_relational_let_evaluates_bound_relation_once() -> None: + bound = QueryExpr(op="source") + body = QueryExpr( + op="materialized_view", + params={"name": "__bound", "columns": ("value",)}, + ) + query = QueryExpr(op="let", inputs=(bound, body), params={"name": "__bound"}) + calls = 0 + + def execute(expr: QueryExpr, inputs: dict[str, Any]) -> pd.DataFrame: + nonlocal calls + if expr.op == "source": + calls += 1 + return pd.DataFrame({"value": [1, 2]}) + return inputs[str(expr.params["name"])] + + result = execute_let(query, {}, execute) + + assert calls == 1 + assert result.to_dict(orient="records") == [{"value": 1}, {"value": 2}] + + def test_union_by_name_aligns_missing_columns_and_deduplicates() -> None: left = QueryExpr(op="materialized_view", params={"name": "left"}) right = QueryExpr(op="materialized_view", params={"name": "right"}) diff --git a/tests/test_sem_topk_join.py b/tests/test_sem_topk_join.py new file mode 100644 index 0000000..5b0b152 --- /dev/null +++ b/tests/test_sem_topk_join.py @@ -0,0 +1,313 @@ +"""Tests for cardinality-bounded semantic joins.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.adapter import LotusAdapter +from agent_memory.adapters.lotus.sem_join import execute_sem_join +from agent_memory.adapters.lotus.sem_topk_join import ( + LISTWISE_JOIN_SYSTEM_PROMPT, + _pairwise_topk, + _parse_listwise_ids, +) +from agent_memory.policy.logical import QueryExpr + + +class _StaticJoinContext: + def __init__(self, config: LotusExecutionConfig) -> None: + self.config = config + self.pair_embedding_provider = None + + def configure(self) -> None: + pass + + +class _BatchListwiseLM: + def __init__(self, outputs: list[str]) -> None: + self.outputs = outputs + self.calls: list[tuple[Any, dict[str, Any]]] = [] + + def __call__(self, messages: Any, **kwargs: Any) -> Any: + self.calls.append((messages, kwargs)) + return SimpleNamespace(outputs=self.outputs) + + +def _query(*, k: int, how: str = "inner", on: tuple[str, ...] = ()) -> QueryExpr: + return QueryExpr( + op="sem_join", + inputs=( + QueryExpr(op="materialized_view", params={"name": "left"}), + QueryExpr(op="materialized_view", params={"name": "right"}), + ), + params={ + "instruction": "The left and right rows describe the same entity.", + "how": how, + "k": k, + **({"on": on} if on else {}), + }, + ) + + +def test_listwise_topk_join_selects_zero_to_k_with_exact_key_domain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + left = pd.DataFrame( + {"tenant": ["a", "b"], "name": ["Melanie", "Poppy"]}, + index=[10, 20], + ) + right = pd.DataFrame( + { + "tenant": ["a", "a", "b"], + "name": ["Mel", "Melanie", "Max"], + }, + index=[100, 101, 200], + ) + lm = _BatchListwiseLM( + [ + '{"selected_ids": ["candidate_1"]}', + '{"selected_ids": []}', + ] + ) + monkeypatch.setattr(lotus.settings, "lm", lm) + + result = execute_sem_join( + _query(k=1, how="outer", on=("tenant",)), + {}, + lambda query, _inputs: left if query.params["name"] == "left" else right, + _StaticJoinContext(LotusExecutionConfig(sem_join_topk_method="listwise")), + ) + + assert len(lm.calls) == 1 + assert len(lm.calls[0][0]) == 2 + assert lm.calls[0][0][0][0]["content"] == LISTWISE_JOIN_SYSTEM_PROMPT + assert "one left row" in LISTWISE_JOIN_SYSTEM_PROMPT + assert "Select at most max_matches candidates" in LISTWISE_JOIN_SYSTEM_PROMPT + assert "Do not select merely related candidates" in LISTWISE_JOIN_SYSTEM_PROMPT + assert "empty selected_ids array when none match" in LISTWISE_JOIN_SYSTEM_PROMPT + assert len(result) == 4 + matched = result[result["name:left"].eq("Melanie") & result["name:right"].notna()] + assert matched["name:right"].tolist() == ["Melanie"] + assert result[result["name:left"].eq("Poppy")]["name:right"].isna().all() + assert set(result.loc[result["name:left"].isna(), "name:right"]) == {"Mel", "Max"} + + +def test_listwise_topk_join_can_select_multiple_matches_up_to_k( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + left = pd.DataFrame({"name": ["Melanie"]}, index=[10]) + right = pd.DataFrame( + {"name": ["Mel", "Melanie", "Max"]}, + index=[100, 101, 102], + ) + lm = _BatchListwiseLM( + ['{"selected_ids": ["candidate_0", "candidate_1"]}'] + ) + monkeypatch.setattr(lotus.settings, "lm", lm) + + result = execute_sem_join( + _query(k=2), + {}, + lambda query, _inputs: left if query.params["name"] == "left" else right, + _StaticJoinContext(LotusExecutionConfig(sem_join_topk_method="listwise")), + ) + + assert result["name:right"].tolist() == ["Mel", "Melanie"] + + +def test_pairwise_topk_join_reuses_lotus_ranking_after_predicate_verification( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import agent_memory.adapters.lotus.sem_join as sem_join_module + + left = pd.DataFrame({"name": ["Melanie {not_a_column}"]}, index=[10]) + right = pd.DataFrame({"name": ["Mel", "Melanie"]}, index=[100, 101]) + captured: dict[str, Any] = {} + + def verify(candidates: pd.DataFrame, **_kwargs: Any) -> list[tuple[Any, Any, None]]: + return [ + (row["_left_pair_id"], row["_right_pair_id"], None) + for _, row in candidates.iterrows() + ] + + def sem_topk(frame: pd.DataFrame, instruction: str, **kwargs: Any) -> pd.DataFrame: + captured["frame"] = frame.copy() + captured["instruction"] = instruction + captured.update(kwargs) + return frame.iloc[[1]] + + monkeypatch.setattr(sem_join_module, "verify_semantic_join_candidates", verify) + monkeypatch.setattr(pd.DataFrame, "sem_topk", sem_topk, raising=False) + + result = execute_sem_join( + _query(k=1), + {}, + lambda query, _inputs: left if query.params["name"] == "left" else right, + _StaticJoinContext( + LotusExecutionConfig(sem_join_topk_method="pairwise-quick") + ), + ) + + assert captured["K"] == 1 + assert captured["method"] == "quick" + assert captured["frame"]["left"].nunique() == 1 + assert "{not_a_column}" in captured["frame"]["left"].iloc[0] + assert captured["frame"]["right"].tolist() == ["name: Mel", "name: Melanie"] + assert "{left}" in captured["instruction"] + assert "{right}" in captured["instruction"] + assert "{not_a_column}" not in captured["instruction"] + assert result["name:right"].tolist() == ["Melanie"] + + +def test_pairwise_topk_join_skips_oracle_when_search_filter_selects_no_pairs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import agent_memory.adapters.lotus.sem_join as sem_join_module + + def unexpected_oracle(*_args: Any, **_kwargs: Any) -> Any: + raise AssertionError("empty candidate sets must not call the oracle") + + monkeypatch.setattr( + sem_join_module, + "verify_semantic_join_candidates", + unexpected_oracle, + ) + candidates = pd.DataFrame( + columns=[ + "_left_pair_id", + "_right_pair_id", + "_left_pair_text", + "_right_pair_text", + ] + ) + + assert ( + _pairwise_topk( + candidates, + instruction="Rows match.", + left_label="left", + right_label="right", + k=1, + method="pairwise-quick", + context=_StaticJoinContext(LotusExecutionConfig()), + ) + == [] + ) + + +def test_exact_key_semantic_join_without_k_restricts_oracle_pairs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import agent_memory.adapters.lotus.sem_join as sem_join_module + + left = pd.DataFrame( + {"tenant": ["a", "b"], "name": ["Melanie", "Poppy"]}, + index=[10, 20], + ) + right = pd.DataFrame( + {"tenant": ["a", "b"], "name": ["Mel", "Max"]}, + index=[100, 200], + ) + captured: dict[str, Any] = {} + + def verify(candidates: pd.DataFrame, **_kwargs: Any) -> list[tuple[Any, Any, None]]: + captured["pairs"] = list( + zip( + candidates["_left_pair_id"], + candidates["_right_pair_id"], + strict=True, + ) + ) + return [] + + monkeypatch.setattr(sem_join_module, "verify_semantic_join_candidates", verify) + query = _query(k=1, on=("tenant",)) + query = QueryExpr( + op=query.op, + inputs=query.inputs, + params={key: value for key, value in query.params.items() if key != "k"}, + ) + + execute_sem_join( + query, + {}, + lambda expr, _inputs: left if expr.params["name"] == "left" else right, + _StaticJoinContext(LotusExecutionConfig()), + ) + + assert captured["pairs"] == [(10, 100), (20, 200)] + + +def test_pairwise_topk_method_changes_physical_identity_without_changing_default() -> None: + assert LotusAdapter().maintenance_execution_fingerprint == "" + assert ( + LotusAdapter( + config=LotusExecutionConfig(sem_join_topk_method="pairwise-quick") + ).maintenance_execution_fingerprint + != "" + ) + + +def test_listwise_topk_join_parser_allows_no_match_and_rejects_overflow() -> None: + assert _parse_listwise_ids( + '{"selected_ids": []}', + valid_ids={"candidate_0"}, + k=1, + ) == () + + with pytest.raises(ValueError, match="more than k=1"): + _parse_listwise_ids( + '{"selected_ids": ["candidate_0", "candidate_1"]}', + valid_ids={"candidate_0", "candidate_1"}, + k=1, + ) + + +def test_semantic_join_internal_ids_preserve_matched_and_unmatched_provenance( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import agent_memory.adapters.lotus.sem_join as sem_join_module + + left = pd.DataFrame({"name": ["Melanie", "Poppy"]}, index=[10, 20]) + right = pd.DataFrame({"name": ["Mel", "Max"]}, index=[100, 200]) + + def verify(*_args: Any, **_kwargs: Any) -> list[tuple[int, int, None]]: + return [(10, 100, None)] + + monkeypatch.setattr(sem_join_module, "evaluate_semantic_join", verify) + query = QueryExpr( + op="sem_join", + inputs=( + QueryExpr(op="materialized_view", params={"name": "left"}), + QueryExpr(op="materialized_view", params={"name": "right"}), + ), + params={ + "instruction": "Rows match.", + "how": "outer", + "id_columns": ("_left_id", "_right_id"), + }, + ) + + result = execute_sem_join( + query, + {}, + lambda expr, _inputs: left if expr.params["name"] == "left" else right, + _StaticJoinContext(LotusExecutionConfig()), + ) + + ids = result[["_left_id", "_right_id"]].astype(object) + ids = ids.where(ids.notna(), None) + assert ids.to_dict(orient="records") == [ + {"_left_id": 10, "_right_id": 100}, + {"_left_id": 20, "_right_id": None}, + {"_left_id": None, "_right_id": 200}, + ] diff --git a/tests/test_semantic_pair_execution.py b/tests/test_semantic_pair_execution.py index 27b7bb5..e7f6384 100644 --- a/tests/test_semantic_pair_execution.py +++ b/tests/test_semantic_pair_execution.py @@ -1041,6 +1041,39 @@ def test_sem_join_profile_rejects_lotus_cascade_before_input_execution( ) +def test_exact_key_sem_join_rejects_lotus_cascade_before_input_execution() -> None: + query = QueryExpr( + op="sem_join", + inputs=( + QueryExpr(op="materialized_view", params={"name": "left"}), + QueryExpr(op="materialized_view", params={"name": "right"}), + ), + params={ + "instruction": "Rows match.", + "how": "inner", + "on": ("tenant_id",), + }, + ) + context = FakeContext( + LotusExecutionConfig( + sem_join_cascade_args={"sampling_percentage": 0.1}, + ) + ) + + with pytest.raises(ValueError, match="exact-key sem_join cannot be combined"): + execute_sem_join( + query, + { + "left": pd.DataFrame({"tenant_id": ["a"], "topic": ["alpha"]}), + "right": pd.DataFrame({"tenant_id": ["a"], "topic": ["alpha"]}), + }, + lambda _expression, _inputs: (_ for _ in ()).throw( + AssertionError("input execution must not run before validation") + ), + context, + ) + + def test_sem_groupby_search_filter_zero_candidates_skips_oracle( monkeypatch: pytest.MonkeyPatch, ) -> None: From ce4a6b5993f613e6fd8eaf72b692ad840d54d3b1 Mon Sep 17 00:00:00 2001 From: Von Date: Mon, 31 Aug 2026 00:01:46 +0800 Subject: [PATCH 2/5] [feat-benchmark] Expose Zep semantic join top-k methods - configure Zep grouping with explicit exclusive membership - expose and validate semantic join top-k methods across benchmark CLIs - propagate the selected physical method through driver configuration and provenance - document and test the supported join-map argument combinations Signed-off-by: Von --- .../evaluation/agent_memory_drivers.py | 42 +++++--- src/agent_memory/evaluation/run.py | 30 +++++- src/agent_memory/memories/zep/policy.py | 5 + .../memories/zep/policy_extended.py | 3 + tests/test_agent_memory_benchmark_drivers.py | 97 +++++++++++++------ tests/test_benchmark_cli.py | 48 ++++++++- tests/test_benchmark_harness.py | 3 - tests/test_zep_memory_policy.py | 51 ++++------ tools/evaluation/README.md | 92 +++++++++++++++++- tools/evaluation/locomo.py | 20 ++++ tools/evaluation/longmemeval.py | 20 ++++ tools/evaluation/memory_agent_bench.py | 20 ++++ 12 files changed, 344 insertions(+), 87 deletions(-) diff --git a/src/agent_memory/evaluation/agent_memory_drivers.py b/src/agent_memory/evaluation/agent_memory_drivers.py index bfa2367..0133c4e 100644 --- a/src/agent_memory/evaluation/agent_memory_drivers.py +++ b/src/agent_memory/evaluation/agent_memory_drivers.py @@ -13,15 +13,12 @@ import pandas as pd from agent_memory.adapters.lotus.pair_execution import ( - PAIR_LEFT_ID_COLUMN, - PAIR_LEFT_TEXT_COLUMN, - PAIR_RIGHT_ID_COLUMN, - PAIR_RIGHT_TEXT_COLUMN, SEMANTIC_PAIR_EXECUTION_MODES, SemanticPairExecutionProfile, SemanticPairSite, semantic_pair_site_contract, semantic_pair_site_id, + semantic_pair_site_physical_contract, ) from agent_memory.evaluation.claude_memory.bindings import event_to_claude_log_row from agent_memory.evaluation.embedding_trace import TracingEmbeddingProvider @@ -68,7 +65,7 @@ def build_operator_semantic_pair_profiles( if top_k is not None or min_similarity is not None: raise ValueError("oracle-only does not accept semantic pair bounds") return {} - supported = {"sem_join", "sem_groupby"} + supported = {"sem_filter", "sem_join", "sem_groupby"} unknown = sorted(set(operators) - supported) if unknown: raise ValueError(f"unsupported semantic pair operators: {unknown}") @@ -76,10 +73,9 @@ def build_operator_semantic_pair_profiles( profiles: dict[str, SemanticPairExecutionProfile] = {} sites = inventory_operator_semantic_pair_sites(policy, operators=operators) for site in sites.values(): - direction = "left-to-right" if site.operator == "sem_join" else "symmetric" profile = _operator_semantic_pair_profile( mode=mode, - direction=direction, + site=site, embedding=embedding, embedding_device=embedding_device, top_k=top_k, @@ -102,7 +98,7 @@ def inventory_operator_semantic_pair_sites( ) -> dict[str, SemanticPairSite]: """Group differential query copies by their semantic predicate contract.""" - supported = {"sem_join", "sem_groupby"} + supported = {"sem_filter", "sem_join", "sem_groupby"} unknown = sorted(set(operators) - supported) if unknown: raise ValueError(f"unsupported semantic pair operators: {unknown}") @@ -129,21 +125,26 @@ def inventory_operator_semantic_pair_sites( continue site_id = semantic_pair_site_id(query) contract = semantic_pair_site_contract(query) + physical_contract = semantic_pair_site_physical_contract(query) group = grouped.setdefault( site_id, { "operator": query.op, "contract": contract, + "physical_contract": physical_contract, "query_digests": set(), }, ) if group["contract"] != contract: raise RuntimeError("semantic pair site digest collision") + if group["physical_contract"] != physical_contract: + raise RuntimeError("semantic pair site physical contract mismatch") group["query_digests"].add(query_digest(query)) result: dict[str, SemanticPairSite] = {} for site_id, group in sorted(grouped.items()): contract = group["contract"] + physical_contract = group["physical_contract"] predicate_sha256 = site_id.split(":", 1)[1] result[site_id] = SemanticPairSite( site_id=site_id, @@ -152,6 +153,11 @@ def inventory_operator_semantic_pair_sites( instruction=str(contract["instruction"]), semantic_columns=tuple(contract["semantic_columns"]), partition_by=tuple(contract["partition_by"]), + direction=str(physical_contract["direction"]), + left_id_columns=tuple(physical_contract["left_id_columns"]), + right_id_columns=tuple(physical_contract["right_id_columns"]), + left_text_columns=tuple(physical_contract["left_text_columns"]), + right_text_columns=tuple(physical_contract["right_text_columns"]), query_digests=tuple(sorted(group["query_digests"])), ) return result @@ -179,10 +185,9 @@ def build_site_semantic_pair_profiles( if binding.mode == "oracle-only": continue site = sites[binding.site_id] - direction = "left-to-right" if site.operator == "sem_join" else "symmetric" profile = _operator_semantic_pair_profile( mode=binding.mode, - direction=direction, + site=site, embedding=embedding, embedding_device=embedding_device, top_k=binding.top_k, @@ -197,7 +202,7 @@ def build_site_semantic_pair_profiles( def _operator_semantic_pair_profile( *, mode: str, - direction: str, + site: SemanticPairSite, embedding: EmbeddingSpec, embedding_device: str, top_k: int | None, @@ -205,11 +210,11 @@ def _operator_semantic_pair_profile( ) -> SemanticPairExecutionProfile: return SemanticPairExecutionProfile( mode=mode, - direction=direction, - 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,), + direction=site.direction, + left_id_columns=site.left_id_columns, + right_id_columns=site.right_id_columns, + left_text_columns=site.left_text_columns, + right_text_columns=site.right_text_columns, embedding=embedding, embedding_device=embedding_device, top_k=top_k, @@ -735,6 +740,7 @@ def __init__( base_namespace: str, model_id: str = "deepseek/deepseek-v4-flash", grouped_agg_rule: str = "rule-re-group", + sem_join_topk_method: str = "listwise", sem_groupby_pair_batch_size: int | None = None, sem_groupby_pair_batch_retries: int = 0, semantic_pair_profiles: dict[str, SemanticPairExecutionProfile] | None = None, @@ -757,6 +763,7 @@ def __init__( self.base_namespace = base_namespace self.model_id = model_id self.grouped_agg_rule = grouped_agg_rule + self.sem_join_topk_method = sem_join_topk_method self.sem_groupby_pair_batch_size = sem_groupby_pair_batch_size self.sem_groupby_pair_batch_retries = sem_groupby_pair_batch_retries self.semantic_pair_profiles = dict(semantic_pair_profiles or {}) @@ -793,6 +800,7 @@ def from_environment( base_namespace: str, model_id: str = "deepseek/deepseek-v4-flash", grouped_agg_rule: str = "rule-re-group", + sem_join_topk_method: str = "listwise", sem_groupby_pair_batch_size: int | None = None, sem_groupby_pair_batch_retries: int = 0, semantic_pair_profiles: dict[str, SemanticPairExecutionProfile] | None = None, @@ -840,6 +848,7 @@ def from_environment( base_namespace=base_namespace, model_id=model_id, grouped_agg_rule=grouped_agg_rule, + sem_join_topk_method=sem_join_topk_method, sem_groupby_pair_batch_size=sem_groupby_pair_batch_size, sem_groupby_pair_batch_retries=sem_groupby_pair_batch_retries, semantic_pair_profiles=semantic_pair_profiles, @@ -908,6 +917,7 @@ def __call__( }, lm_enable_cache=self.lotus_cache_mode == "memory", structured_max_tokens=BENCHMARK_STRUCTURED_MAX_TOKENS, + sem_join_topk_method=self.sem_join_topk_method, sem_groupby_pair_batch_size=self.sem_groupby_pair_batch_size, sem_groupby_pair_batch_retries=self.sem_groupby_pair_batch_retries, semantic_pair_profiles=self.semantic_pair_profiles, diff --git a/src/agent_memory/evaluation/run.py b/src/agent_memory/evaluation/run.py index 63ca6fa..358f4d7 100644 --- a/src/agent_memory/evaluation/run.py +++ b/src/agent_memory/evaluation/run.py @@ -27,6 +27,7 @@ from agent_memory.adapters.lotus.context import ( LOTUS_MEMORY_CACHE_ID, LOTUS_MEMORY_CACHE_MAX_SIZE, + SEM_JOIN_TOPK_METHODS, ) from .artifacts import BenchmarkArtifactStore from .bundle import BenchmarkBundle @@ -99,6 +100,7 @@ def run_agent_memory_bundle( base_namespace: str | None = None, grouped_agg_rule: str = "rule-all-group", sem_topk_method: str | None = None, + sem_join_topk_method: str | None = None, sem_groupby_pair_batch_size: int | None = None, sem_groupby_pair_batch_retries: int = 0, semantic_pair_profile: str = "oracle-only", @@ -122,6 +124,14 @@ def run_agent_memory_bundle( raise ValueError("sem_groupby_pair_batch_size must be positive") if sem_groupby_pair_batch_retries < 0: raise ValueError("sem_groupby_pair_batch_retries cannot be negative") + if ( + sem_join_topk_method is not None + and sem_join_topk_method not in SEM_JOIN_TOPK_METHODS + ): + raise ValueError( + "sem_join_topk_method must be one of: " + + ", ".join(SEM_JOIN_TOPK_METHODS) + ) if semantic_pair_profile not in SEMANTIC_PAIR_PROFILES: raise ValueError( "semantic_pair_profile must be one of: " @@ -181,6 +191,15 @@ def run_agent_memory_bundle( sem_topk_method = ( "pairwise-quick" if system_id == "mem0-enhanced" else "pairwise-naive" ) + uses_sem_join_topk = system_id == "zep-memory" and grouped_agg_rule in { + "join-map", + "rule-join-map", + } + if sem_join_topk_method is not None and not uses_sem_join_topk: + raise ValueError( + "sem_join_topk_method requires Zep with a join-map grouped aggregate rule" + ) + resolved_sem_join_topk_method = sem_join_topk_method or "listwise" site_profile_config = ( load_semantic_pair_profile_config(semantic_pair_profile_config) if semantic_pair_profile_config is not None @@ -214,7 +233,7 @@ def run_agent_memory_bundle( am.ZepMemory.spec(), statements=GRAPHITI_NEO4J_STATEMENTS, ) - operators = ("sem_join", "sem_groupby") + operators = ("sem_filter", "sem_join", "sem_groupby") embedding = GRAPHITI_BGE_M3 semantic_pair_profiles, semantic_pair_sites = ( build_site_semantic_pair_profiles( @@ -301,6 +320,11 @@ def run_agent_memory_bundle( part for part in ( semantic_pair_execution_id, + ( + f"sem-join-topk:{resolved_sem_join_topk_method}" + if uses_sem_join_topk + else "" + ), f"lotus-cache:{LOTUS_MEMORY_CACHE_ID}" if lotus_cache_mode == "memory" else "", @@ -346,6 +370,9 @@ def run_agent_memory_bundle( lotus_execution_provenance = { "sem_groupby_pair_batch_size": sem_groupby_pair_batch_size, "sem_groupby_pair_batch_retries": sem_groupby_pair_batch_retries, + "sem_join_topk_method": ( + resolved_sem_join_topk_method if uses_sem_join_topk else None + ), "semantic_pair_profile": semantic_pair_profile, "semantic_pair_top_k": semantic_pair_top_k, "semantic_pair_min_similarity": semantic_pair_min_similarity, @@ -448,6 +475,7 @@ def run_agent_memory_bundle( base_namespace=base_namespace or _namespace(bundle.benchmark_id, output_dir), model_id=memory_provider_model_id, grouped_agg_rule=grouped_agg_rule, + sem_join_topk_method=resolved_sem_join_topk_method, sem_groupby_pair_batch_size=sem_groupby_pair_batch_size, sem_groupby_pair_batch_retries=sem_groupby_pair_batch_retries, semantic_pair_profiles=semantic_pair_profiles, diff --git a/src/agent_memory/memories/zep/policy.py b/src/agent_memory/memories/zep/policy.py index 323dde1..e616a79 100644 --- a/src/agent_memory/memories/zep/policy.py +++ b/src/agent_memory/memories/zep/policy.py @@ -337,6 +337,7 @@ class ZepMemory(Memory): _extracted_entities.sem_groupby( input_cols=["name"], instruction=_ENTITY_GROUP_INSTRUCTION, + membership="exclusive", ) .agg( sem_agg( @@ -492,6 +493,7 @@ class ZepMemory(Memory): input_cols=["relation_type", "fact"], partition_by=["source_entity_id", "target_entity_id"], instruction=_FACT_GROUP_INSTRUCTION, + membership="exclusive", ) .agg( sem_agg( @@ -542,6 +544,9 @@ class ZepMemory(Memory): _earlier_added_facts = _deduplicated_facts.alias("earlier_added") _later_added_facts = _deduplicated_facts.alias("later_added") + # Contradiction candidates are global within the graph. Native Graphiti + # physically blocks this domain with hybrid top-10 retrieval; that choice + # belongs to an explicit execution profile, not this logical view. # Each row is one pair of facts judged to contradict each other. _contradictory_fact_pairs = _earlier_added_facts.join( _later_added_facts, diff --git a/src/agent_memory/memories/zep/policy_extended.py b/src/agent_memory/memories/zep/policy_extended.py index ec3cd37..16a84da 100644 --- a/src/agent_memory/memories/zep/policy_extended.py +++ b/src/agent_memory/memories/zep/policy_extended.py @@ -352,6 +352,7 @@ class ZepMemoryExtended(Memory): _extracted_entities.sem_groupby( input_cols=["name"], instruction=_ENTITY_GROUP_INSTRUCTION, + membership="exclusive", ) .agg( sem_agg( @@ -507,6 +508,7 @@ class ZepMemoryExtended(Memory): input_cols=["relation_type", "fact"], partition_by=["source_entity_id", "target_entity_id"], instruction=_FACT_GROUP_INSTRUCTION, + membership="exclusive", ) .agg( sem_agg( @@ -668,6 +670,7 @@ class ZepMemoryExtended(Memory): .sem_groupby( input_cols=["name", "summary"], instruction=_COMMUNITY_GROUP_INSTRUCTION, + membership="exclusive", ) .agg( sem_agg( diff --git a/tests/test_agent_memory_benchmark_drivers.py b/tests/test_agent_memory_benchmark_drivers.py index ab423de..cd879c2 100644 --- a/tests/test_agent_memory_benchmark_drivers.py +++ b/tests/test_agent_memory_benchmark_drivers.py @@ -313,7 +313,7 @@ def __init__(self, policy, *, adapter, storage=None) -> None: connector=connector, base_namespace="benchmark-run", model_id="model", - grouped_agg_rule="prefer-join-map", + grouped_agg_rule="rule-join-map", sem_groupby_pair_batch_size=12, sem_groupby_pair_batch_retries=2, semantic_pair_profiles={"query": _search_filter_profile()}, @@ -342,7 +342,7 @@ def __init__(self, policy, *, adapter, storage=None) -> None: } assert created[-1]._runtime.policy == ( am.ZepMemory.spec(), - "prefer-join-map", + "rule-join-map", storage.statements, ) assert isinstance(connector.embedding_provider, TracingEmbeddingProvider) @@ -602,7 +602,8 @@ def run(self, bundle) -> None: system_id="zep-memory", output_dir=output_dir, memory_provider_model_id="provider-model", - grouped_agg_rule="prefer-join-map", + grouped_agg_rule="rule-join-map", + sem_join_topk_method="pairwise-quick", sem_groupby_pair_batch_size=12, sem_groupby_pair_batch_retries=2, memory_thinking_enabled=False, @@ -617,7 +618,8 @@ def run(self, bundle) -> None: assert captured["factory"] == { "base_namespace": f"longmemeval-v1-cleaned-s-{namespace_digest}", "model_id": "provider-model", - "grouped_agg_rule": "prefer-join-map", + "grouped_agg_rule": "rule-join-map", + "sem_join_topk_method": "pairwise-quick", "sem_groupby_pair_batch_size": 12, "sem_groupby_pair_batch_retries": 2, "semantic_pair_profiles": {}, @@ -629,13 +631,17 @@ def run(self, bundle) -> None: system_contract = captured["runner"]["system_contract"] assert system_contract.system_id == "zep-memory" assert system_contract.condition_id == "ZEP-SMOKE" - assert system_contract.maintenance_rule == "prefer-join-map" + assert system_contract.maintenance_rule == "rule-join-map" + assert system_contract.maintenance_execution_id == ( + "sem-join-topk:pairwise-quick" + ) assert system_contract.thinking_enabled is False assert captured["runner"]["runtime_provenance"]["runtime"][ "lotus_execution" ] == { "sem_groupby_pair_batch_size": 12, "sem_groupby_pair_batch_retries": 2, + "sem_join_topk_method": "pairwise-quick", "semantic_pair_profile": "oracle-only", "semantic_pair_top_k": None, "semantic_pair_min_similarity": None, @@ -744,7 +750,7 @@ def test_operator_profiles_target_claude_join_and_zep_groupby(mode: str) -> None assert {profile.mode for profile in zep_profiles.values()} == {mode} -def test_zep_differential_groupby_queries_form_two_predicate_sites() -> None: +def test_zep_differential_queries_expose_grouping_and_contradiction_sites() -> None: import agent_memory as am from agent_memory.memories.zep.storage import ( GRAPHITI_BGE_M3, @@ -760,16 +766,32 @@ def test_zep_differential_groupby_queries_form_two_predicate_sites() -> None: ) sites = inventory_operator_semantic_pair_sites( zep, - operators=("sem_groupby",), + operators=("sem_filter", "sem_groupby"), ) - assert len(sites) == 2 - assert sorted(len(site.query_digests) for site in sites.values()) == [3, 3] - entity_site = next(site for site in sites.values() if not site.partition_by) - fact_site = next(site for site in sites.values() if site.partition_by) + assert len(sites) == 3 + assert sorted(len(site.query_digests) for site in sites.values()) == [1, 3, 3] + entity_site = next( + site + for site in sites.values() + if site.operator == "sem_groupby" and not site.partition_by + ) + fact_site = next( + site + for site in sites.values() + if site.operator == "sem_groupby" and site.partition_by + ) + contradiction_site = next( + site for site in sites.values() if site.operator == "sem_filter" + ) assert entity_site.semantic_columns == ("name",) assert fact_site.semantic_columns == ("relation_type", "fact") assert fact_site.partition_by == ("source_entity_id", "target_entity_id") + assert contradiction_site.direction == "right-to-left" + assert contradiction_site.left_text_columns == ("fact:earlier_added",) + assert contradiction_site.right_text_columns == ("fact:later_added",) + assert "fact_id:earlier_added" in contradiction_site.left_id_columns + assert "fact_id:later_added" in contradiction_site.right_id_columns profiles, resolved_sites = build_site_semantic_pair_profiles( zep, @@ -786,14 +808,20 @@ def test_zep_differential_groupby_queries_form_two_predicate_sites() -> None: top_k=10, min_similarity=None, ), + SemanticPairSiteBinding( + site_id=contradiction_site.site_id, + mode="search-filter", + top_k=10, + min_similarity=None, + ), ), - operators=("sem_groupby",), + operators=("sem_filter", "sem_groupby"), embedding=GRAPHITI_BGE_M3, embedding_device="cuda", ) assert resolved_sites == sites - assert len(profiles) == 6 + assert len(profiles) == 7 assert { (profile.top_k, profile.min_similarity) for digest, profile in profiles.items() @@ -804,6 +832,11 @@ def test_zep_differential_groupby_queries_form_two_predicate_sites() -> None: for digest, profile in profiles.items() if digest in fact_site.query_digests } == {(10, None)} + assert { + (profile.direction, profile.top_k, profile.min_similarity) + for digest, profile in profiles.items() + if digest in contradiction_site.query_digests + } == {("right-to-left", 10, None)} def test_site_profiles_reject_stale_site_before_external_setup() -> None: @@ -893,9 +926,9 @@ def test_run_rejects_stale_site_before_provenance_and_environment( @pytest.mark.parametrize( ("grouped_agg_rule", "expected_profile_count", "expected_site_count"), - ( - ("rule-re-group", 6, 2), - ("prefer-join-map", 4, 3), + ( + ("rule-re-group", 7, 3), + ("rule-join-map", 5, 5), ), ) def test_run_resolves_site_config_into_manifest_and_query_profiles( @@ -918,18 +951,21 @@ def test_run_resolves_site_config_into_manifest_and_query_profiles( ) sites = inventory_operator_semantic_pair_sites( policy, - operators=("sem_join", "sem_groupby"), + operators=("sem_filter", "sem_join", "sem_groupby"), ) bindings = [ { "site_id": site.site_id, "mode": "search-filter", - "top_k": 10 if site.partition_by else 15, - "min_similarity": None if site.partition_by else 0.6, + "top_k": 10 if site.partition_by or site.operator == "sem_filter" else 15, + "min_similarity": ( + None if site.partition_by or site.operator == "sem_filter" else 0.6 + ), } for site in sites.values() if grouped_agg_rule == "rule-re-group" or site.operator == "sem_join" + or site.operator == "sem_filter" or site.partition_by ] profile_path = tmp_path / "profiles.json" @@ -1013,6 +1049,7 @@ def run(self, bundle) -> None: for site in sites.values() if grouped_agg_rule == "rule-re-group" or site.operator == "sem_join" + or site.operator == "sem_filter" or site.partition_by for query_digest in site.query_digests } @@ -1262,7 +1299,7 @@ def test_proxy_only_requires_threshold_before_external_setup( ) -def test_strict_zep_join_map_fails_before_storage_side_effects(tmp_path) -> None: +def test_strict_zep_join_map_compiles_before_storage_side_effects(tmp_path) -> None: class CountingConnector: def __init__(self) -> None: self.prepare_count = 0 @@ -1285,16 +1322,22 @@ def close(self) -> None: connector=connector, base_namespace="strict-join-map", grouped_agg_rule="rule-join-map", + sem_join_topk_method="pairwise-quick", neo4j_image="neo4j:5.26.2", neo4j_image_digest="sha256:image", ) - with pytest.raises(NotImplementedError, match="partition_by.*rule-join-map"): - factory( - "case-1", - tmp_path / "attempt-0001", - tmp_path / "trace", - ) + driver = factory( + "case-1", + tmp_path / "attempt-0001", + tmp_path / "trace", + ) - assert connector.prepare_count == 0 + assert isinstance(driver, ZepMemoryDriver) + assert factory.sem_join_topk_method == "pairwise-quick" + assert ( + driver._memory._runtime._engine.adapter.config.sem_join_topk_method + == "pairwise-quick" + ) + assert connector.prepare_count == 1 assert connector.transaction_count == 0 diff --git a/tests/test_benchmark_cli.py b/tests/test_benchmark_cli.py index 98b584b..f45b4db 100644 --- a/tests/test_benchmark_cli.py +++ b/tests/test_benchmark_cli.py @@ -9,6 +9,52 @@ from tools.evaluation import locomo, longmemeval, memory_agent_bench +@pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) +@pytest.mark.parametrize("method", ("pairwise-quick", "listwise")) +def test_agent_benchmark_clis_select_sem_join_topk_access_path( + module, method, tmp_path +) -> None: + args = module.parse_args( + [ + "run", + "--bundle-dir", + str(tmp_path / "bundle"), + "--output-dir", + str(tmp_path / "output"), + "--system", + "zep-memory", + "--grouped-agg-rule", + "rule-join-map", + "--sem-join-topk-method", + method, + ] + ) + + assert args.sem_join_topk_method == method + + +@pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) +def test_agent_benchmark_clis_reject_sem_join_topk_without_zep_join_map( + module, tmp_path +) -> None: + with pytest.raises(SystemExit): + module.parse_args( + [ + "run", + "--bundle-dir", + str(tmp_path / "bundle"), + "--output-dir", + str(tmp_path / "output"), + "--system", + "zep-memory", + "--grouped-agg-rule", + "rule-re-group", + "--sem-join-topk-method", + "listwise", + ] + ) + + @pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) def test_agent_benchmark_clis_share_sem_groupby_execution_options( module, tmp_path @@ -262,7 +308,7 @@ def test_longmemeval_cli_rejects_explicit_claude_options_for_zep( @pytest.mark.parametrize("system_id", ("claude-memory", "zep-memory")) @pytest.mark.parametrize( "grouped_agg_rule", - ("rule-re-group", "rule-join-map", "prefer-join-map"), + ("rule-re-group", "rule-join-map"), ) def test_longmemeval_cli_accepts_planner_strategies_for_all_agent_policies( tmp_path, system_id, grouped_agg_rule diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py index a2eb093..5e3153b 100644 --- a/tests/test_benchmark_harness.py +++ b/tests/test_benchmark_harness.py @@ -2017,9 +2017,6 @@ def test_maintenance_fingerprint_separates_rules_but_not_retrieval() -> None: retrieval_recipe_id="claude-memory:listwise", ) re_group = replace(join_map, maintenance_rule="rule-re-group") - preferred = replace(join_map, maintenance_rule="prefer-join-map") assert join_map.maintenance_fingerprint == join_map_listwise.maintenance_fingerprint assert join_map.maintenance_fingerprint != re_group.maintenance_fingerprint - assert preferred.maintenance_fingerprint != join_map.maintenance_fingerprint - assert preferred.maintenance_fingerprint != re_group.maintenance_fingerprint diff --git a/tests/test_zep_memory_policy.py b/tests/test_zep_memory_policy.py index ce77956..785175e 100644 --- a/tests/test_zep_memory_policy.py +++ b/tests/test_zep_memory_policy.py @@ -3,7 +3,6 @@ from __future__ import annotations import agent_memory as am -import pytest from agent_memory.policy.aggregates import ( ArrayAggregateSpec, MinAggregateSpec, @@ -194,47 +193,28 @@ def test_zep_retrieval_is_one_two_channel_storage_backed_dag() -> None: ).fingerprint -def test_prefer_join_map_selects_rules_by_grouping_capability() -> None: +def test_strict_join_map_supports_partitioned_zep_facts() -> None: policy = PolicyDifferentiator( - rules=DifferentialRules(grouped_agg_rule="prefer-join-map") + rules=DifferentialRules(grouped_agg_rule="rule-join-map") ).differentiate( ZepMemory.spec(), statements=GRAPHITI_NEO4J_STATEMENTS, ) - grouped_nodes = [ + semantic_joins = [ node - for node in policy.nodes.values() - if node.execution_kind == "semantic_state" - and node.query.op == "agg" - and node.query.inputs[0].op == "sem_groupby" + for plan_node in policy.nodes.values() + if plan_node.maintenance_query is not None + for node in _walk(plan_node.maintenance_query) + if node.op == "sem_join" ] - unpartitioned = next( - node - for node in grouped_nodes - if not node.query.inputs[0].params.get("partition_by") - ) - partitioned = next( - node - for node in grouped_nodes - if node.query.inputs[0].params.get("partition_by") - ) - assert policy.grouped_agg_rule == "prefer-join-map" - assert unpartitioned.maintenance_query is not None - assert "sem_join" in _ops(unpartitioned.maintenance_query) - assert partitioned.maintenance_query is not None - assert "sem_join" not in _ops(partitioned.maintenance_query) - assert "concat" in _ops(partitioned.maintenance_query) - - -def test_strict_join_map_still_rejects_partitioned_zep_facts() -> None: - with pytest.raises(NotImplementedError, match="partition_by.*rule-join-map"): - PolicyDifferentiator( - rules=DifferentialRules(grouped_agg_rule="rule-join-map") - ).differentiate( - ZepMemory.spec(), - statements=GRAPHITI_NEO4J_STATEMENTS, - ) + assert semantic_joins + assert all(join.params["k"] == 1 for join in semantic_joins) + assert any( + join.params.get("on") + == ("source_entity_id", "target_entity_id") + for join in semantic_joins + ) def test_zep_episodes_view_uses_select_not_map() -> None: @@ -302,6 +282,7 @@ def test_zep_entities_view_uses_context_extraction_and_semantic_aggregation() -> assert "Always extract the speaker" in extraction.params["instruction"] assert "When in doubt, do not extract" in extraction.params["instruction"] grouped_agg = next(node for node in _walk(query) if node.op == "agg") + assert grouped_agg.inputs[0].params["membership"] == "exclusive" specs = grouped_agg.params["aggregates"] entity_id = next( specification @@ -411,6 +392,7 @@ def test_zep_facts_view_contains_temporal_self_join_paths() -> None: "source_entity_id", "target_entity_id", ) + assert fact_grouping.params["membership"] == "exclusive" specs = deduplicated_agg.params["aggregates"] assert any(isinstance(specification, SemanticAggregateSpec) for specification in specs) assert any(isinstance(specification, ArrayAggregateSpec) for specification in specs) @@ -516,6 +498,7 @@ def test_zep_communities_view_is_semantic_grouping_not_label_propagation() -> No assert "label_propagation" not in ops groupby = next(node for node in _walk(query) if node.op == "sem_groupby") assert "partition_by" not in groupby.params + assert groupby.params["membership"] == "exclusive" aggregate = next(node for node in _walk(query) if node.op == "agg") aggregate_specs = aggregate.params["aggregates"] community_id = next( diff --git a/tools/evaluation/README.md b/tools/evaluation/README.md index fc2e85e..77b1664 100644 --- a/tools/evaluation/README.md +++ b/tools/evaluation/README.md @@ -14,11 +14,11 @@ System drivers only add events and retrieve context. The runner owns case isolation, resume, answering, grading, trace, and artifacts. All agent-memory runners accept the planner's complete -`--grouped-agg-rule` choice set. In particular, `rule-re-group`, -`rule-join-map`, and `prefer-join-map` have the same meaning for every policy -and benchmark. A strict strategy remains selectable even when a particular -policy cannot compile it; that run fails before insertion instead of being -hidden or silently changed by the CLI. +`--grouped-agg-rule` choice set. In particular, `rule-re-group` and +`rule-join-map` have the same meaning for every policy and benchmark. A +strict strategy remains selectable even when a particular policy cannot +compile it; that run fails before insertion instead of being hidden or silently +changed by the CLI. ## Install @@ -68,6 +68,88 @@ Real runs require `DEEPSEEK_API_KEY`. Zep runs additionally require the `NEO4J_USER`, and `NEO4J_PASSWORD` and should point to a separate empty Neo4j 5.26.2 instance. +### Zep join-map physical execution + +Zep declares exclusive semantic grouping in its policy. With +`rule-join-map`, the compiler compares only changed groups with the current +view and lowers each changed group to a zero-or-one-target semantic join. +Fact groups also carry exact source/target endpoint keys, so semantic matching +never crosses those deterministic partitions. + +The policy does not choose how the bounded semantic join is executed. Select +that physical access path at run time: + +```bash +uv run --extra benchmarks --extra zep \ + python tools/evaluation/locomo.py run \ + --bundle-dir .memory-test/bundles/locomo-s0 \ + --system zep-memory \ + --output-dir .memory-test/runs/locomo-zep-join-map \ + --grouped-agg-rule rule-join-map \ + --sem-join-topk-method listwise \ + --semantic-pair-profile-config /path/to/zep-site-profiles.json \ + --lotus-cache-mode disabled \ + --embedding-device cuda \ + --semantic-trace-snapshot-mode compact +``` + +`--sem-join-topk-method` accepts `listwise`, `pairwise-naive`, +`pairwise-quick`, or `pairwise-heap`. It is currently valid only for Zep when +`--grouped-agg-rule` is `join-map` or `rule-join-map`. The policy still contains +the same `sem_join(k=1)` expression; this flag changes only how the adapter +resolves that join. + +`--semantic-pair-profile-config` binds Search-Filter or Proxy-Only profiles to +stable semantic predicate site IDs. A config has this strict shape: + +```json +{ + "schema_version": 1, + "bindings": [ + { + "site_id": "sem_join:", + "mode": "search-filter", + "top_k": 15, + "min_similarity": 0.6 + }, + { + "site_id": "sem_join:", + "mode": "search-filter", + "top_k": 10, + "min_similarity": null + }, + { + "site_id": "sem_filter:", + "mode": "search-filter", + "top_k": 10, + "min_similarity": null + } + ] +} +``` + +Replace the placeholders with site IDs inventoried from the exact compiled +policy revision; do not copy digests between revisions by hand. Unknown, +duplicate, or drifted sites fail before model and storage initialization. +Site-level config is mutually exclusive with the global +`--semantic-pair-profile`, `--semantic-pair-top-k`, and +`--semantic-pair-min-similarity` options. + +The remaining flags belong to separate physical layers: + +- `--lotus-cache-mode disabled|memory` controls LOTUS's process-local exact + cache. `memory` is an independent experiment condition and is not the same as + DeepSeek provider prompt caching. +- `--embedding-device cpu|cuda` selects where the configured embedding model + runs. It does not choose or change the embedding model. +- `--semantic-trace-snapshot-mode compact|full` controls trace detail. + `compact` keeps pair decisions and accounting without full intermediate + DataFrame snapshots; use `full` only for bounded debugging. + +These physical settings enter run provenance and checkpoint identity. Do not +restore a checkpoint under a different join resolver, site profile, cache mode, +embedding device, or trace contract. + ## LongMemEval v1 Prepare selected complete questions from the pinned cleaned 500-question diff --git a/tools/evaluation/locomo.py b/tools/evaluation/locomo.py index 67d46b0..4b2d74f 100644 --- a/tools/evaluation/locomo.py +++ b/tools/evaluation/locomo.py @@ -21,6 +21,7 @@ AGENT_MEMORY_SYSTEMS, DEFAULT_PROVIDER_MODEL_ID, SEMANTIC_PAIR_PROFILES, + SEM_JOIN_TOPK_METHODS, run_agent_memory_bundle, ) from agent_memory.planner import GROUPED_AGG_RULES # noqa: E402 @@ -73,6 +74,10 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: choices=("pairwise-naive", "pairwise-quick", "pairwise-heap", "listwise"), default=None, ) + run.add_argument( + "--sem-join-topk-method", + choices=SEM_JOIN_TOPK_METHODS, + ) run.add_argument("--sem-groupby-pair-batch-size", type=int) run.add_argument("--sem-groupby-pair-batch-retries", type=int, default=0) run.add_argument( @@ -96,6 +101,20 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ) args = parser.parse_args(raw_args) if args.command == "run": + sem_join_topk_requested = any( + argument == "--sem-join-topk-method" + or argument.startswith("--sem-join-topk-method=") + for argument in raw_args + ) + if sem_join_topk_requested and ( + args.system != "zep-memory" + or args.grouped_agg_rule + not in {"join-map", "rule-join-map"} + ): + parser.error( + "--sem-join-topk-method requires --system zep-memory and a " + "join-map grouped aggregate rule" + ) if args.semantic_pair_profile_config is not None and any( argument == option or argument.startswith(f"{option}=") for option in ( @@ -186,6 +205,7 @@ def main(argv: Sequence[str] | None = None) -> Path: base_namespace=args.namespace, grouped_agg_rule=args.grouped_agg_rule, sem_topk_method=args.sem_topk_method, + sem_join_topk_method=args.sem_join_topk_method, sem_groupby_pair_batch_size=args.sem_groupby_pair_batch_size, sem_groupby_pair_batch_retries=args.sem_groupby_pair_batch_retries, semantic_pair_profile=args.semantic_pair_profile, diff --git a/tools/evaluation/longmemeval.py b/tools/evaluation/longmemeval.py index fb3da72..11fc45b 100644 --- a/tools/evaluation/longmemeval.py +++ b/tools/evaluation/longmemeval.py @@ -25,6 +25,7 @@ from agent_memory.evaluation.run import ( # noqa: E402 AGENT_MEMORY_SYSTEMS, SEMANTIC_PAIR_PROFILES, + SEM_JOIN_TOPK_METHODS, run_agent_memory_bundle, ) from agent_memory.planner import GROUPED_AGG_RULES # noqa: E402 @@ -122,6 +123,10 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: choices=("pairwise-naive", "pairwise-quick", "pairwise-heap", "listwise"), default=None, ) + run.add_argument( + "--sem-join-topk-method", + choices=SEM_JOIN_TOPK_METHODS, + ) run.add_argument("--sem-groupby-pair-batch-size", type=int) run.add_argument("--sem-groupby-pair-batch-retries", type=int, default=0) run.add_argument( @@ -154,6 +159,20 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: if getattr(args, "question_ids", None) is not None: args.question_ids = tuple(args.question_ids) if args.command == "run": + sem_join_topk_requested = any( + argument == "--sem-join-topk-method" + or argument.startswith("--sem-join-topk-method=") + for argument in raw_args + ) + if sem_join_topk_requested and ( + args.system != "zep-memory" + or args.grouped_agg_rule + not in {"join-map", "rule-join-map"} + ): + parser.error( + "--sem-join-topk-method requires --system zep-memory and a " + "join-map grouped aggregate rule" + ) if args.semantic_pair_profile_config is not None and any( argument == option or argument.startswith(f"{option}=") for option in ( @@ -250,6 +269,7 @@ def main(argv: Sequence[str] | None = None) -> Path: base_namespace=args.namespace, grouped_agg_rule=args.grouped_agg_rule, sem_topk_method=args.sem_topk_method, + sem_join_topk_method=args.sem_join_topk_method, sem_groupby_pair_batch_size=args.sem_groupby_pair_batch_size, sem_groupby_pair_batch_retries=args.sem_groupby_pair_batch_retries, semantic_pair_profile=args.semantic_pair_profile, diff --git a/tools/evaluation/memory_agent_bench.py b/tools/evaluation/memory_agent_bench.py index 6693d06..13c7a5a 100644 --- a/tools/evaluation/memory_agent_bench.py +++ b/tools/evaluation/memory_agent_bench.py @@ -26,6 +26,7 @@ from agent_memory.evaluation.run import ( # noqa: E402 AGENT_MEMORY_SYSTEMS, SEMANTIC_PAIR_PROFILES, + SEM_JOIN_TOPK_METHODS, run_agent_memory_bundle, ) from agent_memory.planner import GROUPED_AGG_RULES # noqa: E402 @@ -78,6 +79,10 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: choices=("pairwise-naive", "pairwise-quick", "pairwise-heap", "listwise"), default=None, ) + run.add_argument( + "--sem-join-topk-method", + choices=SEM_JOIN_TOPK_METHODS, + ) run.add_argument("--sem-groupby-pair-batch-size", type=int) run.add_argument("--sem-groupby-pair-batch-retries", type=int, default=0) run.add_argument( @@ -101,6 +106,20 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ) args = parser.parse_args(raw_args) if args.command == "run": + sem_join_topk_requested = any( + argument == "--sem-join-topk-method" + or argument.startswith("--sem-join-topk-method=") + for argument in raw_args + ) + if sem_join_topk_requested and ( + args.system != "zep-memory" + or args.grouped_agg_rule + not in {"join-map", "rule-join-map"} + ): + parser.error( + "--sem-join-topk-method requires --system zep-memory and a " + "join-map grouped aggregate rule" + ) if args.semantic_pair_profile_config is not None and any( argument == option or argument.startswith(f"{option}=") for option in ( @@ -193,6 +212,7 @@ def main(argv: Sequence[str] | None = None) -> Path: base_namespace=args.namespace, grouped_agg_rule=args.grouped_agg_rule, sem_topk_method=args.sem_topk_method, + sem_join_topk_method=args.sem_join_topk_method, sem_groupby_pair_batch_size=args.sem_groupby_pair_batch_size, sem_groupby_pair_batch_retries=args.sem_groupby_pair_batch_retries, semantic_pair_profile=args.semantic_pair_profile, From d4180ef7b93fef2ea84678693da976ec187d99b8 Mon Sep 17 00:00:00 2001 From: Von Date: Mon, 31 Aug 2026 00:02:03 +0800 Subject: [PATCH 3/5] [analysis] Add Mem0 Search-Filter causal audit - compare Native, Oracle, and Search-Filter benchmark evidence - audit omitted Oracle-positive pairs and counterfactual retained memories - decompose retrieval, answer, and grading score changes - reject incomplete event-count evidence instead of assuming a full run Signed-off-by: Von --- tests/test_mem0_search_filter_causal_audit.py | 395 ++++++ .../mem0_search_filter_causal_audit.py | 1227 +++++++++++++++++ 2 files changed, 1622 insertions(+) create mode 100644 tests/test_mem0_search_filter_causal_audit.py create mode 100644 tools/analysis/mem0_search_filter_causal_audit.py diff --git a/tests/test_mem0_search_filter_causal_audit.py b/tests/test_mem0_search_filter_causal_audit.py new file mode 100644 index 0000000..685700f --- /dev/null +++ b/tests/test_mem0_search_filter_causal_audit.py @@ -0,0 +1,395 @@ +"""Tests for the read-only Mem0 Search-Filter causal audit.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import replace +import io +import json +from pathlib import Path +import pickle +import tarfile +from typing import Any + +import pandas as pd +import pytest + +from tools.analysis.mem0_search_filter_causal_audit import ( + AuditError, + DirectoryRunReader, + RunArtifacts, + _combine_maintenance_and_retrieval, + _compare_runs, + _completed_event_count, + analyze_top_k_pairs, + analyze_top_k_pair_parity, + collect_tar_run_readers, + load_maintenance, + load_retrieval, +) +from tools.analysis.semantic_pair_candidates import PairGroup, PairRecord + + +class _FakeScorer: + def __init__(self, scores: Mapping[str, Sequence[float]]) -> None: + self.scores = scores + + @property + def metadata(self) -> Mapping[str, Any]: + return {"kind": "fake"} + + def score(self, group: PairGroup) -> Sequence[float]: + return self.scores[group.group_id] + + +def test_directory_and_tar_condition_loading_match(tmp_path: Path) -> None: + maintenance = tmp_path / "directory-maintenance" + retrieval = tmp_path / "directory-retrieval" + _write_maintenance(maintenance, memories=("memory one", "memory two")) + _write_retrieval(retrieval) + + directory_run = _combine_maintenance_and_retrieval( + load_maintenance(DirectoryRunReader(maintenance), name="maintenance"), + load_retrieval(DirectoryRunReader(retrieval), name="retrieval"), + name="condition", + ) + + archive = tmp_path / "artifacts.tar" + _write_tar( + archive, + { + "runs/condition-maintenance": maintenance, + "runs/condition-retrieval": retrieval, + }, + ) + readers = collect_tar_run_readers( + archive, + { + "condition-maintenance": "runs/condition-maintenance", + "condition-retrieval": "runs/condition-retrieval", + }, + ) + tar_run = _combine_maintenance_and_retrieval( + load_maintenance( + readers["condition-maintenance"], name="maintenance" + ), + load_retrieval(readers["condition-retrieval"], name="retrieval"), + name="condition", + ) + + assert tar_run == directory_run + + +def test_top_k_pair_audit_reports_missed_positive_and_retained_memory() -> None: + group = PairGroup( + group_id="group-1", + operator="sem_filter", + direction="right-to-left", + source="fixture", + case_id="case-1", + session_id="session-1", + event_id="event-1", + query_digest="query-1", + pairs=( + _pair("a", "new-1", "earlier A", "later one", False), + _pair("b", "new-1", "earlier B", "later one", True), + _pair("c", "new-2", "earlier C", "later two", True), + _pair("d", "new-2", "earlier D", "later two", False), + ), + ) + report = analyze_top_k_pairs( + iter((group,)), + _FakeScorer({"group-1": (0.9, 0.8, 0.7, 0.6)}), + top_k=1, + ) + + assert report["pair_count"] == 4 + assert report["selected_pair_count"] == 2 + assert report["selected_oracle_positive_pair_count"] == 1 + assert report["missed_oracle_positive_pair_count"] == 1 + assert report["counterfactual_retained_later_memories"] == ["later one"] + assert report["missed_oracle_positive_pairs"][0][ + "candidate_rank_for_later" + ] == 2 + + +def test_pair_device_parity_compares_candidate_identity_not_only_counts() -> None: + group = PairGroup( + group_id="group-1", + operator="sem_filter", + direction="right-to-left", + source="fixture", + case_id="case-1", + session_id="session-1", + event_id="event-1", + query_digest="query-1", + pairs=( + _pair("a", "new", "earlier A", "later", False), + _pair("b", "new", "earlier B", "later", True), + ), + ) + _, same = analyze_top_k_pair_parity( + iter((group,)), + _FakeScorer({"group-1": (0.9, 0.8)}), + _FakeScorer({"group-1": (0.9, 0.8)}), + top_k=1, + ) + _, changed = analyze_top_k_pair_parity( + iter((group,)), + _FakeScorer({"group-1": (0.9, 0.8)}), + _FakeScorer({"group-1": (0.8, 0.9)}), + top_k=1, + ) + + assert same["selected_pair_set_equal"] is True + assert same["missed_positive_set_equal"] is True + assert changed["selected_pair_set_equal"] is False + assert changed["missed_positive_set_equal"] is False + + +def test_question_decomposition_separates_memory_and_generation_changes() -> None: + baseline = _run( + name="oracle", + memories=("shared",), + retrieval_memories={"case:q1": ("shared",), "case:q2": ("shared",)}, + answers={"case:q1": "old one", "case:q2": "old two"}, + scores={"case:q1": 0.0, "case:q2": 0.0}, + ) + optimized = _run( + name="optimized", + memories=("shared", "retained support"), + retrieval_memories={ + "case:q1": ("shared",), + "case:q2": ("retained support",), + }, + answers={"case:q1": "new one", "case:q2": "new two"}, + scores={"case:q1": 1.0, "case:q2": 1.0}, + ) + + comparison = _compare_runs(baseline, optimized) + + assert comparison["context_answer_decomposition"] == { + "different_context_with_b_only_memory": 1, + "same_context_different_answer": 1, + } + assert comparison["scorers"]["official"]["wins_b"] == 2 + assert comparison["questions"][1]["b_only_retrieved_memories"] == [ + "retained support" + ] + + +def test_question_contract_mismatch_is_rejected() -> None: + baseline = _run( + name="oracle", + memories=("shared",), + retrieval_memories={"case:q1": ("shared",), "case:q2": ("shared",)}, + answers={"case:q1": "one", "case:q2": "two"}, + scores={"case:q1": 0.0, "case:q2": 0.0}, + ) + changed = _run( + name="optimized", + memories=("shared",), + retrieval_memories={"case:q1": ("shared",), "case:q2": ("shared",)}, + answers={"case:q1": "one", "case:q2": "two"}, + scores={"case:q1": 0.0, "case:q2": 0.0}, + second_question="different question", + ) + + from tools.analysis.mem0_search_filter_causal_audit import analyze_conditions + + with pytest.raises(AuditError, match="question contract differs"): + analyze_conditions(baseline, baseline, changed, pair_audit={}) + + +def test_completed_event_count_requires_explicit_summary_evidence() -> None: + run = _run( + name="condition", + memories=("shared",), + retrieval_memories={"case:q1": ("shared",), "case:q2": ("shared",)}, + answers={"case:q1": "one", "case:q2": "two"}, + scores={"case:q1": 1.0, "case:q2": 1.0}, + ) + + assert _completed_event_count(run) == 2 + assert _completed_event_count(replace(run, summary={"event_count": 3})) == 3 + with pytest.raises(AuditError, match="missing completed event count: condition"): + _completed_event_count(replace(run, summary={})) + + +def _pair( + left_id: str, + right_id: str, + left: str, + right: str, + baseline_match: bool, +) -> PairRecord: + return PairRecord( + pair_id=f"{left_id}:{right_id}", + left_id=left_id, + right_id=right_id, + left=left, + right=right, + baseline_match=baseline_match, + ) + + +def _write_maintenance(root: Path, *, memories: tuple[str, ...]) -> None: + (root / "metrics").mkdir(parents=True) + checkpoint = ( + root + / "cases/case-1/checkpoints/snapshots/events-000419-fixture/driver" + ) + checkpoint.mkdir(parents=True) + (root / "manifest.json").write_text( + json.dumps({"condition_id": "fixture"}), encoding="utf-8" + ) + (root / "metrics/summary.json").write_text( + json.dumps(_summary()), encoding="utf-8" + ) + (root / "metrics/operation_usage.csv").write_text( + "phase,operator,logical_call_count,provider_call_count,total_tokens,known_cost_usd\n" + "insertion,sem_flat_map,2,2,20,0.01\n", + encoding="utf-8", + ) + (root / "metrics/provider_usage.csv").write_text( + "phase,operator,total_tokens,known_cost_usd\n" + "insertion,sem_flat_map,20,0.01\n", + encoding="utf-8", + ) + snapshot = {"state": {"memories": pd.DataFrame({"memory": memories})}} + (checkpoint / "runtime.pkl").write_bytes(pickle.dumps(snapshot)) + + +def _write_retrieval(root: Path) -> None: + (root / "input").mkdir(parents=True) + case = root / "cases/case-1" + case.mkdir(parents=True) + (root / "manifest.json").write_text( + json.dumps({"condition_id": "fixture"}), encoding="utf-8" + ) + questions = ( + {"question_id": "case:q1", "question": "one?", "category": "1"}, + {"question_id": "case:q2", "question": "two?", "category": "2"}, + ) + _write_jsonl(root / "input/questions.jsonl", questions) + _write_jsonl( + case / "retrieval.jsonl", + tuple( + { + "question_id": row["question_id"], + "channels": {"memories": [{"memory": "memory one"}]}, + } + for row in questions + ), + ) + _write_jsonl( + case / "answers.jsonl", + tuple( + {"question_id": row["question_id"], "answer": "answer"} + for row in questions + ), + ) + _write_jsonl( + case / "grades.jsonl", + tuple( + { + "question_id": row["question_id"], + "scorer_id": "official", + "score": 1.0, + } + for row in questions + ) + + ( + { + "question_id": "case:q1", + "scorer_id": "judge", + "score": 1.0, + }, + ), + ) + + +def _write_tar(archive: Path, roots: Mapping[str, Path]) -> None: + with tarfile.open(archive, mode="w") as output: + for member_root, source_root in roots.items(): + for path in sorted(source_root.rglob("*")): + if not path.is_file(): + continue + value = path.read_bytes() + info = tarfile.TarInfo( + f"{member_root}/{path.relative_to(source_root).as_posix()}" + ) + info.size = len(value) + output.addfile(info, io.BytesIO(value)) + + +def _write_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None: + path.write_text( + "".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8" + ) + + +def _summary() -> Mapping[str, Any]: + return { + "completed_event_count": 2, + "final_successful_provider_usage": { + "phases": { + "insertion": { + "provider_call_count": 2, + "prompt_tokens": 18, + "cache_hit_tokens": 0, + "cache_miss_tokens": 18, + "completion_tokens": 2, + "total_tokens": 20, + "known_cost_usd": 0.01, + } + } + }, + } + + +def _run( + *, + name: str, + memories: tuple[str, ...], + retrieval_memories: Mapping[str, tuple[str, ...]], + answers: Mapping[str, str], + scores: Mapping[str, float], + second_question: str = "question two", +) -> RunArtifacts: + questions = ( + {"question_id": "case:q1", "question": "question one", "category": "1"}, + {"question_id": "case:q2", "question": second_question, "category": "2"}, + ) + return RunArtifacts( + name=name, + manifest={}, + maintenance_manifest={}, + questions=questions, + retrievals=tuple( + { + "question_id": question_id, + "channels": { + "memories": [{"memory": value} for value in values] + }, + } + for question_id, values in retrieval_memories.items() + ), + answers=tuple( + {"question_id": question_id, "answer": answer} + for question_id, answer in answers.items() + ), + grades=tuple( + { + "question_id": question_id, + "scorer_id": "official", + "score": score, + } + for question_id, score in scores.items() + ), + summary=_summary(), + operation_usage=(), + provider_usage=(), + memories=memories, + evidence_sha256=f"digest-{name}", + ) diff --git a/tools/analysis/mem0_search_filter_causal_audit.py b/tools/analysis/mem0_search_filter_causal_audit.py new file mode 100644 index 0000000..cb3d9f5 --- /dev/null +++ b/tools/analysis/mem0_search_filter_causal_audit.py @@ -0,0 +1,1227 @@ +"""Audit Mem0 Search-Filter quality changes from immutable benchmark artifacts.""" + +from __future__ import annotations + +import argparse +from collections import Counter, defaultdict +from collections.abc import Iterator, Mapping, Sequence +import csv +from dataclasses import dataclass +import hashlib +import io +import json +import math +from pathlib import Path, PurePosixPath +import pickle +import random +import re +import tarfile +from typing import Any, Protocol + +from tools.analysis.semantic_pair_candidates import ( + CandidateStrategy, + PairGroup, + PairScorer, + SentenceTransformerCosineScorer, + TarSource, + _endpoint_pair_id, + _select_candidates, +) + + +class AuditError(ValueError): + """Raised when causal-audit evidence is missing or inconsistent.""" + + +@dataclass(frozen=True) +class MaintenanceEvidence: + """Immutable maintenance evidence for one benchmark condition.""" + + name: str + manifest: Mapping[str, Any] + summary: Mapping[str, Any] + operation_usage: tuple[Mapping[str, str], ...] + provider_usage: tuple[Mapping[str, str], ...] + memories: tuple[str, ...] | None + evidence_sha256: str + + +@dataclass(frozen=True) +class RetrievalEvidence: + """Immutable retrieval, answer, and grade evidence for one condition.""" + + name: str + manifest: Mapping[str, Any] + questions: tuple[Mapping[str, Any], ...] + retrievals: tuple[Mapping[str, Any], ...] + answers: tuple[Mapping[str, Any], ...] + grades: tuple[Mapping[str, Any], ...] + evidence_sha256: str + + +@dataclass(frozen=True) +class RunArtifacts: + """Combined maintenance and retrieval evidence for one condition.""" + + name: str + manifest: Mapping[str, Any] + maintenance_manifest: Mapping[str, Any] + questions: tuple[Mapping[str, Any], ...] + retrievals: tuple[Mapping[str, Any], ...] + answers: tuple[Mapping[str, Any], ...] + grades: tuple[Mapping[str, Any], ...] + summary: Mapping[str, Any] + operation_usage: tuple[Mapping[str, str], ...] + provider_usage: tuple[Mapping[str, str], ...] + memories: tuple[str, ...] | None + evidence_sha256: str + + +class RunReader(Protocol): + """Read named bytes from an immutable run root.""" + + def read(self, relative_path: str) -> bytes: + """Return one required artifact.""" + + ... + + def find_one(self, pattern: re.Pattern[str]) -> tuple[str, bytes]: + """Return the unique artifact whose relative path matches pattern.""" + + ... + + +class DirectoryRunReader: + """Read one benchmark run directly from a directory.""" + + def __init__(self, root: Path) -> None: + self.root = root.resolve() + + def read(self, relative_path: str) -> bytes: + target = (self.root / _safe_relative_path(relative_path)).resolve() + if not target.is_relative_to(self.root): + raise AuditError(f"artifact escapes run root: {relative_path}") + try: + return target.read_bytes() + except OSError as error: + raise AuditError(f"cannot read required artifact: {target}") from error + + def find_one(self, pattern: re.Pattern[str]) -> tuple[str, bytes]: + matches = [ + path + for path in self.root.rglob("*") + if path.is_file() and pattern.fullmatch(path.relative_to(self.root).as_posix()) + ] + if len(matches) != 1: + raise AuditError( + f"expected one artifact matching {pattern.pattern!r} under " + f"{self.root}, found {len(matches)}" + ) + path = matches[0] + return path.relative_to(self.root).as_posix(), path.read_bytes() + + +class CollectedTarRunReader: + """Read one run from bytes collected during a single tar stream.""" + + def __init__(self, *, root: str, members: Mapping[str, bytes]) -> None: + self.root = root.rstrip("/") + self.members = dict(members) + + def read(self, relative_path: str) -> bytes: + relative = _safe_relative_path(relative_path).as_posix() + try: + return self.members[relative] + except KeyError as error: + raise AuditError( + f"missing tar artifact: {self.root}/{relative}" + ) from error + + def find_one(self, pattern: re.Pattern[str]) -> tuple[str, bytes]: + matches = [ + (name, value) + for name, value in self.members.items() + if pattern.fullmatch(name) + ] + if len(matches) != 1: + raise AuditError( + f"expected one tar artifact matching {pattern.pattern!r} under " + f"{self.root}, found {len(matches)}" + ) + return matches[0] + + +_MAINTENANCE_FIXED_MEMBERS = ( + "manifest.json", + "metrics/summary.json", + "metrics/operation_usage.csv", + "metrics/provider_usage.csv", +) +_RETRIEVAL_FIXED_MEMBERS = ("manifest.json", "input/questions.jsonl") +_CASE_MEMBER_NAMES = frozenset( + {"retrieval.jsonl", "answers.jsonl", "grades.jsonl"} +) +_FINAL_RUNTIME_PATTERN = re.compile( + r"cases/[^/]+/checkpoints/snapshots/events-000419-[^/]+/driver/runtime\.pkl" +) + + +def collect_tar_run_readers( + archive_path: Path, + roots: Mapping[str, str], +) -> dict[str, CollectedTarRunReader]: + """Collect only small required evidence for several runs in one tar pass.""" + + archive = archive_path.resolve() + normalized_roots = {name: root.strip("/") for name, root in roots.items()} + collected: dict[str, dict[str, bytes]] = {name: {} for name in roots} + try: + stream = tarfile.open(archive, mode="r|") + except (OSError, tarfile.TarError) as error: + raise AuditError(f"cannot stream tar archive: {archive}") from error + with stream: + for member in stream: + if not member.isfile(): + continue + for name, root in normalized_roots.items(): + prefix = f"{root}/" + if not member.name.startswith(prefix): + continue + relative = member.name.removeprefix(prefix) + if not _wanted_tar_member(relative, source_name=name): + break + handle = stream.extractfile(member) + if handle is None: + raise AuditError(f"cannot read tar member: {member.name}") + if relative in collected[name]: + raise AuditError(f"duplicate tar member: {member.name}") + collected[name][relative] = handle.read() + break + return { + name: CollectedTarRunReader(root=normalized_roots[name], members=members) + for name, members in collected.items() + } + + +def _wanted_tar_member(relative: str, *, source_name: str) -> bool: + if source_name.endswith("-maintenance"): + return relative in _MAINTENANCE_FIXED_MEMBERS or bool( + source_name != "native-maintenance" + and _FINAL_RUNTIME_PATTERN.fullmatch(relative) + ) + if relative in _RETRIEVAL_FIXED_MEMBERS: + return True + path = PurePosixPath(relative) + return ( + len(path.parts) == 3 + and path.parts[0] == "cases" + and path.parts[2] in _CASE_MEMBER_NAMES + ) + + +def load_maintenance( + reader: RunReader, *, name: str, load_memory_view: bool = True +) -> MaintenanceEvidence: + """Parse maintenance-only evidence from one run root.""" + + raw: dict[str, bytes] = { + path: reader.read(path) for path in _MAINTENANCE_FIXED_MEMBERS + } + runtime_path: str | None = None + runtime_bytes: bytes | None = None + if load_memory_view: + runtime_path, runtime_bytes = reader.find_one(_FINAL_RUNTIME_PATTERN) + raw[runtime_path] = runtime_bytes + return MaintenanceEvidence( + name=name, + manifest=_json_object(raw["manifest.json"], source=f"{name}:manifest.json"), + summary=_json_object( + raw["metrics/summary.json"], source=f"{name}:metrics/summary.json" + ), + operation_usage=_csv_rows( + raw["metrics/operation_usage.csv"], + source=f"{name}:metrics/operation_usage.csv", + ), + provider_usage=_csv_rows( + raw["metrics/provider_usage.csv"], + source=f"{name}:metrics/provider_usage.csv", + ), + memories=( + _runtime_memories(runtime_bytes, source=f"{name}:{runtime_path}") + if runtime_bytes is not None and runtime_path is not None + else None + ), + evidence_sha256=_evidence_digest(raw), + ) + + +def load_retrieval(reader: RunReader, *, name: str) -> RetrievalEvidence: + """Parse retrieval, answer, and grade evidence from one run root.""" + + raw: dict[str, bytes] = { + path: reader.read(path) for path in _RETRIEVAL_FIXED_MEMBERS + } + case_rows: dict[str, tuple[Mapping[str, Any], ...]] = {} + for filename in sorted(_CASE_MEMBER_NAMES): + relative, value = reader.find_one( + re.compile(rf"cases/[^/]+/{re.escape(filename)}") + ) + raw[relative] = value + case_rows[filename] = _jsonl(value, source=f"{name}:{relative}") + questions = _jsonl( + raw["input/questions.jsonl"], source=f"{name}:input/questions.jsonl" + ) + question_ids = _unique_rows(questions, "question_id", source=name) + for filename, rows in case_rows.items(): + if filename == "grades.jsonl": + grade_groups = _grades_by_scorer(rows) + grade_ids = set().union(*(set(group) for group in grade_groups.values())) + if not grade_ids <= question_ids or not any( + set(group) == question_ids for group in grade_groups.values() + ): + raise AuditError( + f"grade coverage mismatch in {name}:{filename}" + ) + continue + row_ids = _unique_rows(rows, "question_id", source=f"{name}:{filename}") + if row_ids != question_ids: + missing = sorted(question_ids - row_ids)[:3] + extra = sorted(row_ids - question_ids)[:3] + raise AuditError( + f"question coverage mismatch in {name}:{filename}; " + f"missing={missing}, extra={extra}" + ) + return RetrievalEvidence( + name=name, + manifest=_json_object(raw["manifest.json"], source=f"{name}:manifest.json"), + questions=questions, + retrievals=case_rows["retrieval.jsonl"], + answers=case_rows["answers.jsonl"], + grades=case_rows["grades.jsonl"], + evidence_sha256=_evidence_digest(raw), + ) + + +def _evidence_digest(values: Mapping[str, bytes]) -> str: + digest = hashlib.sha256() + for path, value in sorted(values.items()): + digest.update(path.encode("utf-8")) + digest.update(b"\0") + digest.update(value) + return digest.hexdigest() + + +def analyze_conditions( + native: RunArtifacts | None, + oracle: RunArtifacts, + optimized: RunArtifacts, + *, + pair_audit: Mapping[str, Any], + native_reference: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Return one deterministic three-condition causal-audit report.""" + + compared_runs = (oracle, optimized) if native is None else (native, oracle, optimized) + _validate_question_contract(compared_runs) + conditions: dict[str, Any] = { + run.name: _condition_summary(run) for run in compared_runs + } + retained_memories = set( + _string_sequence(pair_audit.get("counterfactual_retained_later_memories")) + ) + comparisons = { + "optimized_minus_oracle": _compare_runs( + oracle, + optimized, + attributed_memories=retained_memories, + ), + } + if native is not None: + comparisons["oracle_minus_native"] = _compare_runs(native, oracle) + comparisons["optimized_minus_native"] = _compare_runs(native, optimized) + elif native_reference is not None: + conditions["native"] = dict(native_reference) + comparisons["native_pairwise_comparison"] = { + "available": False, + "reason": "formal Native per-question artifacts are absent from the archive", + } + else: + raise AuditError("native artifacts or a frozen Native reference are required") + if oracle.memories is None or optimized.memories is None: + raise AuditError("AM Oracle and Search-Filter memory views are required") + oracle_memories = set(oracle.memories) + optimized_memories = set(optimized.memories) + comparisons["optimized_minus_oracle"]["memory_text_diff"] = { + "oracle_only_count": len(oracle_memories - optimized_memories), + "optimized_only_count": len(optimized_memories - oracle_memories), + "exact_intersection_count": len(oracle_memories & optimized_memories), + "counterfactual_retained_count": len(retained_memories), + "counterfactual_retained_present_in_optimized_count": len( + retained_memories & optimized_memories + ), + "counterfactual_retained_present_in_oracle_count": len( + retained_memories & oracle_memories + ), + } + return { + "schema_version": 1, + "analysis_kind": "mem0-search-filter-causal-audit", + "conditions": conditions, + "contract_comparison": _contract_comparison(compared_runs), + "pair_audit": dict(pair_audit), + "comparisons": comparisons, + "interpretation_inputs": _interpretation_inputs( + comparisons["optimized_minus_oracle"], pair_audit + ), + } + + +def analyze_top_k_pairs( + groups: Iterator[PairGroup], + scorer: PairScorer, + *, + top_k: int, +) -> dict[str, Any]: + """Audit directed Top-k false negatives and counterfactual retained rows.""" + + accumulator = _TopKPairAudit(top_k=top_k, scorer_metadata=scorer.metadata) + for group in groups: + scores = tuple(float(value) for value in scorer.score(group)) + accumulator.add(group, scores) + return accumulator.report() + + +def analyze_top_k_pair_parity( + groups: Iterator[PairGroup], + primary: PairScorer, + comparison: PairScorer, + *, + top_k: int, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Compare two scorers over exactly the same streamed pair groups.""" + + first = _TopKPairAudit(top_k=top_k, scorer_metadata=primary.metadata) + second = _TopKPairAudit(top_k=top_k, scorer_metadata=comparison.metadata) + for group in groups: + first.add(group, tuple(float(value) for value in primary.score(group))) + second.add(group, tuple(float(value) for value in comparison.score(group))) + first_report = first.report() + second_report = second.report() + comparison_report = { + "primary_scorer": first_report["scorer"], + "comparison_scorer": second_report["scorer"], + "selected_pair_set_equal": ( + first_report["selected_pair_ids_sha256"] + == second_report["selected_pair_ids_sha256"] + ), + "missed_positive_set_equal": ( + first_report["missed_positive_pair_ids_sha256"] + == second_report["missed_positive_pair_ids_sha256"] + ), + "primary_selected_pair_count": first_report["selected_pair_count"], + "comparison_selected_pair_count": second_report["selected_pair_count"], + "primary_missed_positive_count": first_report[ + "missed_oracle_positive_pair_count" + ], + "comparison_missed_positive_count": second_report[ + "missed_oracle_positive_pair_count" + ], + } + return first_report, comparison_report + + +class _TopKPairAudit: + def __init__(self, *, top_k: int, scorer_metadata: Mapping[str, Any]) -> None: + self.strategy = CandidateStrategy(top_k=top_k) + self.scorer_metadata = scorer_metadata + self.pair_count = 0 + self.positive_count = 0 + self.selected_count = 0 + self.selected_positive_count = 0 + self.missed: list[dict[str, Any]] = [] + self.positive_later_ids: set[str] = set() + self.selected_positive_later_ids: set[str] = set() + self.later_text: dict[str, str] = {} + self.selected_digest = hashlib.sha256() + self.missed_digest = hashlib.sha256() + + def add(self, group: PairGroup, scores: Sequence[float]) -> None: + if group.operator != "sem_filter" or group.direction != "right-to-left": + raise AuditError( + "Mem0 pair audit requires right-to-left sem_filter groups" + ) + if len(scores) != len(group.pairs) or not all(map(math.isfinite, scores)): + raise AuditError(f"invalid pair scores for group {group.group_id}") + selected = _select_candidates(group, scores, self.strategy) + ranks = _directed_ranks(group, scores) + self.pair_count += len(group.pairs) + self.selected_count += len(selected) + for index in sorted(selected): + self.selected_digest.update( + f"{group.group_id}\0{group.pairs[index].pair_id}\n".encode() + ) + for index, pair in enumerate(group.pairs): + self.later_text.setdefault(pair.right_id, pair.right) + if self.later_text[pair.right_id] != pair.right: + raise AuditError(f"later endpoint text changed: {pair.right_id}") + if not pair.baseline_match: + continue + self.positive_count += 1 + self.positive_later_ids.add(pair.right_id) + if index in selected: + self.selected_positive_count += 1 + self.selected_positive_later_ids.add(pair.right_id) + continue + self.missed_digest.update( + f"{group.group_id}\0{pair.pair_id}\n".encode() + ) + self.missed.append( + { + "group_id": group.group_id, + "event_id": group.event_id, + "pair_id": pair.pair_id, + "earlier_id": pair.left_id, + "later_id": pair.right_id, + "earlier_memory": pair.left, + "later_memory": pair.right, + "cosine_score": scores[index], + "candidate_rank_for_later": ranks[index], + "oracle_duplicate": True, + "researcher_adjudication": None, + } + ) + + def report(self) -> dict[str, Any]: + retained_later_ids = ( + self.positive_later_ids - self.selected_positive_later_ids + ) + return { + "strategy": self.strategy.strategy_id, + "scorer": dict(self.scorer_metadata), + "pair_count": self.pair_count, + "selected_pair_count": self.selected_count, + "selected_pair_ids_sha256": self.selected_digest.hexdigest(), + "pair_reduction": ( + 1.0 - self.selected_count / self.pair_count + if self.pair_count + else 0.0 + ), + "oracle_positive_pair_count": self.positive_count, + "selected_oracle_positive_pair_count": self.selected_positive_count, + "oracle_positive_recall": ( + self.selected_positive_count / self.positive_count + if self.positive_count + else None + ), + "missed_oracle_positive_pair_count": len(self.missed), + "missed_positive_pair_ids_sha256": self.missed_digest.hexdigest(), + "missed_oracle_positive_pairs": self.missed, + "oracle_deleted_later_count": len(self.positive_later_ids), + "counterfactual_retained_later_count": len(retained_later_ids), + "counterfactual_retained_later_memories": sorted( + self.later_text[value] for value in retained_later_ids + ), + } + + +def _directed_ranks(group: PairGroup, scores: Sequence[float]) -> dict[int, int]: + buckets: dict[str, list[int]] = defaultdict(list) + for index, pair in enumerate(group.pairs): + buckets[pair.right_id].append(index) + ranks: dict[int, int] = {} + for indices in buckets.values(): + indices.sort( + key=lambda index: ( + -float(scores[index]), + _endpoint_pair_id( + group.pairs[index].left_id, group.pairs[index].right_id + ), + ) + ) + for rank, index in enumerate(indices, start=1): + ranks[index] = rank + return ranks + + +def _condition_summary(run: RunArtifacts) -> dict[str, Any]: + questions = {str(row["question_id"]): row for row in run.questions} + grades = _grades_by_scorer(run.grades) + return { + "evidence_sha256": run.evidence_sha256, + "question_count": len(questions), + "memory_view_available": run.memories is not None, + "memory_count": len(run.memories) if run.memories is not None else None, + "exact_memory_text_count": ( + len(set(run.memories)) if run.memories is not None else None + ), + "scores": { + scorer: _score_summary(rows, questions) + for scorer, rows in sorted(grades.items()) + }, + "insertion_usage": _insertion_usage(run), + } + + +def _compare_runs( + a: RunArtifacts, + b: RunArtifacts, + *, + attributed_memories: set[str] | None = None, +) -> dict[str, Any]: + questions = {str(row["question_id"]): row for row in a.questions} + retrieval_a = _rows_by_question(a.retrievals, source=a.name) + retrieval_b = _rows_by_question(b.retrievals, source=b.name) + answers_a = _rows_by_question(a.answers, source=a.name) + answers_b = _rows_by_question(b.answers, source=b.name) + grades_a = _grades_by_scorer(a.grades) + grades_b = _grades_by_scorer(b.grades) + scorer_ids = sorted(set(grades_a) | set(grades_b)) + if set(grades_a) != set(grades_b): + raise AuditError(f"scorer mismatch: {a.name} vs {b.name}") + b_only_memories = ( + set(b.memories) - set(a.memories) + if a.memories is not None and b.memories is not None + else set() + ) + attributed = attributed_memories or set() + question_rows: list[dict[str, Any]] = [] + for question_id in sorted(questions, key=_question_sort_key): + context_a = _canonical_context(retrieval_a[question_id]) + context_b = _canonical_context(retrieval_b[question_id]) + answer_a = _canonical_text(answers_a[question_id].get("answer")) + answer_b = _canonical_text(answers_b[question_id].get("answer")) + retrieved_b_texts = set(_retrieved_memory_texts(retrieval_b[question_id])) + retrieved_a_texts = set(_retrieved_memory_texts(retrieval_a[question_id])) + scores = { + scorer: { + "a": float(grades_a[scorer][question_id]["score"]), + "b": float(grades_b[scorer][question_id]["score"]), + "delta_b_minus_a": float(grades_b[scorer][question_id]["score"]) + - float(grades_a[scorer][question_id]["score"]), + } + for scorer in scorer_ids + if question_id in grades_a[scorer] and question_id in grades_b[scorer] + } + question_rows.append( + { + "question_id": question_id, + "category": str(questions[question_id].get("category") or ""), + "question": str(questions[question_id].get("question") or ""), + "context_same": context_a == context_b, + "answer_same": answer_a == answer_b, + "a_answer": answer_a, + "b_answer": answer_b, + "a_retrieved_memories": sorted(retrieved_a_texts), + "b_retrieved_memories": sorted(retrieved_b_texts), + "b_only_memory_retrieved": bool( + retrieved_b_texts & b_only_memories + ), + "b_only_retrieved_memories": sorted( + retrieved_b_texts & b_only_memories + ), + "attributed_memory_retrieved": bool( + retrieved_b_texts & attributed + ), + "attributed_retrieved_memories": sorted( + retrieved_b_texts & attributed + ), + "scores": scores, + } + ) + return { + "a": a.name, + "b": b.name, + "scorers": { + scorer: { + **_paired_score_summary( + [ + row["scores"][scorer]["delta_b_minus_a"] + for row in question_rows + if scorer in row["scores"] + ] + ), + "change_decomposition": _score_change_decomposition( + question_rows, scorer=scorer + ), + } + for scorer in scorer_ids + }, + "context_answer_decomposition": _context_answer_decomposition(question_rows), + "questions": question_rows, + } + + +def _context_answer_decomposition( + rows: Sequence[Mapping[str, Any]], +) -> dict[str, int]: + counts: Counter[str] = Counter() + for row in rows: + if row["attributed_memory_retrieved"]: + label = "search_filter_retained_memory_retrieved" + elif row["context_same"] and row["answer_same"]: + label = "same_context_same_answer" + elif row["context_same"]: + label = "same_context_different_answer" + elif row["b_only_memory_retrieved"]: + label = "different_context_with_b_only_memory" + else: + label = "different_context_without_exact_b_only_memory" + counts[label] += 1 + return dict(sorted(counts.items())) + + +def _score_change_decomposition( + rows: Sequence[Mapping[str, Any]], *, scorer: str +) -> dict[str, Mapping[str, int | float]]: + grouped: dict[str, list[float]] = defaultdict(list) + for row in rows: + score = row["scores"].get(scorer) + if not isinstance(score, Mapping): + continue + if row["attributed_memory_retrieved"]: + label = "search_filter_retained_memory_retrieved" + elif not row["context_same"]: + label = "retrieval_context_changed_without_attribution" + elif not row["answer_same"]: + label = "same_context_answer_changed" + elif float(score["delta_b_minus_a"]) != 0.0: + label = "same_answer_grade_changed" + else: + label = "no_observed_change" + grouped[label].append(float(score["delta_b_minus_a"])) + denominator = sum(len(values) for values in grouped.values()) + return { + label: { + "question_count": len(values), + "score_delta_sum": sum(values), + "overall_mean_contribution": ( + sum(values) / denominator if denominator else 0.0 + ), + } + for label, values in sorted(grouped.items()) + } + + +def _paired_score_summary(deltas: Sequence[float]) -> dict[str, Any]: + if not deltas: + return { + "question_count": 0, + "mean_delta_b_minus_a": None, + "wins_b": 0, + "ties": 0, + "losses_b": 0, + "bootstrap_95_ci": [None, None], + } + epsilon = 1e-12 + return { + "question_count": len(deltas), + "mean_delta_b_minus_a": sum(deltas) / len(deltas), + "wins_b": sum(value > epsilon for value in deltas), + "ties": sum(abs(value) <= epsilon for value in deltas), + "losses_b": sum(value < -epsilon for value in deltas), + "bootstrap_95_ci": list(_bootstrap_mean_interval(deltas)), + } + + +def _bootstrap_mean_interval( + values: Sequence[float], *, samples: int = 20_000 +) -> tuple[float, float]: + rng = random.Random(0) + size = len(values) + means = sorted( + sum(values[rng.randrange(size)] for _ in range(size)) / size + for _ in range(samples) + ) + return means[int(samples * 0.025)], means[int(samples * 0.975) - 1] + + +def _score_summary( + rows: Mapping[str, Mapping[str, Any]], + questions: Mapping[str, Mapping[str, Any]], +) -> dict[str, Any]: + by_category: dict[str, list[float]] = defaultdict(list) + for question_id, row in rows.items(): + by_category[str(questions[question_id].get("category") or "")].append( + float(row["score"]) + ) + all_scores = [score for scores in by_category.values() for score in scores] + return { + "question_count": len(all_scores), + "mean": sum(all_scores) / len(all_scores) if all_scores else None, + "by_category": { + category: { + "question_count": len(scores), + "mean": sum(scores) / len(scores), + } + for category, scores in sorted(by_category.items()) + }, + } + + +def _insertion_usage(run: RunArtifacts) -> dict[str, Any]: + summary_usage = run.summary.get("final_successful_provider_usage") + if not isinstance(summary_usage, Mapping): + summary_usage = run.summary.get("actual_provider_usage") + if not isinstance(summary_usage, Mapping): + raise AuditError(f"missing provider usage summary: {run.name}") + phases = summary_usage.get("phases") + insertion = phases.get("insertion") if isinstance(phases, Mapping) else None + if not isinstance(insertion, Mapping): + insertion = summary_usage + events = _completed_event_count(run) + total_tokens = _number(insertion.get("total_tokens"), default=0.0) + return { + "completed_events": events, + "provider_call_count": int(_number(insertion.get("provider_call_count"), default=0)), + "prompt_tokens": int(_number(insertion.get("prompt_tokens"), default=0)), + "cache_hit_tokens": int(_number(insertion.get("cache_hit_tokens"), default=0)), + "cache_miss_tokens": int(_number(insertion.get("cache_miss_tokens"), default=0)), + "completion_tokens": int(_number(insertion.get("completion_tokens"), default=0)), + "total_tokens": int(total_tokens), + "tokens_per_event": total_tokens / events if events else None, + "known_cost_usd": insertion.get("known_cost_usd"), + "operations": _operation_usage(run.operation_usage), + } + + +def _completed_event_count(run: RunArtifacts) -> int: + for key in ("completed_event_count", "event_count"): + value = run.summary.get(key) + if isinstance(value, int) and not isinstance(value, bool): + return value + raise AuditError(f"missing completed event count: {run.name}") + + +def _operation_usage(rows: Sequence[Mapping[str, str]]) -> list[dict[str, Any]]: + result: list[dict[str, Any]] = [] + for row in rows: + operation = row.get("operation") or row.get("operator") or row.get("op") + if not operation: + continue + result.append( + { + "operation": operation, + "logical_call_count": _optional_number(row.get("logical_call_count")), + "provider_call_count": _optional_number(row.get("provider_call_count")), + "total_tokens": _optional_number(row.get("total_tokens")), + "known_cost_usd": _optional_number(row.get("known_cost_usd")), + } + ) + return result + + +def _contract_comparison(runs: Sequence[RunArtifacts]) -> dict[str, Any]: + fields = ( + "policy_input_fingerprint", + "memory_provider_model_id", + "thinking_enabled", + "answer_prompt_digests", + "official_scorer_digest", + "zep_judge_digest", + "retrieval_recipe_id", + "retrieval_recipe_digest", + ) + rows: dict[str, Any] = {} + for field in fields: + values = { + run.name: run.manifest.get(field) + for run in runs + } + rows[field] = {"values": values, "all_equal": len(_json_values(values)) == 1} + rows["known_maintenance_design_difference"] = { + "native": "message-to-existing-memory top-10 inside one extraction call", + "agent_memory": "fact extraction followed by fact-to-fact semantic dedup", + } + return rows + + +def _interpretation_inputs( + comparison: Mapping[str, Any], pair_audit: Mapping[str, Any] +) -> dict[str, Any]: + decomposition = comparison["context_answer_decomposition"] + return { + "same_context_different_answer_count": decomposition.get( + "same_context_different_answer", 0 + ), + "search_filter_retained_memory_retrieved_count": decomposition.get( + "search_filter_retained_memory_retrieved", 0 + ), + "missed_oracle_positive_pair_count": pair_audit.get( + "missed_oracle_positive_pair_count" + ), + "counterfactual_retained_later_count": pair_audit.get( + "counterfactual_retained_later_count" + ), + "decision": "requires-researcher-adjudication", + } + + +def _validate_question_contract(runs: Sequence[RunArtifacts]) -> None: + baseline = { + str(row["question_id"]): ( + str(row.get("question") or ""), + str(row.get("category") or ""), + json.dumps(row.get("gold_answer"), ensure_ascii=False, sort_keys=True), + ) + for row in runs[0].questions + } + for run in runs[1:]: + current = { + str(row["question_id"]): ( + str(row.get("question") or ""), + str(row.get("category") or ""), + json.dumps(row.get("gold_answer"), ensure_ascii=False, sort_keys=True), + ) + for row in run.questions + } + if current != baseline: + raise AuditError(f"question contract differs for {run.name}") + + +def _runtime_memories(value: bytes, *, source: str) -> tuple[str, ...]: + try: + snapshot = pickle.loads(value) + state = snapshot["state"] + frame = state["memories"] + memories = tuple(str(item) for item in frame["memory"].tolist()) + except Exception as error: + raise AuditError(f"cannot read runtime memories: {source}") from error + if len(memories) != len(set(memories)): + raise AuditError(f"runtime memory view contains exact duplicates: {source}") + return memories + + +def _canonical_context(row: Mapping[str, Any]) -> tuple[str, ...]: + memories = _retrieved_memory_texts(row) + if memories: + return tuple(_canonical_text(value) for value in memories) + context = row.get("context") + return tuple(_canonical_text(context).split("\n")) + + +def _retrieved_memory_texts(row: Mapping[str, Any]) -> tuple[str, ...]: + channels = row.get("channels") + if not isinstance(channels, Mapping): + return () + memories = channels.get("memories") + if not isinstance(memories, list): + return () + result: list[str] = [] + for item in memories: + if isinstance(item, Mapping) and isinstance(item.get("memory"), str): + result.append(item["memory"]) + return tuple(result) + + +def _canonical_text(value: object) -> str: + return " ".join(str(value or "").split()) + + +def _grades_by_scorer( + rows: Sequence[Mapping[str, Any]], +) -> dict[str, dict[str, Mapping[str, Any]]]: + result: dict[str, dict[str, Mapping[str, Any]]] = defaultdict(dict) + for row in rows: + scorer = str(row.get("scorer_id") or "") + question_id = str(row.get("question_id") or "") + if not scorer or not question_id or isinstance(row.get("score"), bool): + raise AuditError("grade row is missing scorer_id, question_id, or score") + if question_id in result[scorer]: + raise AuditError(f"duplicate grade: {scorer}:{question_id}") + result[scorer][question_id] = row + return dict(result) + + +def _rows_by_question( + rows: Sequence[Mapping[str, Any]], *, source: str +) -> dict[str, Mapping[str, Any]]: + result: dict[str, Mapping[str, Any]] = {} + for row in rows: + question_id = str(row.get("question_id") or "") + if not question_id or question_id in result: + raise AuditError(f"invalid or duplicate question row in {source}") + result[question_id] = row + return result + + +def _unique_rows( + rows: Sequence[Mapping[str, Any]], key: str, *, source: str +) -> set[str]: + values = [str(row.get(key) or "") for row in rows] + if any(not value for value in values) or len(values) != len(set(values)): + raise AuditError(f"invalid or duplicate {key} in {source}") + return set(values) + + +def _json_object(value: bytes, *, source: str) -> Mapping[str, Any]: + try: + parsed = json.loads(value) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise AuditError(f"invalid JSON object: {source}") from error + if not isinstance(parsed, Mapping): + raise AuditError(f"expected JSON object: {source}") + return parsed + + +def _jsonl(value: bytes, *, source: str) -> tuple[Mapping[str, Any], ...]: + rows: list[Mapping[str, Any]] = [] + try: + text = value.decode("utf-8") + except UnicodeDecodeError as error: + raise AuditError(f"invalid UTF-8 JSONL: {source}") from error + for line_number, line in enumerate(text.splitlines(), start=1): + try: + row = json.loads(line) + except json.JSONDecodeError as error: + raise AuditError(f"invalid JSONL: {source}:{line_number}") from error + if not isinstance(row, Mapping): + raise AuditError(f"expected JSON object: {source}:{line_number}") + rows.append(row) + return tuple(rows) + + +def _csv_rows(value: bytes, *, source: str) -> tuple[Mapping[str, str], ...]: + try: + text = value.decode("utf-8") + return tuple(dict(row) for row in csv.DictReader(io.StringIO(text))) + except (UnicodeDecodeError, csv.Error) as error: + raise AuditError(f"invalid CSV: {source}") from error + + +def _safe_relative_path(value: str) -> PurePosixPath: + path = PurePosixPath(value) + if path.is_absolute() or ".." in path.parts: + raise AuditError(f"unsafe artifact path: {value}") + return path + + +def _number(value: object, *, default: float) -> float: + if value is None: + return default + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise AuditError(f"expected numeric value, got {value!r}") + return float(value) + + +def _optional_number(value: object) -> int | float | None: + if value in (None, ""): + return None + try: + number = float(str(value)) + except ValueError as error: + raise AuditError(f"expected optional numeric CSV value, got {value!r}") from error + return int(number) if number.is_integer() else number + + +def _json_values(values: Mapping[str, Any]) -> set[str]: + return { + json.dumps(value, ensure_ascii=False, sort_keys=True) + for value in values.values() + } + + +def _string_sequence(value: object) -> tuple[str, ...]: + if value is None: + return () + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise AuditError("expected a sequence of counterfactual memory texts") + result = tuple(str(item) for item in value) + if any(not item for item in result): + raise AuditError("counterfactual memory texts must be non-empty") + return result + + +def _question_sort_key(question_id: str) -> tuple[str, int, str]: + prefix, separator, suffix = question_id.rpartition("q") + return prefix, int(suffix) if separator and suffix.isdigit() else 0, question_id + + +def _write_report(output: Path, report: Mapping[str, Any]) -> None: + target = output.resolve() + if target.exists(): + raise FileExistsError(f"audit report already exists: {target}") + if not target.parent.is_dir(): + raise FileNotFoundError(f"audit report parent does not exist: {target.parent}") + target.write_text( + json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def build_parser() -> argparse.ArgumentParser: + """Build the causal-audit CLI parser.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--archive", type=Path, required=True) + parser.add_argument("--native-maintenance-root") + parser.add_argument("--native-retrieval-root") + parser.add_argument("--native-reference", type=Path) + parser.add_argument("--oracle-maintenance-root", required=True) + parser.add_argument("--oracle-retrieval-root", required=True) + parser.add_argument("--optimized-maintenance-root", type=Path, required=True) + parser.add_argument("--optimized-retrieval-root", type=Path, required=True) + parser.add_argument("--embedding-model", required=True) + parser.add_argument("--embedding-revision", required=True) + parser.add_argument("--embedding-device", default="cpu") + parser.add_argument("--parity-embedding-device") + parser.add_argument("--embedding-batch-size", type=int, default=32) + parser.add_argument("--top-k", type=int, default=10) + parser.add_argument("--output", type=Path, required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + """Run one read-only Mem0 causal audit.""" + + args = build_parser().parse_args(argv) + native_roots_supplied = bool( + args.native_maintenance_root or args.native_retrieval_root + ) + if native_roots_supplied != bool( + args.native_maintenance_root and args.native_retrieval_root + ): + raise AuditError( + "native maintenance and retrieval roots must be provided together" + ) + if native_roots_supplied == bool(args.native_reference): + raise AuditError( + "provide either Native artifact roots or one frozen Native reference" + ) + roots = { + "oracle-maintenance": args.oracle_maintenance_root, + "oracle-retrieval": args.oracle_retrieval_root, + } + if native_roots_supplied: + roots.update( + { + "native-maintenance": args.native_maintenance_root, + "native-retrieval": args.native_retrieval_root, + } + ) + tar_readers = collect_tar_run_readers(args.archive, roots) + native: RunArtifacts | None = None + native_reference: Mapping[str, Any] | None = None + if native_roots_supplied: + native_maintenance = load_maintenance( + tar_readers["native-maintenance"], + name="native-maintenance", + load_memory_view=False, + ) + native_retrieval = load_retrieval( + tar_readers["native-retrieval"], name="native" + ) + native = _combine_maintenance_and_retrieval( + native_maintenance, native_retrieval, name="native" + ) + else: + if args.native_reference is None: + raise AuditError("missing Native reference") + native_reference = _read_native_reference(args.native_reference) + oracle_maintenance = load_maintenance( + tar_readers["oracle-maintenance"], name="oracle-maintenance" + ) + oracle_retrieval = load_retrieval( + tar_readers["oracle-retrieval"], name="oracle" + ) + optimized_maintenance = load_maintenance( + DirectoryRunReader(args.optimized_maintenance_root), + name="optimized-maintenance", + ) + optimized_retrieval = load_retrieval( + DirectoryRunReader(args.optimized_retrieval_root), name="optimized" + ) + oracle = _combine_maintenance_and_retrieval( + oracle_maintenance, oracle_retrieval, name="oracle" + ) + optimized = _combine_maintenance_and_retrieval( + optimized_maintenance, optimized_retrieval, name="optimized" + ) + scorer = SentenceTransformerCosineScorer( + model=args.embedding_model, + revision=args.embedding_revision, + device=args.embedding_device, + batch_size=args.embedding_batch_size, + ) + pair_source = TarSource( + args.archive, + args.oracle_maintenance_root, + read_workers=1, + ) + pair_parity: Mapping[str, Any] | None = None + if args.parity_embedding_device: + parity_scorer = SentenceTransformerCosineScorer( + model=args.embedding_model, + revision=args.embedding_revision, + device=args.parity_embedding_device, + batch_size=args.embedding_batch_size, + ) + pair_audit, pair_parity = analyze_top_k_pair_parity( + pair_source.iter_groups(phase="insertion"), + scorer, + parity_scorer, + top_k=args.top_k, + ) + else: + pair_audit = analyze_top_k_pairs( + pair_source.iter_groups(phase="insertion"), scorer, top_k=args.top_k + ) + report = analyze_conditions( + native, + oracle, + optimized, + pair_audit=pair_audit, + native_reference=native_reference, + ) + report["pair_device_parity"] = pair_parity + report["sources"] = { + "archive": str(args.archive.resolve()), + "native_maintenance_root": args.native_maintenance_root, + "native_retrieval_root": args.native_retrieval_root, + "native_reference": ( + str(args.native_reference.resolve()) if args.native_reference else None + ), + "oracle_maintenance_root": args.oracle_maintenance_root, + "oracle_retrieval_root": args.oracle_retrieval_root, + "optimized_maintenance_root": str(args.optimized_maintenance_root.resolve()), + "optimized_retrieval_root": str(args.optimized_retrieval_root.resolve()), + "pair_trace": pair_source.description, + "pair_source_stats": vars(pair_source.stats), + } + _write_report(args.output, report) + return 0 + + +def _read_native_reference(path: Path) -> Mapping[str, Any]: + value = path.resolve().read_bytes() + parsed = dict(_json_object(value, source=str(path))) + parsed["evidence_sha256"] = hashlib.sha256(value).hexdigest() + parsed["evidence_kind"] = "frozen-report-reference" + return parsed + + +def _combine_maintenance_and_retrieval( + maintenance: MaintenanceEvidence, + retrieval: RetrievalEvidence, + *, + name: str, +) -> RunArtifacts: + return RunArtifacts( + name=name, + manifest=retrieval.manifest, + maintenance_manifest=maintenance.manifest, + questions=retrieval.questions, + retrievals=retrieval.retrievals, + answers=retrieval.answers, + grades=retrieval.grades, + summary=maintenance.summary, + operation_usage=maintenance.operation_usage, + provider_usage=maintenance.provider_usage, + memories=maintenance.memories, + evidence_sha256=hashlib.sha256( + f"{maintenance.evidence_sha256}:{retrieval.evidence_sha256}".encode() + ).hexdigest(), + ) + + +if __name__ == "__main__": + raise SystemExit(main()) From 372a5010f63ac6217b92fa93c901ce1dc1b60ea9 Mon Sep 17 00:00:00 2001 From: Von Date: Mon, 31 Aug 2026 12:53:53 +0800 Subject: [PATCH 4/5] [feat-runtime] Incrementally maintain relational inner joins - compile deterministic inner joins into bag-preserving delta maintenance - combine old results with left, right, and cross-delta join terms - fall back to full recomputation when parent rows are retracted - cover self joins, multiplicity, downstream deltas, and checkpoint restore Signed-off-by: Von --- .../planner/differential_policy.py | 95 +++++- src/agent_memory/runtime/executor.py | 38 +++ tests/test_relational_join_ivm.py | 274 ++++++++++++++++++ 3 files changed, 399 insertions(+), 8 deletions(-) create mode 100644 tests/test_relational_join_ivm.py diff --git a/src/agent_memory/planner/differential_policy.py b/src/agent_memory/planner/differential_policy.py index 791a494..95cbc1c 100644 --- a/src/agent_memory/planner/differential_policy.py +++ b/src/agent_memory/planner/differential_policy.py @@ -243,20 +243,27 @@ def _compile_query(self, query: QueryExpr, *, spec: MemorySpec) -> str: execution_kind = ( "semantic_row" if query.op in _ROW_LOCAL_SEMANTIC_OPS else "deterministic" ) - if query.op == "sem_agg": + maintenance_query: QueryExpr | None = None + if query.op == "join" and str(query.params.get("how", "inner")) == "inner": + execution_kind = "relational_state" + maintenance_query = self._relational_inner_join_maintenance_query( + local_query, + input_node_ids, + current_node_id=self._node_id(query), + ) + elif query.op == "sem_agg": execution_kind = "semantic_state" + maintenance_query = self._semantic_maintenance_query( + local_query, + input_node_ids, + current_node_id=self._node_id(query), + ) node = self._make_node( query=query, local_query=local_query, input_node_ids=input_node_ids, execution_kind=execution_kind, - maintenance_query=self._semantic_maintenance_query( - local_query, - input_node_ids, - current_node_id=self._node_id(query), - ) - if execution_kind == "semantic_state" - else None, + maintenance_query=maintenance_query, ) return self._register(query, node) @@ -389,6 +396,78 @@ def _semantic_maintenance_query( source_input=changed_source, ) + def _relational_inner_join_maintenance_query( + self, + local_query: QueryExpr, + input_node_ids: tuple[str, ...], + *, + current_node_id: str, + ) -> QueryExpr: + """Build bag-preserving append maintenance for one relational inner join.""" + + if local_query.op != "join" or str(local_query.params.get("how", "inner")) != "inner": + raise ValueError("Relational join maintenance requires how='inner'") + if len(input_node_ids) != 2 or len(local_query.inputs) != 2: + raise ValueError("Relational inner join requires exactly two inputs") + + left_node_id, right_node_id = input_node_ids + left_old = self._node_leaf(left_node_id) + right_old = self._node_leaf(right_node_id) + left_inserted = self._inserted_node_leaf(left_node_id) + right_inserted = self._inserted_node_leaf(right_node_id) + + def join(left: QueryExpr, right: QueryExpr) -> QueryExpr: + return QueryExpr( + op="join", + inputs=( + self._preserve_input_alias(local_query.inputs[0], left), + self._preserve_input_alias(local_query.inputs[1], right), + ), + params=local_query.params, + ) + + current_join = QueryExpr( + op="materialized_view", + params={ + "name": current_node_id, + "columns": output_columns(local_query), + }, + ) + delta_join = QueryExpr( + op="concat", + inputs=( + QueryExpr( + op="concat", + inputs=( + join(left_inserted, right_old), + join(left_old, right_inserted), + ), + ), + join(left_inserted, right_inserted), + ), + ) + return QueryExpr(op="concat", inputs=(current_join, delta_join)) + + def _inserted_node_leaf(self, node_id: str) -> QueryExpr: + """Return the runtime binding for one policy node's inserted rows.""" + + source = self._node_leaf(node_id) + return QueryExpr( + op="materialized_view", + params={ + "name": f"{node_id}__inserted", + "columns": source.params["columns"], + }, + ) + + @staticmethod + def _preserve_input_alias(template: QueryExpr, source: QueryExpr) -> QueryExpr: + """Apply a compiled join-side alias to a replacement materialized input.""" + + if template.op != "alias": + return source + return QueryExpr(op="alias", inputs=(source,), params=template.params) + def _bind_direct_inputs( self, query: QueryExpr, diff --git a/src/agent_memory/runtime/executor.py b/src/agent_memory/runtime/executor.py index 3daf2ef..209624b 100644 --- a/src/agent_memory/runtime/executor.py +++ b/src/agent_memory/runtime/executor.py @@ -155,6 +155,13 @@ def allocate(node_id: str) -> str: if node.execution_kind == "deterministic": next_state = self._execute_deterministic(node_id, staged_node_state) + elif node.execution_kind == "relational_state": + next_state = self._execute_relational_state( + node_id, + old_state=old_state, + parent_updates=parent_updates, + staged_node_state=staged_node_state, + ) elif node.execution_kind == "semantic_row": next_state = self._execute_semantic_row( node_id, @@ -601,6 +608,37 @@ def _execute_deterministic( } return self.adapter.execute(node.query, inputs) + def _execute_relational_state( + self, + node_id: str, + *, + old_state: pd.DataFrame, + parent_updates: tuple[NodeOutputUpdate, ...], + staged_node_state: Mapping[str, pd.DataFrame], + ) -> pd.DataFrame: + """Incrementally maintain an inner join or fully recompute on retractions.""" + + node = self.policy.nodes[node_id] + if len(node.input_node_ids) != 2 or len(parent_updates) != 2: + raise RuntimeError("Relational state nodes require exactly two inputs") + if any(not update.retracted_rows.empty for update in parent_updates): + return self._execute_deterministic(node_id, staged_node_state) + if node.maintenance_query is None: + raise RuntimeError(f"Relational state node {node_id} has no maintenance query") + + inputs: dict[str, pd.DataFrame] = {node_id: old_state} + for parent_id, update in zip( + node.input_node_ids, + parent_updates, + strict=True, + ): + inputs[parent_id] = self._node_state.get( + parent_id, + self._empty_node_frame(parent_id), + ) + inputs[f"{parent_id}__inserted"] = update.inserted_rows + return self.adapter.execute(node.maintenance_query, inputs) + def _execute_semantic_row( self, node_id: str, diff --git a/tests/test_relational_join_ivm.py b/tests/test_relational_join_ivm.py new file mode 100644 index 0000000..c43c00c --- /dev/null +++ b/tests/test_relational_join_ivm.py @@ -0,0 +1,274 @@ +"""Incremental maintenance tests for deterministic relational inner joins.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Mapping +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.planner.differential_policy import DifferentialNode +from agent_memory.policy.logical import QueryExpr +from agent_memory.policy.relation import LOG_ADD_SEQ_COLUMN + + +class _BipartiteJoinMemory(am.Memory): + log = am.Log( + { + "side": "Join side.", + "key": "Join key.", + "value": "Row value.", + } + ) + _left = ( + log.filter(log.col("side") == "left") + .select(["key", "value"]) + .alias("left") + ) + _right = ( + log.filter(log.col("side") == "right") + .select(["key", "value"]) + .alias("right") + ) + pairs = _left.join(_right, on="key", how="inner") + + +class _ExplodedSelfJoinMemory(am.Memory): + log = am.Log({"items": "JSON array of items."}) + _items = ( + log.explode(column="items", output_col="item") + .select(["item"]) + .assign(key="all") + ) + _left = _items.alias("left") + _right = _items.alias("right") + pairs = _left.join(_right, on="key", how="inner") + + +class _TemporalSelfJoinMemory(am.Memory): + log = am.Log({"value": "Event value."}, system_columns=True) + _events = log.select([LOG_ADD_SEQ_COLUMN, "value"]) + _earlier = _events.alias("earlier") + _later = _events.alias("later") + pairs = _earlier.join( + _later, + on=_earlier.col(LOG_ADD_SEQ_COLUMN) < _later.col(LOG_ADD_SEQ_COLUMN), + how="inner", + ) + + +class _TemporalFilteredMemory(am.Memory): + log = am.Log({"value": "Event value."}, system_columns=True) + _events = log.select([LOG_ADD_SEQ_COLUMN, "value"]) + _earlier = _events.alias("earlier") + _later = _events.alias("later") + _pairs = _earlier.join( + _later, + on=_earlier.col(LOG_ADD_SEQ_COLUMN) < _later.col(LOG_ADD_SEQ_COLUMN), + how="inner", + ) + links = _pairs.sem_filter( + instruction="The earlier and later events should be linked." + ) + + +class _RetractingParentMemory(am.Memory): + log = am.Log({"key": "Group key.", "value": "Numeric value."}) + _minimum = log.group_by("key").min(column="value", output_col="value") + _left = _minimum.alias("left") + _right = _minimum.alias("right") + pairs = _left.join(_right, on="key", how="inner") + + +class _TrackingAdapter(LotusAdapter): + """Deterministic adapter that records root execution and accepts sem_filter rows.""" + + def __init__(self) -> None: + super().__init__() + self._depth = 0 + self.root_queries: list[QueryExpr] = [] + self.sem_filter_input_sizes: list[int] = [] + + def execute( + self, + query: QueryExpr, + inputs: Mapping[str, Any], + ) -> pd.DataFrame: + is_root = self._depth == 0 + if is_root: + self.root_queries.append(query) + self._depth += 1 + try: + if query.op == "sem_filter": + source = self.execute(query.inputs[0], inputs) + self.sem_filter_input_sizes.append(len(source)) + return source + return super().execute(query, inputs) + finally: + self._depth -= 1 + + +def _only_join_node(memory_type: type[am.Memory]) -> DifferentialNode: + nodes = [ + node + for node in memory_type.differentiate_policy().nodes.values() + if node.query.op == "join" + ] + assert len(nodes) == 1 + return nodes[0] + + +def _query_nodes(query: QueryExpr) -> list[QueryExpr]: + return [query, *(node for item in query.inputs for node in _query_nodes(item))] + + +def _bag(frame: pd.DataFrame) -> Counter[tuple[object, ...]]: + return Counter(tuple(row) for row in frame.itertuples(index=False, name=None)) + + +def _full_view(memory: am.Memory, view_name: str) -> pd.DataFrame: + query = memory.spec().views[view_name].query + return LotusAdapter().execute(query, {"log": memory._runtime._state["log"]}) + + +def test_compiler_builds_bag_preserving_inner_join_maintenance_query() -> None: + node = _only_join_node(_BipartiteJoinMemory) + + assert node.execution_kind == "relational_state" + assert node.maintenance_query is not None + nodes = _query_nodes(node.maintenance_query) + joins = [query for query in nodes if query.op == "join"] + concats = [query for query in nodes if query.op == "concat"] + materialized_names = { + str(query.params["name"]) + for query in nodes + if query.op == "materialized_view" + } + + assert len(joins) == 3 + assert len(concats) == 3 + assert all(query.params == node.query.params for query in joins) + assert all(query.inputs[0].params["name"] == "left" for query in joins) + assert all(query.inputs[1].params["name"] == "right" for query in joins) + assert node.node_id in materialized_names + for parent_id in node.input_node_ids: + assert parent_id in materialized_names + assert f"{parent_id}__inserted" in materialized_names + + +def test_inner_join_maintains_left_only_and_right_only_insertions() -> None: + memory = _BipartiteJoinMemory() + + memory.add({"side": "right", "key": "k", "value": "r1"}) + assert memory._runtime._state["pairs"].empty + + memory.add({"side": "left", "key": "k", "value": "l1"}) + memory.add({"side": "left", "key": "k", "value": "l2"}) + memory.add({"side": "right", "key": "k", "value": "r2"}) + + assert _bag(memory._runtime._state["pairs"]) == Counter( + { + ("k", "l1", "r1"): 1, + ("k", "l2", "r1"): 1, + ("k", "l1", "r2"): 1, + ("k", "l2", "r2"): 1, + } + ) + + +def test_self_join_multirow_delta_preserves_bag_multiplicity() -> None: + memory = _ExplodedSelfJoinMemory() + + memory.add({"items": '["a", "a"]'}) + assert _bag(memory._runtime._state["pairs"]) == Counter( + {("a", "all", "a"): 4} + ) + assert _bag(memory._runtime._state["pairs"]) == _bag(_full_view(memory, "pairs")) + + memory.add({"items": '["a"]'}) + assert _bag(memory._runtime._state["pairs"]) == Counter( + {("a", "all", "a"): 9} + ) + assert _bag(memory._runtime._state["pairs"]) == _bag(_full_view(memory, "pairs")) + + +def test_predicate_self_join_matches_full_recompute_after_each_append() -> None: + memory = _TemporalSelfJoinMemory() + + for value in ("a", "b", "c"): + memory.add({"value": value}) + assert _bag(memory._runtime._state["pairs"]) == _bag( + _full_view(memory, "pairs") + ) + + pairs = memory._runtime._state["pairs"] + assert list( + zip( + pairs[f"{LOG_ADD_SEQ_COLUMN}:earlier"], + pairs[f"{LOG_ADD_SEQ_COLUMN}:later"], + strict=True, + ) + ) == [(0, 1), (0, 2), (1, 2)] + + +def test_downstream_sem_filter_receives_only_new_join_pairs() -> None: + adapter = _TrackingAdapter() + memory = _TemporalFilteredMemory(adapter=adapter) + + for value in ("a", "b", "c"): + memory.add({"value": value}) + + assert adapter.sem_filter_input_sizes == [1, 2] + assert len(memory._runtime._state["links"]) == 3 + + +def test_parent_retraction_uses_full_join_recompute() -> None: + adapter = _TrackingAdapter() + memory = _RetractingParentMemory(adapter=adapter) + node = _only_join_node(_RetractingParentMemory) + + memory.add({"key": "k", "value": 5}) + adapter.root_queries.clear() + memory.add({"key": "k", "value": 3}) + + assert node.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} + ] + + +def test_non_inner_join_keeps_deterministic_full_recompute_path() -> None: + class OuterJoinMemory(am.Memory): + log = am.Log({"key": "Join key."}) + pairs = log.join(log, on="key", how="outer") + + node = _only_join_node(OuterJoinMemory) + + assert node.execution_kind == "deterministic" + assert node.maintenance_query is None + + +def test_inner_join_snapshot_round_trip_and_fingerprint_validation() -> None: + original = _TemporalSelfJoinMemory() + original.add({"value": "a"}) + original.add({"value": "b"}) + snapshot = original._runtime.snapshot_state() + + restored = _TemporalSelfJoinMemory() + restored._runtime.restore_state(snapshot) + original.add({"value": "c"}) + restored.add({"value": "c"}) + + assert _bag(restored._runtime._state["pairs"]) == _bag( + original._runtime._state["pairs"] + ) + + incompatible = dict(snapshot) + incompatible["plan_fingerprint"] = "pre-relational-state-plan" + with pytest.raises(ValueError, match="plan fingerprint"): + _TemporalSelfJoinMemory()._runtime.restore_state(incompatible) From 2766530f2e53f6963350de6380c357f07ecad8f7 Mon Sep 17 00:00:00 2001 From: Von Date: Mon, 31 Aug 2026 13:38:11 +0800 Subject: [PATCH 5/5] [fix-semantic] Preserve exact join key types - use original hashable composite values for exact semantic join partitions - prevent string coercion from merging distinct typed keys - cover mixed datetime/string values and Zep-style tuple endpoints Signed-off-by: Von --- src/agent_memory/adapters/lotus/sem_join.py | 22 +++++---- tests/test_semantic_pair_execution.py | 53 ++++++++++++++++++++- 2 files changed, 65 insertions(+), 10 deletions(-) diff --git a/src/agent_memory/adapters/lotus/sem_join.py b/src/agent_memory/adapters/lotus/sem_join.py index c281a91..92619e5 100644 --- a/src/agent_memory/adapters/lotus/sem_join.py +++ b/src/agent_memory/adapters/lotus/sem_join.py @@ -4,7 +4,6 @@ from collections import defaultdict from collections.abc import Callable, Mapping, Sequence -import json from typing import Any import pandas as pd @@ -431,7 +430,7 @@ def semantic_join_pair_candidates( if missing: raise ValueError(f"sem_join on columns not found on both sides: {missing}") - right_by_key: dict[str, list[Any]] = defaultdict(list) + right_by_key: dict[tuple[object, ...], list[Any]] = defaultdict(list) if on: assert right_frame is not None for right_id in right.index: @@ -471,13 +470,18 @@ def semantic_join_pair_candidates( ) -def _semantic_join_key(row: pd.Series, columns: Sequence[str]) -> str: - return json.dumps( - [row[column] for column in columns], - ensure_ascii=False, - default=str, - separators=(",", ":"), - ) +def _semantic_join_key( + row: pd.Series, + columns: Sequence[str], +) -> tuple[object, ...]: + key = tuple(row[column] for column in columns) + try: + hash(key) + except TypeError as error: + raise TypeError( + "sem_join exact-key columns must contain hashable scalar values" + ) from error + return key def aligned_join_values( diff --git a/tests/test_semantic_pair_execution.py b/tests/test_semantic_pair_execution.py index e7f6384..ee99ff3 100644 --- a/tests/test_semantic_pair_execution.py +++ b/tests/test_semantic_pair_execution.py @@ -3,6 +3,7 @@ from __future__ import annotations from dataclasses import dataclass, replace +from datetime import datetime import json from pathlib import Path @@ -24,7 +25,10 @@ from agent_memory.adapters.lotus.context import LotusExecutionConfig from agent_memory.adapters.lotus.sem_filter import execute_sem_filter from agent_memory.adapters.lotus.sem_groupby import execute_sem_groupby -from agent_memory.adapters.lotus.sem_join import execute_sem_join +from agent_memory.adapters.lotus.sem_join import ( + execute_sem_join, + semantic_join_pair_candidates, +) from agent_memory.policy.logical import ColumnSpec, QueryExpr from agent_memory.planner import PolicyDifferentiator from agent_memory.runtime import MemoryRuntime @@ -1074,6 +1078,53 @@ def test_exact_key_sem_join_rejects_lotus_cascade_before_input_execution() -> No ) +def test_exact_key_sem_join_does_not_match_stringified_values() -> None: + left = pd.Series(["left fact"], index=["left-row"]) + right = pd.Series(["right fact"], index=["right-row"]) + + candidates = semantic_join_pair_candidates( + left, + right, + left_frame=pd.DataFrame( + {"partition": [datetime(2026, 1, 1)]}, + index=left.index, + ), + right_frame=pd.DataFrame( + {"partition": ["2026-01-01 00:00:00"]}, + index=right.index, + ), + on=("partition",), + ) + + assert candidates.empty + + +def test_exact_key_sem_join_supports_zep_style_composite_tuple_keys() -> None: + left = pd.Series(["changed fact"], index=["left-row"]) + right = pd.Series(["current fact"], index=["right-row"]) + partition = { + "source_entity_id": (12, 3), + "target_entity_id": (7, 1), + } + + candidates = semantic_join_pair_candidates( + left, + right, + left_frame=pd.DataFrame([partition], index=left.index), + right_frame=pd.DataFrame([partition], index=right.index), + on=("source_entity_id", "target_entity_id"), + ) + + assert candidates.loc[:, [PAIR_LEFT_ID_COLUMN, PAIR_RIGHT_ID_COLUMN]].to_dict( + orient="records" + ) == [ + { + PAIR_LEFT_ID_COLUMN: "left-row", + PAIR_RIGHT_ID_COLUMN: "right-row", + } + ] + + def test_sem_groupby_search_filter_zero_candidates_skips_oracle( monkeypatch: pytest.MonkeyPatch, ) -> None: