diff --git a/README.md b/README.md index c3062d2..ef2601b 100644 --- a/README.md +++ b/README.md @@ -101,10 +101,9 @@ examples/helloworld/ ## Design Documents -- `docs/design/v0.0-interface.md`: current v0.0 interface boundary - `docs/design/operator-api.md`: DataFrame-style semantic operator API -- `docs/design/groupby-aggregation-api-rationale.md`: chain aggregation rationale -- `docs/design/TODO.md`: unresolved v0.0 -> v0.1 interface and research questions +- `docs/design/lotus-lowering.md`: LOTUS physical lowering and execution contracts +- `docs/design/TODO.md`: unresolved interface, runtime, and research questions - `docs/optimization/incremental-semantic-view-maintenance.tex`: semantic IVM theory draft ## Development diff --git a/docs/design/TODO.md b/docs/design/TODO.md index b56b3da..7c6f21c 100644 --- a/docs/design/TODO.md +++ b/docs/design/TODO.md @@ -1,8 +1,52 @@ # TODO: Semantic IVM Implementation Questions -Status: the current golden line is `docs/optimization/incremental-semantic-view-maintenance.tex` plus `docs/design/operator-api.md`. This file tracks unresolved implementation questions. It is intentionally question-first, not a decision spec. - -## V0.0 -> V0.1 Interface Contracts +Status: the current golden line is +`docs/optimization/incremental-semantic-view-maintenance.tex` plus +`docs/design/operator-api.md`. This is the single current backlog for unresolved +interface, operator, rule, runtime, and optimizer questions. It is intentionally +question-first, not a decision spec. + +## Current Consolidated Backlog + +### Operator execution + +- Audit LOTUS Cascade with real traces before exposing it as a supported physical + profile. Keep examples, raw outputs, explanations, helper decisions, and stats + in a trace side-channel rather than logical result columns. +- Add an indexed or hybrid execution path for `sem_topk` without changing its + logical ranking contract. +- Evaluate candidate pruning and indexed access for `sem_groupby`; pair batching + improves stability but does not reduce the number of semantic pairs by itself. +- Validate large-group structured `sem_agg` for context limits, retries, trace, + and cost before claiming that path is production-stable. +- Add `min_by` / `arg_min` only when a view needs the payload associated with a + minimum ordering value. Removing a current minimum still requires retraction + state or recomputation. + +### Differential rules and runtime + +- Define source deletion and general negative-delta semantics. +- Define old-state binding or an explicit full-recompute fallback for semantic + joins over arbitrary subqueries. +- Define correct maintenance for materialized left, right, and outer semantic + joins, where a new match can retract an old unmatched row. +- Support dependencies that mix base-log input with changed upstream views. +- Add recursive/fixpoint planning only with an explicit recursive-view contract. +- Add physical optimization such as indexed join/aggregate state and shared + arrangements without changing logical or maintenance fingerprints silently. +- Add spec-time placeholder validation and view-aware planner errors once column + scope inference is reliable. +- Define when a standalone semantic aggregate's current output is sufficient + maintenance state; schema compatibility alone is not a proof. + +### Runtime scheduling + +- Define concurrent and multi-writer ordering beyond the current single-process + append sequence. +- Keep maintenance batching, refresh gates, and background consolidation as + runtime scheduling decisions unless they change the logical input scope. + +## Interface Contracts - Shared finite dependency maintenance is represented by `DifferentiatedPolicy` and executed by `PolicyExecutor`; see @@ -52,28 +96,6 @@ Status: the current golden line is `docs/optimization/incremental-semantic-view- v0.1 should decide whether to hide the expression behind `_expr` or promote a stable debug/inspection API. -- The first stateful differential rewrite rule must decide how executable `ΔQ` - references current view state `V` and changed input rows `ΔD`. v0.0 keeps - `ΔD` as a mathematical/runtime concept rather than modeling it as a - `QueryExpr` operator; current row-local fragment rules work by binding the - source log to changed rows at runtime. - -- The first `sem_groupby(...).sem_agg(...)` view-boundary rule is implemented as - a Claude-style full-next-view query shape: - - ```text - ΔC = ΔD.sem_flat_map(...) - ΔT = ΔC.sem_groupby(...).sem_agg(...) - M = ΔT.sem_join(V, how="outer", instruction=) - V' = M.sem_map(instruction=) - ``` - - Current limits: it only applies at a public view boundary, returns full `V'` - rather than minimal `DeltaV`, and is not a generic exact `sem_agg` - maintenance proof. Deterministic placeholder rewrite is separate from future - semantic prompt rewriting. If a later design computes `DeltaV` instead, the - same logical maintenance may become `left join + upsert/delete/skip`. - - Adapter capability and default adapter injection are still unresolved. Before adding non-LOTUS engines, decide what capability contract each execution adapter exposes and where the default adapter is injected. @@ -83,28 +105,6 @@ Status: the current golden line is `docs/optimization/incremental-semantic-view- on the adapter dataclass or are delegated to LOTUS/provider-native configuration. - -## Claude Code native memory path alignment - -Detailed native Claude Code memory analysis now lives in -`docs/design/native-claude-memory-audit.zh.md`. Keep that document as the -canonical design note instead of duplicating long native-path analysis here. - -Current TODOs: - -- Keep native comparison experiments separated by path: main-agent inline write, - `extractMemories` background extraction, and `autoDream` consolidation. -- Do not use a CLI `-p` LOCOMO replay as evidence that `extractMemories` or - `autoDream` ran; it is only a native CLI / main-agent behavior audit unless - debug artifacts prove otherwise. -- Add an extract-only component harness before claiming parity with - `agent-memory` `sem_flat_map -> sem_groupby.sem_agg`. -- Treat `--turn-size 1` native CLI runs as one-row-per-turn behavior, not strict - one-message extractor tests. -- If comparing consolidation, prefer an extract-only + seeded `autoDream` gate - experiment so native `autoDream` itself runs while the scheduler conditions are - explicit. - ## Open Research Questions ### API context fields diff --git a/docs/design/groupby_agg.md b/docs/design/groupby_agg.md index fff756e..2b93b24 100644 --- a/docs/design/groupby_agg.md +++ b/docs/design/groupby_agg.md @@ -7,7 +7,8 @@ `input_cols` 细节。 - 第三部分记录 implementation-wise differential rules，说明实际 lowering 里 `sem_agg(input_cols=...)`、`collect_list(...)`、`flatten(...)`、`min(...)` - 和 `least(...)` 带来的 rule 变体。 + 和 `least(...)` 带来的 rule 变体，并单独说明普通 relational + `count/sum/avg` 的精确增量状态。 这里的 rules 分成 paper-wise 和 implementation-wise 两层。本文档不承诺当前 implementation 已经支持所有 rules，也不决定 optimizer 如何选择 rule。 @@ -114,6 +115,44 @@ input 输出零组。 keys 重名；重名会直接报错。作为 mixed `group_by(keys).agg(...)` 中的 aggregate spec 使用时，则遵循该 mixed `.agg(...)` 的 key precedence contract。 +#### `count`、`sum` 和 `avg` + +这三个是普通 relational aggregates，不调用 LLM： + +```python +overall = rows.agg( + count(output_col="row_count"), + sum(column="amount", output_col="total_amount"), + avg(column="amount", output_col="average_amount"), +) + +by_region = rows.group_by("region").agg( + count(output_col="row_count"), + sum(column="amount", output_col="total_amount"), + avg(column="amount", output_col="average_amount"), +) +``` + +它们的合同是： + +```text +Relation.agg(...): + many rows -> one output row + output schema = aggregate output columns + +group_by(K).agg(...): + many rows per exact key -> one output row per key + output schema = K + aggregate output columns +``` + +`count` 是 `COUNT(*)`，包括含 null 的 row。`sum` 和 `avg` 忽略 null；没有非 null +数值时输出 null。空的 global relation 仍输出一行，其中 count 为 0，sum/avg 为 +null；空的 grouped relation 输出零行。 + +当前 numeric aggregate family 只能包含 `count/sum/avg`，不能在同一次 `.agg(...)` +里混入 `sem_agg`、`array_agg`、`collect_list` 或 `min`。它只适用于普通 `Relation` +和 deterministic `group_by`；`sem_groupby(...).count/sum/avg(...)` 尚未实现。 + ### 1.3 `sem_groupby(keys).sem_agg(...)` 这里的 `keys` 指 `sem_groupby(input_cols=keys, ...)` 里的 semantic keys。它们是 @@ -261,6 +300,12 @@ Rules: - `am.min(column=..., output_col=...)` 或 `am.min(columns=[...], output_col=...)` 可以和 `sem_agg`、`array_agg`、 `collect_list` 一起作为 aggregate spec。 +- 普通 `Relation` 或 deterministic `group_by` 可以使用只包含 + `count/sum/avg` 的 numeric aggregate family。 +- Numeric aggregate family 暂时不能和上面的 semantic/evidence aggregate + family 混用，也不能用于 `sem_groupby`。 +- Numeric aggregate outputs 不能和 deterministic keys 重名；多个 outputs + 之间也不能重名。 例子： @@ -489,12 +534,42 @@ JoinGroups(GΔ, V, Kr, θg, M) = `group_match(θg)`是compiler从原`sem_groupby` instruction改写出的“changed group 和current group是否属于同一组”谓词。`Ks`只是这个semantic predicate读取的内容， -绝不能写成`on=Ks`：`on`只接受必须精确相等的`partition_by` keys `Kr`。 +绝不能写成`on=Ks`。`sem_join.on`本身支持general deterministic `JoinOn`；但在这条 +grouped-aggregate rule里，`on=Kr`有一个更窄、更明确的职责：只用必须精确相等的 +`partition_by` keys排除跨partition candidates。 `MergeByTarget`表示：matched rows按选中的current target合并；没有target的changed group形成新row；未被触及的current row原样保留。如果多个changed groups命中同一个 target，只生成一次合并后的target row。 +### 可复用的 direct `sem_join` 增量规则 + +`rule-join-map`会在自己的maintenance query内部调用`sem_join`，但direct +`sem_join` view的增量维护是一个独立、可复用的binary-state能力。对： + +```text +J = L.sem_join(R, on=φ, instruction=θ, how="inner", k=None) +``` + +append-only rule是： + +```text +ΔJ = SemJoin(ΔL, R0, φ, θ) + bag-union SemJoin(L0, ΔR, φ, θ) + bag-union SemJoin(ΔL, ΔR, φ, θ) + +J1 = J0 bag-union ΔJ +``` + +这里`φ`可以是`on=None`、同名exact keys或general deterministic `JoinOn`。每一项都 +先用`φ`缩小candidate pairs，再执行原semantic predicate `θ`；`ΔL × ΔR`只在第三项 +计算一次。实现使用保留重复行的`concat`，不是会去重的`union`。 + +当前direct differential只支持append-only、`how="inner"`和`k=None`。这不限制 +静态`sem_join`，也不改变本章的grouped-aggregate rules：`rule-join-map`仍然可以在 +其自身state-merge lowering内部使用outer或top-k semantic join；它不是把该内部 +join再次编译成一个独立的direct differential node。 + ### 2.1 `sem_groupby(...).sem_agg(...)` ```text @@ -661,6 +736,38 @@ semantic keys `Ks`仍然只参与`θg`，不要求精确相等。 `Kr`找到受影响的partition；只有这些partition内部的changed semantic groups需要继续 执行join-map或re-group。未被delta触及的partition原样保留。 +### 2.9 Exact relational `count/sum/avg` + +普通 relational aggregate 不需要 semantic rule family。它把每个 exact key 的 +输入压成可加减的状态： + +```text +S(D) = ( + row_count(D), + numeric_sum(D), + non_null_count(D) +) +``` + +设 `Δ+` 是 parent 插入，`Δ-` 是 parent 撤回： + +```text +S' = S + S(Δ+) - S(Δ-) + +count = row_count +sum = null if non_null_count = 0 else numeric_sum +avg = null if non_null_count = 0 else numeric_sum / non_null_count +``` + +有 `group_by(K)` 时，同一公式按 `K` 分开执行： + +```text +S'[k] = S[k] + S(Δ+[k]) - S(Δ-[k]) +``` + +说人话就是：新row属于哪个普通key，就只更新那个key的计数和总和；旧row被上游确定性 +变化撤回时，再从同一个key减掉。旧key之间不做join，也不重新扫描未变化的groups。 + ## 3. Implementation-wise differential rules paper-wise rules 里不写 `input_cols`。implementation 里必须写，因为 @@ -900,18 +1007,43 @@ rule-join-map: semantic outer join，再在结果上filter Kr”，因为那会丢失本应保留的unmatched rows， 改变outer-join语义。 +### 3.10 Exact relational aggregate state + +实现不会把 `avg` 当成一个可以直接相加的输出值。Planner为整组 +`count/sum/avg`建立一个隐藏的 aggregate-state node： + +```text +logical query: + D[.group_by(K)].agg(count, sum, avg) + +hidden state: + K, row_count, numeric_sum, non_null_count + +public finalizer: + K, declared count/sum/avg outputs +``` + +同一个 input column 被 `sum` 和 `avg` 同时使用时，只维护一份 sum 和 non-null +count。Runtime把 parent的`inserted_rows`加进state，把`retracted_rows`减出state； +grouped state的row count变成0时删除该group，global state则始终保留一行。隐藏列会进入 +checkpoint，但不会进入public view schema。 + +这里没有`rule-join-map`、`rule-re-group`或`rule-all-group`选择，因为key和状态更新都是 +确定性的。它也不改变`min`、`array_agg`或semantic aggregates现有的retraction边界。 + ## 4. 本文档不定义什么 本文档不决定： - optimizer 应该选择哪个 rule family； -- delete / negative delta 时保存 raw group、ordered state，还是 full recompute； +- `min`、`array_agg` 和 semantic aggregates 遇到 delete / negative delta 时保存 + raw group、ordered state，还是 full recompute； - DAG runtime 应该使用哪种 indexed physical aggregate state。 当前 implementation 已用真实 `QueryExpr` operator graph 实现本文第三章覆盖的 -`sem_agg`、`array_agg`、mixed `.agg(...)` 和 `min` lowering；不存在 public 或 -internal `agg_merge` operator。Shared `PolicyExecutor` 会在这些 stateful node -boundary 对 insert-only change 执行本章定义的 maintenance query。如果 parent -change 含 internal retraction，executor 会在 parent 的完整 next state 上重算当前 -semantic aggregate node，再把 replacement change 传给下游。Source delete 和更 -高效的 indexed aggregate state 仍未实现。 +`sem_agg`、`array_agg`、mixed `.agg(...)`、`min` 和 numeric `count/sum/avg` +lowering；不存在 public `aggregate_state` 或 `agg_merge` operator。Numeric +aggregates直接维护hidden additive state，可以消费parent产生的insertions和 +retractions。现有semantic aggregate node遇到parent retraction时，仍会在parent完整 +next state上重算，再把replacement change传给下游。Source delete API、numeric与 +semantic aggregate混合，以及`sem_groupby`后的numeric aggregates仍未实现。 diff --git a/docs/design/next-step.md b/docs/design/lotus-lowering.md similarity index 76% rename from docs/design/next-step.md rename to docs/design/lotus-lowering.md index b6eed64..2fbb887 100644 --- a/docs/design/next-step.md +++ b/docs/design/lotus-lowering.md @@ -1,8 +1,7 @@ -# Next-Step 设计：LOTUS Lowering、Storage、Claude Consolidation +# LOTUS Lowering 设计 -本文记录 `agent-memory` 近期下一步的工程设计边界。当前主要包括三件事： -语义 operator 如何接 LOTUS、storage backend 如何保持独立、Claude-style -topic view 如何对齐 markdown memory 文件。 +本文记录 `agent-memory` 如何把逻辑 `QueryExpr` lower 到 LOTUS-backed +operator execution。它是一份稳定的执行层设计说明，不是实施路线图。 其中 LOTUS 部分讨论的是 `QueryExpr -> LOTUS-backed execution`，不是 `Q -> Q'` 的 differential rewrite。 @@ -21,28 +20,7 @@ topic view 如何对齐 markdown memory 文件。 topology、change/state、共享 node 和原子提交；详细边界见 `policy-differentiation-dataflow-runtime.zh.md`。 -## 1. Implementation Roadmap - -近期实现顺序必须先补齐 operator execution，再谈 rules、Claude policy 和 -storage。推荐里程碑如下： - -1. **Complete `QueryExpr -> LOTUS-backed execution`**：先让当前 public - operators 都能从 `QueryExpr` 执行到 dataframe result。没有完整 execution - layer，differential rules 和 Claude policy 都没有可靠执行目标。 -2. **Implement `Q -> Q'` differential rules**：在 full-query operator - execution 完整后，再扩展 `QueryDifferentiator`。Rules 只负责 query - rewrite，不负责 backend lowering。 -3. **Claude policy full-query and differential comparison**：用同一批 log - rows 比较 `Q(D ∪ ΔD)` 和 differential maintenance 的结果，验证 Claude - policy 的语义和 `Q -> Q'` 是否对齐。 -4. **StatementSet + storage connectors**：最后接 durable materialization。 - Storage 持久化 materialized relations，不定义 logical query 语义，也不替代 - operator/rule correctness；Markdown 只是未来的一个 connector/profile。 - -这个顺序避免把四个问题混在一起：operator 是否能执行、`Q -> Q'` 是否正确、 -Claude policy 是否合理、view 是否能持久化。 - -## 2. Operator Implementation Strategy +## 1. Operator Implementation Strategy Operator 实现要分清两类： @@ -95,35 +73,7 @@ Public `Relation` API 只保存 logical query 参数到 `QueryExpr.params`。例 和 `how`，`sem_topk` 保存 instruction 和 `k`。LOTUS method、cascade、examples、 safe mode、stats、trace 等不进入 policy class，也不进入 query tree。 -## 3. Rule / Storage Sequencing - -Differential rules、Claude policy validation 和 storage 要保持顺序和分层： - -- Storage 负责持久化 materialized views，不定义 logical query。 -- `QueryDifferentiator` 负责 `Q -> Q'`，不关心 markdown 文件怎么写。 -- `ExecutionAdapter` 负责执行 `QueryExpr`，不负责 durable persistence。 -- `QueryDifferentiator` 必须在 operator execution 完整后再扩展，否则 - rewrite 出来的 `Q'` 没有完整 execution target。 -- Claude policy validation 必须在 operator execution 和 rules 之后；它是 - full-query 与 differential result 的对照实验，不是 operator layer 的替代。 -- Storage 排在 Claude operator/rule correctness 之后；它只决定 durable - materialization，不决定 logical semantics。 - -在 `apply_delta` / upsert 语义出现前，Claude-style stateful differential 可以 -先计算下一版 full view： - -```text -ΔC = ΔD.sem_flat_map(...) -ΔT = ΔC.sem_groupby(...).sem_agg(...) -M = ΔT.sem_join(V, how="outer", instruction=) -V' = M.sem_map(instruction=) -``` - -后续如果设计出 `ΔV + apply_delta`，可以把这条路径替换成 -`left join + upsert/delete/skip`，前提是 storage contract 明确支持这些 -应用语义。 - -## 4. 分层边界 +## 2. 分层边界 推荐保持如下分层： @@ -149,7 +99,7 @@ differentiated maintenance query `Q'`。它不应该知道 LOTUS 的 pandas acce 细节。`LotusAdapter` 负责执行已经生成好的 `QueryExpr`，它不应该修改 operator 的逻辑语义。 -## 5. 不直接改 LOTUS +## 3. 不直接改 LOTUS 不要直接改 LOTUS 源码，也不要覆盖或 monkeypatch `df.sem_map`、 `df.sem_filter` 等 pandas accessor。 @@ -164,7 +114,7 @@ operator 的逻辑语义。 正确做法是：在 `agent-memory` 里实现自己的 lowering，底层复用 LOTUS 的 LM、cache、templates、pandas accessors、top-k/filter/join 等能力。 -## 6. 推荐模块结构 +## 4. 推荐模块结构 `LotusAdapter` 不应该膨胀成一个包含所有 operator 细节的大文件。LOTUS lowering 使用专门的 adapter package： @@ -197,7 +147,7 @@ src/agent_memory/adapters/lotus/ 这个 package 不是新的 public adapter surface；public 入口仍然是 `LotusAdapter`。 -## 7. Lowering Matrix +## 5. Lowering Matrix | agent_memory op | LOTUS 支持情况 | 状态 | lowering | |---|---|---| @@ -216,7 +166,7 @@ src/agent_memory/adapters/lotus/ | `sem_agg` 多输出 | 无精确等价 | implemented | agent-memory structured aggregate lowering | | `filter(...)` / `assign(...)` / predicate `join(...)` | pandas 本地实现 | implemented | minimal relation-bound expression subset via `relation.col(...)` | -## 8. `sem_map` +## 6. `sem_map` LOTUS 原生 `sem_map` 是： @@ -247,7 +197,7 @@ option 注入结果列。 这不是自动 lower 到 `sem_extract`。`sem_extract` 的语义是抽取字段； multi-output `sem_map` 的语义仍然是按 instruction 生成 / 改写字段。 -## 9. `sem_flat_map` +## 7. `sem_flat_map` `sem_flat_map` 的语义是： @@ -277,7 +227,7 @@ input DataFrame 但 operator contract 属于 `agent-memory`，不是 LOTUS 原生 accessor。 -## 10. `sem_join(how=...)` +## 8. `sem_join(how=...)` `agent-memory sem_join` 支持 DataFrame-style `how`： @@ -294,7 +244,7 @@ left.sem_join(right, instruction=..., how="outer") - `inner` 可以 lower 到 LOTUS native `df.sem_join(...)`。 - `left/right/outer` 不能宣称是 LOTUS native。 -推荐 future lowering： +当前 lowering： ```text inner: @@ -317,7 +267,7 @@ outer: 注意：这只是 join shape 的补全。semantic predicate 本身仍由 LOTUS inner join 执行。unmatched rows 的补全是 deterministic pandas work。 -## 11. `sem_groupby -> sem_agg` +## 9. `sem_groupby -> sem_agg` `agent-memory sem_groupby` 是 semantic partition / assignment，不是 pandas exact groupby。 @@ -331,7 +281,7 @@ exact groupby。 LOTUS 的 `sem_agg(group_by=...)` 只能在已有 exact group key 或 group id 时 做分组聚合。它不能替代 semantic grouping 本身。 -因此 future lowering 应该是两步： +因此当前 lowering 分成两步： ```text rows @@ -350,7 +300,7 @@ rows 但这些都不是 `agent-memory sem_groupby` 的直接等价物。lowering 必须明确 说明选择了哪一种物理策略。 -## 12. `sem_agg` +## 10. `sem_agg` `sem_agg` 有两个 receiver scope： @@ -380,22 +330,7 @@ partition ids 的 tree fold。因此： 原生 `sem_agg`。Structured aggregate 的 schema/prompt contract 是 `agent-memory` operator lowering 的显式部分。 -## 13. Storage Backend and Claude Markdown Materialization - -本节早期提出的 `StorageBackend.load_state/save_view` 草案已被 -[`zep-graphiti-storage-retrieval-benchmark.zh.md`](./zep-graphiti-storage-retrieval-benchmark.zh.md) -中的通用 storage contract 取代，不应实现成第二套平行接口。 - -当前统一边界是：policy writer 通过 immutable `StatementSet` 声明 relation sink； -`StorageDeployment` 绑定 connector 和 namespace；planner 将 sink roots 编译进共享 DAG； -runtime 把该 DAG 产生的 inserted/retracted rows 事务性写入 connector。Storage 仍然不定义 -logical query 语义，也不进入 `LotusAdapter`。 - -未来 Claude Markdown materialization 应实现为同一 contract 下的 connector/profile， -而不是恢复这里的 `load_state/save_view` 接口。Topic files、`MEMORY.md` layout 和 path -derivation 都属于该 physical profile，不属于 `MemoryView`。 - -## 14. 工程原则 +## 11. 工程原则 Adapter lowering 可以做： diff --git a/docs/design/operator-api.md b/docs/design/operator-api.md index 6411174..6d17492 100644 --- a/docs/design/operator-api.md +++ b/docs/design/operator-api.md @@ -100,6 +100,103 @@ episodes = log.assign( field such as `reference_time` or `valid_at` when real-world temporal order matters. +### Count-triggered refresh + +Memory maintenance is eager by default: each `add(...)` publishes one source +row through the differential plan before returning. A caller may instead make +the refresh boundary explicit: + +```python +memory = MyMemory(refresh=am.CountRefresh(every=32)) + +memory.add(event) # accepts and buffers one normalized source row +memory.pending_count # rows accepted but not yet published +memory.flush() # publishes the current tail as one delta relation +``` + +At the 32nd pending row, `add(...)` synchronously applies those rows as one +atomic relation-valued delta. If maintenance fails, the published view is +unchanged and the same normalized pending rows remain available for retry. +`flush()` is a no-op when nothing is pending. + +`query(...)` reads the most recently published view and never flushes +implicitly. Between successful refreshes, `every=32` may therefore omit up to +31 accepted rows; a failed refresh may leave a complete batch pending for +retry. This is a refresh-scheduling choice, not semantic prompt batching: +operators still choose their own prompt construction and provider dispatch. + +Omitting `refresh`, or using `CountRefresh(every=1)`, preserves eager execution +and the existing checkpoint format. `every > 1` uses a separate checkpoint +contract that records both the refresh configuration and normalized pending +rows. The first implementation is synchronous and count-triggered only; it +does not provide a timer, background worker, query-triggered refresh, or a +crash-durable pending log. + +### Prompt batching + +Prompt batching is an opt-in LOTUS execution setting. It does not decide when +source rows enter the dataflow; it decides how one semantic operator sends the +independent tasks that are ready in its current invocation: + +```python +from agent_memory.adapters.lotus import LotusAdapter, PromptBatching +from agent_memory.adapters.lotus.context import LotusExecutionConfig + +adapter = LotusAdapter( + config=LotusExecutionConfig( + prompt_batching=PromptBatching(), + ) +) +memory = MyMemory(adapter=adapter, refresh=am.CountRefresh(every=32)) +``` + +The two settings are deliberately separate: + +```text +CountRefresh(every=32) + -> up to 32 source rows enter the maintenance DAG as one relation-valued delta + +PromptBatching() + -> each semantic operator puts all of its currently ready independent tasks + into one prompt when the model context permits +``` + +`PromptBatching(max_tasks=M)` caps a prompt at `M` tasks. `max_tasks=None` +means no artificial task-count cap, not one task per prompt. The runner still +splits deterministically when a complete prompt would exceed the model context; +it never truncates a task to make it fit. + +The shared setting applies to the task shapes exposed by `sem_filter`, +`sem_map`, `sem_flat_map`, pairwise `sem_join`, multi-anchor top-k `sem_join`, +pairwise `sem_groupby`, and `sem_agg`. Direct listwise `sem_topk` already asks +one ranking question over its input; delegated LOTUS pairwise ranking keeps its +algorithm-specific comparison schedule. Search-Filter remains an earlier +candidate-generation step, so only selected pairs are prompt-batched. + +Prompt batching is not provider request batching. A packed prompt asks one +model request to return several task results. Provider batching sends several +separate prompts together while preserving each prompt byte-for-byte. The +existing `sem_agg_dispatch="provider-batched"` path is therefore mutually +exclusive with `prompt_batching`. + +Every packed output must contain every stable task ID exactly once and satisfy +the operator's normal output schema. The parser may repair recorded, +syntax-only JSON mistakes, but it rejects missing, duplicate, unknown, or +semantically invalid results. Enabling prompt batching changes the physical +execution and checkpoint fingerprint; omitting it preserves the existing +prompts and behavior. + +Prompt batching is not a semantics-preserving transport optimization. Putting +the same tasks together changes the context seen by the model, so valid +structured output can still contain different semantic decisions. A +32-event Claude/Mem0/Zep smoke completed all tested batch sizes but produced +different materialized states; larger Claude batches omitted durable topics, +and the largest Zep batch made an incorrect entity merge. There is therefore +no empirically supported system-wide `max_tasks` value from this smoke. Keep +this setting opt-in and validate it against each policy's state and end-to-end +quality before using it as a production default. The evidence is summarized in +`next-plan-sem-join-and-count-refresh.zh.md`. + ### `select` Projection over columns. @@ -140,7 +237,21 @@ rows = rows.assign( Use `assign` for cheap, deterministic fields. Use `sem_map` when the new fields require semantic interpretation. Assignment values are scalar literals, column -expressions, row-wise `array_cat`, or `am.least(...)` expressions. +expressions, arithmetic expressions, row-wise `array_cat`, or `am.least(...)` +expressions. + +Arithmetic uses the usual Python operators over relation columns: + +```python +metrics = counts.assign( + ratio=counts.col("positive_count") / counts.col("total_count"), + age=2026 - counts.col("year"), +) +``` + +`+`, `-`, `*`, and `/` accept numeric operands but reject booleans and other +non-numeric values. A null operand produces null. Division by zero raises +`ZeroDivisionError` instead of silently returning infinity. `am.least(a, b, ...)` is a scalar expression, not an aggregate. It compares two or more values in the same row, ignores null operands, and returns null only @@ -356,6 +467,55 @@ group, but it is not a permanent UUID: full recomputation may change extraction order, and future group merge/split support needs downstream remapping. A storage backend may map this logical tuple to a physical UUID later. +### `count`, `sum`, and `avg` + +These are exact relational aggregates. They do not call an LLM and do not use +semantic grouping. + +```python +overall = rows.agg( + am.count(output_col="row_count"), + am.sum(column="amount", output_col="total_amount"), + am.avg(column="amount", output_col="average_amount"), +) + +by_region = rows.group_by("region").agg( + am.count(output_col="row_count"), + am.sum(column="amount", output_col="total_amount"), + am.avg(column="amount", output_col="average_amount"), +) +``` + +Convenience methods build the same query when only one aggregate is needed: + +```python +rows.count(output_col="row_count") +rows.sum(column="amount", output_col="total_amount") +rows.avg(column="amount", output_col="average_amount") + +rows.group_by("region").count(output_col="row_count") +``` + +The behavior follows a small SQL-style contract: + +- `count` is `COUNT(*)`: it counts every row, even when `amount` is null. +- `sum` and `avg` ignore null input values. If no non-null value exists, their + result is null. +- `sum` and `avg` require numeric non-null values and reject booleans. +- An empty global input returns one row: count is zero and sum/avg are null. +- An empty grouped input returns no groups. Null grouping keys form an ordinary + group when rows with null keys are present. +- Aggregate output names must be unique. A grouped numeric aggregate output may + not reuse a grouping-key name. + +`Relation.agg(...)` accepts only `count`, `sum`, and `avg`. Ordinary +`group_by(...).agg(...)` also accepts this numeric family, but it cannot mix it +with `sem_agg`, `array_agg`, `collect_list`, or `min` in the same call yet. +`sem_groupby(...).count/sum/avg(...)` is not supported. The current +semantic-grouped aggregate contract requires `sem_agg` to produce the semantic +key columns, and the numeric aggregate lowering does not consume semantic group +assignments. + ### `array_cat` Deterministic combine operator for one JSON array aggregate-state column. @@ -882,10 +1042,23 @@ Supported aggregate specs: column. - `am.min(columns=[...], output_col=...)`: deterministic lexicographic minimum over complete grouped tuples. +- `am.count(output_col=...)`: exact `COUNT(*)` for an ordinary relation or + ordinary `group_by`. +- `am.sum(column=..., output_col=...)`: exact numeric sum for an ordinary + relation or ordinary `group_by`. +- `am.avg(column=..., output_col=...)`: exact numeric average for an ordinary + relation or ordinary `group_by`. + +This is the union of supported aggregate families, not permission to combine +every item. A numeric `count/sum/avg` call must contain only numeric algebraic +aggregate specs. Existing semantic/evidence aggregation may combine +`sem_agg`, `array_agg`, `collect_list`, and `min` under its documented rules. +Numeric algebraic aggregates are not accepted after `sem_groupby`. For deterministic `group_by(K).agg(...)`, output columns are `K + outputs(A*)`. -If aggregate output names overlap deterministic keys, the deterministic key -column wins. +For the existing semantic/evidence aggregate family, a deterministic key wins +when an aggregate output uses the same name. The numeric algebraic family is +stricter and rejects that collision while authoring the query. For semantic `sem_groupby(..., partition_by=P).agg(...)`, output columns are `P + outputs(A*)`. At least one `sem_agg(...)` spec is required, and the @@ -996,8 +1169,15 @@ The canonical documentation form is logical join condition, deterministic keys, and match cardinality in one DataFrame operator. -`on` is an optional exact equality condition evaluated before the semantic -predicate: +`on` is an optional deterministic candidate condition evaluated before the +semantic predicate. It accepts the same `JoinOn` forms as ordinary `join(...)`: + +- one same-named key column; +- a sequence of same-named key columns; +- one relation-bound boolean expression; +- a sequence of relation-bound boolean expressions, interpreted as conjunction. + +Exact-key example: ```python changed_facts.sem_join( @@ -1012,6 +1192,30 @@ 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. +General deterministic-predicate example: + +```python +left = facts.alias("left") +right = facts.alias("right") + +pairs = left.sem_join( + right, + on=[ + left.col("tenant_id") == right.col("tenant_id"), + left.col("valid_at") < right.col("valid_at"), + ], + instruction="{fact:left} and {fact:right} describe the same fact.", + how="inner", +) +``` + +Here, the adapter first keeps pairs from the same tenant whose left timestamp is +earlier. Only those pairs reach embedding-based candidate generation or LLM +verification. `on=None` means that no deterministic restriction is applied, so +the semantic predicate starts from the full left-by-right pair domain. Key +column names and predicate expressions cannot be mixed in the same `on` +sequence. + `k` optionally limits the number of matched right rows for each left row: ```text @@ -1235,6 +1439,40 @@ V_prime = V.union( ) ``` +Exact relational `count`, `sum`, and `avg` use hidden additive state. For one +numeric input column, the runtime keeps only: + +```text +row_count +numeric_sum +non_null_count +``` + +For inserted rows these values are added; for retracted parent rows they are +subtracted. The public result is then finalized as: + +```text +count = row_count +sum = null if non_null_count == 0 else numeric_sum +avg = null if non_null_count == 0 else numeric_sum / non_null_count +``` + +The same update is applied independently per exact `group_by` key. A grouped +state row disappears when its row count reaches zero; the global aggregate +continues to expose its one SQL-style output row. Hidden state columns are an +implementation detail, are saved in checkpoints, and never appear in the +declared view schema. + +This is an exact differential rule for the parent insertions and retractions +that reach the aggregate node. It does not call an LLM and does not select a +`rule-join-map` or `rule-re-group` strategy. + +`rule-join-map` is the default differential strategy for semantic grouped +aggregates. Other supported strategies, including `rule-re-group` and +`rule-all-group`, remain available when a policy or an ablation selects them +explicitly. This default changes the maintenance plan, not the meaning of the +static `sem_groupby(...).agg(...)` query. + 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. @@ -1281,18 +1519,42 @@ 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: +Direct append-only semantic inner join follows the usual bag-preserving +relational delta shape. Let `L0`, `R0`, and `J0` be the committed left, right, +and join states before one update, and let `delta_L` and `delta_R` be the newly +inserted rows: ```python -V = L.sem_join(R, instruction=instruction, how="inner") +J = L.sem_join(R, instruction=instruction, how="inner") -V_prime = ( - delta_L.sem_join(R, instruction=instruction, how="inner") - .union(L.sem_join(delta_R, instruction=instruction, how="inner")) - .union(delta_L.sem_join(delta_R, instruction=instruction, how="inner")) +delta_J = ( + delta_L.sem_join(R0, instruction=instruction, how="inner") + .concat(L0.sem_join(delta_R, instruction=instruction, how="inner")) + .concat(delta_L.sem_join(delta_R, instruction=instruction, how="inner")) ) + +J1 = J0.concat(delta_J) ``` +The same `instruction`, `on`, and semantic execution profile are used in all +three terms. Their pair domains do not overlap, so pairs between `delta_L` and +`delta_R` are evaluated exactly once. `concat` is intentional: join output is a +bag, and exact duplicate rows must retain their multiplicity. + +The current direct `sem_join` differential implementation is deliberately +narrow: + +```text +supported: append-only parent changes, how="inner", k=None +not supported: parent retractions, left/right/outer maintenance, k joins +``` + +An unsupported `how` or `k` is rejected while compiling the differential plan. +A parent retraction is rejected before applying the append-only maintenance +formula. These limits apply only to differential maintenance; static +`sem_join(...)` execution continues to support the documented join types and +`k` values. + For updates, the exact output can be expressed with ordinary set operations: ```python diff --git a/docs/design/operator-todo.md b/docs/design/operator-todo.md deleted file mode 100644 index 82f16fb..0000000 --- a/docs/design/operator-todo.md +++ /dev/null @@ -1,417 +0,0 @@ -# Operator TODO：LOTUS 对齐与优化清单 - -本文只记录 `QueryExpr -> LOTUS-backed execution` 这一层的 operator 后续工作。 -它不讨论 `Q -> ΔQ` differential rules，不讨论 Claude policy，不讨论 storage。 - -当前状态可以概括为： - -- 主要 operator 已经有可运行的 baseline lowering。 -- LOTUS 原生等价的路径优先直接调用 LOTUS。 -- 语义不等价的路径由 `agent-memory` 自己实现 lowering，并尽量复用 LOTUS 的 - LM、prompt formatter、cache、postprocess pattern 和已有 semantic operator。 -- 还没有达到完整 LOTUS parity，尤其是 cascade、strategy、stats、audit trace、 - embedding/index pruning、large-group aggregation 等执行能力。 -- Public policy API 保持只表达逻辑语义；backend、LOTUS、优化策略、debug - 输出都属于 runtime / adapter config，不属于 `MyMemory` class 里的 view query。 - -## 1. 原则 - -- 先保证 operator 语义正确，再谈 cost / latency 优化。 -- 能精确复用 LOTUS 原生 operator 的地方，直接复用。 -- 不能精确复用的地方，不 monkeypatch、不假装等价；在 `agent-memory` 里定义 - lowering contract。 -- 优化能力如果来自 LOTUS，应该通过 runtime / adapter config 显式配置，而不是 - 放进 policy author 的 public operator API。 -- `QueryExpr.params` 只记录 logical query 参数，例如 instruction、columns、join - type、top-k 的 `k`；不要把 backend execution knobs 写进 query tree。 -- 每补一个真实执行能力，都要有 offline unit test 和 gated real LOTUS audit。 - -## 2. 当前 Operator 状态 - -| operator | 当前状态 | LOTUS 对齐程度 | 后续重点 | -|---|---|---|---| -| `select` | 可用 | pandas 本地实现 | 暂无 | -| `concat` / `union` / `subtract` / `drop_duplicates` | 可用 | pandas 本地实现 | 暂无 | -| `filter(predicate)` / `assign(...)` / predicate `join(...)` | 可用 | pandas 本地实现 | 后续只按真实需求扩 deterministic expression subset | -| `sem_filter` | 可用 | native LOTUS path，基础 adapter config 已接入 | 补 side-channel stats / cascade real audit | -| `sem_map` 单输出 | 可用 | native LOTUS `sem_map` | adapter config 对齐 execution options | -| `sem_map` 多输出 | 可用 | custom structured lowering | 对齐 LOTUS execution ergonomics，不把 debug 写入结果列 | -| `sem_flat_map` | 可用 | custom structured lowering | 后续考虑 batching / retry / cost controls | -| `sem_join inner` | 可用 | lower-level LOTUS `sem_join` | 补真实 cascade audit | -| `sem_join left/right/outer` | 可用 | inner + 本地 unmatched rows | 继续验证 column shape / metadata | -| `sem_groupby` | baseline 可用 | custom pairwise lowering | 加 candidate pruning / indexing | -| `sem_agg` 单输出 | 可用 | lower-level LOTUS `sem_agg` | adapter config 接入/验证 aggregation controls | -| `sem_agg` 多输出 | 可用 | agent-memory structured hierarchical lowering | large-group audit；等 PyPI LOTUS 支持 native response_format 后可重新评估 | -| `sem_topk` | 可用 | native LOTUS `sem_topk(...)`，adapter config 控制 execution method | 继续补 hybrid / index path | - -Multi-output `sem_map` 未来可以有两条 optimizer-selectable lowering： - -- structured one-call：一次 structured generation 同时写多个 output columns，成本低， - 字段一致性更好，但不是 LOTUS native `df.sem_map` 的原始 one-column shape。 -- native per-column：每个 output column 调一次 LOTUS native `df.sem_map`，可以更 - 直接复用 LOTUS 原生 execution options，但成本变成 N 倍，字段之间也可能不一致。 - -`sem_flat_map` 当前要求每个 input row 返回 `{"rows": [...]}` wrapper，以提高 -JSON-mode 稳定性，例如： - -```json -{ - "rows": [ - {"topic": "support group", "summary": "Caroline attended an LGBTQ support group."}, - {"topic": "self acceptance", "summary": "The group helped Caroline feel accepted."} - ] -} -``` - -这属于 lowering format，不改变 `sem_flat_map` 的 logical 语义：一行输入仍然 -产生 zero or more output rows。可选 `ordinal_col` 由 executor 在解析成功后按每个 -input row 从 0 编号，不进入 prompt。 - -## 3. `sem_filter` TODO - -当前 `sem_filter` 已经直接使用 LOTUS，所以 LOTUS 的核心执行能力在底层是有的。 -问题不是 policy API 缺参数。当前基础 adapter config 已经能接入常用 LOTUS -execution options，但 stats / raw / explanation 的 side-channel，以及 cascade -真实 audit 还需要继续做。 - -LOTUS 已有但当前没有完整暴露的能力包括： - -- `examples`：few-shot examples，提高 predicate 判断稳定性。 -- `strategy`：例如 COT / ZS_COT。 -- raw outputs / explanations：作为 adapter trace、audit log 或 stats side-channel， - 不默认写入 result DataFrame。 -- `default`：模型输出无法 parse 时的默认 boolean。 -- `safe_mode`：估算成本。 -- `cascade_args` / `helper_examples` / `return_stats`：使用 helper model 或 - embedding proxy 做 cascade filtering，减少大模型调用。 -- `return_all` / `suffix`：LOTUS dataframe shape/debug 选项。当前不进入 - logical result，因为它们会改变输出 shape 或引入 execution/debug columns。 - -这里不是 LOTUS 没实现，也不是要把这些 knobs 加到 `Relation.sem_filter(...)`。 -这些是 execution policy，应进入 adapter/runtime config。当前实现默认保持 -`return_raw_outputs=False`、`return_explanations=False`、`return_stats=False`， -避免 debug 信息混入 logical result DataFrame。 - -待办： - -- 设计 stats / raw / explanation side-channel，不要写入 result DataFrame。 -- 增加 real audit：普通 filter、examples、可选 cascade。 -- 如果后续 runtime 有统一 trace object，把 `sem_filter` 的 parse failure、 - helper-model decisions 和 stats 统一接进去。 - -## 4. `sem_topk` TODO - -当前 `sem_topk` 不应该在 public API 里暴露 `method` 等 backend 参数。Policy -writer 只写： - -```python -df.sem_topk(instruction, k) -``` - -adapter 再根据 execution config 调用 LOTUS： - -```python -df.sem_topk( - instruction, - K=k, - method=method, - strategy=strategy, - cascade_threshold=cascade_threshold, - return_stats=return_stats, - safe_mode=safe_mode, - return_explanations=return_explanations, -) -``` - -LOTUS 支持不同 ranking 方法和执行选项： - -- `method="naive"`：全 pairwise，慢但最直接，适合早期 correctness audit。 -- `method="quick"`：成本较低，但排序路径更启发式。 -- `method="heap"`：heap-based top-k。 -- `method="quick-sem"`：使用 semantic index / embedding 辅助优化。 -- `strategy`：COT / ZS_COT。 -- `cascade_threshold`：小模型 / 大模型 cascade。 -- `return_stats`：返回调用次数、token、排序解释等统计。 -- `return_explanations`：返回解释；后续应进入 trace / stats side-channel，而不是 - logical result DataFrame。 - -当前默认用 `naive`，因为它更接近“不省钱但最直观”的 gold execution。 -`quick`、`heap`、`quick-sem` 和后续 hybrid search 都是可选执行策略。 - -待办： - -- 真实 audit 继续比较 `naive` 和 `quick` 的语义差异、成本差异。 -- 未来设计 agent-memory 自己的 hybrid search，不强行绑定 LOTUS 的单一 - `method`。 -- `quick-sem` 需要 semantic index / retrieval model / vector store 配置；不要在 - 没有 backend config 的情况下假装已完整支持。 -- 如果 `return_stats=True`，不要把 stats 混进 logical result DataFrame；后续应走 - adapter trace / audit side-channel 或 runtime stats object。 - -## 5. `sem_groupby` TODO - -当前 `sem_groupby` 是 baseline： - -```text -exact duplicate input-column collapse --> unique input-column rows pairwise semantic comparison --> union-find group assignment -``` - -这个语义清楚、容易审计，但 pairwise 比较是 `O(N^2)`。数据量变大时，不能把所有 -row pair 都交给 LLM。 - -Public API 使用 `input_cols`，不是 relational `key`。这些 columns 是判断 group -membership 的 evidence；它们不是 exact equality key，也不一定直接成为最终输出 -列。最终 canonical group fields 仍由后续 `sem_agg(...)` 产生。 - -后续优化思路是先做 candidate pruning，再用 LLM predicate 精判： - -```text -rows --> embedding / index / cluster / sim-join 找候选 pairs --> LLM 判断候选 pairs 是否同组 --> union-find 生成 group ids -``` - -可复用的 LOTUS 工具包括： - -- `sem_cluster_by`：候选粗分组，但它不等价于 agent-memory 的 semantic - grouping contract。 -- `sem_sim_join` / vector search / semantic index：找高相似候选 pair。 -- `sem_filter` / `sem_join`：对候选 pair 做 LLM predicate 精判。 -- cascade：先用便宜 proxy，难例交给大模型。 - -当前没有配置 embedding model / vector index backend 时，不应该默认启用这些 -pruning path。先保留清晰的 pairwise baseline 和未来 candidate-generator -接口位置；等 backend config 明确后，再接 `sem_sim_join`、semantic index 或 -cluster-based pruning。 - -另一个 mode 是 closed-world labeled grouping。Open-world grouping 动态发现 -group： - -```python -rows.sem_groupby( - input_cols=["title", "abstract"], - instruction="Rows belong in one group when they discuss the same research topic.", -) -``` - -Closed-world grouping 预先给定 labels： - -```python -rows.sem_groupby( - input_cols=["title", "abstract"], - instruction="Assign each paper to the best matching research area.", - labels={ - "systems": "Systems, infrastructure, distributed systems, and databases.", - "ml": "Machine learning models, training, evaluation, and datasets.", - "hci": "Human-computer interaction and user studies.", - }, -) -``` - -当前最小实现只支持 single-label closed-world assignment： - -- 每行必须输出一个已声明 label。 -- 默认 label column 是 `"_label"`，可用 `label_col` 显式覆盖。 -- 没有隐式 `other` bucket；如果需要 other，policy author 必须把 `"other"` - 写进 `labels`。 -- `_label` 是 visible column，后续 `sem_agg(...)` 可以读取它；内部 - `_agent_memory_group_id` 仍用于 grouped aggregation。 - -Future work 只记录，不在当前实现中做：multi-label assignment、hierarchical -labels、label alias / canonicalization、unmatched confidence threshold、以及 labeled -grouping 的 candidate pruning / batching 优化。 - -除 pairwise baseline 外，未来 lowering strategy 可以包括： - -- `single_batch`：一次 prompt 让模型输出所有 row 的 group assignments。适合 - 小样本 audit，但 JSON 稳定性和 context window 风险更高。 -- `chunked_single_batch`：chunk 内分组，再跨 chunk 合并 representatives。需要 - 明确 representative 生成和跨 chunk merge 语义。 -- `candidate_pairwise`：用 embedding / index / cluster / sim-join 生成候选 - pairs，再用 LLM predicate 精判。这是最贴近 LOTUS 优化能力的优先方向。 - -待办： - -- 设计 `sem_groupby` 的 candidate generator abstraction，但不要提前泛化过度。 -- 第一版优化优先做 `candidate_pairwise`，保留 LLM 精判作为最终语义。 -- 后续完善 closed-world `labels` mode 的优化和 audit；不要把 multi-label 或 - hierarchy 混进当前 single-label 执行路径。 -- audit 要输出候选 pairs、LLM accepted pairs、最终 group ids，方便人工检查。 - -## 6. `sem_agg` TODO - -`sem_agg` 的本质是： - -```text -rows -> aggregate row(s) with declared output schema -``` - -single-output 当前复用 LOTUS lower-level `sem_agg`，因此继承了 LOTUS 的 -hierarchical aggregation / tree fold 能力。 - -multi-output 当前使用 agent-memory structured hierarchical lowering： - -```text -group rows --> LOTUS-main-style hierarchical aggregate --> final LM pass with JSON object response_format --> final JSON object with multiple output fields -``` - -这参考 LOTUS main branch 的方向：hierarchical aggregation 过程中先保持普通文本 -aggregate，最后一轮再应用 structured output contract。区别是当前 PyPI LOTUS 还没 -暴露 lower-level `sem_agg(response_format=...)`，所以 agent-memory 在 adapter -内部保留一个 compatibility helper，而不是 monkeypatch LOTUS 或覆盖 pandas -accessor。 - -更完整的 tree fold 路径未来可以是： - -```text -group rows --> chunked structured partial aggregate --> merge partial aggregate rows --> final JSON object with multiple output fields -``` - -当前 active runtime 不保留 `single_batch` / `lotus_hierarchical` backend -strategy。后续如果要恢复多个 lowering 选择，应该作为 optimizer/lowering 选择重新 -设计，并用 real audit 证明比当前 structured hierarchical baseline 更稳。 - -`sem_agg_model_kwargs` 是 adapter execution knob，只用于 multi-output final LM -pass。当前 final pass 默认保证 `max_tokens >= 1024`，然后再合并 -`sem_agg_model_kwargs`；调用方可以通过 adapter config 覆盖普通 model kwargs， -但不能覆盖 `response_format`。这个字段不是 public `sem_agg(...)` 参数，也不属于 -logical query semantics。 - -tree fold 在大 group 下可能更稳，但要和 LOTUS native `sem_agg` 的优化明确区分： - -- LOTUS native `sem_agg` 会把大 group 分 batch 聚合，再逐层 fold，避免超出 - context window。 -- 当前 multi-output path 依赖 final JSON object response_format。DeepSeek / - LiteLLM 如果结构化输出不稳定，real audit 应该直接失败，不做 retry/fallback - 掩盖问题。 -- 未来 tree fold 在 group 很大时可能出现 partial summary drift、JSON 被截断、 - 输出不稳定等问题，必须有 trace 和 audit 才能进入 active runtime。 -- 暂不把 LOTUS `operator_cache` 直接套到 compatibility helper 上。当前 LOTUS - decorator 面向 pandas accessor method，依赖 `self._obj` 参与 cache key； - plain helper 需要单独设计 hash-safe cache key，不能直接照搬。 - -你说“差距就是 output 行数的多少”只对了一部分。更准确地说： - -- output 行数：whole relation 是 1 行；grouped relation 是 `G` 行。 -- output 列数：single-output 是 1 列；multi-output 是多列。 -- 真正需要继续补齐的是：multi-output 的 tree fold 应该更接近 LOTUS native - aggregation 的工程成熟度，而不是只停留在基础 chunk + merge。 - -待办： - -- 保留 LOCOMO 小样本 real audit 对 whole/grouped、single/multi output 的覆盖。 -- 固定 row-count chunking 暂时不需要做。它不是 LOTUS 原生 `sem_agg` 的 - token-budget hierarchy，而是 agent-memory 曾考虑过的 structured tree-fold - experiment。当前 runtime 已采用 LOTUS-main-style token-budget hierarchy；额外 - long-context chunking 只作为后续 large-group optimization/audit，不混入本次 - reliability fix。 -- 单独增加 large-group optimization audit。当前真实 LOCOMO audit 曾暴露 - experimental tree-fold path 会出现空 JSON / 空 intermediate summary；不要在没有 - trace / retry / cost model 前宣称 large-group path 稳定。 -- 为 tree fold 增加 audit trace：每层 partial rows、merge rows、final rows。 -- 后续根据真实 CSV 审核决定是否需要不同的 leaf instruction / merge instruction。 - -## 7. “接入 LOTUS options”是什么意思 - -这句话不是说重新实现 LOTUS，而是说： - -```text -LOTUS 已经支持某个 execution option -但 agent-memory 的 adapter/runtime config 还没有表达它 -LotusAdapter 也没有把这个参数传下去 -``` - -例如 LOTUS `sem_filter` 支持 `examples` 和 `cascade_args`，但如果 -`LotusExecutionConfig` 或上层 runtime config 没有这些字段，而且 adapter 不转发 -它们，那么系统就无法从 agent-memory 使用 LOTUS 的这些能力。 - -这不代表要把参数加进 `Relation.sem_filter(...)`。Public operator API 的职责是 -表达 logical query；execution config 的职责是选择 backend method、cascade、 -examples、safe mode、trace 等执行策略。 - -这类 TODO 的工作一般是： - -1. adapter/runtime config 增加必要字段。 -2. `QueryExpr.params` 保持只记录 logical query 参数。 -3. `LotusAdapter` lowering 把 config 转成 LOTUS 调用参数。 -4. offline test 检查 logical params 没被污染、config 参数传递正确。 -5. real audit 检查真实效果。 - -## 8. “优化 custom operator”是什么意思 - -custom operator 是指 LOTUS 没有精确等价 operator 的部分，例如： - -- multi-output `sem_map` -- `sem_flat_map` -- `sem_groupby` -- multi-output `sem_agg` - -这些不能假装是 LOTUS native operator。但实现时仍然应该复用 LOTUS 的基础设施： - -- LM configuration -- prompt formatter -- operator cache -- safe mode / cost estimation -- examples / strategy / trace / audit stats -- lower-level `sem_filter` / `sem_join` / `sem_agg` -- semantic index / retrieval / cluster 工具 - -这里的目标不是追求“看起来和 LOTUS API 一样”，而是让 agent-memory 自己的 -operator contract 在执行层也有类似 LOTUS 的成本、稳定性和可审计性。 - -Debug 信息默认不要进入 result DataFrame。比如 raw output、reasoning、 -explanation 更适合作为 adapter trace / audit log / stats object。如果 explanation -本身是业务字段，就应该由 policy writer 显式声明在 `output_cols` 里，例如 -`"memory_summary_rationale"`，而不是通过 execution option 偷偷写入结果列。 - -## 9. Python 代码里是否加 TODO - -暂时不建议在每个 Python 文件里大量加 `TODO`。原因： - -- 当前问题是跨 API、adapter、真实测试的路线问题，不是某个函数里一行能修掉的 - 局部问题。 -- 过多 inline TODO 会让代码噪音变大，而且容易过期。 -- 设计 TODO 集中放在本文档和 `docs/design/TODO.md` 更容易维护。 - -可以加 inline TODO 的情况： - -- 某个函数里确实有明确、局部、下一步就要修的技术债。 -- TODO 能指向本文档的具体章节。 -- 不加 TODO 会导致后续维护者误以为当前行为就是最终设计。 - -推荐格式： - -```python -# TODO(operator-parity): wire sem_filter examples/cascade through adapter config; see docs/design/operator-todo.md. -``` - -不要写模糊 TODO，例如： - -```python -# TODO: improve this later -``` - -## 10. Deterministic Aggregate / Source Metadata TODO - -当前已实现单列/复合 `min` aggregate 和 row-wise `least` expression。复合 -`min(columns=[...])` 返回字典序最小 tuple；`least` 只比较同一 row 的 operands。 -二者都不会返回 ordering key 对应的另一份 payload。 - -待办： - -- `min_by(value, order_by)` / `arg_min`：返回最小 ordering value 对应的整行或 - payload。只有出现真实 policy 需求后再确定 public spelling。 -- delete / correction semantics：append-only `min` 可以只合并 old/new minima； - 当前 minimum 被撤回时，需要 raw-group recompute 或 ordered auxiliary state。 -- concurrent / batch append ordering：当前 `_add_seq` 只定义单进程顺序 add。 - 并发、同 batch 多 rows、跨 restore writer coordination 需要单独 runtime - contract，不能让 policy writer修改 system columns。 diff --git a/docs/design/rules-todo.md b/docs/design/rules-todo.md deleted file mode 100644 index 8ac1d67..0000000 --- a/docs/design/rules-todo.md +++ /dev/null @@ -1,594 +0,0 @@ -# Differential Rules TODO - -本文记录 `Q -> Q'` differential rules 的当前共识和未决点。它不是 -implementation spec，也不表示这些 rules 已经实现。当前黄金线仍然是 -`docs/optimization/incremental-semantic-view-maintenance.tex` 和 -`docs/design/operator-api.md`。 - -## 1. 核心区分 - -必须区分两个层次： - -- **Delta expression**：由 differential rule 推导出来的 changed-output query - fragment，例如 `DeltaD.sem_filter(...)`。 -- **Maintenance query `Q'`**：最终交给 runtime 执行的 logical `QueryExpr`， - 语义是返回下一版 materialized view `V'`。 - -`Q -> Q'` 是 query rewrite，不看真实 data。Planner 阶段只处理 logical -`QueryExpr`。Runtime 阶段才把真实 `changed_rows` 绑定到 differentiated query -的 source input，把当前 materialized view 绑定成 `V`，然后执行 `Q'`。 - -本文里的 `DeltaD` / `DeltaQ` 是数学说明用语。代码主路径不要新增 -`delta_log`、`delta_source`、`delta_query` 这类名字，因为当前 interface 里没有 -单独的 delta schema、delta table、或 delta leaf。当前代码里 differentiated -query 的 source leaf 仍然是 `QueryExpr(op="log")`，runtime 只是在执行 `Q'` 时 -把它绑定到本次 `changed_rows`。 - -因此，primitive operator rule 通常只推导 delta expression；`union`、 -`subtract`、upsert、delete、replace 等 state application 语义属于 -materialization boundary，不应该无脑塞进每个 primitive operator rule。 - -例子： - -```text -Q = D.sem_filter(p).select(cols) - -delta expression: - DeltaQ = DeltaD.sem_filter(p).select(cols) - -maintenance query: - Q' = V.union(DeltaQ) -``` - -## 2. Differential Policy Artifact / DifferentiatedPolicy - -`Q -> Q'` 不应该只能发生在 `memory.add(...)` 时。它不看真实 data，所以更适合 -作为 policy-level differentiation 阶段的一部分： - -```text -Policy Q - -> differentiate 得到 Q' - -> offline optimizer 优化 Q' - -> 保存 DifferentiatedPolicy artifact - -> end-user runtime load artifact 并执行 -``` - -`PolicyDifferentiator` 负责 whole-policy differentiation。它组合 -`QueryDifferentiator`，把完整 policy 转成 `DifferentiatedPolicy`： - -```text -MemorySpec - -> PolicyDifferentiator - -> DifferentiatedPolicy - -> optional offline optimizer - -> DifferentiatedPolicy - -> runtime -``` - -`DifferentiatedPolicy` 的最小方向是保存： - -- `spec: MemorySpec`。 -- `nodes`、`execution_order` 和 `view_outputs`，包含共享 node、node-local - execution query、stateful maintenance query、topological order 和 public - sinks。 -- `retrieval_queries: Mapping[str, QueryExpr]`，parameterized retrieval query - templates。 -- `fingerprint` 和 canonical grouped rule，用于 checkpoint compatibility。 - -不要在 `DifferentiatedPolicy` 里增加 `optimized_dataflow_plan` / -`optimized_retrieval_queries` 这类平行字段。Optimizer 应输入一个 -`DifferentiatedPolicy`，输出一个新的 `DifferentiatedPolicy`，直接替换 nodes、 -retrieval queries 或 execution order。 - -artifact 可以用 JSON、YAML、或其他机器可执行格式表示。Runtime load artifact 后 -只负责把真实 `changed_rows`、当前 materialized views、storage state 绑定进去并 -执行；runtime 可以继续做执行期优化，但不应该承担主要 query rewrite 责任。 - -早期实现曾在 `memory.add(...)` 时调用 planner，适合调试但不是最终架构。 -当前第一版已经采用 in-memory `DifferentiatedPolicy`； -JSON/YAML IO、serializer 和 durable artifact save/load 留到后续。 - -## 3. Declarative Retrieval Query Templates - -Retrieval query 也应该进入 differentiated policy，但它进入的是 parameterized -query template，不是已经绑定具体 user text 的 concrete query。 - -Policy writer 应声明 class-body retrieval query： - -```python -class ClaudeMemory(am.Memory): - ... - retrieval_query = ( - topics.select(["name", "description", "type"]) - .sem_topk(am.UserQuery(), 5) - .join(topics, on="name") - .select(["name", "description:right", "type:right", "body"]) - ) -``` - -`Memory.query(...)` 是 framework 内置的 end-user method，policy writer 不需要 -手写： - -```python -memory.query("design docs") -``` - -`UserQuery` 是 runtime-bound placeholder，语义是 end-user retrieval input。它 -应该和 column placeholder `{message}` 分开，避免把 user query text 和 dataframe -column reference 混在同一套 string placeholder 里。 - -Offline 阶段保存： - -```text -retrieval_queries["default"] = - materialized_view("topics") - .select(["name", "description", "type"]) - .sem_topk(UserQuery("query"), 5) - .join(materialized_view("topics"), on="name") - .select(["name", "description:right", "type:right", "body"]) -``` - -Runtime 阶段再绑定： - -```text -UserQuery("query") -> "design docs" -``` - -Retrieval query 不做 view maintenance rule，也不产生 `V'`。Retrieval-time -exact `join(on=...)` 是确定性 lookup / body fetch，不需要 differential rule。 -它属于 query-time retrieval template，可以被 offline optimizer 改写 template， -也可以被 online optimizer 在绑定 user text 后做 batching、cascade、cache、 -model choice 或 index 相关优化。 - -## 4. Planner / Rules 职责边界 - -代码命名约定： - -- `QueryDifferentiator(...).differentiate(view)`：生成 differentiated query - `Q'`，语义是返回下一版 view `V'`。 -- `DifferentialRules.differentiate(query, *, source_input, current_view, - is_view_boundary, instruction_rewriter)`：执行 operator / pattern rewrite。 -- `source_input`：differentiated query 的 source leaf；当前复用 - `QueryExpr(op="log")`，不新增 `delta_input` / `delta_log` operator。 -- `current_view`：当前 materialized view 的 logical placeholder，例如 - `QueryExpr(op="materialized_view", params={"name": view.name})`。 -- `changed_rows`：runtime 真实 DataFrame，来自 `Memory.add(...)`，不是 - `QueryExpr`。 -- `DifferentialInstructionRewriter`：instruction rewrite hook。第一版只做 - deterministic placeholder rewrite，不做 semantic prompt rewrite。 - -暂不引入 `Context` class。如果以后参数变多，可以考虑 -`DifferentiationScope`，但当前先使用显式 keyword 参数，避免过早抽象。 - -`rules.py` 应负责 operator-level 和 pattern-level rewrite： - -- 推导 `select`、`sem_filter`、`sem_map`、`sem_flat_map` 等 row-local delta - expression。 -- 推导 `sem_join` 这类有明确 relational delta 公式的 operator pattern。 -- 推导 `sem_groupby(...).sem_agg(...)` 这类 grouped aggregate view-boundary - pattern。 - -`differential.py` 应负责 view-level context： - -- 知道当前正在维护哪个 `MemoryView`。 -- 提供 `current_view` / `V` 这类 materialized-view placeholder。 -- 提供 `DeltaD` / changed input placeholder。 -- 把 rule 产生的 delta expression assemble 成返回 `V'` 的 maintenance query。 -- 注入 instruction transformer。第一版 transformer 只做 deterministic - placeholder rewrite；以后如果 `sem_groupby -> sem_join` 或 `sem_agg -> sem_map` - 需要 semantic prompt rewrite，可以替换 transformer，而不是重写整个 planner。 - -`differentiate(...)` 的返回约束： - -- `is_view_boundary=False`：只能返回 changed-output fragment。 -- `is_view_boundary=True`：可以返回 full `V'` query。 -- `sem_groupby(...).sem_agg(...)` 这类 stateful pattern 只能在 view boundary - 返回 full `V'`。如果它出现在普通中间 subtree，第一版应 - `NotImplementedError`，不要假装可以维护未 materialized 的中间 state。 - -## 5. Row-Local Rules - -当前确定的 row-local rules： - -```text -Delta(log) = DeltaD - -Delta(select(Q)) = - select(Delta(Q)) - -Delta(sem_filter(Q, instruction)) = - sem_filter(Delta(Q), instruction) - -Delta(sem_map(Q, input_cols, output_cols, instruction)) = - sem_map(Delta(Q), input_cols, output_cols, instruction) - -Delta(sem_flat_map(Q, input_cols, output_cols, instruction, ordinal_col)) = - sem_flat_map(Delta(Q), input_cols, output_cols, instruction, ordinal_col) -``` - -这些 rule 只产生 changed output。它们不负责 `V.union(...)`。如果这些 -operator 的输出是某个 materialized view 的最终输出，planner 再在 view -boundary assemble： - -```text -Q' = V.union(DeltaQ) -``` - -如果这些 operator 只是中间 subquery，则不需要在中间位置 union。未来如果 -optimizer 选择 materialize 某个中间 table，该中间 table 自己才形成新的 -materialization boundary。 - -## 6. `sem_join` Rules - -`sem_join` 的 inner join delta 公式是确定的： - -```text -Q = L.sem_join(R, instruction, how="inner") - -DeltaQ = - DeltaL.sem_join(R, instruction, how="inner") - .union(L.sem_join(DeltaR, instruction, how="inner")) - .union(DeltaL.sem_join(DeltaR, instruction, how="inner")) -``` - -如果 `Q` 是 materialized view `V` 的 definition，则 maintenance query 可以是： - -```text -Q' = V.union(DeltaQ) -``` - -这个 rule 的前提是 planner / runtime 能绑定 old `L` 和 old `R`。也就是说， -`L` 和 `R` 必须是 view、stored table、或已经 materialized 的 intermediate -state。 - -如果 `L` / `R` 是任意 subquery，当前不要假装可以 generic differential。 -未来有两个可选方向： - -- **要求中间 materialization**：planner 发现 join 需要 old `L/R`，就要求 - runtime / materializer 保存 `L` 和 `R`。 -- **full recompute fallback**：重新计算 `L(D union DeltaD)` 和 - `R(D union DeltaD)`，再执行 join。结果正确，但不是 differential，成本和 - latency tradeoff 必须显式暴露。 - -第一版 generic `sem_join` differential 只支持 `how="inner"`。`how` 为 -`"left"`、`"right"`、或 `"outer"` 时应先 `NotImplementedError`，不要假装 -append-only maintenance 已经正确。 - -原因是 non-inner join 的 old unmatched output 可能会因为新 match 失效。例子： - -```text -old L: - left_row = "Alice likes concise docs" - -old R: - empty - -old V = L.sem_join(R, how="left"): - ("Alice likes concise docs", NULL) - -DeltaR: - right_row = "documentation preference" - -如果 DeltaR 匹配 left_row，正确 V' 应该是： - ("Alice likes concise docs", "documentation preference") - -旧输出： - ("Alice likes concise docs", NULL) - -必须删除或改写。单纯 V.union(DeltaQ) 会同时保留 matched row 和旧 unmatched -row，因此是错误的。 -``` - -未来可以考虑两类解决方向： - -- **full recompute fallback**：重新计算完整 non-inner join，结果正确但不是 - differential。 -- **provenance / action-based maintenance**：记录 row identity、`delta_minus`、 - row action tag、或 `match_type`，从而表达旧 unmatched output 的删除 / - 改写。 - -这不影响 Claude-style grouped aggregate rule。Claude rule 里的 `outer join` 是 -maintenance query `Q'` 内部用来构造 merge input 的一步，随后 `sem_map` 直接 -产出下一版 full view `V'`；它不是一个长期 materialized join view 的 -append-only maintenance。 - -## 7. Instruction Placeholder Rewrite - -`sem_groupby(...).sem_agg(...) -> sem_join(...).sem_map(...)` 这类 rewrite 会改变 -输入 schema。典型情况是 `sem_join` 后重名列会被改成 side-suffixed columns： - -```text -name -> name:left, name:right -description -> description:left, description:right -body -> body:left, body:right -``` - -如果原 instruction 仍然包含 `{name}`、`{body}` 这样的 placeholder，而 join 后 -DataFrame 已经没有 `name` / `body` column，LOTUS 会严格解析 placeholder 并可能 -直接报错。这是 correctness risk，不只是 accuracy risk。 - -第一版 instruction rewrite 只做 deterministic placeholder rewrite，例如： - -```text -{name} -> {name:left} and {name:right} -{body} -> {body:left} and {body:right} -``` - -这只是保证 differentiated query 可执行，不做 create / keep / merge / delete 的 -semantic prompt rewrite。更好的 semantic rewrite 后续再设计，例如显式说明 left -side 是 new candidate、right side 是 existing memory、matched rows 才需要 merge。 - -`rule-re-group` 的 final semantic aggregate 还会消费 aggregate-state rows，而不 -是原始 rows。当前 schema-aware rewrite 保留 state 中仍存在的 `{column}`；如果 -某个 placeholder 属于原始 `input_cols`、但已经不在 state schema 中，只去掉花 -括号变成普通文字。final `sem_agg(input_cols=None)` 读取除 internal semantic -group id 外的所有可见 state columns。这同样是 deterministic schema repair， -不是 semantic rewrite。 - -## 8. `sem_agg` Rules - -`sem_agg` 应该有 maintenance rule，但必须写清楚精度边界。 - -Full view: - -```text -V = D.sem_agg(input_cols, output_cols, instruction) -``` - -已实现的第一版 standalone compressed-state maintenance 形态是： - -```text -DeltaX = differentiate(X) - -Q' = - V.select(input_cols) - .union(DeltaX.select(input_cols)) - .sem_agg(input_cols, output_cols, instruction) -``` - -这里的 hard correctness check 只是 schema-level：`V` 和 `DeltaX` 都必须包含 -显式声明的 `input_cols`。多余 columns 不影响 operator correctness；缺少 -`input_cols` 才是 hard error。 - -这条 rule 是 compressed-state semantic maintenance，可能不同于： - -```text -Q(D union DeltaD) -``` - -除非 `V` 恰好是 sufficient aggregate state。这里的 sufficient 不是 schema -equality。最终 view 的 columns 和原始输入 columns 不同，不代表一定不 -sufficient；关键是维护所需信息是否仍然存在： - -```text -required maintenance inputs - subset of current_view_state + changed_input -``` - -例如原始输入包含 `name, body, timestamp, evidence`，最终 view 只有 `name, body`。 -如果 instruction 只需要 `name/body` 继续 merge，则 current view 可能 sufficient。 -如果 instruction 需要 `timestamp/evidence` 判断 freshness、contradiction、delete， -而这些字段已经被 view 丢掉，则 current view 不 sufficient。 - -第一版不做 semantic sufficient proof，也不做 full recompute fallback。只要 -`input_cols` 存在，就可以生成 compressed-state `Q'`；但文档和 audit 必须明确 -它是 approximate maintenance，不声明 full-recompute equivalence。 - -第一版 instruction 可以直接复用。未来如果发现输入 shape 改变导致语义偏移， -再通过 instruction transformer 自动生成 aggregate-to-merge instruction。 - -## 9. `sem_groupby(...).sem_agg(...)` Grouped Aggregate View-Boundary Rule - -`sem_groupby(...).sem_agg(...)` 应该是 pattern rule，不应该作为 -`differential.py` 里的特殊业务逻辑。 - -Full view pattern: - -```text -V = ( - X - .sem_groupby(input_cols=group_cols, instruction=group_instruction) - .sem_agg(input_cols=agg_input_cols, output_cols=V_columns, instruction=agg_instruction) -) -``` - -Delta branch: - -```text -DeltaT = ( - DeltaX - .sem_groupby(input_cols=group_cols, instruction=group_instruction) - .sem_agg(input_cols=agg_input_cols, output_cols=V_columns, instruction=agg_instruction) -) -``` - -Coarse full-next-view maintenance query: - -```text -Q' = ( - DeltaT - .sem_join(V, instruction=group_to_join(group_instruction), how="outer") - .sem_map(output_cols=V_columns, instruction=agg_to_map(agg_instruction)) - .select(V_columns) -) -``` - -第一版 `group_to_join(...)` 和 `agg_to_map(...)` 只做 deterministic placeholder -rewrite，不做 semantic prompt rewrite： - -- `sem_groupby -> sem_join` 复用 topic identity / grouping predicate。 -- `sem_agg -> sem_map` 复用 aggregate / merge instruction。 - -这条 rule 产出的是 `V'`，不是 minimal `DeltaV`。它适合 Claude-style topic -memory 的第一版实现，因为我们还没有 `delta_plus` / `delta_minus` / -`apply_delta` / storage upsert semantics。 - -代码命名不要把这个 generic rule 叫 `consolidation`。`consolidation` 可以作为 -Claude prompt 里的业务词，但 planner / rules 层应使用 -`grouped aggregate view-boundary rule` 或类似命名。后续实现也应拆小函数，例如： - -- `_differentiate_row_local_unary(...)` -- `_differentiate_binary_fragment(...)` -- `_differentiate_sem_join(...)` -- `_differentiate_sem_agg(...)` -- `_differentiate_grouped_aggregate_view(...)` - -## 10. `sem_groupby.sem_agg` 的优化空间 - -`DeltaT.sem_join(V, how="outer")` 之后，不是每一行都必须做 expensive -semantic merge。 - -可优化情形： - -- existing-only row：delta side 为空，结果可以 deterministic keep。 -- delta-only row：existing side 为空，结果可以 deterministic create。 -- matched row：两边都有内容，才需要 semantic merge / rewrite。 - -要做这个优化，需要 join output 暴露足够的信息，例如： - -- left/right side 是否为空。 -- `match_type`：matched、unmatched_delta、unmatched_existing、contradict、 - delete_target、uncertain 等。 -- row action tag：keep、create、update、delete、skip。 - -这些会影响 query shape、runtime apply semantics、storage API 和 audit trace。 -当前先记录为 future optimization，不在第一版 rules 中定死。 - -## 11. Apply Delta / Future State Effects - -未来可能有多种维护形态： - -- `Q'` 直接返回 full next view `V'`。 -- `Q'` 返回 `delta_plus` / `delta_minus`，runtime 用 `subtract` + `union` - 应用。 -- `Q'` 返回带 action tag 的 rows，由 runtime / storage 执行 keep、create、 - update、delete、skip。 -- 引入 internal-only `apply_delta` operator，但不暴露为 public policy API。 - -当前 Claude-first implementation 应优先使用 full-next-view `V'` 形态，因为它 -最容易审计，也不要求提前设计 storage upsert/delete contract。 - -## 12. Public View Dependency Binding - -如果一个 public view 依赖另一个 public view，planner 应把依赖子树绑定成 -materialized view，而不是在 dependent view 里再次 differentiate 上游 view 的 -内部 query。 - -例子： - -```text -topics = log.sem_flat_map(...).sem_groupby(...).sem_agg(...) -catalog = topics.sem_map(...).select(...) - -catalog 的 differentiated query 应基于： - materialized_view("topics").sem_map(...).select(...) -``` - -这条规则是 general 的，不是 Claude special case。但当前第一版有两个限制： - -- **exact structural match only**：只有 query subtree 和另一个 public - `MemoryView.query` 完全相等时才替换；不识别 logically equivalent query。 -- **declaration-order dependency**：runtime 当前依赖 view declaration order。 - 如果 dependent view 先执行，可能读到 missing / stale materialized state。 -- **mixed log + upstream view unsupported**：当前 fast path 只覆盖 pure - upstream-only dependency，例如 `catalog = topics.sem_map(...)`。如果一个 - dependent view 同时包含 base `log` 和 upstream `materialized_view`，例如 - `union(log.sem_filter(...), topics.sem_map(...))`，就需要 upstream view 在本次 - add 中产生的 changed-view state。v0.0 runtime 目前只暴露 current / next view - state，不暴露 changed upstream view rows，因此这类 mixed tree 暂不支持 - differential。 - -未来 `DifferentiatedPolicy` artifact 应保存 public view dependency graph，并在 -compile 阶段做 topological ordering 和 cycle detection。Runtime load artifact -后按 compiled order 执行，而不是依赖 class body declaration order。 - -## 13. 暂不实现的 Rules - -- generic `filter(predicate)` / `assign(...)` differential：operator execution 已 - 收窄成 minimal serializable `relation.col(...)` expression subset，但通用 - differential rules 仍需逐类定义。 -- generic `subtract` differential：需要 negative delta / delete semantics。 -- generic arbitrary-subquery `sem_join`：需要 old intermediate materialization - 或 full recompute fallback。 -- mixed base log + upstream materialized view dependency：需要 upstream - changed-view state、中间 materialization，或 full recompute fallback。 -- generic `left/right/outer sem_join` delta：旧 unmatched output 的失效和改写 - 还没设计。 -- standalone `sem_groupby`：它是 grouped intermediate，不是最终 view output。 -- `sem_topk` maintenance rule：`sem_topk` 主要是 query-time retrieval / - ranking template。它不维护 materialized view，也不产生 `V'`；但 declarative - retrieval query template 应进入 `DifferentiatedPolicy.retrieval_queries`。 - -这些不是永远不做，而是不能在没有 state / materialization / apply semantics 的 -情况下假装已经有正确 generic rule。 - -## 14. Validation / Error-Message TODO - -当前第一版把 correctness 相关的 placeholder rewrite 做在 differentiation 阶段。 -后续可以补两个 DX 改进，但不应在本轮规则里硬塞完整 validator： - -- **spec-time placeholder validation**：在 `Memory.spec()` / policy compile 阶段 - 提前检查 semantic instruction 中的 `{identifier}` placeholder 是否能解析到 - 当前 operator scope 的 input / output columns。注意这需要可靠 column-scope - 推导，不能只做字符串扫描。 -- **view-name-aware planner errors**：当某个 view 的 grouped aggregate pattern - 缺 output columns、placeholder typo、或 unsupported mixed dependency 报错时， - 错误信息应带上 `MemoryView.name`，例如 `topics: ...`，减少 deep stack debug。 - -## 15. 实现顺序 - -建议按以下顺序实现，避免一次性重写 planner/runtime： - -1. 更新 planner / rules docstring 和局部变量名，不改行为。 -2. 让 `DifferentialRules.differentiate(...)` 接收显式 keyword 参数： - `source_input`、`current_view`、`is_view_boundary`、`instruction_rewriter`。 -3. 实现 row-local changed-output fragment rules：`log`、`select`、 - `sem_filter`、`sem_map`、`sem_flat_map`。 -4. 在 `QueryDifferentiator(...).differentiate(view)` 的 view boundary assemble - `current_view.union(fragment)`，使 planner 返回完整 `Q'`。 -5. 实现 `sem_groupby(...).sem_agg(...)` 的 view-boundary stateful pattern， - 返回 full `V'` query。 -6. 更新 runtime：执行 `Q'`，并把结果写入 `state[view.name]`。Runtime 继续用 - `changed_rows` 表示真实输入 batch，不引入 `delta_log` / `delta_source` - 变量。 -7. 已实现：增加 in-memory policy-level differentiation surface： - `Memory.differentiate_policy() -> DifferentiatedPolicy`。 -8. 已实现：支持 declarative retrieval query：policy writer 写 - native-like `topics` manifest selection + deterministic `join(on="name")` - body fetch，runtime 内置 `Memory.query(...)`。 -9. 已实现：给 `DifferentialInstructionRewriter` 加 deterministic placeholder - rewrite，先不做 semantic prompt rewrite。 -10. 已实现：给 standalone `sem_agg` 加 approximate compressed-state view rule； - 只做 schema-level `input_cols` check，不做 semantic sufficient proof 或 full - recompute fallback。 -11. 已实现：把 grouped aggregate view-boundary rule 和 standalone aggregate view - rule 拆成独立 helper，避免 generic rules 层使用 `consolidation` 作为抽象名。 -12. 已实现：`PolicyDifferentiator` 生成共享 `DifferentiatedPolicy`；runtime 按 - node topology 传播 multiset change，保存 semantic output cache 和 window - state，并原子提交一次 add step。 - -## 16. Grouped Aggregate / DAG Follow-ups - -当前 grouped aggregate implementation 已使用真实 operator graph 表达 -`collect_list(...).flatten(...)`、row-wise `array_cat(...)`、`min(...)` 和 -`least(...)`；没有 `agg_merge` operator 或 planner-only merge mode。 -单列和复合 `min(columns=[...])` 共享相同 rule：第一次聚合产生 scalar/tuple -minimum state，后续只合并该 output state。 - -Shared finite-DAG changed-state propagation 已实现。Compiler 合并结构相同的 -subtree；runtime 保存 private node state，并让 replacement change 继续传给下游。 -完整设计见 `policy-differentiation-dataflow-runtime.zh.md`。 - -仍未完成： - -- **partitioned semantic join-map**：`sem_groupby(partition_by=...).agg(...)` 的 - re-group rule 可保留 deterministic partitions；join-map 当前明确拒绝，直到 - semantic join 能在 ON 语义中同时强制 partition equality。 -- **minimum removal**：append-only `min` merge 是 exact；delete 或 current-min - removal 需要 negative delta、raw-group recompute 或 ordered auxiliary state。 -- **`min_by` / `arg_min`**：只有 view contract 真正需要“最小 ordering value - 对应 payload”时再增加 operator 和 differential rule。 -- **recursive/fixpoint dataflow**：当前 compiler 只接受有限 DAG；recursive view - query 在 compile time 拒绝。 -- **physical optimizer**：deterministic node 先按 next parent state 精确重算； - indexed join/aggregate state 和 arrangement sharing 留给 optimizer。 diff --git a/docs/design/v0.0-interface.md b/docs/design/v0.0-interface.md index 7ead756..e25f982 100644 --- a/docs/design/v0.0-interface.md +++ b/docs/design/v0.0-interface.md @@ -74,6 +74,65 @@ v0.0 message-to-log column convention: - Platform-native message objects are not interpreted by core `Memory`; a future agent-platform adapter should convert them into `am.Message` or a mapping. +Maintenance is eager unless the caller explicitly chooses a count-based +refresh boundary: + +```python +memory = ClaudeMemory(refresh=am.CountRefresh(every=32)) + +memory.add(event) +print(memory.pending_count) +memory.flush() +``` + +With this configuration, `add(...)` normalizes and buffers each source row. +Reaching 32 pending rows synchronously runs the full differential DAG once with +one multi-row delta. `flush()` publishes a shorter tail. The refresh is atomic: +on failure, the current view remains unchanged and the pending rows remain +available for retry. + +`query(...)` does not flush. It reads the most recently published view, so it +may omit accepted-but-pending rows. Omitting `refresh`, or choosing +`CountRefresh(every=1)`, preserves the existing eager behavior. v0.0 supports +only this synchronous count trigger; timers, background scheduling, +query-triggered refresh, and a crash-durable pending log remain future work. + +LOTUS prompt batching is a separate, opt-in physical setting. `CountRefresh` +controls how many source rows enter one maintenance step; `PromptBatching` +controls how the independent semantic tasks that are ready inside one operator +invocation are placed into prompts: + +```python +from agent_memory.adapters.lotus import LotusAdapter, PromptBatching +from agent_memory.adapters.lotus.context import LotusExecutionConfig + +adapter = LotusAdapter( + config=LotusExecutionConfig( + prompt_batching=PromptBatching(), # all ready tasks per prompt + # prompt_batching=PromptBatching(max_tasks=8), + ) +) +memory = ClaudeMemory( + adapter=adapter, + refresh=am.CountRefresh(every=32), +) +``` + +The default `prompt_batching=None` preserves existing prompts. When enabled, +the adapter uses stable task IDs, deterministic context-bounded chunks, strict +output validation, and retries only failed chunks. Search-Filter still runs +before prompt construction. Prompt batching is also distinct from provider +request batching: the former puts several tasks in one prompt, while the latter +sends several unchanged prompts in one provider request. The physical setting +is recorded in the maintenance/checkpoint identity. + +This setting is experimental and is not semantically transparent. Even when +every task payload is unchanged, placing several tasks in the same model +context can change extraction, matching, or aggregation decisions. Current +32-event Claude/Mem0/Zep evidence did not establish a safe system-wide batch +size, so `PromptBatching()` and `max_tasks` must remain explicit physical +experiment conditions rather than v0.0 defaults. + By contrast, `query(...)` is policy-owned retrieval behavior, but policy writers do not override the method. A concrete memory policy declares a parameterized class-body retrieval template. For Claude-style memory, retrieval selects from @@ -284,16 +343,20 @@ Responsibilities: - append new log rows - normalize end-user append input such as string messages into the configured log columns +- optionally retain normalized rows until a configured count boundary or an + explicit `flush()` - propagate relation changes through the shared policy nodes -- execute each changed policy node at most once per add step +- execute each changed policy node at most once per published delta - call an `ExecutionAdapter` - stage and atomically commit public and private node state - bind `UserQuery` placeholders in compiled retrieval templates - support write-side maintenance for end-user methods such as `add(...)` -v0.0 runs one eager, sequential dataflow step per appended log row. Lazy, -batched, query-triggered, freshness-aware, concurrent, and watermark-based -scheduling remain future runtime work. +v0.0 runs one eager, sequential dataflow step per appended log row by default. +`CountRefresh(every=N)` may instead coalesce source rows into one synchronous, +atomic delta step. This is not a general lazy scheduler: time triggers, +background refresh, query-triggered refresh, freshness policies, concurrent +maintenance, and watermark-based scheduling remain future runtime work. `MemoryRuntime` owns materialized memory state. `MemorySpec` remains the logical policy declaration and should not carry runtime dataframe contents. The v0.0 diff --git a/pyproject.toml b/pyproject.toml index 7a099fd..4bbb629 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ authors = [ requires-python = ">=3.12" dependencies = [ "dspy>=3.2.0", + "json5>=0.14.0", "lotus-ai==1.1.4", "nltk>=3.9.1", "pandas>=2.3.3", diff --git a/src/agent_memory/__init__.py b/src/agent_memory/__init__.py index fa2f566..31954a0 100644 --- a/src/agent_memory/__init__.py +++ b/src/agent_memory/__init__.py @@ -1,6 +1,7 @@ """Semantic memory framework experiments for agent developers.""" -from .api import Memory, Message +from .api import CountRefresh, Memory, Message +from .dataflow import SemanticDataflow, Source from .memories import ( ClaudeMemory, Mem0Memory, @@ -18,12 +19,18 @@ Log, RetrievalQuery, RetrievalResult, + Relation, UserQuery, array_agg, + avg, + case_when, collect_list, + count, least, min, sem_agg, + sum, + try_cast, ) __version__ = "0.0.1" @@ -34,6 +41,7 @@ "BFS", "BM25", "CosineSimilarity", + "CountRefresh", "CrossEncoder", "Log", "Mem0Memory", @@ -41,15 +49,23 @@ "Memory", "Message", "RRF", + "Relation", "RetrievalQuery", "RetrievalResult", + "SemanticDataflow", + "Source", "UserQuery", "ZepMemory", "ZepMemoryExtended", "__version__", "array_agg", + "avg", + "case_when", "collect_list", + "count", "least", "min", "sem_agg", + "sum", + "try_cast", ] diff --git a/src/agent_memory/adapters/lotus/__init__.py b/src/agent_memory/adapters/lotus/__init__.py index 69dbd90..c96725c 100644 --- a/src/agent_memory/adapters/lotus/__init__.py +++ b/src/agent_memory/adapters/lotus/__init__.py @@ -1,5 +1,6 @@ """LOTUS adapter package.""" from .adapter import DEFAULT_LOTUS_MODEL, LotusAdapter +from .prompt_batching import PromptBatching -__all__ = ["DEFAULT_LOTUS_MODEL", "LotusAdapter"] +__all__ = ["DEFAULT_LOTUS_MODEL", "LotusAdapter", "PromptBatching"] diff --git a/src/agent_memory/adapters/lotus/adapter.py b/src/agent_memory/adapters/lotus/adapter.py index 3e2b28c..030b3d8 100644 --- a/src/agent_memory/adapters/lotus/adapter.py +++ b/src/agent_memory/adapters/lotus/adapter.py @@ -13,6 +13,13 @@ LotusExecutionConfig, LotusExecutionContext, ) +from agent_memory.adapters.lotus.algebraic_aggregates import ( + execute_aggregate_finalize, + execute_aggregate_state, + execute_aggregate_state_update, + execute_algebraic_aggregate, + is_algebraic_aggregate_query, +) from agent_memory.adapters.lotus.pair_execution import ( semantic_pair_profiles_fingerprint, ) @@ -30,6 +37,7 @@ execute_group_by, execute_join, execute_let, + execute_limit, execute_min, execute_select, execute_subtract, @@ -87,6 +95,41 @@ def maintenance_execution_fingerprint(self) -> str: pair_fingerprint = semantic_pair_profiles_fingerprint( self.config.semantic_pair_profiles ) + context_limit = self.config.lm_model_kwargs.get("max_ctx_len") + sem_agg_configured = self.config.sem_agg_dispatch != "sequential" + transport_configured = ( + self.config.structured_output_transport != "chat-json-object" + ) + if ( + sem_agg_configured + or self.config.prompt_batching is not None + or transport_configured + or context_limit is not None + ): + base_fingerprint = self._base_maintenance_execution_fingerprint( + pair_fingerprint + ) + parts = [f"base_execution_fingerprint={base_fingerprint}"] + if context_limit is not None: + parts.append(f"lm_max_ctx_len={context_limit}") + if sem_agg_configured: + parts.append(f"sem_agg_dispatch={self.config.sem_agg_dispatch}") + if transport_configured: + parts.append( + f"structured_output_transport={self.config.structured_output_transport}" + ) + if self.config.prompt_batching is not None: + parts.append( + "prompt_batching=" + f"{self.config.prompt_batching.fingerprint}" + ) + payload = "\n".join(parts) + return sha256(payload.encode("utf-8")).hexdigest() + return self._base_maintenance_execution_fingerprint(pair_fingerprint) + + def _base_maintenance_execution_fingerprint(self, pair_fingerprint: str) -> str: + """Preserve the pre-prompting physical fingerprint contract.""" + if self.config.sem_join_topk_method == "listwise": return pair_fingerprint payload = ( @@ -107,6 +150,8 @@ def execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: return execute_window_source(inputs) case "select": return self._execute_traced_relational(query, inputs, execute_select) + case "limit": + return self._execute_traced_relational(query, inputs, execute_limit) case "alias": return self._execute_traced_relational(query, inputs, execute_alias) case "assign": @@ -130,7 +175,23 @@ def execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: case "unnest": return self._execute_traced_relational(query, inputs, execute_unnest) case "agg": + if is_algebraic_aggregate_query(query): + return self._execute_traced_relational( + query, inputs, execute_algebraic_aggregate + ) return self._execute_traced_semantic(query, inputs, execute_agg) + case "aggregate_state": + return self._execute_traced_relational( + query, inputs, execute_aggregate_state + ) + case "aggregate_state_update": + return self._execute_traced_relational( + query, inputs, execute_aggregate_state_update + ) + case "aggregate_finalize": + return self._execute_traced_relational( + query, inputs, execute_aggregate_finalize + ) case "concat": return self._execute_traced_relational(query, inputs, execute_concat) case "union": diff --git a/src/agent_memory/adapters/lotus/algebraic_aggregates.py b/src/agent_memory/adapters/lotus/algebraic_aggregates.py new file mode 100644 index 0000000..b6ff031 --- /dev/null +++ b/src/agent_memory/adapters/lotus/algebraic_aggregates.py @@ -0,0 +1,362 @@ +"""Deterministic execution for exact relational aggregate state.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from numbers import Number +from typing import Any + +import numpy as np +import pandas as pd + +from agent_memory.policy.aggregates import ( + AggregateSpec, + AvgAggregateSpec, + CountAggregateSpec, + SumAggregateSpec, + aggregate_output_names, + is_algebraic_aggregate_spec, +) +from agent_memory.policy.logical import QueryExpr + +def is_algebraic_aggregate_query(query: QueryExpr) -> bool: + """Return whether ``query`` contains only exact numeric aggregates.""" + + aggregates = tuple(query.params.get("aggregates", ())) + return ( + query.op == "agg" + and bool(aggregates) + and all(is_algebraic_aggregate_spec(spec) for spec in aggregates) + ) + + +def execute_algebraic_aggregate( + query: QueryExpr, + inputs: Mapping[str, Any], + execute: Callable[[QueryExpr, Mapping[str, Any]], Any], +) -> pd.DataFrame: + """Execute one complete global or ordinary grouped aggregate query.""" + + if not is_algebraic_aggregate_query(query): + raise TypeError("algebraic aggregate execution requires count/sum/avg specs") + source_query = query.inputs[0] + if source_query.op == "sem_groupby": + raise NotImplementedError( + "sem_groupby numeric algebraic aggregates are not supported" + ) + group_keys = ( + tuple(str(key) for key in source_query.params["keys"]) + if source_query.op == "group_by" + else () + ) + source = execute(source_query, inputs) + layout = _direct_layout( + group_keys=group_keys, + aggregates=tuple(query.params["aggregates"]), + ) + empty_source = source.iloc[0:0].copy() + state = _update_state( + pd.DataFrame(columns=layout["state_columns"]), + source, + empty_source, + layout, + ) + return _finalize_state(state, layout) + + +def execute_aggregate_state( + query: QueryExpr, + inputs: Mapping[str, Any], + execute: Callable[[QueryExpr, Mapping[str, Any]], Any], +) -> pd.DataFrame: + """Build complete hidden aggregate state from one input relation.""" + + source = execute(query.inputs[0], inputs) + return _update_state( + pd.DataFrame(columns=query.params["state_columns"]), + source, + source.iloc[0:0].copy(), + query.params, + ) + + +def execute_aggregate_state_update( + query: QueryExpr, + inputs: Mapping[str, Any], + execute: Callable[[QueryExpr, Mapping[str, Any]], Any], +) -> pd.DataFrame: + """Merge parent insertions and retractions into hidden aggregate state.""" + + if len(query.inputs) != 3: + raise ValueError("aggregate_state_update expects current, inserted, and retracted inputs") + current, inserted, retracted = ( + execute(input_query, inputs) for input_query in query.inputs + ) + return _update_state(current, inserted, retracted, query.params) + + +def execute_aggregate_finalize( + query: QueryExpr, + inputs: Mapping[str, Any], + execute: Callable[[QueryExpr, Mapping[str, Any]], Any], +) -> pd.DataFrame: + """Project hidden aggregate state into the declared public columns.""" + + state = execute(query.inputs[0], inputs) + return _finalize_state(state, query.params) + + +def _direct_layout( + *, + group_keys: tuple[str, ...], + aggregates: tuple[AggregateSpec, ...], +) -> dict[str, Any]: + occupied = set(group_keys) + row_count_col = _unique_hidden_name("row_count", occupied) + occupied.add(row_count_col) + numeric_columns: list[str] = [] + for aggregate in aggregates: + if isinstance(aggregate, (SumAggregateSpec, AvgAggregateSpec)): + if aggregate.column not in numeric_columns: + numeric_columns.append(aggregate.column) + numeric_states: list[dict[str, str]] = [] + for index, column in enumerate(numeric_columns): + sum_col = _unique_hidden_name(f"sum_{index}", occupied) + occupied.add(sum_col) + non_null_count_col = _unique_hidden_name( + f"non_null_count_{index}", occupied + ) + occupied.add(non_null_count_col) + numeric_states.append( + { + "input_col": column, + "sum_col": sum_col, + "non_null_count_col": non_null_count_col, + } + ) + state_columns = ( + *group_keys, + row_count_col, + *( + state_column + for numeric_state in numeric_states + for state_column in ( + numeric_state["sum_col"], + numeric_state["non_null_count_col"], + ) + ), + ) + return { + "group_keys": group_keys, + "aggregates": aggregates, + "row_count_col": row_count_col, + "numeric_states": tuple(numeric_states), + "state_columns": state_columns, + } + + +def _unique_hidden_name(suffix: str, occupied: set[str]) -> str: + candidate = f"__am_aggregate_{suffix}" + while candidate in occupied: + candidate = f"_{candidate}" + return candidate + + +def _update_state( + current: pd.DataFrame, + inserted: pd.DataFrame, + retracted: pd.DataFrame, + layout: Mapping[str, Any], +) -> pd.DataFrame: + group_keys = tuple(str(key) for key in layout["group_keys"]) + row_count_col = str(layout["row_count_col"]) + numeric_states = tuple(layout["numeric_states"]) + state_columns = tuple(str(column) for column in layout["state_columns"]) + _require_columns(current, state_columns, relation="aggregate state") + required_input = ( + *group_keys, + *(str(state["input_col"]) for state in numeric_states), + ) + _require_columns(inserted, required_input, relation="inserted rows") + _require_columns(retracted, required_input, relation="retracted rows") + + states: dict[tuple[Any, ...], dict[str, Any]] = {} + for record in current.to_dict(orient="records"): + key = _group_key(record, group_keys) + if key in states: + raise ValueError("aggregate state contains duplicate group keys") + states[key] = dict(record) + if not group_keys and not states: + states[()] = _empty_accumulator(layout, key_values={}) + + _apply_rows(states, inserted, layout, sign=1) + _apply_rows(states, retracted, layout, sign=-1) + + rows: list[dict[str, Any]] = [] + for key, state in states.items(): + row_count = _require_integral_state(state[row_count_col], column=row_count_col) + if row_count < 0: + raise ValueError(f"aggregate state row count became negative for group {key!r}") + for numeric_state in numeric_states: + non_null_col = str(numeric_state["non_null_count_col"]) + non_null_count = _require_integral_state( + state[non_null_col], column=non_null_col + ) + if non_null_count < 0 or non_null_count > row_count: + raise ValueError( + f"aggregate state non-null count is invalid for group {key!r}" + ) + if group_keys and row_count == 0: + continue + rows.append(state) + if not group_keys and not rows: + rows.append(_empty_accumulator(layout, key_values={})) + return pd.DataFrame(rows, columns=state_columns) + + +def _apply_rows( + states: dict[tuple[Any, ...], dict[str, Any]], + rows: pd.DataFrame, + layout: Mapping[str, Any], + *, + sign: int, +) -> None: + group_keys = tuple(str(key) for key in layout["group_keys"]) + row_count_col = str(layout["row_count_col"]) + numeric_states = tuple(layout["numeric_states"]) + for record in rows.to_dict(orient="records"): + key = _group_key(record, group_keys) + state = states.get(key) + if state is None: + state = _empty_accumulator( + layout, + key_values={key_name: record[key_name] for key_name in group_keys}, + ) + states[key] = state + state[row_count_col] = _require_integral_state( + state[row_count_col], column=row_count_col + ) + sign + for numeric_state in numeric_states: + input_col = str(numeric_state["input_col"]) + sum_col = str(numeric_state["sum_col"]) + non_null_col = str(numeric_state["non_null_count_col"]) + value = record[input_col] + if _is_null(value): + continue + _require_numeric(value, column=input_col) + state[sum_col] = state[sum_col] + sign * value + state[non_null_col] = _require_integral_state( + state[non_null_col], column=non_null_col + ) + sign + + +def _empty_accumulator( + layout: Mapping[str, Any], + *, + key_values: Mapping[str, Any], +) -> dict[str, Any]: + state = dict(key_values) + state[str(layout["row_count_col"])] = 0 + for numeric_state in layout["numeric_states"]: + state[str(numeric_state["sum_col"])] = 0 + state[str(numeric_state["non_null_count_col"])] = 0 + return state + + +def _finalize_state( + state: pd.DataFrame, + layout: Mapping[str, Any], +) -> pd.DataFrame: + group_keys = tuple(str(key) for key in layout["group_keys"]) + aggregates = tuple(layout["aggregates"]) + state_columns = tuple(str(column) for column in layout["state_columns"]) + _require_columns(state, state_columns, relation="aggregate state") + numeric_by_input = { + str(numeric_state["input_col"]): numeric_state + for numeric_state in layout["numeric_states"] + } + rows: list[dict[str, Any]] = [] + for state_record in state.to_dict(orient="records"): + output = {key: state_record[key] for key in group_keys} + for aggregate in aggregates: + if isinstance(aggregate, CountAggregateSpec): + output[aggregate.output_col] = _require_integral_state( + state_record[str(layout["row_count_col"])], + column=str(layout["row_count_col"]), + ) + continue + if isinstance(aggregate, (SumAggregateSpec, AvgAggregateSpec)): + numeric_state = numeric_by_input[aggregate.column] + non_null_count = _require_integral_state( + state_record[str(numeric_state["non_null_count_col"])], + column=str(numeric_state["non_null_count_col"]), + ) + value = ( + None + if non_null_count == 0 + else state_record[str(numeric_state["sum_col"])] + ) + if isinstance(aggregate, AvgAggregateSpec) and value is not None: + value = value / non_null_count + output[aggregate.output_col] = value + continue + raise TypeError(f"Unsupported aggregate spec: {type(aggregate).__name__}") + rows.append(output) + columns = [*group_keys, *(name for spec in aggregates for name in aggregate_output_names(spec))] + return pd.DataFrame(rows, columns=columns) + + +def _group_key( + record: Mapping[str, Any], group_keys: Sequence[str] +) -> tuple[Any, ...]: + return tuple(_key_value(record[key]) for key in group_keys) + + +def _key_value(value: Any) -> Any: + if _is_null(value): + return ("null",) + try: + hash(value) + except TypeError as error: + raise TypeError("aggregate group key values must be hashable") from error + return (type(value).__module__, type(value).__qualname__, value) + + +def _require_columns( + frame: pd.DataFrame, + columns: Sequence[str], + *, + relation: str, +) -> None: + missing = [column for column in columns if column not in frame.columns] + if missing: + raise ValueError(f"{relation} is missing columns: {missing}") + + +def _require_numeric(value: Any, *, column: str) -> None: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Number): + raise TypeError( + f"aggregate input column {column!r} requires numeric non-null values" + ) + + +def _require_integral_state(value: Any, *, column: str) -> int: + if isinstance(value, (bool, np.bool_)) or not isinstance(value, (int, np.integer)): + raise TypeError(f"aggregate state column {column!r} must contain integers") + return int(value) + + +def _is_null(value: Any) -> bool: + if value is None or value is pd.NA: + return True + result = pd.isna(value) + return isinstance(result, (bool, np.bool_)) and bool(result) + + +__all__ = [ + "execute_aggregate_finalize", + "execute_aggregate_state", + "execute_aggregate_state_update", + "execute_algebraic_aggregate", + "is_algebraic_aggregate_query", +] diff --git a/src/agent_memory/adapters/lotus/context.py b/src/agent_memory/adapters/lotus/context.py index 2c5fb1d..97171e9 100644 --- a/src/agent_memory/adapters/lotus/context.py +++ b/src/agent_memory/adapters/lotus/context.py @@ -10,6 +10,10 @@ from agent_memory.adapters.lotus.pair_execution import ( SemanticPairExecutionProfile, ) +from agent_memory.adapters.lotus.prompt_batching import ( + PromptBatching, + validate_structured_output_transport, +) from agent_memory.adapters.lotus.provider_usage_lm import provider_usage_tracing_lm_class from agent_memory.adapters.lotus.traced_lm import TracedLM from agent_memory.storage.embedding import EmbeddingProvider @@ -25,6 +29,7 @@ "listwise", ) SEM_JOIN_TOPK_METHODS = SEM_TOPK_METHODS +SEM_AGG_DISPATCH_METHODS = ("sequential", "provider-batched") _LM_OWNED_KWARGS = { "cache", "max_batch_size", @@ -35,6 +40,15 @@ } +def validate_lm_max_ctx_len(value: int | None) -> None: + """Validate an optional client-side context capacity override.""" + + if value is not None and ( + isinstance(value, bool) or not isinstance(value, int) or value <= 0 + ): + raise ValueError("lm_max_ctx_len must be a positive integer or None") + + @dataclass(frozen=True) class LotusExecutionConfig: """Backend execution knobs that are not part of policy query semantics.""" @@ -45,6 +59,7 @@ class LotusExecutionConfig: lm_rate_limit: int | None = None lm_model_kwargs: Mapping[str, Any] = field(default_factory=dict) lm_enable_cache: bool | None = None + # Fixed per-request ceiling for structured output, independent of task count. structured_max_tokens: int = DEFAULT_STRUCTURED_MAX_TOKENS structured_parse_retries: int = DEFAULT_STRUCTURED_PARSE_RETRIES semantic_trace_dir: Path | str | None = None @@ -52,6 +67,8 @@ class LotusExecutionConfig: semantic_pair_profiles: Mapping[str, SemanticPairExecutionProfile] = field( default_factory=dict ) + prompt_batching: PromptBatching | None = None + structured_output_transport: str = "chat-json-object" sem_filter_examples: Sequence[Mapping[str, Any]] | None = None sem_filter_helper_examples: Sequence[Mapping[str, Any]] | None = None @@ -91,10 +108,33 @@ class LotusExecutionConfig: sem_agg_safe_mode: bool = False sem_agg_progress_bar_desc: str = "Aggregating" sem_agg_model_kwargs: Mapping[str, Any] = field(default_factory=dict) + sem_agg_dispatch: str = "sequential" def __post_init__(self) -> None: """Validate bounded semantic execution settings.""" + validate_lm_max_ctx_len(self.lm_model_kwargs.get("max_ctx_len")) + validate_structured_output_transport(self.structured_output_transport) + if self.prompt_batching is not None and not isinstance( + self.prompt_batching, PromptBatching + ): + raise TypeError("prompt_batching must be PromptBatching or None") + if ( + self.prompt_batching is not None + and self.sem_agg_dispatch != "sequential" + ): + raise ValueError( + "prompt_batching cannot be combined with sem_agg_dispatch" + ) + if ( + self.prompt_batching is not None + and self.sem_groupby_pair_batch_size is not None + ): + raise ValueError( + "prompt_batching cannot be combined with " + "sem_groupby_pair_batch_size" + ) + if ( self.sem_groupby_pair_batch_size is not None and self.sem_groupby_pair_batch_size < 1 @@ -111,6 +151,11 @@ def __post_init__(self) -> None: "sem_join_topk_method must be one of: " + ", ".join(SEM_JOIN_TOPK_METHODS) ) + if self.sem_agg_dispatch not in SEM_AGG_DISPATCH_METHODS: + raise ValueError( + "sem_agg_dispatch must be one of: " + + ", ".join(SEM_AGG_DISPATCH_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") @@ -141,6 +186,13 @@ class LotusExecutionContext: repr=False, ) + def __post_init__(self) -> None: + """Reject unsupported transport/model combinations before configuring LOTUS.""" + + validate_structured_output_transport( + self.config.structured_output_transport, model=self.model + ) + def configure(self) -> None: """Configure LOTUS before invoking semantic dataframe operators.""" @@ -174,10 +226,17 @@ def configure(self) -> None: ) trace_dir = self.config.trace_dir() + lm_class = LM + if self.config.structured_output_transport == "responses-json-schema": + from agent_memory.adapters.lotus.deepseek_responses_lm import ( + deepseek_responses_lm_class, + ) + + lm_class = deepseek_responses_lm_class(LM) base_lm = ( - provider_usage_tracing_lm_class(LM)(**lm_kwargs, trace_dir=trace_dir) - if trace_dir is not None - else LM(**lm_kwargs) + provider_usage_tracing_lm_class(lm_class)(**lm_kwargs, trace_dir=trace_dir) + if trace_dir is not None or self.config.prompt_batching is not None + else lm_class(**lm_kwargs) ) lm = TracedLM(base_lm, trace_dir) if trace_dir is not None else base_lm self._lm = lm diff --git a/src/agent_memory/adapters/lotus/deepseek_responses_lm.py b/src/agent_memory/adapters/lotus/deepseek_responses_lm.py new file mode 100644 index 0000000..d6e56fd --- /dev/null +++ b/src/agent_memory/adapters/lotus/deepseek_responses_lm.py @@ -0,0 +1,451 @@ +"""Native DeepSeek Responses transport for LOTUS JSON-schema requests. + +Wrap this factory's result with provider_usage_tracing_lm_class, not vice versa. +The uncached hook returns ordered ModelResponse/OpenAIError items before LOTUS +accounts usage. Every item carries provider_response_metadata with transport, +status, incomplete_reason, finish_reason, and raw_usage. Native response failures +also expose response (the normalized partial ModelResponse), native_response +(the original SDK Response), model, and usage when reported. Normalization +errors retain those fields with finish_reason="normalization_error" and the +original exception as their cause; SDK errors retain their type and identity. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +import math +import os +import time +from typing import Any + +from litellm import ModelResponse +from litellm.types.utils import Usage +from openai import OpenAI, OpenAIError +from openai.types.responses import Response +from tqdm import tqdm + +from .request_hook import next_request_batch + +__all__ = ["deepseek_responses_lm_class", "validate_deepseek_responses_model"] + +_MODELS = {"deepseek-flash", "deepseek-v4-flash", "deepseek-v4-pro"} +_PARAMS = { + "response_format", + "max_tokens", + "temperature", + "top_p", + "timeout", + "num_retries", + "extra_body", + "reasoning_effort", + "stream", + "api_key", + "api_base", + "user", +} +_EFFORTS = { + "none": "none", + "low": "low", + "medium": "high", + "high": "high", + "xhigh": "high", + "max": "max", +} +_TRANSPORT = "responses-json-schema" + + +def validate_deepseek_responses_model(model: str) -> None: + """Reject non-native providers and unsupported native DeepSeek model names.""" + if not isinstance(model, str) or model.removeprefix("deepseek/") not in _MODELS: + raise ValueError( + f"Unsupported DeepSeek Responses model: {model!r}; use " + f"{', '.join(sorted(_MODELS))}, optionally prefixed with deepseek/" + ) + + +def _reasoning_effort(kwargs: dict[str, Any]) -> str: + extra = kwargs.get("extra_body", {}) + if not isinstance(extra, Mapping) or extra.keys() - {"thinking"}: + raise ValueError("Responses extra_body supports only thinking.type") + thinking = extra.get("thinking", {}) + if not isinstance(thinking, Mapping) or thinking.keys() - {"type"}: + raise ValueError("Responses thinking supports only type") + mode = thinking.get("type") + if "thinking" in extra and mode not in {"enabled", "disabled"}: + raise ValueError("thinking.type must be enabled or disabled") + effort = kwargs.get("reasoning_effort", "none" if mode == "disabled" else "high") + if not isinstance(effort, str) or effort not in _EFFORTS: + raise ValueError("Unsupported reasoning_effort for DeepSeek Responses") + if (mode == "disabled" and effort != "none") or ( + mode == "enabled" and effort == "none" + ): + raise ValueError("thinking.type conflicts with reasoning_effort") + return _EFFORTS[effort] + + +def _request_options(kwargs: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: + unsupported = kwargs.keys() - _PARAMS + if unsupported: + raise ValueError( + f"Unsupported DeepSeek Responses parameters: {', '.join(sorted(unsupported))}" + ) + if kwargs.get("stream", False) is not False: + raise ValueError("DeepSeek Responses adapter supports only stream=False") + response_format = kwargs["response_format"] + schema = response_format.get("json_schema") + if ( + response_format.keys() - {"type", "json_schema"} + or not isinstance(schema, Mapping) + or schema.keys() - {"name", "schema", "strict", "description"} + or not isinstance(schema.get("name"), str) + or not schema["name"] + or not isinstance(schema.get("schema"), Mapping) + ): + raise ValueError( + "response_format requires json_schema.name and json_schema.schema" + ) + if "strict" in schema and not isinstance(schema["strict"], bool): + raise ValueError("json_schema.strict must be a boolean") + if "description" in schema and not isinstance(schema["description"], str): + raise ValueError("json_schema.description must be a string") + + retries = kwargs.get("num_retries", 2) + if type(retries) is not int or not 0 <= retries <= 10: + raise ValueError("num_retries must be an integer between 0 and 10") + payload: dict[str, Any] = { + "text": {"format": {"type": "json_schema", **schema}}, + "reasoning": {"effort": _reasoning_effort(kwargs)}, + "stream": False, + } + if "max_tokens" in kwargs: + tokens = kwargs["max_tokens"] + if type(tokens) is not int or tokens <= 0: + raise ValueError("max_tokens must be a positive integer") + payload["max_output_tokens"] = tokens + for key, maximum in (("temperature", 2), ("top_p", 1), ("timeout", None)): + if key not in kwargs: + continue + value = kwargs[key] + if ( + type(value) not in (int, float) + or not math.isfinite(value) + or (value <= 0 if key == "timeout" else value < 0) + or (maximum is not None and value > maximum) + ): + raise ValueError(f"Unsupported {key} value for DeepSeek Responses") + if key != "timeout": + payload[key] = value + if "user" in kwargs: + if not isinstance(kwargs["user"], str): + raise ValueError("user must be a string") + payload["user"] = kwargs["user"] + api_key = kwargs.get("api_key") or os.environ.get("DEEPSEEK_API_KEY") + if not isinstance(api_key, str) or not api_key.strip(): + raise ValueError("Set DEEPSEEK_API_KEY or pass api_key for DeepSeek Responses") + base_url = kwargs.get("api_base", "https://api.deepseek.com") + if not isinstance(base_url, str) or not base_url.startswith( + ("https://", "http://") + ): + raise ValueError("api_base must be an HTTP(S) base URL") + client_options = {"api_key": api_key, "base_url": base_url, "max_retries": retries} + if "timeout" in kwargs: + client_options["timeout"] = kwargs["timeout"] + return payload, client_options + + +def _normalize_usage(raw: dict[str, Any] | None) -> Usage | None: + if raw is None: + return None + details = raw.get("input_tokens_details") or {} + output_details = raw.get("output_tokens_details") or {} + kwargs: dict[str, Any] = { + "prompt_tokens": raw["input_tokens"], + "completion_tokens": raw["output_tokens"], + "total_tokens": raw["total_tokens"], + } + cached = details.get("cached_tokens") + if cached is not None: + kwargs.update( + prompt_cache_hit_tokens=cached, + prompt_cache_miss_tokens=raw["input_tokens"] - cached, + prompt_tokens_details={"cached_tokens": cached}, + ) + if output_details.get("reasoning_tokens") is not None: + kwargs["completion_tokens_details"] = { + "reasoning_tokens": output_details["reasoning_tokens"] + } + return Usage(**kwargs) + + +class _ResponseFailure(OpenAIError): + """A non-success native response that still incurred provider usage.""" + + def __init__( + self, response: ModelResponse, detail: Any, native_response: Response + ) -> None: + dynamic_response: Any = response + metadata: dict[str, Any] = dynamic_response.provider_response_metadata + super().__init__(f"DeepSeek Responses {metadata['finish_reason']}: {detail}") + self.response = response + self.native_response = native_response + self.model = response.model + self.provider_response_metadata = metadata + if hasattr(response, "usage"): + self.usage = dynamic_response.usage + + +def _response_envelope(response: Response, finish: str, content: str) -> ModelResponse: + # Read accounting fields independently of output-item serialization/parsing. + raw_usage = ( + response.usage.model_dump(mode="json") if response.usage is not None else None + ) + return ModelResponse( + id=response.id, + created=int(response.created_at), + model=response.model, + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": finish, + } + ], + usage=_normalize_usage(raw_usage), + provider_response_metadata={ + "transport": _TRANSPORT, + "status": response.status, + "incomplete_reason": ( + response.incomplete_details.reason + if response.incomplete_details is not None + else None + ), + "finish_reason": finish, + "raw_usage": raw_usage, + }, + ) + + +def _normalize_response(response: Response) -> ModelResponse | OpenAIError: + raw = response.model_dump(mode="json") + status = raw.get("status") + reason = (raw.get("incomplete_details") or {}).get("reason") + text: list[str] = [] + refusal: list[str] = [] + invalid_output = False + for item in raw.get("output", []): + if item["type"] == "reasoning": + continue + if item["type"] != "message": + invalid_output = True + continue + if item.get("status") != "completed": + invalid_output = True + for part in item["content"]: + if part["type"] == "output_text": + text.append(part["text"]) + elif part["type"] == "refusal": + refusal.append(part["refusal"]) + else: + invalid_output = True + content = "".join(text) + if status == "incomplete": + finish = { + "max_output_tokens": "length", + "content_filter": "content_filter", + }.get(reason or "", "incomplete") + elif status == "failed" or raw.get("error"): + finish = "error" + elif status != "completed": + finish = status or "unknown" + elif refusal: + finish = "content_filter" + elif invalid_output: + finish = "unsupported_output" + elif not content.strip(): + finish = "empty_output" + else: + finish = "stop" + normalized = _response_envelope(response, finish, content) + if finish != "stop": + return _ResponseFailure( + normalized, raw.get("error") or refusal or reason or status, response + ) + return normalized + + +class _DeepSeekResponsesMixin: + """Replace only the JSON-schema network path of LOTUS 1.1.4.""" + + model: str + max_batch_size: int + rate_limit: int | None + + def __init__(self, model: str, *args: Any, **kwargs: Any) -> None: + validate_deepseek_responses_model(model) + parent: Any = super() + parent.__init__(model, *args, **kwargs) + + def _process_uncached_messages( + self, + uncached_data: list[tuple[list[dict[str, str]], str]], + all_kwargs: dict[str, Any], + show_progress_bar: bool, + progress_bar_desc: str, + ) -> list[Any]: + recorded_batch = next_request_batch() + response_format = all_kwargs.get("response_format") + if ( + not isinstance(response_format, Mapping) + or response_format.get("type") != "json_schema" + ): + if recorded_batch is not None: + raise ValueError( + "durable requests require Responses JSON schema; no fallback send" + ) + parent: Any = super() + return parent._process_uncached_messages( + uncached_data, all_kwargs, show_progress_bar, progress_bar_desc + ) + if not uncached_data: + return [] + validate_deepseek_responses_model(self.model) + payload, client_options = _request_options(all_kwargs) + if recorded_batch is not None and client_options.get("max_retries", 0) != 0: + raise ValueError("durable requests require SDK retries=0") + batch = [messages for messages, _ in uncached_data] + for messages in batch: + if not messages or any( + not isinstance(message, Mapping) + or message.keys() != {"role", "content"} + or message["role"] not in {"system", "developer", "user", "assistant"} + or not isinstance(message["content"], str) + for message in messages + ): + raise ValueError( + "DeepSeek Responses supports text-only role/content messages" + ) + + with ( + OpenAI(**client_options) as client, + ThreadPoolExecutor(max_workers=self.max_batch_size) as executor, + tqdm( + total=len(batch), desc=progress_bar_desc, disable=not show_progress_bar + ) as pbar, + ): + responses_api: Any = client.responses + + def request( + item: tuple[int, list[dict[str, str]]], + ) -> ModelResponse | OpenAIError: + index, messages = item + if recorded_batch is not None: + hook, identity = recorded_batch + body = { + "model": self.model.removeprefix("deepseek/"), + "input": messages, + **payload, + } + + def send() -> dict[str, Any]: + return responses_api.create(**body).model_dump(mode="json") + + # Journal/control failures must reach the caller, not become False. + raw = hook.execute({**identity, "item": index}, body, send) + return _normalize_response(Response.model_validate(raw)) + try: + native_response: Response = responses_api.create( + model=self.model.removeprefix("deepseek/"), + input=messages, + **payload, + ) + except Exception as error: + # Like LiteLLM batch_completion, keep failures in input order. + failure: Any = ( + error + if isinstance(error, OpenAIError) + else OpenAIError(str(error)) + ) + if failure is not error: + failure.__cause__ = error + failure.provider_response_metadata = { + "transport": _TRANSPORT, + "status": "failed", + "incomplete_reason": None, + "finish_reason": "error", + "raw_usage": None, + } + return failure + try: + return _normalize_response(native_response) + except Exception as error: + try: + envelope = _response_envelope( + native_response, "normalization_error", "" + ) + except Exception as accounting_error: + error.add_note( + "Response accounting unavailable: " + f"{type(accounting_error).__name__}: {accounting_error}" + ) + # Do not read malformed native fields again in this fallback. + envelope = ModelResponse( + model=self.model, + created=0, + choices=[], + usage=None, + provider_response_metadata={ + "transport": _TRANSPORT, + "status": None, + "incomplete_reason": None, + "finish_reason": "normalization_error", + "raw_usage": None, + }, + ) + normalization_failure = _ResponseFailure( + envelope, error, native_response + ) + normalization_failure.__cause__ = error + return normalization_failure + + responses: list[Any] = [] + batch_size = ( + self.max_batch_size if self.rate_limit is not None else len(batch) + ) + for start in range(0, len(batch), batch_size): + started = time.monotonic() + sub_batch = batch[start : start + batch_size] + responses.extend(executor.map(request, enumerate(sub_batch, start))) + pbar.update(len(sub_batch)) + elapsed = time.monotonic() - started + if self.rate_limit is not None and start + batch_size < len(batch): + delay = len(sub_batch) * 60 / self.rate_limit - elapsed + if delay > 0: + time.sleep(delay) + return responses + + def _cache_response(self, response: Any, hash: str) -> None: + if ( + isinstance(response, OpenAIError) + and getattr(response, "provider_response_metadata", {}).get("transport") + == _TRANSPORT + ): + # LOTUS must finish accounting the batch before _get_top_choice raises. + return + parent: Any = super() + parent._cache_response(response, hash) + + def _update_stats(self, response: Any, is_cached: bool = False) -> None: + if isinstance(response, _ResponseFailure): + response = response.response + parent: Any = super() + parent._update_stats(response, is_cached=is_cached) + + +def deepseek_responses_lm_class(base: type[Any]) -> type[Any]: + """Return a LOTUS subclass routing only json_schema calls to native Responses.""" + + class DeepSeekResponsesLM(_DeepSeekResponsesMixin, base): + pass + + return DeepSeekResponsesLM diff --git a/src/agent_memory/adapters/lotus/json_output.py b/src/agent_memory/adapters/lotus/json_output.py new file mode 100644 index 0000000..a5675f8 --- /dev/null +++ b/src/agent_memory/adapters/lotus/json_output.py @@ -0,0 +1,245 @@ +"""JSON loading and bounded, value-preserving structural repair.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from dataclasses import dataclass +import heapq +from itertools import count +import json +import re +from typing import Any, Generic, TypeVar + +import json5 + + +JSON_REPAIR_VERSION = "bounded-json-v2" +MAX_REPAIR_CANDIDATES = 4096 +MAX_REPAIR_EDITS = 8 +MAX_REPAIR_CHARACTERS = 8 * 1024 * 1024 +_DELIMITERS = "}],:" +_TOKEN = re.compile( + r'\s+|"(?:[^"\\\x00-\x1f]|\\(?:["\\/bfnrt]|u[0-9a-fA-F]{4}))*"' + r"|-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?" + r"|true|false|null|[{}\[\],:]" +) +ResultT = TypeVar("ResultT") + + +@dataclass(frozen=True) +class StructuredJsonResult: + """A decoded value and any existing JSON5 normalization provenance.""" + + value: Any + repair_method: str | None = None + + +@dataclass(frozen=True) +class RepairedJson(Generic[ResultT]): + """A unique minimum-edit result; edit offsets apply sequentially.""" + + value: ResultT + edits: tuple[tuple[str, int, str], ...] + + +def strict_json_loads(value: str) -> Any: + """Decode JSON without duplicate object keys or non-finite constants.""" + + return json.loads( + value, object_pairs_hook=_unique_object, parse_constant=_reject_constant + ) + + +def load_structured_json_with_syntax_repair( + raw_output: str, *, operator: str, expected_shape: str +) -> StructuredJsonResult: + """Preserve the established strict JSON, fence and complete JSON5 path.""" + + text = raw_output.strip() + try: + return StructuredJsonResult(strict_json_loads(text)) + except ValueError: + pass + candidate, fenced = _strip_fence(text) + try: + value = json5.loads( + candidate, + allow_duplicate_keys=False, + consume_trailing=True, + parse_constant=_reject_constant, + ) + except ValueError as error: + preview = raw_output[:240] + ("..." if len(raw_output) > 240 else "") + raise ValueError( + f"{operator} returned invalid JSON; expected {expected_shape}; " + f"raw_output={preview!r}" + ) from error + return StructuredJsonResult(value, "json5-code-fence" if fenced else "json5") + + +def repair_json_structure( + raw_output: str, *, validator: Callable[[str], ResultT] +) -> RepairedJson[ResultT]: + """Accept a unique minimum-edit result within a bounded, forward-only search.""" + + # Valid JSON/JSON5 with a bad schema is not a syntax-repair problem. + try: + load_structured_json_with_syntax_repair( + raw_output, operator="repair", expected_shape="JSON" + ) + except ValueError: + pass + else: + raise ValueError("valid JSON contract errors cannot be structurally repaired") + + text, _ = _strip_fence(raw_output.strip()) + tokens = _tokens(text) + scalars = _scalars(tokens) + try: + strict_json_loads(text) + except json.JSONDecodeError as error: + position = error.pos + except ValueError as error: + raise ValueError( + "duplicate keys or non-finite values cannot be repaired" + ) from error + else: + raise ValueError("output does not require syntax repair") + + serial = count() + pending: list[tuple[int, int, str, tuple[tuple[str, int, str], ...], int]] = [ + (0, next(serial), text, (), position) + ] + seen = {tuple(token for token, _, _ in tokens): 0} + accepted: dict[tuple[str, ...], RepairedJson[ResultT]] = {} + best_cost = MAX_REPAIR_EDITS + 1 + examined = 0 + characters = 0 + while pending: + cost, _, current, edits, position = heapq.heappop(pending) + if cost >= min(best_cost, MAX_REPAIR_EDITS): + continue + for candidate, edit in _local_candidates(current, position).items(): + candidate_cost = cost + len(edit[2]) + if candidate_cost > min(best_cost, MAX_REPAIR_EDITS): + continue + characters += len(candidate) + if characters > MAX_REPAIR_CHARACTERS: + raise ValueError("JSON repair character-work bound exceeded") + candidate_tokens = _tokens(candidate) + # The same delimiter moved across whitespace is one candidate. + key = tuple(token for token, _, _ in candidate_tokens) + if seen.get(key, MAX_REPAIR_EDITS + 1) <= candidate_cost: + continue + seen[key] = candidate_cost + examined += 1 + if examined > MAX_REPAIR_CANDIDATES: + raise ValueError("JSON repair candidate bound exceeded") + if _scalars(candidate_tokens) != scalars: + continue + candidate_edits = (*edits, edit) + try: + strict_json_loads(candidate) + except json.JSONDecodeError as error: + # Resolve the current error before attempting the next one. + if error.pos > position or ( + error.pos == position and len(candidate) < len(current) + ): + heapq.heappush( + pending, + ( + candidate_cost, + next(serial), + candidate, + candidate_edits, + error.pos, + ), + ) + continue + except ValueError: + continue + try: + value = validator(candidate) + except ValueError: + # Syntactically valid contract errors must never be edited further. + continue + if candidate_cost < best_cost: + best_cost = candidate_cost + accepted.clear() + accepted[key] = RepairedJson(value, candidate_edits) + if len(accepted) != 1: + raise ValueError( + f"JSON repair requires one valid candidate at minimum edits; got {len(accepted)}" + ) + return next(iter(accepted.values())) + + +def _local_candidates(text: str, position: int) -> dict[str, tuple[str, int, str]]: + tokens = _tokens(text) + nearby = [i for i, (_, start, end) in enumerate(tokens) if start <= position <= end] + anchor = nearby[-1] if nearby else len(tokens) - 1 + boundaries = {position} + for _, start, end in tokens[max(0, anchor - 2) : anchor + 2]: + boundaries.update((start, end)) + # Never insert inside a string, number, boolean, or null token. + boundaries = { + p for p in boundaries if not any(start < p < end for _, start, end in tokens) + } + candidates: dict[str, tuple[str, int, str]] = {} + for p in sorted(boundaries): + if p < len(text) and text[p] in _DELIMITERS: + candidates[text[:p] + text[p + 1 :]] = ("delete", p, text[p]) + for delimiter in _DELIMITERS: + candidates[text[:p] + delimiter + text[p:]] = ("insert", p, delimiter) + if position == len(text): + suffix = _closing_suffix(tokens) + if suffix: + candidates[text + suffix] = ("insert", len(text), suffix) + return candidates + + +def _tokens(text: str) -> list[tuple[str, int, int]]: + tokens: list[tuple[str, int, int]] = [] + position = 0 + while position < len(text): + match = _TOKEN.match(text, position) + if match is None: + raise ValueError("incomplete string or scalar cannot be repaired") + token = match.group() + if not token.isspace(): + tokens.append((token, position, match.end())) + position = match.end() + return tokens + + +def _scalars(tokens: Sequence[tuple[str, int, int]]) -> tuple[str, ...]: + return tuple(token for token, _, _ in tokens if token not in "{}[],:") + + +def _closing_suffix(tokens: Sequence[tuple[str, int, int]]) -> str: + stack: list[str] = [] + for token, _, _ in tokens: + if token in ("{", "["): + stack.append("}" if token == "{" else "]") + elif token in ("}", "]"): + if not stack or stack.pop() != token: + return "" + return "".join(reversed(stack)) + + +def _unique_object(pairs: Sequence[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON object key: {key!r}") + result[key] = value + return result + + +def _reject_constant(value: str) -> Any: + raise ValueError(f"non-finite JSON value is not supported: {value}") + + +def _strip_fence(value: str) -> tuple[str, bool]: + match = re.fullmatch(r"```(?:json)?[ \t]*\n(.*)\n```", value, flags=re.DOTALL) + return (value, False) if match is None else (match.group(1).strip(), True) diff --git a/src/agent_memory/adapters/lotus/predicate.py b/src/agent_memory/adapters/lotus/predicate.py new file mode 100644 index 0000000..d69226a --- /dev/null +++ b/src/agent_memory/adapters/lotus/predicate.py @@ -0,0 +1,177 @@ +"""Single-context structured predicates, independent of prompt batching.""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +from typing import Any + +import pandas as pd + +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.structured import execute_structured_lm_retry_result +from agent_memory.policy.logical import ColumnSpec +from agent_memory.tracing.semantic import write_trace_event + + +PREDICATE_SCHEMA = { + "type": "object", + "properties": {"keep": {"type": "boolean"}}, + "required": ["keep"], + "additionalProperties": False, +} +PREDICATE_SYSTEM_PROMPT = ( + "The user will provide a claim and relevant context. Determine whether the " + "claim is true for that context. Return only a JSON object with the boolean " + "field keep: true when the claim is true, false otherwise." +) + + +def single_schema_predicate(config: LotusExecutionConfig) -> bool: + """Select only the explicitly requested non-batching schema path.""" + + return ( + config.prompt_batching is None + and config.structured_output_transport == "responses-json-schema" + ) + + +def validate_schema_predicate(config: LotusExecutionConfig, operator: str) -> None: + """Reject unsupported native options rather than silently dropping them.""" + + names = ["examples", "strategy", "cascade_args", "safe_mode"] + if operator == "sem_filter": + names += ["helper_examples", "additional_cot_instructions"] + unsupported = [] + for name in names: + value = getattr(config, f"{operator}_{name}") + if ( + value is not None + and value is not False + and not (isinstance(value, str) and value == "") + ): + unsupported.append(name) + default = getattr(config, f"{operator}_default") + if default != (operator == "sem_filter"): + unsupported.append("default") + if unsupported: + raise ValueError( + f"single-task schema {operator} does not support: {', '.join(unsupported)}" + ) + + +def _validate_output(raw: str) -> None: + value = json.loads(raw) + if ( + not isinstance(value, dict) + or set(value) != {"keep"} + or type(value["keep"]) is not bool + ): + raise ValueError( + 'predicate output must contain exactly one boolean field "keep"' + ) + + +@dataclass(frozen=True) +class PredicateResult: + """Aligned decisions and raw attempts for existing operator trace assembly.""" + + decisions: tuple[bool, ...] + raw_outputs: tuple[str, ...] + + +def execute_schema_predicate( + source: pd.DataFrame, + *, + instruction: str, + config: LotusExecutionConfig, + operator: str, +) -> PredicateResult: + """Evaluate one context per prompt using the shared structured retry executor.""" + + import lotus + from lotus.nl_expression import parse_cols + from lotus.templates import task_instructions + + validate_schema_predicate(config, operator) + columns = list(parse_cols(instruction)) + if not columns or any(column not in source.columns for column in columns): + raise ValueError( + "schema predicate requires valid instruction column placeholders" + ) + if source.empty: + return PredicateResult((), ()) + docs = task_instructions.df2multimodal_info(source, columns) + if any(doc.get("image") for doc in docs): + raise NotImplementedError("single-task schema predicates support text only") + prompts = [ + [ + {"role": "system", "content": PREDICATE_SYSTEM_PROMPT}, + { + "role": "user", + "content": f"Context:\n{doc['text']}\n\nClaim: {instruction}", + }, + ] + for doc in docs + ] + payload: dict[str, Any] = { + "instruction": instruction, + "structured_output_transport": config.structured_output_transport, + "predicate_prompt_mode": "single-task", + "tasks_per_prompt": 1, + "prompt_count": len(prompts), + } + try: + result = execute_structured_lm_retry_result( + lotus.settings.lm, + prompts, + lm_kwargs={ + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "predicate", + "strict": True, + "schema": PREDICATE_SCHEMA, + }, + }, + "max_tokens": max( + int(getattr(lotus.settings.lm, "max_tokens", 512) or 512), + config.structured_max_tokens, + ), + "progress_bar_desc": getattr(config, f"{operator}_progress_bar_desc"), + }, + output_cols=(ColumnSpec("keep", "Whether the claim is true."),), + shape="object", + require_explanation=False, + operator=operator, + max_retries=config.structured_parse_retries, + output_validator=_validate_output, + ) + except Exception as error: + write_trace_event( + config.trace_dir(), + operator=operator, + event_type="predicate_failure", + payload={**payload, "error": str(error)}, + raw_output=str(error), + ) + raise + write_trace_event( + config.trace_dir(), + operator=operator, + event_type="predicate_generation", + payload={**payload, "invalid_indices": list(result.invalid_indices)}, + raw_output={"prompts": prompts, "attempts": result.raw_output_attempts}, + parsed_output=[ + None if index in result.invalid_indices else json.loads(raw)["keep"] + for index, raw in enumerate(result.raw_outputs) + ], + ) + if result.invalid_indices: + raise ValueError( + f"{operator} schema predicate failed for rows {result.invalid_indices}; artifacts: {result.failure_artifact_paths}" + ) + return PredicateResult( + tuple(json.loads(raw)["keep"] for raw in result.raw_outputs), + tuple(result.raw_outputs), + ) diff --git a/src/agent_memory/adapters/lotus/prompt_batching.py b/src/agent_memory/adapters/lotus/prompt_batching.py new file mode 100644 index 0000000..61451fc --- /dev/null +++ b/src/agent_memory/adapters/lotus/prompt_batching.py @@ -0,0 +1,555 @@ +"""Shared execution mechanics for multi-task semantic prompts.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, field, replace +from hashlib import sha256 +import json +from time import perf_counter +from typing import Any, Generic, TypeVar + +from agent_memory.adapters.lotus.json_output import ( + JSON_REPAIR_VERSION, + repair_json_structure, +) +from agent_memory.tracing.semantic import write_trace_event + + +TaskT = TypeVar("TaskT") +ResultT = TypeVar("ResultT") +STRUCTURED_OUTPUT_TRANSPORTS = ("chat-json-object", "responses-json-schema") + + +def validate_structured_output_transport( + transport: str, *, model: str | None = None +) -> None: + """Validate an explicitly selected physical structured-output path.""" + + if transport not in STRUCTURED_OUTPUT_TRANSPORTS: + raise ValueError(f"unsupported structured_output_transport: {transport!r}") + if transport == "responses-json-schema" and model is not None: + from agent_memory.adapters.lotus.deepseek_responses_lm import ( + validate_deepseek_responses_model, + ) + + validate_deepseek_responses_model(model) + + +@dataclass(frozen=True) +class PromptBatching: + """Bound the ready semantic tasks placed in one prompt.""" + + max_tasks: int | None = None + repair_version: str = field(default=JSON_REPAIR_VERSION, init=False) + + def __post_init__(self) -> None: + """Validate the optional task-count limit.""" + + if self.max_tasks is None: + return + if not isinstance(self.max_tasks, int) or isinstance(self.max_tasks, bool): + raise TypeError("PromptBatching.max_tasks must be an integer or None") + if self.max_tasks < 1: + raise ValueError("PromptBatching.max_tasks must be at least 1") + + def to_dict(self) -> dict[str, int | None]: + """Return the canonical physical execution contract.""" + + return {"max_tasks": self.max_tasks} + + @property + def fingerprint(self) -> str: + """Return a stable identity for checkpoint isolation.""" + + payload = json.dumps( + {**self.to_dict(), "repair_version": self.repair_version}, + sort_keys=True, + separators=(",", ":"), + ) + return sha256(payload.encode("utf-8")).hexdigest() + + +def parse_prompt_batch_size(value: str) -> PromptBatching: + """Parse the benchmark CLI's ``all`` or positive-integer contract.""" + + normalized = value.strip().lower() + if normalized == "all": + return PromptBatching() + try: + max_tasks = int(normalized) + except ValueError as error: + raise ValueError("expected 'all' or a positive integer") from error + try: + return PromptBatching(max_tasks=max_tasks) + except (TypeError, ValueError) as error: + raise ValueError("expected 'all' or a positive integer") from error + + +@dataclass(frozen=True) +class PromptBatchRequest: + """One model request built from a stable sequence of semantic tasks.""" + + task_ids: tuple[str, ...] + prompt: list[dict[str, str]] + max_tokens: int + + +@dataclass(frozen=True) +class PromptBatchItem(Generic[ResultT]): + """One task result returned by an operator-specific parser.""" + + task_id: str + value: ResultT + + +@dataclass(frozen=True) +class ParsedPromptBatch(Generic[ResultT]): + """Parsed task results plus optional syntax-repair provenance.""" + + items: tuple[PromptBatchItem[ResultT], ...] + repair_method: str | None = None + + +@dataclass(frozen=True) +class PromptBatchRunResult(Generic[ResultT]): + """Task outputs and physical prompt execution accounting.""" + + outputs: tuple[ResultT, ...] + raw_output_attempts: tuple[tuple[str, ...], ...] + repair_methods: tuple[str | None, ...] + prompt_count: int + retry_count: int + chunk_sizes: tuple[int, ...] + output_token_limit: int | None = None + repairs: tuple[dict[str, Any], ...] = () + + +def run_prompt_batches( + tasks: Sequence[TaskT], + *, + task_id: Callable[[TaskT], str], + build_request: Callable[[tuple[TaskT, ...]], PromptBatchRequest], + parse_results: Callable[[str], ParsedPromptBatch[ResultT]], + parse_single_result: Callable[[str], ResultT] | None = None, + parse_repaired_results: Callable[[str], ParsedPromptBatch[ResultT]] | None = None, + output_schema: Mapping[str, Any] | None = None, + structured_output_transport: str = "chat-json-object", + model: Any, + config: PromptBatching, + max_retries: int, + progress_bar_desc: str, + operator: str, + trace_dir: Any = None, + model_kwargs: Mapping[str, Any] | None = None, +) -> PromptBatchRunResult[ResultT]: + """Run ready semantic tasks in deterministic, context-bounded prompts.""" + + started = perf_counter() + validate_structured_output_transport(structured_output_transport) + if structured_output_transport == "responses-json-schema" and output_schema is None: + raise ValueError("responses-json-schema requires an operator output schema") + response_format: dict[str, Any] = {"type": "json_object"} + if structured_output_transport == "responses-json-schema": + response_format = { + "type": "json_schema", + "json_schema": { + "name": "semantic_batch", + "schema": dict(output_schema or {}), + }, + } + if max_retries < 0: + raise ValueError("prompt batch max_retries cannot be negative") + + task_sequence = tuple(tasks) + task_ids = tuple(task_id(task) for task in task_sequence) + _validate_unique_task_ids(task_ids) + if not task_sequence: + result = PromptBatchRunResult((), (), (), 0, 0, ()) + _write_prompt_batch_trace( + trace_dir, + operator=operator, + config=config, + task_count=0, + result=result, + latency_ms=(perf_counter() - started) * 1000, + ) + return result + + try: + requests = _build_requests( + task_sequence, + task_ids=task_ids, + build_request=build_request, + model=model, + config=config, + output_schema=( + output_schema + if structured_output_transport == "responses-json-schema" + else None + ), + ) + except ValueError as error: + _write_batch_failure( + trace_dir, + operator=operator, + attempt=0, + task_count=len(task_sequence), + error=error, + phase="preflight", + ) + raise + request_max_tokens = _fixed_output_token_limit(requests) + pending = list(requests) + outputs: dict[str, ResultT] = {} + raw_output_attempts = {identifier: [] for identifier in task_ids} + repair_methods: dict[str, str | None] = {} + prompt_count = 0 + retry_count = 0 + repairs: list[dict[str, Any]] = [] + last_error: ValueError | None = None + kwargs = dict(model_kwargs or {}) + conflicts = sorted({"max_tokens", "progress_bar_desc", "response_format"} & kwargs.keys()) + if conflicts: + raise ValueError( + "prompt batching model_kwargs cannot override runner options: " + f"{conflicts}" + ) + + for _attempt in range(max_retries + 1): + if not pending: + break + prompt_count += len(pending) + try: + model_output = model( + [request.prompt for request in pending], + progress_bar_desc=progress_bar_desc, + max_tokens=request_max_tokens, + **kwargs, + response_format=response_format, + ) + except Exception as error: + _write_batch_failure( + trace_dir, + operator=operator, + attempt=_attempt + 1, + task_count=sum(len(request.task_ids) for request in pending), + error=error, + phase="provider", + ) + raise + raw_outputs = [str(value) for value in getattr(model_output, "outputs", ())] + if len(raw_outputs) != len(pending): + raise ValueError( + "prompt batching returned an unexpected number of outputs: " + f"expected {len(pending)}, got {len(raw_outputs)}" + ) + + retry: list[PromptBatchRequest] = [] + metadata = getattr(model_output, "response_metadata", None) + for index, (request, raw_output) in enumerate( + zip(pending, raw_outputs, strict=True) + ): + for identifier in request.task_ids: + raw_output_attempts[identifier].append(raw_output) + completion = metadata[index] if metadata is not None else {} + try: + if _incomplete(completion): + raise ValueError(f"incomplete structured response: {completion}") + parsed = parse_results(raw_output) + request_outputs = _validate_parsed_items( + parsed.items, + expected_ids=request.task_ids, + ) + except ValueError as error: + try: + if _incomplete(completion): + raise error + parsed, edits = _repair_batch( + raw_output, + expected_ids=request.task_ids, + parser=parse_repaired_results or parse_results, + singleton_parser=parse_single_result, + ) + request_outputs = _validate_parsed_items( + parsed.items, expected_ids=request.task_ids + ) + except ValueError as repair_error: + last_error = error + retry.append(request) + _write_batch_failure( + trace_dir, + operator=operator, + attempt=_attempt + 1, + task_count=len(request.task_ids), + error=error, + phase="parse", + repair_error=repair_error, + ) + continue + repairs.append( + { + "method": parsed.repair_method, + "version": JSON_REPAIR_VERSION, + "attempt": _attempt + 1, + "affected_task_count": len(request.task_ids), + "edits": edits, + } + ) + else: + if parsed.repair_method is not None: + repairs.append( + { + "method": parsed.repair_method, + "version": JSON_REPAIR_VERSION, + "attempt": _attempt + 1, + "affected_task_count": len(request.task_ids), + "edits": (), + } + ) + outputs.update(request_outputs) + repair_methods.update( + dict.fromkeys(request.task_ids, parsed.repair_method) + ) + retry_count += len(retry) + pending = retry + + if pending: + assert last_error is not None + raise ValueError( + "prompt batching returned invalid structured output after " + f"{max_retries + 1} attempt(s): {last_error}" + ) from last_error + + result = PromptBatchRunResult( + outputs=tuple(outputs[identifier] for identifier in task_ids), + raw_output_attempts=tuple( + tuple(raw_output_attempts[identifier]) for identifier in task_ids + ), + repair_methods=tuple(repair_methods[identifier] for identifier in task_ids), + prompt_count=prompt_count, + retry_count=retry_count, + chunk_sizes=tuple(len(request.task_ids) for request in requests), + output_token_limit=request_max_tokens, + repairs=tuple(repairs), + ) + _write_prompt_batch_trace( + trace_dir, + operator=operator, + config=config, + task_count=len(task_sequence), + result=result, + latency_ms=(perf_counter() - started) * 1000, + ) + return result + + +def _repair_batch( + raw: str, + *, + expected_ids: tuple[str, ...], + parser: Callable[[str], ParsedPromptBatch[ResultT]], + singleton_parser: Callable[[str], ResultT] | None, +) -> tuple[ParsedPromptBatch[ResultT], tuple[tuple[str, int, str], ...]]: + if len(expected_ids) == 1 and singleton_parser is not None: + try: + value = singleton_parser(raw) + except ValueError: + pass + else: + return ParsedPromptBatch( + (PromptBatchItem(expected_ids[0], value),), "singleton-envelope" + ), () + + def validate(candidate: str) -> ParsedPromptBatch[ResultT]: + parsed = parser(candidate) + _validate_parsed_items(parsed.items, expected_ids=expected_ids) + return parsed + + repaired = repair_json_structure(raw, validator=validate) + return replace(repaired.value, repair_method="bounded-json-syntax"), repaired.edits + + +def _write_batch_failure( + trace_dir: Any, + *, + operator: str, + attempt: int, + task_count: int, + error: Exception, + phase: str, + repair_error: ValueError | None = None, +) -> None: + try: + write_trace_event( + trace_dir, + operator=operator, + event_type="prompt_batching_failure", + payload={ + "status": "error", + "phase": phase, + "attempt": attempt, + "task_count": task_count, + "error_type": type(error).__name__, + "error": str(error)[:1000], + "repair_result": "not-accepted", + "repair_error": str(repair_error)[:1000] if repair_error else None, + "repair_version": JSON_REPAIR_VERSION, + }, + ) + except Exception as trace_error: + error.add_note(f"Prompt batching failure trace also failed: {trace_error}") + + +def _incomplete(metadata: Mapping[str, Any]) -> bool: + return metadata.get("finish_reason") in { + "length", + "content_filter", + } or metadata.get("status") in {"incomplete", "failed", "cancelled"} + + +def _write_prompt_batch_trace( + trace_dir: Any, + *, + operator: str, + config: PromptBatching, + task_count: int, + result: PromptBatchRunResult[Any], + latency_ms: float, +) -> None: + """Record compact mechanics; provider usage remains the token source of truth.""" + + repair_methods = sorted( + {method for method in result.repair_methods if method is not None} + ) + syntax_repair_methods = [ + method for method in repair_methods if method != "singleton-envelope" + ] + write_trace_event( + trace_dir, + operator=operator, + event_type="prompt_batching", + payload={ + "prompt_batching": config.to_dict(), + "task_count": task_count, + "prompt_count": result.prompt_count, + "chunk_sizes": list(result.chunk_sizes), + "structured_output_token_limit": result.output_token_limit, + "retry_count": result.retry_count, + "structured_output_repair_count": len(result.repairs), + "structured_output_repaired_task_count": sum( + method is not None for method in result.repair_methods + ), + "structured_output_repairs": list(result.repairs), + "structured_output_repair_version": JSON_REPAIR_VERSION, + "structured_output_repair_methods": repair_methods, + "syntax_repair_count": sum( + repair["method"] != "singleton-envelope" for repair in result.repairs + ), + "syntax_repair_methods": syntax_repair_methods, + "prompt_batching_latency_ms": latency_ms, + }, + ) + + +def _build_requests( + tasks: tuple[TaskT, ...], + *, + task_ids: tuple[str, ...], + build_request: Callable[[tuple[TaskT, ...]], PromptBatchRequest], + model: Any, + config: PromptBatching, + output_schema: Mapping[str, Any] | None = None, +) -> tuple[PromptBatchRequest, ...]: + requests: list[PromptBatchRequest] = [] + limit = config.max_tasks or len(tasks) + for start in range(0, len(tasks), limit): + stop = min(start + limit, len(tasks)) + request = build_request(tasks[start:stop]) + expected_ids = task_ids[start:stop] + if request.task_ids != expected_ids: + raise ValueError("prompt batch request task IDs must preserve task order") + if request.max_tokens < 1: + raise ValueError("prompt batch request max_tokens must be positive") + max_ctx_len = getattr(model, "max_ctx_len", None) + if max_ctx_len is None: + requests.append(request) + continue + prompt_tokens = int(model.count_tokens(request.prompt)) + schema_tokens = ( + int( + model.count_tokens( + [{"role": "user", "content": json.dumps(output_schema)}] + ) + ) + if output_schema is not None + else 0 + ) + if prompt_tokens + schema_tokens + request.max_tokens > int(max_ctx_len): + raise ValueError( + "configured prompt batch does not fit the model context; " + f"batch_size={len(expected_ids)}, " + f"estimated_input_tokens={prompt_tokens}, estimated_schema_tokens={schema_tokens}, " + f"reserved_output_tokens={request.max_tokens}, context_limit={max_ctx_len}, " + f"first_task_id={expected_ids[0]!r}, " + f"last_task_id={expected_ids[-1]!r}" + ) + requests.append(request) + return tuple(requests) + + +def _fixed_output_token_limit(requests: Sequence[PromptBatchRequest]) -> int: + """Return the fixed per-request output ceiling for one batch run.""" + + limits = {request.max_tokens for request in requests} + if len(limits) != 1: + raise ValueError( + "all prompt batch requests must use the same output token limit" + ) + return next(iter(limits)) + + +def _validate_unique_task_ids(task_ids: Sequence[str]) -> None: + seen: set[str] = set() + for identifier in task_ids: + if not isinstance(identifier, str) or not identifier: + raise ValueError("prompt batching task IDs must be non-empty strings") + if identifier in seen: + raise ValueError(f"prompt batching input contains duplicate task ID {identifier!r}") + seen.add(identifier) + + +def _validate_parsed_items( + items: Sequence[PromptBatchItem[ResultT]], + *, + expected_ids: Sequence[str], +) -> dict[str, ResultT]: + parsed: dict[str, ResultT] = {} + for item in items: + if item.task_id in parsed: + raise ValueError( + f"prompt batching returned duplicate task ID {item.task_id!r}" + ) + parsed[item.task_id] = item.value + + expected = set(expected_ids) + unknown = sorted(set(parsed) - expected) + missing = sorted(expected - set(parsed)) + if unknown: + raise ValueError(f"prompt batching returned unknown task IDs: {unknown}") + if missing: + raise ValueError( + f"prompt batching omitted or returned missing task IDs: {missing}" + ) + return parsed + + +__all__ = [ + "ParsedPromptBatch", + "PromptBatchItem", + "PromptBatchRequest", + "PromptBatchRunResult", + "PromptBatching", + "parse_prompt_batch_size", + "run_prompt_batches", +] diff --git a/src/agent_memory/adapters/lotus/provider_usage_lm.py b/src/agent_memory/adapters/lotus/provider_usage_lm.py index 256151e..c86fead 100644 --- a/src/agent_memory/adapters/lotus/provider_usage_lm.py +++ b/src/agent_memory/adapters/lotus/provider_usage_lm.py @@ -31,8 +31,37 @@ def __init__( **kwargs: Any, ) -> None: self._provider_usage_trace_dir = trace_dir + self._completion_metadata: list[dict[str, Any]] = [] super().__init__(*args, **kwargs) + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """Keep completion metadata aligned with outputs, including cache hits.""" + + self._completion_metadata = [] + parent: Any = super() + output = parent.__call__(*args, **kwargs) + output.response_metadata = tuple(self._completion_metadata) + return output + + def _get_top_choice(self, response: Any) -> str: + """Preserve the finish state discarded by LOTUS's text-only LMOutput.""" + + parent: Any = super() + text = parent._get_top_choice(response) + metadata = getattr(response, "provider_response_metadata", {}) + choices = getattr(response, "choices", ()) + self._completion_metadata.append( + { + "finish_reason": metadata.get( + "finish_reason", + getattr(choices[0], "finish_reason", None) if choices else None, + ), + "status": metadata.get("status"), + "incomplete_reason": metadata.get("incomplete_reason"), + } + ) + return text + def _process_uncached_messages( self, uncached_data: list[tuple[list[dict[str, str]], str]], diff --git a/src/agent_memory/adapters/lotus/relational.py b/src/agent_memory/adapters/lotus/relational.py index ce03ab0..56db059 100644 --- a/src/agent_memory/adapters/lotus/relational.py +++ b/src/agent_memory/adapters/lotus/relational.py @@ -19,16 +19,21 @@ aggregate_groups_with_keys, aggregate_input_columns, execute_native_sem_agg_group, + execute_native_sem_agg_groups, execute_structured_sem_agg_group, + execute_structured_sem_agg_groups, ) from agent_memory.policy.expressions import ( + ArithmeticExpr, ArrayCatExpr, BooleanExpr, + CaseWhenExpr, ColumnExpr, ComparisonExpr, Expr, LeastExpr, LiteralExpr, + TryCastExpr, expr_from_param, is_scalar, ) @@ -48,6 +53,21 @@ def execute_select( return source.loc[:, list(query.params["columns"])].copy() +def execute_limit( + query: QueryExpr, + inputs: Mapping[str, Any], + execute: Callable[[QueryExpr, Mapping[str, Any]], Any], +) -> Any: + """Select a legal bounded bag; physical position is not a query guarantee.""" + + n = query.params["n"] + if isinstance(n, bool) or not isinstance(n, int): + raise TypeError("limit requires a non-negative integer") + if n < 0: + raise ValueError("limit requires a non-negative integer") + return execute(query.inputs[0], inputs).iloc[:n].copy() + + def execute_alias( query: QueryExpr, inputs: Mapping[str, Any], @@ -129,8 +149,12 @@ def execute_union_by_name( if not isinstance(allow_missing, bool): raise TypeError("union_by_name allow_missing_columns must be a bool") columns = _union_by_name_columns(left, right, allow_missing_columns=allow_missing) - left_aligned = _align_columns_by_name(left, columns, allow_missing_columns=allow_missing) - right_aligned = _align_columns_by_name(right, columns, allow_missing_columns=allow_missing) + left_aligned = _align_columns_by_name( + left, columns, allow_missing_columns=allow_missing + ) + right_aligned = _align_columns_by_name( + right, columns, allow_missing_columns=allow_missing + ) concatenated = _concat_rows(left_aligned, right_aligned) return concatenated.drop_duplicates(ignore_index=True) @@ -152,10 +176,20 @@ def execute_assign( source[column_name] = value continue expr = expr_from_param(value) - if not isinstance(expr, (LiteralExpr, ColumnExpr, ArrayCatExpr, LeastExpr)): + if not isinstance( + expr, + ( + LiteralExpr, + ColumnExpr, + ArithmeticExpr, + ArrayCatExpr, + LeastExpr, + TryCastExpr, + CaseWhenExpr, + ), + ): raise TypeError( - "assign values must be scalar literals, column expressions, " - "array_cat expressions, or least expressions" + "assign values must be scalar literals or supported row expressions" ) source[column_name] = evaluate_expr(expr, source) return source @@ -242,6 +276,56 @@ def execute_join( ).reset_index(drop=True) +def predicate_join_row_ids( + query: QueryExpr, + left: pd.DataFrame, + right: pd.DataFrame, + *, + on: Sequence[Mapping[str, Any]], +) -> tuple[tuple[Any, Any], ...]: + """Return input row-ID pairs that satisfy a deterministic join predicate.""" + + predicates = tuple(on) + if not _is_predicate_join(predicates): + raise ValueError("predicate join row IDs require normalized predicate expressions") + if len(query.inputs) != 2: + raise ValueError("predicate join row IDs require exactly two query inputs") + + left_id_column = _unused_column_name( + left, + right, + base="_agent_memory_join_left_row_id", + ) + right_id_column = _unused_column_name( + left, + right, + base="_agent_memory_join_right_row_id", + ) + left_with_ids = left.copy() + right_with_ids = right.copy() + left_with_ids[left_id_column] = list(left.index) + right_with_ids[right_id_column] = list(right.index) + joined = _execute_predicate_join( + query, + left_with_ids, + right_with_ids, + predicates=tuple(expr_from_param(predicate) for predicate in predicates), + how="inner", + ) + + left_alias = _relation_alias(query.inputs[0]) or "left" + right_alias = _relation_alias(query.inputs[1]) or "right" + left_output = f"{left_id_column}:{left_alias}" + right_output = f"{right_id_column}:{right_alias}" + return tuple( + zip( + joined[left_output], + joined[right_output], + strict=True, + ) + ) + + def execute_drop_duplicates( query: QueryExpr, inputs: Mapping[str, Any], @@ -278,7 +362,9 @@ def execute_array_agg( raise ValueError(f"array_agg input columns not found in DataFrame: {missing}") records = _strict_json_records(source, columns) - value = json.dumps(records, ensure_ascii=False, default=_json_default, allow_nan=False) + value = json.dumps( + records, ensure_ascii=False, default=_json_default, allow_nan=False + ) return pd.DataFrame([{output_col: value}], columns=[output_col]) @@ -295,13 +381,17 @@ def execute_grouped_array_agg( output_col = str(query.params["output_col"]) missing = [column for column in columns if column not in source.columns] if missing: - raise ValueError(f"grouped array_agg input columns not found in DataFrame: {missing}") + raise ValueError( + f"grouped array_agg input columns not found in DataFrame: {missing}" + ) if group_query.op == "group_by": group_keys = tuple(str(key) for key in group_query.params["keys"]) missing_keys = [key for key in group_keys if key not in source.columns] if missing_keys: - raise ValueError(f"group_by key columns not found in DataFrame: {missing_keys}") + raise ValueError( + f"group_by key columns not found in DataFrame: {missing_keys}" + ) return _array_agg_by_keys( source, group_keys=group_keys, @@ -331,9 +421,7 @@ def execute_min( missing = [column for column in columns if column not in source.columns] if missing: if len(missing) == 1: - raise ValueError( - f"min input column not found in DataFrame: {missing[0]!r}" - ) + raise ValueError(f"min input column not found in DataFrame: {missing[0]!r}") raise ValueError(f"min input columns not found in DataFrame: {missing}") if source_query.op == "group_by": @@ -374,15 +462,29 @@ def execute_agg( raise ValueError("agg expects group_by or sem_groupby input") if source.empty: return pd.DataFrame(columns=list(output_columns(query))) + if ( + context.config.sem_agg_dispatch != "sequential" + or context.config.prompt_batching is not None + ): + return _execute_agg_with_group_batching( + query, + source, + aggregates, + context=context, + ) rows: list[dict[str, Any]] = [] - for group_index, (key_values, group) in enumerate(aggregate_groups_with_keys(source)): + for group_index, (key_values, group) in enumerate( + aggregate_groups_with_keys(source) + ): row: dict[str, Any] = dict(key_values) for aggregate in aggregates: if isinstance(aggregate, ArrayAggregateSpec): if aggregate.output_col in row: continue - row[aggregate.output_col] = _array_records_json(group, aggregate.columns) + row[aggregate.output_col] = _array_records_json( + group, aggregate.columns + ) continue if isinstance(aggregate, CollectListAggregateSpec): if aggregate.output_col in row: @@ -399,8 +501,7 @@ def execute_agg( ] if missing: raise ValueError( - "min input columns not found in aggregate group: " - f"{missing}" + f"min input columns not found in aggregate group: {missing}" ) row[aggregate.output_col] = _minimum_value(group, aggregate.columns) continue @@ -427,6 +528,111 @@ def execute_agg( return result +def _execute_agg_with_group_batching( + query: QueryExpr, + source: pd.DataFrame, + aggregates: Sequence[object], + *, + context: Any, +) -> pd.DataFrame: + """Execute semantic aggregate specs across independent groups in batches.""" + + grouped = aggregate_groups_with_keys(source) + groups = [group for _key_values, group in grouped] + rows = [dict(key_values) for key_values, _group in grouped] + for aggregate in aggregates: + if isinstance(aggregate, SemanticAggregateSpec): + semantic_values = _execute_grouped_semantic_aggregate_spec_many( + aggregate, + groups, + context=context, + ) + for row, values in zip(rows, semantic_values, strict=True): + for column, value in values.items(): + if column not in row: + row[column] = value + continue + for row, group in zip(rows, groups, strict=True): + _apply_deterministic_aggregate(aggregate, group, row) + + result = pd.DataFrame(rows, columns=list(output_columns(query))) + for aggregate in aggregates: + if isinstance(aggregate, MinAggregateSpec): + result[aggregate.output_col] = pd.Series( + [row.get(aggregate.output_col) for row in rows], + dtype=object, + ) + return result + + +def _apply_deterministic_aggregate( + aggregate: object, + group: pd.DataFrame, + row: dict[str, Any], +) -> None: + """Apply one non-semantic aggregate spec to one grouped frame.""" + + if isinstance(aggregate, ArrayAggregateSpec): + if aggregate.output_col not in row: + row[aggregate.output_col] = _array_records_json(group, aggregate.columns) + return + if isinstance(aggregate, CollectListAggregateSpec): + if aggregate.output_col not in row: + row[aggregate.output_col] = _collect_list_json(group, aggregate.column) + return + if isinstance(aggregate, MinAggregateSpec): + if aggregate.output_col in row: + return + missing = [ + column for column in aggregate.columns if column not in group.columns + ] + if missing: + raise ValueError( + f"min input columns not found in aggregate group: {missing}" + ) + row[aggregate.output_col] = _minimum_value(group, aggregate.columns) + return + raise TypeError(f"Unsupported aggregate spec: {type(aggregate).__name__}") + + +def _execute_grouped_semantic_aggregate_spec_many( + aggregate: SemanticAggregateSpec, + groups: Sequence[pd.DataFrame], + *, + context: Any, +) -> list[dict[str, Any]]: + """Execute one semantic aggregate spec over independent grouped frames.""" + + if not groups: + return [] + query = QueryExpr( + op="sem_agg", + params={ + "input_cols": aggregate.input_cols, + "output_cols": aggregate.output_cols, + "instruction": aggregate.instruction, + }, + ) + input_cols = aggregate_input_columns(groups[0], query.params.get("input_cols")) + if len(aggregate.output_cols) == 1: + output_col = aggregate.output_cols[0] + values = execute_native_sem_agg_groups( + query, + groups, + input_cols, + context.config, + ) + return [{output_col.name: value} for value in values] + parsed = execute_structured_sem_agg_groups( + query, + groups, + input_cols, + aggregate.output_cols, + context.config, + ) + return [dict(values) for values in parsed] + + def _execute_grouped_semantic_aggregate_spec( aggregate: SemanticAggregateSpec, group: pd.DataFrame, @@ -477,13 +683,21 @@ def execute_over_array_agg( ) columns = tuple(str(column) for column in query.params["columns"]) output_col = str(query.params["output_col"]) - emit_columns = tuple(output_columns(over_query.inputs[0])) or tuple(emit_source.columns) + emit_columns = tuple(output_columns(over_query.inputs[0])) or tuple( + emit_source.columns + ) missing = [column for column in columns if column not in frame_source.columns] if missing: - raise ValueError(f"over array_agg input columns not found in DataFrame: {missing}") - missing_emit = [column for column in emit_columns if column not in emit_source.columns] + raise ValueError( + f"over array_agg input columns not found in DataFrame: {missing}" + ) + missing_emit = [ + column for column in emit_columns if column not in emit_source.columns + ] if missing_emit: - raise ValueError(f"over array_agg emit columns not found in DataFrame: {missing_emit}") + raise ValueError( + f"over array_agg emit columns not found in DataFrame: {missing_emit}" + ) rows: list[dict[str, Any]] = [] for frame in over_frames(emit_source, frame_source, over_query.params): @@ -590,7 +804,9 @@ def execute_unnest( expected_columns = list(output_columns(query)) source = execute(query.inputs[0], inputs).copy() column = str(query.params["column"]) - fields = tuple((str(field), str(output)) for field, output in query.params["fields"]) + fields = tuple( + (str(field), str(output)) for field, output in query.params["fields"] + ) if column not in source.columns: raise ValueError(f"unnest input column not found in DataFrame: {column!r}") @@ -598,7 +814,9 @@ def execute_unnest( output_names = [output for _, output in fields] conflicts = sorted(set(parent_columns).intersection(output_names)) if conflicts: - raise ValueError(f"unnest output columns conflict with existing columns: {conflicts}") + raise ValueError( + f"unnest output columns conflict with existing columns: {conflicts}" + ) output_rows: list[dict[str, Any]] = [] for _, row in source.iterrows(): @@ -700,7 +918,9 @@ def _require_supported_join_how(how: str) -> None: """Require a pandas relational join mode supported by the public API.""" if how not in {"inner", "left", "right", "outer", "left_anti"}: - raise ValueError("join how must be one of: inner, left, right, outer, left_anti") + raise ValueError( + "join how must be one of: inner, left, right, outer, left_anti" + ) def _is_predicate_join(on: tuple[Any, ...]) -> bool: @@ -717,6 +937,20 @@ def _join_suffixes(query: QueryExpr) -> tuple[str, str]: return f":{left_alias}", f":{right_alias}" +def _unused_column_name( + left: pd.DataFrame, + right: pd.DataFrame, + *, + base: str, +) -> str: + """Return a temporary column name absent from both join inputs.""" + + candidate = base + while candidate in left.columns or candidate in right.columns: + candidate = f"_{candidate}" + return candidate + + def _relation_alias(query: QueryExpr) -> str | None: """Return relation alias name if present.""" @@ -862,6 +1096,12 @@ def evaluate_expr(expr: Expr, frame: Any) -> Any: return expr.value if isinstance(expr, ComparisonExpr): return _evaluate_comparison(expr, frame) + if isinstance(expr, ArithmeticExpr): + return _evaluate_arithmetic(expr, frame) + if isinstance(expr, TryCastExpr): + return _evaluate_try_cast(expr, frame) + if isinstance(expr, CaseWhenExpr): + return _evaluate_case_when(expr, frame) if isinstance(expr, BooleanExpr): return _evaluate_boolean(expr, frame) if isinstance(expr, ArrayCatExpr): @@ -871,6 +1111,123 @@ def evaluate_expr(expr: Expr, frame: Any) -> Any: raise TypeError(f"Unsupported expression type: {type(expr).__name__}") +def _evaluate_arithmetic(expr: ArithmeticExpr, frame: Any) -> Any: + """Evaluate numeric arithmetic with null propagation and explicit zero checks.""" + + left = evaluate_expr(expr.left, frame) + right = evaluate_expr(expr.right, frame) + if not isinstance(left, pd.Series) and not isinstance(right, pd.Series): + return _evaluate_arithmetic_value(expr.op, left, right) + left_series = ( + left + if isinstance(left, pd.Series) + else pd.Series([left] * len(frame), index=frame.index, dtype=object) + ) + right_series = ( + right + if isinstance(right, pd.Series) + else pd.Series([right] * len(frame), index=frame.index, dtype=object) + ) + return pd.Series( + [ + _evaluate_arithmetic_value(expr.op, left_value, right_value) + for left_value, right_value in zip(left_series, right_series, strict=True) + ], + index=frame.index, + dtype=object, + ) + + +def _evaluate_try_cast(expr: TryCastExpr, frame: Any) -> Any: + """Convert scalars to floats, returning null for failed conversions.""" + + if expr.target != "float": + raise ValueError("try_cast currently only supports the 'float' target") + value = evaluate_expr(expr.value, frame) + if isinstance(value, pd.Series): + return value.map(_try_cast_float) + return _try_cast_float(value) + + +def _try_cast_float(value: Any) -> float | None: + """Convert one nullable scalar using Python float semantics.""" + + if _is_arithmetic_null(value): + return None + try: + return float(value) + except (TypeError, ValueError, OverflowError): + return None + + +def _evaluate_case_when(expr: CaseWhenExpr, frame: Any) -> pd.Series: + """Evaluate CASE WHEN, treating a null condition as false.""" + + condition = evaluate_predicate(expr.condition, frame) + selected = condition.to_numpy(dtype=bool) + result = np.empty(len(frame), dtype=object) + for branch, mask in ( + (expr.then_value, selected), + (expr.else_value, ~selected), + ): + positions = np.flatnonzero(mask) + if positions.size == 0: + continue + branch_frame = frame.iloc[positions] + branch_values = _as_expression_series( + evaluate_expr(branch, branch_frame), + branch_frame, + ) + result[positions] = branch_values.to_numpy(dtype=object) + return pd.Series(result, index=frame.index, dtype=object) + + +def _as_expression_series(value: Any, frame: pd.DataFrame) -> pd.Series: + """Broadcast a scalar expression result to the current frame.""" + + if isinstance(value, pd.Series): + return value + return pd.Series([value] * len(frame), index=frame.index, dtype=object) + + +def _evaluate_arithmetic_value(op: str, left: Any, right: Any) -> Any: + """Evaluate one pair of nullable numeric scalar operands.""" + + if _is_arithmetic_null(left) or _is_arithmetic_null(right): + return None + _require_arithmetic_numeric(left) + _require_arithmetic_numeric(right) + if op == "add": + return left + right + if op == "subtract": + return left - right + if op == "multiply": + return left * right + if op == "divide": + if right == 0: + raise ZeroDivisionError("arithmetic division by zero") + return left / right + raise ValueError(f"Unsupported arithmetic expression op: {op!r}") + + +def _require_arithmetic_numeric(value: Any) -> None: + """Require one non-null arithmetic operand to be numeric but not boolean.""" + + from numbers import Number + + if isinstance(value, (bool, np.bool_)) or not isinstance(value, Number): + raise TypeError("arithmetic expressions require numeric non-null operands") + + +def _is_arithmetic_null(value: Any) -> bool: + """Return whether one scalar arithmetic operand is null.""" + + if value is None or value is pd.NA: + return True + result = pd.isna(value) + return isinstance(result, (bool, np.bool_)) and bool(result) + + def _column_values(frame: Any, expr: ColumnExpr) -> Any: """Return the DataFrame column referenced by an expression.""" @@ -991,8 +1348,7 @@ def _evaluate_least(expr: LeastExpr, frame: Any) -> pd.Series: values.append(min(candidates)) except TypeError as exc: raise TypeError( - "least operands are not mutually comparable at row position " - f"{position}" + f"least operands are not mutually comparable at row position {position}" ) from exc return pd.Series(values, index=frame.index) @@ -1043,7 +1399,9 @@ def _array_agg_by_keys( ) output_rows: list[dict[str, Any]] = [] if source.empty: - return pd.DataFrame(columns=[*group_keys, output_col] if include_keys else [output_col]) + return pd.DataFrame( + columns=[*group_keys, output_col] if include_keys else [output_col] + ) grouped = source.groupby(list(group_keys), sort=False, dropna=False) for key, group in grouped: @@ -1129,7 +1487,9 @@ def _load_json_array(value: Any, *, op: str, column: str) -> list[Any]: try: parsed = json.loads(value) except (TypeError, json.JSONDecodeError) as error: - raise ValueError(f"{op} value in column {column!r} must be a JSON array") from error + raise ValueError( + f"{op} value in column {column!r} must be a JSON array" + ) from error if not isinstance(parsed, list): raise ValueError(f"{op} value in column {column!r} must be a JSON array") return parsed @@ -1151,7 +1511,9 @@ def _load_json_object(value: Any, *, op: str, column: str) -> dict[str, Any]: try: parsed = json.loads(value) except (TypeError, json.JSONDecodeError) as error: - raise ValueError(f"{op} value in column {column!r} must be a JSON object") from error + raise ValueError( + f"{op} value in column {column!r} must be a JSON object" + ) from error if not isinstance(parsed, dict): raise ValueError(f"{op} value in column {column!r} must be a JSON object") return parsed @@ -1168,10 +1530,14 @@ def _require_join_keys(left: Any, right: Any, keys: tuple[str, ...]) -> None: details.append(f"left missing {missing_left}") if missing_right: details.append(f"right missing {missing_right}") - raise ValueError(f"join key columns must exist on both sides: {', '.join(details)}") + raise ValueError( + f"join key columns must exist on both sides: {', '.join(details)}" + ) -def _require_non_null_join_keys(frame: Any, keys: tuple[str, ...], *, side: str) -> None: +def _require_non_null_join_keys( + frame: Any, keys: tuple[str, ...], *, side: str +) -> None: """Reject null join keys to avoid pandas null-null matching surprises.""" if frame.loc[:, list(keys)].isna().any().any(): diff --git a/src/agent_memory/adapters/lotus/request_hook.py b/src/agent_memory/adapters/lotus/request_hook.py new file mode 100644 index 0000000..b559e2a --- /dev/null +++ b/src/agent_memory/adapters/lotus/request_hook.py @@ -0,0 +1,54 @@ +"""Optional per-request persistence; no benchmark or storage dependency.""" + +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Any, Protocol + +from agent_memory.tracing.semantic import trace_scope_value + + +class RequestHook(Protocol): + """Execute or recover one deterministic request occurrence.""" + + def execute( + self, + identity: Mapping[str, Any], + payload: Mapping[str, Any], + send: Callable[[], dict[str, Any]], + ) -> dict[str, Any]: ... + + +@dataclass +class RequestScope: + hook: RequestHook + identity: Mapping[str, Any] + batch: int = 0 + + +_SCOPE: ContextVar[RequestScope | None] = ContextVar("request_hook", default=None) + + +@contextmanager +def request_scope(hook: RequestHook | None, **identity: Any) -> Iterator[None]: + """Reset deterministic call positions for one replayable execution unit.""" + token = _SCOPE.set(None if hook is None else RequestScope(hook, identity)) + try: + yield + finally: + _SCOPE.reset(token) + + +def next_request_batch() -> tuple[RequestHook, dict[str, Any]] | None: + """Capture scope before dispatching worker threads.""" + scope = _SCOPE.get() + if scope is None: + return None + identity = {**scope.identity, "batch": scope.batch} + for key in ("semantic_operator", "query_digest"): + value = trace_scope_value(key, None) + if value is not None: + identity[key] = value + scope.batch += 1 + return scope.hook, identity diff --git a/src/agent_memory/adapters/lotus/sem_agg.py b/src/agent_memory/adapters/lotus/sem_agg.py index 1d8c192..e2af101 100644 --- a/src/agent_memory/adapters/lotus/sem_agg.py +++ b/src/agent_memory/adapters/lotus/sem_agg.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass import json from pathlib import Path from typing import Any @@ -15,6 +16,10 @@ ) from agent_memory.tracing.semantic import write_structured_generation_trace from agent_memory.adapters.lotus.sem_groupby import GROUP_ID_COLUMN +from agent_memory.adapters.lotus.sem_agg_batch_prompting import ( + execute_batch_prompted_sem_agg, +) +from agent_memory.adapters.lotus.prompt_batching import PromptBatching from agent_memory.adapters.lotus.structured import ( StructuredLMRetryResult, escape_structured_formatter_placeholders, @@ -31,6 +36,24 @@ JSON_OBJECT_RESPONSE_FORMAT = {"type": "json_object"} +@dataclass(frozen=True) +class _SemAggExecution: + """Final output and structured retry metadata for one aggregate group.""" + + raw_output: str + retry_result: StructuredLMRetryResult | None = None + + +@dataclass +class _SemAggGroupState: + """Mutable hierarchical state for one independent aggregate group.""" + + group_index: int + documents: list[str] + tree_level: int = 0 + execution: _SemAggExecution | None = None + + def execute_sem_agg( query: QueryExpr, inputs: Mapping[str, Any], @@ -77,10 +100,16 @@ def execute_over_sem_agg( raise ValueError("over sem_agg requires explicit output_cols") output_cols = tuple(declared_output_cols) emit_source = execute(over_query.inputs[0], inputs) - emit_columns = tuple(output_columns(over_query.inputs[0])) or tuple(emit_source.columns) - missing_emit = [column for column in emit_columns if column not in emit_source.columns] + emit_columns = tuple(output_columns(over_query.inputs[0])) or tuple( + emit_source.columns + ) + missing_emit = [ + column for column in emit_columns if column not in emit_source.columns + ] if missing_emit: - raise ValueError(f"over sem_agg emit columns not found in DataFrame: {missing_emit}") + raise ValueError( + f"over sem_agg emit columns not found in DataFrame: {missing_emit}" + ) frame_source_query = over_query.params.get("frame_source") frame_source = ( execute(frame_source_query, inputs) @@ -96,7 +125,9 @@ def execute_over_sem_agg( rows.append(row) continue - input_cols = aggregate_input_columns(frame.frame, query.params.get("input_cols")) + input_cols = aggregate_input_columns( + frame.frame, query.params.get("input_cols") + ) if len(output_cols) == 1: result = execute_native_sem_agg( query, @@ -136,24 +167,105 @@ def execute_native_sem_agg( return pd.DataFrame(columns=list(output_columns(query))) config = config or LotusExecutionConfig() + grouped = aggregate_groups_with_keys(source) + raw_outputs = execute_native_sem_agg_groups( + query, + [group for _key_values, group in grouped], + input_cols, + config, + ) rows: list[dict[str, Any]] = [] - for group_index, (key_values, group) in enumerate(aggregate_groups_with_keys(source)): - raw_output = execute_native_sem_agg_group(query, group, input_cols, config) + for group_index, ((key_values, group), raw_output) in enumerate( + zip(grouped, raw_outputs, strict=True) + ): row = dict(key_values) if output_col.name not in row: row[output_col.name] = raw_output rows.append(row) - write_sem_agg_audit( - config, - query=query, - group=group, - input_cols=input_cols, + if _prompt_batching(config) is None: + write_sem_agg_audit( + config, + query=query, + group=group, + input_cols=input_cols, + output_cols=(output_col,), + group_index=group_index, + raw_output=raw_output, + parsed_output={output_col.name: raw_output}, + ) + return pd.DataFrame(rows, columns=list(output_columns(query))) + + +def _prompt_batching(config: LotusExecutionConfig) -> PromptBatching | None: + """Return the shared prompt batching contract.""" + + return config.prompt_batching + + +def execute_native_sem_agg_groups( + query: QueryExpr, + groups: Sequence[pd.DataFrame], + input_cols: Sequence[str], + config: LotusExecutionConfig, +) -> list[str]: + """Aggregate independent groups with the configured physical dispatch.""" + + prompt_batching = _prompt_batching(config) + if prompt_batching is not None: + import lotus + + output_col = aggregate_output_columns(query, input_cols)[0] + result = execute_batch_prompted_sem_agg( + [aggregate_group_text(group, input_cols) for group in groups], + instruction=aggregate_instruction(query, input_cols), output_cols=(output_col,), - group_index=group_index, - raw_output=raw_output, - parsed_output={output_col.name: raw_output}, + model=lotus.settings.lm, + prompt_batching=prompt_batching, + structured_output_transport=config.structured_output_transport, + max_retries=config.structured_parse_retries, + model_kwargs=structured_sem_agg_model_kwargs(config), + progress_bar_desc=config.sem_agg_progress_bar_desc, + trace_dir=config.trace_dir(), ) - return pd.DataFrame(rows, columns=list(output_columns(query))) + outputs = [str(values[output_col.name]) for values in result.outputs] + for group_index, (group, output, raw_output_attempts, repair_method) in enumerate( + zip( + groups, + outputs, + result.raw_output_attempts, + result.repair_methods, + strict=True, + ) + ): + write_sem_agg_audit( + config, + query=query, + group=group, + input_cols=input_cols, + output_cols=(output_col,), + group_index=group_index, + raw_output=raw_output_attempts[-1], + raw_output_attempts=raw_output_attempts, + parsed_output={output_col.name: output}, + syntax_repair_method=repair_method, + ) + return outputs + if config.sem_agg_dispatch == "sequential" or len(groups) < 2: + return [ + execute_native_sem_agg_group(query, group, input_cols, config) + for group in groups + ] + + import lotus + + instruction = aggregate_instruction(query, input_cols) + executions = _execute_independent_sem_agg_groups( + [aggregate_group_text(group, input_cols) for group in groups], + lotus.settings.lm, + instruction, + config=config, + ) + return [execution.raw_output for execution in executions] def execute_native_sem_agg_group( @@ -199,16 +311,20 @@ def execute_structured_sem_agg( return pd.DataFrame(columns=list(output_columns(query))) config = config or LotusExecutionConfig() + grouped = aggregate_groups_with_keys(source) + parsed_outputs = execute_structured_sem_agg_groups( + query, + [group for _key_values, group in grouped], + input_cols, + output_cols, + config, + ) rows: list[dict[str, Any]] = [] - for group_index, (key_values, group) in enumerate(aggregate_groups_with_keys(source)): - parsed = execute_structured_sem_agg_group( - query, - group, - input_cols, - output_cols, - config, - group_index=group_index, - ) + for (key_values, _group), parsed in zip( + grouped, + parsed_outputs, + strict=True, + ): row = dict(key_values) for column in output_cols: if column.name not in row: @@ -217,6 +333,103 @@ def execute_structured_sem_agg( return pd.DataFrame(rows, columns=list(output_columns(query))) +def execute_structured_sem_agg_groups( + query: QueryExpr, + groups: Sequence[pd.DataFrame], + input_cols: Sequence[str], + output_cols: Sequence[ColumnSpec], + config: LotusExecutionConfig, +) -> list[Mapping[str, Any]]: + """Aggregate independent groups into declared structured fields.""" + + prompt_batching = _prompt_batching(config) + if prompt_batching is not None: + import lotus + + instruction = aggregate_instruction(query, input_cols) + result = execute_batch_prompted_sem_agg( + [aggregate_group_text(group, input_cols) for group in groups], + instruction=instruction, + output_cols=output_cols, + model=lotus.settings.lm, + prompt_batching=prompt_batching, + structured_output_transport=config.structured_output_transport, + max_retries=config.structured_parse_retries, + model_kwargs=structured_sem_agg_model_kwargs(config), + progress_bar_desc=config.sem_agg_progress_bar_desc, + trace_dir=config.trace_dir(), + ) + parsed_outputs: list[Mapping[str, Any]] = [ + dict(values) for values in result.outputs + ] + for group_index, (group, parsed, repair_method) in enumerate( + zip(groups, parsed_outputs, result.repair_methods, strict=True) + ): + raw_output_attempts = result.raw_output_attempts[group_index] + raw_output = raw_output_attempts[-1] + write_sem_agg_audit( + config, + query=query, + group=group, + input_cols=input_cols, + output_cols=output_cols, + group_index=group_index, + raw_output=raw_output, + raw_output_attempts=raw_output_attempts, + parsed_output=parsed, + syntax_repair_method=repair_method, + ) + return parsed_outputs + if config.sem_agg_dispatch == "sequential" or len(groups) < 2: + return [ + execute_structured_sem_agg_group( + query, + group, + input_cols, + output_cols, + config, + group_index=group_index, + ) + for group_index, group in enumerate(groups) + ] + + import lotus + + instruction = structured_aggregate_instruction(query, input_cols, output_cols) + executions = _execute_independent_sem_agg_groups( + [aggregate_group_text(group, input_cols) for group in groups], + lotus.settings.lm, + instruction, + config=config, + output_cols=output_cols, + failure_extra_by_group=[ + { + "group_index": group_index, + "final_instruction": instruction, + "group_row_preview": group.head(5) + .astype(str) + .to_dict(orient="records"), + } + for group_index, group in enumerate(groups) + ], + ) + return [ + _parse_structured_sem_agg_execution( + execution, + query=query, + group=group, + input_cols=input_cols, + output_cols=output_cols, + config=config, + group_index=group_index, + instruction=instruction, + ) + for group_index, (group, execution) in enumerate( + zip(groups, executions, strict=True) + ) + ] + + def execute_structured_sem_agg_group( query: QueryExpr, group: pd.DataFrame, @@ -237,8 +450,42 @@ def execute_structured_sem_agg_group( config, group_index=group_index, ) - raw_output = str(retry_result.raw_outputs[0]) - raw_output_attempts = tuple(str(value) for value in retry_result.raw_output_attempts[0]) + execution = _SemAggExecution( + raw_output=str(retry_result.raw_outputs[0]), + retry_result=retry_result, + ) + return _parse_structured_sem_agg_execution( + execution, + query=query, + group=group, + input_cols=input_cols, + output_cols=output_cols, + config=config, + group_index=group_index, + instruction=instruction, + ) + + +def _parse_structured_sem_agg_execution( + execution: _SemAggExecution, + *, + query: QueryExpr, + group: pd.DataFrame, + input_cols: Sequence[str], + output_cols: Sequence[ColumnSpec], + config: LotusExecutionConfig, + group_index: int, + instruction: str, +) -> Mapping[str, Any]: + """Parse, audit, and return one structured aggregate execution.""" + + retry_result = execution.retry_result + if retry_result is None: + raise ValueError("structured sem_agg execution is missing retry metadata") + raw_output = execution.raw_output + raw_output_attempts = tuple( + str(value) for value in retry_result.raw_output_attempts[0] + ) if retry_result.invalid_indices: parse_error = structured_parse_error( raw_output, @@ -317,7 +564,9 @@ def execute_lotus_style_structured_sem_agg_group( 0: { "group_index": group_index, "final_instruction": instruction, - "group_row_preview": group.head(5).astype(str).to_dict(orient="records"), + "group_row_preview": group.head(5) + .astype(str) + .to_dict(orient="records"), } }, ) @@ -332,8 +581,7 @@ def structured_aggregate_instruction( instruction = aggregate_instruction(query, input_cols) field_lines = "\n".join( - f"- {column.name}: {column.description or 'string'}" - for column in output_cols + f"- {column.name}: {column.description or 'string'}" for column in output_cols ) shape = json.dumps( {column.name: "string" for column in output_cols}, @@ -378,6 +626,7 @@ def write_sem_agg_audit( raw_output_attempts: Sequence[str] | None = None, parse_error: str = "", failure_artifact: Path | str | None = None, + syntax_repair_method: str | None = None, ) -> None: """Write a structured audit row for one semantic aggregate group.""" @@ -388,8 +637,8 @@ def write_sem_agg_audit( if len(output_cols) > 1 else formatted_instruction ) - input_preview = group.loc[:, list(input_cols)].head(20).astype(str).to_dict( - orient="records" + input_preview = ( + group.loc[:, list(input_cols)].head(20).astype(str).to_dict(orient="records") ) attempts = [str(value) for value in (raw_output_attempts or (raw_output,))] audit_row = { @@ -412,6 +661,22 @@ def write_sem_agg_audit( "parse_error": parse_error, "failure_artifact": "" if failure_artifact is None else str(failure_artifact), } + if ( + config.sem_agg_dispatch != "sequential" + or config.prompt_batching is not None + ): + audit_row.update( + { + "dispatch": config.sem_agg_dispatch, + "prompt_batching": ( + None + if config.prompt_batching is None + else config.prompt_batching.to_dict() + ), + "structured_output_repaired": syntax_repair_method is not None, + "structured_output_repair_method": syntax_repair_method or "", + } + ) write_structured_generation_trace( config.trace_dir(), operator="sem_agg", @@ -420,6 +685,204 @@ def write_sem_agg_audit( ) +def _execute_independent_sem_agg_groups( + documents_by_group: Sequence[Sequence[str]], + model: Any, + instruction: str, + *, + config: LotusExecutionConfig, + output_cols: Sequence[ColumnSpec] | None = None, + failure_extra_by_group: Sequence[Mapping[str, Any]] | None = None, +) -> list[_SemAggExecution]: + """Batch ready prompts while preserving independent aggregate groups.""" + + if failure_extra_by_group is not None and len(failure_extra_by_group) != len( + documents_by_group + ): + raise ValueError("sem_agg failure metadata must match aggregate groups") + states = [ + _SemAggGroupState( + group_index=group_index, + documents=[str(document) for document in documents], + ) + for group_index, documents in enumerate(documents_by_group) + ] + if any(not state.documents for state in states): + raise ValueError("sem_agg aggregate groups cannot be empty") + + while any(state.execution is None for state in states): + regular_prompts: list[list[dict[str, str]]] = [] + regular_refs: list[tuple[_SemAggGroupState, int]] = [] + regular_counts: dict[int, int] = {} + structured_prompts: list[list[dict[str, str]]] = [] + structured_states: list[_SemAggGroupState] = [] + + for state in states: + if state.execution is not None: + continue + prompts = _build_independent_sem_agg_level_prompts( + state.documents, + model, + instruction, + tree_level=state.tree_level, + lotus_native_prompt=output_cols is None, + ) + if output_cols is not None and len(prompts) == 1: + structured_prompts.append(prompts[0]) + structured_states.append(state) + continue + regular_counts[state.group_index] = len(prompts) + for prompt_index, prompt in enumerate(prompts): + regular_prompts.append(prompt) + regular_refs.append((state, prompt_index)) + + if regular_prompts: + regular_output = model( + regular_prompts, + progress_bar_desc=config.sem_agg_progress_bar_desc, + ) + raw_outputs = [str(value) for value in regular_output.outputs] + if len(raw_outputs) != len(regular_prompts): + raise ValueError( + "sem_agg provider batch returned an unexpected number of " + f"outputs: expected {len(regular_prompts)}, got {len(raw_outputs)}" + ) + outputs_by_group: dict[int, list[str]] = { + group_index: [""] * count + for group_index, count in regular_counts.items() + } + for (state, prompt_index), raw_output in zip( + regular_refs, + raw_outputs, + strict=True, + ): + outputs_by_group[state.group_index][prompt_index] = raw_output + for state in states: + outputs = outputs_by_group.get(state.group_index) + if outputs is None: + continue + state.documents = outputs + state.tree_level += 1 + if len(outputs) == 1 and output_cols is None: + state.execution = _SemAggExecution(raw_output=outputs[0]) + + if structured_prompts: + failure_extra = ( + { + prompt_index: dict(failure_extra_by_group[state.group_index]) + for prompt_index, state in enumerate(structured_states) + } + if failure_extra_by_group is not None + else None + ) + retry_result = execute_structured_lm_retry_result( + model, + structured_prompts, + lm_kwargs={ + "progress_bar_desc": config.sem_agg_progress_bar_desc, + **structured_sem_agg_model_kwargs(config), + "response_format": JSON_OBJECT_RESPONSE_FORMAT, + }, + output_cols=output_cols or (), + shape="object", + require_explanation=False, + operator="sem_agg", + max_retries=config.structured_parse_retries, + failure_extra_by_index=failure_extra, + ) + artifact_by_index = dict( + zip( + retry_result.invalid_indices, + retry_result.failure_artifact_paths, + strict=True, + ) + ) + invalid_indices = set(retry_result.invalid_indices) + for prompt_index, state in enumerate(structured_states): + invalid = prompt_index in invalid_indices + state.execution = _SemAggExecution( + raw_output=str(retry_result.raw_outputs[prompt_index]), + retry_result=StructuredLMRetryResult( + raw_outputs=(retry_result.raw_outputs[prompt_index],), + raw_output_attempts=( + retry_result.raw_output_attempts[prompt_index], + ), + invalid_indices=(0,) if invalid else (), + failure_artifact_paths=( + (artifact_by_index[prompt_index],) if invalid else () + ), + ), + ) + + if config.sem_agg_safe_mode: + model.print_total_usage() + + executions = [state.execution for state in states] + if any(execution is None for execution in executions): + raise RuntimeError("sem_agg provider batching did not complete every group") + return [execution for execution in executions if execution is not None] + + +def _build_independent_sem_agg_level_prompts( + documents: Sequence[str], + model: Any, + instruction: str, + *, + tree_level: int, + lotus_native_prompt: bool, +) -> list[list[dict[str, str]]]: + """Build one LOTUS-compatible tree level for one aggregate group.""" + + template = ( + leaf_instruction_template( + instruction, + lotus_native_spacing=lotus_native_prompt, + ) + if tree_level == 0 + else node_instruction_template( + instruction, + lotus_native_spacing=lotus_native_prompt, + ) + ) + template_tokens = model.count_tokens(template) + budget = model.max_ctx_len - model.max_tokens - template_tokens + context_str = "" + context_tokens = 0 + document_counter = 1 + prompts: list[list[dict[str, str]]] = [] + + for document in documents: + formatted = format_aggregate_doc( + tree_level, + str(document), + document_counter, + ) + new_tokens = model.count_tokens(formatted) + if new_tokens + context_tokens > budget and document_counter > 1: + prompt = template.replace("{{docs_str}}", context_str) + prompts.append([{"role": "user", "content": prompt}]) + document_counter = 1 + formatted = format_aggregate_doc( + tree_level, str(document), document_counter + ) + new_tokens = model.count_tokens(formatted) + context_str = "" + context_tokens = 0 + if new_tokens > budget: + raise ValueError( + "sem_agg document does not fit one prompt: " + f"{new_tokens} tokens exceed the {budget}-token budget" + ) + context_str += formatted + context_tokens += new_tokens + document_counter += 1 + + if document_counter > 1 or len(documents) == 1: + prompt = template.replace("{{docs_str}}", context_str) + prompts.append([{"role": "user", "content": prompt}]) + return prompts + + def lotus_style_sem_agg( docs: Sequence[str], model: Any, @@ -519,6 +982,7 @@ def _execute_lotus_style_sem_agg( else node_instruction_template(user_instruction) ) template_tokens = model.count_tokens(template) + budget = model.max_ctx_len - model.max_tokens - template_tokens context_tokens = 0 doc_counter = 1 new_partition_ids: list[int] = [] @@ -529,9 +993,9 @@ def _execute_lotus_style_sem_agg( new_tokens = model.count_tokens(formatted_doc) if ( - new_tokens + context_tokens + template_tokens - > model.max_ctx_len - model.max_tokens - ) or (partition_id != current_partition_id and not do_fold): + new_tokens + context_tokens > budget + or (partition_id != current_partition_id and not do_fold) + ) and doc_counter > 1: prompt = template.replace("{{docs_str}}", context_str) lotus.logger.debug(f"Prompt added to batch: {prompt}") batch.append([{"role": "user", "content": prompt}]) @@ -540,13 +1004,17 @@ def _execute_lotus_style_sem_agg( doc_counter = 1 formatted_doc = format_aggregate_doc(tree_level, doc, doc_counter) - context_str = formatted_doc - context_tokens = new_tokens - doc_counter += 1 - else: - context_str += formatted_doc - context_tokens += new_tokens - doc_counter += 1 + new_tokens = model.count_tokens(formatted_doc) + context_str = "" + context_tokens = 0 + if new_tokens > budget: + raise ValueError( + "sem_agg document does not fit one prompt: " + f"{new_tokens} tokens exceed the {budget}-token budget" + ) + context_str += formatted_doc + context_tokens += new_tokens + doc_counter += 1 if doc_counter > 1 or len(doc_list) == 1: prompt = template.replace("{{docs_str}}", context_str) @@ -636,16 +1104,23 @@ def write_sem_agg_failure_artifact( 0: { "group_index": group_index, "final_instruction": instruction, - "group_row_preview": group.head(5).astype(str).to_dict(orient="records"), + "group_row_preview": group.head(5) + .astype(str) + .to_dict(orient="records"), } }, ) return paths[0] -def leaf_instruction_template(user_instruction: str) -> str: +def leaf_instruction_template( + user_instruction: str, + *, + lotus_native_spacing: bool = False, +) -> str: """Return the LOTUS leaf-level semantic aggregation prompt template.""" + instruction_prefix = "Instruction: " if lotus_native_spacing else "Instruction: " return ( "Your job is to provide an answer to the user's instruction given the context below from multiple documents.\n" "Remember that your job is to answer the user's instruction by combining all relevant information from all provided documents, into a single coherent answer.\n" @@ -654,13 +1129,18 @@ def leaf_instruction_template(user_instruction: str) -> str: "Follow the following format.\n\nContext: relevant facts from multiple documents\n\n" "Instruction: the instruction provided by the user\n\nAnswer: Write your answer\n\n---\n\n" "Context: {{docs_str}}\n\n" - f"Instruction: {user_instruction}\n\nAnswer:\n" + f"{instruction_prefix}{user_instruction}\n\nAnswer:\n" ) -def node_instruction_template(user_instruction: str) -> str: +def node_instruction_template( + user_instruction: str, + *, + lotus_native_spacing: bool = False, +) -> str: """Return the LOTUS intermediate-node semantic aggregation prompt template.""" + instruction_prefix = "Instruction: " if lotus_native_spacing else "Instruction: " return ( "Your job is to provide an answer to the user's instruction given the context below from multiple sources.\n" "Note that each source may be formatted differently and contain information about several different documents.\n" @@ -673,7 +1153,7 @@ def node_instruction_template(user_instruction: str) -> str: "Follow the following format.\n\nContext: relevant facts from multiple sources\n\n" "Instruction: the instruction provided by the user\n\nAnswer: Write your answer\n\n---\n\n" "Context: {{docs_str}}\n\n" - f"Instruction: {user_instruction}\n\nAnswer:\n" + f"{instruction_prefix}{user_instruction}\n\nAnswer:\n" ) @@ -691,7 +1171,9 @@ def parse_structured_sem_agg_output( """Parse and validate one structured LOTUS sem_agg output.""" if isinstance(raw_output, Mapping): - missing = [column.name for column in output_cols if column.name not in raw_output] + missing = [ + column.name for column in output_cols if column.name not in raw_output + ] if missing: raise ValueError(f"sem_agg JSON output is missing required keys: {missing}") return structured_scalar_values( @@ -756,7 +1238,8 @@ def aggregate_groups_with_keys( str(key) for key in source.attrs.get("agent_memory_groupby_keys", ()) ) partition_keys = tuple( - str(key) for key in source.attrs.get("agent_memory_sem_groupby_partition_by", ()) + str(key) + for key in source.attrs.get("agent_memory_sem_groupby_partition_by", ()) ) if deterministic_keys: return _aggregate_groups_by_keys( @@ -786,7 +1269,9 @@ def _aggregate_groups_by_keys( missing = [column for column in group_keys if column not in source.columns] if missing: - raise ValueError(f"aggregate group key columns not found in DataFrame: {missing}") + raise ValueError( + f"aggregate group key columns not found in DataFrame: {missing}" + ) if source.empty: return [] groups: list[tuple[dict[str, Any], pd.DataFrame]] = [] @@ -797,7 +1282,9 @@ def _aggregate_groups_by_keys( groups.append( ( output_key_values, - group.drop(columns=list(drop_columns), errors="ignore").reset_index(drop=True), + group.drop(columns=list(drop_columns), errors="ignore").reset_index( + drop=True + ), ) ) return groups diff --git a/src/agent_memory/adapters/lotus/sem_agg_batch_prompting.py b/src/agent_memory/adapters/lotus/sem_agg_batch_prompting.py new file mode 100644 index 0000000..623284a --- /dev/null +++ b/src/agent_memory/adapters/lotus/sem_agg_batch_prompting.py @@ -0,0 +1,242 @@ +"""Batch-prompting access path for independent semantic aggregate groups.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import json +from typing import Any + +from agent_memory.adapters.lotus.json_output import ( + load_structured_json_with_syntax_repair, +) +from agent_memory.adapters.lotus.prompt_batching import ( + ParsedPromptBatch, + PromptBatchItem, + PromptBatchRequest, + PromptBatching, + run_prompt_batches, +) +from agent_memory.adapters.lotus.structured import ( + structured_scalar_values, +) +from agent_memory.policy.logical import ColumnSpec + + +SEM_AGG_BATCH_SYSTEM_PROMPT = ( + "The user will provide one semantic aggregation instruction and several " + "independent groups. Aggregate every group separately using all documents in " + "that group. Never combine, compare, or copy information across group IDs. " + "Return every supplied group_id exactly once, with exactly the requested " + "output fields. Return only the requested JSON object." +) + + +@dataclass(frozen=True) +class BatchPromptedSemAggResult: + """Per-group outputs plus prompt accounting for one aggregate call.""" + + outputs: tuple[Mapping[str, Any], ...] + raw_output_attempts: tuple[tuple[str, ...], ...] + prompt_count: int + retry_count: int + repair_methods: tuple[str | None, ...] + + +@dataclass(frozen=True) +class _BatchPromptTask: + group_index: int + group_id: str + documents: tuple[str, ...] + + +def execute_batch_prompted_sem_agg( + documents_by_group: Sequence[Sequence[str]], + *, + instruction: str, + output_cols: Sequence[ColumnSpec], + model: Any, + prompt_batching: PromptBatching, + max_retries: int, + model_kwargs: Mapping[str, Any], + progress_bar_desc: str, + trace_dir: Any = None, + structured_output_transport: str = "chat-json-object", +) -> BatchPromptedSemAggResult: + """Aggregate independent groups in bounded shared prompts.""" + + output_token_limit = int(model_kwargs.get("max_tokens", model.max_tokens)) + tasks = tuple( + _BatchPromptTask( + group_index=group_index, + group_id=f"group_{group_index}", + documents=tuple(str(document) for document in documents), + ) + for group_index, documents in enumerate(documents_by_group) + ) + if not tasks: + return BatchPromptedSemAggResult((), (), 0, 0, ()) + + request_model_kwargs = dict(model_kwargs) + request_model_kwargs.pop("max_tokens", None) + execution = run_prompt_batches( + tasks, + task_id=lambda task: task.group_id, + build_request=lambda batch: _build_request( + batch, + instruction=instruction, + output_cols=output_cols, + max_tokens=output_token_limit, + ), + parse_results=lambda raw_output: _parse_outputs( + raw_output, + output_cols=output_cols, + ), + model=model, + config=prompt_batching, + output_schema=_aggregate_schema(output_cols), + structured_output_transport=structured_output_transport, + max_retries=max_retries, + progress_bar_desc=progress_bar_desc, + operator="sem_agg", + trace_dir=trace_dir, + model_kwargs=request_model_kwargs, + ) + return BatchPromptedSemAggResult( + outputs=execution.outputs, + raw_output_attempts=execution.raw_output_attempts, + prompt_count=execution.prompt_count, + retry_count=execution.retry_count, + repair_methods=execution.repair_methods, + ) + + +def _aggregate_schema(output_cols: Sequence[ColumnSpec]) -> dict[str, Any]: + fields = { + column.name: {"type": ["string", "number", "boolean", "null"]} + for column in output_cols + } + return { + "type": "object", + "properties": { + "results": { + "type": "array", + "items": { + "type": "object", + "properties": { + "group_id": {"type": "string"}, + "output": { + "type": "object", + "properties": fields, + "required": list(fields), + "additionalProperties": False, + }, + }, + "required": ["group_id", "output"], + "additionalProperties": False, + }, + } + }, + "required": ["results"], + "additionalProperties": False, + } + + +def _build_request( + tasks: tuple[_BatchPromptTask, ...], + *, + instruction: str, + output_cols: Sequence[ColumnSpec], + max_tokens: int, +) -> PromptBatchRequest: + payload = { + "instruction": instruction, + "groups": [ + { + "group_id": task.group_id, + "documents": list(task.documents), + } + for task in tasks + ], + "output_fields": { + column.name: column.description or "string" for column in output_cols + }, + "output_schema": { + "results": [ + { + "group_id": "group_id", + "output": {column.name: "string" for column in output_cols}, + } + ] + }, + } + return PromptBatchRequest( + task_ids=tuple(task.group_id for task in tasks), + prompt=[ + {"role": "system", "content": SEM_AGG_BATCH_SYSTEM_PROMPT}, + { + "role": "user", + "content": json.dumps(payload, ensure_ascii=False, default=str), + }, + ], + max_tokens=max_tokens, + ) + + +def _parse_outputs( + raw_output: str, + *, + output_cols: Sequence[ColumnSpec], +) -> ParsedPromptBatch[Mapping[str, Any]]: + if not raw_output.strip(): + raise ValueError("batch-prompting sem_agg output is empty") + decoded = load_structured_json_with_syntax_repair( + raw_output, + operator="batch-prompting sem_agg", + expected_shape='JSON object with a "results" array', + ) + payload = decoded.value + if not isinstance(payload, dict) or set(payload) != {"results"}: + raise ValueError("batch-prompting sem_agg output must contain only results") + results = payload["results"] + if not isinstance(results, list): + raise ValueError("batch-prompting sem_agg results must be a list") + + parsed: list[PromptBatchItem[Mapping[str, Any]]] = [] + for result in results: + if not isinstance(result, dict) or set(result) != {"group_id", "output"}: + raise ValueError( + "batch-prompting sem_agg results require group_id and output" + ) + group_id = result["group_id"] + output = result["output"] + if not isinstance(group_id, str): + raise ValueError("batch-prompting sem_agg group_id must be a string") + if not isinstance(output, Mapping): + raise ValueError("batch-prompting sem_agg output must be an object") + expected_fields = {column.name for column in output_cols} + if set(output) != expected_fields: + raise ValueError( + "batch-prompting sem_agg output fields do not match the declared " + f"schema for {group_id!r}" + ) + parsed.append( + PromptBatchItem( + group_id, + structured_scalar_values( + output, + output_cols, + operator="sem_agg", + ), + ) + ) + return ParsedPromptBatch( + items=tuple(parsed), + repair_method=decoded.repair_method, + ) + + +__all__ = [ + "BatchPromptedSemAggResult", + "execute_batch_prompted_sem_agg", +] diff --git a/src/agent_memory/adapters/lotus/sem_filter.py b/src/agent_memory/adapters/lotus/sem_filter.py index ac7fb6a..cb4c5d7 100644 --- a/src/agent_memory/adapters/lotus/sem_filter.py +++ b/src/agent_memory/adapters/lotus/sem_filter.py @@ -11,11 +11,20 @@ import pandas as pd from agent_memory.adapters.lotus.context import LotusExecutionConfig, LotusExecutionContext +from agent_memory.adapters.lotus.predicate import ( + execute_schema_predicate, + single_schema_predicate, + validate_schema_predicate, +) from agent_memory.adapters.lotus.pair_execution import ( PairCandidateSelection, select_semantic_pair_candidates, write_semantic_pair_execution_trace, ) +from agent_memory.adapters.lotus.sem_filter_batch_prompting import ( + execute_batch_prompted_sem_filter, + validate_batch_prompting_sem_filter_config, +) from agent_memory.tracing.semantic import ( query_digest, semantic_trace_scope, @@ -40,6 +49,12 @@ def execute_sem_filter( digest = query_digest(query) profile = context.config.semantic_pair_profiles.get(digest) + prompt_batching = context.config.prompt_batching + batch_prompting = prompt_batching is not None + if batch_prompting: + validate_batch_prompting_sem_filter_config(context.config) + if single_schema_predicate(context.config): + validate_schema_predicate(context.config, "sem_filter") if profile is not None and profile.mode in {"search-filter", "proxy-only"}: if ( profile.mode == "proxy-only" @@ -75,10 +90,37 @@ def execute_sem_filter( oracle_source, str(query.params["instruction"]), ) + prompt_count: int | None = 0 if batch_prompting else None + prompt_retry_count: int | None = ( + 0 if batch_prompting else None + ) + verified_tuple_count: int | None = ( + len(lotus_source) if batch_prompting else None + ) if profile is not None and profile.mode == "proxy-only": result = lotus_source.copy() elif selection is not None and lotus_source.empty: result = lotus_source.copy() + elif single_schema_predicate(context.config): + decision = execute_schema_predicate( + lotus_source, + instruction=instruction, + config=context.config, + operator="sem_filter", + ) + selected_positions = [i for i, keep in enumerate(decision.decisions) if keep] + result = lotus_source.iloc[selected_positions].copy() + elif batch_prompting: + assert prompt_batching is not None + prompted = execute_batch_prompted_sem_filter( + lotus_source, + instruction=instruction, + context=context, + prompt_batching=prompt_batching, + ) + result = prompted.frame + prompt_count = prompted.prompt_count + prompt_retry_count = prompted.retry_count else: result = lotus_source.sem_filter( instruction, @@ -96,25 +138,36 @@ def execute_sem_filter( instruction=str(query.params["instruction"]), profile=profile, ) + trace_payload: dict[str, Any] = { + "instruction": str(query.params["instruction"]), + "lowered_instruction": instruction, + "semantic_pair_profile": ( + "oracle-only" if profile is None else profile.mode + ), + "semantic_pair_profile_fingerprint": ( + None if profile is None else profile.fingerprint + ), + "candidate_pair_count": ( + None if selection is None else selection.candidate_pair_count + ), + } + if prompt_batching is not None: + trace_payload.update( + { + "prompt_batching": prompt_batching.to_dict(), + "prompt_batching_fingerprint": prompt_batching.fingerprint, + "verification_prompt_count": prompt_count, + "verification_retry_count": prompt_retry_count, + "verified_tuple_count": verified_tuple_count, + } + ) write_compact_operator_trace( context.config.trace_dir(), operator="sem_filter", event_type="operator_result", input_frame=source, output_frame=result, - payload={ - "instruction": str(query.params["instruction"]), - "lowered_instruction": instruction, - "semantic_pair_profile": ( - "oracle-only" if profile is None else profile.mode - ), - "semantic_pair_profile_fingerprint": ( - None if profile is None else profile.fingerprint - ), - "candidate_pair_count": ( - None if selection is None else selection.candidate_pair_count - ), - }, + payload=trace_payload, ) return result diff --git a/src/agent_memory/adapters/lotus/sem_filter_batch_prompting.py b/src/agent_memory/adapters/lotus/sem_filter_batch_prompting.py new file mode 100644 index 0000000..b74a0eb --- /dev/null +++ b/src/agent_memory/adapters/lotus/sem_filter_batch_prompting.py @@ -0,0 +1,337 @@ +"""Batch-prompting access path for independent semantic filter tuples.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import json +from typing import TYPE_CHECKING, Any + +import pandas as pd +from lotus.cache import operator_cache + +from agent_memory.adapters.lotus.json_output import ( + load_structured_json_with_syntax_repair, +) +from agent_memory.adapters.lotus.prompt_batching import ( + ParsedPromptBatch, + PromptBatchItem, + PromptBatchRequest, + PromptBatching, + run_prompt_batches, +) + +if TYPE_CHECKING: + from agent_memory.adapters.lotus.context import LotusExecutionContext + + +SEM_FILTER_BATCH_SYSTEM_PROMPT = ( + "The user will provide one semantic filter claim and several independent row " + "contexts. Evaluate the claim independently for every row. Set keep to true " + "exactly when the claim is true for that row's context; otherwise set it to " + "false. Copy every supplied row_id exactly once without adding or changing IDs. " + "Do not compare rows with each other or let one row affect another row's " + "decision. Return only the requested JSON object." +) + + +@dataclass(frozen=True) +class BatchPromptingResult: + """Filtered rows and logical prompt accounting for one operator call.""" + + frame: pd.DataFrame + prompt_count: int + retry_count: int + tuple_count: int + decisions: tuple[bool, ...] = () + raw_output_attempts: tuple[tuple[str, ...], ...] = () + repair_methods: tuple[str | None, ...] = () + + +@dataclass(frozen=True) +class _BatchPromptingTask: + position: int + row_id: str + context: str + + +@dataclass +class _BatchPromptingExecutor: + """Reuse LOTUS whole-operator caching for the opt-in access path.""" + + _obj: pd.DataFrame + + @operator_cache + def __call__( + self, + *, + instruction: str, + prompt_batching: PromptBatching, + structured_parse_retries: int, + structured_max_tokens: int, + progress_bar_desc: str, + trace_dir: Any = None, + operator: str = "sem_filter", + structured_output_transport: str = "chat-json-object", + ) -> BatchPromptingResult: + """Evaluate independent tuples in bounded structured prompts.""" + + import lotus + from lotus.templates import task_instructions + + lm = lotus.settings.lm + if lm is None: + raise ValueError( + "batch-prompting sem_filter requires a configured language model" + ) + columns = tuple(lotus.nl_expression.parse_cols(instruction)) + if not columns: + raise ValueError( + "batch-prompting sem_filter requires instruction column placeholders" + ) + missing = [column for column in columns if column not in self._obj.columns] + if missing: + raise ValueError( + f"batch-prompting sem_filter columns not found in DataFrame: {missing}" + ) + docs = task_instructions.df2multimodal_info(self._obj, list(columns)) + if any(doc.get("image") for doc in docs): + raise NotImplementedError( + "batch-prompting sem_filter currently supports text columns only" + ) + claim = _render_claim(instruction, columns) + tasks = _build_tasks(docs) + if not tasks: + return BatchPromptingResult( + frame=self._obj.copy(), + prompt_count=0, + retry_count=0, + tuple_count=0, + ) + + current_max_tokens = int(getattr(lm, "max_tokens", 512) or 512) + output_token_limit = max(current_max_tokens, structured_max_tokens) + execution = run_prompt_batches( + tasks, + task_id=lambda task: task.row_id, + build_request=lambda batch: _build_request( + batch, + claim=claim, + max_tokens=output_token_limit, + ), + parse_results=_parse_decisions, + output_schema=_decision_schema(), + structured_output_transport=structured_output_transport, + model=lm, + config=prompt_batching, + max_retries=structured_parse_retries, + progress_bar_desc=progress_bar_desc, + operator=operator, + trace_dir=trace_dir, + model_kwargs={ + "show_progress_bar": True, + }, + ) + selected_positions = [ + task.position + for task, keep in zip(tasks, execution.outputs, strict=True) + if keep + ] + return BatchPromptingResult( + frame=self._obj.iloc[selected_positions].copy(), + prompt_count=execution.prompt_count, + retry_count=execution.retry_count, + tuple_count=len(self._obj), + decisions=execution.outputs, + raw_output_attempts=execution.raw_output_attempts, + repair_methods=execution.repair_methods, + ) + + +def execute_batch_prompted_sem_filter( + source: pd.DataFrame, + *, + instruction: str, + context: LotusExecutionContext, + prompt_batching: PromptBatching, +) -> BatchPromptingResult: + """Execute one explicitly configured batch-prompting semantic filter.""" + + return execute_batch_prompted_predicate( + source, + instruction=instruction, + prompt_batching=prompt_batching, + structured_parse_retries=context.config.structured_parse_retries, + structured_max_tokens=context.config.structured_max_tokens, + progress_bar_desc=context.config.sem_filter_progress_bar_desc, + trace_dir=context.config.trace_dir(), + operator="sem_filter", + structured_output_transport=context.config.structured_output_transport, + ) + + +def execute_batch_prompted_predicate( + source: pd.DataFrame, + *, + instruction: str, + prompt_batching: PromptBatching, + structured_parse_retries: int, + structured_max_tokens: int, + progress_bar_desc: str, + trace_dir: Any = None, + operator: str = "sem_filter", + structured_output_transport: str = "chat-json-object", +) -> BatchPromptingResult: + """Evaluate independent boolean predicate tasks in shared prompts.""" + + return _BatchPromptingExecutor(source)( + instruction=instruction, + prompt_batching=prompt_batching, + structured_parse_retries=structured_parse_retries, + structured_max_tokens=structured_max_tokens, + progress_bar_desc=progress_bar_desc, + trace_dir=trace_dir, + operator=operator, + structured_output_transport=structured_output_transport, + ) + + +def validate_batch_prompting_sem_filter_config(config: Any) -> None: + """Reject LOTUS options the new access path cannot preserve.""" + + unsupported: list[str] = [] + if config.sem_filter_examples is not None: + unsupported.append("examples") + if config.sem_filter_helper_examples is not None: + unsupported.append("helper_examples") + if config.sem_filter_strategy is not None: + unsupported.append("strategy") + if config.sem_filter_default is not True: + unsupported.append("default") + if config.sem_filter_cascade_args is not None: + unsupported.append("cascade_args") + if config.sem_filter_safe_mode: + unsupported.append("safe_mode") + if config.sem_filter_additional_cot_instructions: + unsupported.append("additional_cot_instructions") + if unsupported: + raise ValueError( + "batch-prompting sem_filter does not support LOTUS options: " + + ", ".join(unsupported) + ) + + +def _build_tasks( + docs: Sequence[Mapping[str, Any]], +) -> tuple[_BatchPromptingTask, ...]: + return tuple( + _BatchPromptingTask( + position=position, + row_id=f"row_{position}", + context=str(doc.get("text", "")), + ) + for position, doc in enumerate(docs) + ) + + +def _render_claim(instruction: str, columns: Sequence[str]) -> str: + """Render LOTUS placeholders without treating side labels as format specs.""" + + claim = instruction + for column in columns: + claim = claim.replace(f"{{{column}}}", column.capitalize()) + return claim + + +def _build_request( + tasks: tuple[_BatchPromptingTask, ...], + *, + claim: str, + max_tokens: int, +) -> PromptBatchRequest: + payload = { + "claim": claim, + "rows": [ + {"row_id": task.row_id, "context": task.context} for task in tasks + ], + "output_schema": {"decisions": [{"row_id": "row_id", "keep": True}]}, + } + return PromptBatchRequest( + task_ids=tuple(task.row_id for task in tasks), + prompt=[ + {"role": "system", "content": SEM_FILTER_BATCH_SYSTEM_PROMPT}, + { + "role": "user", + "content": json.dumps(payload, ensure_ascii=False, default=str), + }, + ], + max_tokens=max_tokens, + ) + + +def _parse_decisions( + raw_output: str, +) -> ParsedPromptBatch[bool]: + if not raw_output.strip(): + raise ValueError("batch-prompting sem_filter output is empty") + decoded = load_structured_json_with_syntax_repair( + raw_output, + operator="batch-prompting sem_filter", + expected_shape='JSON object with a "decisions" array', + ) + payload = decoded.value + if not isinstance(payload, dict) or set(payload) != {"decisions"}: + raise ValueError( + "batch-prompting sem_filter output must contain only decisions" + ) + raw_decisions = payload["decisions"] + if not isinstance(raw_decisions, list): + raise ValueError("batch-prompting sem_filter decisions must be a list") + + parsed: list[PromptBatchItem[bool]] = [] + for decision in raw_decisions: + if not isinstance(decision, dict) or set(decision) != {"row_id", "keep"}: + raise ValueError( + "batch-prompting sem_filter decisions require row_id and keep" + ) + row_id = decision["row_id"] + keep = decision["keep"] + if not isinstance(row_id, str): + raise ValueError("batch-prompting sem_filter row_id must be a string") + if type(keep) is not bool: + raise ValueError("batch-prompting sem_filter keep must be a boolean") + parsed.append(PromptBatchItem(row_id, keep)) + return ParsedPromptBatch( + items=tuple(parsed), + repair_method=decoded.repair_method, + ) + + +def _decision_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "decisions": { + "type": "array", + "items": { + "type": "object", + "properties": { + "row_id": {"type": "string"}, + "keep": {"type": "boolean"}, + }, + "required": ["row_id", "keep"], + "additionalProperties": False, + }, + } + }, + "required": ["decisions"], + "additionalProperties": False, + } + + +__all__ = [ + "BatchPromptingResult", + "execute_batch_prompted_predicate", + "execute_batch_prompted_sem_filter", + "validate_batch_prompting_sem_filter_config", +] diff --git a/src/agent_memory/adapters/lotus/sem_flat_map.py b/src/agent_memory/adapters/lotus/sem_flat_map.py index 01f237d..614a84d 100644 --- a/src/agent_memory/adapters/lotus/sem_flat_map.py +++ b/src/agent_memory/adapters/lotus/sem_flat_map.py @@ -41,6 +41,8 @@ def execute_sem_flat_map( structured_parse_retries=context.config.structured_parse_retries, semantic_trace_dir=context.config.trace_dir(), operator="sem_flat_map", + prompt_batching=context.config.prompt_batching, + structured_output_transport=context.config.structured_output_transport, ) return apply_flat_map_outputs( source, diff --git a/src/agent_memory/adapters/lotus/sem_groupby.py b/src/agent_memory/adapters/lotus/sem_groupby.py index 2d618b6..ca5b3c6 100644 --- a/src/agent_memory/adapters/lotus/sem_groupby.py +++ b/src/agent_memory/adapters/lotus/sem_groupby.py @@ -19,6 +19,7 @@ write_semantic_pair_execution_trace, ) from agent_memory.adapters.lotus.sem_join import row_text_series +from agent_memory.adapters.lotus.prompt_batching import PromptBatching from agent_memory.adapters.lotus.structured import StructuredLMExecutor from agent_memory.policy.logical import ColumnSpec, QueryExpr from agent_memory.storage.embedding import EmbeddingProvider @@ -85,6 +86,10 @@ def execute_sem_groupby( query_digest_value=digest, profile=profile, embedding_provider=context.pair_embedding_provider, + prompt_batching=context.config.prompt_batching, + structured_output_transport=context.config.structured_output_transport, + structured_parse_retries=context.config.structured_parse_retries, + structured_max_tokens=context.config.structured_max_tokens, ) if labels: return assign_declared_labels( @@ -93,6 +98,8 @@ def execute_sem_groupby( labels=labels, label_col=str(query.params.get("label_col", "_label")), instruction=str(query.params["instruction"]), + prompt_batching=context.config.prompt_batching, + structured_output_transport=context.config.structured_output_transport, ) unique_rows, row_to_unique = exact_unique_key_rows(source, input_cols) @@ -107,6 +114,10 @@ def execute_sem_groupby( query_digest_value=digest, profile=profile, embedding_provider=context.pair_embedding_provider, + prompt_batching=context.config.prompt_batching, + structured_output_transport=context.config.structured_output_transport, + structured_parse_retries=context.config.structured_parse_retries, + structured_max_tokens=context.config.structured_max_tokens, ) result = assign_semantic_group_ids( source, @@ -145,6 +156,10 @@ def execute_partitioned_sem_groupby( query_digest_value: str, profile: SemanticPairExecutionProfile | None, embedding_provider: EmbeddingProvider | None, + prompt_batching: PromptBatching | None, + structured_parse_retries: int, + structured_max_tokens: int, + structured_output_transport: str = "chat-json-object", ) -> pd.DataFrame: """Assign semantic group ids independently within deterministic partitions.""" @@ -168,6 +183,8 @@ def execute_partitioned_sem_groupby( labels=labels, label_col=label_col, instruction=instruction, + prompt_batching=prompt_batching, + structured_output_transport=structured_output_transport, ) else: unique_rows, row_to_unique = exact_unique_key_rows(partition, input_cols) @@ -182,6 +199,10 @@ def execute_partitioned_sem_groupby( query_digest_value=query_digest_value, profile=profile, embedding_provider=embedding_provider, + prompt_batching=prompt_batching, + structured_output_transport=structured_output_transport, + structured_parse_retries=structured_parse_retries, + structured_max_tokens=structured_max_tokens, ) result = assign_semantic_group_ids( partition, @@ -222,6 +243,8 @@ def assign_declared_labels( labels: Sequence[ColumnSpec], label_col: str, instruction: str, + prompt_batching: PromptBatching | None = None, + structured_output_transport: str = "chat-json-object", ) -> pd.DataFrame: """Assign rows to declared closed-world labels with LOTUS structured output.""" @@ -256,6 +279,8 @@ def assign_declared_labels( progress_bar_desc="Grouping labels", model_kwargs={}, operator="sem_groupby", + prompt_batching=prompt_batching, + structured_output_transport=structured_output_transport, ) label_to_group = {name: index for index, name in enumerate(label_names)} @@ -389,6 +414,10 @@ def evaluate_group_matches( query_digest_value: str = "", profile: SemanticPairExecutionProfile | None = None, embedding_provider: EmbeddingProvider | None = None, + prompt_batching: PromptBatching | None = None, + structured_parse_retries: int = 0, + structured_max_tokens: int = 8192, + structured_output_transport: str = "chat-json-object", ) -> list[tuple[int, int]]: """Evaluate candidate row pairs with LOTUS sem_filter.""" @@ -442,34 +471,62 @@ def evaluate_group_matches( "{left} and {right} satisfy this semantic grouping condition: " f"{lowered_instruction}" ) - batch_size = pair_batch_size or len(pairs) - parsed_outputs: list[bool] = [] - raw_outputs: list[Any] = [] - explanations: list[Any] = [] - for start in range(0, len(pairs), batch_size): - pair_batch = pairs.iloc[start : start + batch_size].reset_index(drop=True) - docs = task_instructions.df2multimodal_info(pair_batch, ["left", "right"]) - output = evaluate_group_match_batch( - sem_filter, - docs=docs, - lm=lotus.settings.lm, - instruction=user_instruction, - default=default, - retries=pair_batch_retries, - ) - batch_outputs = list(output.outputs) - if len(batch_outputs) != len(pair_batch): + if prompt_batching is not None: + if default: + raise ValueError( + "prompt-batched sem_groupby does not support default=True" + ) + if pair_batch_size is not None: raise ValueError( - "sem_groupby pair batch returned an unexpected number of outputs: " - f"expected {len(pair_batch)}, got {len(batch_outputs)}" + "prompt_batching cannot be combined with sem_groupby_pair_batch_size" ) - parsed_outputs.extend(bool(value) for value in batch_outputs) - raw_outputs.extend( - aligned_batch_values(output, "raw_outputs", len(pair_batch), "") + from agent_memory.adapters.lotus.sem_filter_batch_prompting import ( + execute_batch_prompted_predicate, ) - explanations.extend( - aligned_batch_values(output, "explanations", len(pair_batch), "") + + prompted = execute_batch_prompted_predicate( + pairs.loc[:, ["left", "right"]], + instruction=user_instruction, + prompt_batching=prompt_batching, + structured_output_transport=structured_output_transport, + structured_parse_retries=structured_parse_retries, + structured_max_tokens=structured_max_tokens, + progress_bar_desc="Grouping comparisons", + trace_dir=trace_dir, + operator="sem_groupby", ) + parsed_outputs = list(prompted.decisions) + raw_outputs = [attempts[-1] for attempts in prompted.raw_output_attempts] + explanations: list[Any] = [""] * len(parsed_outputs) + else: + batch_size = pair_batch_size or len(pairs) + parsed_outputs = [] + raw_outputs = [] + explanations = [] + for start in range(0, len(pairs), batch_size): + pair_batch = pairs.iloc[start : start + batch_size].reset_index(drop=True) + docs = task_instructions.df2multimodal_info(pair_batch, ["left", "right"]) + output = evaluate_group_match_batch( + sem_filter, + docs=docs, + lm=lotus.settings.lm, + instruction=user_instruction, + default=default, + retries=pair_batch_retries, + ) + batch_outputs = list(output.outputs) + if len(batch_outputs) != len(pair_batch): + raise ValueError( + "sem_groupby pair batch returned an unexpected number of outputs: " + f"expected {len(pair_batch)}, got {len(batch_outputs)}" + ) + parsed_outputs.extend(bool(value) for value in batch_outputs) + raw_outputs.extend( + aligned_batch_values(output, "raw_outputs", len(pair_batch), "") + ) + explanations.extend( + aligned_batch_values(output, "explanations", len(pair_batch), "") + ) write_groupby_pair_trace( trace_dir, diff --git a/src/agent_memory/adapters/lotus/sem_join.py b/src/agent_memory/adapters/lotus/sem_join.py index 92619e5..27d0485 100644 --- a/src/agent_memory/adapters/lotus/sem_join.py +++ b/src/agent_memory/adapters/lotus/sem_join.py @@ -4,11 +4,16 @@ from collections import defaultdict from collections.abc import Callable, Mapping, Sequence -from typing import Any +from typing import Any, cast import pandas as pd from agent_memory.adapters.lotus.context import LotusExecutionConfig, LotusExecutionContext +from agent_memory.adapters.lotus.predicate import ( + execute_schema_predicate, + single_schema_predicate, + validate_schema_predicate, +) from agent_memory.adapters.lotus.pair_execution import ( PAIR_LEFT_ID_COLUMN, PAIR_LEFT_TEXT_COLUMN, @@ -39,14 +44,25 @@ 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")) + on = tuple(cast(Sequence[Any], query.params.get("on", ()))) + if single_schema_predicate(context.config): + validate_schema_predicate(context.config, "sem_join") 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: + if ( + context.config.prompt_batching is not None + and context.config.sem_join_cascade_args is not None + ): + raise ValueError( + "prompt-batched sem_join cannot be combined with LOTUS sem_join cascade" + ) + if 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 + # consume deterministically restricted pairs. Keep that boundary explicit # rather than silently dropping `on` semantics or forking LOTUS here. + restriction = "exact-key" if all(isinstance(item, str) for item in on) else "predicate" raise ValueError( - "exact-key sem_join cannot be combined with LOTUS sem_join cascade" + f"{restriction} 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: @@ -98,7 +114,7 @@ def execute_sem_join( context, profile=profile, ) - elif query.params.get("on"): + elif on: left_series, right_series, left_label, right_label, instruction = join_series( left, right, @@ -109,7 +125,8 @@ def execute_sem_join( right_series, left_frame=left, right_frame=right, - on=tuple(str(column) for column in query.params["on"]), + on=on, + query=query, ) if profile is not None and profile.mode in {"search-filter", "proxy-only"}: assert context.pair_embedding_provider is not None @@ -140,6 +157,22 @@ def execute_sem_join( profile=profile, embedding_provider=context.pair_embedding_provider, ) + elif ( + context.config.prompt_batching is not None + or single_schema_predicate(context.config) + ): + left_series, right_series, left_label, right_label, instruction = join_series( + left, + right, + str(query.params["instruction"]), + ) + join_results = verify_semantic_join_candidates( + semantic_join_pair_candidates(left_series, right_series), + left_label=left_label, + right_label=right_label, + instruction=instruction, + config=context.config, + ) else: join_results = evaluate_semantic_join(query, left, right, context.config) result = assemble_join_frame( @@ -343,8 +376,6 @@ def verify_semantic_join_candidates( if candidates.empty: return [] - import lotus - from lotus.sem_ops.sem_filter import sem_filter from lotus.templates import task_instructions oracle_frame = pd.DataFrame( @@ -353,39 +384,76 @@ def verify_semantic_join_candidates( right_label: candidates[PAIR_RIGHT_TEXT_COLUMN], } ) - docs = task_instructions.df2multimodal_info( - oracle_frame, - [left_label, right_label], - ) - output = sem_filter( - docs, - lotus.settings.lm, - instruction, - examples_multimodal_data=examples_multimodal_data( - config.sem_join_examples, - left_label=left_label, - right_label=right_label, - ), - examples_answers=example_answers(config.sem_join_examples), - cot_reasoning=example_reasoning(config.sem_join_examples), - default=config.sem_join_default, - strategy=normalize_strategy(config.sem_join_strategy), - safe_mode=config.sem_join_safe_mode, - progress_bar_desc=config.sem_join_progress_bar_desc, - ) - outputs = list(output.outputs) + if config.prompt_batching is not None: + _validate_prompt_batched_sem_join_config(config) + from agent_memory.adapters.lotus.sem_filter_batch_prompting import ( + execute_batch_prompted_predicate, + ) + + prompted = execute_batch_prompted_predicate( + oracle_frame, + instruction=_batch_prompt_join_instruction( + instruction, + columns=(left_label, right_label), + ), + prompt_batching=config.prompt_batching, + structured_output_transport=config.structured_output_transport, + structured_parse_retries=config.structured_parse_retries, + structured_max_tokens=config.structured_max_tokens, + progress_bar_desc=config.sem_join_progress_bar_desc, + trace_dir=config.trace_dir(), + operator="sem_join", + ) + outputs = list(prompted.decisions) + raw_outputs = [attempts[-1] for attempts in prompted.raw_output_attempts] + explanations: list[str | None] = [None] * len(outputs) + elif single_schema_predicate(config): + decision = execute_schema_predicate( + oracle_frame, + instruction=instruction, + config=config, + operator="sem_join", + ) + outputs = list(decision.decisions) + raw_outputs = list(decision.raw_outputs) + explanations = [None] * len(outputs) + else: + import lotus + from lotus.sem_ops.sem_filter import sem_filter + + docs = task_instructions.df2multimodal_info( + oracle_frame, + [left_label, right_label], + ) + output = sem_filter( + docs, + lotus.settings.lm, + instruction, + examples_multimodal_data=examples_multimodal_data( + config.sem_join_examples, + left_label=left_label, + right_label=right_label, + ), + examples_answers=example_answers(config.sem_join_examples), + cot_reasoning=example_reasoning(config.sem_join_examples), + default=config.sem_join_default, + strategy=normalize_strategy(config.sem_join_strategy), + safe_mode=config.sem_join_safe_mode, + progress_bar_desc=config.sem_join_progress_bar_desc, + ) + outputs = list(output.outputs) + raw_outputs = aligned_join_values(output, "raw_outputs", len(candidates), "") + explanations = aligned_join_values( + output, + "explanations", + len(candidates), + None, + ) if len(outputs) != len(candidates): raise ValueError( "sem_join candidate verification returned an unexpected number of " f"outputs: expected {len(candidates)}, got {len(outputs)}" ) - raw_outputs = aligned_join_values(output, "raw_outputs", len(candidates), "") - explanations = aligned_join_values( - output, - "explanations", - len(candidates), - None, - ) write_selected_join_pair_trace( config.trace_dir(), candidates, @@ -406,22 +474,84 @@ def verify_semantic_join_candidates( ] +def _batch_prompt_join_instruction( + instruction: str, + *, + columns: Sequence[str], +) -> str: + """Describe join fields already embedded in collapsed left/right records.""" + + from lotus.nl_expression import parse_cols + + available = set(columns) + rendered = instruction + for column in parse_cols(instruction): + if column in available: + continue + if column.endswith(":left"): + description = f"the left row's {column[:-5]} field" + elif column.endswith(":right"): + description = f"the right row's {column[:-6]} field" + else: + continue + rendered = rendered.replace(f"{{{column}}}", description) + return rendered + + +def _validate_prompt_batched_sem_join_config(config: LotusExecutionConfig) -> None: + unsupported: list[str] = [] + if config.sem_join_examples is not None: + unsupported.append("examples") + if config.sem_join_strategy is not None: + unsupported.append("strategy") + if config.sem_join_default: + unsupported.append("default=True") + if config.sem_join_safe_mode: + unsupported.append("safe_mode") + if unsupported: + raise ValueError( + "prompt-batched sem_join does not support LOTUS options: " + + ", ".join(unsupported) + ) + + 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] = (), + on: Sequence[Any] = (), + query: QueryExpr | None = None, ) -> pd.DataFrame: - """Return semantic join pairs after optional exact-key restriction.""" + """Return semantic join pairs after optional deterministic restriction.""" if on and (left_frame is None or right_frame is None): - raise ValueError("exact-key semantic join candidates require both input frames") + raise ValueError("restricted semantic join candidates require both input frames") + + exact_keys = tuple(str(item) for item in on if isinstance(item, str)) + predicate_on = tuple(item for item in on if isinstance(item, Mapping)) + if len(exact_keys) + len(predicate_on) != len(on) or ( + exact_keys and predicate_on + ): + raise ValueError("sem_join on must contain exact keys or predicate expressions") + if predicate_on: + if query is None or left_frame is None or right_frame is None: + raise ValueError("predicate semantic join candidates require the join query") + from agent_memory.adapters.lotus.relational import predicate_join_row_ids + + restricted_pair_ids = predicate_join_row_ids( + query, + left_frame, + right_frame, + on=predicate_on, + ) + return _semantic_join_pair_frame(left, right, restricted_pair_ids) + missing = ( sorted( - set(on).difference(left_frame.columns).union( - set(on).difference(right_frame.columns) + set(exact_keys).difference(left_frame.columns).union( + set(exact_keys).difference(right_frame.columns) ) ) if left_frame is not None and right_frame is not None @@ -431,32 +561,42 @@ def semantic_join_pair_candidates( raise ValueError(f"sem_join on columns not found on both sides: {missing}") right_by_key: dict[tuple[object, ...], list[Any]] = defaultdict(list) - if on: + if exact_keys: 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 - ) + key = _semantic_join_key(right_frame.loc[right_id], exact_keys) + right_by_key[key].append(right_id) - rows: list[dict[str, Any]] = [] - for left_id, left_value in left.items(): - if on: + pair_ids: list[tuple[Any, Any]] = [] + for left_id in left.index: + if exact_keys: assert left_frame is not None right_ids: Sequence[Any] = right_by_key.get( - _semantic_join_key(left_frame.loc[left_id], on), + _semantic_join_key(left_frame.loc[left_id], exact_keys), (), ) 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.loc[right_id], - } - for right_id in right_ids - ) + right_ids = tuple(right.index) + pair_ids.extend((left_id, right_id) for right_id in right_ids) + return _semantic_join_pair_frame(left, right, pair_ids) + + +def _semantic_join_pair_frame( + left: pd.Series, + right: pd.Series, + pair_ids: Sequence[tuple[Any, Any]], +) -> pd.DataFrame: + """Materialize canonical semantic-pair rows from stable input IDs.""" + + rows = [ + { + PAIR_LEFT_ID_COLUMN: left_id, + PAIR_RIGHT_ID_COLUMN: right_id, + PAIR_LEFT_TEXT_COLUMN: left.loc[left_id], + PAIR_RIGHT_TEXT_COLUMN: right.loc[right_id], + } + for left_id, right_id in pair_ids + ] return pd.DataFrame( rows, columns=pd.Index( diff --git a/src/agent_memory/adapters/lotus/sem_map.py b/src/agent_memory/adapters/lotus/sem_map.py index 918ea30..799c510 100644 --- a/src/agent_memory/adapters/lotus/sem_map.py +++ b/src/agent_memory/adapters/lotus/sem_map.py @@ -31,7 +31,7 @@ def execute_sem_map( context.configure() source = execute(query.inputs[0], inputs) output_cols = output_columns(query) - if len(output_cols) == 1: + if len(output_cols) == 1 and context.config.prompt_batching is None: return execute_native_sem_map(query, source, output_cols[0], context.config) return execute_structured_sem_map(query, source, output_cols, context.config) @@ -95,6 +95,8 @@ def execute_structured_sem_map( structured_max_tokens=config.structured_max_tokens, structured_parse_retries=config.structured_parse_retries, semantic_trace_dir=config.trace_dir(), + prompt_batching=config.prompt_batching, + structured_output_transport=config.structured_output_transport, ) return apply_structured_map_outputs( source, diff --git a/src/agent_memory/adapters/lotus/sem_topk_join.py b/src/agent_memory/adapters/lotus/sem_topk_join.py index 95122ee..be6fc7a 100644 --- a/src/agent_memory/adapters/lotus/sem_topk_join.py +++ b/src/agent_memory/adapters/lotus/sem_topk_join.py @@ -10,6 +10,12 @@ import pandas as pd from agent_memory.adapters.lotus.context import LotusExecutionContext +from agent_memory.adapters.lotus.prompt_batching import ( + ParsedPromptBatch, + PromptBatchItem, + PromptBatchRequest, + run_prompt_batches, +) from agent_memory.adapters.lotus.pair_execution import ( PAIR_LEFT_ID_COLUMN, PAIR_LEFT_TEXT_COLUMN, @@ -21,6 +27,9 @@ ) from agent_memory.adapters.lotus.sem_topk import LOTUS_PAIRWISE_METHODS from agent_memory.adapters.lotus.structured import normalize_strategy +from agent_memory.adapters.lotus.json_output import ( + load_structured_json_with_syntax_repair, +) from agent_memory.policy.logical import QueryExpr from agent_memory.tracing.semantic import query_digest, write_pair_trace @@ -33,10 +42,18 @@ "match. If too many candidates match, select the strongest matches. Return an " "empty selected_ids array when none match. Return only a JSON object." ) +LISTWISE_JOIN_BATCH_SYSTEM_PROMPT = ( + "The user will provide several independent semantic join tasks. Resolve each " + "task independently using its join condition, left row, candidate rows, and " + "max_matches. Do not use one task as evidence for another. Return every " + "supplied task_id exactly once with its selected_ids. Do not change or invent " + "IDs. Return only the requested JSON object." +) @dataclass(frozen=True) class _ListwiseTask: + task_id: str left_id: Any prompt: list[dict[str, Any]] candidate_positions: Mapping[str, int] @@ -66,7 +83,8 @@ def evaluate_sem_topk_join( right_series, left_frame=left, right_frame=right, - on=tuple(str(column) for column in query.params.get("on", ())), + on=tuple(query.params.get("on", ())), + query=query, ) candidates, scores = _apply_pair_profile( query, @@ -172,11 +190,91 @@ def _listwise_topk( right_label=right_label, k=k, ) + if context.config.prompt_batching is not None: + task_by_id = {task.task_id: task for task in tasks} + execution = run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=_build_listwise_batch_request, + parse_results=lambda raw_output: _parse_listwise_prompt_batch( + raw_output, + tasks=task_by_id, + k=k, + ), + model=lm, + config=context.config.prompt_batching, + output_schema=_listwise_batch_schema(), + structured_output_transport=context.config.structured_output_transport, + max_retries=context.config.structured_parse_retries, + progress_bar_desc="Listwise join resolution", + operator="sem_join", + trace_dir=context.config.trace_dir(), + ) + selected_by_left = { + task.left_id: tuple( + task.candidate_positions[candidate_id] + for candidate_id in selected_ids + ) + for task, selected_ids in zip(tasks, execution.outputs, strict=True) + } + retry_count = execution.retry_count + else: + selected_by_left, retry_count = _execute_listwise_tasks( + tasks, + lm=lm, + k=k, + max_retries=context.config.structured_parse_retries, + ) + + 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 _execute_listwise_tasks( + tasks: Sequence[_ListwiseTask], + *, + lm: Any, + k: int, + max_retries: int, +) -> tuple[dict[Any, tuple[int, ...]], int]: 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): + for _attempt in range(max_retries + 1): if not pending: break lm_call: Any = lm @@ -216,44 +314,10 @@ def _listwise_topk( 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}" + f"{max_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 + return selected_by_left, retry_count def _listwise_tasks( @@ -265,7 +329,9 @@ def _listwise_tasks( k: int, ) -> list[_ListwiseTask]: tasks: list[_ListwiseTask] = [] - for left_id, group in candidates.groupby(PAIR_LEFT_ID_COLUMN, sort=False): + for task_index, (left_id, group) in enumerate( + candidates.groupby(PAIR_LEFT_ID_COLUMN, sort=False) + ): candidate_positions = { f"candidate_{position}": int(position) for position in group.index } @@ -285,6 +351,7 @@ def _listwise_tasks( } tasks.append( _ListwiseTask( + task_id=f"task_{task_index}", left_id=left_id, prompt=[ {"role": "system", "content": LISTWISE_JOIN_SYSTEM_PROMPT}, @@ -299,6 +366,106 @@ def _listwise_tasks( return tasks +def _listwise_batch_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "results": { + "type": "array", + "items": { + "type": "object", + "properties": { + "task_id": {"type": "string"}, + "selected_ids": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["task_id", "selected_ids"], + "additionalProperties": False, + }, + } + }, + "required": ["results"], + "additionalProperties": False, + } + + +def _build_listwise_batch_request( + tasks: tuple[_ListwiseTask, ...], +) -> PromptBatchRequest: + payload = { + "tasks": [ + {"task_id": task.task_id, "messages": task.prompt} for task in tasks + ], + "output_schema": { + "results": [ + {"task_id": "task_id", "selected_ids": ["candidate_id"]} + ] + }, + } + return PromptBatchRequest( + task_ids=tuple(task.task_id for task in tasks), + prompt=[ + {"role": "system", "content": LISTWISE_JOIN_BATCH_SYSTEM_PROMPT}, + { + "role": "user", + "content": json.dumps(payload, ensure_ascii=False, default=str), + }, + ], + max_tokens=LISTWISE_JOIN_MAX_TOKENS, + ) + + +def _parse_listwise_prompt_batch( + raw_output: str, + *, + tasks: Mapping[str, _ListwiseTask], + k: int, +) -> ParsedPromptBatch[tuple[str, ...]]: + decoded = load_structured_json_with_syntax_repair( + raw_output, + operator="prompt-batched listwise sem_join", + expected_shape='JSON object with a "results" array', + ) + payload = decoded.value + if not isinstance(payload, Mapping) or set(payload) != {"results"}: + raise ValueError( + "prompt-batched listwise sem_join output must contain only results" + ) + results = payload["results"] + if not isinstance(results, list): + raise ValueError("prompt-batched listwise sem_join results must be a list") + items: list[PromptBatchItem[tuple[str, ...]]] = [] + for result in results: + if not isinstance(result, Mapping) or set(result) != { + "task_id", + "selected_ids", + }: + raise ValueError( + "prompt-batched listwise sem_join results require task_id and " + "selected_ids" + ) + identifier = result["task_id"] + if not isinstance(identifier, str): + raise ValueError( + "prompt-batched listwise sem_join task_id must be a string" + ) + task = tasks.get(identifier) + if task is None: + selected = tuple( + value for value in result["selected_ids"] if isinstance(value, str) + ) if isinstance(result["selected_ids"], list) else () + else: + selected = _parse_listwise_ids( + json.dumps({"selected_ids": result["selected_ids"]}), + valid_ids=set(task.candidate_positions), + k=k, + ) + items.append(PromptBatchItem(identifier, selected)) + return ParsedPromptBatch( + items=tuple(items), + repair_method=decoded.repair_method, + ) + + def _parse_listwise_ids( raw_output: str, *, diff --git a/src/agent_memory/adapters/lotus/structured.py b/src/agent_memory/adapters/lotus/structured.py index 292cca0..8dd9880 100644 --- a/src/agent_memory/adapters/lotus/structured.py +++ b/src/agent_memory/adapters/lotus/structured.py @@ -4,7 +4,7 @@ import json import re -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -18,7 +18,20 @@ DEFAULT_STRUCTURED_MAX_TOKENS, DEFAULT_STRUCTURED_PARSE_RETRIES, ) -from agent_memory.tracing.semantic import write_structured_generation_trace +from agent_memory.adapters.lotus.json_output import ( + load_structured_json_with_syntax_repair, +) +from agent_memory.adapters.lotus.prompt_batching import ( + ParsedPromptBatch, + PromptBatchItem, + PromptBatchRequest, + PromptBatching, + run_prompt_batches, +) +from agent_memory.tracing.semantic import ( + trace_scope_value, + write_structured_generation_trace, +) from agent_memory.policy.logical import ColumnSpec, QueryExpr EXPLANATION_FIELD = "_explanation" @@ -26,6 +39,13 @@ STRUCTURED_RESERVED_MODEL_KWARGS = {"progress_bar_desc", "response_format"} RAW_OUTPUT_PREVIEW_CHARS = 240 STRUCTURED_FAILURE_DIR = Path(".memory-test") / "structured-failures" / "latest" +STRUCTURED_BATCH_SYSTEM_PROMPT = ( + "The user will provide several independent semantic operator tasks. Follow " + "each task's messages independently. Do not use one task as evidence for " + "another. Return every supplied task_id exactly once without changing or " + "inventing IDs. Put each task's requested JSON result in its output field. " + "Return only the requested JSON object." +) PLACEHOLDER_PATTERN = re.compile( r"(? Any: """Normalize string reasoning strategy names to LOTUS ReasoningStrategy.""" @@ -457,6 +483,8 @@ def __call__( structured_parse_retries: int = DEFAULT_STRUCTURED_PARSE_RETRIES, semantic_trace_dir: Any = None, operator: str = "sem_map", + prompt_batching: PromptBatching | None = None, + structured_output_transport: str = "chat-json-object", ) -> StructuredGenerationResult: """Run a structured LOTUS-backed LM batch and parse JSON outputs.""" @@ -531,9 +559,61 @@ def __call__( show_safe_mode(estimated_cost, len(prompts)) current_max_tokens = int(getattr(lotus.settings.lm, "max_tokens", 512) or 512) + max_tokens = max(current_max_tokens, structured_max_tokens) + if prompt_batching is not None: + execution = _execute_structured_prompt_batches( + prompts, + model=lotus.settings.lm, + prompt_batching=prompt_batching, + structured_output_transport=structured_output_transport, + output_cols=output_cols, + shape=shape, + require_explanation=require_explanation and shape == "object", + operator=operator, + max_tokens=max_tokens, + max_retries=structured_parse_retries, + progress_bar_desc=progress_bar_desc, + model_kwargs=model_kwargs, + trace_dir=semantic_trace_dir, + ) + parsed_outputs = [output[0] for output in execution.outputs] + explanations = [output[1] for output in execution.outputs] + raw_outputs = [attempts[-1] for attempts in execution.raw_output_attempts] + audit_rows = structured_generation_audit_rows( + source=self._obj, + input_cols=input_cols, + output_cols=output_cols, + instruction=instruction, + formatted_instruction=formatted_instruction, + final_instruction=user_instruction, + raw_outputs=raw_outputs, + raw_output_attempts=execution.raw_output_attempts, + shape=shape, + require_explanation=require_explanation and shape == "object", + invalid_indices=(), + failure_artifact_paths=(), + operator=operator, + parsed_outputs_override=parsed_outputs, + syntax_repair_methods=execution.repair_methods, + ) + write_structured_generation_trace( + semantic_trace_dir, + operator=operator, + rows=audit_rows, + snapshots={"input": self._obj.loc[:, list(input_cols)].copy()}, + ) + if safe_mode: + lotus.settings.lm.print_total_usage() + return StructuredGenerationResult( + parsed_outputs=tuple(parsed_outputs), + raw_outputs=tuple(raw_outputs), + explanations=tuple(explanations), + raw_output_attempts=execution.raw_output_attempts, + ) + lm_kwargs: dict[str, Any] = { "progress_bar_desc": progress_bar_desc, - "max_tokens": max(current_max_tokens, structured_max_tokens), + "max_tokens": max_tokens, **dict(model_kwargs), "response_format": {"type": "json_object"}, } @@ -597,6 +677,279 @@ def __call__( ) +def _execute_structured_prompt_batches( + prompts: Sequence[Any], + *, + model: Any, + prompt_batching: PromptBatching, + structured_output_transport: str, + output_cols: Sequence[ColumnSpec], + shape: Literal["object", "array"], + require_explanation: bool, + operator: str, + max_tokens: int, + max_retries: int, + progress_bar_desc: str, + model_kwargs: Mapping[str, Any], + trace_dir: Any, +) -> Any: + tasks = tuple( + _StructuredPromptTask( + task_id=f"task_{index}", + prompt=_text_prompt(prompt, operator=operator), + ) + for index, prompt in enumerate(prompts) + ) + return run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=lambda batch: _build_structured_prompt_batch_request( + batch, + max_tokens=max_tokens, + ), + parse_results=lambda raw_output: _parse_structured_prompt_batch( + raw_output, + output_cols=output_cols, + shape=shape, + require_explanation=require_explanation, + operator=operator, + ), + parse_single_result=lambda raw_output: _parse_single_structured_prompt_output( + raw_output, + output_cols=output_cols, + shape=shape, + require_explanation=require_explanation, + operator=operator, + ), + parse_repaired_results=lambda raw_output: _parse_structured_prompt_batch( + raw_output, + output_cols=output_cols, + shape=shape, + require_explanation=require_explanation, + operator=operator, + strict_fields=True, + ), + output_schema=_structured_batch_schema( + output_cols, shape=shape, require_explanation=require_explanation + ), + structured_output_transport=structured_output_transport, + model=model, + config=prompt_batching, + max_retries=max_retries, + progress_bar_desc=progress_bar_desc, + operator=operator, + trace_dir=trace_dir, + model_kwargs=model_kwargs, + ) + + +def _text_prompt(prompt: Any, *, operator: str) -> list[dict[str, str]]: + if not isinstance(prompt, list) or any( + not isinstance(message, Mapping) + or not isinstance(message.get("role"), str) + or not isinstance(message.get("content"), str) + for message in prompt + ): + raise NotImplementedError( + f"prompt-batched {operator} currently supports text prompts only" + ) + return [ + {"role": str(message["role"]), "content": str(message["content"])} + for message in prompt + ] + + +def _build_structured_prompt_batch_request( + tasks: tuple[_StructuredPromptTask, ...], + *, + max_tokens: int, +) -> PromptBatchRequest: + payload = { + "tasks": [ + {"task_id": task.task_id, "messages": task.prompt} for task in tasks + ], + "output_schema": { + "results": [{"task_id": "task_id", "output": {}}] + }, + } + return PromptBatchRequest( + task_ids=tuple(task.task_id for task in tasks), + prompt=[ + {"role": "system", "content": STRUCTURED_BATCH_SYSTEM_PROMPT}, + { + "role": "user", + "content": json.dumps(payload, ensure_ascii=False), + }, + ], + max_tokens=max_tokens, + ) + + +def _parse_structured_prompt_batch( + raw_output: str, + *, + output_cols: Sequence[ColumnSpec], + shape: Literal["object", "array"], + require_explanation: bool, + operator: str, + strict_fields: bool = False, +) -> ParsedPromptBatch[tuple[Any, str | None]]: + decoded = load_structured_json_with_syntax_repair( + raw_output, + operator=f"prompt-batched {operator}", + expected_shape='JSON object with a "results" array', + ) + payload = decoded.value + if not isinstance(payload, Mapping) or set(payload) != {"results"}: + raise ValueError(f"prompt-batched {operator} output must contain only results") + results = payload["results"] + if not isinstance(results, list): + raise ValueError(f"prompt-batched {operator} results must be a list") + + items: list[PromptBatchItem[tuple[Any, str | None]]] = [] + for result in results: + if not isinstance(result, Mapping) or set(result) != {"task_id", "output"}: + raise ValueError( + f"prompt-batched {operator} results require task_id and output" + ) + identifier = result["task_id"] + output = result["output"] + if not isinstance(identifier, str): + raise ValueError(f"prompt-batched {operator} task_id must be a string") + if not isinstance(output, Mapping): + raise ValueError(f"prompt-batched {operator} output must be an object") + serialized = json.dumps(output, ensure_ascii=False) + if strict_fields: + parsed, explanation = _parse_single_structured_prompt_output( + serialized, + output_cols=output_cols, + shape=shape, + require_explanation=require_explanation, + operator=operator, + ) + items.append(PromptBatchItem(identifier, (parsed, explanation))) + continue + if shape == "object": + parsed, explanation = parse_structured_object_json( + serialized, + output_cols, + require_explanation=require_explanation, + operator=operator, + ) + else: + parsed = parse_structured_array_json( + serialized, + output_cols, + operator=operator, + ) + explanation = None + items.append(PromptBatchItem(identifier, (parsed, explanation))) + return ParsedPromptBatch( + items=tuple(items), + repair_method=decoded.repair_method, + ) + + +def _structured_batch_schema( + output_cols: Sequence[ColumnSpec], *, shape: str, require_explanation: bool +) -> dict[str, Any]: + fields: dict[str, Any] = { + column.name: {"type": ["string", "number", "boolean", "null"]} + for column in output_cols + } + if require_explanation: + fields[EXPLANATION_FIELD] = {"type": "string"} + output: dict[str, Any] = { + "type": "object", + "properties": fields, + "required": list(fields), + "additionalProperties": False, + } + if shape == "array": + output = { + "type": "object", + "properties": {"rows": {"type": "array", "items": output}}, + "required": ["rows"], + "additionalProperties": False, + } + return { + "type": "object", + "properties": { + "results": { + "type": "array", + "items": { + "type": "object", + "properties": {"task_id": {"type": "string"}, "output": output}, + "required": ["task_id", "output"], + "additionalProperties": False, + }, + } + }, + "required": ["results"], + "additionalProperties": False, + } + + +def _parse_single_structured_prompt_output( + raw_output: str, + *, + output_cols: Sequence[ColumnSpec], + shape: Literal["object", "array"], + require_explanation: bool, + operator: str, +) -> tuple[Any, str | None]: + """Parse one unwrapped task result when its task identity is unambiguous.""" + + decoded = load_structured_json_with_syntax_repair( + raw_output, + operator=f"singleton prompt-batched {operator}", + expected_shape="the operator's exact JSON output object", + ) + if decoded.repair_method is not None: + raise ValueError( + f"singleton prompt-batched {operator} output must be strict JSON" + ) + payload = decoded.value + if not isinstance(payload, Mapping): + raise ValueError( + f"singleton prompt-batched {operator} output must be an object" + ) + + if shape == "object": + expected_fields = {column.name for column in output_cols} + if require_explanation: + expected_fields.add(EXPLANATION_FIELD) + if set(payload) != expected_fields: + raise ValueError( + f"singleton prompt-batched {operator} output fields do not match " + "the declared schema" + ) + return parse_structured_object_json( + raw_output, + output_cols, + require_explanation=require_explanation, + operator=operator, + ) + + if set(payload) != {FLAT_MAP_ROWS_FIELD}: + raise ValueError( + f"singleton prompt-batched {operator} output must contain only " + f"{FLAT_MAP_ROWS_FIELD!r}" + ) + rows = payload[FLAT_MAP_ROWS_FIELD] + expected_fields = {column.name for column in output_cols} + if not isinstance(rows, list) or any( + not isinstance(row, Mapping) or set(row) != expected_fields for row in rows + ): + raise ValueError( + f"singleton prompt-batched {operator} rows do not match the declared schema" + ) + return ( + parse_structured_array_json(raw_output, output_cols, operator=operator), + None, + ) + + def execute_structured_lm_with_retries( model: Any, prompts: Sequence[Any], @@ -642,6 +995,7 @@ def execute_structured_lm_retry_result( operator: str, max_retries: int, failure_extra_by_index: Mapping[int, Mapping[str, Any]] | None = None, + output_validator: Callable[[str], None] | None = None, ) -> StructuredLMRetryResult: """Call the LM and return retry metadata without hiding parse failures.""" @@ -649,9 +1003,12 @@ def execute_structured_lm_retry_result( output: LMOutput = model(prompts, **dict(lm_kwargs)) raw_outputs = list(output.outputs) + if output_validator is not None and len(raw_outputs) != len(prompts): + raise ValueError(f"{operator} expected {len(prompts)} outputs, got {raw_outputs!r}") raw_output_attempts = [[raw_output] for raw_output in raw_outputs] invalid = invalid_structured_output_indices( raw_outputs, + output_validator=output_validator, output_cols=output_cols, shape=shape, require_explanation=require_explanation, @@ -668,12 +1025,18 @@ def execute_structured_lm_retry_result( f"{lm_kwargs.get('progress_bar_desc', 'Structured generation')} retry" ) retry_output: LMOutput = model(retry_prompts, **retry_kwargs) + if output_validator is not None and len(retry_output.outputs) != len(invalid): + raise ValueError( + f"{operator} expected {len(invalid)} retry outputs, " + f"got {retry_output.outputs!r}" + ) for index, raw_output in zip(invalid, retry_output.outputs): raw_outputs[index] = raw_output raw_output_attempts[index].append(raw_output) retries_left -= 1 invalid = invalid_structured_output_indices( raw_outputs, + output_validator=output_validator, output_cols=output_cols, shape=shape, require_explanation=require_explanation, @@ -690,6 +1053,7 @@ def execute_structured_lm_retry_result( require_explanation=require_explanation, operator=operator, extra_by_index=failure_extra_by_index, + output_validator=output_validator, ) _STRUCTURED_RETRY_STATS.failure_artifacts += len(artifact_paths) else: @@ -743,6 +1107,8 @@ def structured_generation_audit_rows( invalid_indices: Sequence[int], failure_artifact_paths: Sequence[Path], operator: str, + parsed_outputs_override: Sequence[Any] | None = None, + syntax_repair_methods: Sequence[str | None] | None = None, ) -> list[dict[str, Any]]: """Build structured generation audit rows without mutating source data.""" @@ -764,6 +1130,8 @@ def structured_generation_audit_rows( require_explanation=require_explanation, operator=operator, ) + elif parsed_outputs_override is not None: + parsed_output = parsed_outputs_override[index] else: parsed, _explanations = parse_structured_outputs( [raw_output], @@ -799,6 +1167,16 @@ def structured_generation_audit_rows( "parsed_output": parsed_output, "parse_error": parse_error, "failure_artifact": artifact_by_index.get(index, ""), + "structured_output_repaired": bool( + syntax_repair_methods + and index < len(syntax_repair_methods) + and syntax_repair_methods[index] + ), + "structured_output_repair_method": ( + syntax_repair_methods[index] + if syntax_repair_methods and index < len(syntax_repair_methods) + else None + ), } ) return rows @@ -814,9 +1192,13 @@ def write_structured_failure_artifacts( require_explanation: bool, operator: str, extra_by_index: Mapping[int, Mapping[str, Any]] | None = None, + output_validator: Callable[[str], None] | None = None, ) -> list[Path]: """Write structured generation failure artifacts for local inspection.""" + if trace_scope_value("provider_journal", False): + # Raw responses are already durable; do not duplicate failed rows as files. + return [] STRUCTURED_FAILURE_DIR.mkdir(parents=True, exist_ok=True) timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ") paths: list[Path] = [] @@ -824,6 +1206,7 @@ def write_structured_failure_artifacts( final_raw_output = raw_output_attempts[index][-1] parse_error = structured_parse_error( final_raw_output, + output_validator=output_validator, output_cols=output_cols, shape=shape, require_explanation=require_explanation, @@ -865,10 +1248,13 @@ def structured_parse_error( shape: Literal["object", "array"], require_explanation: bool, operator: str, + output_validator: Callable[[str], None] | None = None, ) -> str: """Return the parser error message for one structured raw output.""" try: + if output_validator is not None: + output_validator(raw_output) parse_one_structured_output( raw_output, output_cols=output_cols, @@ -888,12 +1274,15 @@ def invalid_structured_output_indices( shape: Literal["object", "array"], require_explanation: bool, operator: str, + output_validator: Callable[[str], None] | None = None, ) -> list[int]: """Return output indices that fail the structured parser.""" invalid: list[int] = [] for index, raw_output in enumerate(raw_outputs): try: + if output_validator is not None: + output_validator(raw_output) parse_one_structured_output( raw_output, output_cols=output_cols, diff --git a/src/agent_memory/api.py b/src/agent_memory/api.py index 330af56..17d8e11 100644 --- a/src/agent_memory/api.py +++ b/src/agent_memory/api.py @@ -74,6 +74,19 @@ def __hash__(self) -> int: MessageInput = str | Message | Mapping[str, Any] +@dataclass(frozen=True) +class CountRefresh: + """Refresh memory after accepting a fixed number of source rows.""" + + every: int + + def __post_init__(self) -> None: + if isinstance(self.every, bool) or not isinstance(self.every, int): + raise TypeError("CountRefresh.every must be an integer") + if self.every < 1: + raise ValueError("CountRefresh.every must be at least 1") + + class Memory: """Base class for declarative memory policies.""" @@ -95,6 +108,7 @@ def __init__( *, adapter: Any | None = None, storage: StorageDeployment | None = None, + refresh: CountRefresh | None = None, ) -> None: from .runtime import MemoryRuntime @@ -110,6 +124,7 @@ def __init__( policy, adapter=adapter, storage=storage, + refresh=refresh, ) def add(self, message: MessageInput) -> None: @@ -127,6 +142,17 @@ def query(self, text: str) -> Any: return self._runtime.execute_retrieval_query("default", text) + @property + def pending_count(self) -> int: + """Return the number of accepted rows awaiting maintenance.""" + + return self._runtime.pending_count + + def flush(self) -> None: + """Apply all pending rows as one atomic relation-valued delta.""" + + self._runtime.flush() + @classmethod def spec(cls) -> MemorySpec: """Return the collected memory specification for this Memory subclass. diff --git a/src/agent_memory/dataflow.py b/src/agent_memory/dataflow.py new file mode 100644 index 0000000..c7b6d9d --- /dev/null +++ b/src/agent_memory/dataflow.py @@ -0,0 +1,79 @@ +"""Public facade for executing declarative semantic dataflows.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import pandas as pd + +from .planner.differential_policy import PolicyDifferentiator +from .policy.logical import MemorySpec, MemoryView +from .policy.relation import Log, Relation +from .runtime.executor import PolicyExecutor + +Source = Log + + +class SemanticDataflow: + """Compile and execute named semantic views over one source relation.""" + + def __init__( + self, + *, + source: Source, + views: Mapping[str, Relation], + adapter: Any | None = None, + ) -> None: + if not isinstance(source, Log): + raise TypeError("source must be a Source") + if not views: + raise ValueError("views cannot be empty") + + declared_views: dict[str, MemoryView] = {} + for name, relation in views.items(): + if not isinstance(name, str) or not name: + raise ValueError("view names must be non-empty strings") + if not isinstance(relation, Relation): + raise TypeError("views values must be Relation instances") + pending = [relation.expr] + visited: set[int] = set() + while pending: + query = pending.pop() + if id(query) in visited: + continue + visited.add(id(query)) + if query.op == "log" and query is not source.expr: + raise ValueError( + f"view {name!r} must derive from the declared source" + ) + pending.extend(query.inputs) + declared_views[name] = MemoryView(name=name, query=relation.expr) + + spec = MemorySpec( + log=source, + views=declared_views, + private_relations={}, + ) + self._executor = PolicyExecutor( + PolicyDifferentiator().differentiate(spec), + adapter=adapter, + ) + + def apply(self, rows: pd.DataFrame) -> None: + """Apply one relation-valued source batch to all declared views.""" + + self._executor.apply_delta(rows) + + def view(self, name: str) -> pd.DataFrame: + """Return a defensive copy of one named public view.""" + + return self._executor.read_view(name) + + def snapshot_state(self) -> dict[str, Any]: + """Return the executor's existing snapshot for trusted persistence.""" + return self._executor.snapshot_state() + + def restore_state(self, snapshot: Mapping[str, Any]) -> None: + """Restore a snapshot, including the executor's compatibility checks.""" + self._executor.restore_state(snapshot) diff --git a/src/agent_memory/evaluation/agent_memory_drivers.py b/src/agent_memory/evaluation/agent_memory_drivers.py index 0133c4e..3c1ce53 100644 --- a/src/agent_memory/evaluation/agent_memory_drivers.py +++ b/src/agent_memory/evaluation/agent_memory_drivers.py @@ -20,6 +20,11 @@ semantic_pair_site_id, semantic_pair_site_physical_contract, ) +from agent_memory.adapters.lotus.context import validate_lm_max_ctx_len +from agent_memory.adapters.lotus.prompt_batching import ( + PromptBatching, + validate_structured_output_transport, +) from agent_memory.evaluation.claude_memory.bindings import event_to_claude_log_row from agent_memory.evaluation.embedding_trace import TracingEmbeddingProvider from agent_memory.evaluation.harness import RetrievalOutput @@ -27,6 +32,7 @@ from agent_memory.evaluation.types import BenchmarkEvent, RetrievalRequest from agent_memory.evaluation.zep.answering import format_retrieval_context from agent_memory.policy.retrieval import RetrievalResult +from agent_memory.planner import DEFAULT_GROUPED_AGG_RULE from agent_memory.storage import EmbeddingSpec from agent_memory.tracing.semantic import trace_event_count @@ -247,6 +253,11 @@ def _validate_lotus_cache_mode(mode: str) -> None: ) +def _validate_refresh_every(every: int) -> None: + if isinstance(every, bool) or not isinstance(every, int) or every < 1: + raise ValueError("refresh_every must be a positive integer") + + def build_mem0_semantic_pair_profiles( memory_type: type[Any], *, @@ -442,6 +453,12 @@ def finish_session(self, session_id: str) -> dict[str, Any]: del session_id return {} + def flush(self) -> dict[str, int]: + """Publish any pending source rows before the final checkpoint.""" + + self._memory.flush() + return _claude_memory_shape(self._memory) + def retrieve(self, request: RetrievalRequest) -> RetrievalOutput: """Run ClaudeMemory's declared semantic top-k retrieval unchanged.""" @@ -498,6 +515,12 @@ def finish_session(self, session_id: str) -> dict[str, Any]: del session_id return {} + def flush(self) -> dict[str, int]: + """Publish any pending source rows before the final checkpoint.""" + + self._memory.flush() + return _view_counts(self._memory) + def retrieve(self, request: RetrievalRequest) -> RetrievalOutput: """Run the declared entity/fact retrieval DAG without an answer model.""" @@ -558,6 +581,12 @@ def finish_session(self, session_id: str) -> dict[str, Any]: del session_id return {} + def flush(self) -> dict[str, int]: + """Publish any pending source rows before the final checkpoint.""" + + self._memory.flush() + return _view_counts(self._memory) + def retrieve(self, request: RetrievalRequest) -> RetrievalOutput: """Run indexed cosine retrieval without a generative model call.""" @@ -638,18 +667,27 @@ def __init__( self, *, model_id: str = "deepseek/deepseek-v4-flash", - grouped_agg_rule: str = "rule-all-group", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, sem_topk_method: str = "pairwise-naive", sem_groupby_pair_batch_size: int | None = None, sem_groupby_pair_batch_retries: int = 0, semantic_pair_profiles: dict[str, SemanticPairExecutionProfile] | None = None, + prompt_batching: PromptBatching | None = None, + sem_agg_dispatch: str = "sequential", + structured_output_transport: str = "chat-json-object", + lm_max_ctx_len: int | None = None, semantic_trace_snapshot_mode: str = "compact", lotus_cache_mode: str = "disabled", + refresh_every: int = 1, thinking_enabled: bool = True, ) -> None: from agent_memory.adapters.lotus.context import SEM_TOPK_METHODS from agent_memory.planner.rules import GROUPED_AGG_RULES + validate_lm_max_ctx_len(lm_max_ctx_len) + validate_structured_output_transport( + structured_output_transport, model=model_id + ) if grouped_agg_rule not in GROUPED_AGG_RULES: raise ValueError( "grouped_agg_rule must be one of: " + ", ".join(GROUPED_AGG_RULES) @@ -664,10 +702,16 @@ def __init__( 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 {}) + self.prompt_batching = prompt_batching + self.sem_agg_dispatch = sem_agg_dispatch + self.structured_output_transport = structured_output_transport + self.lm_max_ctx_len = lm_max_ctx_len _semantic_pair_embedding_contract(self.semantic_pair_profiles) self.semantic_trace_snapshot_mode = semantic_trace_snapshot_mode _validate_lotus_cache_mode(lotus_cache_mode) self.lotus_cache_mode = lotus_cache_mode + _validate_refresh_every(refresh_every) + self.refresh_every = refresh_every self.thinking_enabled = thinking_enabled def __call__( @@ -706,6 +750,10 @@ def __call__( semantic_trace_dir=trace_dir, lm_num_retries=BENCHMARK_LM_NUM_RETRIES, lm_model_kwargs={ + **( + {"max_ctx_len": self.lm_max_ctx_len} + if self.lm_max_ctx_len is not None else {} + ), "extra_body": { "thinking": { "type": "enabled" if self.thinking_enabled else "disabled" @@ -718,6 +766,9 @@ def __call__( 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, + prompt_batching=self.prompt_batching, + sem_agg_dispatch=self.sem_agg_dispatch, + structured_output_transport=self.structured_output_transport, semantic_trace_snapshot_mode=self.semantic_trace_snapshot_mode, ), pair_embedding_provider=pair_embedding_provider, @@ -726,7 +777,14 @@ def __call__( policy = PolicyDifferentiator( rules=DifferentialRules(grouped_agg_rule=self.grouped_agg_rule) ).differentiate(am.ClaudeMemory.spec()) - memory._runtime = MemoryRuntime(policy, adapter=adapter) + if self.refresh_every > 1: + memory._runtime = MemoryRuntime( + policy, + adapter=adapter, + refresh=am.CountRefresh(every=self.refresh_every), + ) + else: + memory._runtime = MemoryRuntime(policy, adapter=adapter) return ClaudeMemoryDriver(memory) @@ -739,20 +797,29 @@ def __init__( connector: Any, base_namespace: str, model_id: str = "deepseek/deepseek-v4-flash", - grouped_agg_rule: str = "rule-re-group", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, 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, + prompt_batching: PromptBatching | None = None, + sem_agg_dispatch: str = "sequential", + structured_output_transport: str = "chat-json-object", + lm_max_ctx_len: int | None = None, embedding_device: str = "cpu", semantic_trace_snapshot_mode: str = "compact", lotus_cache_mode: str = "disabled", + refresh_every: int = 1, thinking_enabled: bool = True, neo4j_image: str, neo4j_image_digest: str, ) -> None: from agent_memory.planner.rules import GROUPED_AGG_RULES + validate_lm_max_ctx_len(lm_max_ctx_len) + validate_structured_output_transport( + structured_output_transport, model=model_id + ) if not base_namespace: raise ValueError("Zep benchmark base_namespace must be non-empty") if grouped_agg_rule not in GROUPED_AGG_RULES: @@ -767,6 +834,10 @@ def __init__( 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 {}) + self.prompt_batching = prompt_batching + self.sem_agg_dispatch = sem_agg_dispatch + self.structured_output_transport = structured_output_transport + self.lm_max_ctx_len = lm_max_ctx_len embedding_contract = _semantic_pair_embedding_contract( self.semantic_pair_profiles ) @@ -788,6 +859,8 @@ def __init__( self.semantic_trace_snapshot_mode = semantic_trace_snapshot_mode _validate_lotus_cache_mode(lotus_cache_mode) self.lotus_cache_mode = lotus_cache_mode + _validate_refresh_every(refresh_every) + self.refresh_every = refresh_every self.thinking_enabled = thinking_enabled self.neo4j_image = neo4j_image self.neo4j_image_digest = neo4j_image_digest @@ -799,18 +872,27 @@ def from_environment( *, base_namespace: str, model_id: str = "deepseek/deepseek-v4-flash", - grouped_agg_rule: str = "rule-re-group", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, 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, + prompt_batching: PromptBatching | None = None, + sem_agg_dispatch: str = "sequential", + structured_output_transport: str = "chat-json-object", + lm_max_ctx_len: int | None = None, embedding_device: str = "cpu", semantic_trace_snapshot_mode: str = "compact", lotus_cache_mode: str = "disabled", + refresh_every: int = 1, thinking_enabled: bool = True, ) -> ZepMemoryDriverFactory: """Create the pinned Graphiti-compatible deployment connector.""" + validate_lm_max_ctx_len(lm_max_ctx_len) + validate_structured_output_transport( + structured_output_transport, model=model_id + ) neo4j_image = os.getenv("AGENT_MEMORY_NEO4J_IMAGE") neo4j_image_digest = os.getenv("AGENT_MEMORY_NEO4J_IMAGE_DIGEST") if not neo4j_image or not neo4j_image_digest: @@ -852,9 +934,14 @@ def from_environment( 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, + prompt_batching=prompt_batching, + sem_agg_dispatch=sem_agg_dispatch, + structured_output_transport=structured_output_transport, + lm_max_ctx_len=lm_max_ctx_len, embedding_device=embedding_device, semantic_trace_snapshot_mode=semantic_trace_snapshot_mode, lotus_cache_mode=lotus_cache_mode, + refresh_every=refresh_every, thinking_enabled=thinking_enabled, neo4j_image=neo4j_image, neo4j_image_digest=neo4j_image_digest, @@ -909,6 +996,10 @@ def __call__( semantic_trace_dir=trace_dir, lm_num_retries=BENCHMARK_LM_NUM_RETRIES, lm_model_kwargs={ + **( + {"max_ctx_len": self.lm_max_ctx_len} + if self.lm_max_ctx_len is not None else {} + ), "extra_body": { "thinking": { "type": "enabled" if self.thinking_enabled else "disabled" @@ -921,6 +1012,9 @@ def __call__( 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, + prompt_batching=self.prompt_batching, + sem_agg_dispatch=self.sem_agg_dispatch, + structured_output_transport=self.structured_output_transport, semantic_trace_snapshot_mode=self.semantic_trace_snapshot_mode, ), pair_embedding_provider=( @@ -934,11 +1028,19 @@ def __call__( statements=storage.statements, ) memory = am.ZepMemory(adapter=adapter) - memory._runtime = MemoryRuntime( - policy, - adapter=adapter, - storage=storage, - ) + if self.refresh_every > 1: + memory._runtime = MemoryRuntime( + policy, + adapter=adapter, + storage=storage, + refresh=am.CountRefresh(every=self.refresh_every), + ) + else: + memory._runtime = MemoryRuntime( + policy, + adapter=adapter, + storage=storage, + ) return ZepMemoryDriver( memory, trace_dir=trace_dir, @@ -967,11 +1069,20 @@ def __init__( sem_groupby_pair_batch_size: int | None = None, sem_groupby_pair_batch_retries: int = 0, semantic_pair_profiles: dict[str, SemanticPairExecutionProfile] | None = None, + prompt_batching: PromptBatching | None = None, + sem_agg_dispatch: str = "sequential", + structured_output_transport: str = "chat-json-object", + lm_max_ctx_len: int | None = None, embedding_device: str = "cpu", semantic_trace_snapshot_mode: str = "compact", lotus_cache_mode: str = "disabled", + refresh_every: int = 1, thinking_enabled: bool = False, ) -> None: + validate_lm_max_ctx_len(lm_max_ctx_len) + validate_structured_output_transport( + structured_output_transport, model=model_id + ) if not base_namespace: raise ValueError("Mem0 benchmark base_namespace must be non-empty") self.base_namespace = base_namespace @@ -979,12 +1090,18 @@ def __init__( 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 {}) + self.prompt_batching = prompt_batching + self.sem_agg_dispatch = sem_agg_dispatch + self.structured_output_transport = structured_output_transport + self.lm_max_ctx_len = lm_max_ctx_len if embedding_device not in {"cpu", "cuda"}: raise ValueError("Mem0 embedding_device must be 'cpu' or 'cuda'") self.embedding_device = embedding_device self.semantic_trace_snapshot_mode = semantic_trace_snapshot_mode _validate_lotus_cache_mode(lotus_cache_mode) self.lotus_cache_mode = lotus_cache_mode + _validate_refresh_every(refresh_every) + self.refresh_every = refresh_every self.thinking_enabled = thinking_enabled self.sem_topk_method = "pairwise-naive" @@ -1058,12 +1175,14 @@ def __call__( semantic_trace_dir=trace_dir, lm_num_retries=BENCHMARK_LM_NUM_RETRIES, lm_model_kwargs={ + **( + {"max_ctx_len": self.lm_max_ctx_len} + if self.lm_max_ctx_len is not None else {} + ), "extra_body": { "thinking": { "type": ( - "enabled" - if self.thinking_enabled - else "disabled" + "enabled" if self.thinking_enabled else "disabled" ) } } @@ -1074,6 +1193,9 @@ def __call__( 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, + prompt_batching=self.prompt_batching, + sem_agg_dispatch=self.sem_agg_dispatch, + structured_output_transport=self.structured_output_transport, semantic_trace_snapshot_mode=self.semantic_trace_snapshot_mode, ), pair_embedding_provider=embedding_provider, @@ -1083,11 +1205,19 @@ def __call__( statements=storage.statements, ) memory = memory_type(adapter=adapter) - memory._runtime = MemoryRuntime( - policy, - adapter=adapter, - storage=storage, - ) + if self.refresh_every > 1: + memory._runtime = MemoryRuntime( + policy, + adapter=adapter, + storage=storage, + refresh=am.CountRefresh(every=self.refresh_every), + ) + else: + memory._runtime = MemoryRuntime( + policy, + adapter=adapter, + storage=storage, + ) except BaseException: connector.close() raise @@ -1114,9 +1244,14 @@ def __init__( sem_groupby_pair_batch_size: int | None = None, sem_groupby_pair_batch_retries: int = 0, semantic_pair_profiles: dict[str, SemanticPairExecutionProfile] | None = None, + prompt_batching: PromptBatching | None = None, + sem_agg_dispatch: str = "sequential", + structured_output_transport: str = "chat-json-object", + lm_max_ctx_len: int | None = None, embedding_device: str = "cpu", semantic_trace_snapshot_mode: str = "compact", lotus_cache_mode: str = "disabled", + refresh_every: int = 1, thinking_enabled: bool = False, ) -> None: from agent_memory.adapters.lotus.context import SEM_TOPK_METHODS @@ -1131,9 +1266,14 @@ def __init__( 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, + prompt_batching=prompt_batching, + sem_agg_dispatch=sem_agg_dispatch, + structured_output_transport=structured_output_transport, + lm_max_ctx_len=lm_max_ctx_len, embedding_device=embedding_device, semantic_trace_snapshot_mode=semantic_trace_snapshot_mode, lotus_cache_mode=lotus_cache_mode, + refresh_every=refresh_every, thinking_enabled=thinking_enabled, ) self.sem_topk_method = sem_topk_method diff --git a/src/agent_memory/evaluation/artifacts.py b/src/agent_memory/evaluation/artifacts.py index b89e0c5..5c877ca 100644 --- a/src/agent_memory/evaluation/artifacts.py +++ b/src/agent_memory/evaluation/artifacts.py @@ -198,6 +198,7 @@ def _latency_stats(values: list[float]) -> dict[str, float | None]: "mean_ms": round(mean(values), 3) if values else None, "median_ms": round(median(values), 3) if values else None, "p95_ms": _percentile(values, 0.95), + "max_ms": round(max(values), 3) if values else None, } @@ -2120,6 +2121,18 @@ def finalize_metrics(self) -> None: "insertion_median_latency_ms": _latency_stats( insertion_latencies )["median_ms"], + "insertion_excluding_semantic_trace_mean_ms": summary[ + "insertion_wall_latency_excluding_trace_io" + ]["mean_ms"], + "insertion_excluding_semantic_trace_median_ms": summary[ + "insertion_wall_latency_excluding_trace_io" + ]["median_ms"], + "insertion_excluding_semantic_trace_p95_ms": summary[ + "insertion_wall_latency_excluding_trace_io" + ]["p95_ms"], + "insertion_excluding_semantic_trace_max_ms": summary[ + "insertion_wall_latency_excluding_trace_io" + ]["max_ms"], "retrieval_mean_latency_ms": _latency_stats( retrieval_latencies )["mean_ms"], diff --git a/src/agent_memory/evaluation/claude_memory/locomo.py b/src/agent_memory/evaluation/claude_memory/locomo.py index 621dbcf..25c49a3 100644 --- a/src/agent_memory/evaluation/claude_memory/locomo.py +++ b/src/agent_memory/evaluation/claude_memory/locomo.py @@ -43,6 +43,7 @@ ) from agent_memory.evaluation.types import BenchmarkEvent, BenchmarkQuestion from agent_memory.planner import ( + DEFAULT_GROUPED_AGG_RULE, GROUPED_AGG_RULES, DifferentialRules, PolicyDifferentiator, @@ -91,7 +92,7 @@ class ClaudeMemoryLocomoRunConfig: trace: bool = False answer: bool = False resume: bool = False - grouped_agg_rule: str = "compressed" + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE sem_topk_method: str = "pairwise-naive" restore_csv_state: bool = False restore_artifact_csv_state: bool = False @@ -549,7 +550,7 @@ def checkpoint_manifest( question_limit: int, model: str, answer: bool, - grouped_agg_rule: str = "compressed", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, sem_topk_method: str = "pairwise-naive", maintenance_mode: str = "ingest", source_run_dir: str = "", @@ -604,7 +605,7 @@ def save_checkpoint( question_limit: int, model: str, answer: bool, - grouped_agg_rule: str = "compressed", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, sem_topk_method: str = "pairwise-naive", maintenance_mode: str = "ingest", source_run_dir: str = "", @@ -666,7 +667,7 @@ def load_checkpoint( question_limit: int, model: str, answer: bool, - grouped_agg_rule: str = "compressed", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, sem_topk_method: str = "pairwise-naive", maintenance_mode: str = "ingest", source_run_dir: str = "", @@ -1159,7 +1160,7 @@ def create_memory( *, model: str, trace_dir: Path | None, - grouped_agg_rule: str = "compressed", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, sem_topk_method: str = "pairwise-naive", ) -> am.ClaudeMemory: """Create the ClaudeMemory instance used by one benchmark run.""" @@ -1180,9 +1181,6 @@ def create_memory( ), ) memory = am.ClaudeMemory(adapter=adapter) - if grouped_agg_rule == "compressed": - return memory - policy = PolicyDifferentiator( rules=DifferentialRules(grouped_agg_rule=grouped_agg_rule), ).differentiate(am.ClaudeMemory.spec()) @@ -1322,7 +1320,7 @@ def summary_frame( *, run_mode: str, model: str, - grouped_agg_rule: str = "compressed", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, sem_topk_method: str = "pairwise-naive", sample_index: int, maintenance_mode: str, @@ -1432,7 +1430,7 @@ def write_run_artifacts( output_dir: Path, run_mode: str, model: str, - grouped_agg_rule: str = "compressed", + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, sem_topk_method: str = "pairwise-naive", sample_index: int, maintenance_mode: str, diff --git a/src/agent_memory/evaluation/harness.py b/src/agent_memory/evaluation/harness.py index 27cea33..8ffd862 100644 --- a/src/agent_memory/evaluation/harness.py +++ b/src/agent_memory/evaluation/harness.py @@ -210,6 +210,11 @@ def finish_session(self, session_id: str) -> Mapping[str, Any]: ... + def flush(self) -> Mapping[str, Any]: + """Publish any case-final pending source rows.""" + + ... + def save_state(self, directory: Path) -> Mapping[str, Any]: """Write opaque driver state into a new checkpoint directory.""" @@ -657,7 +662,9 @@ def _run_case( execution_attempt=attempt, unit_attempt=unit_attempt, ): - metrics = driver.add(event) + metrics = dict(driver.add(event)) + if event_index + 1 == len(case.events): + metrics.update(driver.flush()) trace_io_latency_ms = round(trace_io.latency_ms, 3) trace_bytes_written = trace_io.bytes_written insertion_latency_ms = round( diff --git a/src/agent_memory/evaluation/run.py b/src/agent_memory/evaluation/run.py index 358f4d7..a95ef2c 100644 --- a/src/agent_memory/evaluation/run.py +++ b/src/agent_memory/evaluation/run.py @@ -8,6 +8,7 @@ import os from pathlib import Path import re +from typing import Any from .agent_memory_drivers import ( ClaudeMemoryDriverFactory, @@ -25,10 +26,18 @@ semantic_pair_profiles_fingerprint, ) from agent_memory.adapters.lotus.context import ( + validate_lm_max_ctx_len, LOTUS_MEMORY_CACHE_ID, LOTUS_MEMORY_CACHE_MAX_SIZE, + SEM_AGG_DISPATCH_METHODS, SEM_JOIN_TOPK_METHODS, ) +from agent_memory.adapters.lotus.json_output import JSON_REPAIR_VERSION +from agent_memory.adapters.lotus.prompt_batching import ( + PromptBatching, + validate_structured_output_transport, +) +from agent_memory.planner import DEFAULT_GROUPED_AGG_RULE from .artifacts import BenchmarkArtifactStore from .bundle import BenchmarkBundle from .harness import BenchmarkRunner, MemorySystemContract, TaskContract @@ -67,6 +76,19 @@ } +def resolve_grouped_agg_rule( + system_id: str, + grouped_agg_rule: str | None, +) -> str: + """Resolve the grouped-aggregate default for one built-in system.""" + + if grouped_agg_rule is not None: + return grouped_agg_rule + if system_id in {"mem0-memory", "mem0-enhanced"}: + return "rule-all-group" + return DEFAULT_GROUPED_AGG_RULE + + def _namespace(benchmark_id: str, output_dir: Path) -> str: label = re.sub(r"[^a-z0-9]+", "-", benchmark_id.lower()).strip("-") digest = sha256(str(output_dir.resolve()).encode("utf-8")).hexdigest()[:12] @@ -98,11 +120,15 @@ def run_agent_memory_bundle( answer_model_id: str = DEFAULT_PROVIDER_MODEL_ID, judge_model_id: str = DEFAULT_PROVIDER_MODEL_ID, base_namespace: str | None = None, - grouped_agg_rule: str = "rule-all-group", + grouped_agg_rule: str | None = None, 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, + sem_agg_dispatch: str = "sequential", + prompt_batching: PromptBatching | None = None, + structured_output_transport: str = "chat-json-object", + lm_max_ctx_len: int | None = None, semantic_pair_profile: str = "oracle-only", semantic_pair_top_k: int | None = None, semantic_pair_min_similarity: float | None = None, @@ -110,6 +136,7 @@ def run_agent_memory_bundle( lotus_cache_mode: str = "disabled", embedding_device: str = "cpu", semantic_trace_snapshot_mode: str = "compact", + refresh_every: int = 1, memory_thinking_enabled: bool = True, condition_id: str = "", maintenance_only: bool = False, @@ -120,10 +147,39 @@ def run_agent_memory_bundle( if system_id not in AGENT_MEMORY_SYSTEMS: raise ValueError(f"unsupported agent-memory benchmark system {system_id!r}") + validate_lm_max_ctx_len(lm_max_ctx_len) + validate_structured_output_transport( + structured_output_transport, model=memory_provider_model_id + ) + grouped_agg_rule = resolve_grouped_agg_rule(system_id, grouped_agg_rule) + if ( + isinstance(refresh_every, bool) + or not isinstance(refresh_every, int) + or refresh_every < 1 + ): + raise ValueError("refresh_every must be a positive integer") if sem_groupby_pair_batch_size is not None and sem_groupby_pair_batch_size < 1: 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_agg_dispatch not in SEM_AGG_DISPATCH_METHODS: + raise ValueError( + "sem_agg_dispatch must be one of: " + + ", ".join(SEM_AGG_DISPATCH_METHODS) + ) + if prompt_batching is not None and not isinstance( + prompt_batching, PromptBatching + ): + raise TypeError("prompt_batching must be PromptBatching or None") + if prompt_batching is not None and sem_agg_dispatch != "sequential": + raise ValueError( + "prompt_batching cannot be combined with sem_agg_dispatch" + ) + if prompt_batching is not None and sem_groupby_pair_batch_size is not None: + raise ValueError( + "prompt_batching cannot be combined with " + "sem_groupby_pair_batch_size" + ) if ( sem_join_topk_method is not None and sem_join_topk_method not in SEM_JOIN_TOPK_METHODS @@ -325,6 +381,23 @@ def run_agent_memory_bundle( if uses_sem_join_topk else "" ), + ( + f"sem-agg-dispatch:{sem_agg_dispatch}" + if sem_agg_dispatch != "sequential" + else "" + ), + ( + f"prompt-batching:{prompt_batching.fingerprint}" + if prompt_batching is not None + else "" + ), + ( + f"structured-output-transport:{structured_output_transport}" + if structured_output_transport != "chat-json-object" + else "" + ), + f"lm-max-ctx-len:{lm_max_ctx_len}" if lm_max_ctx_len is not None else "", + f"refresh:count:{refresh_every}" if refresh_every > 1 else "", f"lotus-cache:{LOTUS_MEMORY_CACHE_ID}" if lotus_cache_mode == "memory" else "", @@ -367,6 +440,11 @@ def run_agent_memory_bundle( ), ), ) + if refresh_every > 1: + runtime_provenance["runtime"]["refresh"] = { + "type": "count", + "every": refresh_every, + } lotus_execution_provenance = { "sem_groupby_pair_batch_size": sem_groupby_pair_batch_size, "sem_groupby_pair_batch_retries": sem_groupby_pair_batch_retries, @@ -386,6 +464,23 @@ def run_agent_memory_bundle( for digest, profile in sorted(semantic_pair_profiles.items()) }, } + if sem_agg_dispatch != "sequential": + lotus_execution_provenance.update( + { + "sem_agg_dispatch": sem_agg_dispatch, + } + ) + if prompt_batching is not None: + lotus_execution_provenance["prompt_batching"] = ( + prompt_batching.to_dict() + ) + lotus_execution_provenance["json_repair_version"] = JSON_REPAIR_VERSION + if structured_output_transport != "chat-json-object": + lotus_execution_provenance["structured_output_transport"] = ( + structured_output_transport + ) + if lm_max_ctx_len is not None: + lotus_execution_provenance["lm_max_ctx_len"] = lm_max_ctx_len if site_profile_config is not None: lotus_execution_provenance.update( { @@ -458,6 +553,29 @@ def run_agent_memory_bundle( framework_cache_mode=framework_cache_mode, checkpoint_enabled=True, ) + sem_agg_execution_options = ( + { + "sem_agg_dispatch": sem_agg_dispatch, + } + if sem_agg_dispatch != "sequential" + else {} + ) + prompt_batching_options = ( + {"prompt_batching": prompt_batching} + if prompt_batching is not None + else {} + ) + structured_transport_options = ( + {"structured_output_transport": structured_output_transport} + if structured_output_transport != "chat-json-object" + else {} + ) + lm_context_options = ( + {"lm_max_ctx_len": lm_max_ctx_len} if lm_max_ctx_len is not None else {} + ) + refresh_execution_options: dict[str, Any] = {} + if refresh_every > 1: + refresh_execution_options["refresh_every"] = refresh_every if system_id == "claude-memory": driver_factory = ClaudeMemoryDriverFactory( model_id=memory_provider_model_id, @@ -469,10 +587,16 @@ def run_agent_memory_bundle( semantic_trace_snapshot_mode=semantic_trace_snapshot_mode, lotus_cache_mode=lotus_cache_mode, thinking_enabled=memory_thinking_enabled, + **sem_agg_execution_options, + **prompt_batching_options, + **structured_transport_options, + **lm_context_options, + **refresh_execution_options, ) elif system_id == "zep-memory": driver_factory = ZepMemoryDriverFactory.from_environment( - base_namespace=base_namespace or _namespace(bundle.benchmark_id, output_dir), + 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, @@ -483,10 +607,16 @@ def run_agent_memory_bundle( semantic_trace_snapshot_mode=semantic_trace_snapshot_mode, lotus_cache_mode=lotus_cache_mode, thinking_enabled=memory_thinking_enabled, + **sem_agg_execution_options, + **prompt_batching_options, + **structured_transport_options, + **lm_context_options, + **refresh_execution_options, ) elif system_id == "mem0-memory": driver_factory = Mem0MemoryDriverFactory( - base_namespace=base_namespace or _namespace(bundle.benchmark_id, output_dir), + base_namespace=base_namespace + or _namespace(bundle.benchmark_id, output_dir), model_id=memory_provider_model_id, sem_groupby_pair_batch_size=sem_groupby_pair_batch_size, sem_groupby_pair_batch_retries=sem_groupby_pair_batch_retries, @@ -495,10 +625,16 @@ def run_agent_memory_bundle( semantic_trace_snapshot_mode=semantic_trace_snapshot_mode, lotus_cache_mode=lotus_cache_mode, thinking_enabled=memory_thinking_enabled, + **sem_agg_execution_options, + **prompt_batching_options, + **structured_transport_options, + **lm_context_options, + **refresh_execution_options, ) else: driver_factory = Mem0MemoryEnhancedDriverFactory( - base_namespace=base_namespace or _namespace(bundle.benchmark_id, output_dir), + base_namespace=base_namespace + or _namespace(bundle.benchmark_id, output_dir), model_id=memory_provider_model_id, sem_topk_method=sem_topk_method, sem_groupby_pair_batch_size=sem_groupby_pair_batch_size, @@ -508,6 +644,11 @@ def run_agent_memory_bundle( semantic_trace_snapshot_mode=semantic_trace_snapshot_mode, lotus_cache_mode=lotus_cache_mode, thinking_enabled=memory_thinking_enabled, + **sem_agg_execution_options, + **prompt_batching_options, + **structured_transport_options, + **lm_context_options, + **refresh_execution_options, ) try: storage_provenance = ( diff --git a/src/agent_memory/evaluation/semantic_pair_config.py b/src/agent_memory/evaluation/semantic_pair_config.py index af8c68e..083b002 100644 --- a/src/agent_memory/evaluation/semantic_pair_config.py +++ b/src/agent_memory/evaluation/semantic_pair_config.py @@ -12,8 +12,6 @@ from agent_memory.adapters.lotus.pair_execution import ( SEMANTIC_PAIR_EXECUTION_MODES, ) - - @dataclass(frozen=True) class SemanticPairSiteBinding: """One user-selected physical profile for a semantic predicate site.""" @@ -55,7 +53,6 @@ def __post_init__(self) -> None: raise ValueError("search-filter site bindings require a candidate bound") if self.mode == "proxy-only" and self.min_similarity is None: raise ValueError("proxy-only site bindings require min_similarity") - def to_dict(self) -> dict[str, object]: """Return the canonical manifest representation.""" @@ -73,12 +70,13 @@ class SemanticPairProfileConfig: bindings: tuple[SemanticPairSiteBinding, ...] source_sha256: str + schema_version: int = 1 def to_dict(self) -> dict[str, object]: """Return normalized provenance without the machine-local path.""" return { - "schema_version": 1, + "schema_version": self.schema_version, "source_sha256": self.source_sha256, "bindings": [binding.to_dict() for binding in self.bindings], } @@ -100,7 +98,8 @@ def load_semantic_pair_profile_config(path: Path) -> SemanticPairProfileConfig: if not isinstance(payload, dict): raise TypeError("semantic pair profile config must be a JSON object") _require_exact_keys(payload, {"schema_version", "bindings"}, "config") - if payload["schema_version"] != 1: + schema_version = payload["schema_version"] + if schema_version != 1: raise ValueError("semantic pair profile config schema_version must be 1") raw_bindings = payload["bindings"] if not isinstance(raw_bindings, list) or not raw_bindings: @@ -111,11 +110,8 @@ def load_semantic_pair_profile_config(path: Path) -> SemanticPairProfileConfig: for index, item in enumerate(raw_bindings): if not isinstance(item, dict): raise TypeError(f"semantic pair binding {index} must be an object") - _require_exact_keys( - item, - {"site_id", "mode", "top_k", "min_similarity"}, - f"binding {index}", - ) + expected_keys = {"site_id", "mode", "top_k", "min_similarity"} + _require_exact_keys(item, expected_keys, f"binding {index}") site_id = item["site_id"] mode = item["mode"] if not isinstance(site_id, str) or not isinstance(mode, str): @@ -134,6 +130,7 @@ def load_semantic_pair_profile_config(path: Path) -> SemanticPairProfileConfig: return SemanticPairProfileConfig( bindings=tuple(bindings), source_sha256=sha256(content).hexdigest(), + schema_version=schema_version, ) diff --git a/src/agent_memory/evaluation/trace_metrics.py b/src/agent_memory/evaluation/trace_metrics.py index 4bfca51..ea0a877 100644 --- a/src/agent_memory/evaluation/trace_metrics.py +++ b/src/agent_memory/evaluation/trace_metrics.py @@ -27,6 +27,7 @@ def normalize_provider_calls( *, output_dir: Path, pricing: PricingSnapshot | None = None, + include_cost: bool = True, ) -> list[dict[str, Any]]: """Return one normalized row per provider response. @@ -34,9 +35,11 @@ def normalize_provider_calls( provider event is authoritative when both exist, avoiding double counting. Responses retain their framework batch identity so reports can distinguish logical operator calls, framework batches, and provider calls. + With include_cost=False, both cost fields remain unknown; no default price + is applied. Existing callers retain the historical pricing behavior. """ - pricing = pricing or PricingSnapshot.deepseek_2026_07_17() + pricing = (pricing or PricingSnapshot.deepseek_2026_07_17()) if include_cost else None pending_semantic_calls: dict[ tuple[str, int], deque[Mapping[str, Any]] ] = defaultdict(deque) @@ -191,7 +194,7 @@ def summarize_framework_cache_usage( def _agent_provider_row( event: Mapping[str, Any], output_dir: Path, - pricing: PricingSnapshot, + pricing: PricingSnapshot | None, *, semantic_event: Mapping[str, Any] | None, ) -> dict[str, Any]: @@ -223,7 +226,7 @@ def _agent_provider_row( def _agent_runner_row( event: Mapping[str, Any], - pricing: PricingSnapshot, + pricing: PricingSnapshot | None, ) -> dict[str, Any]: error = event.get("event_type") == "llm_batch_error" return _row( @@ -249,7 +252,7 @@ def _agent_runner_row( def _native_graphiti_row( event: Mapping[str, Any], - pricing: PricingSnapshot, + pricing: PricingSnapshot | None, ) -> dict[str, Any]: usage = _mapping(event.get("usage")) provider_cache = _mapping(event.get("provider_cache")) @@ -272,7 +275,7 @@ def _native_graphiti_row( def _native_claude_row( event: Mapping[str, Any], - pricing: PricingSnapshot, + pricing: PricingSnapshot | None, ) -> dict[str, Any]: usage = _mapping(event.get("usage")) error = event.get("event_type") == "llm_call_error" @@ -309,10 +312,10 @@ def _row( completion_tokens: int | None, reasoning_tokens: int | None, usage_available: bool, - pricing: PricingSnapshot, + pricing: PricingSnapshot | None, ) -> dict[str, Any]: total = _sum_optional(prompt_tokens, completion_tokens) - cost = pricing.estimate_cost_usd( + cost = None if pricing is None else pricing.estimate_cost_usd( cache_hit_input_tokens=cache_hit_tokens, cache_miss_input_tokens=cache_miss_tokens, output_tokens=completion_tokens, @@ -362,7 +365,7 @@ def _row( "total_tokens": total, "usage_available": usage_available, "known_cost_usd": ( - float(round(cost, 12)) if cost is not None else 0.0 + None if pricing is None else float(round(cost, 12)) if cost is not None else 0.0 ), "estimated_cost_usd": ( float(round(cost, 12)) if cost is not None else None diff --git a/src/agent_memory/evaluation/zep/locomo.py b/src/agent_memory/evaluation/zep/locomo.py index 5eafd62..97aba57 100644 --- a/src/agent_memory/evaluation/zep/locomo.py +++ b/src/agent_memory/evaluation/zep/locomo.py @@ -37,6 +37,7 @@ ) from agent_memory.evaluation.zep.artifacts import ArtifactStore, PricingSnapshot from agent_memory.policy.retrieval import RetrievalResult +from agent_memory.planner import DEFAULT_GROUPED_AGG_RULE from agent_memory.tracing.semantic import semantic_trace_scope _LOCOMO_TIMESTAMP_PATTERN = re.compile( @@ -105,7 +106,7 @@ class ZepLocomoRunConfig: question_numbers: tuple[int, ...] | None = None include_adversarial: bool = True model: str = DEFAULT_MODEL - grouped_agg_rule: str = "rule-re-group" + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE namespace: str | None = None diff --git a/src/agent_memory/memories/claude.py b/src/agent_memory/memories/claude.py index 6828b92..0169676 100644 --- a/src/agent_memory/memories/claude.py +++ b/src/agent_memory/memories/claude.py @@ -177,9 +177,13 @@ class ClaudeMemory(Memory): "session_id", ] ) + # Membership is intentionally unspecified, so rule-join-map preserves the + # current multi-target matching behavior for Claude topics. .sem_groupby( input_cols=["name", "description", "type"], instruction=_GROUP_INSTRUCTION, + # Future single-target experiment: + # membership="exclusive", ) .sem_agg( input_cols=[ diff --git a/src/agent_memory/planner/__init__.py b/src/agent_memory/planner/__init__.py index 0020844..6cff3d6 100644 --- a/src/agent_memory/planner/__init__.py +++ b/src/agent_memory/planner/__init__.py @@ -8,12 +8,14 @@ from .differential_query import QueryDifferentiator from .retrieval import RetrievalNode, RetrievalPlan, RetrievalPlanner from .rules import ( + DEFAULT_GROUPED_AGG_RULE, GROUPED_AGG_RULES, DifferentialInstructionRewriter, DifferentialRules, ) __all__ = [ + "DEFAULT_GROUPED_AGG_RULE", "DifferentialInstructionRewriter", "DifferentialNode", "DifferentialRules", diff --git a/src/agent_memory/planner/algebraic_aggregates.py b/src/agent_memory/planner/algebraic_aggregates.py new file mode 100644 index 0000000..21ad9fb --- /dev/null +++ b/src/agent_memory/planner/algebraic_aggregates.py @@ -0,0 +1,139 @@ +"""Lower exact relational aggregates into hidden state and public finalization.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from agent_memory.policy.aggregates import ( + AvgAggregateSpec, + SumAggregateSpec, + is_algebraic_aggregate_spec, +) +from agent_memory.policy.logical import QueryExpr + +_HIDDEN_PREFIX = "__am_aggregate_" + + +def is_algebraic_aggregate_query(query: QueryExpr) -> bool: + """Return whether a logical query is an exact count/sum/avg aggregate.""" + + aggregates = tuple(query.params.get("aggregates", ())) + return ( + query.op == "agg" + and bool(aggregates) + and all(is_algebraic_aggregate_spec(spec) for spec in aggregates) + ) + + +def aggregate_layout(query: QueryExpr) -> dict[str, Any]: + """Return a stable hidden-state layout for one logical aggregate query.""" + + if not is_algebraic_aggregate_query(query): + raise TypeError("aggregate lowering requires only count/sum/avg specs") + aggregate_input = query.inputs[0] + if aggregate_input.op == "sem_groupby": + raise NotImplementedError( + "sem_groupby numeric algebraic aggregates are not supported" + ) + group_keys = ( + tuple(str(key) for key in aggregate_input.params["keys"]) + if aggregate_input.op == "group_by" + else () + ) + aggregates = tuple(query.params["aggregates"]) + occupied = set(group_keys) + row_count_col = _unique_hidden_name("row_count", occupied) + occupied.add(row_count_col) + + numeric_columns: list[str] = [] + for aggregate in aggregates: + if isinstance(aggregate, (SumAggregateSpec, AvgAggregateSpec)): + if aggregate.column not in numeric_columns: + numeric_columns.append(aggregate.column) + numeric_states: list[dict[str, str]] = [] + for index, column in enumerate(numeric_columns): + sum_col = _unique_hidden_name(f"sum_{index}", occupied) + occupied.add(sum_col) + non_null_count_col = _unique_hidden_name( + f"non_null_count_{index}", occupied + ) + occupied.add(non_null_count_col) + numeric_states.append( + { + "input_col": column, + "sum_col": sum_col, + "non_null_count_col": non_null_count_col, + } + ) + state_columns = ( + *group_keys, + row_count_col, + *( + state_column + for numeric_state in numeric_states + for state_column in ( + numeric_state["sum_col"], + numeric_state["non_null_count_col"], + ) + ), + ) + return { + "group_keys": group_keys, + "aggregates": aggregates, + "row_count_col": row_count_col, + "numeric_states": tuple(numeric_states), + "state_columns": state_columns, + } + + +def aggregate_state_query( + source: QueryExpr, + *, + layout: Mapping[str, Any], +) -> QueryExpr: + """Build the complete hidden-state query used as a reference computation.""" + + return QueryExpr(op="aggregate_state", inputs=(source,), params=layout) + + +def aggregate_state_update_query( + *, + current_state: QueryExpr, + inserted_rows: QueryExpr, + retracted_rows: QueryExpr, + layout: Mapping[str, Any], +) -> QueryExpr: + """Build exact signed parent-delta maintenance for hidden aggregate state.""" + + return QueryExpr( + op="aggregate_state_update", + inputs=(current_state, inserted_rows, retracted_rows), + params=layout, + ) + + +def aggregate_finalize_query( + state: QueryExpr, + *, + layout: Mapping[str, Any], +) -> QueryExpr: + """Build the deterministic public projection from hidden aggregate state.""" + + return QueryExpr(op="aggregate_finalize", inputs=(state,), params=layout) + + +def _unique_hidden_name(suffix: str, occupied: set[str]) -> str: + candidate = f"{_HIDDEN_PREFIX}{suffix}" + while candidate in occupied: + candidate = f"_{candidate}" + return candidate + + +__all__ = [ + "aggregate_finalize_query", + "aggregate_layout", + "aggregate_state_query", + "aggregate_state_update_query", + "is_algebraic_aggregate_query", +] diff --git a/src/agent_memory/planner/differential_policy.py b/src/agent_memory/planner/differential_policy.py index 95cbc1c..6175b67 100644 --- a/src/agent_memory/planner/differential_policy.py +++ b/src/agent_memory/planner/differential_policy.py @@ -12,6 +12,13 @@ from agent_memory.policy.retrieval import RetrievalQuery from agent_memory.policy.schema import output_columns from agent_memory.planner.differential_query import QueryDifferentiator +from agent_memory.planner.algebraic_aggregates import ( + aggregate_finalize_query, + aggregate_layout, + aggregate_state_query, + aggregate_state_update_query, + is_algebraic_aggregate_query, +) from agent_memory.planner.retrieval import RetrievalPlan, RetrievalPlanner from agent_memory.planner.serialization import stable_json, stable_value from agent_memory.planner.rules import ( @@ -201,10 +208,6 @@ def _compile_query(self, query: QueryExpr, *, spec: MemorySpec) -> str: self._query_nodes[query] = node_id return node_id - if query.op == "sem_join": - raise NotImplementedError( - "Direct view-time sem_join is not supported by policy differentiation" - ) if query.op == "sem_topk": raise NotImplementedError( "View-time sem_topk is not supported by policy differentiation" @@ -230,6 +233,9 @@ def _compile_query(self, query: QueryExpr, *, spec: MemorySpec) -> str: if query.op == "process_window": return self._compile_process_window(query, spec=spec) + if is_algebraic_aggregate_query(query): + return self._compile_algebraic_aggregate(query, spec=spec) + if self._is_grouped_aggregate(query): return self._compile_grouped_aggregate(query, spec=spec) @@ -246,10 +252,28 @@ def _compile_query(self, query: QueryExpr, *, spec: MemorySpec) -> str: 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( + maintenance_query = self._append_inner_join_maintenance_query( + local_query, + input_node_ids, + current_node_id=self._node_id(query), + operator="join", + ) + elif query.op == "sem_join": + how = str(query.params.get("how", "inner")).lower() + if how != "inner": + raise NotImplementedError( + "Direct sem_join differential currently supports only how='inner'" + ) + if query.params.get("k") is not None: + raise NotImplementedError( + "Direct sem_join differential currently requires k=None" + ) + execution_kind = "semantic_binary_state" + maintenance_query = self._append_inner_join_maintenance_query( local_query, input_node_ids, current_node_id=self._node_id(query), + operator="sem_join", ) elif query.op == "sem_agg": execution_kind = "semantic_state" @@ -267,6 +291,65 @@ def _compile_query(self, query: QueryExpr, *, spec: MemorySpec) -> str: ) return self._register(query, node) + def _compile_algebraic_aggregate( + self, + query: QueryExpr, + *, + spec: MemorySpec, + ) -> str: + """Compile count/sum/avg into hidden algebraic state and a finalizer.""" + + aggregate_input = query.inputs[0] + source_query = ( + aggregate_input.inputs[0] + if aggregate_input.op == "group_by" + else aggregate_input + ) + source_node_id = self._compile_query(source_query, spec=spec) + source_leaf = self._node_leaf(source_node_id) + layout = aggregate_layout(query) + state_query = aggregate_state_query(source_leaf, layout=layout) + state_node_id = self._node_id(state_query) + state_columns = tuple(str(column) for column in layout["state_columns"]) + current_state = QueryExpr( + op="materialized_view", + params={"name": state_node_id, "columns": state_columns}, + ) + inserted_rows = self._inserted_node_leaf(source_node_id) + retracted_rows = QueryExpr( + op="materialized_view", + params={ + "name": f"{source_node_id}__retracted", + "columns": source_leaf.params["columns"], + }, + ) + state_node = DifferentialNode( + node_id=state_node_id, + query=state_query, + input_node_ids=(source_node_id,), + execution_kind="algebraic_state", + output_columns=state_columns, + maintenance_query=aggregate_state_update_query( + current_state=current_state, + inserted_rows=inserted_rows, + retracted_rows=retracted_rows, + layout=layout, + ), + ) + self._register(state_query, state_node) + + final_query = aggregate_finalize_query( + self._node_leaf(state_node_id), + layout=layout, + ) + public_node = self._make_node( + query=query, + local_query=final_query, + input_node_ids=(state_node_id,), + execution_kind="deterministic", + ) + return self._register(query, public_node) + def _compile_grouped_aggregate(self, query: QueryExpr, *, spec: MemorySpec) -> str: """Compile one aggregate together with its grouping carrier.""" @@ -396,19 +479,22 @@ def _semantic_maintenance_query( source_input=changed_source, ) - def _relational_inner_join_maintenance_query( + def _append_inner_join_maintenance_query( self, local_query: QueryExpr, input_node_ids: tuple[str, ...], *, current_node_id: str, + operator: str, ) -> QueryExpr: - """Build bag-preserving append maintenance for one relational inner join.""" + """Build bag-preserving append maintenance for one inner join node.""" - if local_query.op != "join" or str(local_query.params.get("how", "inner")) != "inner": - raise ValueError("Relational join maintenance requires how='inner'") + if local_query.op != operator: + raise ValueError(f"{operator} maintenance requires a {operator} query") + if str(local_query.params.get("how", "inner")).lower() != "inner": + raise ValueError(f"{operator} 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") + raise ValueError(f"{operator} inner join requires exactly two inputs") left_node_id, right_node_id = input_node_ids left_old = self._node_leaf(left_node_id) @@ -418,7 +504,7 @@ def _relational_inner_join_maintenance_query( def join(left: QueryExpr, right: QueryExpr) -> QueryExpr: return QueryExpr( - op="join", + op=operator, inputs=( self._preserve_input_alias(local_query.inputs[0], left), self._preserve_input_alias(local_query.inputs[1], right), diff --git a/src/agent_memory/planner/rules.py b/src/agent_memory/planner/rules.py index 3c5178d..2e65dbb 100644 --- a/src/agent_memory/planner/rules.py +++ b/src/agent_memory/planner/rules.py @@ -10,6 +10,7 @@ CollectListAggregateSpec, MinAggregateSpec, SemanticAggregateSpec, + aggregate_output_names, ) from agent_memory.policy.expressions import ColumnExpr, least from agent_memory.policy.logical import ColumnSpec, QueryExpr @@ -43,6 +44,7 @@ "rule-all-group", "rule-all-group-optimized", ) +DEFAULT_GROUPED_AGG_RULE = "rule-join-map" PLACEHOLDER_PATTERN = re.compile( r"(? str: class DifferentialRules: """Operator-level and pattern-level differential rules.""" - def __init__(self, *, grouped_agg_rule: str = "compressed") -> None: + def __init__( + self, + *, + grouped_agg_rule: str = DEFAULT_GROUPED_AGG_RULE, + ) -> None: """Configure differential rule strategy experiments.""" if grouped_agg_rule not in GROUPED_AGG_RULES: @@ -662,7 +668,6 @@ def _build_grouped_agg_re_group_candidate( merged = self._build_aggregate_remerge( grouped=regrouped, aggregate=aggregate, - state_cols=final_columns, instruction_rewriter=instruction_rewriter, ) return self._select_after_array_remerge_flatten(merged, aggregate, final_columns) @@ -672,7 +677,6 @@ def _build_aggregate_remerge( *, grouped: QueryExpr, aggregate: QueryExpr, - state_cols: Sequence[str], instruction_rewriter: DifferentialInstructionRewriter, ) -> QueryExpr: """Build grouped aggregate remerge using ordinary aggregate specs.""" @@ -683,7 +687,6 @@ def _build_aggregate_remerge( params={ "aggregates": self._remerge_aggregate_specs( aggregate, - state_cols=state_cols, instruction_rewriter=instruction_rewriter, ) }, @@ -693,7 +696,6 @@ def _remerge_aggregate_specs( self, aggregate: QueryExpr, *, - state_cols: Sequence[str], instruction_rewriter: DifferentialInstructionRewriter, ) -> tuple[object, ...]: """Return aggregate specs for remerging grouped aggregate rows.""" @@ -701,9 +703,10 @@ def _remerge_aggregate_specs( specs: list[object] = [] for spec in aggregate.params.get("aggregates", ()): if isinstance(spec, SemanticAggregateSpec): + state_cols = aggregate_output_names(spec) specs.append( SemanticAggregateSpec( - input_cols=None, + input_cols=state_cols, output_cols=spec.output_cols, instruction=instruction_rewriter.state_reaggregation( spec.instruction, @@ -1063,7 +1066,6 @@ def _reaggregate_join_map_states( 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( @@ -1075,10 +1077,11 @@ def _reaggregate_join_map_states( raise TypeError(f"Unsupported join-map aggregate: {aggregate.op}") reaggregate_params = dict(aggregate.params) - reaggregate_params["input_cols"] = None + state_cols = self._output_column_names(aggregate.params.get("output_cols")) + reaggregate_params["input_cols"] = state_cols reaggregate_params["instruction"] = instruction_rewriter.state_reaggregation( str(aggregate.params["instruction"]), - state_cols=final_columns, + state_cols=state_cols, raw_input_cols=tuple( str(column) for column in (aggregate.params.get("input_cols") or ()) ), @@ -1217,10 +1220,11 @@ def _build_grouped_sem_agg_re_group_candidate( params=groupby.params, ) reaggregate_params = dict(aggregate.params) - reaggregate_params["input_cols"] = None + state_cols = self._output_column_names(aggregate.params.get("output_cols")) + reaggregate_params["input_cols"] = state_cols reaggregate_params["instruction"] = instruction_rewriter.state_reaggregation( str(aggregate.params["instruction"]), - state_cols=final_columns, + state_cols=state_cols, raw_input_cols=tuple( str(column) for column in (aggregate.params.get("input_cols") or ()) ), diff --git a/src/agent_memory/policy/__init__.py b/src/agent_memory/policy/__init__.py index b0d5cd4..3c8d588 100644 --- a/src/agent_memory/policy/__init__.py +++ b/src/agent_memory/policy/__init__.py @@ -1,7 +1,7 @@ """Logical policy authoring objects and immutable query descriptors.""" -from .aggregates import array_agg, collect_list, min, sem_agg -from .expressions import least +from .aggregates import array_agg, avg, collect_list, count, min, sem_agg, sum +from .expressions import case_when, least, try_cast from .logical import ColumnSpec, MemorySpec, MemoryView, QueryExpr, UserQuery from .relation import Log, Relation, SearchRelation from .retrieval import ( @@ -26,10 +26,15 @@ "SearchRelation", "UserQuery", "array_agg", + "avg", + "case_when", "collect_list", + "count", "least", "min", "sem_agg", + "sum", + "try_cast", "BFS", "BM25", "CosineSimilarity", diff --git a/src/agent_memory/policy/aggregates.py b/src/agent_memory/policy/aggregates.py index 14b78ca..357dc85 100644 --- a/src/agent_memory/policy/aggregates.py +++ b/src/agent_memory/policy/aggregates.py @@ -42,11 +42,38 @@ class MinAggregateSpec: output_col: str +@dataclass(frozen=True) +class CountAggregateSpec: + """Serializable descriptor for relational ``COUNT(*)``.""" + + output_col: str + + +@dataclass(frozen=True) +class SumAggregateSpec: + """Serializable descriptor for one nullable numeric sum.""" + + column: str + output_col: str + + +@dataclass(frozen=True) +class AvgAggregateSpec: + """Serializable descriptor for one nullable numeric average.""" + + column: str + output_col: str + + +AlgebraicAggregateSpec: TypeAlias = ( + CountAggregateSpec | SumAggregateSpec | AvgAggregateSpec +) AggregateSpec: TypeAlias = ( SemanticAggregateSpec | ArrayAggregateSpec | CollectListAggregateSpec | MinAggregateSpec + | AlgebraicAggregateSpec ) ColumnOutput: TypeAlias = Sequence[str] | Mapping[str, str] @@ -110,6 +137,30 @@ def min( return MinAggregateSpec(columns=normalized_columns, output_col=str(output_col)) +def count(*, output_col: str) -> CountAggregateSpec: + """Declare relational ``COUNT(*)`` for global or ordinary grouped ``agg``.""" + + return CountAggregateSpec(output_col=_require_column_name(output_col, field="output_col")) + + +def sum(*, column: str, output_col: str) -> SumAggregateSpec: + """Declare a nullable numeric sum for global or ordinary grouped ``agg``.""" + + return SumAggregateSpec( + column=_require_column_name(column, field="column"), + output_col=_require_column_name(output_col, field="output_col"), + ) + + +def avg(*, column: str, output_col: str) -> AvgAggregateSpec: + """Declare a nullable numeric average for global or ordinary grouped ``agg``.""" + + return AvgAggregateSpec( + column=_require_column_name(column, field="column"), + output_col=_require_column_name(output_col, field="output_col"), + ) + + def normalize_min_columns( *, column: str | None, @@ -136,7 +187,17 @@ def aggregate_output_names(spec: AggregateSpec) -> tuple[str, ...]: if isinstance(spec, SemanticAggregateSpec): return tuple(column.name for column in spec.output_cols) - if isinstance(spec, (ArrayAggregateSpec, CollectListAggregateSpec, MinAggregateSpec)): + if isinstance( + spec, + ( + ArrayAggregateSpec, + CollectListAggregateSpec, + MinAggregateSpec, + CountAggregateSpec, + SumAggregateSpec, + AvgAggregateSpec, + ), + ): return (spec.output_col,) raise TypeError(f"Unsupported aggregate spec: {type(spec).__name__}") @@ -157,13 +218,18 @@ def normalize_aggregate_specs(specs: Sequence[AggregateSpec]) -> tuple[Aggregate ArrayAggregateSpec, CollectListAggregateSpec, MinAggregateSpec, + CountAggregateSpec, + SumAggregateSpec, + AvgAggregateSpec, ), ) ] if invalid: raise TypeError( "agg accepts only agent_memory.sem_agg(...), agent_memory.array_agg(...), " - "agent_memory.collect_list(...), or agent_memory.min(...) " + "agent_memory.collect_list(...), agent_memory.min(...), " + "agent_memory.count(...), agent_memory.sum(...), or " + "agent_memory.avg(...) " f"aggregate specs; got {invalid}" ) outputs = [ @@ -177,6 +243,22 @@ def normalize_aggregate_specs(specs: Sequence[AggregateSpec]) -> tuple[Aggregate return normalized +def is_algebraic_aggregate_spec(spec: object) -> bool: + """Return whether a descriptor has exact additive aggregate state.""" + + return isinstance(spec, (CountAggregateSpec, SumAggregateSpec, AvgAggregateSpec)) + + +def _require_column_name(value: object, *, field: str) -> str: + """Validate one aggregate input or output column name.""" + + if not isinstance(value, str): + raise TypeError(f"aggregate {field} must be a string") + if not value: + raise ValueError(f"aggregate {field} cannot be empty") + return value + + def _normalize_output_cols(output_cols: ColumnOutput) -> tuple[ColumnSpec, ...]: """Normalize output column declarations into ColumnSpec tuples.""" diff --git a/src/agent_memory/policy/expressions.py b/src/agent_memory/policy/expressions.py index e20246c..a54a4ae 100644 --- a/src/agent_memory/policy/expressions.py +++ b/src/agent_memory/policy/expressions.py @@ -31,6 +31,28 @@ def least(*values: object) -> "LeastExpr": return LeastExpr(operands=tuple(ensure_expr(value) for value in values)) +def try_cast(value: object, *, to: str) -> "TryCastExpr": + """Build a nullable scalar conversion expression.""" + + if to != "float": + raise ValueError("try_cast currently only supports the 'float' target") + return TryCastExpr(value=ensure_expr(value), target=to) + + +def case_when( + condition: object, + then_value: object, + else_value: object, +) -> "CaseWhenExpr": + """Build a SQL-style conditional scalar expression.""" + + return CaseWhenExpr( + condition=ensure_boolean_expr(condition), + then_value=ensure_expr(then_value), + else_value=ensure_expr(else_value), + ) + + def ensure_expr(value: object) -> "Expr": """Normalize supported public expression values into expression objects.""" @@ -41,12 +63,16 @@ def ensure_expr(value: object) -> "Expr": raise TypeError(f"Expected a relational expression or scalar literal, got {type(value).__name__}") -def ensure_boolean_expr(value: object) -> "BooleanExpr | ColumnExpr | ComparisonExpr": +def ensure_boolean_expr( + value: object, +) -> "BooleanExpr | ColumnExpr | ComparisonExpr | LiteralExpr": """Normalize and require a boolean-valued expression.""" expr = ensure_expr(value) if isinstance(expr, (BooleanExpr, ColumnExpr, ComparisonExpr)): return expr + if isinstance(expr, LiteralExpr) and isinstance(expr.value, bool): + return expr raise TypeError(f"Expected a boolean relational expression, got {type(expr).__name__}") @@ -86,6 +112,29 @@ def expr_from_param(value: object) -> "Expr": left=expr_from_param(value.get("left")), right=expr_from_param(value.get("right")), ) + if kind == "arithmetic": + op = value.get("op") + if op not in {"add", "subtract", "multiply", "divide"}: + raise ValueError(f"Unsupported arithmetic expression op: {op!r}") + return ArithmeticExpr( + op=op, # type: ignore[arg-type] + left=expr_from_param(value.get("left")), + right=expr_from_param(value.get("right")), + ) + if kind == "try_cast": + target = value.get("target") + if target != "float": + raise ValueError("try_cast currently only supports the 'float' target") + return TryCastExpr( + value=expr_from_param(value.get("value")), + target=target, + ) + if kind == "case_when": + return CaseWhenExpr( + condition=ensure_boolean_expr(expr_from_param(value.get("condition"))), + then_value=expr_from_param(value.get("then")), + else_value=expr_from_param(value.get("else")), + ) if kind == "boolean": op = value.get("op") operands = value.get("operands") @@ -114,10 +163,12 @@ def expr_from_param(value: object) -> "Expr": raise ValueError(f"Unsupported expression kind: {kind!r}") -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class Expr: """Base class for serializable relational expressions.""" + __hash__ = None # type: ignore[assignment] + def to_param(self) -> ExprParam: """Return a QueryExpr parameter representation.""" @@ -131,6 +182,24 @@ def __bool__(self) -> bool: "Pass them to filter/join, or compare to_param() for structural checks." ) + def __eq__(self, other: object) -> "ComparisonExpr": # type: ignore[override] + return ComparisonExpr(op="eq", left=self, right=ensure_expr(other)) + + def __ne__(self, other: object) -> "ComparisonExpr": # type: ignore[override] + return ComparisonExpr(op="ne", left=self, right=ensure_expr(other)) + + def __lt__(self, other: object) -> "ComparisonExpr": + return ComparisonExpr(op="lt", left=self, right=ensure_expr(other)) + + def __le__(self, other: object) -> "ComparisonExpr": + return ComparisonExpr(op="le", left=self, right=ensure_expr(other)) + + def __gt__(self, other: object) -> "ComparisonExpr": + return ComparisonExpr(op="gt", left=self, right=ensure_expr(other)) + + def __ge__(self, other: object) -> "ComparisonExpr": + return ComparisonExpr(op="ge", left=self, right=ensure_expr(other)) + def __and__(self, other: object) -> "BooleanExpr": """Build a boolean AND expression.""" @@ -171,13 +240,51 @@ def array_cat(self, other: object) -> "ArrayCatExpr": return ArrayCatExpr(left=self, right=ensure_expr(other)) + def __add__(self, other: object) -> "ArithmeticExpr": + """Build an addition expression.""" + + return ArithmeticExpr(op="add", left=self, right=ensure_expr(other)) + + def __radd__(self, other: object) -> "ArithmeticExpr": + """Build an addition expression with a scalar left operand.""" + + return ArithmeticExpr(op="add", left=ensure_expr(other), right=self) + + def __sub__(self, other: object) -> "ArithmeticExpr": + """Build a subtraction expression.""" + + return ArithmeticExpr(op="subtract", left=self, right=ensure_expr(other)) + + def __rsub__(self, other: object) -> "ArithmeticExpr": + """Build a subtraction expression with a scalar left operand.""" + + return ArithmeticExpr(op="subtract", left=ensure_expr(other), right=self) + + def __mul__(self, other: object) -> "ArithmeticExpr": + """Build a multiplication expression.""" + + return ArithmeticExpr(op="multiply", left=self, right=ensure_expr(other)) + + def __rmul__(self, other: object) -> "ArithmeticExpr": + """Build a multiplication expression with a scalar left operand.""" + + return ArithmeticExpr(op="multiply", left=ensure_expr(other), right=self) + + def __truediv__(self, other: object) -> "ArithmeticExpr": + """Build a division expression.""" + + return ArithmeticExpr(op="divide", left=self, right=ensure_expr(other)) + + def __rtruediv__(self, other: object) -> "ArithmeticExpr": + """Build a division expression with a scalar left operand.""" + + return ArithmeticExpr(op="divide", left=ensure_expr(other), right=self) + @dataclass(frozen=True, eq=False) class ColumnExpr(Expr): """Reference to one relation column, optionally qualified by an alias.""" - __hash__ = None # type: ignore[assignment] - name: str qualifier: str | None = None @@ -192,26 +299,8 @@ def to_param(self) -> ExprParam: return {"kind": "column", "name": self.name, "qualifier": self.qualifier} - def __eq__(self, other: object) -> "ComparisonExpr": # type: ignore[override] - return ComparisonExpr(op="eq", left=self, right=ensure_expr(other)) - def __ne__(self, other: object) -> "ComparisonExpr": # type: ignore[override] - return ComparisonExpr(op="ne", left=self, right=ensure_expr(other)) - - def __lt__(self, other: object) -> "ComparisonExpr": - return ComparisonExpr(op="lt", left=self, right=ensure_expr(other)) - - def __le__(self, other: object) -> "ComparisonExpr": - return ComparisonExpr(op="le", left=self, right=ensure_expr(other)) - - def __gt__(self, other: object) -> "ComparisonExpr": - return ComparisonExpr(op="gt", left=self, right=ensure_expr(other)) - - def __ge__(self, other: object) -> "ComparisonExpr": - return ComparisonExpr(op="ge", left=self, right=ensure_expr(other)) - - -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class LiteralExpr(Expr): """Scalar literal expression.""" @@ -223,7 +312,7 @@ def to_param(self) -> ExprParam: return {"kind": "literal", "value": self.value} -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class ComparisonExpr(Expr): """Binary comparison expression.""" @@ -242,7 +331,62 @@ def to_param(self) -> ExprParam: } -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) +class ArithmeticExpr(Expr): + """Binary numeric expression with SQL-style null propagation.""" + + op: Literal["add", "subtract", "multiply", "divide"] + left: Expr + right: Expr + + def to_param(self) -> ExprParam: + """Return a QueryExpr parameter representation.""" + + return { + "kind": "arithmetic", + "op": self.op, + "left": self.left.to_param(), + "right": self.right.to_param(), + } + + +@dataclass(frozen=True, eq=False) +class TryCastExpr(Expr): + """Nullable scalar conversion expression.""" + + value: Expr + target: Literal["float"] + + def to_param(self) -> ExprParam: + """Return a QueryExpr parameter representation.""" + + return { + "kind": "try_cast", + "value": self.value.to_param(), + "target": self.target, + } + + +@dataclass(frozen=True, eq=False) +class CaseWhenExpr(Expr): + """SQL-style conditional expression with an explicit else branch.""" + + condition: Expr + then_value: Expr + else_value: Expr + + def to_param(self) -> ExprParam: + """Return a QueryExpr parameter representation.""" + + return { + "kind": "case_when", + "condition": self.condition.to_param(), + "then": self.then_value.to_param(), + "else": self.else_value.to_param(), + } + + +@dataclass(frozen=True, eq=False) class BooleanExpr(Expr): """Boolean expression over relational expressions.""" @@ -259,7 +403,7 @@ def to_param(self) -> ExprParam: } -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class ArrayCatExpr(Expr): """Row-wise JSON array-state concatenation expression.""" @@ -276,7 +420,7 @@ def to_param(self) -> ExprParam: } -@dataclass(frozen=True) +@dataclass(frozen=True, eq=False) class LeastExpr(Expr): """Row-wise minimum expression over two or more nullable operands.""" diff --git a/src/agent_memory/policy/relation.py b/src/agent_memory/policy/relation.py index c1e5147..e9b6eca 100644 --- a/src/agent_memory/policy/relation.py +++ b/src/agent_memory/policy/relation.py @@ -3,19 +3,29 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence -from typing import Any +from typing import Any, cast from .aggregates import ( AggregateSpec, + AlgebraicAggregateSpec, + AvgAggregateSpec, SemanticAggregateSpec, + SumAggregateSpec, aggregate_output_names, + avg as avg_aggregate, + count as count_aggregate, + is_algebraic_aggregate_spec, normalize_aggregate_specs, normalize_min_columns, + sum as sum_aggregate, ) from .expressions import ( + ArithmeticExpr, ArrayCatExpr, + CaseWhenExpr, ColumnExpr, LeastExpr, + TryCastExpr, ensure_boolean_expr, ensure_expr, expr_to_param, @@ -134,20 +144,6 @@ 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, ...]: @@ -206,11 +202,20 @@ def _normalize_assignment_value(value: object) -> Mapping[str, Any]: """Normalize assign values to literal, column, or supported array expression params.""" expr = ensure_expr(value) - if isinstance(expr, (ColumnExpr, ArrayCatExpr, LeastExpr)) or is_scalar(value): + if isinstance( + expr, + ( + ColumnExpr, + ArithmeticExpr, + ArrayCatExpr, + LeastExpr, + TryCastExpr, + CaseWhenExpr, + ), + ) or is_scalar(value): return expr.to_param() raise TypeError( - "assign values must be scalar literals, column expressions, array_cat expressions, " - "or least expressions" + "assign values must be scalar literals or supported row expressions" ) @@ -297,6 +302,15 @@ def select(self, columns: Sequence[str]) -> "Relation": return self._derive("select", columns=tuple(columns)) + def limit(self, n: int) -> "Relation": + """Return at most n rows, without promising an order or chosen subset.""" + + if isinstance(n, bool) or not isinstance(n, int): + raise TypeError("limit requires a non-negative integer") + if n < 0: + raise ValueError("limit requires a non-negative integer") + return self._derive("limit", n=n) + def alias(self, name: str) -> "Relation": """Assign a temporary relation qualifier for self-joins.""" @@ -397,6 +411,30 @@ def group_by(self, keys: str | Sequence[str]) -> "GroupedRelation": ) ) + def agg(self, *aggregates: AggregateSpec) -> "Relation": + """Apply one or more exact global relational aggregates.""" + + specs = _normalize_algebraic_aggregates( + aggregates, + source_columns=output_columns(self.expr), + ) + return self._derive("agg", aggregates=specs) + + def count(self, *, output_col: str) -> "Relation": + """Count all rows in this relation.""" + + return self.agg(count_aggregate(output_col=output_col)) + + def sum(self, *, column: str, output_col: str) -> "Relation": + """Sum one numeric column while ignoring null values.""" + + return self.agg(sum_aggregate(column=column, output_col=output_col)) + + def avg(self, *, column: str, output_col: str) -> "Relation": + """Average one numeric column while ignoring null values.""" + + return self.agg(avg_aggregate(column=column, output_col=output_col)) + def drop_duplicates( self, *, @@ -635,7 +673,7 @@ def sem_join( *, instruction: str, how: str = "inner", - on: str | Sequence[str] | None = None, + on: JoinOn | None = None, k: int | None = None, ) -> "Relation": """Add a semantic join expression.""" @@ -646,7 +684,7 @@ def sem_join( "how": how, } if on is not None: - params["on"] = _normalize_sem_join_keys(on) + params["on"] = _normalize_join_on(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") @@ -730,10 +768,9 @@ def __init__( class GroupedRelation(RelationHandle): - """Intermediate grouped expression returned by sem_groupby. + """Intermediate expression returned by group_by or sem_groupby. - A GroupedRelation represents semantic partition state before aggregation. It - becomes a normal Relation only after sem_agg is called. + It becomes a normal Relation after applying a supported aggregate. """ def sem_agg( @@ -815,11 +852,35 @@ def min( ) ) + def count(self, *, output_col: str) -> Relation: + """Count all rows in each ordinary relational group.""" + + return self.agg(count_aggregate(output_col=output_col)) + + def sum(self, *, column: str, output_col: str) -> Relation: + """Sum one numeric column in each ordinary relational group.""" + + return self.agg(sum_aggregate(column=column, output_col=output_col)) + + def avg(self, *, column: str, output_col: str) -> Relation: + """Average one numeric column in each ordinary relational group.""" + + return self.agg(avg_aggregate(column=column, output_col=output_col)) + def agg(self, *aggregates: AggregateSpec) -> Relation: """Aggregate each group with one or more declared aggregate functions.""" specs = normalize_aggregate_specs(aggregates) + algebraic = tuple( + specification + for specification in specs + if is_algebraic_aggregate_spec(specification) + ) if self.expr.op == "sem_groupby": + if algebraic: + raise NotImplementedError( + "sem_groupby numeric algebraic aggregates are not supported" + ) semantic_outputs = set(_semantic_aggregate_output_names(specs)) if not semantic_outputs: raise NotImplementedError( @@ -833,6 +894,18 @@ def agg(self, *aggregates: AggregateSpec) -> Relation: "sem_groupby(...).agg(...) requires semantic keys to appear " f"in sem_agg output_cols; missing {missing_keys}." ) + elif algebraic: + if len(algebraic) != len(specs): + raise NotImplementedError( + "ordinary group_by cannot mix numeric algebraic and existing " + "aggregate specs" + ) + keys = tuple(str(key) for key in self.expr.params["keys"]) + _normalize_algebraic_aggregates( + algebraic, + source_columns=output_columns(self.expr.inputs[0]), + group_keys=keys, + ) return Relation( QueryExpr( op="agg", @@ -842,6 +915,40 @@ def agg(self, *aggregates: AggregateSpec) -> Relation: ) +def _normalize_algebraic_aggregates( + aggregates: Sequence[AggregateSpec], + *, + source_columns: Sequence[str], + group_keys: Sequence[str] = (), +) -> tuple[AlgebraicAggregateSpec, ...]: + """Validate exact numeric aggregate descriptors against one relation schema.""" + + specs = normalize_aggregate_specs(aggregates) + if not all(is_algebraic_aggregate_spec(specification) for specification in specs): + raise TypeError( + "Relation.agg accepts only agent_memory.count(...), " + "agent_memory.sum(...), and agent_memory.avg(...)" + ) + algebraic = cast(tuple[AlgebraicAggregateSpec, ...], specs) + inputs = { + specification.column + for specification in algebraic + if isinstance(specification, (SumAggregateSpec, AvgAggregateSpec)) + } + missing = sorted(inputs.difference(source_columns)) + if missing: + raise ValueError(f"aggregate input columns not found: {missing}") + output_names = tuple( + output + for specification in algebraic + for output in aggregate_output_names(specification) + ) + conflicts = sorted(set(group_keys).intersection(output_names)) + if conflicts: + raise ValueError(f"aggregate output columns conflict with group keys: {conflicts}") + return algebraic + + class WindowedRelation(RelationHandle): """Count-window expression that must be closed by process_window.""" diff --git a/src/agent_memory/policy/schema.py b/src/agent_memory/policy/schema.py index ea7f2e7..bd6dc47 100644 --- a/src/agent_memory/policy/schema.py +++ b/src/agent_memory/policy/schema.py @@ -6,9 +6,12 @@ from .aggregates import ( ArrayAggregateSpec, + AvgAggregateSpec, CollectListAggregateSpec, + CountAggregateSpec, MinAggregateSpec, SemanticAggregateSpec, + SumAggregateSpec, ) from .logical import ColumnSpec, QueryExpr @@ -44,7 +47,7 @@ def output_columns( f"{conflicts}" ) return source_columns + metadata_columns - if query.op in {"alias", "group_by"}: + if query.op in {"alias", "group_by", "limit"}: 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) @@ -165,7 +168,24 @@ def output_columns( ) if grouped.op == "sem_groupby": return _append_unique(_sem_groupby_partition_keys(grouped), outputs) - raise NotImplementedError("agg output columns require group_by or sem_groupby input") + if all( + isinstance( + aggregate, + (CountAggregateSpec, SumAggregateSpec, AvgAggregateSpec), + ) + for aggregate in query.params.get("aggregates", ()) + ): + return outputs + raise NotImplementedError( + "agg output columns require group_by, sem_groupby, or algebraic aggregates" + ) + if query.op in {"aggregate_state", "aggregate_state_update"}: + return tuple(str(column) for column in query.params["state_columns"]) + if query.op == "aggregate_finalize": + return _append_unique( + tuple(str(key) for key in query.params.get("group_keys", ())), + _aggregate_output_columns(query.params.get("aggregates", ())), + ) if query.op in {"sem_map", "sem_flat_map"}: columns = _append_unique( output_columns(query.inputs[0], window_source_columns=window_source_columns), @@ -312,7 +332,14 @@ def _aggregate_output_columns(aggregates: object) -> tuple[str, ...]: names = tuple(column.name for column in aggregate.output_cols) elif isinstance( aggregate, - (ArrayAggregateSpec, CollectListAggregateSpec, MinAggregateSpec), + ( + ArrayAggregateSpec, + CollectListAggregateSpec, + MinAggregateSpec, + CountAggregateSpec, + SumAggregateSpec, + AvgAggregateSpec, + ), ): names = (aggregate.output_col,) else: diff --git a/src/agent_memory/runtime/executor.py b/src/agent_memory/runtime/executor.py index 209624b..1bf7023 100644 --- a/src/agent_memory/runtime/executor.py +++ b/src/agent_memory/runtime/executor.py @@ -4,6 +4,7 @@ from collections import Counter, defaultdict, deque from collections.abc import Callable, Hashable, Mapping +from copy import deepcopy from dataclasses import dataclass from datetime import date, datetime, timezone from typing import Any @@ -100,10 +101,49 @@ def node_state(self) -> Mapping[str, pd.DataFrame]: return self._node_state + @property + def source_columns(self) -> tuple[str, ...]: + """Return the declared source relation columns in stable order.""" + + return tuple( + column.name for column in self.spec.log.expr.params.get("columns", ()) + ) + + @property + def source_row_count(self) -> int: + """Return the number of source rows in the committed state.""" + + return len(self._state.get("log", ())) + + def read_view(self, name: str) -> pd.DataFrame: + """Return a defensive copy of one declared public view.""" + + if name not in self.spec.views: + raise KeyError(f"unknown public view: {name!r}") + value = self._state.get(name) + if value is None: + return self._empty_view_frame(self.spec.views[name]) + return _copy_frame_with_objects(value) + def add(self, message: MessageInput) -> None: """Append one source row and propagate its change through the shared DAG.""" - row = self._normalize_message(message) + row = self.normalize_source_row(message) + self.apply_delta(pd.DataFrame([row], columns=self.source_columns)) + + def apply_delta(self, changed_rows: pd.DataFrame) -> None: + """Propagate one relation-valued source delta and publish it atomically.""" + + if not isinstance(changed_rows, pd.DataFrame): + raise TypeError("changed_rows must be a pandas DataFrame") + if changed_rows.empty: + return + expected_columns = list(self.source_columns) + if list(changed_rows.columns) != expected_columns: + raise ValueError( + "source delta columns must exactly match the declared log schema" + ) + changed_rows = _copy_frame_with_objects(changed_rows) staged_next_occurrence = self._next_occurrence def allocate(node_id: str) -> str: @@ -112,14 +152,16 @@ def allocate(node_id: str) -> str: staged_next_occurrence += 1 return occurrence - changed_rows = pd.DataFrame([row]) source_node_ids = tuple( node_id for node_id in self.policy.execution_order if self.policy.nodes[node_id].execution_kind == "source" ) source_prefix = source_node_ids[0] if source_node_ids else "log" - changed_rows.index = pd.Index([allocate(source_prefix)], dtype="object") + changed_rows.index = pd.Index( + [allocate(source_prefix) for _ in range(len(changed_rows))], + dtype="object", + ) current_log = ( self._node_state.get(source_node_ids[0]) if source_node_ids @@ -148,15 +190,32 @@ def allocate(node_id: str) -> str: parent_updates = tuple(updates[parent_id] for parent_id in node.input_node_ids) old_state = self._node_state.get(node_id, self._empty_node_frame(node_id)) - if all(update.is_empty for update in parent_updates): + needs_global_aggregate_identity = ( + node.execution_kind == "algebraic_state" + and old_state.empty + and not tuple(node.query.params["group_keys"]) + ) + if ( + all(update.is_empty for update in parent_updates) + and not needs_global_aggregate_identity + ): staged_node_state.setdefault(node_id, old_state) updates[node_id] = NodeOutputUpdate.between(old_state, old_state.copy()) continue 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( + elif node.execution_kind == "algebraic_state": + next_state = self._execute_algebraic_state( + node_id, + old_state=old_state, + parent_update=parent_updates[0], + ) + elif node.execution_kind in { + "relational_state", + "semantic_binary_state", + }: + next_state = self._execute_binary_state( node_id, old_state=old_state, parent_updates=parent_updates, @@ -608,7 +667,7 @@ def _execute_deterministic( } return self.adapter.execute(node.query, inputs) - def _execute_relational_state( + def _execute_binary_state( self, node_id: str, *, @@ -616,15 +675,19 @@ def _execute_relational_state( parent_updates: tuple[NodeOutputUpdate, ...], staged_node_state: Mapping[str, pd.DataFrame], ) -> pd.DataFrame: - """Incrementally maintain an inner join or fully recompute on retractions.""" + """Incrementally maintain a binary inner join from old and inserted rows.""" 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") + raise RuntimeError("Binary state nodes require exactly two inputs") if any(not update.retracted_rows.empty for update in parent_updates): + if node.execution_kind == "semantic_binary_state": + raise NotImplementedError( + "Semantic binary state maintenance requires append-only input changes" + ) return self._execute_deterministic(node_id, staged_node_state) if node.maintenance_query is None: - raise RuntimeError(f"Relational state node {node_id} has no maintenance query") + raise RuntimeError(f"Binary state node {node_id} has no maintenance query") inputs: dict[str, pd.DataFrame] = {node_id: old_state} for parent_id, update in zip( @@ -639,6 +702,32 @@ def _execute_relational_state( inputs[f"{parent_id}__inserted"] = update.inserted_rows return self.adapter.execute(node.maintenance_query, inputs) + def _execute_algebraic_state( + self, + node_id: str, + *, + old_state: pd.DataFrame, + parent_update: NodeOutputUpdate, + ) -> pd.DataFrame: + """Apply exact parent insertions and retractions to aggregate state.""" + + node = self.policy.nodes[node_id] + if len(node.input_node_ids) != 1: + raise RuntimeError("Algebraic aggregate state requires exactly one input") + if node.maintenance_query is None: + raise RuntimeError( + f"Algebraic aggregate state node {node_id} has no maintenance query" + ) + parent_id = node.input_node_ids[0] + return self.adapter.execute( + node.maintenance_query, + { + node_id: old_state, + f"{parent_id}__inserted": parent_update.inserted_rows, + f"{parent_id}__retracted": parent_update.retracted_rows, + }, + ) + def _execute_semantic_row( self, node_id: str, @@ -798,7 +887,14 @@ def _empty_node_frame(self, node_id: str) -> pd.DataFrame: def _empty_view_frame(self, view: MemoryView) -> pd.DataFrame: return pd.DataFrame(columns=self.query_output_columns(view.query)) - def _normalize_message(self, message: MessageInput) -> dict[str, Any]: + def normalize_source_row( + self, + message: MessageInput, + *, + add_seq: int | None = None, + ) -> dict[str, Any]: + """Normalize one accepted message into the declared source schema.""" + row: dict[str, Any] if isinstance(message, str): row = {"message": message} @@ -825,13 +921,12 @@ def _normalize_message(self, message: MessageInput) -> dict[str, Any]: { LOG_ROW_ID_COLUMN: str(uuid4()), LOG_ADDED_AT_COLUMN: datetime.now(timezone.utc), - LOG_ADD_SEQ_COLUMN: len(self._state.get("log", ())), + LOG_ADD_SEQ_COLUMN: ( + self.source_row_count if add_seq is None else add_seq + ), } ) - declared_columns = tuple( - column.name for column in self.spec.log.expr.params.get("columns", ()) - ) - return {column: row.get(column) for column in declared_columns} + return {column: row.get(column) for column in self.source_columns} def _bind_user_query(self, query: QueryExpr, text: str) -> QueryExpr: return QueryExpr( @@ -889,6 +984,15 @@ def _default_adapter(self) -> Any: return LotusAdapter() +def _copy_frame_with_objects(frame: pd.DataFrame) -> pd.DataFrame: + """Detach caller-owned nested values at public input/output boundaries.""" + + result = frame.copy(deep=True) + for column in result.select_dtypes(include="object").columns: + result[column] = result[column].map(deepcopy) + return result + + def _multiset_difference( old: pd.DataFrame, new: pd.DataFrame, diff --git a/src/agent_memory/runtime/runtime.py b/src/agent_memory/runtime/runtime.py index 4d6196d..76ee256 100644 --- a/src/agent_memory/runtime/runtime.py +++ b/src/agent_memory/runtime/runtime.py @@ -5,10 +5,13 @@ from collections.abc import Mapping from typing import Any -from agent_memory.api import MessageInput +import pandas as pd + +from agent_memory.api import CountRefresh, MessageInput from agent_memory.planner.differential_policy import DifferentiatedPolicy from agent_memory.planner.legacy import differentiate_legacy_policy from agent_memory.policy.logical import QueryExpr +from agent_memory.policy.relation import LOG_ADD_SEQ_COLUMN, LOG_ROW_ID_COLUMN from agent_memory.runtime.executor import PolicyExecutor from agent_memory.runtime.legacy import LegacyViewRuntime from agent_memory.storage.deployment import StorageDeployment @@ -23,9 +26,12 @@ def __init__( *, adapter: Any | None = None, storage: StorageDeployment | None = None, + refresh: CountRefresh | None = None, ) -> None: self.policy = policy self.storage = storage + self.refresh = refresh + self._pending_rows: list[dict[str, Any]] = [] self._engine: PolicyExecutor | LegacyViewRuntime = PolicyExecutor( policy, adapter=adapter, @@ -48,7 +54,40 @@ def _state(self, state: Mapping[str, Any]) -> None: def add(self, message: MessageInput) -> None: """Append one input through the selected execution engine.""" - self._engine.add(message) + if self.refresh is None or self.refresh.every == 1: + self._engine.add(message) + return + if not isinstance(self._engine, PolicyExecutor): + raise RuntimeError("count refresh requires the schema-v2 policy executor") + + row = self._engine.normalize_source_row( + message, + add_seq=self._engine.source_row_count + len(self._pending_rows), + ) + self._pending_rows.append(row) + if len(self._pending_rows) >= self.refresh.every: + self.flush() + + @property + def pending_count(self) -> int: + """Return the number of normalized source rows awaiting refresh.""" + + return len(self._pending_rows) + + def flush(self) -> None: + """Apply all pending rows in one atomic policy-execution step.""" + + if not self._pending_rows: + return + if not isinstance(self._engine, PolicyExecutor): + raise RuntimeError("count refresh requires the schema-v2 policy executor") + + changed_rows = pd.DataFrame( + self._pending_rows, + columns=self._engine.source_columns, + ) + self._engine.apply_delta(changed_rows) + self._pending_rows.clear() def execute_retrieval_query(self, name: str, text: str) -> Any: """Execute one policy-owned retrieval query.""" @@ -58,12 +97,32 @@ def execute_retrieval_query(self, name: str, text: str) -> Any: def snapshot_state(self) -> dict[str, Any]: """Return a checkpoint in the active engine's schema.""" - return self._engine.snapshot_state() + engine_state = self._engine.snapshot_state() + if self.refresh is None or self.refresh.every == 1: + return engine_state + if not isinstance(self._engine, PolicyExecutor): + raise RuntimeError("count refresh requires the schema-v2 policy executor") + return { + "schema_version": 3, + "refresh": {"type": "count", "every": self.refresh.every}, + "pending_rows": pd.DataFrame( + self._pending_rows, + columns=self._engine.source_columns, + ), + "engine_state": engine_state, + } def restore_state(self, snapshot: Mapping[str, Any]) -> None: """Restore v2 directly or route schema-v1 state to the legacy engine.""" schema_version = snapshot.get("schema_version") + if schema_version == 3: + self._restore_count_refresh_state(snapshot) + return + if self.refresh is not None and self.refresh.every > 1: + raise ValueError( + "count refresh contract requires a schema-v3 runtime checkpoint" + ) if schema_version == 1: if self.storage is not None: raise NotImplementedError( @@ -88,6 +147,73 @@ def restore_state(self, snapshot: Mapping[str, Any]) -> None: return raise ValueError("Unsupported runtime snapshot schema_version") + def _restore_count_refresh_state(self, snapshot: Mapping[str, Any]) -> None: + """Restore a schema-v3 count-refresh envelope atomically.""" + + if self.refresh is None or self.refresh.every == 1: + raise ValueError( + "schema-v3 checkpoint requires a matching count refresh contract" + ) + if not isinstance(self._engine, PolicyExecutor): + raise ValueError("count refresh requires the schema-v2 policy executor") + + refresh = snapshot.get("refresh") + expected_refresh = {"type": "count", "every": self.refresh.every} + if refresh != expected_refresh: + raise ValueError("runtime snapshot refresh contract does not match") + engine_state = snapshot.get("engine_state") + if not isinstance(engine_state, Mapping): + raise ValueError("schema-v3 checkpoint requires engine_state") + pending_rows = snapshot.get("pending_rows") + if not isinstance(pending_rows, pd.DataFrame): + raise ValueError("schema-v3 checkpoint requires DataFrame pending_rows") + if list(pending_rows.columns) != list(self._engine.source_columns): + raise ValueError( + "runtime snapshot pending columns do not match the log schema" + ) + + pending_records = self._validate_pending_rows( + pending_rows, + engine_state=engine_state, + ) + self._engine.restore_state(engine_state) + self._pending_rows = pending_records + + def _validate_pending_rows( + self, + pending_rows: pd.DataFrame, + *, + engine_state: Mapping[str, Any], + ) -> list[dict[str, Any]]: + """Validate pending source identities before mutating engine state.""" + + state = engine_state.get("state") + if not isinstance(state, Mapping): + raise ValueError("schema-v3 engine_state requires state") + committed_rows = state.get("log") + if committed_rows is None: + committed_rows = pd.DataFrame(columns=self._engine.source_columns) + elif not isinstance(committed_rows, pd.DataFrame): + raise ValueError("schema-v3 engine_state requires DataFrame log state") + + if LOG_ADD_SEQ_COLUMN in pending_rows.columns: + expected = list( + range(len(committed_rows), len(committed_rows) + len(pending_rows)) + ) + if pending_rows[LOG_ADD_SEQ_COLUMN].tolist() != expected: + raise ValueError("runtime snapshot pending add sequence is invalid") + if LOG_ROW_ID_COLUMN in pending_rows.columns: + pending_ids = pending_rows[LOG_ROW_ID_COLUMN] + if pending_ids.isna().any() or not pending_ids.is_unique: + raise ValueError("runtime snapshot pending row IDs must be unique") + if LOG_ROW_ID_COLUMN in committed_rows.columns: + committed_ids = set(committed_rows[LOG_ROW_ID_COLUMN].dropna()) + if any(row_id in committed_ids for row_id in pending_ids): + raise ValueError( + "runtime snapshot pending row IDs overlap committed rows" + ) + return pending_rows.to_dict(orient="records") + def _query_output_columns(self, query: QueryExpr) -> list[str]: """Infer columns through the selected execution engine.""" diff --git a/src/agent_memory/tracing/semantic.py b/src/agent_memory/tracing/semantic.py index a1cc2f0..2ebaa4f 100644 --- a/src/agent_memory/tracing/semantic.py +++ b/src/agent_memory/tracing/semantic.py @@ -109,8 +109,7 @@ def trace_event_count( root = _event_root(trace_dir) with _TRACE_WRITE_LOCK: return sum( - _TRACE_EVENT_COUNTS.get((root, event_type), 0) - for event_type in event_types + _TRACE_EVENT_COUNTS.get((root, event_type), 0) for event_type in event_types ) @@ -129,6 +128,13 @@ def write_trace_event( if trace_dir is None: return None + journaled = trace_scope_value("provider_journal", False) + if journaled: + if event_type == "pair_decision": + return None + raw_output = parsed_output = None + snapshots = None + root = _event_root(trace_dir) root.mkdir(parents=True, exist_ok=True) trace_id = _trace_id(operator) @@ -182,7 +188,7 @@ def write_llm_call_trace( ) -> list[dict[str, Any]]: """Write actual LOTUS LM request/response trace events.""" - if trace_dir is None: + if trace_dir is None or trace_scope_value("provider_journal", False): return [] root = _event_root(trace_dir) @@ -251,20 +257,24 @@ def write_provider_usage_trace( ) -> list[dict[str, Any]]: """Write raw provider usage payloads observed at the ModelResponse boundary.""" - if trace_dir is None: + if trace_dir is None or trace_scope_value("provider_journal", False): return [] root = _event_root(trace_dir) root.mkdir(parents=True, exist_ok=True) scope = active_trace_scope() operator = str(scope.get("operator", scope.get("semantic_operator", "llm"))) - trace_operator = "provider-usage" if operator == "llm" else f"{operator}-provider-usage" + trace_operator = ( + "provider-usage" if operator == "llm" else f"{operator}-provider-usage" + ) batch_id = _trace_id(trace_operator) rows: list[dict[str, Any]] = [] for index, response in enumerate(responses): trace_id = f"{batch_id}-{index:04d}" usage = getattr(response, "usage", None) raw_usage = _usage_payload(usage) + response_metadata = getattr(response, "provider_response_metadata", {}) + choices = getattr(response, "choices", ()) event: dict[str, Any] = { "trace_id": trace_id, "timestamp": _timestamp(), @@ -275,6 +285,13 @@ def write_provider_usage_trace( "provider_batch_size": len(responses), "model": model, "provider_usage_available": raw_usage is not None, + "provider_finish_reason": response_metadata.get( + "finish_reason", + getattr(choices[0], "finish_reason", None) if choices else None, + ), + "provider_response_status": response_metadata.get("status"), + "provider_incomplete_reason": response_metadata.get("incomplete_reason"), + "provider_response_model": getattr(response, "model", None), } event.update(scope) if request_metadata: @@ -287,7 +304,16 @@ def write_provider_usage_trace( root / TRACE_OUTPUTS_DIR, trace_id, "provider-usage.json", - raw_usage, + response_metadata.get("raw_usage") or raw_usage, + ) + partial_response = getattr(response, "response", None) + partial_choices = getattr(partial_response, "choices", ()) + if partial_choices: + event["provider_partial_output_path"] = _write_json_artifact( + root / TRACE_OUTPUTS_DIR, + trace_id, + "partial-output.json", + {"output": getattr(partial_choices[0].message, "content", None)}, ) _append_event(root, event) rows.append(event) @@ -487,7 +513,9 @@ def _json_safe(value: Any) -> Any: except TypeError: if isinstance(value, Mapping): return {str(key): _json_safe(item) for key, item in value.items()} - if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + if isinstance(value, Sequence) and not isinstance( + value, (str, bytes, bytearray) + ): return [_json_safe(item) for item in value] return str(value) @@ -523,7 +551,9 @@ def _normalized_provider_usage(usage: Any, raw_usage: Any) -> dict[str, Any]: return { "provider_prompt_tokens": _usage_value(usage, raw_usage, "prompt_tokens"), - "provider_completion_tokens": _usage_value(usage, raw_usage, "completion_tokens"), + "provider_completion_tokens": _usage_value( + usage, raw_usage, "completion_tokens" + ), "provider_total_tokens": _usage_value(usage, raw_usage, "total_tokens"), "provider_prompt_cache_hit_tokens": _usage_value( usage, diff --git a/tests/test_agent_memory_benchmark_drivers.py b/tests/test_agent_memory_benchmark_drivers.py index cd879c2..71a11e6 100644 --- a/tests/test_agent_memory_benchmark_drivers.py +++ b/tests/test_agent_memory_benchmark_drivers.py @@ -16,6 +16,7 @@ PAIR_RIGHT_TEXT_COLUMN, SemanticPairExecutionProfile, ) +from agent_memory.adapters.lotus.prompt_batching import PromptBatching from agent_memory.evaluation.agent_memory_drivers import ( ClaudeMemoryDriverFactory, ClaudeMemoryDriver, @@ -657,6 +658,139 @@ def run(self, bundle) -> None: assert captured["closed"] is True +def test_run_records_and_propagates_count_refresh_contract( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import agent_memory.evaluation.run as run_module + + captured: dict[str, Any] = {} + + class FakeFactory: + def __init__(self, **kwargs: Any) -> None: + captured["factory"] = kwargs + + class FakeRunner: + def __init__(self, **kwargs: Any) -> None: + captured["runner"] = kwargs + + def run(self, bundle: BenchmarkBundle) -> None: + captured["bundle"] = bundle + + monkeypatch.setattr(run_module, "_require_environment", lambda _system: None) + monkeypatch.setattr( + run_module, + "collect_runtime_provenance", + lambda *args, **kwargs: {"source": {}, "runtime": {}}, + ) + monkeypatch.setattr( + run_module, + "validate_run_provenance", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr(run_module, "ClaudeMemoryDriverFactory", FakeFactory) + monkeypatch.setattr(run_module, "BenchmarkRunner", FakeRunner) + bundle = BenchmarkBundle( + "locomo", + "revision", + "sha256", + ( + BenchmarkCase( + case_id="case-1", + task_id="locomo", + events=(_event(),), + questions=(BenchmarkQuestion("q1", "case-1", "?", "a", ()),), + ), + ), + {"run_mode": "integration-smoke"}, + ) + + run_agent_memory_bundle( + bundle=bundle, + contracts={}, + system_id="claude-memory", + output_dir=tmp_path / "output", + refresh_every=4, + memory_thinking_enabled=False, + ) + + assert captured["factory"]["refresh_every"] == 4 + assert captured["runner"]["system_contract"].maintenance_execution_id == ( + "refresh:count:4" + ) + assert captured["runner"]["runtime_provenance"]["runtime"]["refresh"] == { + "type": "count", + "every": 4, + } + + +def test_run_records_and_propagates_prompt_batching_contract( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import agent_memory.evaluation.run as run_module + + captured: dict[str, Any] = {} + + class FakeFactory: + def __init__(self, **kwargs: Any) -> None: + captured["factory"] = kwargs + + class FakeRunner: + def __init__(self, **kwargs: Any) -> None: + captured["runner"] = kwargs + + def run(self, bundle: BenchmarkBundle) -> None: + captured["bundle"] = bundle + + monkeypatch.setattr(run_module, "_require_environment", lambda _system: None) + monkeypatch.setattr( + run_module, + "collect_runtime_provenance", + lambda *args, **kwargs: {"source": {}, "runtime": {}}, + ) + monkeypatch.setattr( + run_module, + "validate_run_provenance", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr(run_module, "ClaudeMemoryDriverFactory", FakeFactory) + monkeypatch.setattr(run_module, "BenchmarkRunner", FakeRunner) + bundle = BenchmarkBundle( + "locomo", + "revision", + "sha256", + ( + BenchmarkCase( + case_id="case-1", + task_id="locomo", + events=(_event(),), + questions=(BenchmarkQuestion("q1", "case-1", "?", "a", ()),), + ), + ), + {"run_mode": "integration-smoke"}, + ) + prompt_batching = PromptBatching(max_tasks=8) + + run_agent_memory_bundle( + bundle=bundle, + contracts={}, + system_id="claude-memory", + output_dir=tmp_path / "output", + prompt_batching=prompt_batching, + memory_thinking_enabled=False, + ) + + assert captured["factory"]["prompt_batching"] == prompt_batching + contract = captured["runner"]["system_contract"] + assert contract.maintenance_execution_id == ( + f"prompt-batching:{prompt_batching.fingerprint}" + ) + assert captured["runner"]["runtime_provenance"]["runtime"][ + "lotus_execution" + ]["prompt_batching"] == {"max_tasks": 8} + + def test_zep_run_closes_factory_when_storage_provenance_fails( monkeypatch, tmp_path ) -> None: @@ -1219,6 +1353,7 @@ def unexpected_provenance(*args, **kwargs): contracts={}, system_id="claude-memory", output_dir=tmp_path / "output", + grouped_agg_rule="rule-all-group", semantic_pair_profile="search-filter", semantic_pair_min_similarity=0.6, ) diff --git a/tests/test_aggregate_ivm.py b/tests/test_aggregate_ivm.py new file mode 100644 index 0000000..4c74cd5 --- /dev/null +++ b/tests/test_aggregate_ivm.py @@ -0,0 +1,215 @@ +"""Exact incremental-maintenance tests for relational algebraic aggregates.""" + +from __future__ import annotations + +from collections import Counter + +import pandas as pd + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.policy.schema import output_columns + + +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"]}) + + +class _GlobalAggregateMemory(am.Memory): + log = am.Log({"amount": "Numeric amount."}) + metrics = log.agg( + am.count(output_col="row_count"), + am.sum(column="amount", output_col="total"), + am.avg(column="amount", output_col="average"), + ) + + +class _FilteredGlobalAggregateMemory(am.Memory): + log = am.Log({"kind": "Row kind.", "amount": "Numeric amount."}) + _selected = log.filter(log.col("kind") == "selected") + metrics = _selected.agg( + am.count(output_col="row_count"), + am.sum(column="amount", output_col="total"), + am.avg(column="amount", output_col="average"), + ) + + +class _GroupedAggregateMemory(am.Memory): + log = am.Log({"region": "Region.", "amount": "Numeric amount."}) + metrics = log.group_by("region").agg( + am.count(output_col="row_count"), + am.sum(column="amount", output_col="total"), + am.avg(column="amount", output_col="average"), + ) + + +class _ReplacementAggregateMemory(am.Memory): + log = am.Log({"key": "Key.", "amount": "Numeric amount."}) + _minimums = log.group_by("key").min(column="amount", output_col="amount") + metrics = _minimums.agg( + am.count(output_col="row_count"), + am.sum(column="amount", output_col="total"), + am.avg(column="amount", output_col="average"), + ) + + +class _GroupMovementMemory(am.Memory): + log = am.Log({"group": "Ordered group."}) + _current_group = log.min(column="group", output_col="group") + metrics = _current_group.group_by("group").count(output_col="row_count") + + +def test_planner_lowers_aggregate_to_hidden_state_and_public_finalizer() -> None: + policy = _GlobalAggregateMemory.differentiate_policy() + state_nodes = [ + node for node in policy.nodes.values() if node.execution_kind == "algebraic_state" + ] + + assert len(state_nodes) == 1 + state = state_nodes[0] + assert state.query.op == "aggregate_state" + assert state.maintenance_query is not None + assert state.maintenance_query.op == "aggregate_state_update" + assert any(column.startswith("__am_aggregate_") for column in state.output_columns) + assert len(tuple(state.query.params["numeric_states"])) == 1 + + public_node = policy.nodes[policy.view_outputs["metrics"]] + assert public_node.query.op == "aggregate_finalize" + assert public_node.execution_kind == "deterministic" + assert public_node.output_columns == ("row_count", "total", "average") + assert all( + not column.startswith("__am_aggregate_") + for column in public_node.output_columns + ) + + grouped_policy = _GroupedAggregateMemory.differentiate_policy() + grouped_public = grouped_policy.nodes[grouped_policy.view_outputs["metrics"]] + assert tuple(grouped_public.query.params["group_keys"]) == ("region",) + assert tuple(grouped_public.output_columns) == ( + "region", + "row_count", + "total", + "average", + ) + assert output_columns(grouped_public.query) == grouped_public.output_columns + + +def test_global_incremental_result_matches_full_recomputation_after_each_batch() -> None: + memory = _GlobalAggregateMemory() + batches = ( + pd.DataFrame([{"amount": 10}, {"amount": None}]), + pd.DataFrame([{"amount": 20}, {"amount": 20}]), + ) + + for batch in batches: + memory._runtime._engine.apply_delta(batch) + assert _bag(memory._runtime._state["metrics"]) == _bag( + _full_view(memory, "metrics") + ) + + assert memory._runtime._state["metrics"].to_dict(orient="records") == [ + {"row_count": 4, "total": 50, "average": 50 / 3} + ] + + +def test_global_aggregate_materializes_identity_when_parent_delta_is_empty() -> None: + memory = _FilteredGlobalAggregateMemory() + + memory.add({"kind": "ignored", "amount": 10}) + + assert memory._runtime._state["metrics"].to_dict(orient="records") == [ + {"row_count": 0, "total": None, "average": None} + ] + assert _bag(memory._runtime._state["metrics"]) == _bag( + _full_view(memory, "metrics") + ) + + +def test_grouped_incremental_result_matches_full_recomputation() -> None: + memory = _GroupedAggregateMemory() + + for row in ( + {"region": "north", "amount": 10}, + {"region": "south", "amount": 5}, + {"region": "north", "amount": None}, + {"region": "north", "amount": 20}, + ): + memory.add(row) + assert _bag(memory._runtime._state["metrics"]) == _bag( + _full_view(memory, "metrics") + ) + + +def test_parent_row_replacement_subtracts_old_values_without_full_recompute() -> None: + memory = _ReplacementAggregateMemory() + + memory.add({"key": "a", "amount": 5}) + memory.add({"key": "b", "amount": 7}) + memory.add({"key": "a", "amount": 3}) + + assert memory._runtime._state["metrics"].to_dict(orient="records") == [ + {"row_count": 2, "total": 10, "average": 5} + ] + assert _bag(memory._runtime._state["metrics"]) == _bag( + _full_view(memory, "metrics") + ) + + +def test_group_disappears_when_deterministic_parent_moves_its_only_row() -> None: + memory = _GroupMovementMemory() + + memory.add({"group": "b"}) + assert memory._runtime._state["metrics"].to_dict(orient="records") == [ + {"group": "b", "row_count": 1} + ] + + memory.add({"group": "a"}) + assert memory._runtime._state["metrics"].to_dict(orient="records") == [ + {"group": "a", "row_count": 1} + ] + assert _bag(memory._runtime._state["metrics"]) == _bag( + _full_view(memory, "metrics") + ) + + +def test_aggregate_snapshot_restore_preserves_hidden_state() -> None: + original = _GroupedAggregateMemory() + original.add({"region": "north", "amount": 10}) + original.add({"region": "south", "amount": 5}) + snapshot = original._runtime.snapshot_state() + + restored = _GroupedAggregateMemory() + restored._runtime.restore_state(snapshot) + original.add({"region": "north", "amount": 20}) + restored.add({"region": "north", "amount": 20}) + + assert _bag(restored._runtime._state["metrics"]) == _bag( + original._runtime._state["metrics"] + ) + assert _bag(restored._runtime._state["metrics"]) == _bag( + _full_view(restored, "metrics") + ) + + +def test_aggregate_plan_fingerprint_is_stable_and_builtin_defaults_are_pinned() -> None: + assert ( + _GlobalAggregateMemory.differentiate_policy().fingerprint + == _GlobalAggregateMemory.differentiate_policy().fingerprint + ) + assert ( + am.ClaudeMemory.differentiate_policy().fingerprint + == "625a53555e82e283d04576eacf7e937c590b546a485bcccc6224d3b8c4dae5fe" + ) + assert ( + am.ZepMemory.differentiate_policy().fingerprint + == "9d1cb107ef677e9624dffeb22399dd6edd7eb6d058e4a37da26c61c2a8e64059" + ) + assert ( + am.Mem0Memory.differentiate_policy().fingerprint + == "d904fdb2fbd64ffdbaef65de51fb7f366c29b7ebbd05ba5194e909725b7f2d28" + ) diff --git a/tests/test_benchmark_cli.py b/tests/test_benchmark_cli.py index f45b4db..c71fcc1 100644 --- a/tests/test_benchmark_cli.py +++ b/tests/test_benchmark_cli.py @@ -5,10 +5,40 @@ import pytest +from agent_memory.adapters.lotus.prompt_batching import PromptBatching + from agent_memory.planner.rules import GROUPED_AGG_RULES from tools.evaluation import locomo, longmemeval, memory_agent_bench +@pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) +@pytest.mark.parametrize( + ("system_id", "expected_rule"), + ( + ("claude-memory", "rule-join-map"), + ("zep-memory", "rule-join-map"), + ("mem0-memory", "rule-all-group"), + ("mem0-enhanced", "rule-all-group"), + ), +) +def test_agent_benchmark_clis_resolve_grouped_agg_default_by_system( + module, system_id, expected_rule, tmp_path +) -> None: + args = module.parse_args( + [ + "run", + "--bundle-dir", + str(tmp_path / "bundle"), + "--output-dir", + str(tmp_path / "output"), + "--system", + system_id, + ] + ) + + assert args.grouped_agg_rule == expected_rule + + @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( @@ -80,6 +110,50 @@ def test_agent_benchmark_clis_share_sem_groupby_execution_options( assert args.semantic_trace_snapshot_mode == "compact" +@pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) +@pytest.mark.parametrize( + ("value", "expected"), + (("all", PromptBatching()), ("4", PromptBatching(max_tasks=4))), +) +def test_agent_benchmark_clis_share_prompt_batching_option( + module, value, expected, tmp_path +) -> None: + args = module.parse_args( + [ + "run", + "--bundle-dir", + str(tmp_path / "bundle"), + "--output-dir", + str(tmp_path / "output"), + "--system", + "zep-memory", + "--prompt-batch-size", + value, + ] + ) + + assert args.prompt_batch_size == expected + + +@pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) +def test_agent_benchmark_clis_disable_prompt_batching_by_default( + module, tmp_path +) -> None: + args = module.parse_args( + [ + "run", + "--bundle-dir", + str(tmp_path / "bundle"), + "--output-dir", + str(tmp_path / "output"), + "--system", + "zep-memory", + ] + ) + + assert args.prompt_batch_size is None + + @pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) def test_agent_benchmark_clis_accept_full_semantic_trace_snapshots( module, tmp_path @@ -157,6 +231,25 @@ def test_agent_benchmark_clis_share_site_profiles_and_cache_mode( assert args.lotus_cache_mode == "memory" +@pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) +def test_agent_benchmark_clis_share_count_refresh_option(module, tmp_path) -> None: + args = module.parse_args( + [ + "run", + "--bundle-dir", + str(tmp_path / "bundle"), + "--output-dir", + str(tmp_path / "output"), + "--system", + "zep-memory", + "--refresh-every", + "4", + ] + ) + + assert args.refresh_every == 4 + + @pytest.mark.parametrize("module", (locomo, longmemeval, memory_agent_bench)) def test_agent_benchmark_clis_reject_site_and_global_profiles( module, tmp_path @@ -208,7 +301,7 @@ def test_longmemeval_cli_separates_bundle_preparation_from_system_run(tmp_path) assert run.memory_model_id == "deepseek-v4-flash" assert run.answer_model == "deepseek/deepseek-v4-flash" assert run.judge_model == "deepseek/deepseek-v4-flash" - assert run.grouped_agg_rule == "rule-all-group" + assert run.grouped_agg_rule == "rule-join-map" assert run.sem_topk_method == "pairwise-naive" assert run.memory_thinking == "disabled" @@ -270,8 +363,8 @@ def test_longmemeval_cli_selects_zep_memory_without_claude_options(tmp_path) -> assert args.system == "zep-memory" assert args.namespace == "longmemeval-zep-smoke" assert args.memory_thinking == "enabled" - assert args.grouped_agg_rule == "rule-all-group" - assert longmemeval._condition_id(args) == "zep-memory-rule-all-group" + assert args.grouped_agg_rule == "rule-join-map" + assert longmemeval._condition_id(args) == "zep-memory-rule-join-map" @pytest.mark.parametrize( @@ -396,7 +489,7 @@ def test_longmemeval_cli_allows_derived_namespace_for_zep(tmp_path) -> None: ) assert args.namespace is None - assert longmemeval._condition_id(args) == "zep-memory-rule-all-group" + assert longmemeval._condition_id(args) == "zep-memory-rule-join-map" def test_longmemeval_zep_run_forwards_existing_runner_contract( diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py index 5e3153b..67dd1fc 100644 --- a/tests/test_benchmark_harness.py +++ b/tests/test_benchmark_harness.py @@ -194,6 +194,7 @@ class _Driver: fail_event_id: str | None = None fail_finish: bool = False restored_event_ids: tuple[str, ...] = () + flush_count: int = 0 def __post_init__(self) -> None: self.events = [] @@ -223,6 +224,10 @@ def finish_session(self, session_id: str) -> dict[str, str]: raise RuntimeError("consolidation failed") return {"session_id": session_id} + def flush(self) -> dict[str, int]: + self.flush_count += 1 + return {"flushed": self.flush_count} + def close(self) -> None: self.closed = True @@ -426,11 +431,13 @@ def test_runner_answers_once_and_records_primary_and_secondary_graders( "mean_ms": 0.7, "median_ms": 0.7, "p95_ms": 0.7, + "max_ms": 0.7, }, "secondary": { "mean_ms": 860.0, "median_ms": 860.0, "p95_ms": 860.0, + "max_ms": 860.0, }, } with (tmp_path / "metrics" / "per_question.csv").open( @@ -600,6 +607,10 @@ def factory(case_id: str, state_dir: Path, trace_dir: Path) -> _Driver: newline="", encoding="utf-8" ) as stream: case_row = next(csv.DictReader(stream)) + with (tmp_path / "metrics" / "overview.csv").open( + newline="", encoding="utf-8" + ) as stream: + overview_row = next(csv.DictReader(stream)) summary = json.loads((tmp_path / "metrics" / "summary.json").read_text()) assert len(event_rows) == 2 @@ -617,7 +628,72 @@ def factory(case_id: str, state_dir: Path, trace_dir: Path) -> _Driver: case_row["semantic_trace_bytes_written"] ) assert summary["semantic_trace_io_wall_latency"]["mean_ms"] is not None - assert summary["insertion_wall_latency_excluding_trace_io"]["mean_ms"] is not None + adjusted = summary["insertion_wall_latency_excluding_trace_io"] + adjusted_values = [ + float(row["insertion_latency_excluding_trace_io_ms"]) + for row in event_rows + ] + expected_middle = round(sum(adjusted_values) / 2, 3) + assert adjusted["mean_ms"] == expected_middle + assert adjusted["median_ms"] == expected_middle + assert adjusted["p95_ms"] == pytest.approx( + round(min(adjusted_values) * 0.05 + max(adjusted_values) * 0.95, 3) + ) + assert adjusted["max_ms"] == round(max(adjusted_values), 3) + insertion = summary["insertion_wall_latency"] + assert float(overview_row["insertion_mean_latency_ms"]) == pytest.approx( + insertion["mean_ms"] + ) + assert float(overview_row["insertion_median_latency_ms"]) == pytest.approx( + insertion["median_ms"] + ) + assert float(overview_row["insertion_p95_latency_ms"]) == pytest.approx( + insertion["p95_ms"] + ) + assert float( + overview_row["insertion_excluding_semantic_trace_mean_ms"] + ) == pytest.approx(adjusted["mean_ms"]) + assert float( + overview_row["insertion_excluding_semantic_trace_median_ms"] + ) == pytest.approx(adjusted["median_ms"]) + assert float( + overview_row["insertion_excluding_semantic_trace_p95_ms"] + ) == pytest.approx(adjusted["p95_ms"]) + assert float( + overview_row["insertion_excluding_semantic_trace_max_ms"] + ) == pytest.approx(adjusted["max_ms"]) + + events_path = tmp_path / "trace" / "events.jsonl" + trace_rows = [json.loads(line) for line in events_path.read_text().splitlines()] + for row in trace_rows: + if row.get("event_type") != "insertion_result": + continue + row.pop("semantic_trace_io_latency_ms", None) + row.pop("semantic_trace_bytes_written", None) + row.pop("insertion_latency_excluding_trace_io_ms", None) + events_path.write_text( + "".join(f"{json.dumps(row)}\n" for row in trace_rows), + encoding="utf-8", + ) + BenchmarkArtifactStore(tmp_path).finalize_metrics() + + legacy_summary = json.loads( + (tmp_path / "metrics" / "summary.json").read_text() + ) + with (tmp_path / "metrics" / "overview.csv").open( + newline="", encoding="utf-8" + ) as stream: + legacy_overview = next(csv.DictReader(stream)) + assert legacy_summary["insertion_wall_latency_excluding_trace_io"] == { + "mean_ms": None, + "median_ms": None, + "p95_ms": None, + "max_ms": None, + } + assert legacy_overview["insertion_excluding_semantic_trace_mean_ms"] == "" + assert legacy_overview["insertion_excluding_semantic_trace_median_ms"] == "" + assert legacy_overview["insertion_excluding_semantic_trace_p95_ms"] == "" + assert legacy_overview["insertion_excluding_semantic_trace_max_ms"] == "" def test_operation_usage_separates_calls_batches_and_provider_responses( @@ -1428,6 +1504,55 @@ def _session_bundle() -> BenchmarkBundle: return BenchmarkBundle("benchmark", "revision", "sha", (case,)) +def test_runner_flushes_once_after_the_last_event_before_final_checkpoint( + tmp_path: Path, +) -> None: + operations: list[str] = [] + + class RecordingDriver(_Driver): + def add(self, event: BenchmarkEvent) -> dict[str, int]: + operations.append(f"add:{event.event_id}") + return super().add(event) + + def finish_session(self, session_id: str) -> dict[str, str]: + operations.append(f"finish:{session_id}") + return super().finish_session(session_id) + + def flush(self) -> dict[str, int]: + operations.append("flush") + return super().flush() + + def save_state(self, directory: Path) -> dict[str, Any]: + operations.append("checkpoint") + return super().save_state(directory) + + drivers: list[RecordingDriver] = [] + runner = BenchmarkRunner( + system_contract=_system_contract(), + contracts={"task-1": _contract()}, + driver_factory=lambda case_id, state_dir, trace_dir: drivers.append( + RecordingDriver(state_dir) + ) + or drivers[-1], + answer_model=_Model(), + judge_model=_Model(), + artifacts=BenchmarkArtifactStore(tmp_path), + maintenance_only=True, + ) + + runner.run(_session_bundle()) + + assert drivers[0].flush_count == 1 + assert operations.count("flush") == 1 + assert operations.index("finish:session-1") < operations.index("flush") + assert operations.index("add:event-4") < operations.index("flush") + assert operations.index("flush") < len(operations) - 1 + assert operations[operations.index("flush") + 1 :] == [ + "finish:session-2", + "checkpoint", + ] + + def test_session_checkpoint_resumes_without_readding_completed_session( tmp_path: Path, ) -> None: diff --git a/tests/test_benchmark_trace_metrics.py b/tests/test_benchmark_trace_metrics.py index d129752..a9a120b 100644 --- a/tests/test_benchmark_trace_metrics.py +++ b/tests/test_benchmark_trace_metrics.py @@ -11,6 +11,19 @@ ) +def test_normalize_provider_calls_can_omit_pricing(tmp_path: Path) -> None: + events = [{ + "event_type": "provider_usage", "provider_usage_available": True, + "provider_prompt_tokens": 10, "provider_completion_tokens": 2, + "provider_prompt_cache_hit_tokens": 4, "provider_prompt_cache_miss_tokens": 6, + }] + priced = normalize_provider_calls(events, output_dir=tmp_path)[0] + unpriced = normalize_provider_calls(events, output_dir=tmp_path, include_cost=False)[0] + assert priced["estimated_cost_usd"] is not None + for key, value in priced.items(): + assert unpriced[key] == (None if key.endswith("cost_usd") else value) + + def test_normalize_provider_calls_deduplicates_agent_semantic_usage( tmp_path: Path, ) -> None: diff --git a/tests/test_count_refresh.py b/tests/test_count_refresh.py new file mode 100644 index 0000000..222e962 --- /dev/null +++ b/tests/test_count_refresh.py @@ -0,0 +1,321 @@ +"""Focused tests for count-triggered relation-valued refresh.""" + +from __future__ import annotations + +import json +from typing import Any + +import agent_memory as am +import pandas as pd +import pytest +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.pair_execution import ( + PAIR_LEFT_ID_COLUMN, + PAIR_RIGHT_ID_COLUMN, +) +from agent_memory.policy.logical import QueryExpr + + +class _CountingAdapter(LotusAdapter): + """Record root operator executions while preserving normal adapter behavior.""" + + def __init__(self) -> None: + super().__init__() + self.calls: list[str] = [] + self.fail_assign = False + + def execute(self, query: QueryExpr, inputs: dict[str, Any]) -> Any: + self.calls.append(query.op) + if query.op == "assign" and self.fail_assign: + raise RuntimeError("injected refresh failure") + return super().execute(query, inputs) + + +class _BufferedMemory(am.Memory): + log = am.Log({"message": "Message body."}, system_columns=True) + rows = log.assign(published=True).select( + ["message", "published", "_row_id", "_add_seq"] + ) + retrieval_query = rows + + +class _SemanticFilterAdapter(LotusAdapter): + """Accept every semantic-filter row and record the relation delta size.""" + + def __init__(self) -> None: + super().__init__() + self.input_sizes: list[int] = [] + + def execute(self, query: QueryExpr, inputs: dict[str, Any]) -> Any: + if query.op == "sem_filter": + source = inputs[str(query.inputs[0].params["name"])] + self.input_sizes.append(len(source)) + return source.copy() + return super().execute(query, inputs) + + +class _SemanticFilterMemory(am.Memory): + log = am.Log({"message": "Message body."}) + rows = log.sem_filter(instruction="Keep {message}.") + + +class _SemanticSelfJoinMemory(am.Memory): + log = am.Log({"key": "Exact partition.", "value": "Semantic value."}) + _left = log.alias("left") + _right = log.alias("right") + pairs = _left.sem_join( + _right, + on="key", + instruction="{value:left} and {value:right} mean the same thing.", + ) + + +class _WindowMemory(am.Memory): + log = am.Log({"message": "Message body."}) + blocks = log.count_window(size=2, slide=1).process_window( + lambda window: window.array_agg( + columns=("message",), + output_col="messages", + ) + ) + + +def test_count_refresh_publishes_one_relation_delta_at_threshold() -> None: + adapter = _CountingAdapter() + memory = _BufferedMemory(adapter=adapter, refresh=am.CountRefresh(every=3)) + + memory.add({"message": "one"}) + memory.add({"message": "two"}) + + assert memory.pending_count == 2 + assert memory._runtime._state == {} + assert "assign" not in adapter.calls + + memory.add({"message": "three"}) + + assert memory.pending_count == 0 + assert adapter.calls.count("assign") == 1 + assert memory._runtime._state["rows"]["message"].tolist() == [ + "one", + "two", + "three", + ] + assert memory._runtime._state["rows"]["_add_seq"].tolist() == [0, 1, 2] + assert memory._runtime._state["rows"]["_row_id"].is_unique + + +def test_query_reads_committed_view_and_flush_publishes_the_tail() -> None: + memory = _BufferedMemory(refresh=am.CountRefresh(every=3)) + for message in ("one", "two", "three"): + memory.add({"message": message}) + memory.add({"message": "four"}) + + result = memory.query("unused") + + assert result["message"].tolist() == ["one", "two", "three"] + assert memory.pending_count == 1 + + memory.flush() + + assert memory.pending_count == 0 + assert memory._runtime._state["rows"]["message"].tolist() == [ + "one", + "two", + "three", + "four", + ] + + +def test_failed_refresh_retains_pending_rows_and_publishes_nothing() -> None: + adapter = _CountingAdapter() + adapter.fail_assign = True + memory = _BufferedMemory(adapter=adapter, refresh=am.CountRefresh(every=2)) + + memory.add({"message": "one"}) + pending_ids = [memory._runtime._pending_rows[0]["_row_id"]] + with pytest.raises(RuntimeError, match="injected refresh failure"): + memory.add({"message": "two"}) + pending_ids.append(memory._runtime._pending_rows[1]["_row_id"]) + + assert memory.pending_count == 2 + assert memory._runtime._state == {} + assert memory._runtime._engine.node_state == {} + + adapter.fail_assign = False + memory.flush() + + assert memory.pending_count == 0 + assert memory._runtime._state["rows"]["_row_id"].tolist() == pending_ids + + +def test_count_refresh_snapshot_restores_pending_rows_with_matching_config() -> None: + original = _BufferedMemory(refresh=am.CountRefresh(every=3)) + original.add({"message": "one"}) + original.add({"message": "two"}) + + snapshot = original._runtime.snapshot_state() + restored = _BufferedMemory(refresh=am.CountRefresh(every=3)) + restored._runtime.restore_state(snapshot) + + assert snapshot["schema_version"] == 3 + assert snapshot["refresh"] == {"type": "count", "every": 3} + assert restored.pending_count == 2 + + restored.flush() + + assert restored._runtime._state["rows"]["message"].tolist() == ["one", "two"] + assert restored._runtime._state["rows"]["_row_id"].tolist() == snapshot[ + "pending_rows" + ]["_row_id"].tolist() + + +def test_count_refresh_restore_rejects_a_different_refresh_contract() -> None: + original = _BufferedMemory(refresh=am.CountRefresh(every=3)) + original.add({"message": "one"}) + snapshot = original._runtime.snapshot_state() + restored = _BufferedMemory(refresh=am.CountRefresh(every=4)) + + with pytest.raises(ValueError, match="refresh contract"): + restored._runtime.restore_state(snapshot) + + +def test_deferred_runtime_rejects_an_eager_checkpoint() -> None: + eager = _BufferedMemory() + eager.add({"message": "one"}) + deferred = _BufferedMemory(refresh=am.CountRefresh(every=3)) + + with pytest.raises(ValueError, match="schema-v3"): + deferred._runtime.restore_state(eager._runtime.snapshot_state()) + + +def test_eager_runtime_rejects_a_deferred_checkpoint() -> None: + deferred = _BufferedMemory(refresh=am.CountRefresh(every=3)) + deferred.add({"message": "one"}) + eager = _BufferedMemory() + + with pytest.raises(ValueError, match="matching count refresh contract"): + eager._runtime.restore_state(deferred._runtime.snapshot_state()) + + +def test_count_refresh_restore_rejects_invalid_pending_sequence() -> None: + original = _BufferedMemory(refresh=am.CountRefresh(every=3)) + original.add({"message": "one"}) + snapshot = original._runtime.snapshot_state() + snapshot["pending_rows"].loc[0, "_add_seq"] = 7 + restored = _BufferedMemory(refresh=am.CountRefresh(every=3)) + + with pytest.raises(ValueError, match="pending add sequence"): + restored._runtime.restore_state(snapshot) + + assert restored.pending_count == 0 + assert restored._runtime._state == {} + + +def test_count_refresh_restore_rejects_duplicate_pending_row_ids() -> None: + original = _BufferedMemory(refresh=am.CountRefresh(every=3)) + original.add({"message": "one"}) + original.add({"message": "two"}) + snapshot = original._runtime.snapshot_state() + snapshot["pending_rows"].loc[1, "_row_id"] = snapshot["pending_rows"].loc[ + 0, "_row_id" + ] + restored = _BufferedMemory(refresh=am.CountRefresh(every=3)) + + with pytest.raises(ValueError, match="row IDs must be unique"): + restored._runtime.restore_state(snapshot) + + +def test_count_refresh_every_one_is_exactly_eager() -> None: + adapter = _CountingAdapter() + memory = _BufferedMemory(adapter=adapter, refresh=am.CountRefresh(every=1)) + + memory.add({"message": "one"}) + + assert memory.pending_count == 0 + assert adapter.calls.count("assign") == 1 + assert memory._runtime.snapshot_state()["schema_version"] == 2 + + +def test_eager_runtime_keeps_schema_v2_checkpoint_contract() -> None: + memory = _BufferedMemory() + memory.add({"message": "one"}) + + assert memory.pending_count == 0 + assert memory._runtime.snapshot_state()["schema_version"] == 2 + + +def test_relation_delta_reaches_semantic_row_operator_in_one_execution() -> None: + adapter = _SemanticFilterAdapter() + memory = _SemanticFilterMemory( + adapter=adapter, + refresh=am.CountRefresh(every=3), + ) + + for message in ("one", "two", "three"): + memory.add({"message": message}) + + assert adapter.input_sizes == [3] + assert memory._runtime._state["rows"]["message"].tolist() == [ + "one", + "two", + "three", + ] + + +def test_relation_delta_drives_semantic_self_join_once_without_missing_pairs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import agent_memory.adapters.lotus.sem_join as sem_join_module + + candidate_sizes: list[int] = [] + + def verify(candidates: pd.DataFrame, **_kwargs: object) -> list[tuple[Any, Any, None]]: + if not candidates.empty: + candidate_sizes.append(len(candidates)) + return [ + ( + row[PAIR_LEFT_ID_COLUMN], + row[PAIR_RIGHT_ID_COLUMN], + None, + ) + for _index, row in candidates.iterrows() + ] + + monkeypatch.setattr(sem_join_module, "verify_semantic_join_candidates", verify) + memory = _SemanticSelfJoinMemory(refresh=am.CountRefresh(every=2)) + + memory.add({"key": "all", "value": "one"}) + memory.add({"key": "all", "value": "two"}) + + assert candidate_sizes == [4] + assert len(memory._runtime._state["pairs"]) == 4 + + +def test_relation_delta_completes_all_ready_count_windows() -> None: + memory = _WindowMemory(refresh=am.CountRefresh(every=3)) + + for message in ("one", "two", "three"): + memory.add({"message": message}) + + blocks = memory._runtime._state["blocks"] + assert len(blocks) == 2 + assert json.loads(blocks.loc[0, "messages"]) == [ + {"message": "one"}, + {"message": "two"}, + ] + assert json.loads(blocks.loc[1, "messages"]) == [ + {"message": "two"}, + {"message": "three"}, + ] + + +@pytest.mark.parametrize("every", [0, -1]) +def test_count_refresh_rejects_non_positive_thresholds(every: int) -> None: + with pytest.raises(ValueError, match="at least 1"): + am.CountRefresh(every=every) + + +@pytest.mark.parametrize("every", [True, 1.5, "2"]) +def test_count_refresh_rejects_non_integer_thresholds(every: Any) -> None: + with pytest.raises(TypeError, match="must be an integer"): + am.CountRefresh(every=every) diff --git a/tests/test_dataflow_recovery_hook.py b/tests/test_dataflow_recovery_hook.py new file mode 100644 index 0000000..1a09d6f --- /dev/null +++ b/tests/test_dataflow_recovery_hook.py @@ -0,0 +1,85 @@ +"""Offline tests of public recovery and deterministic request identities.""" + +from typing import Any +import pandas as pd +import pytest +import agent_memory as am +from agent_memory.adapters.lotus.request_hook import next_request_batch, request_scope + + +def test_dataflow_snapshot_continues_count() -> None: + source = am.Source({"x": "number"}) + flow = am.SemanticDataflow(source=source, views={"v": source.count(output_col="n")}) + flow.apply(pd.DataFrame({"x": [1, 2]})) + state = flow.snapshot_state() + restored = am.SemanticDataflow( + source=source, views={"v": source.count(output_col="n")} + ) + restored.restore_state(state) + restored.apply(pd.DataFrame({"x": [3]})) + assert restored.view("v").iloc[0]["n"] == 3 + wrong = am.SemanticDataflow(source=source, views={"v": source.limit(1)}) + with pytest.raises(ValueError, match="fingerprint"): + wrong.restore_state(state) + + +def test_request_positions_reset_per_unit() -> None: + class Hook: + def execute(self, identity: Any, payload: Any, send: Any) -> dict[str, Any]: + return send() + + hook = Hook() + for _ in range(2): + with request_scope(hook, mode="full", event_id="1"): + assert next_request_batch() == ( + hook, + {"mode": "full", "event_id": "1", "batch": 0}, + ) + assert next_request_batch() == ( + hook, + {"mode": "full", "event_id": "1", "batch": 1}, + ) + assert next_request_batch() is None + + +def test_hook_refuses_unjournaled_text_fallback() -> None: + from agent_memory.adapters.lotus.deepseek_responses_lm import ( + deepseek_responses_lm_class, + ) + + class Base: + def __init__(self, *args: Any, **kwargs: Any) -> None: + pass + + def _process_uncached_messages(self, *args: Any) -> Any: + pytest.fail("must not fall through to unjournaled provider") + + class Hook: + def execute(self, *args: Any) -> Any: + pytest.fail("unsupported request must be rejected before send") + + lm = deepseek_responses_lm_class(Base)("deepseek/deepseek-flash") + with request_scope(Hook(), mode="full"): + with pytest.raises(ValueError, match="no fallback"): + lm._process_uncached_messages([], {}, False, "") + + +def test_compact_failure_keeps_no_per_row_files() -> None: + from agent_memory.adapters.lotus.structured import ( + write_structured_failure_artifacts, + ) + from agent_memory.tracing.semantic import semantic_trace_scope + + with semantic_trace_scope(provider_journal=True): + assert ( + write_structured_failure_artifacts( + ["prompt"] * 1000, + [["invalid"]] * 1000, + list(range(1000)), + output_cols=(), + shape="object", + require_explanation=False, + operator="sem_join", + ) + == [] + ) diff --git a/tests/test_deepseek_responses_lm.py b/tests/test_deepseek_responses_lm.py new file mode 100644 index 0000000..44ff427 --- /dev/null +++ b/tests/test_deepseek_responses_lm.py @@ -0,0 +1,899 @@ +"""Offline contract tests for the native DeepSeek Responses LOTUS adapter.""" + +from __future__ import annotations + +from collections.abc import Callable +from copy import deepcopy +import json +from threading import Event, Lock +from typing import Any + +import httpx +import lotus +from litellm import ModelResponse +from lotus.cache import InMemoryCache +from lotus.models import LM +from openai import APIConnectionError, OpenAI, OpenAIError +from openai.types.responses import Response +import pytest + +from agent_memory.adapters.lotus import deepseek_responses_lm as adapter + + +SCHEMA = { + "type": "json_schema", + "json_schema": { + "name": "answer", + "schema": { + "type": "object", + "properties": {"answer": {"type": "integer"}}, + "required": ["answer"], + "additionalProperties": False, + }, + "strict": True, + }, +} +MODEL = "deepseek/deepseek-v4-flash" + + +def _native(text: str = '{"answer":1}', **overrides: Any) -> Response: + data = { + "id": "resp-test", + "object": "response", + "created_at": 123, + "model": "deepseek-v4-flash-actual", + "status": "completed", + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "output": [ + { + "id": "msg-test", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ], + "parallel_tool_calls": True, + "temperature": 0, + "tool_choice": "auto", + "tools": [], + "top_p": 1, + "usage": { + "input_tokens": 12, + "input_tokens_details": {"cached_tokens": 8}, + "output_tokens": 7, + "output_tokens_details": {"reasoning_tokens": 3}, + "total_tokens": 19, + }, + } + data.update(overrides) + return Response.model_validate(data) + + +class FakeClient: + """Record SDK boundary calls without making network requests.""" + + def __init__(self, **options: Any) -> None: + self.options = options + self.responses = self + self.calls: list[dict[str, Any]] = [] + self.handler: Callable[[dict[str, Any]], Response] = lambda _: _native() + self.closed = False + + def create(self, **payload: Any) -> Response: + """Return the configured native response or raise its failure.""" + self.calls.append(deepcopy(payload)) + return self.handler(payload) + + def __enter__(self) -> FakeClient: + return self + + def __exit__(self, *args: Any) -> None: + self.closed = True + + +@pytest.fixture +def clients(monkeypatch: pytest.MonkeyPatch) -> list[FakeClient]: + made: list[FakeClient] = [] + + def create(**options: Any) -> FakeClient: + client = FakeClient(**options) + made.append(client) + return client + + def no_network(*args: Any, **kwargs: Any) -> Any: + raise AssertionError("Tests must not make network requests") + + monkeypatch.setattr(adapter, "OpenAI", create) + monkeypatch.setattr(httpx.Client, "send", no_network) + monkeypatch.setenv("DEEPSEEK_API_KEY", "fake-deepseek-key") + monkeypatch.setattr(lotus.settings, "enable_cache", False) + monkeypatch.setattr("lotus.models.lm.completion_cost", lambda **_: 0.0) + return made + + +def _lm(*, model: str = MODEL, **kwargs: Any) -> Any: + return adapter.deepseek_responses_lm_class(LM)( + model=model, cache=InMemoryCache(max_size=64), **kwargs + ) + + +def _messages(count: int = 1) -> list[list[dict[str, str]]]: + return [[{"role": "user", "content": str(i)}] for i in range(count)] + + +def _process(lm: Any, count: int = 1, **kwargs: Any) -> list[Any]: + return lm._process_uncached_messages( + [(messages, str(i)) for i, messages in enumerate(_messages(count))], + {**lm.kwargs, "response_format": SCHEMA, **kwargs}, + False, + "Test responses", + ) + + +@pytest.mark.parametrize( + "model", + [ + "deepseek-flash", + "deepseek/deepseek-flash", + "deepseek-v4-flash", + "deepseek/deepseek-v4-flash", + "deepseek-v4-pro", + "deepseek/deepseek-v4-pro", + ], +) +def test_supported_model_names(model: str) -> None: + assert adapter.validate_deepseek_responses_model(model) is None + + +@pytest.mark.parametrize( + "model", + [ + "openai/deepseek-flash", + "openrouter/deepseek/deepseek-flash", + "openai/deepseek-v4-flash", + "openrouter/deepseek/deepseek-v4-flash", + "deepseek/deepseek-chat", + "deepseek-reasoner", + "gpt-4o", + "", + "deepseek-v5", + ], +) +def test_invalid_model_rejected_before_parent_resources(model: str) -> None: + class ResourceBase: + def __init__(self, **kwargs: Any) -> None: + pytest.fail("Parent resources initialized for unsupported model") + + with pytest.raises(ValueError, match="model"): + adapter.validate_deepseek_responses_model(model) + with pytest.raises(ValueError, match="model"): + adapter.deepseek_responses_lm_class(ResourceBase)(model=model) + + +@pytest.mark.parametrize("name", ["deepseek-flash", "deepseek-v4-flash", "deepseek-v4-pro"]) +@pytest.mark.parametrize("prefix", ["", "deepseek/"]) +def test_schema_payload_and_nonthinking_translation( + clients: list[FakeClient], name: str, prefix: str +) -> None: + schema = deepcopy(SCHEMA) + kwargs = { + "max_tokens": 1000, + "temperature": 0.3, + "top_p": 0.8, + "timeout": 9.5, + "num_retries": 4, + "extra_body": {"thinking": {"type": "disabled"}}, + "stream": False, + } + lm = _lm(model=f"{prefix}{name}", **kwargs) + messages = [ + [ + {"role": "system", "content": "Answer with JSON."}, + {"role": "user", "content": "one"}, + {"role": "assistant", "content": '{"answer":0}'}, + {"role": "user", "content": "two"}, + ] + ] + result = lm(messages, response_format=schema, show_progress_bar=False) + + assert result.outputs == ['{"answer":1}'] + assert clients[0].options == { + "base_url": "https://api.deepseek.com", + "api_key": "fake-deepseek-key", + "max_retries": 4, + "timeout": 9.5, + } + assert clients[0].calls == [ + { + "model": name, + "input": messages[0], + "text": {"format": {"type": "json_schema", **SCHEMA["json_schema"]}}, + "reasoning": {"effort": "none"}, + "max_output_tokens": 1000, + "temperature": 0.3, + "top_p": 0.8, + "stream": False, + } + ] + assert clients[0].closed + assert schema == SCHEMA + assert lm.kwargs == kwargs + + +@pytest.mark.parametrize( + ("thinking", "effort", "expected"), + [ + (None, None, "high"), + ("enabled", None, "high"), + ("enabled", "low", "low"), + ("enabled", "medium", "high"), + ("enabled", "high", "high"), + ("enabled", "xhigh", "high"), + ("enabled", "max", "max"), + ("disabled", None, "none"), + (None, "none", "none"), + ], +) +def test_explicit_reasoning_mapping( + clients: list[FakeClient], thinking: str | None, effort: str | None, expected: str +) -> None: + kwargs: dict[str, Any] = {} + if thinking is not None: + kwargs["extra_body"] = {"thinking": {"type": thinking}} + if effort is not None: + kwargs["reasoning_effort"] = effort + _process(_lm(**kwargs)) + assert clients[0].calls[0]["reasoning"] == {"effort": expected} + + +@pytest.mark.parametrize( + "kwargs", + [ + {"tools": []}, + {"tool_choice": "none"}, + {"stream": True}, + {"stream": 1}, + {"logprobs": True}, + {"top_logprobs": 3}, + {"n": 2}, + {"seed": 2}, + {"stop": ["END"]}, + {"frequency_penalty": 0}, + {"made_up": True}, + {"extra_body": {"other": 1}}, + {"extra_body": {"thinking": {"type": "maybe"}}}, + {"extra_body": {"thinking": {"type": "enabled", "budget_tokens": 100}}}, + {"extra_body": {"thinking": {"type": "disabled"}}, "reasoning_effort": "high"}, + {"extra_body": {"thinking": {"type": "enabled"}}, "reasoning_effort": "none"}, + {"reasoning_effort": "minimal"}, + {"num_retries": -1}, + {"num_retries": 11}, + {"num_retries": 1.5}, + {"num_retries": True}, + {"max_tokens": 0}, + {"max_tokens": True}, + {"temperature": 3}, + {"top_p": -0.1}, + {"timeout": 0}, + {"timeout": float("inf")}, + { + "response_format": { + "type": "json_schema", + "json_schema": {"name": "missing"}, + } + }, + ], +) +def test_unsupported_parameters_fail_before_client_creation( + clients: list[FakeClient], kwargs: dict[str, Any] +) -> None: + with pytest.raises((TypeError, ValueError)): + _process(_lm(), **kwargs) + assert clients == [] + + +@pytest.mark.parametrize( + "message", + [ + {"role": "tool", "content": "x"}, + {"role": "user", "content": [{"type": "image_url", "image_url": "x"}]}, + {"role": "assistant", "content": "x", "tool_calls": []}, + ], +) +def test_entire_input_is_validated_before_any_request( + clients: list[FakeClient], message: dict[str, Any] +) -> None: + with pytest.raises(ValueError, match="text-only"): + _lm()( + _messages() + [[message]], response_format=SCHEMA, show_progress_bar=False + ) + assert clients == [] + + +def test_client_overrides_and_zero_retries(clients: list[FakeClient]) -> None: + _process( + _lm(), + api_key="override-key", + api_base="https://gateway.example/v1", + num_retries=0, + ) + assert clients[0].options["api_key"] == "override-key" + assert clients[0].options["base_url"] == "https://gateway.example/v1" + assert clients[0].options["max_retries"] == 0 + assert "api_key" not in clients[0].calls[0] + + +def test_missing_deepseek_key_does_not_use_openai_key( + clients: list[FakeClient], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DEEPSEEK_API_KEY") + monkeypatch.setenv("OPENAI_API_KEY", "wrong-provider-key") + with pytest.raises(ValueError, match="DEEPSEEK_API_KEY"): + _process(_lm()) + assert clients == [] + + +def test_normalizes_usage_actual_model_and_metadata(clients: list[FakeClient]) -> None: + response = _process(_lm())[0] + assert isinstance(response, ModelResponse) + assert response.model == "deepseek-v4-flash-actual" + assert response.id == "resp-test" + assert response.created == 123 + assert response.choices[0].message.content == '{"answer":1}' + assert response.choices[0].finish_reason == "stop" + assert response.usage.prompt_tokens == 12 + assert response.usage.completion_tokens == 7 + assert response.usage.total_tokens == 19 + assert response.usage.prompt_cache_hit_tokens == 8 + assert response.usage.prompt_cache_miss_tokens == 4 + assert response.usage.prompt_tokens_details.cached_tokens == 8 + assert response.usage.completion_tokens_details.reasoning_tokens == 3 + assert response.provider_response_metadata == { + "transport": "responses-json-schema", + "status": "completed", + "incomplete_reason": None, + "finish_reason": "stop", + "raw_usage": _native().usage.model_dump(mode="json"), + } + assert ( + json.loads(response.model_dump_json())["provider_response_metadata"] + == response.provider_response_metadata + ) + + +def test_concurrency_is_bounded_and_output_order_is_preserved( + clients: list[FakeClient], monkeypatch: pytest.MonkeyPatch +) -> None: + client = FakeClient() + second_started = Event() + lock = Lock() + active = 0 + peak = 0 + finished: list[int] = [] + + def handler(payload: dict[str, Any]) -> Response: + nonlocal active, peak + index = int(payload["input"][0]["content"]) + with lock: + active += 1 + peak = max(peak, active) + if index == 0: + assert second_started.wait(timeout=5) + if index == 1: + second_started.set() + with lock: + finished.append(index) + active -= 1 + return _native(str(index)) + + client.handler = handler + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + results = _process(_lm(max_batch_size=2), count=6) + assert [r.choices[0].message.content for r in results] == [ + "0", + "1", + "2", + "3", + "4", + "5", + ] + assert peak == 2 + assert sorted(finished) == list(range(6)) + + +@pytest.mark.parametrize( + ("elapsed", "expected_sleeps"), [(0.25, [1.75, 1.75]), (3, [])] +) +def test_rate_limit_wait_matches_lotus_batches( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, + elapsed: float, + expected_sleeps: list[float], +) -> None: + ticks = iter([0, elapsed, 10, 10 + elapsed, 20, 20 + elapsed]) + sleeps: list[float] = [] + monkeypatch.setattr(adapter.time, "monotonic", lambda: next(ticks)) + monkeypatch.setattr(adapter.time, "sleep", sleeps.append) + results = _process(_lm(max_batch_size=2, rate_limit=60), count=5) + assert len(results) == 5 + assert sleeps == expected_sleeps + + +def test_rate_limit_caps_batch_size(clients: list[FakeClient]) -> None: + lm = _lm(max_batch_size=10, rate_limit=2) + assert lm.max_batch_size == 2 + + +@pytest.mark.parametrize( + ("overrides", "finish_reason"), + [ + ( + { + "status": "incomplete", + "incomplete_details": {"reason": "max_output_tokens"}, + }, + "length", + ), + ( + { + "status": "incomplete", + "incomplete_details": {"reason": "content_filter"}, + }, + "content_filter", + ), + ({"status": "incomplete"}, "incomplete"), + ( + { + "status": "failed", + "error": {"code": "server_error", "message": "failed"}, + }, + "error", + ), + ({"status": "in_progress"}, "in_progress"), + ({"output": []}, "empty_output"), + ( + { + "output": [ + { + "id": "m", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "refusal", "refusal": "No."}], + } + ] + }, + "content_filter", + ), + ], +) +def test_non_success_is_an_exception_with_usage_not_empty_success( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, + overrides: dict[str, Any], + finish_reason: str, +) -> None: + client = FakeClient() + client.handler = lambda _: _native(**overrides) + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + response = _process(_lm())[0] + assert isinstance(response, OpenAIError) + assert response.usage.total_tokens == 19 + assert response.model == "deepseek-v4-flash-actual" + assert response.provider_response_metadata["finish_reason"] == finish_reason + assert response.provider_response_metadata["status"] == overrides.get( + "status", "completed" + ) + assert response.provider_response_metadata["raw_usage"]["input_tokens"] == 12 + lm = _lm() + with pytest.raises(OpenAIError): + lm(_messages(), response_format=SCHEMA, show_progress_bar=False) + assert lm.stats.physical_usage.total_tokens == 19 + + +@pytest.mark.parametrize("cache_enabled", [False, True]) +@pytest.mark.parametrize("failure_index", [0, 1]) +@pytest.mark.parametrize("native_failure", [False, True]) +def test_partial_failures_preserve_all_usage_and_successful_cache_entries( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, + cache_enabled: bool, + failure_index: int, + native_failure: bool, +) -> None: + client = FakeClient() + network_error = APIConnectionError( + request=httpx.Request("POST", "https://api.deepseek.com/responses") + ) + + def handler(payload: dict[str, Any]) -> Response: + index = int(payload["input"][0]["content"]) + if index == failure_index: + if native_failure: + return _native( + status="incomplete", + incomplete_details={"reason": "max_output_tokens"}, + ) + raise network_error + return _native(str(index)) + + client.handler = handler + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + monkeypatch.setattr(lotus.settings, "enable_cache", cache_enabled) + lm = _lm(num_retries=2) + results = _process(lm, count=3) + assert isinstance(results[failure_index], OpenAIError) + assert len(results) == 3 + assert len(client.calls) == 3 # No outer retry loop. + assert ( + results[failure_index].provider_response_metadata["transport"] + == "responses-json-schema" + ) + assert lm.stats.physical_usage.total_tokens == 0 # Parent owns accounting. + client.calls.clear() + with pytest.raises(OpenAIError) as caught: + lm(_messages(3), response_format=SCHEMA, show_progress_bar=False) + if not native_failure: + assert caught.value is network_error + assert len(client.calls) == 3 + assert lm.stats.physical_usage.total_tokens == (57 if native_failure else 38) + assert lm.stats.virtual_usage.total_tokens == (57 if native_failure else 38) + client.handler = lambda p: _native(p["input"][0]["content"]) + client.calls.clear() + assert lm( + _messages(3), response_format=SCHEMA, show_progress_bar=False + ).outputs == ["0", "1", "2"] + assert len(client.calls) == (1 if cache_enabled else 3) + assert lm.stats.cache_hits == (2 if cache_enabled else 0) + + +@pytest.mark.parametrize( + "response_format", [None, {"type": "json_object"}, {"type": "text"}] +) +@pytest.mark.parametrize("model", [MODEL, "deepseek/deepseek-flash"]) +def test_non_schema_delegates_native_chat_without_sdk_client( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, + response_format: Any, + model: str, +) -> None: + calls: list[dict[str, Any]] = [] + + def chat(model: str, batch: Any, **kwargs: Any) -> list[ModelResponse]: + calls.append({"model": model, "batch": batch, **kwargs}) + return [ + ModelResponse( + choices=[{"message": {"role": "assistant", "content": "chat"}}] + ) + ] + + monkeypatch.setattr("lotus.models.lm.batch_completion", chat) + monkeypatch.delenv("DEEPSEEK_API_KEY") + result = _lm(model=model)( + _messages(), + response_format=response_format, + stop=["END"], + show_progress_bar=False, + ) + assert result.outputs == ["chat"] + assert clients == [] + assert calls[0]["response_format"] == response_format + assert calls[0]["stop"] == ["END"] + assert calls[0]["model"] == model + + +def test_cache_hits_keep_virtual_usage_without_new_client( + clients: list[FakeClient], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(lotus.settings, "enable_cache", True) + lm = _lm() + assert lm(_messages(), response_format=SCHEMA, show_progress_bar=False).outputs == [ + '{"answer":1}' + ] + assert lm(_messages(), response_format=SCHEMA, show_progress_bar=False).outputs == [ + '{"answer":1}' + ] + assert len(clients) == 1 + assert lm.stats.physical_usage.total_tokens == 19 + assert lm.stats.virtual_usage.total_tokens == 38 + assert lm.stats.cache_hits == 1 + altered_schema = deepcopy(SCHEMA) + altered_schema["json_schema"]["name"] = "different" + lm(_messages(), response_format=altered_schema, show_progress_bar=False) + assert len(clients) == 2 + + +def test_empty_batch_does_not_initialize_client(clients: list[FakeClient]) -> None: + assert _process(_lm(), count=0) == [] + assert clients == [] + + +def test_schema_without_optional_fields_and_default_retry_budget( + clients: list[FakeClient], +) -> None: + schema = { + "type": "json_schema", + "json_schema": {"name": "answer", "schema": {"type": "object"}}, + } + _process(_lm(), response_format=schema) + assert clients[0].calls[0]["text"] == { + "format": { + "type": "json_schema", + "name": "answer", + "schema": {"type": "object"}, + } + } + assert clients[0].options["max_retries"] == 2 + + +def test_reasoning_is_not_mixed_into_text_and_usage_extensions_survive( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, +) -> None: + raw = _native().model_dump(mode="json") + raw["output"].insert( + 0, + { + "id": "r", + "type": "reasoning", + "summary": [], + "content": [{"type": "reasoning_text", "text": "private reasoning"}], + }, + ) + raw["output"][1]["content"] = [ + {"type": "output_text", "text": '{"answer":', "annotations": []}, + {"type": "output_text", "text": "1}", "annotations": []}, + ] + raw["usage"]["provider_counter"] = 42 + raw["usage"]["output_tokens_details"]["reasoning_tokens"] = 0 + client = FakeClient() + client.handler = lambda _: Response.model_validate(raw) + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + response = _process(_lm())[0] + assert response.choices[0].message.content == '{"answer":1}' + assert response.usage.completion_tokens_details.reasoning_tokens == 0 + assert response.provider_response_metadata["raw_usage"]["provider_counter"] == 42 + + +def test_missing_usage_is_not_fabricated( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, +) -> None: + client = FakeClient() + client.handler = lambda _: _native(usage=None, status="failed") + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + response = _process(_lm())[0] + assert isinstance(response, OpenAIError) + assert getattr(response, "usage", None) is None + assert response.provider_response_metadata["raw_usage"] is None + lm = _lm() + with pytest.raises(OpenAIError): + lm(_messages(), response_format=SCHEMA, show_progress_bar=False) + assert lm.stats.physical_usage.total_tokens == 0 + + +@pytest.mark.parametrize("cache_enabled", [False, True]) +@pytest.mark.parametrize("bad_field", ["created_at", "usage"]) +def test_malformed_accounting_preserves_order_and_other_usage( + clients: list[FakeClient], monkeypatch: pytest.MonkeyPatch, + bad_field: str, cache_enabled: bool, +) -> None: + native = _native().model_copy(update={ + bad_field: None if bad_field == "created_at" else {"invalid": "usage"}, + }) + error = ValueError("original normalization failure") + original_normalizer = adapter._normalize_response + + def normalize(response: Response) -> Any: + if response is native: + raise error + return original_normalizer(response) + + client = FakeClient() + client.handler = lambda p: native if p["input"][0]["content"] == "1" else _native() + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + monkeypatch.setattr(adapter, "_normalize_response", normalize) + monkeypatch.setattr(lotus.settings, "enable_cache", cache_enabled) + lm = _lm() + responses = _process(lm, count=3) + assert isinstance(responses[0], ModelResponse) + assert isinstance(responses[2], ModelResponse) + failure = responses[1] + assert isinstance(failure, OpenAIError) + assert failure.native_response is native + assert failure.__cause__ is error + assert any("accounting" in note for note in error.__notes__) + assert failure.provider_response_metadata["finish_reason"] == "normalization_error" + assert failure.provider_response_metadata["raw_usage"] is None + assert getattr(failure, "usage", None) is None + with pytest.raises(OpenAIError, match="original normalization failure"): + lm(_messages(3), response_format=SCHEMA, show_progress_bar=False) + assert lm.stats.physical_usage.total_tokens == 38 + assert lm.stats.virtual_usage.total_tokens == 38 + + +@pytest.mark.parametrize("cache_enabled", [False, True]) +@pytest.mark.parametrize("usage_available", [False, True]) +@pytest.mark.parametrize("failure_stage", ["normalizer", "output_serialization"]) +def test_normalizer_failure_preserves_native_response_and_batch_usage( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, + cache_enabled: bool, + usage_available: bool, + failure_stage: str, +) -> None: + native = _native( + id="resp-normalization-failure", + status="incomplete", + incomplete_details={"reason": "max_output_tokens"}, + usage=_native().usage if usage_available else None, + ) + client = FakeClient() + client.handler = lambda p: native if p["input"][0]["content"] == "0" else _native() + original_normalizer = adapter._normalize_response + normalization_error = ValueError("Unsupported native item") + + def normalize(response: Response) -> Any: + if response is native: + raise normalization_error + return original_normalizer(response) + + if failure_stage == "normalizer": + monkeypatch.setattr(adapter, "_normalize_response", normalize) + else: + original_dump = Response.model_dump + + def dump(response: Response, *args: Any, **kwargs: Any) -> dict[str, Any]: + if response is native: + raise normalization_error + return original_dump(response, *args, **kwargs) + + monkeypatch.setattr(Response, "model_dump", dump) + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + monkeypatch.setattr(lotus.settings, "enable_cache", cache_enabled) + lm = _lm() + responses = _process(lm, count=3) + failure = responses[0] + assert isinstance(failure, OpenAIError) + assert failure.native_response is native + assert failure.__cause__ is normalization_error + assert failure.model == "deepseek-v4-flash-actual" + assert isinstance(failure.response, ModelResponse) + assert failure.provider_response_metadata == { + "transport": "responses-json-schema", + "status": "incomplete", + "incomplete_reason": "max_output_tokens", + "finish_reason": "normalization_error", + "raw_usage": native.usage.model_dump(mode="json") if usage_available else None, + } + if usage_available: + assert failure.usage.total_tokens == 19 + assert failure.usage.prompt_cache_hit_tokens == 8 + assert failure.usage.prompt_cache_miss_tokens == 4 + assert failure.usage.completion_tokens_details.reasoning_tokens == 3 + else: + assert not hasattr(failure, "usage") + assert lm.stats.physical_usage.total_tokens == 0 + client.calls.clear() + with pytest.raises(OpenAIError, match="Unsupported native item"): + lm(_messages(3), response_format=SCHEMA, show_progress_bar=False) + assert len(client.calls) == 3 + assert lm.stats.physical_usage.total_tokens == (57 if usage_available else 38) + assert lm.stats.virtual_usage.total_tokens == (57 if usage_available else 38) + client.calls.clear() + client.handler = lambda _: _native() + assert ( + len(lm(_messages(3), response_format=SCHEMA, show_progress_bar=False).outputs) + == 3 + ) + assert len(client.calls) == (1 if cache_enabled else 3) + + +def test_unsupported_native_output_keeps_original_response( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, +) -> None: + native = _native( + output=[ + { + "type": "function_call", + "id": "fc-test", + "call_id": "call-test", + "name": "unsupported", + "arguments": "{}", + "status": "completed", + } + ] + ) + client = FakeClient() + client.handler = lambda _: native + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + failure = _process(_lm())[0] + assert isinstance(failure, OpenAIError) + assert failure.native_response is native + assert failure.usage.total_tokens == 19 + assert failure.provider_response_metadata["finish_reason"] == "unsupported_output" + + +def test_outer_hook_receives_ordered_failure_metadata_before_parent_accounting( + clients: list[FakeClient], + monkeypatch: pytest.MonkeyPatch, +) -> None: + seen: list[Any] = [] + + class OuterHook(adapter.deepseek_responses_lm_class(LM)): + def _process_uncached_messages(self, *args: Any, **kwargs: Any) -> list[Any]: + responses = super()._process_uncached_messages(*args, **kwargs) + assert self.stats.physical_usage.total_tokens == 0 + seen.extend(responses) + return responses + + client = FakeClient() + client.handler = lambda p: _native( + p["input"][0]["content"], + status="incomplete" if p["input"][0]["content"] == "0" else "completed", + incomplete_details={"reason": "max_output_tokens"} + if p["input"][0]["content"] == "0" + else None, + ) + monkeypatch.setattr(adapter, "OpenAI", lambda **_: client) + monkeypatch.setattr(lotus.settings, "enable_cache", True) + lm = OuterHook(model=MODEL, cache=InMemoryCache(max_size=8)) + with pytest.raises(OpenAIError): + lm(_messages(2), response_format=SCHEMA, show_progress_bar=False) + assert [r.provider_response_metadata["finish_reason"] for r in seen] == [ + "length", + "stop", + ] + assert ( + seen[0].provider_response_metadata["incomplete_reason"] == "max_output_tokens" + ) + assert lm.stats.physical_usage.total_tokens == 38 + + +@pytest.mark.parametrize("always_fail", [False, True]) +def test_installed_sdk_serializes_responses_and_owns_bounded_retries( + monkeypatch: pytest.MonkeyPatch, + always_fail: bool, +) -> None: + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + if always_fail or len(requests) == 1: + return httpx.Response( + 429, + headers={"retry-after-ms": "1"}, + json={"error": {"message": "retry"}}, + ) + return httpx.Response(200, json=_native().model_dump(mode="json")) + + def sdk_client(**kwargs: Any) -> OpenAI: + return OpenAI( + **kwargs, http_client=httpx.Client(transport=httpx.MockTransport(handler)) + ) + + monkeypatch.setattr(adapter, "OpenAI", sdk_client) + monkeypatch.setenv("DEEPSEEK_API_KEY", "fake-deepseek-key") + result = _process( + _lm(num_retries=1, extra_body={"thinking": {"type": "disabled"}}) + )[0] + assert len(requests) == 2 + assert str(requests[0].url) == "https://api.deepseek.com/responses" + assert requests[0].headers["authorization"] == "Bearer fake-deepseek-key" + payload = json.loads(requests[0].content) + assert payload["reasoning"] == {"effort": "none"} + assert payload["max_output_tokens"] == 512 + assert payload["text"]["format"]["schema"] == SCHEMA["json_schema"]["schema"] + assert "max_tokens" not in payload + assert "num_retries" not in payload + if always_fail: + assert isinstance(result, OpenAIError) + assert result.provider_response_metadata["status"] == "failed" + else: + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == '{"answer":1}' diff --git a/tests/test_differential_runtime.py b/tests/test_differential_runtime.py index ea2eed1..ec8f17c 100644 --- a/tests/test_differential_runtime.py +++ b/tests/test_differential_runtime.py @@ -11,6 +11,7 @@ import agent_memory.planner as planner from agent_memory.adapters import LotusAdapter from agent_memory.adapters.lotus.sem_flat_map import apply_flat_map_outputs +from agent_memory.adapters.lotus.sem_join import assemble_join_frame from agent_memory.policy.aggregates import ( ArrayAggregateSpec, CollectListAggregateSpec, @@ -27,6 +28,45 @@ from agent_memory.runtime.legacy import LegacyViewRuntime +@pytest.fixture(autouse=True) +def _block_provider_configuration(monkeypatch: pytest.MonkeyPatch) -> None: + """Runtime tests must stub every semantic operator before provider setup.""" + + from agent_memory.adapters.lotus.context import LotusExecutionContext + + def reject_provider(self: LotusExecutionContext) -> None: + raise AssertionError("Runtime test reached an unstubbed semantic operator") + + monkeypatch.setattr(LotusExecutionContext, "configure", reject_provider) + + +def _join_stub_keys( + query: am.QueryExpr, + left: pd.DataFrame, + right: pd.DataFrame, + keys: tuple[str, ...], +) -> pd.DataFrame: + """Use fixture equality as the semantic oracle, retaining join assembly.""" + + keys = (*tuple(query.params.get("on", ())), *keys) + matches = [] + for left_index, left_row in left.iterrows(): + selected = [ + (left_index, right_index, None) + for right_index, right_row in right.iterrows() + if all(left_row[key] == right_row[key] for key in keys) + ] + k = query.params.get("k") + matches.extend(selected if k is None else selected[:k]) + return assemble_join_frame( + left, + right, + matches, + how=str(query.params["how"]), + id_columns=tuple(query.params["id_columns"]), + ) + + def test_planner_exports_differentiator_components_not_free_functions() -> None: assert planner.QueryDifferentiator.__name__ == "QueryDifferentiator" assert planner.PolicyDifferentiator.__name__ == "PolicyDifferentiator" @@ -141,17 +181,16 @@ class GroupedMemory(am.Memory): assert not any(node.query.op == "sem_groupby" for node in plan.nodes.values()) -def test_policy_compiler_rejects_view_time_semantic_join() -> None: - class UnsupportedMemory(am.Memory): +def test_policy_compiler_supports_append_only_inner_semantic_join() -> None: + class SemanticJoinMemory(am.Memory): log = am.Log({"message": "Message."}) joined = log.sem_join(log, instruction="{message:left} matches {message:right}.") - try: - UnsupportedMemory.differentiate_policy() - except NotImplementedError as error: - assert "view-time sem_join" in str(error) - else: - raise AssertionError("view-time sem_join should be rejected") + plan = SemanticJoinMemory.differentiate_policy() + node = next(node for node in plan.nodes.values() if node.query.op == "sem_join") + + assert node.execution_kind == "semantic_binary_state" + assert node.maintenance_query is not None def test_policy_compiler_rejects_view_time_semantic_topk() -> None: @@ -277,6 +316,13 @@ def execute( return source if query.op == "sem_agg": return self._execute_sem_agg(query, inputs) + if query.op == "sem_join": + return _join_stub_keys( + query, + self.execute(query.inputs[0], inputs), + self.execute(query.inputs[1], inputs), + ("key",), + ) return super().execute(query, inputs) finally: self._depth -= 1 @@ -417,6 +463,10 @@ def test_semantic_state_uses_maintenance_query_for_insert_only_change() -> None: assert any(query is node.maintenance_query for query in adapter.root_queries) assert not any(query is node.query for query in adapter.root_queries) + assert memory._runtime._state["summaries"].to_dict("records") == [ + {"key": "a", "summary": "5", "propagated": True}, + {"key": "b", "summary": "7", "propagated": True}, + ] def test_semantic_state_recomputes_node_on_parent_replacement_and_propagates() -> None: @@ -744,9 +794,27 @@ def execute( source = self.execute(query.inputs[0], inputs).copy() source[GROUP_ID_COLUMN] = pd.factorize(source["name"], sort=False)[0] return source + if query.op == "sem_join": + left = self.execute(query.inputs[0], inputs) + right = self.execute(query.inputs[1], inputs) + matches = [ + (left_index, right_index, None) + for left_index, left_row in left.iterrows() + for right_index, right_row in right.iterrows() + if left_row["name"] == right_row["name"] + ] + return assemble_join_frame( + left, + right, + matches, + how=str(query.params["how"]), + id_columns=tuple(query.params["id_columns"]), + ) if query.op == "sem_agg": source = self.execute(query.inputs[0], inputs) output_names = [column.name for column in query.params["output_cols"]] + if source.empty: + return pd.DataFrame(columns=output_names) groups = ( source.groupby(GROUP_ID_COLUMN, sort=False, dropna=False) if GROUP_ID_COLUMN in source.columns @@ -814,6 +882,16 @@ def execute( ) -> pd.DataFrame: if query.op == "sem_flat_map": return self._flat_map(query, inputs) + if query.op == "sem_join": + left = self.execute(query.inputs[0], inputs) + right = self.execute(query.inputs[1], inputs) + if "fact" in left.columns: + keys = ("relation_type", "fact") + elif "members" in left.columns: + keys = ("name", "summary") + else: + keys = ("name",) + return _join_stub_keys(query, left, right, keys) if query.op == "sem_groupby": source = self.execute(query.inputs[0], inputs).copy() grouping_columns = ( diff --git a/tests/test_interface_scaffold.py b/tests/test_interface_scaffold.py index 010297b..48765f0 100644 --- a/tests/test_interface_scaffold.py +++ b/tests/test_interface_scaffold.py @@ -1831,7 +1831,7 @@ class TopicSummaryMemory(am.Memory): assert differentiated.params["columns"] == ("topic", "summary") final_aggregate = differentiated.inputs[0] assert final_aggregate.op == "sem_agg" - assert final_aggregate.params["input_cols"] is None + assert final_aggregate.params["input_cols"] == ("summary",) assert final_aggregate.inputs[0].op == "group_by" combined = final_aggregate.inputs[0].inputs[0] assert combined.op == "union_by_name" @@ -1876,7 +1876,7 @@ class TopicEvidenceMemory(am.Memory): final_aggregate = flattened.inputs[0] assert final_aggregate.op == "agg" semantic_spec, collect_spec, min_spec = final_aggregate.params["aggregates"] - assert semantic_spec.input_cols is None + assert semantic_spec.input_cols == ("summary",) assert semantic_spec.instruction == "Summarize body." assert collect_spec.column == "evidence" assert collect_spec.output_col == "evidence" @@ -1889,6 +1889,34 @@ class TopicEvidenceMemory(am.Memory): assert combined.inputs[1].op == "agg" +def test_re_group_reaggregates_each_semantic_state_from_its_own_outputs() -> None: + class TopicStateMemory(am.Memory): + log = am.Log({"topic": "Topic.", "body": "Evidence body."}) + summaries = log.group_by("topic").agg( + am.sem_agg( + input_cols=["body"], + output_cols={"summary": "Topic summary."}, + instruction="Summarize {body} as {summary}.", + ), + am.sem_agg( + input_cols=["body"], + output_cols={"keywords": "Topic keywords."}, + instruction="Extract keywords from {body} as {keywords}.", + ), + ) + + differentiated = QueryDifferentiator( + rules=DifferentialRules(grouped_agg_rule="rule-re-group"), + ).differentiate(TopicStateMemory.spec().views["summaries"]) + + final_aggregate = differentiated.inputs[0] + summary_spec, keyword_spec = final_aggregate.params["aggregates"] + assert summary_spec.input_cols == ("summary",) + assert summary_spec.instruction == "Summarize body as {summary}." + assert keyword_spec.input_cols == ("keywords",) + assert keyword_spec.instruction == "Extract keywords from body as {keywords}." + + def test_differential_query_planner_supports_group_by_mixed_agg_join_map_rule() -> None: class TopicEvidenceMemory(am.Memory): log = am.Log( @@ -2073,10 +2101,12 @@ def test_differential_rules_support_explode_and_unnest_fragments() -> None: assert differentiated.inputs[0].inputs[0].op == "log" -def test_differential_query_planner_builds_claude_topics_full_next_view() -> None: +def test_compressed_rule_builds_claude_topics_full_next_view() -> None: view = am.ClaudeMemory.spec().views["topics"] - differentiated = QueryDifferentiator().differentiate(view) + differentiated = QueryDifferentiator( + rules=DifferentialRules(grouped_agg_rule="compressed"), + ).differentiate(view) assert differentiated.op == "select" assert differentiated.params["columns"] == ("name", "description", "type", "body") @@ -2106,18 +2136,18 @@ def test_differential_query_planner_builds_claude_topics_full_next_view() -> Non ) -def test_claude_topics_differentiated_query_no_longer_uses_join_map_merge() -> None: +def test_default_grouped_agg_rule_uses_join_map() -> None: view = am.ClaudeMemory.spec().views["topics"] differentiated = QueryDifferentiator().differentiate(view) ops = _query_ops(differentiated) assert "union_by_name" in ops - assert "sem_join" not in ops + assert "sem_join" in ops assert "sem_map" not in ops -def test_default_grouped_agg_rule_stays_compressed() -> None: +def test_explicit_compressed_grouped_agg_rule_remains_available() -> None: view = am.ClaudeMemory.spec().views["topics"] differentiated = QueryDifferentiator( @@ -6965,7 +6995,7 @@ class Model: def count_tokens(self, value: Any) -> int: text = str(value) if "doc two" in text: - return 100 + return 8 return 1 def __call__(self, batch: list[Any], **kwargs: Any) -> Any: @@ -7038,7 +7068,7 @@ class Model: max_tokens = 1 def count_tokens(self, value: Any) -> int: - return 100 if "doc two" in str(value) else 1 + return 8 if "doc two" in str(value) else 1 def __call__(self, batch: list[Any], **kwargs: Any) -> Any: calls.append((len(batch), dict(kwargs))) @@ -8474,5 +8504,7 @@ def test_runtime_query_output_columns_reject_unknown_ops() -> None: memory._runtime._query_output_columns(QueryExpr(op="unknown")) -def test_relation_is_not_top_level_public_api() -> None: - assert not hasattr(am, "Relation") +def test_relation_is_top_level_public_api() -> None: + from agent_memory.policy import Relation + + assert am.Relation is Relation diff --git a/tests/test_json_output.py b/tests/test_json_output.py new file mode 100644 index 0000000..04046e6 --- /dev/null +++ b/tests/test_json_output.py @@ -0,0 +1,204 @@ +"""Value-preserving JSON structural repair and its refusal boundaries.""" + +from __future__ import annotations + +import json +from typing import Any + +import pytest + +from agent_memory.adapters.lotus import json_output +from agent_memory.adapters.lotus.json_output import ( + load_structured_json_with_syntax_repair, + repair_json_structure, + strict_json_loads, +) + + +def _validate(raw: str) -> Any: + value = strict_json_loads(raw) + if value != {"a": ['unchanged } , : " text', 42, True, None]}: + raise ValueError("contract mismatch") + return value + + +@pytest.mark.parametrize( + "damaged", + [ + '{"a" ["unchanged } , : \\" text",42,true,null]}', + '{"a"::["unchanged } , : \\" text",42,true,null]}', + '{"a":["unchanged } , : \\" text" 42,true,null]}', + '{"a":["unchanged } , : \\" text",42,true,null]}}', + '{"a":["unchanged } , : \\" text",42,true,null', + ], +) +def test_repair_preserves_every_scalar(damaged: str) -> None: + repaired = repair_json_structure(damaged, validator=_validate) + assert repaired.value == _validate(json.dumps(repaired.value)) + assert len(repaired.edits) == 1 + + +def test_missing_result_envelope_closure() -> None: + expected = { + "results": [ + {"task_id": "task_0", "output": {"rows": []}}, + {"task_id": "task_1", "output": {"rows": []}}, + ] + } + damaged = ( + '{"results":[{"task_id":"task_0","output":{"rows":[]},' + '{"task_id":"task_1","output":{"rows":[]}}]}' + ) + + def validate(raw: str) -> Any: + value = strict_json_loads(raw) + if value != expected: + raise ValueError("task contract") + return value + + assert repair_json_structure(damaged, validator=validate).value == expected + + +@pytest.mark.parametrize( + "damaged", + [ + '{"a": "cut off', + '{"a": 1e', + '{"a": tru', + '{"a": NaN}', + '{"a":1,"a":2}', + '{"a":1,"a":2}}', + '{"a":1 "b":2 "c":3}', + '{"a":[1,2]}}}', + '{"a":1}', # Valid JSON with a bad operator contract is not syntax damage. + ], +) +def test_refuses_guessing_values_or_bad_contracts(damaged: str) -> None: + with pytest.raises(ValueError): + repair_json_structure(damaged, validator=_validate) + + +def test_rejects_ambiguous_repairs() -> None: + # Closing the nested array here or closing it at EOF gives different values. + with pytest.raises(ValueError, match="one valid candidate"): + repair_json_structure("[[1,2]", validator=strict_json_loads) + + +def test_repairs_repeated_batch_envelope_damage() -> None: + raw = ( + '{"results":[{"task_id":"a","output":{"rows":[{"body":"A"}]},' + '{"task_id":"b","output":{"rows":[{"body":"B"}]}]}' + ) + expected = { + "results": [ + {"task_id": "a", "output": {"rows": [{"body": "A"}]}}, + {"task_id": "b", "output": {"rows": [{"body": "B"}]}}, + ] + } + + def validate(value: str) -> Any: + parsed = strict_json_loads(value) + if parsed != expected: + raise ValueError("batch contract") + return parsed + + repaired = repair_json_structure(raw, validator=validate) + assert repaired.value == expected + assert len(repaired.edits) == 2 + + +@pytest.mark.parametrize("count", [3, 5, 8]) +def test_multiple_closures_preserve_all_items(count: int) -> None: + expected = { + "items": [{"id": str(i), "value": {"text": f"item {i}"}} for i in range(count)] + } + raw = ( + '{"items":[' + + ",".join(json.dumps(item)[:-1] for item in expected["items"]) + + "]}" + ) + + def validate(text: str) -> Any: + value = strict_json_loads(text) + if value != expected: + raise ValueError("item contract") + return value + + repaired = repair_json_structure(raw, validator=validate) + assert repaired.value == expected + assert len(repaired.edits) == count + assert all(kind == "insert" and token == "}" for kind, _, token in repaired.edits) + + +@pytest.mark.parametrize("raw", ['{"a":1 "b":2 "c":3}', '{"a":1}}}']) +def test_multiple_separator_edits(raw: str) -> None: + expected = {"a": 1, "b": 2, "c": 3} if "b" in raw else {"a": 1} + result = repair_json_structure(raw, validator=strict_json_loads) + assert result.value == expected + assert len(result.edits) == 2 + + +def test_multiple_edits_still_reject_ambiguity() -> None: + with pytest.raises(ValueError, match="one valid candidate"): + repair_json_structure("[[1 2]", validator=strict_json_loads) + + +def test_does_not_accept_first_candidate_before_search_finishes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(json_output, "MAX_REPAIR_CANDIDATES", 1) + monkeypatch.setattr(json_output, "_local_candidates", lambda *_: { + "[[1,2]]": ("insert", 6, "]"), + "[[1],2]": ("insert", 3, "]"), + }) + checked = [] + + def validate(raw: str) -> Any: + result = strict_json_loads(raw) + checked.append(result) + return result + + with pytest.raises(ValueError, match="bound"): + repair_json_structure("[[1,2]", validator=validate) + assert checked == [[[1, 2]]] + + +def test_edit_limit_is_not_silently_exceeded(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(json_output, "MAX_REPAIR_EDITS", 1) + with pytest.raises(ValueError): + repair_json_structure('{"a":1 "b":2 "c":3}', validator=strict_json_loads) + + +def test_large_output_work_is_bounded(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(json_output, "MAX_REPAIR_CHARACTERS", 1) + with pytest.raises(ValueError, match="bound"): + repair_json_structure('{"a":1}}', validator=strict_json_loads) + + +def test_recorded_edits_reproduce_the_accepted_text() -> None: + raw = '{"a":1 "b":2 "c":3}}' + repaired = repair_json_structure(raw, validator=strict_json_loads) + candidate = raw + for kind, position, text in repaired.edits: + if kind == "insert": + candidate = candidate[:position] + text + candidate[position:] + else: + assert candidate[position : position + len(text)] == text + candidate = candidate[:position] + candidate[position + len(text) :] + assert strict_json_loads(candidate) == repaired.value == {"a": 1, "b": 2, "c": 3} + + +@pytest.mark.parametrize( + ("raw", "method"), + [ + ('{"a":1}', None), + ("{a:1,}", "json5"), + ('```json\n{"a":1}\n```', "json5-code-fence"), + ], +) +def test_existing_json5_and_fence_contract(raw: str, method: str | None) -> None: + result = load_structured_json_with_syntax_repair( + raw, operator="test", expected_shape="object" + ) + assert result.value == {"a": 1} + assert result.repair_method == method diff --git a/tests/test_prompt_batching.py b/tests/test_prompt_batching.py new file mode 100644 index 0000000..e7266a8 --- /dev/null +++ b/tests/test_prompt_batching.py @@ -0,0 +1,599 @@ +"""Tests for the shared semantic prompt batching runner.""" + +from __future__ import annotations + +from dataclasses import dataclass +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +from agent_memory.adapters.lotus.prompt_batching import ( + ParsedPromptBatch, + PromptBatchItem, + PromptBatchRequest, + PromptBatching, + run_prompt_batches, +) + + +@dataclass(frozen=True) +class _Task: + task_id: str + text: str + + +class _Model: + max_ctx_len = 1_000 + + def __init__(self, output_batches: list[list[str]]) -> None: + self._output_batches = iter(output_batches) + self.calls: list[tuple[Any, dict[str, Any]]] = [] + + def count_tokens(self, prompt: Any) -> int: + return len(prompt[1]["content"]) + + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + self.calls.append((prompts, kwargs)) + return SimpleNamespace(outputs=next(self._output_batches)) + + +def _request(tasks: tuple[_Task, ...]) -> PromptBatchRequest: + return PromptBatchRequest( + task_ids=tuple(task.task_id for task in tasks), + prompt=[ + {"role": "system", "content": "Process each task independently."}, + {"role": "user", "content": "|".join(task.text for task in tasks)}, + ], + max_tokens=100, + ) + + +def _parse(raw_output: str) -> ParsedPromptBatch[str]: + items = tuple( + PromptBatchItem(task_id, value) + for task_id, value in ( + part.split("=", 1) for part in raw_output.split(",") if part + ) + ) + return ParsedPromptBatch(items=items) + + +def test_prompt_batching_validates_max_tasks() -> None: + assert PromptBatching().max_tasks is None + assert PromptBatching(max_tasks=4).max_tasks == 4 + with pytest.raises(TypeError, match="integer"): + PromptBatching(max_tasks=True) + with pytest.raises(ValueError, match="at least 1"): + PromptBatching(max_tasks=0) + + +def test_execution_config_accepts_one_prompt_batching_contract() -> None: + from agent_memory.adapters.lotus.context import LotusExecutionConfig + + config = LotusExecutionConfig(prompt_batching=PromptBatching(max_tasks=8)) + + assert config.prompt_batching == PromptBatching(max_tasks=8) + with pytest.raises(TypeError, match="PromptBatching"): + LotusExecutionConfig(prompt_batching="all") # type: ignore[arg-type] + + +def test_prompt_batching_changes_only_enabled_maintenance_fingerprint() -> None: + from agent_memory.adapters.lotus import LotusAdapter + from agent_memory.adapters.lotus.context import LotusExecutionConfig + + baseline = LotusAdapter(config=LotusExecutionConfig()) + explicit_disabled = LotusAdapter(config=LotusExecutionConfig(prompt_batching=None)) + all_ready = LotusAdapter( + config=LotusExecutionConfig(prompt_batching=PromptBatching()) + ) + bounded = LotusAdapter( + config=LotusExecutionConfig(prompt_batching=PromptBatching(max_tasks=8)) + ) + + assert ( + explicit_disabled.maintenance_execution_fingerprint + == baseline.maintenance_execution_fingerprint + ) + assert all_ready.maintenance_execution_fingerprint != baseline.maintenance_execution_fingerprint + assert bounded.maintenance_execution_fingerprint != all_ready.maintenance_execution_fingerprint + + +def test_all_ready_tasks_share_one_prompt_and_restore_input_order() -> None: + model = _Model([["task_2=C,task_0=A,task_1=B"]]) + tasks = tuple(_Task(f"task_{index}", value) for index, value in enumerate("ABC")) + + result = run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=_request, + parse_results=_parse, + model=model, + config=PromptBatching(), + max_retries=0, + progress_bar_desc="Testing", + operator="test", + ) + + assert result.outputs == ("A", "B", "C") + assert result.prompt_count == 1 + assert result.retry_count == 0 + assert result.chunk_sizes == (3,) + assert len(model.calls) == 1 + assert len(model.calls[0][0]) == 1 + + +def test_shared_runner_records_compact_execution_mechanics(tmp_path) -> None: + model = _Model([["task_0=A,task_1=B"]]) + tasks = (_Task("task_0", "A"), _Task("task_1", "B")) + + run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=_request, + parse_results=_parse, + model=model, + config=PromptBatching(), + max_retries=0, + progress_bar_desc="Testing", + operator="sem_filter", + trace_dir=tmp_path, + ) + + [event] = [ + json.loads(line) + for line in (tmp_path / "events.jsonl").read_text().splitlines() + ] + assert event["operator"] == "sem_filter" + assert event["event_type"] == "prompt_batching" + assert event["task_count"] == 2 + assert event["prompt_count"] == 1 + assert event["chunk_sizes"] == [2] + assert event["structured_output_token_limit"] == 100 + assert event["retry_count"] == 0 + + +def test_explicit_limit_chunks_tasks_deterministically() -> None: + model = _Model([["task_0=A,task_1=B", "task_2=C,task_3=D", "task_4=E"]]) + tasks = tuple(_Task(f"task_{index}", value) for index, value in enumerate("ABCDE")) + + result = run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=_request, + parse_results=_parse, + model=model, + config=PromptBatching(max_tasks=2), + max_retries=0, + progress_bar_desc="Testing", + operator="test", + ) + + assert result.outputs == tuple("ABCDE") + assert result.chunk_sizes == (2, 2, 1) + + +def test_configured_batch_fails_instead_of_silently_shrinking_for_context() -> None: + model = _Model([]) + model.max_ctx_len = 109 + tasks = ( + _Task("task_0", "AAAA"), + _Task("task_1", "BBBB"), + _Task("task_2", "CCCC"), + ) + + with pytest.raises(ValueError, match="configured prompt batch"): + run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=_request, + parse_results=_parse, + model=model, + config=PromptBatching(), + max_retries=0, + progress_bar_desc="Testing", + operator="test", + ) + + assert model.calls == [] + + +def test_prompt_batches_require_one_fixed_output_token_limit() -> None: + model = _Model([]) + tasks = (_Task("task_0", "A"), _Task("task_1", "B")) + + def varying_request(batch: tuple[_Task, ...]) -> PromptBatchRequest: + request = _request(batch) + return PromptBatchRequest( + task_ids=request.task_ids, + prompt=request.prompt, + max_tokens=100 + int(batch[0].task_id.removeprefix("task_")), + ) + + with pytest.raises(ValueError, match="same output token limit"): + run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=varying_request, + parse_results=_parse, + model=model, + config=PromptBatching(max_tasks=1), + max_retries=0, + progress_bar_desc="Testing", + operator="test", + ) + + assert model.calls == [] + + +def test_only_invalid_prompt_is_retried() -> None: + model = _Model( + [ + ["invalid", "task_2=C,task_3=D"], + ["task_0=A,task_1=B"], + ] + ) + tasks = tuple(_Task(f"task_{index}", value) for index, value in enumerate("ABCD")) + + def parse(raw_output: str) -> ParsedPromptBatch[str]: + if raw_output == "invalid": + raise ValueError("invalid output") + return _parse(raw_output) + + result = run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=_request, + parse_results=parse, + model=model, + config=PromptBatching(max_tasks=2), + max_retries=1, + progress_bar_desc="Testing", + operator="test", + ) + + assert result.outputs == ("A", "B", "C", "D") + assert result.prompt_count == 3 + assert result.retry_count == 1 + assert result.raw_output_attempts == ( + ("invalid", "task_0=A,task_1=B"), + ("invalid", "task_0=A,task_1=B"), + ("task_2=C,task_3=D",), + ("task_2=C,task_3=D",), + ) + assert len(model.calls) == 2 + assert len(model.calls[1][0]) == 1 + + +@pytest.mark.parametrize( + ("raw_output", "message"), + ( + ("task_0=A", "missing"), + ("task_0=A,task_1=B,task_9=X", "unknown"), + ("task_0=A,task_0=B,task_1=C", "duplicate"), + ), +) +def test_shared_runner_rejects_invalid_task_id_coverage( + raw_output: str, + message: str, +) -> None: + model = _Model([[raw_output]]) + tasks = (_Task("task_0", "A"), _Task("task_1", "B")) + + with pytest.raises(ValueError, match=message): + run_prompt_batches( + tasks, + task_id=lambda task: task.task_id, + build_request=_request, + parse_results=_parse, + model=model, + config=PromptBatching(), + max_retries=0, + progress_bar_desc="Testing", + operator="test", + ) + + +def test_empty_input_skips_model() -> None: + model = _Model([]) + + result = run_prompt_batches( + (), + task_id=lambda task: task.task_id, + build_request=_request, + parse_results=_parse, + model=model, + config=PromptBatching(), + max_retries=0, + progress_bar_desc="Testing", + operator="test", + ) + + assert result.outputs == () + assert result.prompt_count == 0 + assert model.calls == [] + + +def test_structured_transform_batches_rows_and_preserves_row_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + from agent_memory.adapters.lotus.structured import StructuredLMExecutor + from agent_memory.policy.logical import ColumnSpec + + raw_output = ( + '{"results":[' + '{"task_id":"task_1","output":{"label":"B"}},' + '{"task_id":"task_0","output":{"label":"A"}}]}' + ) + model = _Model([[raw_output]]) + model.max_tokens = 128 + model.max_ctx_len = 100_000 + model.cache = None + monkeypatch.setattr(lotus.settings, "lm", model) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + + generation = StructuredLMExecutor( + pd.DataFrame({"text": ["alpha", "beta"]}) + )( + input_cols=("text",), + output_cols=(ColumnSpec("label"),), + instruction="Label {text}.", + shape="object", + progress_bar_desc="Mapping", + model_kwargs={}, + structured_max_tokens=128, + prompt_batching=PromptBatching(), + operator="sem_map", + ) + + assert generation.parsed_outputs == ({"label": "A"}, {"label": "B"}) + assert generation.raw_output_attempts == ((raw_output,), (raw_output,)) + assert len(model.calls) == 1 + assert len(model.calls[0][0]) == 1 + shared_prompt = str(model.calls[0][0][0]) + assert "task_0" in shared_prompt and "task_1" in shared_prompt + + +def test_structured_flat_map_keeps_emitted_rows_with_their_source_task( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + from agent_memory.adapters.lotus.structured import StructuredLMExecutor + from agent_memory.policy.logical import ColumnSpec + + model = _Model( + [[ + '{"results":[' + '{"task_id":"task_0","output":{"rows":[' + '{"fact":"A1"},{"fact":"A2"}]}},' + '{"task_id":"task_1","output":{"rows":[]}}]}' + ]] + ) + model.max_tokens = 128 + model.max_ctx_len = 100_000 + model.cache = None + monkeypatch.setattr(lotus.settings, "lm", model) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + + generation = StructuredLMExecutor( + pd.DataFrame({"text": ["alpha", "beta"]}) + )( + input_cols=("text",), + output_cols=(ColumnSpec("fact"),), + instruction="Extract facts from {text}.", + shape="array", + progress_bar_desc="Flat mapping", + model_kwargs={}, + structured_max_tokens=128, + prompt_batching=PromptBatching(), + operator="sem_flat_map", + ) + + assert generation.parsed_outputs == ( + [{"fact": "A1"}, {"fact": "A2"}], + [], + ) + + +@pytest.mark.parametrize( + ("raw_output", "expected"), + ( + ('{"rows":[]}', []), + ('{"rows":[{"fact":"A1"},{"fact":"A2"}]}', [ + {"fact": "A1"}, + {"fact": "A2"}, + ]), + ), +) +def test_structured_flat_map_repairs_a_valid_singleton_output_envelope( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + raw_output: str, + expected: list[dict[str, str]], +) -> None: + import lotus + + from agent_memory.adapters.lotus.structured import StructuredLMExecutor + from agent_memory.policy.logical import ColumnSpec + + model = _Model([[raw_output]]) + model.max_tokens = 128 + model.max_ctx_len = 100_000 + model.cache = None + monkeypatch.setattr(lotus.settings, "lm", model) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + + generation = StructuredLMExecutor(pd.DataFrame({"text": ["alpha"]}))( + input_cols=("text",), + output_cols=(ColumnSpec("fact"),), + instruction="Extract facts from {text}.", + shape="array", + progress_bar_desc="Flat mapping", + model_kwargs={}, + structured_max_tokens=128, + structured_parse_retries=0, + semantic_trace_dir=tmp_path, + prompt_batching=PromptBatching(max_tasks=1), + operator="sem_flat_map", + ) + + assert generation.parsed_outputs == (expected,) + assert generation.raw_output_attempts == ((raw_output,),) + assert len(model.calls) == 1 + events = [ + json.loads(line) + for line in (tmp_path / "events.jsonl").read_text().splitlines() + ] + [batch_event] = [ + event for event in events if event["event_type"] == "prompt_batching" + ] + assert batch_event["retry_count"] == 0 + assert batch_event["structured_output_repair_count"] == 1 + assert batch_event["structured_output_repair_methods"] == [ + "singleton-envelope" + ] + + +def test_structured_map_repairs_a_valid_singleton_output_envelope( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + from agent_memory.adapters.lotus.structured import StructuredLMExecutor + from agent_memory.policy.logical import ColumnSpec + + raw_output = '{"label":"A"}' + model = _Model([[raw_output]]) + model.max_tokens = 128 + model.max_ctx_len = 100_000 + model.cache = None + monkeypatch.setattr(lotus.settings, "lm", model) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + + generation = StructuredLMExecutor(pd.DataFrame({"text": ["alpha"]}))( + input_cols=("text",), + output_cols=(ColumnSpec("label"),), + instruction="Label {text}.", + shape="object", + progress_bar_desc="Mapping", + model_kwargs={}, + structured_max_tokens=128, + structured_parse_retries=0, + prompt_batching=PromptBatching(max_tasks=1), + operator="sem_map", + ) + + assert generation.parsed_outputs == ({"label": "A"},) + assert generation.raw_output_attempts == ((raw_output,),) + assert len(model.calls) == 1 + + +def test_structured_prompt_batch_does_not_guess_a_multi_task_envelope( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + from agent_memory.adapters.lotus.structured import StructuredLMExecutor + from agent_memory.policy.logical import ColumnSpec + + model = _Model([['{"rows":[]}']]) + model.max_tokens = 128 + model.max_ctx_len = 100_000 + model.cache = None + monkeypatch.setattr(lotus.settings, "lm", model) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + + with pytest.raises(ValueError, match="after 1 attempt"): + StructuredLMExecutor(pd.DataFrame({"text": ["alpha", "beta"]}))( + input_cols=("text",), + output_cols=(ColumnSpec("fact"),), + instruction="Extract facts from {text}.", + shape="array", + progress_bar_desc="Flat mapping", + model_kwargs={}, + structured_max_tokens=128, + structured_parse_retries=0, + prompt_batching=PromptBatching(), + operator="sem_flat_map", + ) + + +@pytest.mark.parametrize( + "raw_output", + ( + '{"rows":[],"extra":true}', + '{"rows":[],"rows":[]}', + '{"rows":[{"wrong":"value"}]}', + '{"rows":', + ), +) +def test_structured_prompt_batch_rejects_invalid_singleton_inner_outputs( + monkeypatch: pytest.MonkeyPatch, + raw_output: str, +) -> None: + import lotus + + from agent_memory.adapters.lotus.structured import StructuredLMExecutor + from agent_memory.policy.logical import ColumnSpec + + model = _Model([[raw_output]]) + model.max_tokens = 128 + model.max_ctx_len = 100_000 + model.cache = None + monkeypatch.setattr(lotus.settings, "lm", model) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + + with pytest.raises(ValueError, match="after 1 attempt"): + StructuredLMExecutor(pd.DataFrame({"text": ["alpha"]}))( + input_cols=("text",), + output_cols=(ColumnSpec("fact"),), + instruction="Extract facts from {text}.", + shape="array", + progress_bar_desc="Flat mapping", + model_kwargs={}, + structured_max_tokens=128, + structured_parse_retries=0, + prompt_batching=PromptBatching(max_tasks=1), + operator="sem_flat_map", + ) + + +@pytest.mark.parametrize("context_limit", (128_000, 1_000_000)) +def test_p64_budget_uses_configured_context_without_reducing_output( + context_limit: int, +) -> None: + tasks = tuple(_Task(f"task_{i}", "x") for i in range(64)) + model = _Model([[",".join(f"{task.task_id}=ok" for task in tasks)]]) + model.max_ctx_len = context_limit + + def request(batch: tuple[_Task, ...]) -> PromptBatchRequest: + return PromptBatchRequest( + task_ids=tuple(task.task_id for task in batch), + prompt=[{"role": "system", "content": ""}, + {"role": "user", "content": "x" * 150_049}], + max_tokens=32_768, + ) + + def execute() -> Any: + return run_prompt_batches( + tasks, task_id=lambda task: task.task_id, + build_request=request, parse_results=_parse, model=model, + config=PromptBatching(max_tasks=64), max_retries=0, + progress_bar_desc="Testing", operator="test", + ) + + if context_limit == 128_000: + with pytest.raises(ValueError, match="context_limit=128000"): + execute() + assert model.calls == [] + else: + result = execute() + assert result.outputs == ("ok",) * 64 + assert result.chunk_sizes == (64,) + assert model.calls[0][1]["max_tokens"] == 32_768 diff --git a/tests/test_prompt_batching_scale.py b/tests/test_prompt_batching_scale.py new file mode 100644 index 0000000..66372a4 --- /dev/null +++ b/tests/test_prompt_batching_scale.py @@ -0,0 +1,613 @@ +from __future__ import annotations + +from datetime import datetime +from decimal import Decimal +import json +from pathlib import Path +from typing import Any +from unittest.mock import Mock, call + +import pytest + +from tools.analysis.prompt_batching_scale_report import generate_report +from tools.evaluation import prompt_batching_scale as controller + + +EXPECTED_LEVELS = ( + ("original", 1, None), + ("packed-1", 1, 1), + ("packed-2", 2, 2), + ("packed-4", 4, 4), + ("packed-8", 8, 8), + ("packed-16", 16, 16), + ("packed-32", 32, 32), + ("packed-64", 64, 64), + ("packed-128", 128, 128), + ("packed-all", 128, "all"), +) + + +def test_condition_matrix_runs_two_identical_forward_passes() -> None: + matrix = controller.condition_matrix() + + assert len(matrix) == 20 + for pass_index in (1, 2): + rows = [condition for condition in matrix if condition.pass_index == pass_index] + assert tuple( + (row.label, row.refresh_every, row.prompt_batch) for row in rows + ) == EXPECTED_LEVELS + assert [row.run_id for row in rows] == [ + f"pass-{pass_index}-{label}" for label, _refresh, _prompt in EXPECTED_LEVELS + ] + + +@pytest.mark.parametrize( + ("label", "refresh", "prompt", "expected_prompt"), + ( + ("original", 1, None, None), + ("packed-1", 1, 1, "1"), + ("packed-128", 128, 128, "128"), + ("packed-all", 128, "all", "all"), + ), +) +def test_condition_command_preserves_the_fixed_experiment_contract( + label: str, + refresh: int, + prompt: int | str | None, + expected_prompt: str | None, + tmp_path: Path, +) -> None: + condition = controller.ConditionSpec(1, label, refresh, prompt) + command = controller.build_condition_command( + condition, + source=tmp_path / "source", + venv=tmp_path / "venv", + bundle=tmp_path / "bundle", + output=tmp_path / "output", + ) + + assert command[command.index("--system") + 1] == "claude-memory" + assert command[command.index("--grouped-agg-rule") + 1] == "rule-join-map" + assert command[command.index("--sem-topk-method") + 1] == "listwise" + assert command[command.index("--semantic-pair-profile") + 1] == "search-filter" + assert command[command.index("--semantic-pair-top-k") + 1] == "20" + assert command[command.index("--semantic-pair-min-similarity") + 1] == "0.5" + assert command[command.index("--embedding-device") + 1] == "cuda" + assert command[command.index("--lotus-cache-mode") + 1] == "disabled" + assert command[command.index("--refresh-every") + 1] == str(refresh) + if expected_prompt is None: + assert "--prompt-batch-size" not in command + else: + assert command[command.index("--prompt-batch-size") + 1] == expected_prompt + + +def test_deepseek_cost_uses_request_time_periods() -> None: + assert controller._is_peak(datetime(2026, 9, 4, 9, 30)) + assert not controller._is_peak(datetime(2026, 9, 4, 13, 30)) + assert not controller._is_peak(datetime(2026, 9, 5, 10, 30)) + usage = { + "off_peak": {"cache_hit": 1_000_000, "cache_miss": 1_000_000, "completion": 1_000_000}, + "peak": {"cache_hit": 1_000_000, "cache_miss": 1_000_000, "completion": 1_000_000}, + } + + assert controller._usage_cost_cny(usage) == Decimal("18.150000") + + +def test_initialize_root_allows_only_prepared_source_package(tmp_path: Path) -> None: + (tmp_path / "source").mkdir() + (tmp_path / "package").mkdir() + + controller._validate_initialization_root(tmp_path) + + (tmp_path / "conditions").mkdir() + with pytest.raises(FileExistsError, match="unexpected entries"): + controller._validate_initialization_root(tmp_path) + + +def test_completed_questions_uses_current_artifact_layout(tmp_path: Path) -> None: + current = tmp_path / "cases" / "case-1" / "question-results" / "q1" + legacy = tmp_path / "cases" / "case-1" / "questions" / "q2" + current.mkdir(parents=True) + legacy.mkdir(parents=True) + (current / "complete.json").write_text("{}") + (legacy / "complete.json").write_text("{}") + + assert controller._completed_questions(tmp_path) == 1 + + +def test_monitor_counts_structured_output_repairs(tmp_path: Path) -> None: + trace = tmp_path / "events.jsonl" + trace.write_text( + json.dumps( + { + "operator": "sem_flat_map", + "event_type": "prompt_batching", + "query_digest": "site", + "task_count": 1, + "prompt_count": 1, + "retry_count": 0, + "chunk_sizes": [1], + "structured_output_repair_count": 1, + "syntax_repair_count": 0, + } + ) + + "\n" + ) + + state = controller._update_monitor(trace, tmp_path / "monitor-state.json") + + assert state["prompt_batching"]["sem_flat_map:site"]["repairs"] == 1 + + +def _contract(tmp_path: Path, conditions: list[dict[str, Any]]) -> dict[str, Any]: + return { + "experiment": "claude-prompt-batching-scale-128e", + "source": {"directory": str(tmp_path / "source")}, + "runtime": { + "venv": str(tmp_path / "venv"), + "deps_dir": str(tmp_path / "deps"), + "hf_home": str(tmp_path / "hf-home"), + }, + "conditions": conditions, + } + + +def test_condition_environment_uses_source_before_dependency_overlay( + tmp_path: Path, +) -> None: + paths = controller.ExperimentPaths.from_contract(tmp_path, _contract(tmp_path, [])) + + environment = controller._condition_environment(paths) + + assert environment["PYTHONPATH"].split(":") == [ + str((tmp_path / "source" / "src").resolve()), + str((tmp_path / "deps").resolve()), + ] + + +class _FailedProcess: + returncode = 1 + pid = 123 + + def poll(self) -> int: + return 1 + + +def test_first_condition_failure_stops_all_remaining_conditions( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + first = controller.ConditionSpec(1, "original", 1, None) + later = controller.ConditionSpec(1, "packed-1", 1, 1) + second_pass = controller.ConditionSpec(2, "original", 1, None) + control = tmp_path / "control" + control.mkdir() + (tmp_path / "conditions").mkdir() + (control / "experiment-contract.json").write_text( + json.dumps(_contract(tmp_path, [row.to_dict() for row in (first, later, second_pass)])) + ) + (control / "experiment-state").write_text("prepared\n") + monkeypatch.setattr(controller, "_before_condition", lambda _paths: None) + monkeypatch.setattr(controller.subprocess, "Popen", lambda *_args, **_kwargs: _FailedProcess()) + + with pytest.raises(RuntimeError, match="condition failed: pass-1-original"): + controller.run_experiment(tmp_path) + + assert (tmp_path / "conditions" / first.run_id / "state").read_text().strip() == "failed" + assert not (tmp_path / "conditions" / later.run_id).exists() + assert not (tmp_path / "conditions" / second_pass.run_id).exists() + + +class _SuccessfulProcess: + returncode = 0 + pid = 456 + + def poll(self) -> int: + return 0 + + +@pytest.mark.parametrize( + "failure_site", + ["pid", "_update_monitor", "_total_cost", "_directory_bytes", "_disk_available_gib", "interrupt"], +) +@pytest.mark.parametrize("stop_write_failure", [None, "log", "state"]) +def test_supervision_failure_stops_child_before_propagating( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + failure_site: str, + stop_write_failure: str | None, +) -> None: + conditions = [ + controller.ConditionSpec(1, "original", 1, None), + controller.ConditionSpec(2, "original", 1, None), + ] + control = tmp_path / "control" + control.mkdir() + (control / "experiment-contract.json").write_text( + json.dumps(_contract(tmp_path, [c.to_dict() for c in conditions])) + ) + (control / "experiment-state").write_text("prepared\n") + directory = tmp_path / "conditions" / conditions[0].run_id + trace = directory / "output" / "trace" / "events.jsonl" + trace.parent.mkdir(parents=True) + trace.write_text("{}\n") + error = ( + KeyboardInterrupt("interrupted") + if failure_site == "interrupt" + else OSError("monitor failed") + ) + process = Mock(pid=456, returncode=None) + process.poll.side_effect = lambda: process.returncode + + def wait(**_kwargs: Any) -> int: + process.returncode = 0 + return 0 + + process.wait.side_effect = wait + killpg = Mock() + monkeypatch.setattr(controller.os, "killpg", killpg) + monkeypatch.setattr(controller.subprocess, "Popen", lambda *_a, **_k: process) + monkeypatch.setattr(controller, "_before_condition", lambda _p: None) + monkeypatch.setattr( + controller, "_update_monitor", + lambda *_a: {"reasoning_tokens": 0, "provider_contract_violations": []}, + ) + monkeypatch.setattr(controller, "_total_cost", lambda _p: 0) + monkeypatch.setattr(controller, "_directory_bytes", lambda _p: 0) + monkeypatch.setattr(controller, "_disk_available_gib", lambda _p: 200) + if failure_site == "pid": + original_write = controller._atomic_text + + def write(path: Path, text: str) -> None: + if path.name == "runner.pid": + raise error + original_write(path, text) + + monkeypatch.setattr(controller, "_atomic_text", write) + elif failure_site == "interrupt": + monkeypatch.setattr(controller.time, "sleep", Mock(side_effect=error)) + else: + monkeypatch.setattr(controller, failure_site, Mock(side_effect=error)) + if stop_write_failure == "log": + original_log = controller._log + + def log(paths: Any, message: str) -> None: + if message.startswith("SAFETY_STOP"): + raise OSError("stop log failed") + original_log(paths, message) + + monkeypatch.setattr(controller, "_log", log) + elif stop_write_failure == "state": + original_set_state = controller._set_state + + def set_state(paths: Any, key: str, value: str) -> None: + if value == "safety-stopped": + raise OSError("stop state failed") + original_set_state(paths, key, value) + + monkeypatch.setattr(controller, "_set_state", set_state) + + with pytest.raises(type(error)) as raised: + controller.run_experiment(tmp_path) + + assert raised.value is error + killpg.assert_called_once_with(process.pid, controller.signal.SIGINT) + process.wait.assert_called_once_with(timeout=60) + assert not (tmp_path / "conditions" / conditions[1].run_id).exists() + if stop_write_failure is not None: + assert any(f"stop {stop_write_failure} failed" in note for note in error.__notes__) + + +def test_safe_stop_preserves_timeout_escalation( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + paths = controller.ExperimentPaths.from_contract(tmp_path, _contract(tmp_path, [])) + process = Mock(pid=456) + process.wait.side_effect = [controller.subprocess.TimeoutExpired("worker", 60), 0] + killpg = Mock() + monkeypatch.setattr(controller.os, "killpg", killpg) + monkeypatch.setattr(controller, "_log", Mock()) + monkeypatch.setattr(controller, "_set_state", Mock()) + + controller._safe_stop(paths, process, reason="test") + + assert killpg.call_args_list == [ + call(process.pid, controller.signal.SIGINT), + call(process.pid, controller.signal.SIGTERM), + ] + assert process.wait.call_args_list == [call(timeout=60), call(timeout=30)] + + +def test_second_pass_starts_only_after_the_first_pass_completes( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + conditions = [ + controller.ConditionSpec(1, "original", 1, None), + controller.ConditionSpec(1, "packed-1", 1, 1), + controller.ConditionSpec(2, "original", 1, None), + controller.ConditionSpec(2, "packed-1", 1, 1), + ] + control = tmp_path / "control" + control.mkdir() + (tmp_path / "conditions").mkdir() + (control / "experiment-contract.json").write_text( + json.dumps(_contract(tmp_path, [row.to_dict() for row in conditions])) + ) + (control / "experiment-state").write_text("prepared\n") + observed: list[str] = [] + monkeypatch.setattr(controller, "_before_condition", lambda _paths: None) + monkeypatch.setattr( + controller.subprocess, "Popen", lambda *_args, **_kwargs: _SuccessfulProcess() + ) + + def validate( + _paths: controller.ExperimentPaths, condition: controller.ConditionSpec + ) -> dict[str, float]: + observed.append(condition.run_id) + return {"cost_cny": 0.0} + + monkeypatch.setattr(controller, "_validate_condition", validate) + monkeypatch.setattr(controller, "summarize", lambda _root: {}) + + controller.run_experiment(tmp_path) + + assert observed == [row.run_id for row in conditions] + assert (control / "experiment-state").read_text().strip() == "completed" + + +def _validation(condition: controller.ConditionSpec) -> dict[str, Any]: + scale = float(condition.refresh_every) + return { + **condition.to_dict(), + "provider_call_count": int(100 / scale), + "physical_tokens": int(10_000 / scale), + "cache_hit_tokens": 100, + "cache_miss_tokens": 200, + "completion_tokens": 300, + "cost_cny": 1 / scale, + "insertion_excluding_trace_mean_ms": 1000 / scale, + "insertion_excluding_trace_median_ms": 900 / scale, + "insertion_excluding_trace_p95_ms": 1500 / scale, + "insertion_excluding_trace_max_ms": 2000 / scale, + "condition_end_to_end_wall_seconds": 300 / scale, + "retrieval_mean_ms": 50, + "official_locomo_score": 0.55, + "zep_judge_score": 0.7, + "candidate_pairs": { + "sem_join:site": {"eligible": 20, "selected": 10} + }, + "operator_usage": { + "sem_join:site": {"responses": 2, "tokens": 500} + }, + "prompt_batching": { + "sem_join:site": { + "tasks": 10, + "prompts": 2, + "max_chunk": 5, + "retries": 0, + "repairs": 0, + } + }, + "final_state": { + "topics": {"row_count": 5, "multiset_digest": "topic-digest"}, + "catalog": {"row_count": 5, "multiset_digest": "catalog-digest"}, + }, + "artifact_bytes": 1234, + } + + +def test_report_is_bounded_and_preserves_both_passes(tmp_path: Path) -> None: + conditions = [ + controller.ConditionSpec(1, "original", 1, None), + controller.ConditionSpec(2, "original", 1, None), + ] + control = tmp_path / "control" + control.mkdir() + (control / "experiment-contract.json").write_text( + json.dumps(_contract(tmp_path, [row.to_dict() for row in conditions])) + ) + for condition in conditions: + directory = tmp_path / "conditions" / condition.run_id + directory.mkdir(parents=True) + (directory / "validation.json").write_text( + json.dumps(_validation(condition)) + ) + + summary = generate_report(tmp_path) + + assert summary["completed_conditions"] == 2 + assert summary["complete"] is False + assert summary["cross_pass"]["original"]["topic_state_equal"] is True + assert {path.name for path in (tmp_path / "analysis").iterdir()} == { + "conditions.csv", + "site_metrics.csv", + "summary.json", + "efficiency.svg", + "quality.svg", + "pareto.svg", + } + assert "pass-1-original" in (tmp_path / "analysis" / "conditions.csv").read_text() + assert "pass-2-original" in (tmp_path / "analysis" / "conditions.csv").read_text() + assert (tmp_path / "analysis" / "efficiency.svg").read_text().startswith(" dict[str, Any]: + condition = controller.condition_matrix()[0] + contract = _contract(root, [condition.to_dict()]) + contract.update(input={"bundle_fingerprint": "input"}, fixed_execution={"model": "model"}) + controller._atomic_json(root / "control/experiment-contract.json", contract) + controller._atomic_text(root / "control/experiment-state", "failed") + controller._atomic_json(root / "control/source-evidence.json", {"source": "old"}) + directory = root / "conditions" / condition.run_id + controller._atomic_text(directory / "state", "completed") + controller._atomic_json(directory / "validation.json", _validation(condition)) + controller._atomic_json(directory / "output/manifest.json", {"source": "old"}) + controller._atomic_json(directory / "output/metrics/summary.json", { + "completed_cases": 1, "failed_cases": 0, "question_count": 71, + "memory_system_error_question_count": 0, "retrieval_system_error_count": 0, + "scores_by_scorer": { + controller.OFFICIAL_SCORER_ID: {"grade_count": 71}, + controller.JUDGE_SCORER_ID: {"grade_count": 54}, + }, + }) + controller._atomic_json(directory / "monitor-state.json", {"insertion_results": 128}) + for number in range(71): + controller._atomic_json(directory / f"output/cases/case/question-results/{number}/complete.json", {}) + trace = directory / "output/trace/events.jsonl" + controller._atomic_text(trace, "") + # Failed attempts must contribute to the carried budget as well. + controller._atomic_text(root / "conditions/pass-1-packed-1/output/trace/events.jsonl", json.dumps({ + "event_type": "provider_usage", "timestamp": "20260906T080000000000Z", + "provider_prompt_cache_hit_tokens": 0, "provider_prompt_cache_miss_tokens": 1_000_000, + "provider_completion_tokens": 0, + }) + "\n") + return contract + + +@pytest.mark.parametrize("damage", ["missing", "malformed"]) +def test_continuation_requires_cost_evidence(tmp_path: Path, damage: str) -> None: + contract = _completed_origin(tmp_path) + trace = tmp_path / "conditions/pass-1-original/output/trace/events.jsonl" + if damage == "missing": + trace.rename(trace.with_suffix(".saved")) + with pytest.raises(RuntimeError, match="reuse.*cost.*pass-1-original"): + controller._capture_continuation(tmp_path, ("pass-1-original",), contract) + else: + trace.write_text("{invalid JSON\n") + with pytest.raises(json.JSONDecodeError): + controller._capture_continuation(tmp_path, ("pass-1-original",), contract) + + +def test_chained_continuation_keeps_original_results_and_all_cost(tmp_path: Path) -> None: + old = tmp_path / "old" + original = _completed_origin(old) + middle = tmp_path / "middle" + contract = _completed_origin(middle) + contract["continuation"] = controller._capture_continuation( + old, ("pass-1-original",), original, + ) + controller._atomic_json(middle / "control/experiment-contract.json", contract) + result = controller._capture_continuation(middle, ("pass-1-original",), contract) + assert result["reused_conditions"]["pass-1-original"] == str(old / "conditions/pass-1-original") + assert Decimal(result["prior_cost_cny"]) == Decimal("3.0") + assert result["reused_sources"]["pass-1-original"] == original["source"] + controller._validate_continuation({"continuation": result}) + (old / "conditions/pass-1-original/validation.json").write_text("{}") + with pytest.raises(RuntimeError, match="changed"): + controller._validate_continuation({"continuation": result}) + + +def test_continuation_preserves_origin_and_failed_cost(tmp_path: Path) -> None: + origin = tmp_path / "old" + contract = _completed_origin(origin) + before = {str(p): p.read_bytes() for p in origin.rglob("*") if p.is_file()} + continuation = controller._capture_continuation(origin, ("pass-1-original",), contract) + assert Decimal(continuation["prior_condition_costs_cny"]["pass-1-original"]) == 0 + assert Decimal(continuation["prior_cost_cny"]) == Decimal("1.5") + assert {str(p): p.read_bytes() for p in origin.rglob("*") if p.is_file()} == before + new = tmp_path / "new" + contract = {**contract, "continuation": continuation} + controller._atomic_json(new / "control/experiment-contract.json", contract) + paths = controller.ExperimentPaths.from_contract(new, contract) + controller._atomic_json(new / "conditions/pass-1-packed-16/monitor-state.json", {"cost_cny": 2}) + assert controller._total_cost(paths) == Decimal("3.5") + controller._validate_continuation(contract) + summary = controller.summarize(new) + assert summary["runs"][0]["result_origin"]["status"] == "reused-completed" + assert summary["runs"][0]["result_origin"]["source"] == contract["source"] + report = generate_report(new) + assert report["result_origins"]["pass-1-original"]["status"] == "reused-completed" + assert not (new / "conditions/pass-1-original").exists() + (origin / "conditions/pass-1-original/validation.json").write_text("{}") + with pytest.raises(RuntimeError, match="changed"): + controller._validate_continuation(contract) + + +@pytest.mark.parametrize("damage", ("state", "questions", "summary", "input", "duplicate", "unknown")) +def test_continuation_rejects_invalid_reuse(tmp_path: Path, damage: str) -> None: + contract = _completed_origin(tmp_path) + directory = tmp_path / "conditions/pass-1-original" + ids = ("pass-1-original",) + if damage == "state": + (directory / "state").write_text("failed") + elif damage == "questions": + controller._atomic_json(directory / "output/cases/case/question-results/extra/complete.json", {}) + elif damage == "summary": + (directory / "output/metrics/summary.json").write_text("{}") + elif damage == "input": + contract = {**contract, "input": {"bundle_fingerprint": "different"}} + elif damage == "duplicate": + ids = ("pass-1-original", "pass-1-original") + else: + ids = ("pass-1-unknown",) + with pytest.raises((RuntimeError, ValueError)): + controller._capture_continuation(tmp_path, ids, contract) + + +def test_reused_condition_is_skipped_but_failure_blocks_second_pass( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + old = tmp_path / "old" + old_contract = _completed_origin(old) + root = tmp_path / "new" + conditions = [controller.condition_matrix()[i] for i in (0, 5, 10)] + contract = _contract(root, [c.to_dict() for c in conditions]) + contract["continuation"] = controller._capture_continuation(old, ("pass-1-original",), old_contract) + controller._atomic_json(root / "control/experiment-contract.json", contract) + controller._atomic_text(root / "control/experiment-state", "prepared") + monkeypatch.setattr(controller, "_before_condition", lambda _paths: None) + monkeypatch.setattr(controller.subprocess, "Popen", lambda *_args, **_kwargs: _FailedProcess()) + with pytest.raises(RuntimeError, match="condition failed: pass-1-packed-16"): + controller.run_experiment(root) + assert not (root / "conditions/pass-1-original").exists() + assert not (root / "conditions/pass-2-original").exists() + + +def test_monitor_includes_answer_judge_without_counting_lotus_mirrors(tmp_path: Path) -> None: + direct = { + "event_type": "llm_call", "operator": "llm", "phase": "answering", + "timestamp": "20260906T080000000000Z", "usage_scope": "batch", "llm_batch_size": 1, + "usage_prompt_cache_hit_tokens": 0, "usage_prompt_cache_miss_tokens": 1_000_000, + "usage_completion_tokens": 0, "llm_kwargs": {"thinking": {"type": "disabled"}}, + } + trace = tmp_path / "events.jsonl" + trace.write_text("\n".join(json.dumps(e) for e in ( + direct, {**direct, "phase": "grading"}, {**direct, "operator": "sem_map", "phase": "insertion"}, + )) + "\n") + state = controller._update_monitor(trace, tmp_path / "monitor.json") + assert state["cost_cny"] == 3 + assert state["provider_responses"] == 2 + assert state["provider_contract_violations"] == [] + assert controller._update_monitor(trace, tmp_path / "monitor.json")["cost_cny"] == 3 + + +def test_successful_continuation_enters_second_pass( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, +) -> None: + old = tmp_path / "old" + previous = _completed_origin(old) + root = tmp_path / "new" + conditions = [controller.condition_matrix()[i] for i in (0, 5, 10)] + contract = _contract(root, [c.to_dict() for c in conditions]) + contract["continuation"] = controller._capture_continuation(old, ("pass-1-original",), previous) + controller._atomic_json(root / "control/experiment-contract.json", contract) + controller._atomic_text(root / "control/experiment-state", "prepared") + monkeypatch.setattr(controller, "_before_condition", lambda _paths: None) + monkeypatch.setattr(controller.subprocess, "Popen", lambda *_args, **_kwargs: _SuccessfulProcess()) + observed = [] + + def validate(_paths: controller.ExperimentPaths, condition: controller.ConditionSpec) -> dict[str, float]: + observed.append(condition.run_id) + return {"cost_cny": 0} + + monkeypatch.setattr(controller, "_validate_condition", validate) + controller.run_experiment(root) + assert observed == ["pass-1-packed-16", "pass-2-original"] + + +def test_scale_commands_explicitly_set_memory_context() -> None: + for condition in controller.condition_matrix(): + command = controller.build_condition_command( + condition, source=Path("/source"), venv=Path("/venv"), + bundle=Path("/bundle"), output=Path("/output"), + ) + assert command[command.index("--lm-max-ctx-len") + 1] == "1000000" diff --git a/tests/test_relational_aggregates.py b/tests/test_relational_aggregates.py new file mode 100644 index 0000000..81a76c6 --- /dev/null +++ b/tests/test_relational_aggregates.py @@ -0,0 +1,343 @@ +"""Public API and full-query tests for deterministic relational aggregates.""" + +from __future__ import annotations + +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.policy.aggregates import ( + AvgAggregateSpec, + CountAggregateSpec, + SumAggregateSpec, +) +from agent_memory.policy.expressions import ArithmeticExpr, expr_from_param +from agent_memory.policy.logical import MemorySpec, MemoryView, QueryExpr +from agent_memory.policy.schema import output_columns +from agent_memory.planner.differential_policy import PolicyDifferentiator + + +def _execute(relation: object, rows: list[dict[str, object]]) -> pd.DataFrame: + expr = relation.expr # type: ignore[attr-defined] + columns = list(output_columns(expr.inputs[0])) + source = pd.DataFrame(rows, columns=columns) + return LotusAdapter().execute(expr, {"log": source}) + + +def test_aggregate_descriptors_and_convenience_methods_build_the_same_query() -> None: + rows = am.Log({"price": "Numeric price."}) + + composed = rows.agg( + am.count(output_col="row_count"), + am.sum(column="price", output_col="total_price"), + am.avg(column="price", output_col="average_price"), + ) + + specs = tuple(composed.expr.params["aggregates"]) + assert isinstance(specs[0], CountAggregateSpec) + assert isinstance(specs[1], SumAggregateSpec) + assert isinstance(specs[2], AvgAggregateSpec) + assert rows.count(output_col="row_count").expr.params["aggregates"] == specs[:1] + assert rows.sum(column="price", output_col="total_price").expr.params[ + "aggregates" + ] == specs[1:2] + assert rows.avg(column="price", output_col="average_price").expr.params[ + "aggregates" + ] == specs[2:] + assert output_columns(composed.expr) == ( + "row_count", + "total_price", + "average_price", + ) + + +def test_global_aggregate_uses_count_star_and_ignores_null_numeric_values() -> None: + rows = am.Log({"price": "Numeric price."}) + aggregate = rows.agg( + am.count(output_col="row_count"), + am.sum(column="price", output_col="total_price"), + am.avg(column="price", output_col="average_price"), + ) + + result = _execute( + aggregate, + [{"price": 10}, {"price": None}, {"price": 20}, {"price": 20}], + ) + + assert result.to_dict(orient="records") == [ + {"row_count": 4, "total_price": 50, "average_price": 50 / 3} + ] + + +def test_empty_global_aggregate_returns_one_sql_style_row() -> None: + rows = am.Log({"price": "Numeric price."}) + aggregate = rows.agg( + am.count(output_col="row_count"), + am.sum(column="price", output_col="total_price"), + am.avg(column="price", output_col="average_price"), + ) + + result = _execute(aggregate, []) + + assert result["row_count"].tolist() == [0] + assert result["total_price"].isna().tolist() == [True] + assert result["average_price"].isna().tolist() == [True] + + null_only = _execute(aggregate, [{"price": None}, {"price": None}]) + assert null_only["row_count"].tolist() == [2] + assert null_only["total_price"].isna().tolist() == [True] + assert null_only["average_price"].isna().tolist() == [True] + + +def test_grouped_multi_key_aggregate_returns_no_rows_for_empty_input() -> None: + rows = am.Log( + {"region": "Region.", "category": "Category.", "price": "Price."} + ) + aggregate = rows.group_by(["region", "category"]).agg( + am.count(output_col="row_count"), + am.sum(column="price", output_col="total_price"), + am.avg(column="price", output_col="average_price"), + ) + + result = _execute(aggregate, []) + + assert result.empty + assert list(result.columns) == [ + "region", + "category", + "row_count", + "total_price", + "average_price", + ] + + +def test_grouped_aggregate_preserves_duplicate_rows_and_null_groups() -> None: + rows = am.Log({"group": "Group.", "price": "Price."}) + aggregate = rows.group_by("group").agg( + am.count(output_col="row_count"), + am.sum(column="price", output_col="total_price"), + am.avg(column="price", output_col="average_price"), + ) + + result = _execute( + aggregate, + [ + {"group": "a", "price": 2}, + {"group": "a", "price": 2}, + {"group": "a", "price": None}, + {"group": None, "price": 3}, + ], + ) + + assert result.iloc[0].to_dict() == { + "group": "a", + "row_count": 3, + "total_price": 4, + "average_price": 2, + } + assert pd.isna(result.iloc[1]["group"]) + assert result.iloc[1]["row_count"] == 1 + assert result.iloc[1]["total_price"] == 3 + assert result.iloc[1]["average_price"] == 3 + + +def test_grouped_aggregate_does_not_collide_with_internal_column_names() -> None: + rows = am.Log( + { + "__am_aggregate_row_count": "A caller-owned group key.", + "price": "Price.", + } + ) + aggregate = rows.group_by("__am_aggregate_row_count").count( + output_col="count" + ) + + result = _execute( + aggregate, + [{"__am_aggregate_row_count": "group-a", "price": 1}], + ) + + assert result.to_dict(orient="records") == [ + {"__am_aggregate_row_count": "group-a", "count": 1} + ] + + +@pytest.mark.parametrize("value", ["12", True, object()]) +def test_sum_and_avg_reject_non_numeric_non_null_values(value: object) -> None: + rows = am.Log({"price": "Numeric price."}) + aggregate = rows.agg(am.sum(column="price", output_col="total")) + + with pytest.raises(TypeError, match="numeric"): + _execute(aggregate, [{"price": value}]) + + +def test_aggregate_authoring_rejects_invalid_columns_and_output_collisions() -> None: + rows = am.Log({"group": "Group.", "price": "Price."}) + + with pytest.raises(ValueError, match="input column"): + rows.sum(column="missing", output_col="total") + with pytest.raises(ValueError, match="group key"): + rows.group_by("group").count(output_col="group") + with pytest.raises(ValueError, match="unique"): + rows.agg( + am.count(output_col="result"), + am.sum(column="price", output_col="result"), + ) + + +def test_sem_groupby_rejects_numeric_aggregates_without_changing_semantic_agg() -> None: + rows = am.Log({"text": "Text."}) + grouped = rows.sem_groupby( + input_cols=["text"], instruction="Group equivalent text." + ) + + with pytest.raises(NotImplementedError, match="numeric algebraic"): + grouped.count(output_col="row_count") + with pytest.raises(NotImplementedError, match="numeric algebraic"): + grouped.agg(am.avg(column="text", output_col="average")) + + existing = grouped.agg( + am.sem_agg( + input_cols=["text"], + output_cols={"text": "Canonical text."}, + instruction="Choose the canonical text.", + ) + ) + assert existing.expr.op == "agg" + + +def test_arithmetic_expressions_round_trip_and_execute_both_operand_orders() -> None: + rows = am.Log({"left": "Left number.", "right": "Right number."}) + expression = (rows.col("left") + 2) * (10 - rows.col("right")) / 4 + rebuilt = expr_from_param(expression.to_param()) + + assert isinstance(expression, ArithmeticExpr) + assert rebuilt.to_param() == expression.to_param() + result = LotusAdapter().execute( + rows.assign( + add=rows.col("left") + rows.col("right"), + subtract=10 - rows.col("left"), + multiply=rows.col("left") * 3, + divide=rows.col("left") / 2, + reverse_add=2 + rows.col("left"), + reverse_multiply=3 * rows.col("left"), + reverse_divide=12 / rows.col("left"), + composed=expression, + ).expr, + {"log": pd.DataFrame([{"left": 6, "right": 2}])}, + ) + assert result.iloc[0].to_dict() == { + "left": 6, + "right": 2, + "add": 8, + "subtract": 4, + "multiply": 18, + "divide": 3, + "reverse_add": 8, + "reverse_multiply": 18, + "reverse_divide": 2, + "composed": 16, + } + + +def test_arithmetic_null_propagates_and_division_by_zero_is_explicit() -> None: + rows = am.Log({"numerator": "Numerator.", "denominator": "Denominator."}) + divided = rows.assign( + result=rows.col("numerator") / rows.col("denominator") + ) + + null_result = LotusAdapter().execute( + divided.expr, + {"log": pd.DataFrame([{"numerator": None, "denominator": 2}])}, + ) + assert null_result["result"].isna().tolist() == [True] + + with pytest.raises(ZeroDivisionError, match="division by zero"): + LotusAdapter().execute( + divided.expr, + {"log": pd.DataFrame([{"numerator": 1, "denominator": 0}])}, + ) + + +def test_sembench_representative_views_are_expressible_without_special_operators() -> None: + reviews = am.Log( + {"review": "Review text.", "movie": "Movie identifier.", "year": "Year."} + ) + positive = reviews.select(["review", "movie", "year"]).sem_filter( + instruction="The review expresses a positive opinion: {review}" + ) + + movie_q3 = positive.count(output_col="positive_count") + positive_count = movie_q3.assign(_scalar_key=1) + total_count = reviews.count(output_col="total_count").assign(_scalar_key=1) + movie_q4 = ( + positive_count.join(total_count, on="_scalar_key") + .assign( + positive_ratio=positive_count.col("positive_count") + / total_count.col("total_count") + ) + .select(["positive_ratio"]) + ) + movie_q8 = ( + reviews.sem_map( + input_cols=["review"], + output_cols={"sentiment": "POSITIVE or NEGATIVE."}, + instruction="Classify the review sentiment.", + ) + .group_by("sentiment") + .count(output_col="review_count") + ) + cars_q4 = ( + reviews.assign(age=2026 - reviews.col("year")) + .avg(column="age", output_col="average_age") + ) + + assert output_columns(movie_q3.expr) == ("positive_count",) + assert output_columns(movie_q4.expr) == ("positive_ratio",) + assert output_columns(movie_q8.expr) == ("sentiment", "review_count") + assert output_columns(cars_q4.expr) == ("average_age",) + + for name, relation in { + "movie_q3": movie_q3, + "movie_q4": movie_q4, + "movie_q8": movie_q8, + "cars_q4": cars_q4, + }.items(): + policy = PolicyDifferentiator().differentiate( + MemorySpec( + log=reviews, + views={name: MemoryView(name=name, query=relation.expr)}, + private_relations={}, + ) + ) + assert policy.nodes[policy.view_outputs[name]].output_columns == output_columns( + relation.expr + ) + + +def test_multi_aggregate_executes_its_input_only_once() -> None: + class CountingAdapter(LotusAdapter): + def __init__(self) -> None: + super().__init__() + self.log_executions = 0 + + def execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: + if query.op == "log": + self.log_executions += 1 + return super().execute(query, inputs) + + rows = am.Log({"price": "Price."}) + aggregate = rows.agg( + am.count(output_col="row_count"), + am.sum(column="price", output_col="total"), + am.avg(column="price", output_col="average"), + ) + adapter = CountingAdapter() + + adapter.execute(aggregate.expr, {"log": pd.DataFrame([{"price": 2}])}) + + assert adapter.log_executions == 1 diff --git a/tests/test_relational_limit.py b/tests/test_relational_limit.py new file mode 100644 index 0000000..a17aecb --- /dev/null +++ b/tests/test_relational_limit.py @@ -0,0 +1,73 @@ +"""Unordered LIMIT contracts and local materialized-parent refresh.""" + +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.policy.logical import QueryExpr +from agent_memory.policy.schema import output_columns +from agent_memory.tracing.semantic import query_digest + + +@pytest.mark.parametrize("n", [0, 1, 2, 10]) +@pytest.mark.parametrize("values", [[], [1, 1, 2]]) +def test_limit_bounds_preserves_schema_and_bag(n: int, values: list[int]) -> None: + source = am.Source({"value": "Number."}) + view = source.limit(n) + rows = pd.DataFrame({"value": pd.Series(values, dtype="Int64")}) + result = LotusAdapter().execute(view.expr, {"log": rows}) + assert len(result) == min(n, len(rows)) + assert result.dtypes.equals(rows.dtypes) + assert not (Counter(result.value) - Counter(rows.value)) + assert output_columns(view.expr) == ("value",) + assert view.expr.params == {"n": n} + assert query_digest(view.expr) == query_digest(source.limit(n).expr) + assert query_digest(view.expr) != query_digest(source.limit(n + 1).expr) + + +@pytest.mark.parametrize("n", [True, False, 1.0, "1", None]) +def test_limit_rejects_non_integer(n: Any) -> None: + with pytest.raises(TypeError, match="non-negative integer"): + am.Source({"value": "Number."}).limit(n) + + +def test_limit_rejects_negative() -> None: + with pytest.raises(ValueError, match="non-negative integer"): + am.Source({"value": "Number."}).limit(-1) + + +def test_limit_refreshes_replaced_parent_and_empty_parent() -> None: + source = am.Source({"value": "Number."}) + minimum = source.min(column="value", output_col="value") + view = minimum.filter(minimum.col("value") > 0).limit(2) + flow = am.SemanticDataflow(source=source, views={"result": view}) + for value, expected in [(4, [4]), (2, [2]), (-1, [])]: + flow.apply(pd.DataFrame({"value": [value]})) + assert flow.view("result").value.tolist() == expected + + +def test_limit_does_not_reexecute_semantic_parent() -> None: + class Oracle(LotusAdapter): + seen: list[int] + + def execute(self, query: QueryExpr, inputs: Mapping[str, Any]) -> Any: + if query.op == "sem_filter": + rows = self.execute(query.inputs[0], inputs) + self.seen.extend(rows.value.tolist()) + return rows.copy() + return super().execute(query, inputs) + + adapter = Oracle() + adapter.seen = [] + source = am.Source({"value": "Number."}) + view = source.sem_filter(instruction="Keep {value}.").limit(1) + flow = am.SemanticDataflow(source=source, views={"result": view}, adapter=adapter) + for value in [1, 2, 3]: + flow.apply(pd.DataFrame({"value": [value]})) + assert adapter.seen == [1, 2, 3] + assert len(flow.view("result")) == 1 diff --git a/tests/test_scalar_expressions.py b/tests/test_scalar_expressions.py new file mode 100644 index 0000000..ea34e70 --- /dev/null +++ b/tests/test_scalar_expressions.py @@ -0,0 +1,169 @@ +"""Tests for deterministic scalar expressions used by semantic views.""" + +from __future__ import annotations + +from collections.abc import Callable +import operator +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.policy.expressions import CaseWhenExpr, ComparisonExpr, TryCastExpr, expr_from_param + + +@pytest.mark.parametrize("kind", ["arithmetic", "cast", "case"]) +@pytest.mark.parametrize( + ("compare", "expected"), + [(operator.eq, [2]), (operator.ne, [1, 3]), (operator.lt, [1]), + (operator.le, [1, 2]), (operator.gt, [3]), (operator.ge, [2, 3])], +) +@pytest.mark.parametrize("reversed_operands", [False, True]) +def test_composite_comparisons_filter_rows( + kind: str, compare: Callable[[Any, Any], Any], expected: list[int], reversed_operands: bool, +) -> None: + source = am.Source({"value": "Value."}) + value = source.col("value") + expr = { + "arithmetic": value + 0, + "cast": am.try_cast(value, to="float"), + "case": am.case_when(True, value, 0), + }[kind] + comparison = compare(2, expr) if reversed_operands else compare(expr, 2) + assert isinstance(comparison, ComparisonExpr) + assert expr_from_param(comparison.to_param()).to_param() == comparison.to_param() + with pytest.raises(TypeError): + bool(comparison) + with pytest.raises(TypeError): + hash(expr) + result = LotusAdapter().execute( + source.filter(comparison).expr, + {"log": pd.DataFrame({"value": [1, 2, 3]})}, + ) + assert result["value"].tolist() == ( + sorted(4 - item for item in expected) if reversed_operands else expected + ) + + +def test_try_cast_and_case_when_execute_inside_assign() -> None: + source = am.Source({"raw_score": "Untrusted scalar score."}) + parsed_scores = source.assign( + _parsed_score=am.try_cast(source.col("raw_score"), to="float") + ) + normalized_scores = parsed_scores.assign( + review_score=am.case_when( + (parsed_scores.col("_parsed_score") >= 1) + & (parsed_scores.col("_parsed_score") <= 5), + parsed_scores.col("_parsed_score"), + 3.0, + ) + ).select(["review_score"]) + flow = am.SemanticDataflow( + source=source, + views={ + "parsed": parsed_scores.select(["_parsed_score"]), + "scores": normalized_scores, + }, + adapter=LotusAdapter(), + ) + + flow.apply( + pd.DataFrame( + {"raw_score": ["5", "2.5", "not-a-number", None, "0", "6"]} + ) + ) + + parsed = flow.view("parsed")["_parsed_score"].tolist() + assert parsed[:2] == [5.0, 2.5] + assert all(pd.isna(value) for value in parsed[2:4]) + pd.testing.assert_frame_equal( + flow.view("scores"), + pd.DataFrame({"review_score": [5.0, 2.5, 3.0, 3.0, 3.0, 3.0]}), + check_dtype=False, + ) + + +def test_case_when_only_evaluates_the_selected_branch_rows() -> None: + source = am.Source( + { + "numerator": "Numerator.", + "denominator": "Denominator.", + } + ) + quotient = source.col("numerator") / source.col("denominator") + result = source.assign( + guarded_then=am.case_when( + source.col("denominator") != 0, + quotient, + 0, + ), + guarded_else=am.case_when( + source.col("denominator") == 0, + 0, + quotient, + ), + ) + + frame = LotusAdapter().execute( + result.expr, + { + "log": pd.DataFrame( + [ + {"numerator": 10, "denominator": 0}, + {"numerator": 10, "denominator": 2}, + ], + index=pd.Index([7, 7]), + ) + }, + ) + + assert frame["guarded_then"].tolist() == [0, 5] + assert frame["guarded_else"].tolist() == [0, 5] + assert frame.index.tolist() == [7, 7] + + +@pytest.mark.parametrize(("condition", "expected"), [(True, 7), (False, 0)]) +def test_case_when_accepts_boolean_literal_conditions( + condition: bool, + expected: int, +) -> None: + source = am.Source({"value": "Value."}) + case_expr = am.case_when(condition, source.col("value"), 0) + query = source.assign(result=case_expr).select(["result"]) + + result = LotusAdapter().execute( + query.expr, + {"log": pd.DataFrame([{"value": 7}])}, + ) + + assert result["result"].tolist() == [expected] + assert expr_from_param(case_expr.to_param()).to_param() == case_expr.to_param() + + +@pytest.mark.parametrize("condition", [0, 1, "true"]) +def test_case_when_rejects_non_boolean_literal_conditions(condition: object) -> None: + with pytest.raises(TypeError, match="Expected a boolean relational expression"): + am.case_when(condition, 1, 0) + + +def test_try_cast_and_case_when_round_trip_through_query_params() -> None: + source = am.Source({"raw_score": "Untrusted scalar score."}) + cast_expr = am.try_cast(source.col("raw_score"), to="float") + case_expr = am.case_when(cast_expr.is_not_null(), cast_expr, 3.0) + + rebuilt_cast = expr_from_param(cast_expr.to_param()) + rebuilt_case = expr_from_param(case_expr.to_param()) + + assert isinstance(rebuilt_cast, TryCastExpr) + assert rebuilt_cast.to_param() == cast_expr.to_param() + assert isinstance(rebuilt_case, CaseWhenExpr) + assert rebuilt_case.to_param() == case_expr.to_param() + + +def test_try_cast_rejects_unsupported_target_types() -> None: + source = am.Source({"raw_score": "Untrusted scalar score."}) + + with pytest.raises(ValueError, match="only supports the 'float' target"): + am.try_cast(source.col("raw_score"), to="integer") diff --git a/tests/test_sem_agg_batching.py b/tests/test_sem_agg_batching.py new file mode 100644 index 0000000..22e624b --- /dev/null +++ b/tests/test_sem_agg_batching.py @@ -0,0 +1,782 @@ +"""Physical batching tests for independent semantic aggregate groups.""" + +from __future__ import annotations + +from collections.abc import Sequence +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +from agent_memory.adapters.lotus.adapter import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.relational import execute_agg +from agent_memory.adapters.lotus.sem_agg import ( + _execute_independent_sem_agg_groups, + _execute_lotus_style_sem_agg, + execute_native_sem_agg_groups, + execute_structured_sem_agg_groups, +) +from agent_memory.policy.logical import ColumnSpec, QueryExpr +from agent_memory.policy.aggregates import ArrayAggregateSpec, SemanticAggregateSpec + + +class _RecordingModel: + """Return deterministic outputs while retaining every LOTUS batch.""" + + max_ctx_len = 32_768 + max_tokens = 512 + + def __init__(self, outputs: Sequence[Sequence[str]] | None = None) -> None: + self.calls: list[tuple[list[Any], dict[str, Any]]] = [] + self._outputs = iter(outputs or ()) + + def count_tokens(self, _value: Any) -> int: + return 1 + + def __call__(self, prompts: list[Any], **kwargs: Any) -> Any: + self.calls.append((prompts, dict(kwargs))) + try: + outputs = tuple(next(self._outputs)) + except StopIteration: + outputs = tuple(_native_output(prompt) for prompt in prompts) + if len(outputs) != len(prompts): + raise AssertionError("fake output count must match prompts") + return SimpleNamespace(outputs=outputs) + + +def _native_output(prompt: Any) -> str: + text = str(prompt) + if "alpha" in text: + return "alpha summary" + if "beta" in text: + return "beta summary" + return "summary" + + +class _BudgetModel(_RecordingModel): + max_ctx_len = 10 + max_tokens = 1 + + def count_tokens(self, value: Any) -> int: + text = str(value) + if "{{docs_str}}" in text: + return 1 + return 100 if "oversized" in text else 4 + + +def _run_local_aggregate(path: str, documents: list[str], model: _RecordingModel) -> None: + if path == "independent": + _execute_independent_sem_agg_groups( + [documents], model, "Summarize.", config=LotusExecutionConfig(), + ) + else: + _execute_lotus_style_sem_agg( + documents, model, "Summarize.", [0] * len(documents), + safe_mode=False, progress_bar_desc="Test", response_format=None, + final_model_kwargs=None, + ) + + +@pytest.mark.parametrize("path", ["independent", "lotus-style"]) +@pytest.mark.parametrize("documents", [["oversized"], ["small", "oversized"]]) +def test_local_aggregate_rejects_oversized_document_before_provider( + path: str, documents: list[str], +) -> None: + model = _BudgetModel() + with pytest.raises(ValueError, match="sem_agg document does not fit"): + _run_local_aggregate(path, documents, model) + assert model.calls == [] + + +@pytest.mark.parametrize("path", ["independent", "lotus-style"]) +def test_local_aggregate_rejects_oversized_summary_at_next_level(path: str) -> None: + model = _BudgetModel(outputs=[("oversized", "summary")]) + with pytest.raises(ValueError, match="sem_agg document does not fit"): + _run_local_aggregate(path, ["one", "two", "three"], model) + assert len(model.calls) == 1 + assert len(model.calls[0][0]) == 2 + + +@pytest.mark.parametrize("path", ["independent", "lotus-style"]) +def test_local_aggregate_accepts_exact_budget_without_empty_prompts(path: str) -> None: + model = _BudgetModel() + _run_local_aggregate(path, ["one", "two"], model) + assert len(model.calls) == 1 + assert len(model.calls[0][0]) == 1 + prompt = model.calls[0][0][0][0]["content"] + assert "one" in prompt and "two" in prompt + + +@pytest.mark.parametrize("path", ["independent", "lotus-style"]) +def test_local_aggregate_recounts_after_document_number_reset(path: str) -> None: + class NumberSensitiveModel(_BudgetModel): + def count_tokens(self, value: Any) -> int: + text = str(value) + if "{{docs_str}}" in text or "Source" in text: + return 1 + return 3 if "Document 1: second" in text else 5 + + model = NumberSensitiveModel() + _run_local_aggregate(path, ["first", "second", "third"], model) + assert [len(prompts) for prompts, _kwargs in model.calls] == [2, 1] + second_prompt = model.calls[0][0][1][0]["content"] + assert "Document 1: second" in second_prompt + assert "Document 2: third" in second_prompt + + +def _query(*, structured: bool = False) -> QueryExpr: + output_cols = ( + (ColumnSpec("name"), ColumnSpec("summary")) + if structured + else (ColumnSpec("summary"),) + ) + return QueryExpr( + op="sem_agg", + params={ + "input_cols": ("body",), + "output_cols": output_cols, + "instruction": "Summarize {body}.", + }, + ) + + +def _groups() -> list[pd.DataFrame]: + return [ + pd.DataFrame({"body": ["alpha"]}), + pd.DataFrame({"body": ["beta"]}), + ] + + +def test_provider_batched_sem_agg_preserves_per_group_prompts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + sequential_model = _RecordingModel() + monkeypatch.setattr(lotus.settings, "lm", sequential_model) + sequential = execute_native_sem_agg_groups( + _query(), + _groups(), + ("body",), + LotusExecutionConfig(), + ) + + batched_model = _RecordingModel() + monkeypatch.setattr(lotus.settings, "lm", batched_model) + batched = execute_native_sem_agg_groups( + _query(), + _groups(), + ("body",), + LotusExecutionConfig(sem_agg_dispatch="provider-batched"), + ) + + sequential_prompts = [ + prompt for prompts, _kwargs in sequential_model.calls for prompt in prompts + ] + batched_prompts = [ + prompt for prompts, _kwargs in batched_model.calls for prompt in prompts + ] + assert batched == sequential == ["alpha summary", "beta summary"] + assert batched_prompts == sequential_prompts + assert [len(prompts) for prompts, _kwargs in sequential_model.calls] == [1, 1] + assert [len(prompts) for prompts, _kwargs in batched_model.calls] == [2] + + +def test_provider_batched_sem_agg_preserves_independent_hierarchies( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + class HierarchicalModel(_RecordingModel): + max_ctx_len = 15 + max_tokens = 1 + + def count_tokens(self, value: Any) -> int: + text = str(value) + if "Your job" in text: + return 1 + if "Source" in text: + return 2 + return 8 + + def __call__(self, prompts: list[Any], **kwargs: Any) -> Any: + self.calls.append((prompts, dict(kwargs))) + outputs = [] + for prompt in prompts: + text = str(prompt) + group = "alpha" if "alpha" in text else "beta" + prefix = "final" if "Source" in text else "partial" + outputs.append(f"{prefix} {group}") + return SimpleNamespace(outputs=outputs) + + model = HierarchicalModel() + monkeypatch.setattr(lotus.settings, "lm", model) + groups = [ + pd.DataFrame({"body": ["alpha one", "alpha two"]}), + pd.DataFrame({"body": ["beta one", "beta two"]}), + ] + + result = execute_native_sem_agg_groups( + _query(), + groups, + ("body",), + LotusExecutionConfig(sem_agg_dispatch="provider-batched"), + ) + + assert result == ["final alpha", "final beta"] + assert [len(prompts) for prompts, _kwargs in model.calls] == [4, 2] + + +def test_provider_batched_structured_sem_agg_retries_only_invalid_group( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + model = _RecordingModel( + outputs=( + ("", '{"name":"beta","summary":"B"}'), + ('{"name":"alpha","summary":"A"}',), + ) + ) + monkeypatch.setattr(lotus.settings, "lm", model) + + result = execute_structured_sem_agg_groups( + _query(structured=True), + _groups(), + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig( + sem_agg_dispatch="provider-batched", + structured_parse_retries=1, + ), + ) + + assert result == [ + {"name": "alpha", "summary": "A"}, + {"name": "beta", "summary": "B"}, + ] + assert [len(prompts) for prompts, _kwargs in model.calls] == [2, 1] + assert model.calls[0][1]["response_format"] == {"type": "json_object"} + assert model.calls[1][1]["response_format"] == {"type": "json_object"} + + +def test_provider_batched_structured_sem_agg_preserves_prompts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + outputs = ( + '{"name":"alpha","summary":"A"}', + '{"name":"beta","summary":"B"}', + ) + sequential_model = _RecordingModel(outputs=((outputs[0],), (outputs[1],))) + monkeypatch.setattr(lotus.settings, "lm", sequential_model) + sequential = execute_structured_sem_agg_groups( + _query(structured=True), + _groups(), + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig(), + ) + + batched_model = _RecordingModel(outputs=(outputs,)) + monkeypatch.setattr(lotus.settings, "lm", batched_model) + batched = execute_structured_sem_agg_groups( + _query(structured=True), + _groups(), + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig(sem_agg_dispatch="provider-batched"), + ) + + sequential_prompts = [ + prompt for prompts, _kwargs in sequential_model.calls for prompt in prompts + ] + batched_prompts = [ + prompt for prompts, _kwargs in batched_model.calls for prompt in prompts + ] + assert batched == sequential + assert batched_prompts == sequential_prompts + assert [len(prompts) for prompts, _kwargs in batched_model.calls] == [2] + + +def test_batch_prompted_sem_agg_keeps_group_outputs_independent( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + import agent_memory.adapters.lotus.sem_agg as sem_agg_module + + audit_rows: list[dict[str, Any]] = [] + monkeypatch.setattr( + sem_agg_module, + "write_sem_agg_audit", + lambda _config, **kwargs: audit_rows.append(kwargs), + ) + raw_batch_output = ( + '{"results":[' + '{"group_id":"group_0","output":' + '{"name":"alpha","summary":"A"}},' + '{"group_id":"group_1","output":' + '{"name":"beta","summary":"B"}}]}' + ) + model = _RecordingModel(outputs=((raw_batch_output,),)) + monkeypatch.setattr(lotus.settings, "lm", model) + + result = execute_structured_sem_agg_groups( + _query(structured=True), + _groups(), + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=2), + ), + ) + + assert result == [ + {"name": "alpha", "summary": "A"}, + {"name": "beta", "summary": "B"}, + ] + assert [len(prompts) for prompts, _kwargs in model.calls] == [1] + prompt = model.calls[0][0][0] + assert prompt[0]["role"] == "system" + assert "Never combine" in prompt[0]["content"] + assert '"group_id": "group_0"' in prompt[1]["content"] + assert '"group_id": "group_1"' in prompt[1]["content"] + assert len(audit_rows) == 2 + assert all(row["raw_output"] == raw_batch_output for row in audit_rows) + assert all(row["raw_output_attempts"] == (raw_batch_output,) for row in audit_rows) + + +def test_batch_prompted_sem_agg_keeps_one_fixed_output_limit_for_many_groups( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + group_count = 128 + groups = [pd.DataFrame({"body": [f"body-{index}"]}) for index in range(group_count)] + outputs = ",".join( + f'{{"group_id":"group_{index}","output":' + f'{{"name":"name-{index}","summary":"summary-{index}"}}}}' + for index in range(group_count) + ) + model = _RecordingModel(outputs=((f'{{"results":[{outputs}]}}',),)) + monkeypatch.setattr(lotus.settings, "lm", model) + config = LotusExecutionConfig(prompt_batching=PromptBatching(max_tasks=group_count)) + + result = execute_structured_sem_agg_groups( + _query(structured=True), + groups, + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + config, + ) + + assert len(result) == group_count + assert model.calls[0][1]["max_tokens"] == config.structured_max_tokens + + +def test_batch_prompted_sem_agg_accepts_one_extra_closing_brace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + import agent_memory.adapters.lotus.sem_agg as sem_agg_module + + raw_batch_output = ( + '{"results":[' + '{"group_id":"group_0","output":' + '{"name":"alpha","summary":"A"}},' + '{"group_id":"group_1","output":' + '{"name":"beta","summary":"B"}}]}}' + ) + model = _RecordingModel(outputs=((raw_batch_output,),)) + monkeypatch.setattr(lotus.settings, "lm", model) + audit_rows: list[dict[str, Any]] = [] + monkeypatch.setattr( + sem_agg_module, + "write_sem_agg_audit", + lambda _config, **kwargs: audit_rows.append(kwargs), + ) + + result = execute_structured_sem_agg_groups( + _query(structured=True), + _groups(), + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=2), + ), + ) + + assert result == [ + {"name": "alpha", "summary": "A"}, + {"name": "beta", "summary": "B"}, + ] + assert len(model.calls) == 1 + assert len(audit_rows) == 2 + assert all( + row["syntax_repair_method"] == "bounded-json-syntax" for row in audit_rows + ) + + +@pytest.mark.parametrize( + "raw_batch_output, expected_repair", + [ + ( + '{"results": [{"group_id": "group_0", "output": ' + '{"name": "alpha", "summary": "A"}},]}', + "json5", + ), + ( + "```json\n" + '{"results": [{"group_id": "group_0", "output": ' + '{"name": "alpha", "summary": "A"}}]}\n' + "```", + "json5-code-fence", + ), + ], +) +def test_batch_prompted_sem_agg_accepts_complete_tolerant_json_syntax( + monkeypatch: pytest.MonkeyPatch, + raw_batch_output: str, + expected_repair: str, +) -> None: + import lotus + import agent_memory.adapters.lotus.sem_agg as sem_agg_module + + model = _RecordingModel(outputs=((raw_batch_output,),)) + monkeypatch.setattr(lotus.settings, "lm", model) + audit_rows: list[dict[str, Any]] = [] + monkeypatch.setattr( + sem_agg_module, + "write_sem_agg_audit", + lambda _config, **kwargs: audit_rows.append(kwargs), + ) + + result = execute_structured_sem_agg_groups( + _query(structured=True), + _groups()[:1], + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=1), + ), + ) + + assert result == [{"name": "alpha", "summary": "A"}] + assert audit_rows[0]["syntax_repair_method"] == expected_repair + + +def test_sem_agg_audit_records_syntax_repair( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import agent_memory.adapters.lotus.sem_agg as sem_agg_module + + traced_rows: list[dict[str, Any]] = [] + monkeypatch.setattr( + sem_agg_module, + "write_structured_generation_trace", + lambda _trace_dir, *, operator, rows, snapshots: traced_rows.extend(rows), + ) + + sem_agg_module.write_sem_agg_audit( + LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=4), + ), + query=_query(structured=True), + group=_groups()[0], + input_cols=("body",), + output_cols=(ColumnSpec("name"), ColumnSpec("summary")), + group_index=0, + raw_output='{"results": []}}', + parsed_output={"name": "alpha", "summary": "A"}, + syntax_repair_method="json5-extra-closing-brace", + ) + + assert traced_rows[0]["structured_output_repaired"] is True + assert ( + traced_rows[0]["structured_output_repair_method"] + == "json5-extra-closing-brace" + ) + + +def test_batch_prompted_sem_agg_repairs_multiple_extra_closing_braces( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + model = _RecordingModel( + outputs=( + ( + '{"results":[{"group_id":"group_0","output":' + '{"name":"alpha","summary":"A"}}]}}}', + ), + ) + ) + monkeypatch.setattr(lotus.settings, "lm", model) + + result = execute_structured_sem_agg_groups( + _query(structured=True), + _groups()[:1], + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=1), + structured_parse_retries=0, + ), + ) + assert result == [{"name": "alpha", "summary": "A"}] + assert len(model.calls) == 1 + + +@pytest.mark.parametrize( + "raw_batch_output", + [ + ( + '{"results":[{"group_id":"group_0","group_id":"group_0",' + '"output":{"name":"alpha","summary":"A"}}]}' + ), + ( + '{"results":[{"group_id":"group_0","output":' + '{"name":"alpha","summary":"truncated}}]}' + ), + ], +) +def test_batch_prompted_sem_agg_rejects_ambiguous_json_repair( + monkeypatch: pytest.MonkeyPatch, + raw_batch_output: str, +) -> None: + import lotus + + model = _RecordingModel(outputs=((raw_batch_output,),)) + monkeypatch.setattr(lotus.settings, "lm", model) + + with pytest.raises(ValueError, match="invalid JSON"): + execute_structured_sem_agg_groups( + _query(structured=True), + _groups()[:1], + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=1), + structured_parse_retries=0, + ), + ) + + +def test_batch_prompted_native_sem_agg_audits_actual_shared_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + import agent_memory.adapters.lotus.sem_agg as sem_agg_module + + raw_batch_output = ( + '{"results":[' + '{"group_id":"group_0","output":{"summary":"A"}},' + '{"group_id":"group_1","output":{"summary":"B"}}]}' + ) + model = _RecordingModel(outputs=((raw_batch_output,),)) + monkeypatch.setattr(lotus.settings, "lm", model) + audit_rows: list[dict[str, Any]] = [] + monkeypatch.setattr( + sem_agg_module, + "write_sem_agg_audit", + lambda _config, **kwargs: audit_rows.append(kwargs), + ) + + result = execute_native_sem_agg_groups( + _query(), + _groups(), + ("body",), + LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=2), + ), + ) + + assert result == ["A", "B"] + assert len(audit_rows) == 2 + assert all(row["raw_output"] == raw_batch_output for row in audit_rows) + + +def test_batch_prompted_sem_agg_retries_only_invalid_prompt_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + groups = [ + pd.DataFrame({"body": ["alpha"]}), + pd.DataFrame({"body": ["beta"]}), + pd.DataFrame({"body": ["gamma"]}), + ] + model = _RecordingModel( + outputs=( + ( + "", + ( + '{"results":[{"group_id":"group_2","output":' + '{"name":"gamma","summary":"C"}}]}' + ), + ), + ( + ( + '{"results":[' + '{"group_id":"group_0","output":' + '{"name":"alpha","summary":"A"}},' + '{"group_id":"group_1","output":' + '{"name":"beta","summary":"B"}}]}' + ), + ), + ) + ) + monkeypatch.setattr(lotus.settings, "lm", model) + + result = execute_structured_sem_agg_groups( + _query(structured=True), + groups, + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=2), + structured_parse_retries=1, + ), + ) + + assert [values["name"] for values in result] == ["alpha", "beta", "gamma"] + assert [len(prompts) for prompts, _kwargs in model.calls] == [2, 1] + + +def test_batch_prompted_sem_agg_rejects_oversized_group( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + class TinyContextModel(_RecordingModel): + max_ctx_len = 20 + max_tokens = 5 + + def count_tokens(self, _value: Any) -> int: + return 100 + + monkeypatch.setattr(lotus.settings, "lm", TinyContextModel()) + + with pytest.raises(ValueError, match="does not fit the model context"): + execute_structured_sem_agg_groups( + _query(structured=True), + _groups()[:1], + ("body",), + (ColumnSpec("name"), ColumnSpec("summary")), + LotusExecutionConfig( + structured_max_tokens=5, + prompt_batching=PromptBatching(max_tasks=1), + ), + ) + + +def test_mixed_aggregate_batches_only_semantic_specs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + model = _RecordingModel( + outputs=( + ( + '{"name":"alpha","summary":"A"}', + '{"name":"beta","summary":"B"}', + ), + ) + ) + monkeypatch.setattr(lotus.settings, "lm", model) + source = pd.DataFrame( + { + "partition": ["a", "a", "b"], + "body": ["alpha one", "alpha two", "beta"], + "evidence": [1, 2, 3], + } + ) + source.attrs["agent_memory_groupby_keys"] = ("partition",) + group_query = QueryExpr(op="group_by", params={"keys": ("partition",)}) + query = QueryExpr( + op="agg", + inputs=(group_query,), + params={ + "aggregates": ( + SemanticAggregateSpec( + input_cols=("body",), + output_cols=(ColumnSpec("name"), ColumnSpec("summary")), + instruction="Summarize {body}.", + ), + ArrayAggregateSpec( + columns=("evidence",), + output_col="evidence_rows", + ), + ) + }, + ) + + result = execute_agg( + query, + {}, + lambda _query, _inputs: source, + SimpleNamespace( + config=LotusExecutionConfig(sem_agg_dispatch="provider-batched") + ), + ) + + assert result[["partition", "name", "summary"]].to_dict(orient="records") == [ + {"partition": "a", "name": "alpha", "summary": "A"}, + {"partition": "b", "name": "beta", "summary": "B"}, + ] + assert [len(prompts) for prompts, _kwargs in model.calls] == [2] + assert all(result["evidence_rows"]) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + ( + ({"sem_agg_dispatch": "parallel"}, "sem_agg_dispatch"), + ( + { + "sem_agg_dispatch": "provider-batched", + "prompt_batching": PromptBatching(max_tasks=4), + }, + "cannot be combined", + ), + ( + { + "sem_groupby_pair_batch_size": 4, + "prompt_batching": PromptBatching(max_tasks=4), + }, + "cannot be combined", + ), + ), +) +def test_sem_agg_physical_config_rejects_invalid_combinations( + kwargs: dict[str, Any], + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + LotusExecutionConfig(**kwargs) + + +def test_sem_agg_physical_methods_isolate_maintenance_identity() -> None: + baseline = LotusAdapter().maintenance_execution_fingerprint + provider_batched = LotusAdapter( + config=LotusExecutionConfig(sem_agg_dispatch="provider-batched") + ).maintenance_execution_fingerprint + batch_prompted_two = LotusAdapter( + config=LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=2), + ) + ).maintenance_execution_fingerprint + batch_prompted_four = LotusAdapter( + config=LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=4), + ) + ).maintenance_execution_fingerprint + + assert baseline == "" + assert ( + len({baseline, provider_batched, batch_prompted_two, batch_prompted_four}) == 4 + ) diff --git a/tests/test_sem_filter_batch_prompting.py b/tests/test_sem_filter_batch_prompting.py new file mode 100644 index 0000000..8bac223 --- /dev/null +++ b/tests/test_sem_filter_batch_prompting.py @@ -0,0 +1,224 @@ +"""Tests for semantic-filter batch prompting.""" + +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.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.sem_filter_batch_prompting import ( + execute_batch_prompted_sem_filter, +) + + +class FakeBatchLM: + """Return one configured output batch for each LM invocation.""" + + max_tokens = 512 + cache = None + + def __init__(self, output_batches: list[list[str]]) -> None: + self._output_batches = iter(output_batches) + self.calls: list[tuple[Any, dict[str, Any]]] = [] + + def __call__(self, messages: Any, **kwargs: Any) -> Any: + self.calls.append((messages, kwargs)) + return SimpleNamespace(outputs=next(self._output_batches)) + + +def _context( + *, + batch_size: int = 2, + retries: int = 0, +) -> SimpleNamespace: + return SimpleNamespace( + config=LotusExecutionConfig( + structured_parse_retries=retries, + prompt_batching=PromptBatching(max_tasks=batch_size), + ), + ) + + +def test_batch_prompting_preserves_order_multiplicity_and_uses_stable_ids( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + source = pd.DataFrame( + { + "left": ["same", "other", "same"], + "right": ["candidate", "candidate", "candidate"], + "untouched": [1, 2, 3], + }, + index=[7, 7, 9], + ) + lm = FakeBatchLM( + [ + [ + '{"decisions":[{"row_id":"row_1","keep":false},' + '{"row_id":"row_0","keep":true}]}', + '{"decisions":[{"row_id":"row_2","keep":true}]}', + ] + ] + ) + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + context = _context() + + result = execute_batch_prompted_sem_filter( + source, + instruction="{left} matches {right}.", + prompt_batching=context.config.prompt_batching, + context=context, + ) + + pd.testing.assert_frame_equal(result.frame, source.iloc[[0, 2]]) + assert result.prompt_count == 2 + assert result.retry_count == 0 + assert result.tuple_count == 3 + messages, kwargs = lm.calls[0] + assert len(messages) == 2 + prompt_text = str(messages) + assert all(row_id in prompt_text for row_id in ("row_0", "row_1", "row_2")) + assert "untouched" not in prompt_text + assert kwargs["response_format"] == {"type": "json_object"} + + +def test_batch_prompting_keeps_one_fixed_output_limit_for_many_tasks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + task_count = 128 + source = pd.DataFrame( + { + "left": [f"left-{index}" for index in range(task_count)], + "right": [f"right-{index}" for index in range(task_count)], + } + ) + decisions = ",".join( + f'{{"row_id":"row_{index}","keep":true}}' + for index in range(task_count) + ) + lm = FakeBatchLM([[f'{{"decisions":[{decisions}]}}']]) + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + context = _context(batch_size=task_count) + + result = execute_batch_prompted_sem_filter( + source, + instruction="{left} matches {right}.", + prompt_batching=context.config.prompt_batching, + context=context, + ) + + assert len(result.frame) == task_count + assert lm.calls[0][1]["max_tokens"] == context.config.structured_max_tokens + + +def test_batch_prompting_retries_only_invalid_batches( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + source = pd.DataFrame( + { + "left": ["a", "b", "c", "d"], + "right": ["A", "B", "C", "D"], + } + ) + lm = FakeBatchLM( + [ + [ + '{"decisions":[{"row_id":"row_0","keep":true},' + '{"row_id":"row_1","keep":false}]}', + '{"decisions":[{"row_id":"row_2","keep":true}]}', + ], + [ + '{"decisions":[{"row_id":"row_2","keep":false},' + '{"row_id":"row_3","keep":true}]}', + ], + ] + ) + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + context = _context(retries=1) + + result = execute_batch_prompted_sem_filter( + source, + instruction="{left} matches {right}.", + prompt_batching=context.config.prompt_batching, + context=context, + ) + + pd.testing.assert_frame_equal(result.frame, source.iloc[[0, 3]]) + assert result.prompt_count == 3 + assert result.retry_count == 1 + assert len(lm.calls) == 2 + retry_messages, _kwargs = lm.calls[1] + assert len(retry_messages) == 1 + retry_prompt = str(retry_messages) + assert "row_2" in retry_prompt and "row_3" in retry_prompt + assert "row_0" not in retry_prompt and "row_1" not in retry_prompt + + +@pytest.mark.parametrize( + ("raw_output", "message"), + ( + ('{"decisions":[{"row_id":"row_0","keep":1}]}', "boolean"), + ( + '{"decisions":[{"row_id":"row_0","keep":true},' + '{"row_id":"row_0","keep":false}]}', + "duplicate", + ), + ('{"decisions":[{"row_id":"row_9","keep":true}]}', "unknown"), + ('{"decisions":[]}', "omitted"), + ), +) +def test_batch_prompting_rejects_invalid_decision_contract( + monkeypatch: pytest.MonkeyPatch, + raw_output: str, + message: str, +) -> None: + import lotus + + lm = FakeBatchLM([[raw_output]]) + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + context = _context(batch_size=1) + + with pytest.raises(ValueError, match=message): + execute_batch_prompted_sem_filter( + pd.DataFrame({"left": ["a"], "right": ["A"]}), + instruction="{left} matches {right}.", + prompt_batching=context.config.prompt_batching, + context=context, + ) + + +def test_batch_prompting_empty_input_skips_lm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + lm = FakeBatchLM([]) + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + source = pd.DataFrame(columns=["left", "right"]) + context = _context() + + result = execute_batch_prompted_sem_filter( + source, + instruction="{left} matches {right}.", + prompt_batching=context.config.prompt_batching, + context=context, + ) + + pd.testing.assert_frame_equal(result.frame, source) + assert result.prompt_count == 0 + assert result.tuple_count == 0 + assert lm.calls == [] diff --git a/tests/test_sem_topk_join.py b/tests/test_sem_topk_join.py index 5b0b152..f6efa49 100644 --- a/tests/test_sem_topk_join.py +++ b/tests/test_sem_topk_join.py @@ -16,6 +16,7 @@ _pairwise_topk, _parse_listwise_ids, ) +from agent_memory.adapters.lotus.prompt_batching import PromptBatching from agent_memory.policy.logical import QueryExpr @@ -29,6 +30,10 @@ def configure(self) -> None: class _BatchListwiseLM: + max_ctx_len = 100_000 + max_tokens = 1024 + cache = None + def __init__(self, outputs: list[str]) -> None: self.outputs = outputs self.calls: list[tuple[Any, dict[str, Any]]] = [] @@ -37,6 +42,9 @@ def __call__(self, messages: Any, **kwargs: Any) -> Any: self.calls.append((messages, kwargs)) return SimpleNamespace(outputs=self.outputs) + def count_tokens(self, messages: Any) -> int: + return len(str(messages)) + def _query(*, k: int, how: str = "inner", on: tuple[str, ...] = ()) -> QueryExpr: return QueryExpr( @@ -124,6 +132,44 @@ def test_listwise_topk_join_can_select_multiple_matches_up_to_k( assert result["name:right"].tolist() == ["Mel", "Melanie"] +def test_prompt_batching_packs_multiple_listwise_join_anchors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + left = pd.DataFrame({"name": ["Melanie", "Poppy"]}, index=[10, 20]) + right = pd.DataFrame({"name": ["Mel", "Poppy"]}, index=[100, 200]) + lm = _BatchListwiseLM( + [ + '{"results":[' + '{"task_id":"task_1","selected_ids":["candidate_3"]},' + '{"task_id":"task_0","selected_ids":["candidate_0"]}]}' + ] + ) + monkeypatch.setattr(lotus.settings, "lm", lm) + + result = execute_sem_join( + _query(k=1), + {}, + lambda query, _inputs: left if query.params["name"] == "left" else right, + _StaticJoinContext( + LotusExecutionConfig( + sem_join_topk_method="listwise", + prompt_batching=PromptBatching(), + ) + ), + ) + + assert list(zip(result["name:left"], result["name:right"], strict=True)) == [ + ("Melanie", "Mel"), + ("Poppy", "Poppy"), + ] + assert len(lm.calls) == 1 + assert len(lm.calls[0][0]) == 1 + assert "task_0" in str(lm.calls[0][0][0]) + assert "task_1" in str(lm.calls[0][0][0]) + + def test_pairwise_topk_join_reuses_lotus_ranking_after_predicate_verification( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/test_semantic_dataflow.py b/tests/test_semantic_dataflow.py new file mode 100644 index 0000000..827d341 --- /dev/null +++ b/tests/test_semantic_dataflow.py @@ -0,0 +1,158 @@ +"""Tests for the public semantic dataflow facade.""" + +from __future__ import annotations + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter + + +def _build_flow() -> am.SemanticDataflow: + source = am.Source( + { + "category": "Row category.", + "value": "Numeric value.", + } + ) + selected = source.filter(source.col("category") == "keep").select( + ["category", "value"] + ) + counts = selected.count(output_col="row_count") + return am.SemanticDataflow( + source=source, + views={"selected": selected, "counts": counts}, + adapter=LotusAdapter(), + ) + + +def test_semantic_dataflow_applies_relation_rows_to_shared_views() -> None: + flow = _build_flow() + + flow.apply( + pd.DataFrame( + [ + {"category": "keep", "value": 1}, + {"category": "drop", "value": 2}, + {"category": "keep", "value": 3}, + ] + ) + ) + + pd.testing.assert_frame_equal( + flow.view("selected"), + pd.DataFrame( + [ + {"category": "keep", "value": 1}, + {"category": "keep", "value": 3}, + ] + ), + check_dtype=False, + ) + pd.testing.assert_frame_equal( + flow.view("counts"), + pd.DataFrame({"row_count": [2]}), + check_dtype=False, + ) + + +def test_semantic_dataflow_requires_exact_source_schema() -> None: + flow = _build_flow() + + with pytest.raises(ValueError, match="source delta columns must exactly match"): + flow.apply(pd.DataFrame([{"value": 1, "category": "keep"}])) + + +def test_semantic_dataflow_rejects_unknown_view() -> None: + flow = _build_flow() + + with pytest.raises(KeyError, match="unknown"): + flow.view("unknown") + + +def test_semantic_dataflow_view_is_a_defensive_copy() -> None: + flow = _build_flow() + flow.apply(pd.DataFrame([{"category": "keep", "value": 1}])) + + first = flow.view("selected") + first.loc[0, "value"] = 999 + + assert flow.view("selected").loc[0, "value"] == 1 + + +def test_semantic_dataflow_view_copies_nested_object_values() -> None: + flow = _build_flow() + flow.apply( + pd.DataFrame( + [{"category": "keep", "value": {"evidence": ["original"]}}] + ) + ) + + first = flow.view("selected") + first.loc[0, "value"]["evidence"].append("changed") + + assert flow.view("selected").loc[0, "value"] == { + "evidence": ["original"] + } + + +def test_semantic_dataflow_validates_declared_views() -> None: + source = am.Source({"value": "Numeric value."}) + + with pytest.raises(TypeError, match="views values must be Relation"): + am.SemanticDataflow( + source=source, + views={"result": object()}, # type: ignore[dict-item] + adapter=LotusAdapter(), + ) + + +@pytest.mark.parametrize("same_schema", [False, True]) +@pytest.mark.parametrize("mixed", [False, True]) +def test_semantic_dataflow_rejects_foreign_source( + same_schema: bool, mixed: bool, +) -> None: + source = am.Source({"value": "Numeric value."}) + foreign = am.Source( + {"value": "Numeric value."} if same_schema else {"other": "Other value."} + ) + view = foreign.select(["value"] if same_schema else ["other"]) + if mixed: + view = source.union_by_name(view) + + with pytest.raises(ValueError, match="view 'result'.*declared source"): + am.SemanticDataflow( + source=source, views={"result": view}, adapter=LotusAdapter(), + ) + + +def test_semantic_dataflow_accepts_self_join_and_window_source() -> None: + source = am.Source({"value": "Numeric value."}) + joined = source.alias("left").join(source.alias("right"), on="value") + windowed = source.count_window(size=2).process_window( + lambda window: window.select(["value"]) + ) + flow = am.SemanticDataflow( + source=source, views={"joined": joined, "windowed": windowed}, + adapter=LotusAdapter(), + ) + flow.apply(pd.DataFrame({"value": [1, 2]})) + assert len(flow.view("joined")) == 2 + assert flow.view("windowed")["value"].tolist() == [1, 2] + + +def test_semantic_dataflow_owns_nested_input_after_apply() -> None: + flow = _build_flow() + payload = {"evidence": ["original"]} + batch = pd.DataFrame([{"category": "keep", "value": payload}]) + flow.apply(batch) + + payload["evidence"].append("changed") + batch.loc[0, "value"]["extra"] = ["caller"] + assert flow.view("selected").loc[0, "value"] == {"evidence": ["original"]} + + # A later update must still consume the detached, original source row. + flow.apply(pd.DataFrame([{"category": "keep", "value": 2}])) + assert flow.view("selected").loc[0, "value"] == {"evidence": ["original"]} + assert flow.view("counts").loc[0, "row_count"] == 2 diff --git a/tests/test_semantic_join_ivm.py b/tests/test_semantic_join_ivm.py new file mode 100644 index 0000000..81eeae0 --- /dev/null +++ b/tests/test_semantic_join_ivm.py @@ -0,0 +1,420 @@ +"""Incremental maintenance tests for direct semantic inner joins.""" + +from __future__ import annotations + +from collections import Counter +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.pair_execution import ( + PAIR_LEFT_ID_COLUMN, + PAIR_LEFT_TEXT_COLUMN, + PAIR_RIGHT_ID_COLUMN, + PAIR_RIGHT_TEXT_COLUMN, +) +from agent_memory.adapters.lotus.sem_join import execute_sem_join +from agent_memory.evaluation.agent_memory_drivers import ( + inventory_operator_semantic_pair_sites, +) +from agent_memory.planner.differential_policy import DifferentialNode +from agent_memory.policy import QueryExpr, Relation + + +class _GeneralSemanticJoinMemory(am.Memory): + log = am.Log( + { + "side": "Join side.", + "tenant": "Exact tenant partition.", + "position": "Event order.", + "value": "Semantic value.", + } + ) + _left = ( + log.filter(log.col("side") == "left") + .select(["tenant", "position", "value"]) + .alias("left") + ) + _right = ( + log.filter(log.col("side") == "right") + .select(["tenant", "position", "value"]) + .alias("right") + ) + pairs = _left.sem_join( + _right, + on=[ + _left.col("tenant") == _right.col("tenant"), + _left.col("position") < _right.col("position"), + ], + instruction="{value:left} and {value:right} mean the same thing.", + how="inner", + ) + + +class _RetractingSemanticJoinMemory(am.Memory): + log = am.Log({"side": "Join side.", "key": "Key.", "value": "Value."}) + _minimum = log.group_by(["side", "key"]).min( + column="value", + output_col="value", + ) + _left = ( + _minimum.filter(_minimum.col("side") == "left") + .select(["key", "value"]) + .alias("left") + ) + _right = ( + _minimum.filter(_minimum.col("side") == "right") + .select(["key", "value"]) + .alias("right") + ) + pairs = _left.sem_join( + _right, + on="key", + instruction="{value:left} and {value:right} mean the same thing.", + how="inner", + ) + + +class _SemanticSelfJoinMemory(am.Memory): + log = am.Log({"key": "Exact partition.", "value": "Semantic value."}) + _left = log.select(["key", "value"]).alias("left") + _right = log.select(["key", "value"]).alias("right") + pairs = _left.sem_join( + _right, + on="key", + instruction="{value:left} and {value:right} mean the same thing.", + how="inner", + ) + + +class _TrackingSemanticAdapter(LotusAdapter): + """Use deterministic semantic decisions while recording candidate counts.""" + + def __init__(self) -> None: + super().__init__() + self.verified_pair_counts: list[int] = [] + + def verify_candidates(self, candidates: pd.DataFrame) -> list[tuple[Any, Any, None]]: + """Accept candidates whose projected semantic values are equal.""" + + if candidates.empty: + return [] + self.verified_pair_counts.append(len(candidates)) + return [ + ( + row[PAIR_LEFT_ID_COLUMN], + row[PAIR_RIGHT_ID_COLUMN], + None, + ) + for _index, row in candidates.iterrows() + if row[PAIR_LEFT_TEXT_COLUMN] == row[PAIR_RIGHT_TEXT_COLUMN] + ] + + +def _only_semantic_join_node( + memory_type: type[am.Memory], +) -> DifferentialNode: + nodes = [ + node + for node in memory_type.differentiate_policy().nodes.values() + if node.query.op == "sem_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 memory._runtime._engine.adapter.execute( + query, + {"log": memory._runtime._state["log"]}, + ) + + +def _patch_semantic_verifier( + monkeypatch: pytest.MonkeyPatch, + adapter: _TrackingSemanticAdapter, +) -> None: + """Replace only the LLM verdict while retaining production join execution.""" + + import agent_memory.adapters.lotus.sem_join as sem_join_module + + def verify( + candidates: pd.DataFrame, + *, + left_label: str, + right_label: str, + instruction: str, + config: LotusExecutionConfig, + ) -> list[tuple[Any, Any, None]]: + del left_label, right_label, instruction, config + return adapter.verify_candidates(candidates) + + monkeypatch.setattr(sem_join_module, "verify_semantic_join_candidates", verify) + + +def test_sem_join_reuses_general_join_on_normalization() -> None: + left = Relation( + QueryExpr( + op="materialized_view", + params={"name": "left", "columns": ("tenant", "position", "value")}, + ) + ).alias("left") + right = Relation( + QueryExpr( + op="materialized_view", + params={"name": "right", "columns": ("tenant", "position", "value")}, + ) + ).alias("right") + on = [ + left.col("tenant") == right.col("tenant"), + left.col("position") < right.col("position"), + ] + + semantic = left.sem_join(right, on=on, instruction="Rows match.") + relational = left.join(right, on=on) + + assert semantic.expr.params["on"] == relational.expr.params["on"] + + +def test_general_join_on_restricts_candidates_before_semantic_verification( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import agent_memory.adapters.lotus.sem_join as sem_join_module + + left = Relation( + QueryExpr( + op="materialized_view", + params={"name": "left", "columns": ("tenant", "position", "value")}, + ) + ).alias("left") + right = Relation( + QueryExpr( + op="materialized_view", + params={"name": "right", "columns": ("tenant", "position", "value")}, + ) + ).alias("right") + query = left.sem_join( + right, + on=[ + left.col("tenant") == right.col("tenant"), + left.col("position") < right.col("position"), + ], + instruction="{value:left} and {value:right} mean the same thing.", + how="outer", + ).expr + captured: list[pd.DataFrame] = [] + + def verify( + candidates: pd.DataFrame, + **_kwargs: object, + ) -> list[tuple[Any, Any, None]]: + captured.append(candidates.copy()) + return [ + ( + row[PAIR_LEFT_ID_COLUMN], + row[PAIR_RIGHT_ID_COLUMN], + None, + ) + for _index, row in candidates.iterrows() + ] + + monkeypatch.setattr(sem_join_module, "verify_semantic_join_candidates", verify) + inputs = { + "left": pd.DataFrame( + { + "tenant": ["a", "b"], + "position": [1, 1], + "value": ["left-a", "left-b"], + }, + index=pd.Index([10, 20]), + ), + "right": pd.DataFrame( + { + "tenant": ["a", "a", "c"], + "position": [2, 0, 3], + "value": ["right-a-new", "right-a-old", "right-c"], + }, + index=pd.Index([100, 200, 300]), + ), + } + + class Context: + config = LotusExecutionConfig() + pair_embedding_provider = None + + def configure(self) -> None: + pass + + context: Any = Context() + result = execute_sem_join(query, inputs, LotusAdapter().execute, context) + + assert len(captured) == 1 + assert captured[0].loc[:, [PAIR_LEFT_ID_COLUMN, PAIR_RIGHT_ID_COLUMN]].to_dict( + orient="records" + ) == [{PAIR_LEFT_ID_COLUMN: 10, PAIR_RIGHT_ID_COLUMN: 100}] + assert len(result) == 4 + + +def test_compiler_builds_bag_preserving_semantic_inner_join_maintenance() -> None: + node = _only_semantic_join_node(_GeneralSemanticJoinMemory) + + assert node.execution_kind == "semantic_binary_state" + assert node.maintenance_query is not None + nodes = _query_nodes(node.maintenance_query) + joins = [query for query in nodes if query.op == "sem_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_semantic_pair_site_inventory_includes_all_maintenance_copies() -> None: + policy = _GeneralSemanticJoinMemory.differentiate_policy() + + sites = inventory_operator_semantic_pair_sites( + policy, + operators=("sem_join",), + ) + + assert len(sites) == 1 + [site] = sites.values() + assert len(site.query_digests) == 4 + + +def test_semantic_inner_join_maintains_only_new_pairs_and_matches_full_recompute( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter = _TrackingSemanticAdapter() + _patch_semantic_verifier(monkeypatch, adapter) + memory = _GeneralSemanticJoinMemory(adapter=adapter) + + events = ( + {"side": "right", "tenant": "a", "position": 2, "value": "same"}, + {"side": "left", "tenant": "a", "position": 1, "value": "same"}, + {"side": "left", "tenant": "b", "position": 1, "value": "other"}, + {"side": "right", "tenant": "b", "position": 3, "value": "different"}, + {"side": "right", "tenant": "b", "position": 4, "value": "other"}, + ) + incremental_pair_counts: list[int] = [] + for event in events: + before = len(adapter.verified_pair_counts) + memory.add(event) + incremental_pair_counts.extend(adapter.verified_pair_counts[before:]) + assert _bag(memory._runtime._state["pairs"]) == _bag( + _full_view(memory, "pairs") + ) + + assert _bag(memory._runtime._state["pairs"]) == Counter( + { + ("a", 1, "same", "a", 2, "same"): 1, + ("b", 1, "other", "b", 4, "other"): 1, + } + ) + assert incremental_pair_counts == [1, 1, 1] + + +def test_semantic_inner_join_rejects_retracting_parent_changes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter = _TrackingSemanticAdapter() + _patch_semantic_verifier(monkeypatch, adapter) + memory = _RetractingSemanticJoinMemory(adapter=adapter) + + memory.add({"side": "right", "key": "k", "value": 5}) + memory.add({"side": "left", "key": "k", "value": 5}) + + with pytest.raises(NotImplementedError, match="append-only"): + memory.add({"side": "left", "key": "k", "value": 3}) + + +def test_semantic_self_join_counts_simultaneous_delta_pairs_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter = _TrackingSemanticAdapter() + _patch_semantic_verifier(monkeypatch, adapter) + memory = _SemanticSelfJoinMemory(adapter=adapter) + + memory.add({"key": "all", "value": "same"}) + assert len(memory._runtime._state["pairs"]) == 1 + + before = len(adapter.verified_pair_counts) + memory.add({"key": "all", "value": "same"}) + + assert adapter.verified_pair_counts[before:] == [1, 1, 1] + assert _bag(memory._runtime._state["pairs"]) == Counter( + {("all", "same", "all", "same"): 4} + ) + assert _bag(memory._runtime._state["pairs"]) == _bag( + _full_view(memory, "pairs") + ) + + +def test_semantic_inner_join_snapshot_round_trip( + monkeypatch: pytest.MonkeyPatch, +) -> None: + adapter = _TrackingSemanticAdapter() + _patch_semantic_verifier(monkeypatch, adapter) + original = _SemanticSelfJoinMemory(adapter=adapter) + original.add({"key": "all", "value": "same"}) + snapshot = original._runtime.snapshot_state() + + restored = _SemanticSelfJoinMemory(adapter=adapter) + restored._runtime.restore_state(snapshot) + restored.add({"key": "all", "value": "same"}) + + assert _bag(restored._runtime._state["pairs"]) == Counter( + {("all", "same", "all", "same"): 4} + ) + assert _bag(restored._runtime._state["pairs"]) == _bag( + _full_view(restored, "pairs") + ) + + +@pytest.mark.parametrize( + ("how", "k", "message"), + [ + ("outer", None, "how='inner'"), + ("inner", 1, "k=None"), + ], +) +def test_semantic_join_differential_rejects_unsupported_contracts( + how: str, + k: int | None, + message: str, +) -> None: + class UnsupportedSemanticJoinMemory(am.Memory): + log = am.Log({"value": "Value."}) + pairs = log.sem_join( + log, + instruction="{value:left} and {value:right} match.", + how=how, + k=k, + ) + + with pytest.raises(NotImplementedError, match=message): + UnsupportedSemanticJoinMemory.differentiate_policy() diff --git a/tests/test_semantic_pair_config.py b/tests/test_semantic_pair_config.py index 5728dc8..031a51a 100644 --- a/tests/test_semantic_pair_config.py +++ b/tests/test_semantic_pair_config.py @@ -46,6 +46,32 @@ def test_load_semantic_pair_profile_config_is_strict_and_canonical( "sem_groupby:fact", ] assert config.to_dict()["source_sha256"] == config.source_sha256 + assert config.schema_version == 1 + + +def test_load_semantic_pair_profile_config_rejects_unknown_schema( + tmp_path: Path, +) -> None: + path = tmp_path / "profiles.json" + path.write_text( + json.dumps( + { + "schema_version": 2, + "bindings": [ + { + "site_id": "sem_groupby:entity", + "mode": "search-filter", + "top_k": 15, + "min_similarity": 0.6, + } + ], + } + ), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="schema_version must be 1"): + load_semantic_pair_profile_config(path) @pytest.mark.parametrize( diff --git a/tests/test_semantic_pair_execution.py b/tests/test_semantic_pair_execution.py index ee99ff3..195f643 100644 --- a/tests/test_semantic_pair_execution.py +++ b/tests/test_semantic_pair_execution.py @@ -23,6 +23,7 @@ ) from agent_memory.evaluation.harness import MemorySystemContract from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.prompt_batching import PromptBatching 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 ( @@ -534,6 +535,110 @@ def test_sem_filter_search_filter_sends_only_candidates_to_oracle(tmp_path) -> N assert "pairs_snapshot_path" not in decisions[0] +def test_sem_filter_search_filter_batch_prompting_verifies_only_candidates( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import lotus + + class FakeLM: + max_tokens = 512 + cache = None + + def __init__(self) -> None: + self.calls: list[tuple[object, dict[str, object]]] = [] + + def __call__(self, messages: object, **kwargs: object) -> object: + self.calls.append((messages, kwargs)) + return type( + "Output", + (), + { + "outputs": [ + '{"decisions":[{"row_id":"row_0","keep":true}]}' + ] + }, + )() + + query = _filter_query() + source = _filter_source() + provider = FakeEmbeddingProvider( + { + "memory: old alpha": [1.0, 0.0], + "memory: old beta": [0.0, 1.0], + "memory: new alpha": [0.8, 0.6], + } + ) + digest = query_digest(query) + lm = FakeLM() + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + config = LotusExecutionConfig( + semantic_trace_dir=tmp_path, + semantic_pair_profiles={digest: _filter_profile(0.75)}, + prompt_batching=PromptBatching(max_tasks=4), + ) + FakeFilterFrame.oracle_batches = [] + + result = execute_sem_filter( + query, + {}, + lambda _query, _inputs: source, + FakeContext(config, provider), + ) + + assert list(result["memory:earlier"]) == ["old alpha"] + assert FakeFilterFrame.oracle_batches == [] + assert len(lm.calls) == 1 + prompt_text = json.dumps(lm.calls[0][0], ensure_ascii=False) + assert "old alpha" in prompt_text + assert "old beta" not in prompt_text + events = [ + json.loads(line) + for line in (tmp_path / "events.jsonl").read_text().splitlines() + ] + operator_result = next( + event + for event in events + if event["event_type"] == "operator_result" + ) + assert operator_result["prompt_batching"] == {"max_tasks": 4} + assert operator_result["verification_prompt_count"] == 1 + assert operator_result["verification_retry_count"] == 0 + assert operator_result["verified_tuple_count"] == 1 + + +def test_sem_filter_batch_prompting_rejects_lotus_cascade_before_input() -> None: + query = _filter_query() + config = LotusExecutionConfig( + sem_filter_cascade_args={"sampling_percentage": 0.1}, + prompt_batching=PromptBatching(max_tasks=4), + ) + + with pytest.raises(ValueError, match="does not support LOTUS options: cascade_args"): + execute_sem_filter( + query, + {}, + lambda _query, _inputs: (_ for _ in ()).throw( + AssertionError("input execution must not run before validation") + ), + FakeContext(config), + ) + + +def test_sem_filter_batch_prompting_changes_only_physical_fingerprint() -> None: + baseline = LotusAdapter().maintenance_execution_fingerprint + prompted = LotusAdapter( + config=LotusExecutionConfig( + prompt_batching=PromptBatching(max_tasks=4), + ) + ).maintenance_execution_fingerprint + + assert baseline == "" + assert prompted + assert prompted != baseline + + def test_sem_filter_without_matching_profile_keeps_oracle_path() -> None: query = _filter_query() source = _filter_source() @@ -768,6 +873,99 @@ def sem_filter(docs: list[object], *args: object, **kwargs: object) -> Output: assert "pairs" not in candidate +def test_sem_join_prompt_batching_verifies_ready_pairs_in_one_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + class FakeLM: + max_tokens = 512 + cache = None + + def __init__(self) -> None: + self.calls: list[object] = [] + + def __call__(self, messages: object, **_kwargs: object) -> object: + self.calls.append(messages) + return type( + "Output", + (), + { + "outputs": [ + '{"decisions":[' + '{"row_id":"row_0","keep":true},' + '{"row_id":"row_1","keep":false},' + '{"row_id":"row_2","keep":false},' + '{"row_id":"row_3","keep":true}]}' + ] + }, + )() + + lm = FakeLM() + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + left = pd.DataFrame( + { + "topic": ["alpha", "tea"], + "description": ["software project", "hot drink"], + } + ) + right = pd.DataFrame( + { + "topic": ["alpha project", "tea preference"], + "description": ["software initiative", "preferred hot drink"], + } + ) + query = QueryExpr( + op="sem_join", + inputs=( + QueryExpr(op="materialized_view", params={"name": "left"}), + QueryExpr(op="materialized_view", params={"name": "right"}), + ), + params={ + "instruction": ( + "{topic:left} and {topic:right} are the same topic when " + "{description:left} and {description:right} agree." + ), + "how": "inner", + }, + ) + + result = execute_sem_join( + query, + {"left": left, "right": right}, + lambda expression, inputs: inputs[str(expression.params["name"])], + FakeContext( + LotusExecutionConfig(prompt_batching=PromptBatching()), + ), + ) + + assert result.to_dict(orient="records") == [ + { + "topic:left": "alpha", + "description:left": "software project", + "topic:right": "alpha project", + "description:right": "software initiative", + }, + { + "topic:left": "tea", + "description:left": "hot drink", + "topic:right": "tea preference", + "description:right": "preferred hot drink", + }, + ] + assert len(lm.calls) == 1 + assert len(lm.calls[0]) == 1 + prompt = str(lm.calls[0]) + assert "topic: alpha" in prompt + assert "description: software project" in prompt + assert "left row" in prompt + assert "topic field" in prompt + assert "right row" in prompt + assert "description field" in prompt + assert "{topic:left}" not in prompt + + def test_sem_groupby_search_filter_verifies_only_candidate_edges( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -841,6 +1039,62 @@ def sem_filter(docs: list[object], *args: object, **kwargs: object) -> Output: assert "pairs" not in candidate +def test_sem_groupby_prompt_batching_verifies_ready_pairs_in_one_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + + class FakeLM: + max_tokens = 512 + cache = None + + def __init__(self) -> None: + self.calls: list[object] = [] + + def __call__(self, messages: object, **_kwargs: object) -> object: + self.calls.append(messages) + return type( + "Output", + (), + { + "outputs": [ + '{"decisions":[' + '{"row_id":"row_0","keep":true},' + '{"row_id":"row_1","keep":false},' + '{"row_id":"row_2","keep":false}]}' + ] + }, + )() + + lm = FakeLM() + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + source = pd.DataFrame( + {"topic": ["alpha", "alpha project", "tea preference"]} + ) + query = QueryExpr( + op="sem_groupby", + inputs=(QueryExpr(op="materialized_view", params={"name": "rows"}),), + params={ + "input_cols": ("topic",), + "instruction": "Group rows that describe the same topic: {topic}.", + }, + ) + + result = execute_sem_groupby( + query, + {"rows": source}, + lambda expression, inputs: inputs[str(expression.params["name"])], + FakeContext( + LotusExecutionConfig(prompt_batching=PromptBatching()), + ), + ) + + assert list(result["_agent_memory_group_id"]) == [0, 0, 1] + assert len(lm.calls) == 1 + assert len(lm.calls[0]) == 1 + + def test_sem_join_proxy_only_assembles_selected_pairs_without_oracle( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1078,6 +1332,36 @@ def test_exact_key_sem_join_rejects_lotus_cascade_before_input_execution() -> No ) +def test_prompt_batched_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"}, + ) + context = FakeContext( + LotusExecutionConfig( + prompt_batching=PromptBatching(), + sem_join_cascade_args={"sampling_percentage": 0.1}, + ) + ) + + with pytest.raises(ValueError, match="cannot be combined with LOTUS sem_join cascade"): + execute_sem_join( + query, + { + "left": pd.DataFrame({"topic": ["alpha"]}), + "right": pd.DataFrame({"topic": ["alpha"]}), + }, + lambda _expression, _inputs: (_ for _ in ()).throw( + AssertionError("input execution must not run before validation") + ), + context, + ) + + 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"]) diff --git a/tests/test_single_schema_predicate.py b/tests/test_single_schema_predicate.py new file mode 100644 index 0000000..5896096 --- /dev/null +++ b/tests/test_single_schema_predicate.py @@ -0,0 +1,216 @@ +"""Provider-free tests of genuinely single-context structured predicates.""" + +from collections import Counter +import json +from pathlib import Path +import re +from types import SimpleNamespace +from typing import Any + +import pandas as pd +import pytest + +import agent_memory as am +from agent_memory.adapters import LotusAdapter +from agent_memory.adapters.lotus.context import ( + LotusExecutionConfig, + LotusExecutionContext, +) +from agent_memory.adapters.lotus.predicate import ( + execute_schema_predicate, + single_schema_predicate, +) +from agent_memory.adapters.lotus.prompt_batching import PromptBatching + + +class FakeLM: + max_tokens = 128 + + def __init__(self, outputs: list[list[str]] | None = None) -> None: + self.outputs = iter(outputs) if outputs is not None else None + self.calls: list[tuple[Any, Any]] = [] + + def __call__(self, prompts: Any, **kwargs: Any) -> Any: + self.calls.append((prompts, kwargs)) + if self.outputs is not None: + return SimpleNamespace(outputs=next(self.outputs)) + results = [] + for prompt in prompts: + texts = re.findall("«(.*?)»", prompt[1]["content"], re.DOTALL) + keep = texts[0] != texts[1] if len(texts) == 2 else texts[0] == "positive" + results.append(json.dumps({"keep": keep})) + return SimpleNamespace(outputs=results) + + +@pytest.fixture +def config(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> LotusExecutionConfig: + import agent_memory.adapters.lotus.structured as structured + + monkeypatch.setattr(structured, "STRUCTURED_FAILURE_DIR", tmp_path / "failures") + return LotusExecutionConfig( + structured_output_transport="responses-json-schema", + structured_parse_retries=0, + semantic_trace_dir=tmp_path / "trace", + ) + + +def test_single_prompt_schema_and_failed_item_retry( + config: LotusExecutionConfig, monkeypatch: pytest.MonkeyPatch +) -> None: + import lotus + from dataclasses import replace + + lm = FakeLM([['{"keep":true}', '{"keep":"false"}'], ['{"keep":false}']]) + monkeypatch.setattr(lotus.settings, "lm", lm) + result = execute_schema_predicate( + pd.DataFrame({"text": ["positive", "negative"]}), + instruction="{text} is positive.", + config=replace(config, structured_parse_retries=1), + operator="sem_filter", + ) + assert result.decisions == (True, False) + assert [len(call[0]) for call in lm.calls] == [2, 1] + assert lm.calls[1][0] == [lm.calls[0][0][1]] + assert config.prompt_batching is None + for prompts, kwargs in lm.calls: + schema = kwargs["response_format"]["json_schema"]["schema"] + assert schema["properties"] == {"keep": {"type": "boolean"}} + assert schema["additionalProperties"] is False + for prompt in prompts: + assert len(re.findall("«(.*?)»", prompt[1]["content"])) == 1 + assert all( + word not in str(prompt) + for word in ("row_id", "decisions", "independent row") + ) + + +@pytest.mark.parametrize( + "raw", + [ + '{"keep":1}', + '{"keep":null}', + '{"keep":true,"extra":0}', + "[]", + "garbage", + '{"keep":"false"}', + ], +) +def test_invalid_output_is_not_false( + raw: str, + config: LotusExecutionConfig, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import lotus + + monkeypatch.setattr(lotus.settings, "lm", FakeLM([[raw]])) + with pytest.raises(ValueError, match="failed for rows"): + execute_schema_predicate( + pd.DataFrame({"text": ["x"]}), + instruction="{text} matches.", + config=config, + operator="sem_filter", + ) + artifacts = list((tmp_path / "failures").glob("*.json")) + assert len(artifacts) == 1 + assert json.loads(artifacts[0].read_text())["raw_outputs"] == [raw] + + +@pytest.mark.parametrize("outputs", [[], ['{"keep":true}', '{"keep":false}']]) +def test_wrong_output_count_fails( + outputs: list[str], config: LotusExecutionConfig, monkeypatch: pytest.MonkeyPatch +) -> None: + import lotus + + monkeypatch.setattr(lotus.settings, "lm", FakeLM([outputs])) + with pytest.raises(ValueError, match="expected 1 outputs"): + execute_schema_predicate( + pd.DataFrame({"text": ["x"]}), + instruction="{text} matches.", + config=config, + operator="sem_filter", + ) + + +def test_empty_and_default_paths(config: LotusExecutionConfig) -> None: + assert ( + execute_schema_predicate( + pd.DataFrame({"text": []}), + instruction="{text} matches.", + config=config, + operator="sem_filter", + ).decisions + == () + ) + assert not single_schema_predicate(LotusExecutionConfig()) + assert not single_schema_predicate( + LotusExecutionConfig( + structured_output_transport="responses-json-schema", + prompt_batching=PromptBatching(max_tasks=1), + ) + ) + + +@pytest.mark.parametrize( + "option", + [ + "sem_join_examples", + "sem_join_strategy", + "sem_join_cascade_args", + "sem_join_safe_mode", + ], +) +def test_join_rejects_unsupported_before_configure( + option: str, config: LotusExecutionConfig, monkeypatch: pytest.MonkeyPatch +) -> None: + from dataclasses import replace + from agent_memory.adapters.lotus.sem_join import execute_sem_join + + source = am.Source({"text": "Text."}) + query = source.sem_join( + source, instruction="{text:left} differs from {text:right}." + ) + context = LotusExecutionContext( + model="deepseek/deepseek-v4-flash", config=replace(config, **{option: True}) + ) + monkeypatch.setattr( + LotusExecutionContext, "configure", lambda _: pytest.fail("configured model") + ) + with pytest.raises(ValueError, match="does not support"): + execute_sem_join( + query.expr, {}, lambda *_: pytest.fail("executed input"), context + ) + + +def test_real_planner_join_and_filter( + config: LotusExecutionConfig, monkeypatch: pytest.MonkeyPatch +) -> None: + import lotus + + lm = FakeLM() + monkeypatch.setattr(lotus.settings, "lm", lm) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + monkeypatch.setattr(LotusExecutionContext, "configure", lambda _: None) + adapter = LotusAdapter(model="deepseek/deepseek-v4-flash", config=config) + source = am.Source({"id": "ID.", "text": "Text."}) + pairs = source.sem_join( + source, instruction="{text:left} differs from {text:right}." + ) + pairs = pairs.filter(pairs.col("id:left") != pairs.col("id:right")) + pairs = pairs.select(["id:left", "id:right"]) + positive = source.sem_filter(instruction="{text} is positive.") + flow = am.SemanticDataflow( + source=source, views={"pairs": pairs, "positive": positive}, adapter=adapter + ) + rows = pd.DataFrame({"id": [1, 2, 3], "text": ["positive", "negative", "positive"]}) + for i in range(len(rows)): + flow.apply(rows.iloc[[i]]) + incremental_tasks = sum(len(p) for p, _ in lm.calls) + full = adapter.execute(pairs.expr, {"log": rows}) + assert Counter(full.itertuples(index=False, name=None)) == Counter( + flow.view("pairs").itertuples(index=False, name=None) + ) + assert len(full) == 4 + assert incremental_tasks == 9 + 3 + assert sum(len(p) for p, _ in lm.calls) - incremental_tasks == 9 + assert list(flow.view("positive")["id"]) == [1, 3] diff --git a/tests/test_storage_contract.py b/tests/test_storage_contract.py index 8658516..5dc2045 100644 --- a/tests/test_storage_contract.py +++ b/tests/test_storage_contract.py @@ -366,7 +366,11 @@ def rebuild( self.markers[namespace] = next_commit -def _stored_memory(connector: _RecordingConnector) -> am.Memory: +def _stored_memory( + connector: _RecordingConnector, + *, + refresh: am.CountRefresh | None = None, +) -> am.Memory: class StoredMemory(am.Memory): log = am.Log({"value": "Value."}) rows = log.select(["value"]) @@ -379,6 +383,7 @@ class StoredMemory(am.Memory): statements=statements, namespace="test", ), + refresh=refresh, ) @@ -396,6 +401,42 @@ def test_storage_receives_exact_sink_changelog_after_successful_step() -> None: assert memory._runtime._state["rows"].to_dict("records") == [{"value": 7}] +def test_count_refresh_commits_one_storage_transaction_for_the_delta_batch() -> None: + connector = _RecordingConnector() + memory = _stored_memory(connector, refresh=am.CountRefresh(every=3)) + + memory.add({"value": 1}) + memory.add({"value": 2}) + + assert connector.commits == [] + + memory.add({"value": 3}) + + assert len(connector.commits) == 1 + _statement, inserted, retracted = connector.commits[0][0] + assert inserted.to_dict("records") == [ + {"value": 1}, + {"value": 2}, + {"value": 3}, + ] + assert retracted.empty + assert connector.markers["test"].source_row_count == 3 + + +def test_count_refresh_storage_failure_keeps_the_whole_pending_batch() -> None: + connector = _RecordingConnector(fail_writes=True) + memory = _stored_memory(connector, refresh=am.CountRefresh(every=2)) + + memory.add({"value": 1}) + with pytest.raises(RuntimeError, match="storage write failed"): + memory.add({"value": 2}) + + assert connector.commits == [] + assert connector.rollbacks == 1 + assert memory.pending_count == 2 + assert memory._runtime._state == {} + + def test_storage_receives_retractions_when_a_sink_row_is_replaced() -> None: class EarliestMemory(am.Memory): log = am.Log({"key": "Key.", "value": "Value."}) diff --git a/tests/test_structured_output_reliability.py b/tests/test_structured_output_reliability.py new file mode 100644 index 0000000..4f8fbdf --- /dev/null +++ b/tests/test_structured_output_reliability.py @@ -0,0 +1,556 @@ +"""Provider-free integration contracts for structured batch transport and repair.""" + +from __future__ import annotations + +from copy import deepcopy +from hashlib import sha256 +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Literal + +import lotus +import pandas as pd +import pytest + +from agent_memory.adapters.lotus import LotusAdapter +from agent_memory.adapters.lotus.context import LotusExecutionConfig +from agent_memory.adapters.lotus.pair_execution import ( + PAIR_LEFT_ID_COLUMN, + PAIR_LEFT_TEXT_COLUMN, + PAIR_RIGHT_ID_COLUMN, + PAIR_RIGHT_TEXT_COLUMN, +) +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.adapters.lotus.sem_agg_batch_prompting import ( + execute_batch_prompted_sem_agg, +) +from agent_memory.adapters.lotus.sem_filter_batch_prompting import ( + execute_batch_prompted_sem_filter, +) +from agent_memory.adapters.lotus.sem_topk_join import _listwise_topk +from agent_memory.adapters.lotus.structured import ( + StructuredLMExecutor, + _parse_structured_prompt_batch, +) +from agent_memory.policy.logical import ColumnSpec +from agent_memory.runtime.executor import PolicyExecutor + + +TRANSPORTS = ("chat-json-object", "responses-json-schema") +PROTOCOLS = ("predicate", "aggregate", "object", "array", "listwise") +OUTPUT_COLS = (ColumnSpec("label"),) + + +class _Model: + """Capture physical requests and return local output/metadata fixtures only.""" + + max_tokens = 256 + max_ctx_len = 100_000 + cache = None + + def __init__( + self, + outputs: list[str], + metadata: list[dict[str, Any]] | None = None, + ) -> None: + self.outputs = outputs + self.metadata = metadata + self.calls: list[tuple[Any, dict[str, Any]]] = [] + + def count_tokens(self, messages: list[dict[str, str]]) -> int: + """Use a deterministic, provider-free context estimate.""" + return sum(len(message["content"]) for message in messages) + + def __call__(self, messages: Any, **kwargs: Any) -> SimpleNamespace: + """Never retry implicitly or invoke an external model.""" + self.calls.append((deepcopy(messages), deepcopy(kwargs))) + return SimpleNamespace(outputs=self.outputs, response_metadata=self.metadata) + + +def _payload(protocol: str, indices: tuple[int, ...] = (0, 1)) -> dict[str, Any]: + if protocol == "predicate": + return { + "decisions": [ + {"row_id": f"row_{index}", "keep": index % 2 == 0} + for index in indices + ] + } + if protocol == "listwise": + return { + "results": [ + {"task_id": f"task_{index}", "selected_ids": [f"candidate_{index * 3}"]} + for index in indices + ] + } + id_field, prefix = ("group_id", "group") if protocol == "aggregate" else ( + "task_id", "task" + ) + return { + "results": [ + { + id_field: f"{prefix}_{index}", + "output": ( + {"rows": [{"label": f"value_{index}"}]} + if protocol == "array" + else {"label": f"value_{index}"} + ), + } + for index in indices + ] + } + + +def _execute( + protocol: str, + model: _Model, + monkeypatch: pytest.MonkeyPatch, + *, + transport: str = "responses-json-schema", + count: int = 2, + batch_size: int | None = None, + trace_dir: Path | None = None, +) -> Any: + monkeypatch.setattr(lotus.settings, "lm", model) + monkeypatch.setattr(lotus.settings, "enable_cache", False) + batching = PromptBatching(max_tasks=batch_size) + config = LotusExecutionConfig( + prompt_batching=batching, + structured_output_transport=transport, + structured_parse_retries=0, + structured_max_tokens=256, + semantic_trace_dir=trace_dir, + ) + context: Any = SimpleNamespace(config=config) + source = pd.DataFrame({"text": [f'Input {index}: "quoted"\nnext line' for index in range(count)]}) + if protocol == "predicate": + return execute_batch_prompted_sem_filter( + source, + instruction="Keep relevant {text}.", + context=context, + prompt_batching=batching, + ).decisions + if protocol == "aggregate": + return execute_batch_prompted_sem_agg( + [[text] for text in source["text"]], + instruction="Summarize each group.", + output_cols=OUTPUT_COLS, + model=model, + prompt_batching=batching, + structured_output_transport=transport, + max_retries=0, + model_kwargs={}, + progress_bar_desc="Testing aggregates", + trace_dir=trace_dir, + ).outputs + if protocol in {"object", "array"}: + return StructuredLMExecutor(source)( + input_cols=("text",), + output_cols=OUTPUT_COLS, + instruction="Label {text}.", + shape=protocol, + progress_bar_desc="Testing structured outputs", + model_kwargs={}, + structured_max_tokens=256, + structured_parse_retries=0, + semantic_trace_dir=trace_dir, + prompt_batching=batching, + structured_output_transport=transport, + operator="sem_map" if protocol == "object" else "sem_flat_map", + ).parsed_outputs + candidates = pd.DataFrame({ + PAIR_LEFT_ID_COLUMN: [index for index in range(count) for _ in range(3)], + PAIR_RIGHT_ID_COLUMN: list(range(count * 3)), + PAIR_LEFT_TEXT_COLUMN: [text for text in source["text"] for _ in range(3)], + PAIR_RIGHT_TEXT_COLUMN: [f"candidate text {index}" for index in range(count * 3)], + }) + matches, retries = _listwise_topk( + candidates, + instruction="Match the same entity.", + left_label="left", + right_label="right", + k=2, + context=context, + ) + assert retries == 0 + return matches + + +def _object_schema(fields: dict[str, Any]) -> dict[str, Any]: + return { + "type": "object", + "properties": fields, + "required": list(fields), + "additionalProperties": False, + } + + +def _expected_schema(protocol: str) -> dict[str, Any]: + scalar = {"type": ["string", "number", "boolean", "null"]} + if protocol == "predicate": + item = _object_schema({"row_id": {"type": "string"}, "keep": {"type": "boolean"}}) + return _object_schema({"decisions": {"type": "array", "items": item}}) + if protocol == "listwise": + item = _object_schema({ + "task_id": {"type": "string"}, + "selected_ids": {"type": "array", "items": {"type": "string"}}, + }) + else: + output = _object_schema({"label": scalar}) + if protocol == "array": + output = _object_schema({"rows": {"type": "array", "items": output}}) + item = _object_schema({ + "group_id" if protocol == "aggregate" else "task_id": {"type": "string"}, + "output": output, + }) + return _object_schema({"results": {"type": "array", "items": item}}) + + +@pytest.mark.parametrize( + ("protocol", "expected"), + ( + ("predicate", (True, False)), + ("aggregate", ({"label": "value_0"}, {"label": "value_1"})), + ("object", ({"label": "value_0"}, {"label": "value_1"})), + ("array", ([{"label": "value_0"}], [{"label": "value_1"}])), + ("listwise", [(0, 0, None), (1, 3, None)]), + ), +) +def test_protocol_schema_reaches_model_without_changing_prompt_bytes( + protocol: str, expected: Any, monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only response_format changes across transports, not prompts or task results.""" + + raw = json.dumps(_payload(protocol, (1, 0))) + models = [_Model([raw]), _Model([raw])] + for transport, model in zip(TRANSPORTS, models, strict=True): + assert _execute(protocol, model, monkeypatch, transport=transport) == expected + assert len(model.calls) == 1 + (chat_messages, chat_kwargs), (responses_messages, responses_kwargs) = ( + model.calls[0] for model in models + ) + assert json.dumps(chat_messages, ensure_ascii=False).encode("utf-8") == ( + json.dumps(responses_messages, ensure_ascii=False).encode("utf-8") + ) + assert chat_kwargs["response_format"] == {"type": "json_object"} + assert responses_kwargs["response_format"] == { + "type": "json_schema", + "json_schema": {"name": "semantic_batch", "schema": _expected_schema(protocol)}, + } + assert {key: value for key, value in chat_kwargs.items() if key != "response_format"} == { + key: value for key, value in responses_kwargs.items() if key != "response_format" + } + + +@pytest.mark.parametrize("protocol", PROTOCOLS) +@pytest.mark.parametrize("transport", TRANSPORTS) +@pytest.mark.parametrize("defect", ("missing", "duplicate", "unknown", "extra")) +@pytest.mark.parametrize("needs_repair", (False, True)) +def test_full_task_id_coverage_is_required_even_after_delimiter_repair( + protocol: str, transport: str, defect: str, needs_repair: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Syntax repair cannot invent, discard, or reassign task identities.""" + + indices = {"missing": (0,), "duplicate": (0, 0), "unknown": (0, 9), "extra": (0, 1, 9)}[defect] + raw = json.dumps(_payload(protocol, indices)) + model = _Model([raw[:-1] if needs_repair else raw]) + with pytest.raises(ValueError, match="invalid structured output after 1 attempt"): + _execute(protocol, model, monkeypatch, transport=transport) + assert len(model.calls) == 1 + + +@pytest.mark.parametrize("protocol", PROTOCOLS) +@pytest.mark.parametrize("transport", TRANSPORTS) +def test_task_ids_must_match_each_request_not_just_the_global_batch( + protocol: str, transport: str, monkeypatch: pytest.MonkeyPatch, +) -> None: + """Globally valid IDs cannot be swapped across independently chunked prompts.""" + + model = _Model([json.dumps(_payload(protocol, (1,))), json.dumps(_payload(protocol, (0,)))]) + with pytest.raises(ValueError, match="unknown task IDs"): + _execute(protocol, model, monkeypatch, transport=transport, batch_size=1) + assert len(model.calls[0][0]) == 2 + + +@pytest.mark.parametrize("transport", TRANSPORTS) +@pytest.mark.parametrize("needs_repair", (False, True)) +@pytest.mark.parametrize( + ("selected_ids", "message"), + ( + (["candidate_99"], "unknown IDs"), + (["candidate_3"], "unknown IDs"), + (["candidate_0", "candidate_0"], "unique"), + (["candidate_0", "candidate_1", "candidate_2"], "more than k=2"), + ([0], "list of strings"), + ), +) +def test_listwise_candidates_remain_task_local_unique_and_cardinality_bounded( + transport: str, needs_repair: bool, selected_ids: list[Any], message: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A repaired selection must still match the complete candidate contract.""" + + payload = _payload("listwise") + payload["results"][0]["selected_ids"] = selected_ids + raw = json.dumps(payload) + if not needs_repair: + with pytest.raises(ValueError, match=message): + _execute("listwise", _Model([raw]), monkeypatch, transport=transport) + else: + with pytest.raises(ValueError, match="invalid structured output after 1 attempt"): + _execute("listwise", _Model([raw[:-1]]), monkeypatch, transport=transport) + + +@pytest.mark.parametrize("transport", TRANSPORTS) +@pytest.mark.parametrize("shape", ("object", "array")) +def test_bounded_repair_rejects_extra_fields_without_tightening_valid_parser( + transport: str, shape: Literal["object", "array"], monkeypatch: pytest.MonkeyPatch, +) -> None: + """The strict repair validator does not replace the established normal parser.""" + + payload = _payload(shape) + output = payload["results"][0]["output"] + if shape == "object": + output["extra"] = "must not be silently dropped during repair" + else: + output["rows"][0]["extra"] = "must not be silently dropped during repair" + raw = json.dumps(payload) + parsed = _parse_structured_prompt_batch( + raw, output_cols=OUTPUT_COLS, shape=shape, + require_explanation=False, operator="sem_map" if shape == "object" else "sem_flat_map", + ) + expected = {"label": "value_0"} if shape == "object" else [{"label": "value_0"}] + assert parsed.items[0].value == (expected, None) + assert _execute(shape, _Model([raw]), monkeypatch, transport=transport)[0] == expected + with pytest.raises(ValueError, match="invalid structured output after 1 attempt"): + _execute(shape, _Model([raw[:-1]]), monkeypatch, transport=transport) + + +@pytest.mark.parametrize("protocol", PROTOCOLS) +@pytest.mark.parametrize("transport", TRANSPORTS) +@pytest.mark.parametrize( + "completion", + ({"finish_reason": "length"}, {"status": "incomplete", "incomplete_reason": "max_output_tokens"}), +) +def test_known_truncation_metadata_blocks_otherwise_valid_structural_repair( + protocol: str, transport: str, completion: dict[str, Any], + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + """Complete scalar tokens alone are not permission to repair known truncation.""" + + import agent_memory.adapters.lotus.prompt_batching as batching_module + + raw = json.dumps(_payload(protocol))[:-1] + _execute(protocol, _Model([raw]), monkeypatch, transport=transport) + + def forbidden_repair(*args: Any, **kwargs: Any) -> Any: + pytest.fail("known truncation reached structural repair") + + monkeypatch.setattr(batching_module, "repair_json_structure", forbidden_repair) + model = _Model([raw], [completion]) + with pytest.raises(ValueError, match="incomplete structured response"): + _execute(protocol, model, monkeypatch, transport=transport, trace_dir=tmp_path) + events = [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()] + failures = [event for event in events if event["event_type"] == "prompt_batching_failure"] + assert len(failures) == 1 + assert failures[0]["task_count"] == 2 + assert failures[0]["repair_result"] == "not-accepted" + assert not any(event["event_type"] == "prompt_batching" for event in events) + + +@pytest.mark.parametrize("protocol", PROTOCOLS) +@pytest.mark.parametrize("transport", TRANSPORTS) +def test_repair_accounting_counts_requests_separately_from_affected_tasks( + protocol: str, transport: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + """Two repaired requests covering four tasks must not be reported as four repairs.""" + + model = _Model([ + json.dumps(_payload(protocol, (0, 1)))[:-1], + json.dumps(_payload(protocol, (2, 3)))[:-1], + json.dumps(_payload(protocol, (4,))), + ]) + result = _execute( + protocol, model, monkeypatch, transport=transport, + count=5, batch_size=2, trace_dir=tmp_path, + ) + assert len(result) == 5 + events = [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()] + [event] = [event for event in events if event["event_type"] == "prompt_batching"] + assert event["task_count"] == 5 + assert event["prompt_count"] == 3 + assert event["chunk_sizes"] == [2, 2, 1] + assert event["retry_count"] == 0 + assert event["structured_output_repair_count"] == 2 + assert event["structured_output_repaired_task_count"] == 4 + assert event["syntax_repair_count"] == 2 + assert event["structured_output_repair_version"] == "bounded-json-v2" + repairs = event["structured_output_repairs"] + assert len(repairs) == 2 + for repair in repairs: + assert repair["affected_task_count"] == 2 + assert repair["attempt"] == 1 + assert repair["method"] == "bounded-json-syntax" + assert repair["version"] == "bounded-json-v2" + assert repair["edits"] + + +@pytest.mark.parametrize("protocol", PROTOCOLS) +@pytest.mark.parametrize("transport", TRANSPORTS) +def test_multiple_closures_share_the_existing_protocol_validator( + protocol: str, transport: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + payload = _payload(protocol) + field = next(iter(payload)) + raw = "{" + json.dumps(field) + ":[" + ",".join( + json.dumps(item)[:-1] for item in payload[field] + ) + "]}" + expected = _execute(protocol, _Model([json.dumps(payload)]), monkeypatch, + transport=transport) + model = _Model([raw]) + actual = _execute(protocol, model, monkeypatch, transport=transport, + trace_dir=tmp_path) + assert actual == expected + assert len(model.calls) == 1 + events = [json.loads(line) for line in (tmp_path / "events.jsonl").read_text().splitlines()] + [event] = [event for event in events if event["event_type"] == "prompt_batching"] + assert event["retry_count"] == 0 + assert event["structured_output_repair_count"] == 1 + assert event["structured_output_repaired_task_count"] == 2 + assert len(event["structured_output_repairs"][0]["edits"]) == 2 + + +def test_provider_usage_trace_preserves_truncated_exception_response(tmp_path: Path) -> None: + """Persist partial failure output and native usage without fabricating response fields.""" + + from agent_memory.tracing.semantic import write_provider_usage_trace + + partial_output = '{"results":[\n{"task_id":"task_0","output":{"label":"unfinished' + native_usage = { + "input_tokens": 11, + "input_tokens_details": {"cached_tokens": 3}, + "output_tokens": 7, + "output_tokens_details": {"reasoning_tokens": 2}, + "total_tokens": 18, + } + failure = SimpleNamespace( + response=SimpleNamespace(choices=[ + SimpleNamespace(message=SimpleNamespace(content=partial_output)), + ]), + provider_response_metadata={ + "transport": "responses-json-schema", + "status": "incomplete", + "finish_reason": "length", + "incomplete_reason": "max_output_tokens", + "raw_usage": native_usage, + }, + usage=SimpleNamespace( + prompt_tokens=11, + completion_tokens=7, + total_tokens=18, + prompt_cache_hit_tokens=3, + ), + ) + rows = write_provider_usage_trace( + tmp_path, + model="deepseek/deepseek-v4-flash", + responses=[failure, SimpleNamespace()], + ) + persisted = [ + json.loads(line) + for line in (tmp_path / "events.jsonl").read_text().splitlines() + ] + assert persisted == rows + event, missing = persisted + assert event["event_type"] == "provider_usage" + assert event["provider_finish_reason"] == "length" + assert event["provider_response_status"] == "incomplete" + assert event["provider_incomplete_reason"] == "max_output_tokens" + assert event["provider_response_model"] is None + assert event["provider_usage_available"] is True + assert event["provider_prompt_tokens"] == 11 + assert event["provider_completion_tokens"] == 7 + assert event["provider_total_tokens"] == 18 + assert event["provider_prompt_cache_hit_tokens"] == 3 + assert json.loads(Path(event["provider_raw_usage_path"]).read_text()) == native_usage + assert json.loads(Path(event["provider_partial_output_path"]).read_text()) == { + "output": partial_output, + } + assert "raw_output_path" not in event + for field in ( + "provider_finish_reason", "provider_response_status", + "provider_incomplete_reason", "provider_response_model", + ): + assert missing[field] is None + assert missing["provider_usage_available"] is False + assert "provider_raw_usage_path" not in missing + assert "provider_partial_output_path" not in missing + + +def _executor(config: LotusExecutionConfig) -> PolicyExecutor: + import agent_memory as am + from agent_memory.planner import PolicyDifferentiator + + policy = PolicyDifferentiator().differentiate(am.ClaudeMemory.spec()) + return PolicyExecutor(policy, adapter=LotusAdapter(config=config)) + + +def test_default_nonbatched_checkpoint_identity_remains_compatible() -> None: + """Explicit default transport retains the old empty physical identity.""" + + snapshot = _executor(LotusExecutionConfig()).snapshot_state() + assert "adapter_execution_fingerprint" not in snapshot + restored = _executor(LotusExecutionConfig(structured_output_transport="chat-json-object")) + restored.restore_state(snapshot) + assert "adapter_execution_fingerprint" not in restored.snapshot_state() + + +@pytest.mark.parametrize("max_tasks", (None, 2)) +@pytest.mark.parametrize("legacy_version", (None, "bounded-json-v1")) +def test_pre_repair_prompt_batching_checkpoint_is_rejected( + max_tasks: int | None, legacy_version: str | None, +) -> None: + """Old parsing contracts must not share a checkpoint with multi-edit repair.""" + + batching = PromptBatching(max_tasks=max_tasks) + config = LotusExecutionConfig(prompt_batching=batching) + executor = _executor(config) + snapshot = executor.snapshot_state() + legacy_contract: dict[str, Any] = {"max_tasks": max_tasks} + if legacy_version is not None: + legacy_contract["repair_version"] = legacy_version + legacy_batch_fingerprint = sha256( + json.dumps(legacy_contract, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + legacy_physical_fingerprint = sha256( + f"base_execution_fingerprint=\nprompt_batching={legacy_batch_fingerprint}".encode("utf-8") + ).hexdigest() + assert batching.fingerprint != legacy_batch_fingerprint + assert snapshot["adapter_execution_fingerprint"] != legacy_physical_fingerprint + _executor(config).restore_state(snapshot) + legacy_snapshot = {**snapshot, "adapter_execution_fingerprint": legacy_physical_fingerprint} + with pytest.raises(ValueError, match="adapter execution fingerprint does not match"): + executor.restore_state(legacy_snapshot) + + +@pytest.mark.parametrize("batching", (None, PromptBatching(max_tasks=2))) +@pytest.mark.parametrize("source_transport", TRANSPORTS) +def test_chat_and_responses_checkpoints_cannot_cross_physical_transports( + batching: PromptBatching | None, source_transport: str, +) -> None: + """Transport changes isolate physical state even when the logical policy is identical.""" + + target_transport = next(transport for transport in TRANSPORTS if transport != source_transport) + source_config = LotusExecutionConfig( + prompt_batching=batching, structured_output_transport=source_transport, + ) + source = _executor(source_config) + target = _executor(LotusExecutionConfig( + prompt_batching=batching, structured_output_transport=target_transport, + )) + snapshot = source.snapshot_state() + assert source.policy.fingerprint == target.policy.fingerprint + _executor(source_config).restore_state(snapshot) + with pytest.raises(ValueError, match="adapter execution fingerprint does not match"): + target.restore_state(snapshot) diff --git a/tests/test_structured_transport_benchmark.py b/tests/test_structured_transport_benchmark.py new file mode 100644 index 0000000..23fbba3 --- /dev/null +++ b/tests/test_structured_transport_benchmark.py @@ -0,0 +1,431 @@ +"""Provider-free benchmark transport wiring and checkpoint identity tests.""" + +from __future__ import annotations + +from pathlib import Path +from types import ModuleType, SimpleNamespace +from typing import Any + +import pytest + +from agent_memory.adapters.lotus.prompt_batching import PromptBatching +from agent_memory.evaluation import agent_memory_drivers as drivers +from agent_memory.evaluation import run as run_module +from agent_memory.evaluation.bundle import BenchmarkBundle +from agent_memory.evaluation.types import BenchmarkCase, BenchmarkEvent, BenchmarkQuestion +from tools.evaluation import locomo, longmemeval, memory_agent_bench + + +DEFAULT_TRANSPORT = "chat-json-object" +RESPONSES_TRANSPORT = "responses-json-schema" +CLI_MODULES = (locomo, longmemeval, memory_agent_bench) +FACTORIES = ( + drivers.ClaudeMemoryDriverFactory, + drivers.ZepMemoryDriverFactory, + drivers.Mem0MemoryDriverFactory, + drivers.Mem0MemoryEnhancedDriverFactory, +) + + +def _bundle(benchmark_id: str = "locomo") -> BenchmarkBundle: + return BenchmarkBundle( + benchmark_id, + "revision", + "sha256", + (BenchmarkCase( + "case", "task", + (BenchmarkEvent("case", "event", "user", "I like tea."),), + (BenchmarkQuestion("question", "case", "What drink?", "tea", ("event",)),), + ),), + {"run_mode": "integration-smoke"}, + ) + + +def _cli_args(tmp_path: Path) -> list[str]: + return [ + "run", "--bundle-dir", str(tmp_path / "bundle"), + "--output-dir", str(tmp_path / "output"), + "--system", "claude-memory", + ] + + +def _factory_options(factory: type[Any]) -> dict[str, Any]: + if factory is drivers.ClaudeMemoryDriverFactory: + return {} + options: dict[str, Any] = {"base_namespace": "transport-test"} + if factory is drivers.ZepMemoryDriverFactory: + options.update( + connector=SimpleNamespace(embedding_provider=None), + neo4j_image="neo4j:5.26.2", + neo4j_image_digest="sha256:test", + ) + return options + + +@pytest.mark.parametrize("module", CLI_MODULES) +@pytest.mark.parametrize("transport", (None, DEFAULT_TRANSPORT, RESPONSES_TRANSPORT)) +def test_cli_forwards_transport_without_changing_default_calls( + module: ModuleType, + transport: str | None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Each CLI forwards opt-in transport while keeping old runner signatures valid.""" + + captured: dict[str, Any] = {} + + def run(**kwargs: Any) -> Path: + captured.update(kwargs) + return kwargs["output_dir"] + + benchmark_id = { + locomo: "locomo", + longmemeval: "longmemeval-v1-cleaned-s", + memory_agent_bench: "memory-agent-bench", + }[module] + monkeypatch.setattr(module, "read_bundle", lambda _path: _bundle(benchmark_id)) + monkeypatch.setattr(module, "run_agent_memory_bundle", run) + if module is memory_agent_bench: + monkeypatch.setattr(module, "download_movie_entity_mapping", lambda path: path) + monkeypatch.setattr(module, "load_movie_entity_mapping", lambda _path: {}) + monkeypatch.setattr(module, "memory_agent_bench_task_contracts", lambda **kw: {}) + args = _cli_args(tmp_path) + if module is longmemeval: + args.append("--maintenance-only") + if transport is not None: + args.extend(("--structured-output-transport", transport)) + + parsed = module.parse_args(args) + assert parsed.structured_output_transport == (transport or DEFAULT_TRANSPORT) + assert module.main(args) == tmp_path / "output" + if transport == RESPONSES_TRANSPORT: + assert captured["structured_output_transport"] == RESPONSES_TRANSPORT + else: + assert "structured_output_transport" not in captured + + +@pytest.mark.parametrize("module", CLI_MODULES) +@pytest.mark.parametrize( + ("transport", "model"), + (("invalid", "deepseek/deepseek-v4-flash"), + (RESPONSES_TRANSPORT, "openai/gpt-4o"), + (RESPONSES_TRANSPORT, "deepseek/deepseek-chat")), +) +def test_cli_rejects_transport_before_loading_inputs( + module: ModuleType, + transport: str, + model: str, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Invalid transport/provider combinations fail before bundle or dataset I/O.""" + + def unexpected(*args: Any, **kwargs: Any) -> Any: + pytest.fail("invalid transport reached input loading") + + monkeypatch.setattr(module, "read_bundle", unexpected) + model_flag = "--model" if module is memory_agent_bench else "--memory-model" + with pytest.raises(SystemExit) as error: + module.main([ + *_cli_args(tmp_path), "--structured-output-transport", transport, + model_flag, model, + ]) + assert error.value.code == 2 + + +@pytest.mark.parametrize("factory", FACTORIES) +@pytest.mark.parametrize("transport", (None, DEFAULT_TRANSPORT, RESPONSES_TRANSPORT)) +@pytest.mark.parametrize("context_limit", (None, 1_000_000)) +def test_factory_passes_transport_into_execution_config( + context_limit: int | None, + factory: type[Any], + transport: str | None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """All four factory paths deliver physical settings to the adapter config.""" + + import agent_memory.adapters.lotus as lotus_module + import agent_memory.storage.qdrant as qdrant_module + + class ConfigCaptured(Exception): + """Stop before any adapter or runtime can initialize a model.""" + + captured: dict[str, Any] = {} + + def adapter(**kwargs: Any) -> Any: + captured.update(kwargs) + raise ConfigCaptured + + monkeypatch.setattr(lotus_module, "LotusAdapter", adapter) + monkeypatch.setattr( + qdrant_module, "SentenceTransformerEmbeddingProvider", + lambda *args, **kwargs: SimpleNamespace(), + ) + monkeypatch.setattr( + qdrant_module, "QdrantConnector", + lambda **kwargs: SimpleNamespace(close=lambda: None), + ) + options = _factory_options(factory) + if context_limit is not None: + options["lm_max_ctx_len"] = context_limit + if transport is not None: + options["structured_output_transport"] = transport + with pytest.raises(ConfigCaptured): + factory(**options)("case", tmp_path / "state", tmp_path / "trace") + assert captured["config"].lm_model_kwargs.get("max_ctx_len") == context_limit + assert captured["config"].structured_max_tokens == 32_768 + assert captured["config"].structured_output_transport == ( + transport or DEFAULT_TRANSPORT + ) + + +@pytest.mark.parametrize("factory", FACTORIES) +@pytest.mark.parametrize( + ("transport", "model", "message"), + (("invalid", "deepseek/deepseek-v4-flash", "structured_output_transport"), + (RESPONSES_TRANSPORT, "openai/gpt-4o", "DeepSeek Responses model"), + (RESPONSES_TRANSPORT, "deepseek/deepseek-chat", "DeepSeek Responses model")), +) +def test_factory_validates_transport_at_construction( + factory: type[Any], transport: str, model: str, message: str, +) -> None: + """Bad execution modes must not be deferred until case-local storage setup.""" + + with pytest.raises(ValueError, match=message): + factory( + **_factory_options(factory), model_id=model, + structured_output_transport=transport, + ) + + +@pytest.mark.parametrize("transport", (DEFAULT_TRANSPORT, RESPONSES_TRANSPORT)) +def test_zep_from_environment_forwards_transport( + transport: str, monkeypatch: pytest.MonkeyPatch, +) -> None: + """The environment constructor retains transport when handing off its connector.""" + + import agent_memory.storage.neo4j as neo4j_module + + for name, value in { + "AGENT_MEMORY_NEO4J_IMAGE": "neo4j:5.26.2", + "AGENT_MEMORY_NEO4J_IMAGE_DIGEST": "sha256:test", + "AGENT_MEMORY_NEO4J_URI": "bolt://unused", + "AGENT_MEMORY_NEO4J_PASSWORD": "unused", + }.items(): + monkeypatch.setenv(name, value) + for name in ( + "SentenceTransformerEmbeddingProvider", "SentenceTransformerCrossEncoderProvider", + ): + monkeypatch.setattr(neo4j_module, name, lambda *args, **kwargs: None) + monkeypatch.setattr(neo4j_module, "Neo4jConnector", lambda **kw: SimpleNamespace(**kw)) + + factory = drivers.ZepMemoryDriverFactory.from_environment( + base_namespace="transport-test", structured_output_transport=transport, + ) + assert factory.structured_output_transport == transport + + +@pytest.mark.parametrize( + ("transport", "model", "message"), + (("invalid", "deepseek/deepseek-v4-flash", "structured_output_transport"), + (RESPONSES_TRANSPORT, "openai/gpt-4o", "DeepSeek Responses model")), +) +def test_zep_from_environment_validates_before_external_initialization( + transport: str, model: str, message: str, monkeypatch: pytest.MonkeyPatch, +) -> None: + """Transport errors take precedence even when deployment credentials are absent.""" + + import agent_memory.storage.neo4j as neo4j_module + + def unexpected(*args: Any, **kwargs: Any) -> Any: + pytest.fail("invalid transport initialized external storage or models") + + monkeypatch.delenv("AGENT_MEMORY_NEO4J_IMAGE", raising=False) + monkeypatch.delenv("AGENT_MEMORY_NEO4J_IMAGE_DIGEST", raising=False) + for name in ( + "Neo4jConnector", "SentenceTransformerEmbeddingProvider", + "SentenceTransformerCrossEncoderProvider", + ): + monkeypatch.setattr(neo4j_module, name, unexpected) + with pytest.raises(ValueError, match=message): + drivers.ZepMemoryDriverFactory.from_environment( + base_namespace="transport-test", model_id=model, + structured_output_transport=transport, + ) + + +@pytest.mark.parametrize("system_id", run_module.AGENT_MEMORY_SYSTEMS) +@pytest.mark.parametrize("transport", (None, DEFAULT_TRANSPORT, RESPONSES_TRANSPORT)) +@pytest.mark.parametrize("batching", (None, PromptBatching(max_tasks=8))) +@pytest.mark.parametrize("context_limit", (None, 1_000_000)) +def test_run_records_only_enabled_execution_contracts( + context_limit: int | None, + system_id: str, + transport: str | None, + batching: PromptBatching | None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + """Manifest provenance and maintenance identity isolate opt-in behavior only.""" + + captured: dict[str, Any] = {} + + class FakeFactory: + """Capture factory inputs without opening benchmark storage.""" + + def __init__(self, **kwargs: Any) -> None: + if transport != RESPONSES_TRANSPORT: + assert "structured_output_transport" not in kwargs + captured["factory"] = kwargs + + @classmethod + def from_environment(cls, **kwargs: Any) -> FakeFactory: + """Capture the Zep constructor path as well.""" + return cls(**kwargs) + + def runtime_provenance(self) -> dict[str, Any]: + """Avoid provider-specific version and storage queries.""" + return {} + + class FakeRunner: + """Capture the manifest inputs without executing benchmark cases.""" + + def __init__(self, **kwargs: Any) -> None: + captured["runner"] = kwargs + + def run(self, bundle: BenchmarkBundle) -> None: + """No model calls are needed to inspect execution identity.""" + + for factory in FACTORIES: + monkeypatch.setattr(run_module, factory.__name__, FakeFactory) + monkeypatch.setattr(run_module, "BenchmarkRunner", FakeRunner) + monkeypatch.setattr(run_module, "LiteLLMBenchmarkModel", lambda **kw: None) + monkeypatch.setattr(run_module, "_require_environment", lambda _system: None) + monkeypatch.setattr( + run_module, "collect_runtime_provenance", + lambda *args, **kwargs: {"source": {}, "runtime": {}}, + ) + monkeypatch.setattr( + run_module, "validate_run_provenance", lambda *args, **kwargs: None, + ) + options: dict[str, Any] = {} if transport is None else {"structured_output_transport": transport} + if context_limit is not None: + options["lm_max_ctx_len"] = context_limit + run_module.run_agent_memory_bundle( + bundle=_bundle(), contracts={}, system_id=system_id, + output_dir=tmp_path / "output", memory_thinking_enabled=False, + prompt_batching=batching, **options, + ) + + provenance = captured["runner"]["runtime_provenance"]["runtime"]["lotus_execution"] + identity = captured["runner"]["system_contract"].maintenance_execution_id + parts = ["sem-join-topk:listwise"] if system_id == "zep-memory" else [] + if batching is not None: + assert provenance["prompt_batching"] == {"max_tasks": 8} + assert provenance["json_repair_version"] == "bounded-json-v2" + parts.append(f"prompt-batching:{batching.fingerprint}") + else: + assert "prompt_batching" not in provenance + assert "json_repair_version" not in provenance + if transport == RESPONSES_TRANSPORT: + assert captured["factory"]["structured_output_transport"] == transport + assert provenance["structured_output_transport"] == transport + parts.append("structured-output-transport:responses-json-schema") + else: + assert "structured_output_transport" not in provenance + if context_limit is not None: + assert captured["factory"]["lm_max_ctx_len"] == context_limit + assert provenance["lm_max_ctx_len"] == context_limit + parts.append(f"lm-max-ctx-len:{context_limit}") + else: + assert "lm_max_ctx_len" not in captured["factory"] + assert "lm_max_ctx_len" not in provenance + assert identity == "|".join(parts) + + +@pytest.mark.parametrize("system_id", run_module.AGENT_MEMORY_SYSTEMS) +@pytest.mark.parametrize( + ("transport", "model", "message"), + (("invalid", "deepseek/deepseek-v4-flash", "structured_output_transport"), + (RESPONSES_TRANSPORT, "openai/gpt-4o", "DeepSeek Responses model")), +) +def test_run_validates_transport_before_environment_or_factories( + system_id: str, transport: str, model: str, message: str, + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + """Invalid modes cannot reach provenance collection, model or storage setup.""" + + def unexpected(*args: Any, **kwargs: Any) -> Any: + pytest.fail("invalid transport reached external initialization") + + for name in ( + "_require_environment", "collect_runtime_provenance", "LiteLLMBenchmarkModel", + *(factory.__name__ for factory in FACTORIES), + ): + monkeypatch.setattr(run_module, name, unexpected) + with pytest.raises(ValueError, match=message): + run_module.run_agent_memory_bundle( + bundle=_bundle(), contracts={}, system_id=system_id, + output_dir=tmp_path / "output", memory_provider_model_id=model, + structured_output_transport=transport, + ) + + +@pytest.mark.parametrize("limit", (None, 1_000_000)) +def test_locomo_context_limit_is_opt_in( + limit: int | None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, +) -> None: + captured: dict[str, Any] = {} + + def run(**kwargs: Any) -> Path: + captured.update(kwargs) + return kwargs["output_dir"] + + monkeypatch.setattr(locomo, "read_bundle", lambda _path: _bundle()) + monkeypatch.setattr(locomo, "run_agent_memory_bundle", run) + args = _cli_args(tmp_path) + if limit is not None: + args.extend(("--lm-max-ctx-len", str(limit))) + locomo.main(args) + assert captured.get("lm_max_ctx_len") == limit + if limit is None: + assert "lm_max_ctx_len" not in captured + + +@pytest.mark.parametrize("limit", (0, -1, True, 1.5)) +def test_invalid_context_limit_fails_before_resources(limit: Any) -> None: + from agent_memory.adapters.lotus.context import LotusExecutionConfig + + with pytest.raises(ValueError, match="lm_max_ctx_len"): + LotusExecutionConfig(lm_model_kwargs={"max_ctx_len": limit}) + for factory in FACTORIES: + with pytest.raises(ValueError, match="lm_max_ctx_len"): + factory(**_factory_options(factory), lm_max_ctx_len=limit) + with pytest.raises(ValueError, match="lm_max_ctx_len"): + drivers.ZepMemoryDriverFactory.from_environment( + base_namespace="unused", lm_max_ctx_len=limit, + ) + + +def test_context_override_reaches_lotus_and_changes_execution_identity( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import lotus + from agent_memory.adapters.lotus import LotusAdapter + from agent_memory.adapters.lotus.context import LotusExecutionConfig, LotusExecutionContext + + monkeypatch.setattr(lotus.settings, "lm", lotus.settings.lm) + monkeypatch.setattr(lotus.settings, "enable_cache", lotus.settings.enable_cache) + config = LotusExecutionConfig(lm_model_kwargs={"max_ctx_len": 1_000_000}) + context = LotusExecutionContext(model="test/model", config=config) + context.configure() + assert lotus.settings.lm.max_ctx_len == 1_000_000 + assert "max_ctx_len" not in lotus.settings.lm.kwargs + fingerprints = { + LotusAdapter(model="test/model", config=LotusExecutionConfig( + lm_model_kwargs=kwargs, + )).maintenance_execution_fingerprint + for kwargs in ({}, {"max_ctx_len": 128_000}, {"max_ctx_len": 1_000_000}) + } + assert len(fingerprints) == 3 diff --git a/tests/test_zep_locomo_benchmark.py b/tests/test_zep_locomo_benchmark.py index db313b0..a097a59 100644 --- a/tests/test_zep_locomo_benchmark.py +++ b/tests/test_zep_locomo_benchmark.py @@ -96,7 +96,7 @@ def test_cli_has_explicit_smoke_and_full_benchmark_defaults(tmp_path: Path) -> N assert (benchmark.start_row, benchmark.row_limit) == (1, None) assert (benchmark.question_start, benchmark.question_limit) == (1, None) assert benchmark.include_adversarial is True - assert smoke.grouped_agg_rule == "rule-re-group" + assert smoke.grouped_agg_rule == "rule-join-map" @pytest.mark.parametrize("grouped_agg_rule", GROUPED_AGG_RULES) diff --git a/tools/analysis/prompt_batching_scale_report.py b/tools/analysis/prompt_batching_scale_report.py new file mode 100644 index 0000000..fd431d4 --- /dev/null +++ b/tools/analysis/prompt_batching_scale_report.py @@ -0,0 +1,466 @@ +"""Summarize and plot the 128-event prompt-batching scale experiment.""" + +from __future__ import annotations + +import argparse +import csv +from dataclasses import dataclass, replace +import html +import json +import math +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence + + +EVENT_COUNT = 128 +QUESTION_COUNT = 71 +EXPECTED_LABELS = ( + "original", + "packed-1", + "packed-2", + "packed-4", + "packed-8", + "packed-16", + "packed-32", + "packed-64", + "packed-128", + "packed-all", +) +PASS_COLORS = {1: "#0F766E", 2: "#D97706"} + + +@dataclass(frozen=True) +class RunRow: + """Normalized metrics for one completed condition.""" + + run_id: str + pass_index: int + label: str + refresh_every: int + prompt_batch: int | str | None + provider_calls: int + physical_tokens: int + cache_hit_tokens: int + cache_miss_tokens: int + completion_tokens: int + cost_cny: float + insertion_mean_ms: float + insertion_median_ms: float + insertion_p95_ms: float + insertion_max_ms: float + condition_wall_seconds: float + retrieval_mean_ms: float + official_score: float + zep_score: float + topic_count: int | None + catalog_count: int | None + topic_digest: str | None + catalog_digest: str | None + artifact_bytes: int + + def to_csv(self) -> dict[str, Any]: + """Return one flat CSV row with explicit units.""" + + return { + "run_id": self.run_id, + "pass": self.pass_index, + "condition": self.label, + "refresh_every": self.refresh_every, + "prompt_batch": self.prompt_batch, + "provider_calls": self.provider_calls, + "physical_tokens": self.physical_tokens, + "physical_tokens_per_event": self.physical_tokens / EVENT_COUNT, + "cache_hit_tokens": self.cache_hit_tokens, + "cache_miss_tokens": self.cache_miss_tokens, + "completion_tokens": self.completion_tokens, + "cost_cny": self.cost_cny, + "cost_cny_per_event": self.cost_cny / EVENT_COUNT, + "insertion_excluding_trace_mean_s": self.insertion_mean_ms / 1000, + "insertion_excluding_trace_median_s": self.insertion_median_ms / 1000, + "insertion_excluding_trace_p95_s": self.insertion_p95_ms / 1000, + "insertion_excluding_trace_max_s": self.insertion_max_ms / 1000, + "condition_end_to_end_wall_s": self.condition_wall_seconds, + "condition_end_to_end_wall_s_per_event": ( + self.condition_wall_seconds / EVENT_COUNT + ), + "retrieval_mean_s": self.retrieval_mean_ms / 1000, + "official_locomo_percent": self.official_score * 100, + "zep_judge_percent": self.zep_score * 100, + "topic_count": self.topic_count, + "catalog_count": self.catalog_count, + "topic_digest": self.topic_digest, + "catalog_digest": self.catalog_digest, + "artifact_bytes": self.artifact_bytes, + } + + +def _required_number(value: Any, *, field: str) -> float: + if not isinstance(value, int | float) or isinstance(value, bool): + raise ValueError(f"{field} must be numeric") + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{field} must be finite") + return result + + +def _state_count(state: Mapping[str, Any], relation: str) -> int | None: + row = state.get(relation) + if not isinstance(row, Mapping): + return None + value = row.get("row_count") + return int(value) if isinstance(value, int) and not isinstance(value, bool) else None + + +def _state_digest(state: Mapping[str, Any], relation: str) -> str | None: + row = state.get(relation) + if not isinstance(row, Mapping): + return None + value = row.get("multiset_digest") + return str(value) if isinstance(value, str) else None + + +def _load_run(path: Path) -> tuple[RunRow, dict[str, Any]]: + payload = json.loads(path.read_text()) + state = payload.get("final_state") + if not isinstance(state, Mapping): + raise ValueError(f"missing final_state: {path}") + run_id = str(payload.get("run_id") or "") + label = str(payload.get("label") or "") + pass_index = int(payload.get("pass_index") or 0) + if not run_id or label not in EXPECTED_LABELS or pass_index not in (1, 2): + raise ValueError(f"invalid run identity: {path}") + row = RunRow( + run_id=run_id, + pass_index=pass_index, + label=label, + refresh_every=int(payload["refresh_every"]), + prompt_batch=payload.get("prompt_batch"), + provider_calls=int(payload["provider_call_count"]), + physical_tokens=int(payload["physical_tokens"]), + cache_hit_tokens=int(payload["cache_hit_tokens"]), + cache_miss_tokens=int(payload["cache_miss_tokens"]), + completion_tokens=int(payload["completion_tokens"]), + cost_cny=_required_number(payload["cost_cny"], field="cost_cny"), + insertion_mean_ms=_required_number( + payload["insertion_excluding_trace_mean_ms"], + field="insertion_excluding_trace_mean_ms", + ), + insertion_median_ms=_required_number( + payload["insertion_excluding_trace_median_ms"], + field="insertion_excluding_trace_median_ms", + ), + insertion_p95_ms=_required_number( + payload["insertion_excluding_trace_p95_ms"], + field="insertion_excluding_trace_p95_ms", + ), + insertion_max_ms=_required_number( + payload["insertion_excluding_trace_max_ms"], + field="insertion_excluding_trace_max_ms", + ), + condition_wall_seconds=_required_number( + payload["condition_end_to_end_wall_seconds"], + field="condition_end_to_end_wall_seconds", + ), + retrieval_mean_ms=_required_number( + payload["retrieval_mean_ms"], field="retrieval_mean_ms" + ), + official_score=_required_number( + payload["official_locomo_score"], field="official_locomo_score" + ), + zep_score=_required_number(payload["zep_judge_score"], field="zep_judge_score"), + topic_count=_state_count(state, "topics"), + catalog_count=_state_count(state, "catalog"), + topic_digest=_state_digest(state, "topics"), + catalog_digest=_state_digest(state, "catalog"), + artifact_bytes=int(payload["artifact_bytes"]), + ) + return row, payload + + +def load_completed_runs(root: Path) -> tuple[tuple[RunRow, ...], dict[str, dict[str, Any]]]: + """Load completed runs in the immutable experiment-contract order.""" + + contract_path = root / "control" / "experiment-contract.json" + contract = json.loads(contract_path.read_text()) + if contract.get("experiment") != "claude-prompt-batching-scale-128e": + raise ValueError(f"not a prompt-batching scale experiment: {root}") + rows: list[RunRow] = [] + raw: dict[str, dict[str, Any]] = {} + continuation = contract.get("continuation") or {} + reused = continuation.get("reused_conditions", {}) + for expected in contract.get("conditions") or (): + run_id = str(expected["run_id"]) + directory = Path(reused[run_id]) if run_id in reused else root / "conditions" / run_id + validation = directory / "validation.json" + if not validation.exists(): + continue + row, payload = _load_run(validation) + if run_id in reused: + row = replace(row, cost_cny=float(continuation["prior_condition_costs_cny"][run_id])) + payload["result_origin"] = { + "status": "reused-completed" if run_id in reused else "completed", + "directory": str(directory), + "source": ( + continuation.get("reused_sources", {}).get(run_id, continuation["source"]) + if run_id in reused else contract["source"] + ), + "repair_version": None if run_id in reused else (contract.get("structured_output") or {}).get("repair_version"), + } + if row.run_id != run_id: + raise ValueError(f"run ID mismatch: {validation}") + rows.append(row) + raw[run_id] = payload + return tuple(rows), raw + + +def _write_csv(path: Path, rows: Sequence[Mapping[str, Any]]) -> None: + if not rows: + path.write_text("", encoding="utf-8") + return + fields = list(rows[0]) + with path.open("w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + writer.writerows(rows) + + +def _site_rows(raw_runs: Mapping[str, Mapping[str, Any]]) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for run_id, run in raw_runs.items(): + prompt = run.get("prompt_batching") or {} + candidates = run.get("candidate_pairs") or {} + usage = run.get("operator_usage") or {} + for site in sorted(set(prompt) | set(candidates) | set(usage)): + batching = prompt.get(site) or {} + candidate = candidates.get(site) or {} + operator = usage.get(site) or {} + eligible = int(candidate.get("eligible") or 0) + selected = int(candidate.get("selected") or 0) + tasks = int(batching.get("tasks") or 0) + prompts = int(batching.get("prompts") or 0) + rows.append( + { + "run_id": run_id, + "pass": run["pass_index"], + "condition": run["label"], + "site": site, + "ready_tasks": tasks, + "physical_prompts": prompts, + "tasks_per_prompt": tasks / prompts if prompts else None, + "max_chunk": int(batching.get("max_chunk") or 0), + "retries": int(batching.get("retries") or 0), + "syntax_repairs": int(batching.get("repairs") or 0), + "provider_responses": int(operator.get("responses") or 0), + "physical_tokens": int(operator.get("tokens") or 0), + "tokens_per_task": ( + int(operator.get("tokens") or 0) / tasks if tasks else None + ), + "eligible_pairs": eligible, + "selected_pairs": selected, + "candidate_reduction": ( + 1 - selected / eligible if eligible else None + ), + } + ) + return rows + + +def _svg_document(width: int, height: int, body: str) -> str: + return ( + f'' + '' + '' + f"{body}\n" + ) + + +def _line_panels( + rows: Sequence[RunRow], + panels: Sequence[tuple[str, Callable[[RunRow], float], str]], +) -> str: + width = 1180 + panel_height = 245 + height = 45 + panel_height * len(panels) + parts = [ + '128-Message Batching Scale', + 'Pass 1', + 'Pass 2', + ] + for panel_index, (title, value, unit) in enumerate(panels): + top = 45 + panel_index * panel_height + left, right = 78, width - 30 + chart_top, bottom = top + 32, top + 190 + values = [value(row) for row in rows] + maximum = max(values, default=1.0) or 1.0 + parts.extend( + ( + f'{html.escape(title)}', + f'', + f'', + f'{maximum:.3g} {html.escape(unit)}', + f'0', + ) + ) + for pass_index in (1, 2): + pass_rows = [row for row in rows if row.pass_index == pass_index] + points: list[str] = [] + for index, label in enumerate(EXPECTED_LABELS): + x = left + index * (right - left) / (len(EXPECTED_LABELS) - 1) + matching = [row for row in pass_rows if row.label == label] + if matching: + y = bottom - value(matching[0]) / maximum * (bottom - chart_top) + points.append(f"{x:.2f},{y:.2f}") + parts.append( + f'' + ) + if panel_index == len(panels) - 1: + parts.append( + f'' + f"{html.escape(label)}" + ) + if len(points) > 1: + parts.append( + f'' + ) + return _svg_document(width, height, "".join(parts)) + + +def _pareto_svg(rows: Sequence[RunRow]) -> str: + width, height = 1000, 620 + left, right, top, bottom = 90, 960, 55, 535 + costs = [row.cost_cny / EVENT_COUNT for row in rows] + scores = [row.official_score * 100 for row in rows] + max_cost = max(costs, default=1.0) or 1.0 + low_score = min(scores, default=0.0) + high_score = max(scores, default=1.0) + score_span = max(high_score - low_score, 1.0) + body = [ + 'Cost-Latency-Quality Pareto View', + f'', + f'', + f'CNY per source event', + f'Official LOCOMO (%)', + ] + for row in rows: + x = left + (row.cost_cny / EVENT_COUNT) / max_cost * (right - left) + y = bottom - ((row.official_score * 100) - low_score) / score_span * (bottom - top) + radius = 5 + min(13, math.sqrt(max(row.insertion_mean_ms, 0)) / 5) + body.append( + f'' + ) + body.append( + f'' + f"{html.escape(row.label)} p{row.pass_index}" + ) + return _svg_document(width, height, "".join(body)) + + +def generate_report(root: Path) -> dict[str, Any]: + """Generate bounded reports without modifying benchmark artifacts.""" + + root = root.resolve() + output = root / "analysis" + output.mkdir(parents=True, exist_ok=True) + rows, raw = load_completed_runs(root) + csv_rows = [ + {**row.to_csv(), "result_status": raw[row.run_id]["result_origin"]["status"], + "source_snapshot_sha256": raw[row.run_id]["result_origin"]["source"].get("source_snapshot_sha256"), + "repair_version": raw[row.run_id]["result_origin"]["repair_version"]} + for row in rows + ] + _write_csv(output / "conditions.csv", csv_rows) + _write_csv(output / "site_metrics.csv", _site_rows(raw)) + + cross_pass: dict[str, Any] = {} + for label in EXPECTED_LABELS: + matches = [row for row in rows if row.label == label] + if len(matches) != 2: + continue + first, second = matches + cross_pass[label] = { + "same_source_snapshot": ( + raw[first.run_id]["result_origin"]["source"].get("source_snapshot_sha256") + == raw[second.run_id]["result_origin"]["source"].get("source_snapshot_sha256") + ), + "topic_state_equal": first.topic_digest == second.topic_digest, + "catalog_state_equal": first.catalog_digest == second.catalog_digest, + "official_score_delta": second.official_score - first.official_score, + "cost_cny_delta": second.cost_cny - first.cost_cny, + "insertion_mean_ms_delta": ( + second.insertion_mean_ms - first.insertion_mean_ms + ), + } + summary = { + "schema_version": 1, + "completed_conditions": len(rows), + "expected_conditions": 20, + "complete": len(rows) == 20, + "total_cost_cny": sum(row.cost_cny for row in rows), + "cross_pass": cross_pass, + "result_origins": {key: value["result_origin"] for key, value in raw.items()}, + "interpretation_boundary": ( + "R and P co-vary. Each pass is shown separately; two passes are an " + "exploratory repeatability check, not a variance estimate. Reused results " + "retain their original source; cross-version differences are not pure run noise." + ), + } + (output / "summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + (output / "efficiency.svg").write_text( + _line_panels( + rows, + ( + ( + "Physical tokens per source message", + lambda row: row.physical_tokens / EVENT_COUNT, + "tokens", + ), + ( + "Insertion mean excluding semantic trace I/O", + lambda row: row.insertion_mean_ms / 1000, + "seconds", + ), + ("Actual condition cost", lambda row: row.cost_cny, "CNY"), + ), + ), + encoding="utf-8", + ) + (output / "quality.svg").write_text( + _line_panels( + rows, + ( + ("Official LOCOMO", lambda row: row.official_score * 100, "%"), + ("Zep judge C1-C4", lambda row: row.zep_score * 100, "%"), + ("Topic count", lambda row: float(row.topic_count or 0), "rows"), + ("Catalog count", lambda row: float(row.catalog_count or 0), "rows"), + ), + ), + encoding="utf-8", + ) + (output / "pareto.svg").write_text(_pareto_svg(rows), encoding="utf-8") + return summary + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-root", type=Path, required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> None: + """Generate the requested report.""" + + args = _parser().parse_args(argv) + print(json.dumps(generate_report(args.run_root), sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/evaluation/README.md b/tools/evaluation/README.md index 77b1664..ae78fb8 100644 --- a/tools/evaluation/README.md +++ b/tools/evaluation/README.md @@ -13,12 +13,15 @@ Dataset adapters prepare evidence. Task contracts define prompts and scoring. 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` 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. +Claude and Zep runners default to `rule-join-map`. They continue to accept the +planner's other grouped-aggregate rules as explicit compatibility or ablation +conditions. If a selected rule cannot compile a policy, the run fails before +insertion instead of silently changing the requested rule. + +Mem0 does not contain a grouped aggregate, so its public benchmark CLI rejects +an explicit `--grouped-agg-rule`. The runner preserves its existing internal +`rule-all-group` identity only for benchmark compatibility and fingerprints; +Mem0 does not execute that grouped-aggregate rule. ## Install @@ -100,7 +103,7 @@ 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: +stable semantic predicate site IDs: ```json { @@ -135,6 +138,58 @@ Site-level config is mutually exclusive with the global `--semantic-pair-profile`, `--semantic-pair-top-k`, and `--semantic-pair-min-similarity` options. +Prompt construction is a separate physical layer from candidate selection. +Enable it for the whole LOTUS execution path with one option: + +```bash +--prompt-batch-size all +--prompt-batch-size 8 +``` + +Omitting the option preserves the existing LOTUS prompts. `all` places all +currently ready, independent tasks from one operator invocation into one +prompt, unless the model context requires deterministic chunks. A positive +integer caps each prompt at that many tasks. Search-Filter still runs first, so +only selected candidates enter a prompt. The setting covers semantic filters, +maps, flat maps, pairwise joins, pairwise group comparisons, multi-anchor +top-k joins, and semantic aggregate groups. A direct listwise `sem_topk` is +already one ranking task; delegated LOTUS pairwise ranking keeps its own +algorithm-specific comparison schedule. + +This option changes prompt construction, so it is explicit opt-in and enters +the maintenance/checkpoint identity. Operators validate every returned task ID +and output schema. Syntax-only JSON repairs are recorded; missing, duplicate, +unknown, or malformed task results fail instead of being guessed. + +Structural validation does not make prompt batching semantics-preserving. A +packed request can return valid JSON while making different semantic decisions +from the same tasks sent in smaller prompts. The current Claude/Mem0/Zep smoke +did not establish a universal best batch size, and larger tested batches caused +material quality loss in at least one policy. Treat every +`--prompt-batch-size` value as a separate experiment condition; do not use +`all` as a general default. + +Semantic aggregate provider dispatch remains a separate control: + +- `--sem-agg-dispatch provider-batched` keeps every per-group prompt unchanged + but submits ready prompts together through the LOTUS LM batch interface. + +`--sem-agg-dispatch provider-batched` and `--prompt-batch-size` are mutually +exclusive. The former sends several unchanged prompts in one provider call; +the latter puts several tasks inside one prompt. Both are disabled by default. + +`--refresh-every N` controls when source rows enter the maintenance DAG. Its +default is `1`, which preserves one eager refresh per event. A value greater +than one buffers normalized rows and publishes each full batch as one atomic +delta; the benchmark flushes the final partial batch after the last event and +before retrieval. That final flush remains inside the last event's insertion +timing, so it is not hidden as unreported cleanup work. + +Count refresh is independent of prompt batching and provider request batching. +It changes the physical execution identity and checkpoint contract. Resume +must use the same `--refresh-every` value; intermediate dataset session +boundaries do not force an early refresh. + The remaining flags belong to separate physical layers: - `--lotus-cache-mode disabled|memory` controls LOTUS's process-local exact diff --git a/tools/evaluation/claude_memory_locomo.py b/tools/evaluation/claude_memory_locomo.py index 852c8b7..ad80b44 100644 --- a/tools/evaluation/claude_memory_locomo.py +++ b/tools/evaluation/claude_memory_locomo.py @@ -20,6 +20,7 @@ SEM_TOPK_METHODS, run_claude_memory_locomo, ) +from agent_memory.planner import DEFAULT_GROUPED_AGG_RULE # noqa: E402 DEFAULT_OUTPUT_DIR = PROJECT_ROOT / ".memory-test" / "locomo-benchmark" / "latest" LOCOMO_CACHE_PATH = PROJECT_ROOT / ".cache" / "agent-memory" / "locomo10.json" @@ -55,7 +56,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--grouped-agg-rule", choices=GROUPED_AGG_RULES, - default="compressed", + default=DEFAULT_GROUPED_AGG_RULE, help="Grouped aggregate differential rule strategy for this evaluation run.", ) parser.add_argument( diff --git a/tools/evaluation/locomo.py b/tools/evaluation/locomo.py index 4b2d74f..0f21156 100644 --- a/tools/evaluation/locomo.py +++ b/tools/evaluation/locomo.py @@ -20,10 +20,17 @@ from agent_memory.evaluation.run import ( # noqa: E402 AGENT_MEMORY_SYSTEMS, DEFAULT_PROVIDER_MODEL_ID, + SEM_AGG_DISPATCH_METHODS, SEMANTIC_PAIR_PROFILES, SEM_JOIN_TOPK_METHODS, + resolve_grouped_agg_rule, run_agent_memory_bundle, ) +from agent_memory.adapters.lotus.prompt_batching import ( # noqa: E402 + STRUCTURED_OUTPUT_TRANSPORTS, + parse_prompt_batch_size, + validate_structured_output_transport, +) from agent_memory.planner import GROUPED_AGG_RULES # noqa: E402 DEFAULT_DATASET = ( @@ -64,10 +71,11 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: run.add_argument("--condition-id") run.add_argument("--maintenance-checkpoint-output-dir", type=Path) run.add_argument("--maintenance-only", action="store_true") + run.add_argument("--refresh-every", type=int, default=1) run.add_argument( "--grouped-agg-rule", choices=GROUPED_AGG_RULES, - default="rule-all-group", + default=None, ) run.add_argument( "--sem-topk-method", @@ -80,6 +88,21 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ) 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( + "--sem-agg-dispatch", + choices=SEM_AGG_DISPATCH_METHODS, + default="sequential", + ) + run.add_argument("--prompt-batch-size", type=parse_prompt_batch_size) + run.add_argument( + "--lm-max-ctx-len", type=int, + help="Memory LM context capacity in tokens; omitted uses LOTUS default.", + ) + run.add_argument( + "--structured-output-transport", + choices=STRUCTURED_OUTPUT_TRANSPORTS, + default="chat-json-object", + ) run.add_argument( "--semantic-pair-profile", choices=SEMANTIC_PAIR_PROFILES, @@ -101,6 +124,16 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ) args = parser.parse_args(raw_args) if args.command == "run": + try: + validate_structured_output_transport( + args.structured_output_transport, model=args.memory_model + ) + except ValueError as error: + parser.error(str(error)) + args.grouped_agg_rule = resolve_grouped_agg_rule( + args.system, + args.grouped_agg_rule, + ) sem_join_topk_requested = any( argument == "--sem-join-topk-method" or argument.startswith("--sem-join-topk-method=") @@ -194,9 +227,7 @@ def main(argv: Sequence[str] | None = None) -> Path: raise ValueError("bundle is not the pinned LOCOMO dataset") return run_agent_memory_bundle( bundle=bundle, - contracts={ - "locomo": locomo_task_contract(judge_model_id=args.judge_model) - }, + contracts={"locomo": locomo_task_contract(judge_model_id=args.judge_model)}, system_id=args.system, output_dir=args.output_dir, memory_provider_model_id=args.memory_model, @@ -208,6 +239,8 @@ def main(argv: Sequence[str] | None = None) -> Path: 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, + sem_agg_dispatch=args.sem_agg_dispatch, + prompt_batching=args.prompt_batch_size, semantic_pair_profile=args.semantic_pair_profile, semantic_pair_top_k=args.semantic_pair_top_k, semantic_pair_min_similarity=args.semantic_pair_min_similarity, @@ -215,10 +248,20 @@ def main(argv: Sequence[str] | None = None) -> Path: lotus_cache_mode=args.lotus_cache_mode, embedding_device=args.embedding_device, semantic_trace_snapshot_mode=args.semantic_trace_snapshot_mode, + refresh_every=args.refresh_every, memory_thinking_enabled=False, condition_id=args.condition_id or "", maintenance_only=args.maintenance_only, maintenance_checkpoint_output_dir=args.maintenance_checkpoint_output_dir, + **( + {"lm_max_ctx_len": args.lm_max_ctx_len} + if args.lm_max_ctx_len is not None else {} + ), + **( + {"structured_output_transport": args.structured_output_transport} + if args.structured_output_transport != "chat-json-object" + else {} + ), ) diff --git a/tools/evaluation/longmemeval.py b/tools/evaluation/longmemeval.py index 11fc45b..22c5318 100644 --- a/tools/evaluation/longmemeval.py +++ b/tools/evaluation/longmemeval.py @@ -24,10 +24,17 @@ ) from agent_memory.evaluation.run import ( # noqa: E402 AGENT_MEMORY_SYSTEMS, + SEM_AGG_DISPATCH_METHODS, SEMANTIC_PAIR_PROFILES, SEM_JOIN_TOPK_METHODS, + resolve_grouped_agg_rule, run_agent_memory_bundle, ) +from agent_memory.adapters.lotus.prompt_batching import ( # noqa: E402 + STRUCTURED_OUTPUT_TRANSPORTS, + parse_prompt_batch_size, + validate_structured_output_transport, +) from agent_memory.planner import GROUPED_AGG_RULES # noqa: E402 DEFAULT_MEMORY_MODEL_ID = "deepseek-v4-flash" @@ -116,7 +123,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: run.add_argument( "--grouped-agg-rule", choices=GROUPED_AGG_RULES, - default="rule-all-group", + default=None, ) run.add_argument( "--sem-topk-method", @@ -129,6 +136,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ) 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( + "--sem-agg-dispatch", + choices=SEM_AGG_DISPATCH_METHODS, + default="sequential", + ) + run.add_argument("--prompt-batch-size", type=parse_prompt_batch_size) + run.add_argument( + "--structured-output-transport", + choices=STRUCTURED_OUTPUT_TRANSPORTS, + default="chat-json-object", + ) run.add_argument( "--semantic-pair-profile", choices=SEMANTIC_PAIR_PROFILES, @@ -155,10 +173,21 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ) run.add_argument("--maintenance-checkpoint-output-dir", type=Path) run.add_argument("--maintenance-only", action="store_true") + run.add_argument("--refresh-every", type=int, default=1) args = parser.parse_args(raw_args) if getattr(args, "question_ids", None) is not None: args.question_ids = tuple(args.question_ids) if args.command == "run": + try: + validate_structured_output_transport( + args.structured_output_transport, model=args.memory_model + ) + except ValueError as error: + parser.error(str(error)) + args.grouped_agg_rule = resolve_grouped_agg_rule( + args.system, + args.grouped_agg_rule, + ) sem_join_topk_requested = any( argument == "--sem-join-topk-method" or argument.startswith("--sem-join-topk-method=") @@ -272,6 +301,8 @@ def main(argv: Sequence[str] | None = None) -> Path: 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, + sem_agg_dispatch=args.sem_agg_dispatch, + prompt_batching=args.prompt_batch_size, semantic_pair_profile=args.semantic_pair_profile, semantic_pair_top_k=args.semantic_pair_top_k, semantic_pair_min_similarity=args.semantic_pair_min_similarity, @@ -279,11 +310,17 @@ def main(argv: Sequence[str] | None = None) -> Path: lotus_cache_mode=args.lotus_cache_mode, embedding_device=args.embedding_device, semantic_trace_snapshot_mode=args.semantic_trace_snapshot_mode, + refresh_every=args.refresh_every, memory_thinking_enabled=args.memory_thinking == "enabled", condition_id=_condition_id(args), maintenance_only=args.maintenance_only, maintenance_checkpoint_output_dir=args.maintenance_checkpoint_output_dir, max_new_cases=args.max_new_cases, + **( + {"structured_output_transport": args.structured_output_transport} + if args.structured_output_transport != "chat-json-object" + else {} + ), ) if not args.maintenance_only: store = BenchmarkArtifactStore(output) diff --git a/tools/evaluation/memory_agent_bench.py b/tools/evaluation/memory_agent_bench.py index 13c7a5a..e106f42 100644 --- a/tools/evaluation/memory_agent_bench.py +++ b/tools/evaluation/memory_agent_bench.py @@ -25,10 +25,17 @@ ) from agent_memory.evaluation.run import ( # noqa: E402 AGENT_MEMORY_SYSTEMS, + SEM_AGG_DISPATCH_METHODS, SEMANTIC_PAIR_PROFILES, SEM_JOIN_TOPK_METHODS, + resolve_grouped_agg_rule, run_agent_memory_bundle, ) +from agent_memory.adapters.lotus.prompt_batching import ( # noqa: E402 + STRUCTURED_OUTPUT_TRANSPORTS, + parse_prompt_batch_size, + validate_structured_output_transport, +) from agent_memory.planner import GROUPED_AGG_RULES # noqa: E402 DEFAULT_MODEL = "deepseek/deepseek-v4-flash" @@ -72,7 +79,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: run.add_argument( "--grouped-agg-rule", choices=GROUPED_AGG_RULES, - default="rule-all-group", + default=None, ) run.add_argument( "--sem-topk-method", @@ -85,6 +92,17 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ) 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( + "--sem-agg-dispatch", + choices=SEM_AGG_DISPATCH_METHODS, + default="sequential", + ) + run.add_argument("--prompt-batch-size", type=parse_prompt_batch_size) + run.add_argument( + "--structured-output-transport", + choices=STRUCTURED_OUTPUT_TRANSPORTS, + default="chat-json-object", + ) run.add_argument( "--semantic-pair-profile", choices=SEMANTIC_PAIR_PROFILES, @@ -104,8 +122,19 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: choices=("compact", "full"), default="compact", ) + run.add_argument("--refresh-every", type=int, default=1) args = parser.parse_args(raw_args) if args.command == "run": + try: + validate_structured_output_transport( + args.structured_output_transport, model=args.model + ) + except ValueError as error: + parser.error(str(error)) + args.grouped_agg_rule = resolve_grouped_agg_rule( + args.system, + args.grouped_agg_rule, + ) sem_join_topk_requested = any( argument == "--sem-join-topk-method" or argument.startswith("--sem-join-topk-method=") @@ -215,6 +244,8 @@ def main(argv: Sequence[str] | None = None) -> Path: 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, + sem_agg_dispatch=args.sem_agg_dispatch, + prompt_batching=args.prompt_batch_size, semantic_pair_profile=args.semantic_pair_profile, semantic_pair_top_k=args.semantic_pair_top_k, semantic_pair_min_similarity=args.semantic_pair_min_similarity, @@ -222,8 +253,14 @@ def main(argv: Sequence[str] | None = None) -> Path: lotus_cache_mode=args.lotus_cache_mode, embedding_device=args.embedding_device, semantic_trace_snapshot_mode=args.semantic_trace_snapshot_mode, + refresh_every=args.refresh_every, memory_thinking_enabled=False, condition_id=args.condition_id or "", + **( + {"structured_output_transport": args.structured_output_transport} + if args.structured_output_transport != "chat-json-object" + else {} + ), ) diff --git a/tools/evaluation/prompt_batching_scale.py b/tools/evaluation/prompt_batching_scale.py new file mode 100644 index 0000000..60acf93 --- /dev/null +++ b/tools/evaluation/prompt_batching_scale.py @@ -0,0 +1,1525 @@ +"""Run the controlled 128-event Claude prompt-batching scale experiment.""" + +from __future__ import annotations + +import argparse +from collections import Counter +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from decimal import Decimal +from hashlib import sha256 +import inspect +import json +import math +import os +from pathlib import Path +import pickle +import shlex +import signal +import subprocess +import sys +import tempfile +import time +import traceback +from typing import Any +from urllib.request import Request, urlopen +from zoneinfo import ZoneInfo + +import pandas as pd + + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(PROJECT_ROOT / "src")) + +from agent_memory.evaluation.bundle import write_bundle # noqa: E402 +from agent_memory.evaluation.locomo import ( # noqa: E402 + LOCOMO_COMMIT, + LOCOMO_SHA256, + eligible_questions, + load_locomo_sample, + locomo_bundle, +) +from agent_memory.evaluation.provenance import ( # noqa: E402 + SOURCE_EVIDENCE_ENV, + build_source_evidence, + collect_runtime_provenance, + validate_run_provenance, +) + + +HONG_KONG = ZoneInfo("Asia/Hong_Kong") +EVENT_COUNT = 128 +QUESTION_COUNT = 71 +JUDGED_QUESTION_COUNT = 54 +FIRST_EVENT_ID = "D1:1" +LAST_EVENT_ID = "D7:20" +CLIENT_REQUEST_MAX_WORKERS = 64 +LM_MAX_CTX_LEN = 1_000_000 +BGE_M3_REVISION = "5617a9f61b028005a4858fdac845db406aefb181" +SOFT_STOP_CNY = Decimal("48") +HARD_STOP_CNY = Decimal("50") +MIN_DISK_GIB = 100.0 +MAX_CONDITION_BYTES = 2 * 1024**3 +MAX_CONDITION_SECONDS = 6 * 60 * 60 +POLL_SECONDS = 15 +OFF_PEAK_PRICES = { + "cache_hit": Decimal("0.05"), + "cache_miss": Decimal("1.5"), + "completion": Decimal("4.5"), +} +PEAK_MULTIPLIER = Decimal("2") +OFFICIAL_SCORER_ID = "locomo_official:v1" +JUDGE_SCORER_ID = "locomo_zep_judge:deepseek/deepseek-v4-flash" + + +@dataclass(frozen=True) +class ConditionSpec: + """One independent end-to-end batching-scale condition.""" + + pass_index: int + label: str + refresh_every: int + prompt_batch: int | str | None + + @property + def run_id(self) -> str: + """Return the stable directory identifier for this condition.""" + + return f"pass-{self.pass_index}-{self.label}" + + def to_dict(self) -> dict[str, Any]: + """Return the JSON-safe condition contract.""" + + return {"run_id": self.run_id, **asdict(self)} + + +@dataclass(frozen=True) +class ExperimentPaths: + """Resolved filesystem paths for one experiment run.""" + + root: Path + control: Path + bundle: Path + conditions: Path + analysis: Path + source: Path + venv: Path + deps: Path | None + hf_home: Path + + @classmethod + def from_contract(cls, root: Path, contract: Mapping[str, Any]) -> ExperimentPaths: + """Build paths from one persisted experiment contract.""" + + deps = contract["runtime"].get("deps_dir") + return cls( + root=root.resolve(), + control=root.resolve() / "control", + bundle=root.resolve() / "bundle", + conditions=root.resolve() / "conditions", + analysis=root.resolve() / "analysis", + source=Path(str(contract["source"]["directory"])).resolve(), + venv=Path(str(contract["runtime"]["venv"])).resolve(), + deps=Path(str(deps)).resolve() if deps is not None else None, + hf_home=Path(str(contract["runtime"]["hf_home"])).resolve(), + ) + + +def condition_matrix() -> tuple[ConditionSpec, ...]: + """Return the two identical ordered passes of the ten-condition sweep.""" + + levels: tuple[tuple[str, int, int | str | None], ...] = ( + ("original", 1, None), + ("packed-1", 1, 1), + ("packed-2", 2, 2), + ("packed-4", 4, 4), + ("packed-8", 8, 8), + ("packed-16", 16, 16), + ("packed-32", 32, 32), + ("packed-64", 64, 64), + ("packed-128", 128, 128), + ("packed-all", 128, "all"), + ) + return tuple( + ConditionSpec(pass_index, label, refresh, prompt) + for pass_index in (1, 2) + for label, refresh, prompt in levels + ) + + +def build_condition_command( + condition: ConditionSpec, + *, + source: Path, + venv: Path, + bundle: Path, + output: Path, +) -> list[str]: + """Build the canonical LOCOMO command for one condition.""" + + command = [ + str(venv / "bin" / "python"), + str(source / "tools" / "evaluation" / "locomo.py"), + "run", + "--bundle-dir", + str(bundle), + "--output-dir", + str(output), + "--system", + "claude-memory", + "--condition-id", + f"Claude-BatchingScale-{condition.run_id}-128e", + "--lm-max-ctx-len", + str(LM_MAX_CTX_LEN), + "--memory-model", + "deepseek/deepseek-v4-flash", + "--answer-model", + "deepseek/deepseek-v4-flash", + "--judge-model", + "deepseek/deepseek-v4-flash", + "--grouped-agg-rule", + "rule-join-map", + "--sem-topk-method", + "listwise", + "--semantic-pair-profile", + "search-filter", + "--semantic-pair-top-k", + "20", + "--semantic-pair-min-similarity", + "0.5", + "--embedding-device", + "cuda", + "--lotus-cache-mode", + "disabled", + "--semantic-trace-snapshot-mode", + "compact", + "--refresh-every", + str(condition.refresh_every), + ] + if condition.prompt_batch is not None: + command.extend(("--prompt-batch-size", str(condition.prompt_batch))) + return command + + +def _atomic_json(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + delete=False, + ) as handle: + json.dump(payload, handle, indent=2, sort_keys=True, ensure_ascii=False) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + temporary = Path(handle.name) + os.replace(temporary, path) + + +def _atomic_text(path: Path, value: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + "w", + encoding="utf-8", + dir=path.parent, + prefix=f".{path.name}.", + delete=False, + ) as handle: + handle.write(value) + handle.flush() + os.fsync(handle.fileno()) + temporary = Path(handle.name) + os.replace(temporary, path) + + +def _sha256_file(path: Path) -> str: + digest = sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _git(repository: Path, *arguments: str) -> str: + return subprocess.check_output( + ("git", "-C", str(repository), *arguments), text=True + ).strip() + + +def _directory_bytes(path: Path) -> int: + total = 0 + if not path.exists(): + return total + for root, _directories, files in os.walk(path): + for name in files: + try: + total += (Path(root) / name).stat().st_size + except FileNotFoundError: + continue + return total + + +def _disk_available_gib(path: Path) -> float: + stats = os.statvfs(path) + return stats.f_bavail * stats.f_frsize / 1024**3 + + +def _is_peak(local: datetime) -> bool: + if local.weekday() >= 5: + return False + minute = local.hour * 60 + local.minute + return 9 * 60 <= minute < 12 * 60 or 14 * 60 <= minute < 18 * 60 + + +def _pricing_period(timestamp: str) -> str: + observed = datetime.strptime(timestamp, "%Y%m%dT%H%M%S%fZ").replace( + tzinfo=timezone.utc + ) + return "peak" if _is_peak(observed.astimezone(HONG_KONG)) else "off_peak" + + +def _usage_cost_cny(pricing_usage: Mapping[str, Mapping[str, int]]) -> Decimal: + total = Decimal(0) + million = Decimal(1_000_000) + for period, usage in pricing_usage.items(): + multiplier = PEAK_MULTIPLIER if period == "peak" else Decimal(1) + total += multiplier * ( + Decimal(usage.get("cache_hit", 0)) * OFF_PEAK_PRICES["cache_hit"] + + Decimal(usage.get("cache_miss", 0)) + * OFF_PEAK_PRICES["cache_miss"] + + Decimal(usage.get("completion", 0)) + * OFF_PEAK_PRICES["completion"] + ) / million + return total.quantize(Decimal("0.000001")) + + +def _empty_monitor() -> dict[str, Any]: + return { + "schema_version": 1, + "offset": 0, + "provider_responses": 0, + "reasoning_tokens": 0, + "reasoning_usage_reports": 0, + "provider_contract_violations": [], + "pricing_usage": { + "off_peak": {"cache_hit": 0, "cache_miss": 0, "completion": 0}, + "peak": {"cache_hit": 0, "cache_miss": 0, "completion": 0}, + }, + "phase_usage": {}, + "insertion_results": 0, + "insertion_latency_ms": 0.0, + "trace_io_latency_ms": 0.0, + "trace_bytes_written": 0, + "retrieval_results": 0, + "answer_results": 0, + "grade_results": 0, + "candidate_pairs": {}, + "operator_usage": {}, + "prompt_batching": {}, + } + + +def _integer_token(event: Mapping[str, Any], key: str) -> int: + value = event.get(key) + if value in (None, ""): + raise ValueError(f"provider usage is missing {key}") + return int(value) + + +def _number_equals(value: Any, expected: float) -> bool: + return ( + isinstance(value, int | float) + and not isinstance(value, bool) + and math.isfinite(float(value)) + and float(value) == expected + ) + + +def _counter_mapping(value: Mapping[str, Mapping[str, Any]]) -> dict[str, Counter[str]]: + return {key: Counter(row) for key, row in value.items()} + + +def _trace_site(event: Mapping[str, Any]) -> str: + operator = str(event.get("operator") or "unknown") + digest = str( + event.get("semantic_pair_site_id") + or event.get("query_digest") + or event.get("query_fingerprint") + or "default" + ) + return f"{operator}:{digest}" + + +def _physical_usage_event(event: dict[str, Any]) -> dict[str, Any]: + """Normalize direct answer/judge usage without counting LOTUS mirrors twice.""" + + if ( + event.get("event_type") != "llm_call" + or event.get("operator") != "llm" + or event.get("phase") not in {"answering", "grading"} + ): + return event + if event.get("usage_scope") != "batch" or event.get("llm_batch_size") != 1: + raise RuntimeError("unsupported direct answer/judge usage scope") + details = event.get("usage_completion_tokens_details") or {} + return { + **event, + "event_type": "provider_usage", + "provider_prompt_cache_hit_tokens": event.get("usage_prompt_cache_hit_tokens"), + "provider_prompt_cache_miss_tokens": event.get("usage_prompt_cache_miss_tokens"), + "provider_completion_tokens": event.get("usage_completion_tokens"), + "provider_reasoning_tokens": details.get("reasoning_tokens"), + "provider_request": {"provider_kwargs": event.get("llm_kwargs") or {}}, + } + + +def _update_monitor(trace: Path, state_path: Path) -> dict[str, Any]: + state = json.loads(state_path.read_text()) if state_path.exists() else _empty_monitor() + if state.get("schema_version") != 1: + raise RuntimeError(f"invalid monitor state: {state_path}") + offset = int(state["offset"]) + if trace.stat().st_size < offset: + raise RuntimeError(f"semantic trace shrank: {trace}") + candidates = _counter_mapping(state["candidate_pairs"]) + operators = _counter_mapping(state["operator_usage"]) + batching = _counter_mapping(state["prompt_batching"]) + phases = _counter_mapping(state["phase_usage"]) + violations = list(state["provider_contract_violations"]) + with trace.open("rb") as handle: + handle.seek(offset) + while True: + line_start = handle.tell() + line = handle.readline() + if not line: + break + if not line.endswith(b"\n"): + handle.seek(line_start) + break + event = _physical_usage_event(json.loads(line)) + event_type = str(event.get("event_type") or "") + site = _trace_site(event) + if event_type == "provider_usage": + try: + hit = _integer_token(event, "provider_prompt_cache_hit_tokens") + miss = _integer_token(event, "provider_prompt_cache_miss_tokens") + completion = _integer_token(event, "provider_completion_tokens") + period = _pricing_period(str(event["timestamp"])) + except (KeyError, TypeError, ValueError) as error: + violations.append(f"incomplete_provider_usage:{error}") + hit = miss = completion = 0 + period = "off_peak" + response_count = int(event.get("response_count") or 1) + state["provider_responses"] += response_count + state["pricing_usage"][period]["cache_hit"] += hit + state["pricing_usage"][period]["cache_miss"] += miss + state["pricing_usage"][period]["completion"] += completion + phase = str(event.get("phase") or "unknown") + phases.setdefault(phase, Counter()) + phases[phase].update( + responses=response_count, + cache_hit=hit, + cache_miss=miss, + completion=completion, + ) + operators.setdefault(site, Counter()) + operators[site].update( + responses=response_count, + tokens=hit + miss + completion, + ) + reasoning = event.get("provider_reasoning_tokens") + if reasoning not in (None, ""): + state["reasoning_usage_reports"] += 1 + state["reasoning_tokens"] += int(reasoning) + request = event.get("provider_request") or {} + kwargs = request.get("provider_kwargs") or {} + if float(kwargs.get("temperature", 0.0)) != 0.0: + violations.append(f"temperature:{kwargs.get('temperature')}") + if kwargs.get("thinking") != {"type": "disabled"}: + violations.append(f"thinking:{kwargs.get('thinking')!r}") + elif event_type == "insertion_result": + state["insertion_results"] += 1 + state["insertion_latency_ms"] += float(event.get("latency_ms") or 0) + state["trace_io_latency_ms"] += float( + event.get("semantic_trace_io_latency_ms") or 0 + ) + state["trace_bytes_written"] += int( + event.get("semantic_trace_bytes_written") or 0 + ) + elif event_type == "retrieval_result": + state["retrieval_results"] += 1 + elif event_type == "answer_result": + state["answer_results"] += 1 + elif event_type == "grade_result": + state["grade_results"] += 1 + elif event_type == "candidate_generation": + candidates.setdefault(site, Counter()) + candidates[site].update( + events=1, + eligible=int(event.get("total_pair_count") or 0), + selected=int(event.get("candidate_pair_count") or 0), + ) + elif event_type == "prompt_batching": + batching.setdefault(site, Counter()) + chunks = [int(item) for item in event.get("chunk_sizes") or ()] + batching[site].update( + events=1, + tasks=int(event.get("task_count") or 0), + prompts=int(event.get("prompt_count") or 0), + retries=int(event.get("retry_count") or 0), + repairs=int( + event.get( + "structured_output_repair_count", + event.get("syntax_repair_count"), + ) + or 0 + ), + ) + if chunks: + batching[site]["max_chunk"] = max( + batching[site]["max_chunk"], max(chunks) + ) + offset = handle.tell() + state["offset"] = offset + state["provider_contract_violations"] = violations + state["candidate_pairs"] = { + key: dict(value) for key, value in sorted(candidates.items()) + } + state["operator_usage"] = { + key: dict(value) for key, value in sorted(operators.items()) + } + state["prompt_batching"] = { + key: dict(value) for key, value in sorted(batching.items()) + } + state["phase_usage"] = { + key: dict(value) for key, value in sorted(phases.items()) + } + state["cost_cny"] = float(_usage_cost_cny(state["pricing_usage"])) + _atomic_json(state_path, state) + return state + + +def _read_contract(root: Path) -> dict[str, Any]: + return json.loads((root / "control" / "experiment-contract.json").read_text()) + + +def _paths(root: Path) -> ExperimentPaths: + contract = _read_contract(root) + return ExperimentPaths.from_contract(root, contract) + + +def _log(paths: ExperimentPaths, message: str) -> None: + timestamp = datetime.now().astimezone().isoformat(timespec="seconds") + with (paths.control / "controller.log").open("a", encoding="utf-8") as handle: + handle.write(f"[{timestamp}] {message}\n") + handle.flush() + + +def _state(paths: ExperimentPaths, name: str) -> str: + path = paths.control / name + return path.read_text().strip() if path.exists() else "" + + +def _set_state(paths: ExperimentPaths, name: str, value: str) -> None: + _atomic_text(paths.control / name, f"{value}\n") + + +def _condition_from_dict(value: Mapping[str, Any]) -> ConditionSpec: + return ConditionSpec( + pass_index=int(value["pass_index"]), + label=str(value["label"]), + refresh_every=int(value["refresh_every"]), + prompt_batch=value.get("prompt_batch"), + ) + + +def _total_cost(paths: ExperimentPaths) -> Decimal: + continuation = _read_contract(paths.root).get("continuation") or {} + total = Decimal(str(continuation.get("prior_cost_cny", "0"))) + for path in paths.conditions.glob("*/monitor-state.json"): + payload = json.loads(path.read_text()) + total += Decimal(str(payload.get("cost_cny") or 0)) + return total + + +def _capture_continuation( + origin: Path, run_ids: Sequence[str], contract: Mapping[str, Any], +) -> dict[str, Any]: + """Pin completed results and all prior physical spending without writing to origin.""" + + origin = origin.resolve() + previous = _read_contract(origin) + _validate_continuation(previous) + inherited = previous.get("continuation") or {} + if (origin / "control/experiment-state").read_text().strip() not in { + "failed", "safety-stopped", "completed", + }: + raise RuntimeError("origin experiment must be stopped") + for key in ("experiment", "input", "fixed_execution", "runtime", "pricing", "conditions"): + if previous.get(key) != contract.get(key): + raise RuntimeError(f"continuation contract mismatch: {key}") + specs = {row["run_id"]: row for row in previous["conditions"]} + if not run_ids or len(set(run_ids)) != len(run_ids) or not set(run_ids) <= specs.keys(): + raise ValueError("reuse IDs must be known, unique, and nonempty") + files = { + origin / "control/experiment-contract.json", + origin / "control/source-evidence.json", + } + if inherited: + files.update( + Path(inherited["origin_root"]) / name + for name in inherited["files_sha256"] + ) + reused = {} + reused_sources = {} + for run_id in run_ids: + directory = _result_directory(origin, previous, run_id) + output = directory / "output" + if (directory / "state").read_text().strip() != "completed": + raise RuntimeError(f"reuse condition is not completed: {run_id}") + trace = output / "trace/events.jsonl" + if not trace.is_file(): + raise RuntimeError(f"reuse cost evidence is missing for {run_id}: {trace}") + result = json.loads((directory / "validation.json").read_text()) + summary = json.loads((output / "metrics/summary.json").read_text()) + monitor = json.loads((directory / "monitor-state.json").read_text()) + scores = summary.get("scores_by_scorer") or {} + if ( + any(result.get(key) != value for key, value in specs[run_id].items()) + or monitor.get("insertion_results") != EVENT_COUNT + or summary.get("completed_cases") != 1 + or summary.get("failed_cases") != 0 + or summary.get("question_count") != QUESTION_COUNT + or summary.get("memory_system_error_question_count") != 0 + or summary.get("retrieval_system_error_count") != 0 + or _completed_questions(output) != QUESTION_COUNT + or (scores.get(OFFICIAL_SCORER_ID) or {}).get("grade_count") != QUESTION_COUNT + or (scores.get(JUDGE_SCORER_ID) or {}).get("grade_count") != JUDGED_QUESTION_COUNT + ): + raise RuntimeError(f"reuse result is incomplete or incompatible: {run_id}") + # Hash the small completed artifacts, including checkpoint and question payloads. + files.update(p for p in directory.rglob("*") if p.is_file()) + reused[run_id] = str(directory) + reused_sources[run_id] = ( + inherited.get("reused_sources", {}).get(run_id, inherited.get("source")) + if run_id in inherited.get("reused_conditions", {}) else previous["source"] + ) + condition_costs = {} + for trace in sorted(origin.glob("conditions/*/output/trace/events.jsonl")): + files.add(trace) + usage = {period: Counter() for period in ("off_peak", "peak")} + with trace.open() as handle: + for line in handle: + event = _physical_usage_event(json.loads(line)) + if event.get("event_type") == "provider_usage": + usage[_pricing_period(event["timestamp"])].update( + cache_hit=_integer_token(event, "provider_prompt_cache_hit_tokens"), + cache_miss=_integer_token(event, "provider_prompt_cache_miss_tokens"), + completion=_integer_token(event, "provider_completion_tokens"), + ) + condition_costs[trace.parents[2].name] = str(_usage_cost_cny(usage)) + return { + "origin_root": str(origin), + "source": previous["source"], + "reused_conditions": reused, + "reused_sources": reused_sources, + "files_sha256": {os.path.relpath(p, origin): _sha256_file(p) for p in sorted(files)}, + "prior_cost_cny": str(sum( + (Decimal(value) for value in condition_costs.values()), + Decimal(str(inherited.get("prior_cost_cny", "0"))), + )), + "prior_condition_costs_cny": { + **inherited.get("prior_condition_costs_cny", {}), **condition_costs, + }, + "interpretation": "Exploratory mixed-source first pass; second pass uses the new source.", + } + + +def _validate_continuation(contract: Mapping[str, Any]) -> None: + continuation = contract.get("continuation") + if not continuation: + return + origin = Path(continuation["origin_root"]) + for name, digest in continuation["files_sha256"].items(): + path = origin / name + if not path.is_file() or _sha256_file(path) != digest: + raise RuntimeError(f"reused evidence changed: {path}") + + +def _result_directory(root: Path, contract: Mapping[str, Any], run_id: str) -> Path: + reused = (contract.get("continuation") or {}).get("reused_conditions", {}) + return Path(reused[run_id]) if run_id in reused else root / "conditions" / run_id + + +def _completed_questions(output: Path) -> int: + return sum( + 1 for _path in output.glob("cases/*/question-results/*/complete.json") + ) + + +def _question_numbers(dataset: Path) -> tuple[int, ...]: + sample = load_locomo_sample(dataset, sample_index=0) + events = sample.events[:EVENT_COUNT] + questions = eligible_questions( + sample.questions, + ingested_event_ids=[event.event_id for event in events], + ) + if ( + len(events) != EVENT_COUNT + or events[0].event_id != FIRST_EVENT_ID + or events[-1].event_id != LAST_EVENT_ID + or len(questions) != QUESTION_COUNT + ): + raise RuntimeError("the pinned LOCOMO 128-event prefix contract changed") + categories = Counter(int(question.category) for question in questions) + if categories != Counter({1: 7, 2: 16, 3: 6, 4: 25, 5: 17}): + raise RuntimeError(f"unexpected question categories: {categories}") + return tuple(int(question.metadata["question_number"]) for question in questions) + + +def _validate_initialization_root(root: Path) -> None: + if not root.exists(): + return + unexpected = sorted( + path.name for path in root.iterdir() if path.name not in {"package", "source"} + ) + if unexpected: + raise FileExistsError( + f"experiment root contains unexpected entries: {root}: {unexpected}" + ) + + +def initialize( + *, + root: Path, + source: Path, + venv: Path, + deps: Path | None, + hf_home: Path, + dataset: Path, + source_package_inventory: Path, + reuse_root: Path | None = None, + reuse_conditions: Sequence[str] = (), +) -> None: + """Create the immutable experiment contract and canonical input bundle.""" + + root = root.resolve() + source = source.resolve() + _validate_initialization_root(root) + control = root / "control" + bundle_dir = root / "bundle" + conditions_dir = root / "conditions" + analysis_dir = root / "analysis" + for path in (control, bundle_dir, conditions_dir, analysis_dir): + path.mkdir(parents=True, exist_ok=True) + numbers = _question_numbers(dataset) + bundle = locomo_bundle( + dataset, + sample_index=0, + start_row=1, + row_limit=EVENT_COUNT, + question_numbers=numbers, + include_adversarial=True, + run_mode="prompt-batching-scale", + ) + write_bundle(bundle, bundle_dir) + manifest = json.loads((bundle_dir / "manifest.json").read_text()) + evidence = build_source_evidence(source) + _atomic_json(control / "source-evidence.json", evidence) + inventory_copy = control / "source-package-inventory.json" + _atomic_text(inventory_copy, source_package_inventory.read_text()) + contract = { + "schema_version": 1, + "experiment": "claude-prompt-batching-scale-128e", + "source": { + "directory": str(source), + "base_commit": _git(source, "rev-parse", "HEAD"), + "source_snapshot_sha256": evidence["source_snapshot_sha256"], + "dirty_patch_sha256": evidence["dirty_patch_sha256"], + "source_evidence_sha256": _sha256_file(control / "source-evidence.json"), + "source_package_inventory_sha256": _sha256_file(inventory_copy), + "lock_sha256": _sha256_file(source / "uv.lock"), + }, + "runtime": { + "venv": str(venv.resolve()), + "deps_dir": str(deps.resolve()) if deps is not None else None, + "hf_home": str(hf_home.resolve()), + "python": "3.12", + "lotus": "1.1.4", + "sentence_transformers": "3.4.1", + "torch": "2.7.0+cu118", + "client_request_max_workers": CLIENT_REQUEST_MAX_WORKERS, + }, + "input": { + "dataset_commit": LOCOMO_COMMIT, + "dataset_sha256": LOCOMO_SHA256, + "bundle_fingerprint": manifest["fingerprint"], + "policy_input_fingerprint": manifest["policy_input_fingerprint"], + "case_id": "conv-26", + "first_event_id": FIRST_EVENT_ID, + "last_event_id": LAST_EVENT_ID, + "event_count": EVENT_COUNT, + "question_count": QUESTION_COUNT, + "question_numbers": list(numbers), + "category_counts": {"1": 7, "2": 16, "3": 6, "4": 25, "5": 17}, + }, + "fixed_execution": { + "system": "claude-memory", + "grouped_agg_rule": "rule-join-map", + "semantic_pair_profile": "search-filter", + "semantic_pair_top_k": 20, + "semantic_pair_min_similarity": 0.5, + "embedding_model": "BAAI/bge-m3", + "embedding_revision": BGE_M3_REVISION, + "embedding_device": "cuda", + "retrieval": "listwise", + "model": "deepseek/deepseek-v4-flash", + "temperature": 0, + "thinking": "disabled", + "lotus_cache": "disabled", + "semantic_trace": "compact", + }, + "pricing": { + "currency": "CNY", + "per_million_tokens_off_peak": { + key: str(value) for key, value in OFF_PEAK_PRICES.items() + }, + "peak_multiplier": str(PEAK_MULTIPLIER), + "soft_stop_cny": str(SOFT_STOP_CNY), + "hard_stop_cny": str(HARD_STOP_CNY), + }, + "conditions": [condition.to_dict() for condition in condition_matrix()], + } + from agent_memory.adapters.lotus.json_output import JSON_REPAIR_VERSION + + contract["structured_output"] = { + "transport": "chat-json-object", "repair_version": JSON_REPAIR_VERSION, + "lm_max_ctx_len": LM_MAX_CTX_LEN, + } + if (reuse_root is None) != (not reuse_conditions): + raise ValueError("reuse-root and reuse-condition must be supplied together") + if reuse_root is not None: + contract["continuation"] = _capture_continuation(reuse_root, reuse_conditions, contract) + _atomic_json(control / "experiment-contract.json", contract) + _set_state( + ExperimentPaths.from_contract(root, contract), "experiment-state", "initialized" + ) + _set_state( + ExperimentPaths.from_contract(root, contract), "current-condition", "none" + ) + status_script = f"""#!/usr/bin/env bash +set -euo pipefail +export PYTHONPATH={shlex.quote(str(source / 'src'))} +exec {shlex.quote(str(venv / 'bin' / 'python'))} \\ + {shlex.quote(str(source / 'tools' / 'evaluation' / 'prompt_batching_scale.py'))} \\ + status --run-root {shlex.quote(str(root))} +""" + _atomic_text(control / "status.sh", status_script) + (control / "status.sh").chmod(0o755) + + +def preflight(root: Path) -> dict[str, Any]: + """Validate source, input, runtime, model cache, credentials, and capacity.""" + + contract = _read_contract(root) + paths = ExperimentPaths.from_contract(root, contract) + _validate_continuation(contract) + evidence_path = paths.control / "source-evidence.json" + os.environ[SOURCE_EVIDENCE_ENV] = str(evidence_path) + provenance = collect_runtime_provenance( + paths.source, + lockfile="uv.lock", + dependencies=("agent-memory", "lotus-ai", "sentence-transformers", "torch"), + ) + validate_run_provenance(provenance, run_mode="prompt-batching-scale") + source = contract["source"] + if _git(paths.source, "rev-parse", "HEAD") != source["base_commit"]: + raise RuntimeError("source base commit changed") + if _sha256_file(paths.source / "uv.lock") != source["lock_sha256"]: + raise RuntimeError("source lock digest changed") + if provenance["source"].get("source_snapshot_sha256") != source[ + "source_snapshot_sha256" + ]: + raise RuntimeError("source snapshot does not match the experiment contract") + if _sha256_file(evidence_path) != source["source_evidence_sha256"]: + raise RuntimeError("source evidence digest changed") + if _sha256_file(paths.control / "source-package-inventory.json") != source[ + "source_package_inventory_sha256" + ]: + raise RuntimeError("source package inventory digest changed") + manifest = json.loads((paths.bundle / "manifest.json").read_text()) + expected_input = contract["input"] + if ( + manifest.get("fingerprint") != expected_input["bundle_fingerprint"] + or manifest.get("event_count") != EVENT_COUNT + or manifest.get("question_count") != QUESTION_COUNT + ): + raise RuntimeError(f"bundle manifest mismatch: {manifest}") + events = [json.loads(line) for line in (paths.bundle / "events.jsonl").read_text().splitlines()] + questions = [ + json.loads(line) for line in (paths.bundle / "questions.jsonl").read_text().splitlines() + ] + if ( + len(events) != EVENT_COUNT + or events[0].get("event_id") != FIRST_EVENT_ID + or events[-1].get("event_id") != LAST_EVENT_ID + or len(questions) != QUESTION_COUNT + ): + raise RuntimeError("bundle contents do not match the 128-event contract") + + from importlib.metadata import version + + import agent_memory + import torch + from agent_memory.adapters.lotus.context import LotusExecutionConfig + from lotus.models import LM + + if sys.version_info[:2] != (3, 12): + raise RuntimeError(f"unexpected Python runtime: {sys.version}") + expected_versions = contract["runtime"] + if version("lotus-ai") != expected_versions["lotus"]: + raise RuntimeError(f"unexpected LOTUS version: {version('lotus-ai')}") + if version("sentence-transformers") != expected_versions["sentence_transformers"]: + raise RuntimeError( + f"unexpected SentenceTransformers version: {version('sentence-transformers')}" + ) + if torch.__version__ != expected_versions["torch"] or not torch.cuda.is_available(): + raise RuntimeError(f"CUDA contract failed: torch={torch.__version__}") + if paths.source not in Path(agent_memory.__file__).resolve().parents: + raise RuntimeError(f"wrong agent_memory source: {agent_memory.__file__}") + config = LotusExecutionConfig(lm_model_kwargs={"max_ctx_len": LM_MAX_CTX_LEN}) + from agent_memory.adapters.lotus.json_output import JSON_REPAIR_VERSION + + if contract.get("structured_output") != { + "transport": config.structured_output_transport, "repair_version": JSON_REPAIR_VERSION, + "lm_max_ctx_len": config.lm_model_kwargs["max_ctx_len"], + }: + raise RuntimeError("structured output contract changed") + signature = inspect.signature(LM.__init__) + if ( + config.lm_max_batch_size != CLIENT_REQUEST_MAX_WORKERS + or signature.parameters["max_batch_size"].default + != CLIENT_REQUEST_MAX_WORKERS + or signature.parameters["temperature"].default != 0.0 + ): + raise RuntimeError("LOTUS LM defaults do not match the experiment contract") + bge_path = ( + paths.hf_home + / "hub" + / "models--BAAI--bge-m3" + / "snapshots" + / BGE_M3_REVISION + ) + if not bge_path.is_dir(): + raise RuntimeError(f"pinned BGE-M3 revision is not cached: {bge_path}") + if not (paths.venv / "bin" / "python").is_file(): + raise RuntimeError(f"validated Python environment is missing: {paths.venv}") + if paths.deps is not None and not paths.deps.is_dir(): + raise RuntimeError(f"dependency overlay is missing: {paths.deps}") + if _disk_available_gib(Path("/mnt/data")) < MIN_DISK_GIB: + raise RuntimeError("insufficient /mnt/data free space") + if _is_peak(datetime.now(tz=HONG_KONG)): + raise RuntimeError("new experiments start only in a DeepSeek off-peak window") + active = subprocess.run( + ("pgrep", "-af", "tools/evaluation/locomo.py run"), + capture_output=True, + text=True, + check=False, + ).stdout.strip() + if active: + raise RuntimeError(f"another LOCOMO worker is active:\n{active}") + api_key = os.environ.get("DEEPSEEK_API_KEY") + if not api_key: + raise RuntimeError("DEEPSEEK_API_KEY is unavailable") + request = Request( + "https://api.deepseek.com/user/balance", + headers={"Accept": "application/json", "Authorization": f"Bearer {api_key}"}, + ) + with urlopen(request, timeout=30) as response: + balance_payload = json.loads(response.read()) + balances = [ + Decimal(str(item["total_balance"])) + for item in balance_payload.get("balance_infos", ()) + if item.get("currency") == "CNY" + ] + if balance_payload.get("is_available") is not True or len(balances) != 1: + raise RuntimeError("DeepSeek CNY balance is unavailable") + if balances[0] < HARD_STOP_CNY: + raise RuntimeError(f"DeepSeek balance is below 50 CNY: {balances[0]}") + result = { + "timestamp": datetime.now().astimezone().isoformat(timespec="seconds"), + "source_commit": source["base_commit"], + "source_dirty": provenance["source"]["dirty"], + "source_snapshot_sha256": source["source_snapshot_sha256"], + "bundle_fingerprint": manifest["fingerprint"], + "event_count": len(events), + "question_count": len(questions), + "python": sys.version.split()[0], + "lotus": version("lotus-ai"), + "sentence_transformers": version("sentence-transformers"), + "torch": torch.__version__, + "cuda_device": torch.cuda.get_device_name(0), + "client_request_max_workers": config.lm_max_batch_size, + "deepseek_balance_cny": str(balances[0]), + "data_disk_available_gib": round(_disk_available_gib(Path("/mnt/data")), 3), + } + _atomic_json(paths.control / "preflight.json", result) + _set_state(paths, "experiment-state", "prepared") + _log(paths, "PREFLIGHT_PASSED " + json.dumps(result, sort_keys=True)) + return result + + +def _runtime_state(output: Path) -> tuple[str, dict[str, Any]]: + cases = [path for path in (output / "cases").iterdir() if path.is_dir()] + if len(cases) != 1: + raise RuntimeError(f"expected one case directory: {cases}") + current = json.loads((cases[0] / "checkpoints" / "current.json").read_text()) + checkpoint_id = str(current["checkpoint_id"]) + runtime_path = ( + cases[0] + / "checkpoints" + / "snapshots" + / checkpoint_id + / "driver" + / "runtime.pkl" + ) + return checkpoint_id, pickle.loads(runtime_path.read_bytes()) + + +def _json_safe(value: Any) -> Any: + if value is None or isinstance(value, (str, int, bool)): + return value + if isinstance(value, float): + return None if not math.isfinite(value) else value + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + if hasattr(value, "item"): + return _json_safe(value.item()) + if hasattr(value, "isoformat"): + return value.isoformat() + return str(value) + + +def _frame_digest(frame: pd.DataFrame) -> dict[str, Any]: + rows = [ + {str(column): _json_safe(row[column]) for column in frame.columns} + for _, row in frame.iterrows() + ] + signatures = sorted( + json.dumps(row, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + for row in rows + ) + return { + "row_count": len(rows), + "multiset_digest": sha256("\n".join(signatures).encode()).hexdigest(), + } + + +def _peak_rss_kib(time_path: Path) -> int | None: + if not time_path.is_file(): + return None + for line in time_path.read_text(errors="replace").splitlines(): + if "Maximum resident set size (kbytes)" in line: + return int(line.rsplit(":", 1)[1].strip()) + return None + + +def _expected_prompt_contract(condition: ConditionSpec) -> dict[str, Any] | None: + if condition.prompt_batch is None: + return None + return { + "max_tasks": None + if condition.prompt_batch == "all" + else int(condition.prompt_batch) + } + + +def _validate_condition( + paths: ExperimentPaths, condition: ConditionSpec +) -> dict[str, Any]: + directory = paths.conditions / condition.run_id + output = directory / "output" + trace = output / "trace" / "events.jsonl" + monitor = _update_monitor(trace, directory / "monitor-state.json") + if monitor["provider_contract_violations"]: + raise RuntimeError(str(monitor["provider_contract_violations"])) + if monitor["reasoning_tokens"] != 0: + raise RuntimeError("reasoning tokens were observed") + if monitor["insertion_results"] != EVENT_COUNT: + raise RuntimeError(f"expected {EVENT_COUNT} insertion results: {monitor}") + summary = json.loads((output / "metrics" / "summary.json").read_text()) + if ( + summary.get("completed_cases") != 1 + or summary.get("failed_cases") != 0 + or summary.get("question_count") != QUESTION_COUNT + or summary.get("memory_system_error_question_count") != 0 + or summary.get("retrieval_system_error_count") != 0 + ): + raise RuntimeError(f"benchmark completion mismatch: {summary}") + scores = summary.get("scores_by_scorer") or {} + if (scores.get(OFFICIAL_SCORER_ID) or {}).get("grade_count") != QUESTION_COUNT: + raise RuntimeError(f"official LOCOMO grades are incomplete: {scores}") + if (scores.get(JUDGE_SCORER_ID) or {}).get("grade_count") != JUDGED_QUESTION_COUNT: + raise RuntimeError(f"Zep judge grades are incomplete: {scores}") + manifest = json.loads((output / "manifest.json").read_text()) + if manifest.get("maintenance_rule") != "rule-join-map": + raise RuntimeError("maintenance rule drifted") + if manifest.get("retrieval_recipe_id") != "claude-memory-declared-sem-topk:v1:listwise": + raise RuntimeError("retrieval method is not listwise") + if manifest.get("framework_cache_mode") != "disabled": + raise RuntimeError("LOTUS cache must remain disabled") + runtime = manifest.get("runtime_provenance") or {} + execution = runtime.get("lotus_execution") or {} + if ( + execution.get("semantic_pair_profile") != "search-filter" + or execution.get("semantic_pair_top_k") != 20 + or not _number_equals(execution.get("semantic_pair_min_similarity"), 0.5) + or execution.get("embedding_device") != "cuda" + or execution.get("semantic_trace_snapshot_mode") != "compact" + or execution.get("prompt_batching") != _expected_prompt_contract(condition) + ): + raise RuntimeError(f"physical execution contract drifted: {execution}") + if execution.get("lm_max_ctx_len") != LM_MAX_CTX_LEN: + raise RuntimeError("memory LM context capacity changed") + profiles = execution.get("semantic_pair_query_profiles") + if not isinstance(profiles, Mapping) or len(profiles) != 1: + raise RuntimeError(f"expected one Claude Search-Filter site: {profiles}") + profile = next(iter(profiles.values())) + embedding = profile.get("embedding") or {} + if ( + profile.get("mode") != "search-filter" + or profile.get("direction") != "left-to-right" + or profile.get("top_k") != 20 + or not _number_equals(profile.get("min_similarity"), 0.5) + or profile.get("embedding_device") != "cuda" + or embedding.get("model") != "BAAI/bge-m3" + or embedding.get("revision") != BGE_M3_REVISION + ): + raise RuntimeError(f"resolved Search-Filter profile drifted: {profile}") + refresh = runtime.get("refresh") + expected_refresh = ( + None + if condition.refresh_every == 1 + else {"type": "count", "every": condition.refresh_every} + ) + if refresh != expected_refresh: + raise RuntimeError(f"refresh contract mismatch: {refresh}") + batching = monitor["prompt_batching"] + if condition.prompt_batch is None and batching: + raise RuntimeError("Original emitted prompt-batching trace events") + if condition.prompt_batch is not None and not batching: + raise RuntimeError("Packed condition emitted no prompt-batching trace events") + if condition.prompt_batch not in (None, "all"): + oversized = { + site: row + for site, row in batching.items() + if int(row.get("max_chunk") or 0) > int(condition.prompt_batch) + } + if oversized: + raise RuntimeError(f"prompt chunk exceeded configured P: {oversized}") + if not any(site.startswith("sem_join:") for site in monitor["candidate_pairs"]): + raise RuntimeError("Claude Search-Filter did not exercise sem_join") + + checkpoint_id, runtime_state = _runtime_state(output) + engine_state = ( + runtime_state + if condition.refresh_every == 1 + else runtime_state.get("engine_state") + ) + if condition.refresh_every > 1: + pending = runtime_state.get("pending_rows") + if not isinstance(pending, pd.DataFrame) or not pending.empty: + raise RuntimeError("final refresh checkpoint contains pending rows") + state = engine_state.get("state") if isinstance(engine_state, Mapping) else None + if not isinstance(state, Mapping) or not all( + isinstance(value, pd.DataFrame) for value in state.values() + ): + raise RuntimeError("runtime checkpoint lacks relational state") + if "log" not in state or len(state["log"]) != EVENT_COUNT: + raise RuntimeError("final log relation does not contain 128 rows") + final_state = {name: _frame_digest(frame) for name, frame in sorted(state.items())} + _atomic_json(directory / "final-state.json", final_state) + timing = json.loads((directory / "timing.json").read_text()) + insertion = summary["insertion_wall_latency_excluding_trace_io"] + provider = summary["actual_provider_usage"] + result = { + "schema_version": 1, + **condition.to_dict(), + "checkpoint_id": checkpoint_id, + "condition_end_to_end_wall_seconds": float(timing["wall_seconds"]), + "peak_rss_kib": _peak_rss_kib(directory / "time.txt"), + "insertion_excluding_trace_mean_ms": insertion["mean_ms"], + "insertion_excluding_trace_median_ms": insertion["median_ms"], + "insertion_excluding_trace_p95_ms": insertion["p95_ms"], + "insertion_excluding_trace_max_ms": insertion["max_ms"], + "retrieval_mean_ms": summary["retrieval_wall_latency"]["mean_ms"], + "provider_call_count": provider["provider_call_count"], + "cache_hit_tokens": provider["cache_hit_tokens"], + "cache_miss_tokens": provider["cache_miss_tokens"], + "completion_tokens": provider["completion_tokens"], + "physical_tokens": provider["total_tokens"], + "cost_cny": monitor["cost_cny"], + "official_locomo_score": scores[OFFICIAL_SCORER_ID]["mean_score"], + "zep_judge_score": scores[JUDGE_SCORER_ID]["mean_score"], + "candidate_pairs": monitor["candidate_pairs"], + "operator_usage": monitor["operator_usage"], + "prompt_batching": batching, + "provider_usage_by_phase": monitor["phase_usage"], + "final_state": final_state, + "artifact_bytes": _directory_bytes(output), + } + _atomic_json(directory / "validation.json", result) + return result + + +def _safe_stop( + paths: ExperimentPaths, process: subprocess.Popen[Any], *, reason: str +) -> None: + try: + _log(paths, f"SAFETY_STOP reason={reason} pid={process.pid}") + _set_state(paths, "experiment-state", "safety-stopped") + finally: + # A full disk or broken status file must not prevent stopping paid work. + try: + os.killpg(process.pid, signal.SIGINT) + except ProcessLookupError: + pass + else: + try: + process.wait(timeout=60) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGTERM) + process.wait(timeout=30) + + +def _before_condition(paths: ExperimentPaths) -> None: + cost = _total_cost(paths) + if cost >= SOFT_STOP_CNY: + raise RuntimeError(f"48 CNY stop before next condition: {cost}") + if _disk_available_gib(Path("/mnt/data")) < MIN_DISK_GIB: + raise RuntimeError("disk safety boundary reached") + + +def _condition_environment(paths: ExperimentPaths) -> dict[str, str]: + environment = dict(os.environ) + python_paths = [str(paths.source / "src")] + if paths.deps is not None: + python_paths.append(str(paths.deps)) + environment["PYTHONPATH"] = os.pathsep.join(python_paths) + environment["HF_HOME"] = str(paths.hf_home) + environment[SOURCE_EVIDENCE_ENV] = str(paths.control / "source-evidence.json") + return environment + + +def run_experiment(root: Path) -> None: + """Run pass one, then pass two only after pass one fully succeeds.""" + + contract = _read_contract(root) + paths = ExperimentPaths.from_contract(root, contract) + if _state(paths, "experiment-state") != "prepared": + raise RuntimeError("experiment must be prepared by preflight") + conditions = tuple(_condition_from_dict(item) for item in contract["conditions"]) + _validate_continuation(contract) + reused = (contract.get("continuation") or {}).get("reused_conditions", {}) + _set_state(paths, "experiment-state", "running") + _log(paths, f"EXPERIMENT_START conditions={len(conditions)}") + for condition in conditions: + if condition.run_id in reused: + _log(paths, f"REUSED_COMPLETE run_id={condition.run_id} origin={reused[condition.run_id]}") + continue + if condition.pass_index == 2: + incomplete = [ + item.run_id + for item in conditions + if item.pass_index == 1 + and item.run_id not in reused + and ( + not (paths.conditions / item.run_id / "state").is_file() + or (paths.conditions / item.run_id / "state").read_text().strip() + != "completed" + ) + ] + if incomplete: + raise RuntimeError(f"second pass cannot start: {incomplete}") + _before_condition(paths) + directory = paths.conditions / condition.run_id + directory.mkdir(parents=True, exist_ok=True) + state_path = directory / "state" + if state_path.exists(): + raise RuntimeError(f"condition already has state: {condition.run_id}") + _atomic_json( + directory / "condition-contract.json", + { + "condition": condition.to_dict(), + "experiment_contract_sha256": _sha256_file( + paths.control / "experiment-contract.json" + ), + "client_request_max_workers": CLIENT_REQUEST_MAX_WORKERS, + }, + ) + _atomic_text(state_path, "running\n") + _set_state(paths, "current-condition", condition.run_id) + output = directory / "output" + command = build_condition_command( + condition, + source=paths.source, + venv=paths.venv, + bundle=paths.bundle, + output=output, + ) + _atomic_json(directory / "command.json", command) + _log( + paths, + f"CONDITION_START run_id={condition.run_id} " + f"R={condition.refresh_every} P={condition.prompt_batch}", + ) + started = time.monotonic() + with (directory / "condition.log").open("ab") as log: + process = subprocess.Popen( + ("/usr/bin/time", "-v", "-o", str(directory / "time.txt"), *command), + stdout=log, + stderr=subprocess.STDOUT, + env=_condition_environment(paths), + start_new_session=True, + ) + try: + _atomic_text(paths.control / "runner.pid", f"{process.pid}\n") + while process.poll() is None: + trace = output / "trace" / "events.jsonl" + if trace.exists() and trace.stat().st_size: + monitor = _update_monitor(trace, directory / "monitor-state.json") + if monitor["reasoning_tokens"] != 0: + _safe_stop(paths, process, reason="reasoning_tokens_nonzero") + raise RuntimeError("reasoning tokens were observed") + if monitor["provider_contract_violations"]: + _safe_stop(paths, process, reason="provider_contract_violation") + raise RuntimeError(str(monitor["provider_contract_violations"])) + if _total_cost(paths) >= HARD_STOP_CNY: + _safe_stop(paths, process, reason="50_cny_hard_stop") + raise RuntimeError("50 CNY safety boundary reached") + if _directory_bytes(output) > MAX_CONDITION_BYTES: + _safe_stop(paths, process, reason="artifact_size_limit") + raise RuntimeError("condition artifact exceeded 2 GiB") + if _disk_available_gib(Path("/mnt/data")) < MIN_DISK_GIB: + _safe_stop(paths, process, reason="disk_limit") + raise RuntimeError("disk safety boundary reached") + if time.monotonic() - started > MAX_CONDITION_SECONDS: + _safe_stop(paths, process, reason="condition_timeout") + raise RuntimeError("condition exceeded its six-hour limit") + time.sleep(POLL_SECONDS) + except BaseException as error: + if process.poll() is None: + try: + _safe_stop(paths, process, reason="supervision_error") + except Exception as stop_error: + error.add_note( + f"Child cleanup failed: {type(stop_error).__name__}: {stop_error}" + ) + raise + _atomic_json( + directory / "timing.json", + {"wall_seconds": time.monotonic() - started}, + ) + if process.returncode != 0: + _atomic_text(state_path, "failed\n") + raise RuntimeError( + f"condition failed: {condition.run_id} status={process.returncode}" + ) + try: + result = _validate_condition(paths, condition) + except Exception: + _atomic_text(state_path, "validation-failed\n") + raise + _atomic_text(state_path, "completed\n") + _log( + paths, + f"CONDITION_COMPLETE run_id={condition.run_id} " + f"cost_cny={result['cost_cny']:.6f}", + ) + if condition.run_id == "pass-1-packed-all": + _log(paths, "PASS_COMPLETE pass=1") + summarize(root) + _set_state(paths, "current-condition", "none") + _set_state(paths, "experiment-state", "completed") + _log(paths, f"EXPERIMENT_COMPLETE total_cost_cny={_total_cost(paths)}") + + +def summarize(root: Path) -> dict[str, Any]: + """Write a compact machine-readable summary of all completed conditions.""" + + contract = _read_contract(root) + paths = ExperimentPaths.from_contract(root, contract) + results = [] + for item in contract["conditions"]: + run_id = str(item["run_id"]) + directory = _result_directory(root, contract, run_id) + path = directory / "validation.json" + if path.exists(): + result = json.loads(path.read_text()) + is_reused = directory != paths.conditions / run_id + if is_reused: + result["cost_cny"] = float(contract["continuation"]["prior_condition_costs_cny"][run_id]) + result["result_origin"] = { + "status": "reused-completed" if is_reused else "completed", + "directory": str(directory), + "source": ( + contract["continuation"].get("reused_sources", {}).get( + run_id, contract["continuation"]["source"], + ) if is_reused else contract["source"] + ), + "repair_version": None if is_reused else (contract.get("structured_output") or {}).get("repair_version"), + } + results.append(result) + payload = { + "schema_version": 1, + "experiment": contract["experiment"], + "completed_conditions": len(results), + "expected_conditions": len(contract["conditions"]), + "total_cost_cny": float(_total_cost(paths)), + "continuation": contract.get("continuation"), + "runs": results, + "interpretation_boundary": ( + "The sweep co-varies CountRefresh R and PromptBatching P. The two passes " + "are exploratory repeatability evidence, not a variance estimate or a " + "claim of a universal optimum." + ), + } + _atomic_json(paths.control / "sweep-summary.json", payload) + return payload + + +def status(root: Path) -> None: + """Print the authoritative, read-only experiment status.""" + + contract = _read_contract(root) + paths = ExperimentPaths.from_contract(root, contract) + current = _state(paths, "current-condition") or "none" + print(f"time_hkt={datetime.now(tz=HONG_KONG).isoformat(timespec='seconds')}") + print(f"run_root={paths.root}") + print(f"experiment_state={_state(paths, 'experiment-state') or 'unknown'}") + print(f"current_condition={current}") + print( + f"total_cost_cny={_total_cost(paths):.6f} " + f"soft_stop={SOFT_STOP_CNY} hard_stop={HARD_STOP_CNY}" + ) + completed = failed = 0 + for item in contract["conditions"]: + condition = _condition_from_dict(item) + directory = _result_directory(root, contract, condition.run_id) + state = (directory / "state").read_text().strip() if (directory / "state").exists() else "waiting" + if directory != paths.conditions / condition.run_id: + state = "reused-completed" + completed += state in {"completed", "reused-completed"} + failed += state in {"failed", "validation-failed"} + monitor_path = directory / "monitor-state.json" + suffix = "" + if state == "reused-completed": + suffix = f" events={EVENT_COUNT}/{EVENT_COUNT} questions={QUESTION_COUNT}/{QUESTION_COUNT} origin={directory}" + print(f"{condition.run_id}: {state}{suffix}") + continue + if monitor_path.exists(): + monitor = json.loads(monitor_path.read_text()) + output = directory / "output" + suffix = ( + f" events={monitor.get('insertion_results', 0)}/{EVENT_COUNT}" + f" questions={_completed_questions(output)}/{QUESTION_COUNT}" + f" responses={monitor.get('provider_responses', 0)}" + f" cost_cny={float(monitor.get('cost_cny') or 0):.6f}" + ) + print(f"{condition.run_id}: {state}{suffix}") + print(f"progress={completed}/{len(contract['conditions'])} failed={failed}") + print(f"data_disk_available_gib={_disk_available_gib(Path('/mnt/data')):.3f}") + gpu = subprocess.run( + ( + "nvidia-smi", + "--query-gpu=name,utilization.gpu,memory.used,memory.total,power.draw", + "--format=csv,noheader,nounits", + ), + capture_output=True, + text=True, + check=False, + ) + if gpu.stdout.strip(): + print(f"gpu={gpu.stdout.strip()}") + log_path = paths.control / "controller.log" + if log_path.exists(): + print("recent_controller_log:") + for line in log_path.read_text(errors="replace").splitlines()[-12:]: + print(line) + if current != "none": + condition_log = paths.conditions / current / "condition.log" + if condition_log.exists(): + errors = [ + line + for line in condition_log.read_text(errors="replace").splitlines() + if any(token in line for token in ("Traceback", "Error", "Exception")) + ] + if errors: + print("recent_error:") + for line in errors[-6:]: + print(line) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + initialize_parser = commands.add_parser("initialize") + initialize_parser.add_argument("--run-root", type=Path, required=True) + initialize_parser.add_argument("--source-dir", type=Path, required=True) + initialize_parser.add_argument("--venv-dir", type=Path, required=True) + initialize_parser.add_argument("--deps-dir", type=Path) + initialize_parser.add_argument("--hf-home", type=Path, required=True) + initialize_parser.add_argument("--dataset-path", type=Path, required=True) + initialize_parser.add_argument("--reuse-root", type=Path) + initialize_parser.add_argument("--reuse-condition", action="append", default=[]) + initialize_parser.add_argument( + "--source-package-inventory", type=Path, required=True + ) + for name in ("preflight", "run", "status", "summarize"): + command = commands.add_parser(name) + command.add_argument("--run-root", type=Path, required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> None: + """Run the requested experiment-controller command.""" + + args = _parser().parse_args(argv) + if args.command == "initialize": + initialize( + root=args.run_root, + source=args.source_dir, + venv=args.venv_dir, + deps=args.deps_dir, + hf_home=args.hf_home, + dataset=args.dataset_path, + source_package_inventory=args.source_package_inventory, + reuse_root=args.reuse_root, + reuse_conditions=args.reuse_condition, + ) + return + try: + if args.command == "preflight": + print(json.dumps(preflight(args.run_root), sort_keys=True)) + elif args.command == "run": + run_experiment(args.run_root) + elif args.command == "status": + status(args.run_root) + else: + print(json.dumps(summarize(args.run_root), sort_keys=True)) + except BaseException as error: + if args.command == "run": + paths = _paths(args.run_root) + if _state(paths, "experiment-state") != "safety-stopped": + _set_state(paths, "experiment-state", "failed") + _log(paths, f"EXPERIMENT_FAILED {type(error).__name__}: {error}") + traceback.print_exc() + raise + + +if __name__ == "__main__": + main() diff --git a/tools/evaluation/zep_locomo.py b/tools/evaluation/zep_locomo.py index 18b8514..f356ade 100644 --- a/tools/evaluation/zep_locomo.py +++ b/tools/evaluation/zep_locomo.py @@ -20,7 +20,10 @@ ZepLocomoRunConfig, run_zep_locomo, ) -from agent_memory.planner import GROUPED_AGG_RULES # noqa: E402 +from agent_memory.planner import ( # noqa: E402 + DEFAULT_GROUPED_AGG_RULE, + GROUPED_AGG_RULES, +) LOCOMO_CACHE_PATH = PROJECT_ROOT / ".cache" / "agent-memory" / "locomo10.json" @@ -66,7 +69,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument( "--grouped-agg-rule", choices=GROUPED_AGG_RULES, - default="rule-re-group", + default=DEFAULT_GROUPED_AGG_RULE, ) parser.add_argument("--namespace", default=None) parser.add_argument("--output-dir", type=Path, default=default_output_dir()) diff --git a/uv.lock b/uv.lock index 56022aa..6a6a5e0 100644 --- a/uv.lock +++ b/uv.lock @@ -12,6 +12,7 @@ version = "0.0.1" source = { editable = "." } dependencies = [ { name = "dspy" }, + { name = "json5" }, { name = "lotus-ai" }, { name = "nltk" }, { name = "pandas" }, @@ -48,6 +49,7 @@ dev = [ requires-dist = [ { name = "dspy", specifier = ">=3.2.0" }, { name = "editdistance", marker = "extra == 'benchmarks'", specifier = ">=0.8.1" }, + { name = "json5", specifier = ">=0.14.0" }, { name = "lotus-ai", specifier = "==1.1.4" }, { name = "neo4j", marker = "extra == 'zep'", specifier = "==6.1.0" }, { name = "nltk", specifier = ">=3.9.1" },