2009YU commited on
Commit
490f3fe
·
verified ·
1 Parent(s): 4b67d61

Release TMCRA Agent Memory v2026.07.18

Browse files

Publish the Apache-2.0 TMCRA Agent Memory release, LongMemEval S500 scorecard (82.2%), bilingual model card, reproducibility pipeline, and sanitized inference-only graph scoring weights. The previous default snapshot is preserved as legacy-2026.05.25.

This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. AUTHORS.md +25 -0
  2. CHANGELOG.md +16 -0
  3. CITATION.cff +27 -0
  4. LICENSE +201 -21
  5. NOTICE +9 -0
  6. README.md +81 -173
  7. README.zh-CN.md +66 -173
  8. SECURITY.md +12 -0
  9. benchmarks/longmemeval/README.md +216 -0
  10. benchmarks/longmemeval/README.zh-CN.md +216 -0
  11. benchmarks/longmemeval/THIRD_PARTY_NOTICES.md +52 -0
  12. benchmarks/longmemeval/configs/answer.env.example +6 -0
  13. benchmarks/longmemeval/configs/assets.lock.json +51 -0
  14. benchmarks/longmemeval/configs/judge.env.example +9 -0
  15. benchmarks/longmemeval/configs/run.env.example +18 -0
  16. benchmarks/longmemeval/configs/writer.env.example +12 -0
  17. benchmarks/longmemeval/fixtures/tiny_longmemeval.json +22 -0
  18. benchmarks/longmemeval/pyproject.toml +30 -0
  19. benchmarks/longmemeval/requirements.txt +8 -0
  20. benchmarks/longmemeval/results/benchmark.json +16 -0
  21. benchmarks/longmemeval/scripts/fetch_assets.py +95 -0
  22. benchmarks/longmemeval/scripts/prepare_longmemeval.py +70 -0
  23. benchmarks/longmemeval/scripts/reproduce_s500.sh +152 -0
  24. benchmarks/longmemeval/scripts/reproduce_smoke.sh +18 -0
  25. benchmarks/longmemeval/scripts/summarize_official_judge.py +82 -0
  26. benchmarks/longmemeval/scripts/validate_stage.py +237 -0
  27. benchmarks/longmemeval/src/tmcra_benchmark/__init__.py +3 -0
  28. benchmarks/longmemeval/src/tmcra_benchmark/harness/__init__.py +1 -0
  29. benchmarks/longmemeval/src/tmcra_benchmark/harness/evaluate_qa_official_judge.py +295 -0
  30. benchmarks/longmemeval/src/tmcra_benchmark/harness/run_lme_s10_native_tmcra.py +0 -0
  31. benchmarks/longmemeval/src/tmcra_benchmark/legacy/__init__.py +1 -0
  32. benchmarks/longmemeval/src/tmcra_benchmark/legacy/audit_tmcra_v4_chain.py +0 -0
  33. benchmarks/longmemeval/src/tmcra_benchmark/legacy/build_v3_runtime_dataset.py +791 -0
  34. benchmarks/longmemeval/src/tmcra_benchmark/legacy/ops/audit_tmcra_v4_subject_attribution.py +877 -0
  35. benchmarks/longmemeval/src/tmcra_benchmark/legacy/prepare_tmcra_v4_e2e_data.py +261 -0
  36. benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_build.py +984 -0
  37. benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_compile_evidence.py +1015 -0
  38. benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_evaluate.py +725 -0
  39. benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_gpt54_answers.py +847 -0
  40. benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_parallel_official_judge.py +273 -0
  41. benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_retrieve.py +524 -0
  42. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_model_assets.py +29 -0
  43. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v2_lme_pipeline.py +855 -0
  44. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_online_runtime.py +1867 -0
  45. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_product_writer.py +0 -0
  46. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_recall_planner.py +406 -0
  47. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_reranker.py +1014 -0
  48. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_schema.py +168 -0
  49. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_slow_graph.py +2138 -0
  50. benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v4_batch_writer.py +0 -0
AUTHORS.md ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Developers and attribution
2
+
3
+ ## Core project
4
+
5
+ - **Yu Haoxin** ([@reshuibuduo](https://github.com/reshuibuduo)) — project creator, lead developer, and TMCRA algorithm engineering.
6
+
7
+ ## Development assistance
8
+
9
+ - **OpenAI Codex** — development and reproducibility engineering assistant, including code organization, benchmark packaging, documentation, and release validation support.
10
+
11
+ OpenAI Codex is credited here as an AI development assistant; project ownership and release responsibility remain with Yu Haoxin.
12
+
13
+ ---
14
+
15
+ # 开发者与署名
16
+
17
+ ## 核心项目
18
+
19
+ - **Yu Haoxin**([@reshuibuduo](https://github.com/reshuibuduo))— 项目创建者、主要开发者与 TMCRA 算法工程。
20
+
21
+ ## 开发协助
22
+
23
+ - **OpenAI Codex** — 开发与可复现性工程助手,协助完成代码整理、benchmark 封装、文档编写和发布校验。
24
+
25
+ OpenAI Codex 在此作为 AI 开发助手署名;项目所有权与发布责任归 Yu Haoxin 所有。
CHANGELOG.md ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Changelog
2
+
3
+ ## 2026.07.18
4
+
5
+ - Repositioned TMCRA as an Agent Memory Engine with a clear Source / Fast /
6
+ Slow layered architecture and scope isolation.
7
+ - Added the maintained LongMemEval-S500 reproduction package.
8
+ - Published the 82.2% scorecard with six task-level results.
9
+ - Added the latest node scorer, path scorer, and retrieval reranker weights.
10
+ - Hardened checkpoint loading with `weights_only=True`, published an
11
+ inference-only reranker checkpoint without training-machine paths, pinned
12
+ third-party model revisions, and added offline CI and gold-isolation checks.
13
+ - Versioned the maintained answer contract as `evidence_operation_bound_v6`
14
+ and made compatible interrupted benchmark stages resumable.
15
+ - Switched the project license to Apache-2.0.
16
+ - Credited Yu Haoxin and OpenAI Codex for their respective development roles.
CITATION.cff ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ cff-version: 1.2.0
2
+ message: "If you use TMCRA, please cite this software. OpenAI Codex assisted with development and reproducibility engineering."
3
+ title: "TMCRA — Agent Memory Engine"
4
+ type: software
5
+ version: "2026.07.18"
6
+ date-released: 2026-07-18
7
+ license: Apache-2.0
8
+ authors:
9
+ - family-names: Yu
10
+ given-names: Haoxin
11
+ alias: reshuibuduo
12
+ website: "https://github.com/reshuibuduo"
13
+ - name: "OpenAI Codex"
14
+ website: "https://openai.com/codex/"
15
+ repository-code: "https://github.com/reshuibuduo/TMCRA-Agent-Memory"
16
+ url: "https://github.com/reshuibuduo/TMCRA-Agent-Memory"
17
+ abstract: >-
18
+ TMCRA is an agent memory engine that builds scope-isolated, source-traceable
19
+ memory from long-running conversations and compiles relevant evidence for
20
+ downstream agents. This release includes the algorithm snapshot, learned
21
+ graph-scoring artifacts, and a reproducible LongMemEval pipeline.
22
+ keywords:
23
+ - agent memory
24
+ - long-term memory
25
+ - memory graph
26
+ - retrieval
27
+ - LongMemEval
LICENSE CHANGED
@@ -1,21 +1,201 @@
1
- MIT License
2
-
3
- Copyright (c) 2026 TMCRA contributors
4
-
5
- Permission is hereby granted, free of charge, to any person obtaining a copy
6
- of this software and associated documentation files (the "Software"), to deal
7
- in the Software without restriction, including without limitation the rights
8
- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
- copies of the Software, and to permit persons to whom the Software is
10
- furnished to do so, subject to the following conditions:
11
-
12
- The above copyright notice and this permission notice shall be included in all
13
- copies or substantial portions of the Software.
14
-
15
- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
- SOFTWARE.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
6
+
7
+ 1. Definitions.
8
+
9
+ "License" shall mean the terms and conditions for use, reproduction,
10
+ and distribution as defined by Sections 1 through 9 of this document.
11
+
12
+ "Licensor" shall mean the copyright owner or entity authorized by
13
+ the copyright owner that is granting the License.
14
+
15
+ "Legal Entity" shall mean the union of the acting entity and all
16
+ other entities that control, are controlled by, or are under common
17
+ control with that entity. For the purposes of this definition,
18
+ "control" means (i) the power, direct or indirect, to cause the
19
+ direction or management of such entity, whether by contract or
20
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
21
+ outstanding shares, or (iii) beneficial ownership of such entity.
22
+
23
+ "You" (or "Your") shall mean an individual or Legal Entity
24
+ exercising permissions granted by this License.
25
+
26
+ "Source" form shall mean the preferred form for making modifications,
27
+ including but not limited to software source code, documentation
28
+ source, and configuration files.
29
+
30
+ "Object" form shall mean any form resulting from mechanical
31
+ transformation or translation of a Source form, including but
32
+ not limited to compiled object code, generated documentation,
33
+ and conversions to other media types.
34
+
35
+ "Work" shall mean the work of authorship, whether in Source or
36
+ Object form, made available under the License, as indicated by a
37
+ copyright notice that is included in or attached to the work
38
+ (an example is provided in the Appendix below).
39
+
40
+ "Derivative Works" shall mean any work, whether in Source or Object
41
+ form, that is based on (or derived from) the Work and for which the
42
+ editorial revisions, annotations, elaborations, or other modifications
43
+ represent, as a whole, an original work of authorship. For the purposes
44
+ of this License, Derivative Works shall not include works that remain
45
+ separable from, or merely link (or bind by name) to the interfaces of,
46
+ the Work and Derivative Works thereof.
47
+
48
+ "Contribution" shall mean any work of authorship, including
49
+ the original version of the Work and any modifications or additions
50
+ to that Work or Derivative Works thereof, that is intentionally
51
+ submitted to Licensor for inclusion in the Work by the copyright owner
52
+ or by an individual or Legal Entity authorized to submit on behalf of
53
+ the copyright owner. For the purposes of this definition, "submitted"
54
+ means any form of electronic, verbal, or written communication sent
55
+ to the Licensor or its representatives, including but not limited to
56
+ communication on electronic mailing lists, source code control systems,
57
+ and issue tracking systems that are managed by, or on behalf of, the
58
+ Licensor for the purpose of discussing and improving the Work, but
59
+ excluding communication that is conspicuously marked or otherwise
60
+ designated in writing by the copyright owner as "Not a Contribution."
61
+
62
+ "Contributor" shall mean Licensor and any individual or Legal Entity
63
+ on behalf of whom a Contribution has been received by Licensor and
64
+ subsequently incorporated within the Work.
65
+
66
+ 2. Grant of Copyright License. Subject to the terms and conditions of
67
+ this License, each Contributor hereby grants to You a perpetual,
68
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
69
+ copyright license to reproduce, prepare Derivative Works of,
70
+ publicly display, publicly perform, sublicense, and distribute the
71
+ Work and such Derivative Works in Source or Object form.
72
+
73
+ 3. Grant of Patent License. Subject to the terms and conditions of
74
+ this License, each Contributor hereby grants to You a perpetual,
75
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
76
+ (except as stated in this section) patent license to make, have made,
77
+ use, offer to sell, sell, import, and otherwise transfer the Work,
78
+ where such license applies only to those patent claims licensable
79
+ by such Contributor that are necessarily infringed by their
80
+ Contribution(s) alone or by combination of their Contribution(s)
81
+ with the Work to which such Contribution(s) was submitted. If You
82
+ institute patent litigation against any entity (including a
83
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
84
+ or a Contribution incorporated within the Work constitutes direct
85
+ or contributory patent infringement, then any patent licenses
86
+ granted to You under this License for that Work shall terminate
87
+ as of the date such litigation is filed.
88
+
89
+ 4. Redistribution. You may reproduce and distribute copies of the
90
+ Work or Derivative Works thereof in any medium, with or without
91
+ modifications, and in Source or Object form, provided that You
92
+ meet the following conditions:
93
+
94
+ (a) You must give any other recipients of the Work or
95
+ Derivative Works a copy of this License; and
96
+
97
+ (b) You must cause any modified files to carry prominent notices
98
+ stating that You changed the files; and
99
+
100
+ (c) You must retain, in the Source form of any Derivative Works
101
+ that You distribute, all copyright, patent, trademark, and
102
+ attribution notices from the Source form of the Work,
103
+ excluding those notices that do not pertain to any part of
104
+ the Derivative Works; and
105
+
106
+ (d) If the Work includes a "NOTICE" text file as part of its
107
+ distribution, then any Derivative Works that You distribute must
108
+ include a readable copy of the attribution notices contained
109
+ within such NOTICE file, excluding those notices that do not
110
+ pertain to any part of the Derivative Works, in at least one
111
+ of the following places: within a NOTICE text file distributed
112
+ as part of the Derivative Works; within the Source form or
113
+ documentation, if provided along with the Derivative Works; or,
114
+ within a display generated by the Derivative Works, if and
115
+ wherever such third-party notices normally appear. The contents
116
+ of the NOTICE file are for informational purposes only and
117
+ do not modify the License. You may add Your own attribution
118
+ notices within Derivative Works that You distribute, alongside
119
+ or as an addendum to the NOTICE text from the Work, provided
120
+ that such additional attribution notices cannot be construed
121
+ as modifying the License.
122
+
123
+ You may add Your own copyright statement to Your modifications and
124
+ may provide additional or different license terms and conditions
125
+ for use, reproduction, or distribution of Your modifications, or
126
+ for any such Derivative Works as a whole, provided Your use,
127
+ reproduction, and distribution of the Work otherwise complies with
128
+ the conditions stated in this License.
129
+
130
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
131
+ any Contribution intentionally submitted for inclusion in the Work
132
+ by You to the Licensor shall be under the terms and conditions of
133
+ this License, without any additional terms or conditions.
134
+ Notwithstanding the above, nothing herein shall supersede or modify
135
+ the terms of any separate license agreement you may have executed
136
+ with Licensor regarding such Contributions.
137
+
138
+ 6. Trademarks. This License does not grant permission to use the trade
139
+ names, trademarks, service marks, or product names of the Licensor,
140
+ except as required for reasonable and customary use in describing the
141
+ origin of the Work and reproducing the content of the NOTICE file.
142
+
143
+ 7. Disclaimer of Warranty. Unless required by applicable law or
144
+ agreed to in writing, Licensor provides the Work (and each
145
+ Contributor provides its Contributions) on an "AS IS" BASIS,
146
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
147
+ implied, including, without limitation, any warranties or conditions
148
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
149
+ PARTICULAR PURPOSE. You are solely responsible for determining the
150
+ appropriateness of using or redistributing the Work and assume any
151
+ risks associated with Your exercise of permissions under this License.
152
+
153
+ 8. Limitation of Liability. In no event and under no legal theory,
154
+ whether in tort (including negligence), contract, or otherwise,
155
+ unless required by applicable law (such as deliberate and grossly
156
+ negligent acts) or agreed to in writing, shall any Contributor be
157
+ liable to You for damages, including any direct, indirect, special,
158
+ incidental, or consequential damages of any character arising as a
159
+ result of this License or out of the use or inability to use the
160
+ Work (including but not limited to damages for loss of goodwill,
161
+ work stoppage, computer failure or malfunction, or any and all
162
+ other commercial damages or losses), even if such Contributor
163
+ has been advised of the possibility of such damages.
164
+
165
+ 9. Accepting Warranty or Additional Liability. While redistributing
166
+ the Work or Derivative Works thereof, You may choose to offer,
167
+ and charge a fee for, acceptance of support, warranty, indemnity,
168
+ or other liability obligations and/or rights consistent with this
169
+ License. However, in accepting such obligations, You may act only
170
+ on Your own behalf and on Your sole responsibility, not on behalf
171
+ of any other Contributor, and only if You agree to indemnify,
172
+ defend, and hold each Contributor harmless for any liability
173
+ incurred by, or claims asserted against, such Contributor by reason
174
+ of your accepting any such warranty or additional liability.
175
+
176
+ END OF TERMS AND CONDITIONS
177
+
178
+ APPENDIX: How to apply the Apache License to your work.
179
+
180
+ To apply the Apache License to your work, attach the following
181
+ boilerplate notice, with the fields enclosed by brackets "[]"
182
+ replaced with your own identifying information. (Don't include
183
+ the brackets!) The text should be enclosed in the appropriate
184
+ comment syntax for the file format. We also recommend that a
185
+ file or class name and description of purpose be included on the
186
+ same "printed page" as the copyright notice for easier
187
+ identification within third-party archives.
188
+
189
+ Copyright 2026 Yu Haoxin and TMCRA contributors
190
+
191
+ Licensed under the Apache License, Version 2.0 (the "License");
192
+ you may not use this file except in compliance with the License.
193
+ You may obtain a copy of the License at
194
+
195
+ http://www.apache.org/licenses/LICENSE-2.0
196
+
197
+ Unless required by applicable law or agreed to in writing, software
198
+ distributed under the License is distributed on an "AS IS" BASIS,
199
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
200
+ See the License for the specific language governing permissions and
201
+ limitations under the License.
NOTICE ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ TMCRA — Agent Memory Engine
2
+ Copyright 2026 Yu Haoxin and TMCRA contributors
3
+
4
+ Development and reproducibility engineering assistance was provided by
5
+ OpenAI Codex. See AUTHORS.md for contributor roles.
6
+
7
+ This product includes or adapts third-party software and model interfaces.
8
+ The corresponding copyright notices and license terms are documented in
9
+ benchmarks/longmemeval/THIRD_PARTY_NOTICES.md.
README.md CHANGED
@@ -1,207 +1,115 @@
1
- # TMCRA Long-Memory Runtime
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
  <p align="center">
4
  <img src="assets/tmcra-logo.png" alt="TMCRA logo" width="420">
5
  </p>
6
 
7
- [中文版本](README.zh-CN.md)
8
-
9
- TMCRA is a runtime layer that gives agent systems long-term memory. It helps an LLM retrieve, connect, and reason over long dialogue histories without exposing the full conversation context on every turn.
10
-
11
- This repository contains a frozen TMCRA baseline package with model weights, runtime code snapshots, training metadata, and LongMemEval S500 benchmark results.
12
-
13
- License: MIT.
14
-
15
- ## What TMCRA Does
16
-
17
- TMCRA adds a dedicated memory runtime between an agent application and its answer model.
18
-
19
- At write time, TMCRA turns dialogue into memory nodes, event units, profile signals, and graph paths. This lets the system preserve not only isolated facts, but also relationships between facts across turns and sessions.
20
 
21
- At retrieval time, TMCRA scores graph nodes and paths, selects compact evidence, and injects only the most relevant memory context into the answer model. The answer model still performs natural-language reasoning, while TMCRA handles long-memory organization, recall, and evidence surfacing.
22
 
23
- The current runtime focuses on:
24
 
25
- - user fact memory
26
- - assistant-response memory
27
- - profile and preference memory
28
- - temporal memory
29
- - cross-session graph tunneling
30
- - learned node/path scoring
31
- - compact evidence selection for downstream LLMs
32
 
33
- ## Algorithm Structure
34
 
35
- ```mermaid
36
- flowchart LR
37
- A["Dialogue Turns"] --> B["Writer Layer"]
38
- B --> C["Memory Nodes"]
39
- B --> D["Event Units"]
40
- B --> E["Profile Signals"]
41
- B --> F["Temporal Signals"]
 
 
42
 
43
- C --> G["Graph Memory Store"]
44
- D --> G
45
- E --> G
46
- F --> G
47
 
48
- G --> H["Learned Node Scorer"]
49
- G --> I["Learned Path Scorer"]
50
- H --> J["Evidence Selection"]
51
- I --> J
52
 
53
- J --> K["Compact Memory Context"]
54
- K --> L["Answer Model"]
55
- L --> M["Agent Response"]
56
 
57
- N["User Query"] --> H
58
- N --> I
59
- N --> J
60
- N --> L
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
61
  ```
62
 
63
- The writer layer produces memory units from dialogue. The graph memory store preserves facts, events, profile signals, temporal signals, and cross-session links. The learned node/path scorers select relevant evidence for the current query, and the answer model uses that compact evidence to produce the final response.
64
 
65
- ## Why TMCRA
66
 
67
- Long-running agents need more than simple vector recall. They need to preserve user facts, preferences, timeline changes, cross-session events, and multi-step evidence chains.
68
 
69
- TMCRA organizes memory into graph nodes and learned retrieval paths, then surfaces compact evidence to the answer model. The goal is to let external agents use long-term memory through a runtime/API layer while keeping the memory algorithm and model weights independently deployable.
 
 
 
 
70
 
71
- ## How to Use
72
 
73
- For inference/runtime use, load the graph scorer weights from:
74
-
75
- ```text
76
- models/action_frame_tunnel_graph548_tunnel_fusion_train_20260524_042557/
77
- ```
78
 
79
- The main runtime files are:
80
-
81
- ```text
82
- node_scorer.pt
83
- path_scorer.pt
84
- export_manifest.json
85
- ```
86
-
87
- A typical runtime configuration points TMCRA to these weights:
88
 
89
  ```bash
90
- export TMCRA_NODE_MODEL_PATH="models/action_frame_tunnel_graph548_tunnel_fusion_train_20260524_042557/node_scorer.pt"
91
- export TMCRA_PATH_MODEL_PATH="models/action_frame_tunnel_graph548_tunnel_fusion_train_20260524_042557/path_scorer.pt"
92
- export TMCRA_RETRIEVAL_MODE="hybrid_node_scored"
93
- export TMCRA_REQUIRE_LEARNED_SCORER="1"
94
- ```
95
-
96
- The evaluation entrypoint snapshot is:
97
 
98
- ```text
99
- code/run_lme_s10_native_tmcra.py
100
- ```
101
-
102
- The core adapter snapshot is:
103
-
104
- ```text
105
- code/memory_adapters.py
106
  ```
107
 
108
- For a deployment build, load the two scorer files into the TMCRA adapter and point the agent's memory middleware to the TMCRA retrieval API. The answer model can be any OpenAI-compatible or local LLM endpoint; TMCRA supplies the selected memory evidence, and the answer model produces the final response.
109
-
110
- ## Dependency Environment
111
 
112
- The included code snapshot is Python-based. A practical runtime environment should include:
113
 
114
- - Python 3.10 or newer
115
- - PyTorch, with CUDA recommended for model inference
116
- - NumPy and standard Python data-processing libraries
117
- - an OpenAI-compatible or local LLM endpoint for the answer layer and writer layer
118
- - optional Git LFS support when pulling the full model package from GitHub
119
 
120
- The benchmark scripts expect LongMemEval-format input data and write JSONL predictions and judge outputs. Runtime deployments can use the same model files without running the benchmark harness.
121
 
122
- ## Optional Modules
 
123
 
124
- TMCRA also keeps optional extension points for retrieval and planning experiments. These modules can be enabled in deployment or evaluation builds when the target use case needs them.
125
 
126
- - **Embedder interface**: an optional semantic embedding channel that can run alongside the graph-memory scorer. It is intended to provide additional dense semantic recall before or during graph evidence selection, without replacing the learned graph node/path scorers.
127
- - **LLM planner interface**: an optional planner hook that can use an external LLM to organize evidence, expand query intent, or create an answer plan before the final answer call. This is useful for experiments and higher-cost deployments, while the default baseline keeps the core graph scorer path independently measurable.
128
 
129
- These interfaces are integration points, not required dependencies for the frozen S500 baseline. They are designed so downstream deployments can decide whether to run a lighter scorer-only path or a heavier path with embedder/planner assistance.
130
-
131
- ## Included Artifacts
132
-
133
- - `code/`: runtime and evaluation code snapshot for this baseline.
134
- - `models/action_frame_tunnel_graph548_tunnel_fusion_train_20260524_042557/`: full trained graph-model output directory.
135
- - `results/`: predictions, judge output, summary metrics, and compressed run artifacts.
136
- - `docs/`: baseline record and result notes.
137
-
138
- Additional documentation:
139
-
140
- - `docs/BASELINE_S500_20260525.md`: benchmark record and subtask metrics.
141
- - `docs/TRAINING.md`: model training direction and released training artifacts.
142
- - `docs/OPTIONAL_MODULES_AND_PARALLEL.md`: optional embedder/planner integration points and parallel evaluation plan.
143
- - `docs/OPTIONAL_MODULES_AND_PARALLEL.zh-CN.md`: Chinese version of the optional module and parallel evaluation guide.
144
-
145
- ## Model Package
146
-
147
- The included model package preserves the full training output for the graph scorer stack:
148
-
149
- - `node_scorer.pt` and `path_scorer.pt`: runtime graph scoring weights.
150
- - `node_scorer_best.pt` and `path_scorer_best.pt`: best checkpoint aliases.
151
- - `node_scorer_last.pt` and `path_scorer_last.pt`: final training aliases.
152
- - `checkpoints/`: epoch and step checkpoints.
153
- - `export_manifest.json`, `train_summary.json`, and `train.log`: model metadata and training trace.
154
-
155
- ## Current Strengths
156
-
157
- - Strong direct user-fact recall in single-session settings.
158
- - Strong assistant-detail recall.
159
- - Competitive knowledge-update behavior for changing facts.
160
- - Working temporal and preference retrieval layers with clear room for further specialization.
161
-
162
- ## Active Improvement Areas
163
-
164
- - Multi-session aggregation and unit coverage.
165
- - Deeper time-graph reasoning.
166
- - Preference-profile abstraction and cross-session tunneling.
167
- - Query-graph to memory-graph matching for complex questions.
168
-
169
- ## Intended Use
170
-
171
- This repository is a public-facing evidence package for TMCRA's long-memory runtime work. It is suitable for:
172
-
173
- - Benchmark review.
174
- - Model and result inspection.
175
- - Reproducing the frozen baseline.
176
- - Demonstrating how TMCRA can be packaged as an external memory runtime for agents.
177
-
178
- ## Benchmark Result
179
-
180
- This package includes a full LongMemEval S500 run.
181
-
182
- - Benchmark: LongMemEval S set, 500 samples
183
- - Evaluation: official-compatible LongMemEval judge prompt
184
- - Judge model: `gpt-4o`, resolved as `gpt-4o-2024-08-06`
185
- - Writer layer used in this run: DeepSeek v4 Flash
186
- - Answer layer used in this run: GPT5.4-compatible API
187
- - Overall accuracy: `310 / 500 = 62.00%`
188
-
189
- ## Results by Task Type
190
-
191
- | task type | accuracy | count |
192
- | --- | ---: | ---: |
193
- | single-session-user | 81.43% | 70 |
194
- | single-session-assistant | 78.57% | 56 |
195
- | knowledge-update | 70.51% | 78 |
196
- | temporal-reasoning | 63.16% | 133 |
197
- | single-session-preference | 56.67% | 30 |
198
- | multi-session | 39.85% | 133 |
199
-
200
- The benchmark outputs are available in:
201
-
202
- ```text
203
- results/predictions.jsonl
204
- results/judge_gpt4o_alias_vectorengine.jsonl
205
- results/judge_gpt4o_alias_vectorengine.jsonl.summary.json
206
- results/lme_s500_frozen_baseline38_full10_20260525_results.tar.gz
207
- ```
 
1
+ ---
2
+ license: apache-2.0
3
+ library_name: pytorch
4
+ language:
5
+ - en
6
+ - zh
7
+ tags:
8
+ - agent-memory
9
+ - long-term-memory
10
+ - memory-graph
11
+ - information-retrieval
12
+ - longmemeval
13
+ - pytorch
14
+ ---
15
+
16
+ # TMCRA — Agent Memory Engine
17
 
18
  <p align="center">
19
  <img src="assets/tmcra-logo.png" alt="TMCRA logo" width="420">
20
  </p>
21
 
22
+ <p align="center">
23
+ <a href="README.zh-CN.md">简体中文</a> ·
24
+ <a href="https://github.com/reshuibuduo/TMCRA-Agent-Memory">GitHub</a> ·
25
+ <a href="https://github.com/reshuibuduo/TMCRA-Agent-Memory/releases/tag/v2026.07.18">Release v2026.07.18</a>
26
+ </p>
 
 
 
 
 
 
 
 
27
 
28
+ TMCRA is a memory engine for long-running agents. It converts multi-turn and multi-session conversations into scope-isolated, source-traceable memory, then recalls compact evidence after the user's next question arrives.
29
 
30
+ This Hugging Face repository contains the Apache-2.0 release of the TMCRA algorithm, its LongMemEval reproduction pipeline, and the inference-only graph scoring weights used by the released benchmark chain. Hosted APIs, accounts, billing, and production control-plane services are not included.
31
 
32
+ ## LongMemEval S500
 
 
 
 
 
 
33
 
34
+ TMCRA achieved **411 / 500 = 82.2%**.
35
 
36
+ | Task category | Correct / Total | Accuracy |
37
+ | --- | ---: | ---: |
38
+ | Knowledge Update | 71 / 78 | 91.0% |
39
+ | Multi-session | 90 / 133 | 67.7% |
40
+ | Single-session Assistant | 55 / 56 | 98.2% |
41
+ | Single-session Preference | 27 / 30 | 90.0% |
42
+ | Single-session User | 67 / 70 | 95.7% |
43
+ | Temporal Reasoning | 101 / 133 | 75.9% |
44
+ | **Overall** | **411 / 500** | **82.2%** |
45
 
46
+ The machine-readable scorecard is [`results/latest_benchmark.json`](results/latest_benchmark.json). The maintained reproduction entrypoint is [`benchmarks/longmemeval/`](benchmarks/longmemeval/README.md). The benchmark answer layer is fixed to GPT-5.4; fresh runs that call external model services can vary slightly.
 
 
 
47
 
48
+ ## Architecture
 
 
 
49
 
50
+ TMCRA separates memory construction from downstream answer generation.
 
 
51
 
52
+ ```text
53
+ Conversation history
54
+ │
55
+ ▼
56
+ Writer ─────────────── subject / entity attribution
57
+ │
58
+ ├── Source layer ─── immutable, traceable evidence
59
+ ├── Fast layer ───── atomic current-state memory
60
+ └── Slow layer ───── durable semantic graph
61
+ │
62
+ └────────── scope-bound indexes
63
+
64
+ New question
65
+ │
66
+ ▼
67
+ Recall Planner → layered retrieval → node/path scoring → reranker
68
+ │
69
+ ▼
70
+ Source-bound evidence packet → downstream agent
71
  ```
72
 
73
+ The TMCRA core ends at the evidence packet. The answer model and benchmark judge are harness components, not part of the online memory engine. User statements and assistant progress remain separately attributed, while both can be recalled when relevant.
74
 
75
+ ## Released weights
76
 
77
+ All three checkpoints contain only their public inference contracts and model state. Training state, machine paths, and private run metadata are excluded. Load them with `torch.load(..., weights_only=True)` through the repository runtime.
78
 
79
+ | File | Purpose | Size | SHA-256 |
80
+ | --- | --- | ---: | --- |
81
+ | `node_scorer.pt` | Graph node scoring | 207,356,243 bytes | `d2318aafb07f9a15a1d95e6eac1b3e09afa07c570dd9d15d3f75544c9950f201` |
82
+ | `path_scorer.pt` | Graph path scoring | 207,356,243 bytes | `ecdc2bca51b7646a8a56c7db9cea4aeec581e64c3b61e6e06f7e9fa2a93b3bf5` |
83
+ | `tmcra_v3_reranker.pt` | Learned retrieval reranking | 4,232,213 bytes | `09a285b484ca857b24b53ad19d5998302f16d224ca9e4fd73a4eb5a52f022942` |
84
 
85
+ The files and their machine-readable contract are under [`models/tmcra_v4_longmemeval_s500_20260715/`](models/tmcra_v4_longmemeval_s500_20260715/README.md).
86
 
87
+ ## Download
 
 
 
 
88
 
89
+ Install the current Hugging Face CLI and download only the released inference package:
 
 
 
 
 
 
 
 
90
 
91
  ```bash
92
+ python -m pip install -U huggingface_hub
 
 
 
 
 
 
93
 
94
+ hf download 2009YU/TMCRA-Agent-Memory \
95
+ --include "models/tmcra_v4_longmemeval_s500_20260715/*" \
96
+ --local-dir ./TMCRA-Agent-Memory
 
 
 
 
 
97
  ```
98
 
99
+ For source, benchmark setup, environment examples, asset checksums, and stage validation, use the [GitHub release](https://github.com/reshuibuduo/TMCRA-Agent-Memory/releases/tag/v2026.07.18) or [`benchmarks/longmemeval/README.md`](benchmarks/longmemeval/README.md).
 
 
100
 
101
+ ## Version history
102
 
103
+ - Current release: [`v2026.07.18`](https://huggingface.co/2009YU/TMCRA-Agent-Memory/tree/v2026.07.18)
104
+ - Preserved pre-release snapshot: [`legacy-2026.05.25`](https://huggingface.co/2009YU/TMCRA-Agent-Memory/tree/legacy-2026.05.25)
 
 
 
105
 
106
+ ## Developers
107
 
108
+ - **Yu Haoxin** — creator, lead developer, and TMCRA algorithm engineering.
109
+ - **OpenAI Codex** — development and reproducibility engineering assistant.
110
 
111
+ See [`AUTHORS.md`](AUTHORS.md) and [`CITATION.cff`](CITATION.cff) for attribution and citation details.
112
 
113
+ ## License
 
114
 
115
+ TMCRA is released under the [Apache License 2.0](LICENSE). Third-party datasets, models, and components retain their own licenses.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
README.zh-CN.md CHANGED
@@ -1,207 +1,100 @@
1
- # TMCRA 长记忆运行时
2
 
3
  <p align="center">
4
  <img src="assets/tmcra-logo.png" alt="TMCRA logo" width="420">
5
  </p>
6
 
7
- [English version](README.md)
8
-
9
- TMCRA 是给 Agent 系统提供长期记忆能力的运行时层。它帮助大语言模型在长对话历史中检索、连接和推理相关记忆,而不需要在每一轮都暴露完整上下文。
10
-
11
- 本仓库包含一个冻结的 TMCRA 基线包,包括模型权重、运行时代码快照、训练元数据和 LongMemEval S500 测评结果。
12
-
13
- 许可证:MIT。
14
-
15
- ## TMCRA 的作用
16
-
17
- TMCRA 在 Agent 应用和回答模型之间增加一层专门的记忆运行时。
18
-
19
- 在写入阶段,TMCRA 会把对话转成记忆节点、事件单元、profile 信号和图路径。这让系统不仅能保存孤立事实,也能保存不同轮次、不同会话之间的事实关系。
20
 
21
- 在召回阶段,TMCRA 会对图节点和路径进行打分,选择紧凑的证据,并只把最相关的记忆上下文注入给回答模型。回答模型仍然负责自然语言推理和最终表达,TMCRA 负责长记忆组织、召回和证据呈现。
22
 
23
- 当前运行时重点覆盖:
24
 
25
- - 用户事实记忆
26
- - 助手回答细节记忆
27
- - profile 与偏好记忆
28
- - 时间记忆
29
- - 跨会话图隧穿
30
- - 学习式节点/路径打分
31
- - 面向下游 LLM 的紧凑证据选择
32
 
33
- ## 算法结构图
34
 
35
- ```mermaid
36
- flowchart LR
37
- A["对话轮次"] --> B["写入层"]
38
- B --> C["记忆节点"]
39
- B --> D["事件单元"]
40
- B --> E["Profile 信号"]
41
- B --> F["时间信号"]
 
 
42
 
43
- C --> G["图记忆库"]
44
- D --> G
45
- E --> G
46
- F --> G
47
 
48
- G --> H["学习式节点 Scorer"]
49
- G --> I["学习式路径 Scorer"]
50
- H --> J["证据选择"]
51
- I --> J
52
 
53
- J --> K["紧凑记忆上下文"]
54
- K --> L["回答模型"]
55
- L --> M["Agent 回复"]
56
 
57
- N["用户问题"] --> H
58
- N --> I
59
- N --> J
60
- N --> L
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
61
  ```
62
 
63
- 写入层从对话中生成记忆单元。图记忆库保存事实、事件、profile 信号、时间信号和跨会话连接。学习式节点/路径 scorer 会根据当前问题选择相关证据,回答模型使用这些紧凑证据生成最终回复。
64
 
65
- ## 为什么需要 TMCRA
66
 
67
- 长期运行的 Agent 不能只依赖简单向量召回。它需要保存用户事实、偏好、时间线变化、跨会话事件,以及多步证据链。
68
 
69
- TMCRA 将记忆组织成图节点和学习得到的召回路径,再把压缩后的证据提供给回答模型。目标是让外部 Agent 可以通过运行时/API 层使用长期记忆,同时保持记忆算法和模型权重可以独立部署。
 
 
 
 
70
 
71
- ## 如何使用
72
 
73
- 推理或运行时使用时,加载下面目录中的图 scorer 权重:
74
-
75
- ```text
76
- models/action_frame_tunnel_graph548_tunnel_fusion_train_20260524_042557/
77
- ```
78
 
79
- 主要运行时文件是:
80
-
81
- ```text
82
- node_scorer.pt
83
- path_scorer.pt
84
- export_manifest.json
85
- ```
86
-
87
- 典型运行配置如下:
88
 
89
  ```bash
90
- export TMCRA_NODE_MODEL_PATH="models/action_frame_tunnel_graph548_tunnel_fusion_train_20260524_042557/node_scorer.pt"
91
- export TMCRA_PATH_MODEL_PATH="models/action_frame_tunnel_graph548_tunnel_fusion_train_20260524_042557/path_scorer.pt"
92
- export TMCRA_RETRIEVAL_MODE="hybrid_node_scored"
93
- export TMCRA_REQUIRE_LEARNED_SCORER="1"
94
- ```
95
-
96
- 测评入口代码快照:
97
 
98
- ```text
99
- code/run_lme_s10_native_tmcra.py
100
- ```
101
-
102
- 核心适配器代码快照:
103
-
104
- ```text
105
- code/memory_adapters.py
106
  ```
107
 
108
- 部署时,将两个 scorer 文件加载到 TMCRA adapter 中,并让 Agent 的记忆中间件调用 TMCRA 召回 API。回答模型可以是任意 OpenAI 兼容接口或本地 LLM;TMCRA 提供选中的记忆证据,回答模型生成最终回复。
109
-
110
- ## 依赖环境
111
 
112
- 当前代码快照基于 Python。建议运行环境包括:
113
 
114
- - Python 3.10 或更高版本
115
- - PyTorch,推荐使用 CUDA 做模型推理
116
- - NumPy 及常见 Python 数据处理库
117
- - 用于回答层和写入层的 OpenAI 兼容接口或本地 LLM endpoint
118
- - 从 GitHub 拉取完整模型包时建议支持 Git LFS
119
 
120
- benchmark 脚本使用 LongMemEval 格式输入数据,并输出 JSONL 格式的预测和 judge 结果。实际运行时部署可以直接使用同一套模型文件,不需要运行 benchmark harness。
121
 
122
- ## 可开启模块
 
123
 
124
- TMCRA 当前也保留了用于召回和规划实验的可选接入口。部署或评估时,如果目标场景需要,可以按需开启这些模块。
125
 
126
- - **Embedder 接入口**:可选的语义 embedding 通道,可以和图记忆 scorer 并行工作。它用于在图证据选择前或选择过程中提供额外的稠密语义召回,但不替代学习式图节点/路径 scorer。
127
- - **LLM planner 接入口**:可选的规划 hook,可以调用外部 LLM 来整理证据、扩展 query intent,或在最终回答前生成 answer plan。它适合实验和更高成本的部署场景;默认基线仍保持核心图 scorer 链路可以独立测量。
128
 
129
- 这些接口是可集成的扩展点,不是冻结 S500 基线的必需依赖。下游部署可以根据成本和效果选择轻量的 scorer-only 路径,或开启 embedder/planner 辅助的增强路径。
130
-
131
- ## 包含内容
132
-
133
- - `code/`:该基线使用的运行时和测评代码快照。
134
- - `models/action_frame_tunnel_graph548_tunnel_fusion_train_20260524_042557/`:完整训练后的图模型输出目录。
135
- - `results/`:预测结果、judge 输出、汇总指标和压缩后的运行结果。
136
- - `docs/`:基线记录和结果说明。
137
-
138
- 补充文档:
139
-
140
- - `docs/BASELINE_S500_20260525.md`:benchmark 记录和子任务指标。
141
- - `docs/TRAINING.md`:模型训练方向和随包发布的训练产物说明。
142
- - `docs/OPTIONAL_MODULES_AND_PARALLEL.zh-CN.md`���embedder/planner 可选接入口和并行评估推进方案。
143
- - `docs/OPTIONAL_MODULES_AND_PARALLEL.md`:上述文档的英文版本。
144
-
145
- ## 模型包
146
-
147
- 随包发布的模型目录保留了图 scorer 栈的完整训练输出:
148
-
149
- - `node_scorer.pt` 和 `path_scorer.pt`:运行时图打分权重。
150
- - `node_scorer_best.pt` 和 `path_scorer_best.pt`:验证表现最好的 checkpoint 别名。
151
- - `node_scorer_last.pt` 和 `path_scorer_last.pt`:训练结束时的 checkpoint 别名。
152
- - `checkpoints/`:按 epoch 和 step 保存的中间 checkpoint。
153
- - `export_manifest.json`、`train_summary.json` 和 `train.log`:模型元数据和训练轨迹。
154
-
155
- ## 当前优势
156
-
157
- - 单会话直接用户事实召回较强。
158
- - 助手细节召回较强。
159
- - 对变化事实的 knowledge-update 任务有较稳定表现。
160
- - 已具备可工作的时间记忆和偏好召回层,后续仍有专项优化空间。
161
-
162
- ## 当前改进方向
163
-
164
- - multi-session 聚合和 unit coverage。
165
- - 更深的时间图推理。
166
- - 偏好 profile 抽象和跨会话隧穿。
167
- - query graph 与 memory graph 的复杂问题匹配。
168
-
169
- ## 适用场景
170
-
171
- 本仓库是 TMCRA 长记忆运行时工作的对外展示包,适合用于:
172
-
173
- - benchmark 审阅
174
- - 模型和结果检查
175
- - 复现冻结基线
176
- - 展示 TMCRA 如何作为外部 Agent 的记忆运行时进行封装
177
-
178
- ## Benchmark 结果
179
-
180
- 本包包含一次完整 LongMemEval S500 测试。
181
-
182
- - Benchmark:LongMemEval S set,500 条样本
183
- - Evaluation:official-compatible LongMemEval judge prompt
184
- - Judge 模型:`gpt-4o`,解析为 `gpt-4o-2024-08-06`
185
- - 本轮使用的写入层:DeepSeek v4 Flash
186
- - 本轮使用的回答层:GPT5.4-compatible API
187
- - 总准确率:`310 / 500 = 62.00%`
188
-
189
- ## 各任务类型结果
190
-
191
- | 任务类型 | 准确率 | 数量 |
192
- | --- | ---: | ---: |
193
- | single-session-user | 81.43% | 70 |
194
- | single-session-assistant | 78.57% | 56 |
195
- | knowledge-update | 70.51% | 78 |
196
- | temporal-reasoning | 63.16% | 133 |
197
- | single-session-preference | 56.67% | 30 |
198
- | multi-session | 39.85% | 133 |
199
-
200
- benchmark 输出文件位于:
201
-
202
- ```text
203
- results/predictions.jsonl
204
- results/judge_gpt4o_alias_vectorengine.jsonl
205
- results/judge_gpt4o_alias_vectorengine.jsonl.summary.json
206
- results/lme_s500_frozen_baseline38_full10_20260525_results.tar.gz
207
- ```
 
1
+ # TMCRA — 面向 Agent 的记忆引擎
2
 
3
  <p align="center">
4
  <img src="assets/tmcra-logo.png" alt="TMCRA logo" width="420">
5
  </p>
6
 
7
+ <p align="center">
8
+ <a href="README.md">English</a> ·
9
+ <a href="https://github.com/reshuibuduo/TMCRA-Agent-Memory">GitHub</a> ·
10
+ <a href="https://github.com/reshuibuduo/TMCRA-Agent-Memory/releases/tag/v2026.07.18">v2026.07.18 发布页</a>
11
+ </p>
 
 
 
 
 
 
 
 
12
 
13
+ TMCRA 是面向长期运行 Agent 的记忆引擎。它把多轮、跨会话的对话整理成作用域隔离、来源可追溯的记忆,并在用户提出下一轮问题后,召回与当前问题相关的紧凑证据。
14
 
15
+ 这个 Hugging Face 仓库包含采用 Apache-2.0 许可发布的 TMCRA 算法、LongMemEval 复现链路,以及本次 benchmark 使用的推理权重。线上 API、账号、计费和生产控制台不在本仓库的开源范围内。
16
 
17
+ ## LongMemEval S500 成绩
 
 
 
 
 
 
18
 
19
+ TMCRA 的成绩为 **411 / 500 = 82.2%**。
20
 
21
+ | 任务类型 | 正确数 / 总数 | 准确率 |
22
+ | --- | ---: | ---: |
23
+ | 信息更新(Knowledge Update) | 71 / 78 | 91.0% |
24
+ | 跨会话整合(Multi-session) | 90 / 133 | 67.7% |
25
+ | 单会话助手信息(Single-session Assistant) | 55 / 56 | 98.2% |
26
+ | 单会话偏好(Single-session Preference) | 27 / 30 | 90.0% |
27
+ | 单会话用户信息(Single-session User) | 67 / 70 | 95.7% |
28
+ | 时间推理(Temporal Reasoning) | 101 / 133 | 75.9% |
29
+ | **总成绩** | **411 / 500** | **82.2%** |
30
 
31
+ 机器可读的成绩文件位于 [`results/latest_benchmark.json`](results/latest_benchmark.json),完整复现入口位于 [`benchmarks/longmemeval/`](benchmarks/longmemeval/README.zh-CN.md)。复现链路固定使用 GPT-5.4 作为回答层;重新运行时如果调用外部模型服务,结果可能有小幅波动。
 
 
 
32
 
33
+ ## 架构
 
 
 
34
 
35
+ TMCRA 把记忆构建与最终回答生成明确分开。
 
 
36
 
37
+ ```text
38
+ 历史对话
39
+ │
40
+ ▼
41
+ Writer ─────────────── 说话者与实体归因
42
+ │
43
+ ├── Source 层 ────── 不可变、可追溯的原始证据
44
+ ├── Fast 层 ──────── 反映当前状态的原子记忆
45
+ └── Slow 层 ──────── 可跨会话保留的语义图
46
+ │
47
+ └──────── 作用域内索引
48
+
49
+ 用户的新问题
50
+ │
51
+ ▼
52
+ 召回规划 → 分层检索 → 节点/路径打分 → 重排
53
+ │
54
+ ▼
55
+ 带来源绑定的证据包 → 下游 Agent
56
  ```
57
 
58
+ TMCRA 的核心边界止于证据包;回答模型和 benchmark Judge 属于评测链路,不是在线记忆引擎本身。用户陈述和 Agent 完成的工作进度会分别保留各自的主体归属,同时都能在相关时被召回。
59
 
60
+ ## 本次发布的权重
61
 
62
+ 三份权重只保留公开推理所需的结构定义和模型状态,不包含训练状态、训练机器路径或私有运行信息。请通过仓库运行时使用 `torch.load(..., weights_only=True)` 安全加载。
63
 
64
+ | 文件 | 用途 | 大小 | SHA-256 |
65
+ | --- | --- | ---: | --- |
66
+ | `node_scorer.pt` | 图节点打分 | 207,356,243 bytes | `d2318aafb07f9a15a1d95e6eac1b3e09afa07c570dd9d15d3f75544c9950f201` |
67
+ | `path_scorer.pt` | 图路径打分 | 207,356,243 bytes | `ecdc2bca51b7646a8a56c7db9cea4aeec581e64c3b61e6e06f7e9fa2a93b3bf5` |
68
+ | `tmcra_v3_reranker.pt` | 召回重排 | 4,232,213 bytes | `09a285b484ca857b24b53ad19d5998302f16d224ca9e4fd73a4eb5a52f022942` |
69
 
70
+ 权重和机器可读的发布清单位于 [`models/tmcra_v4_longmemeval_s500_20260715/`](models/tmcra_v4_longmemeval_s500_20260715/README.md)。
71
 
72
+ ## 下载
 
 
 
 
73
 
74
+ 安装最新版 Hugging Face CLI,然后只下载本次发布的推理包:
 
 
 
 
 
 
 
 
75
 
76
  ```bash
77
+ python -m pip install -U huggingface_hub
 
 
 
 
 
 
78
 
79
+ hf download 2009YU/TMCRA-Agent-Memory \
80
+ --include "models/tmcra_v4_longmemeval_s500_20260715/*" \
81
+ --local-dir ./TMCRA-Agent-Memory
 
 
 
 
 
82
  ```
83
 
84
+ 源代码、benchmark 配置、环境示例、资源校验和分阶段检查说明可查看 [GitHub 发布页](https://github.com/reshuibuduo/TMCRA-Agent-Memory/releases/tag/v2026.07.18) 或 [`benchmarks/longmemeval/README.zh-CN.md`](benchmarks/longmemeval/README.zh-CN.md)。
 
 
85
 
86
+ ## 版本记录
87
 
88
+ - 当前版本:[`v2026.07.18`](https://huggingface.co/2009YU/TMCRA-Agent-Memory/tree/v2026.07.18)
89
+ - 更新前保留版本:[`legacy-2026.05.25`](https://huggingface.co/2009YU/TMCRA-Agent-Memory/tree/legacy-2026.05.25)
 
 
 
90
 
91
+ ## 开发者
92
 
93
+ - **Yu Haoxin**:项目创建者、主要开发者与 TMCRA 算法工程。
94
+ - **OpenAI Codex**:开发与可复现性工程协作。
95
 
96
+ 署名与引用信息见 [`AUTHORS.md`](AUTHORS.md) 和 [`CITATION.cff`](CITATION.cff)。
97
 
98
+ ## 许可证
 
99
 
100
+ TMCRA 采用 [Apache License 2.0](LICENSE) 开源;第三方数据集、模型和组件仍遵循各自的许可证。
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
SECURITY.md ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Security policy
2
+
3
+ Do not open a public issue containing credentials, private memory data, or a
4
+ working exploit. Use GitHub's private vulnerability reporting for this
5
+ repository when available.
6
+
7
+ TMCRA checkpoints must be loaded with `torch.load(..., weights_only=True)` and
8
+ verified against the published SHA-256 manifest. The maintainers do not
9
+ support disabling these checks for untrusted checkpoints.
10
+
11
+ Provider API keys belong only in local `*.env` files. The example files contain
12
+ placeholders and real environment files are excluded by `.gitignore`.
benchmarks/longmemeval/README.md ADDED
@@ -0,0 +1,216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TMCRA Agent Memory System — LongMemEval Benchmark
2
+
3
+ [简体中文](README.zh-CN.md)
4
+
5
+ TMCRA is a scope-isolated, layered memory engine for agents. It keeps durable evidence outside the answer model and returns a compact, traceable memory context for each query. This directory contains the LongMemEval pipeline used to build memory, retrieve across multiple memory layers, compile evidence, generate answers, and evaluate the predictions.
6
+
7
+ ## Benchmark result
8
+
9
+ TMCRA achieves **82.2%** on LongMemEval.
10
+
11
+ | Task category | Correct / Total | Accuracy |
12
+ | --- | ---: | ---: |
13
+ | Knowledge Update | 71 / 78 | 91.0% |
14
+ | Multi-session | 90 / 133 | 67.7% |
15
+ | Single-session Assistant | 55 / 56 | 98.2% |
16
+ | Single-session Preference | 27 / 30 | 90.0% |
17
+ | Single-session User | 67 / 70 | 95.7% |
18
+ | Temporal Reasoning | 101 / 133 | 75.9% |
19
+ | **Overall** | **411 / 500** | **82.2%** |
20
+
21
+ The machine-readable aggregate result is available in [`results/benchmark.json`](results/benchmark.json).
22
+
23
+ ## Architecture
24
+
25
+ TMCRA separates immutable source evidence from two derived memory layers. The Writer stores conversation messages in Source and extracts normalized, current assertions into Fast. Subject Attribution prevents quoted or third-party facts from being promoted under the wrong identity. Eligible records are reorganized into the durable Slow semantic graph, while every derived memory remains traceable to Source.
26
+
27
+ At query time, the Recall Planner resolves the question and assigns layer roles. Source, Fast, and Slow candidates are retrieved, graph-ranked, reranked, and packed into a bounded evidence set. The Evidence Compiler converts this set into a structured, provenance-bound packet. TMCRA's online memory boundary ends at prompt-ready evidence; answer generation and the official Judge belong to the benchmark harness.
28
+
29
+ ```mermaid
30
+ flowchart LR
31
+ subgraph BENCH["Benchmark Adapter"]
32
+ D["LongMemEval Dataset"] --> H["Data Adapter / Harness"]
33
+ Q["Query"]
34
+ G["Gold Answer"]
35
+ D --> G
36
+ end
37
+
38
+ subgraph CORE["TMCRA Memory and Retrieval Core"]
39
+ W["Writer"]
40
+ S["Source Layer<br/>immutable records"]
41
+ W --> F["Fast Layer<br/>atomic current memory"]
42
+ W --> S
43
+ F --> A["Subject Attribution<br/>promotion gate"]
44
+ A --> L["Slow Layer<br/>durable semantic graph"]
45
+ S --> I["Index<br/>searchable snapshot"]
46
+ F --> I
47
+ L --> I
48
+ P["Recall Planner<br/>query and layer roles"]
49
+ P --> R["Layered Retrieval<br/>Source + Fast + Slow"]
50
+ I --> R
51
+ R --> K["Graph / Reranker<br/>Top-K evidence packing"]
52
+ K --> C["Evidence Compiler<br/>traceable evidence packet"]
53
+ end
54
+
55
+ subgraph EVAL["Benchmark Evaluation Only"]
56
+ C --> M["Answer Model"]
57
+ M --> J["Official Judge"]
58
+ J --> O["Overall and task-level scores"]
59
+ end
60
+
61
+ H -->|"conversation history"| W
62
+ H -->|"question only"| Q
63
+ Q --> P
64
+ G -->|"reference used only here"| J
65
+ ```
66
+
67
+ ### Component responsibilities
68
+
69
+ - **Source** retains the immutable conversation record and is the final evidence authority.
70
+ - **Fast** stores source-bound atomic facts, relations, updates, and current state.
71
+ - **Subject Attribution** stops quoted or third-party information from being assigned to the wrong subject.
72
+ - **Slow** organizes eligible cross-session information into a durable semantic graph; it supplements rather than replaces Source.
73
+ - **Recall Planner** interprets the question and memory-layer roles without reading the gold answer.
74
+ - **Layered Retrieval** searches Source, Fast, and Slow, then performs graph ranking, reranking, deduplication, temporal handling, and fixed Top-K packing.
75
+ - **Evidence Compiler** produces a structured evidence packet whose claims remain bound to source IDs.
76
+ - **Answer Model and Judge** are benchmark-only components and are not part of the TMCRA online recall service.
77
+
78
+ Gold answers are isolated from Writer, Planner, Retrieval, Evidence Compiler, and Answer Model. They are only read by the final Judge.
79
+
80
+ ## Viewing and reproducing the benchmark
81
+
82
+ - **View the published result:** [`results/benchmark.json`](results/benchmark.json) contains the public overall and task-level summary shown above.
83
+ - **Rerun the pipeline:** the commands below execute memory construction, layered retrieval, evidence compilation, answer generation, and evaluation. The published 82.2% scorecard is a frozen v5 result; this release uses the security-hardened v6 answer protocol, and fresh external model calls can produce a different score.
84
+
85
+ ### Tested environment
86
+
87
+ - Linux
88
+ - Python 3.12.3
89
+ - NVIDIA GPU with 24 GiB VRAM or more
90
+ - CUDA 12.8
91
+ - PyTorch 2.11.0 + cu128
92
+ - Transformers 5.10.2
93
+
94
+ The full S500 run also needs sufficient system memory and disk space for Writer journals, graphs, indexes, checkpoints, and intermediate evidence artifacts.
95
+
96
+ ### 1. Clone and install
97
+
98
+ ```bash
99
+ git lfs install
100
+ GIT_LFS_SKIP_SMUDGE=1 git clone https://github.com/reshuibuduo/TMCRA-Agent-Memory.git
101
+ cd TMCRA-Agent-Memory/benchmarks/longmemeval
102
+ git lfs pull --include="models/tmcra_v4_longmemeval_s500_20260715/*.pt"
103
+
104
+ python3.12 -m venv .venv
105
+ source .venv/bin/activate
106
+ python -m pip install --upgrade pip
107
+ python -m pip install -r requirements.txt
108
+ python -m pip install -e .
109
+ ```
110
+
111
+ `GIT_LFS_SKIP_SMUDGE=1` prevents the clone from downloading unrelated historical LFS objects; the next command pulls only the three checkpoints required by this benchmark.
112
+
113
+ ### 2. Verify checkpoints and download external assets
114
+
115
+ The three TMCRA checkpoints are tracked in the repository with Git LFS under `models/tmcra_v4_longmemeval_s500_20260715/`. The asset script verifies their byte sizes and SHA-256 hashes in place; it does not create a second copy. It then downloads the official cleaned LongMemEval-S dataset and the revision-pinned BGE models:
116
+
117
+ ```bash
118
+ python scripts/fetch_assets.py --manifest configs/assets.lock.json --kind all
119
+ ```
120
+
121
+ The dataset comes from the [official LongMemEval release](https://github.com/xiaowu0162/LongMemEval) and [cleaned dataset repository](https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned). BGE-M3 and the reranker are downloaded from the official [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3) and [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) model pages.
122
+
123
+ File assets are verified against `configs/assets.lock.json`. If the checkout contains only LFS pointer files, the script stops with the exact `git lfs pull` command required to fetch the weights.
124
+
125
+ ### 3. Configure model providers
126
+
127
+ ```bash
128
+ cp configs/writer.env.example configs/writer.env
129
+ cp configs/answer.env.example configs/answer.env
130
+ cp configs/judge.env.example configs/judge.env
131
+ cp configs/run.env.example configs/run.env
132
+ ```
133
+
134
+ Edit the four generated files:
135
+
136
+ - `writer.env` configures one OpenAI-compatible DeepSeek endpoint and a comma-separated key pool used by the Writer, subject-attribution pass, and evidence planner. The endpoint must accept `deepseek-v4-flash` and `deepseek-v4-pro`, or map those names to compatible deployments.
137
+ - `answer.env` must expose the exact model identifier `gpt-5.4` through an OpenAI-compatible endpoint.
138
+ - `judge.env` must expose `gpt-4o-2024-08-06`. The public command passes the `gpt-4o` alias, and the bundled evaluator resolves that alias to the fixed version.
139
+ - `run.env` selects local paths, concurrency, and PyTorch devices.
140
+
141
+ Keep API keys local; the repository `.gitignore` excludes the real configuration files.
142
+
143
+ The default Judge wire format is OpenAI Chat Completions. Set `OPENAI_WIRE_API=responses` in `judge.env` only if the configured endpoint implements the Responses API.
144
+
145
+ ### 4. Run the offline smoke test
146
+
147
+ This test uses a synthetic fixture and does not call external model providers:
148
+
149
+ ```bash
150
+ bash scripts/reproduce_smoke.sh
151
+ ```
152
+
153
+ ### 5. Run LongMemEval-S500
154
+
155
+ ```bash
156
+ bash scripts/reproduce_s500.sh
157
+ ```
158
+
159
+ The script executes and validates the following stages:
160
+
161
+ ```text
162
+ prepare data and ordered QIDs
163
+ -> build Writer journals, Source/Fast/Slow memory, and indexes
164
+ -> retrieve Source + Fast + Slow evidence
165
+ -> compile operation-bound evidence packets
166
+ -> generate answers with GPT-5.4
167
+ -> evaluate with the official GPT-4o Judge
168
+ -> export overall and task-level scores
169
+ ```
170
+
171
+ Important outputs are written under `runs/longmemeval_s500/`:
172
+
173
+ ```text
174
+ BUILD_COMPLETE
175
+ retrieval_s500/evidence_windows.jsonl
176
+ retrieval_s500/retrieval_debug.jsonl
177
+ retrieval_s500_compiled/COMPILE_COMPLETE
178
+ retrieval_s500_compiled/evidence_windows.jsonl
179
+ answers_s500/answers.jsonl
180
+ official_judge_gpt4o_20240806.jsonl
181
+ scorecard.json
182
+ ```
183
+
184
+ Before the next stage starts, the runner requires completion markers or reports, exact row counts, and the same ordered QID set. Judge rows are also checked for boolean labels from `gpt-4o-2024-08-06`. Compatible incomplete retrieval, compilation, answer, and judge artifacts are resumed. Build recovery is fail-closed and may require the explicit review flags printed by the build runner; an incompatible answer protocol requires a new answer output directory.
185
+
186
+ ## Repository layout
187
+
188
+ ```text
189
+ configs/ provider examples and hash-pinned asset manifest
190
+ fixtures/ synthetic offline test data
191
+ results/ published aggregate score
192
+ scripts/ asset download, validation, smoke, and S500 runners
193
+ src/ TMCRA V4 benchmark core, LongMemEval harness, and minimal adapter
194
+ tests/ offline schema, gold-isolation, and score aggregation tests
195
+ ../../models/tmcra_v4_longmemeval_s500_20260715/
196
+ Git LFS checkpoints used by this benchmark
197
+ ```
198
+
199
+ Production HTTP APIs, account systems, billing, tenant authorization, deployment files, credentials, runtime databases, raw provider logs, and control-plane components are intentionally excluded.
200
+
201
+ ## Validation
202
+
203
+ ```bash
204
+ python -m compileall -q src scripts tests
205
+ python -m pytest
206
+ ```
207
+
208
+ Before publishing a release, also run a secret scan, dependency vulnerability scan, SBOM generation, and a manual review of staged files and release assets.
209
+
210
+ ## Citation
211
+
212
+ When reporting these results, identify the repository version and cite the [LongMemEval project](https://github.com/xiaowu0162/LongMemEval) and its paper.
213
+
214
+ ## License
215
+
216
+ TMCRA benchmark code is released under the [Apache License 2.0](../../LICENSE). LongMemEval and the downloaded BGE models retain their own licenses; see [`THIRD_PARTY_NOTICES.md`](THIRD_PARTY_NOTICES.md).
benchmarks/longmemeval/README.zh-CN.md ADDED
@@ -0,0 +1,216 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TMCRA 智能体记忆系统 — LongMemEval Benchmark
2
+
3
+ [English](README.md)
4
+
5
+ TMCRA 是面向智能体的作用域隔离、分层记忆引擎。它将可追溯的长期证据保存在回答模型之外,并在每次查询时返回紧凑、可核验的记忆上下文。本目录包含 TMCRA 的 LongMemEval benchmark 链路,覆盖记忆构建、多层召回、证据整理、回答生成与结果评测。
6
+
7
+ ## Benchmark 成绩
8
+
9
+ TMCRA 在 LongMemEval benchmark 上取得 **82.2%**。
10
+
11
+ | 任务类型 | 正确数 / 样本数 | 准确率 |
12
+ | --- | ---: | ---: |
13
+ | 信息更新(Knowledge Update) | 71 / 78 | 91.0% |
14
+ | 跨会话信息整合(Multi-session) | 90 / 133 | 67.7% |
15
+ | 单会话助手信息(Single-session Assistant) | 55 / 56 | 98.2% |
16
+ | 单会话偏好理解(Single-session Preference) | 27 / 30 | 90.0% |
17
+ | 单会话用户信息(Single-session User) | 67 / 70 | 95.7% |
18
+ | 时间推理(Temporal Reasoning) | 101 / 133 | 75.9% |
19
+ | **总成绩** | **411 / 500** | **82.2%** |
20
+
21
+ 机器可读的汇总成绩位于 [`results/benchmark.json`](results/benchmark.json)。
22
+
23
+ ## 架构
24
+
25
+ TMCRA 将不可变证据与两类派生记忆分离。Writer 把对话消息保存到 Source,同时将可复用、反映当前状态的事实提取并规范化到 Fast。Subject Attribution 用于阻止引用内容或第三方事实被错误归属;通过门控的记录会进一步组织为稳定的 Slow 语义图,而所有派生记忆始终能够追溯到 Source。
26
+
27
+ 查询时,Recall Planner 先解析问题并为各记忆层分配职责。系统随后从 Source、Fast、Slow 生成候选,经过图排序、重排、去重、时间处理和有界 Top-K 打包,再由 Evidence Compiler 生成带来源绑定的结构化证据包。TMCRA 线上记忆服务的边界止于可注入提示词的证据;回答生成和 Official Judge 属于 benchmark harness。
28
+
29
+ ```mermaid
30
+ flowchart LR
31
+ subgraph BENCH["Benchmark 适配层"]
32
+ D["LongMemEval 数据集"] --> H["数据适配器 / 测试框架"]
33
+ Q["查询问题"]
34
+ G["标准答案"]
35
+ D --> G
36
+ end
37
+
38
+ subgraph CORE["TMCRA 记忆与召回核心"]
39
+ W["记忆写入器(Writer)"]
40
+ S["原始证据层(Source)<br/>不可变记录"]
41
+ W --> F["快速记忆层(Fast)<br/>当前原子记忆"]
42
+ W --> S
43
+ F --> A["主体归属判断<br/>晋升门控"]
44
+ A --> L["慢速记忆层(Slow)<br/>持久语义图"]
45
+ S --> I["检索索引<br/>可搜索快照"]
46
+ F --> I
47
+ L --> I
48
+ P["召回规划器<br/>解析查询与各层职责"]
49
+ P --> R["分层召回<br/>Source + Fast + Slow"]
50
+ I --> R
51
+ R --> K["图排序 / 重排<br/>Top-K 证据打包"]
52
+ K --> C["证据编译器<br/>可追溯证据包"]
53
+ end
54
+
55
+ subgraph EVAL["仅用于 Benchmark 评测"]
56
+ C --> M["回答模型"]
57
+ M --> J["官方评测器(Judge)"]
58
+ J --> O["总成绩与分项成绩"]
59
+ end
60
+
61
+ H -->|"对话历史"| W
62
+ H -->|"仅问题"| Q
63
+ Q --> P
64
+ G -->|"只在此处作为评测参考"| J
65
+ ```
66
+
67
+ ### 核心组件职责
68
+
69
+ - **Source** 保存不可变的对话原文,是最终证据来源。
70
+ - **Fast** 保存与 Source 绑定的原子事实、关系、更新和当前状态。
71
+ - **Subject Attribution** 防止引用内容或第三方信息被归属到错误主体。
72
+ - **Slow** 把符合条件的跨会话信息组织为稳定的语义图;它补充 Source,但不能替代 Source。
73
+ - **Recall Planner** 解析问题和各层职责,但不读取 gold answer。
74
+ - **Layered Retrieval** 同时搜索 Source、Fast、Slow,并执行图排序、重排、去重、时间处理与固定 Top-K 打包。
75
+ - **Evidence Compiler** 输出带 Source ID 绑定的结构化证据包。
76
+ - **Answer Model 与 Judge** 只属于 benchmark,不属于 TMCRA 线上召回服务。
77
+
78
+ Gold answer 与 Writer、Planner、Retrieval、Evidence Compiler、Answer Model 完全隔离,只允许最终 Judge 读取。
79
+
80
+ ## 查看与复现 benchmark
81
+
82
+ - **查看公开成绩:** [`results/benchmark.json`](results/benchmark.json) 包含上表所示的总成绩与六类分项汇总。
83
+ - **重跑完整链路:** 下面的命令会重新执行记忆构建、分层召回、证据整理、回答生成和评测。公开的 82.2% 是冻结的 v5 成绩;本次发布使用经过安全加固的 v6 回答协议,重新调用外部模型时成绩可能发生波动。
84
+
85
+ ### 已验证环境
86
+
87
+ - Linux
88
+ - Python 3.12.3
89
+ - 24 GiB 或更高显存的 NVIDIA GPU
90
+ - CUDA 12.8
91
+ - PyTorch 2.11.0 + cu128
92
+ - Transformers 5.10.2
93
+
94
+ 完整 S500 运行还需要足够的内存与磁盘空间,用于 Writer 日志、图、索引、checkpoint 和中间证据产物。
95
+
96
+ ### 1. 克隆并安装
97
+
98
+ ```bash
99
+ git lfs install
100
+ GIT_LFS_SKIP_SMUDGE=1 git clone https://github.com/reshuibuduo/TMCRA-Agent-Memory.git
101
+ cd TMCRA-Agent-Memory/benchmarks/longmemeval
102
+ git lfs pull --include="models/tmcra_v4_longmemeval_s500_20260715/*.pt"
103
+
104
+ python3.12 -m venv .venv
105
+ source .venv/bin/activate
106
+ python -m pip install --upgrade pip
107
+ python -m pip install -r requirements.txt
108
+ python -m pip install -e .
109
+ ```
110
+
111
+ `GIT_LFS_SKIP_SMUDGE=1` 会阻止 clone 阶段下载无关的历史 LFS 对象;后续命令只拉取本 benchmark 所需的三个 checkpoint。
112
+
113
+ ### 2. 校验 checkpoint 并下载外部资产
114
+
115
+ 三个 TMCRA checkpoint 已通过 Git LFS 保存在仓库的 `models/tmcra_v4_longmemeval_s500_20260715/` 目录。资产脚本会直接校验这些文件的字节数与 SHA-256,不会重复复制;随后下载官方清洗版 LongMemEval-S 和固定 revision 的 BGE 模型:
116
+
117
+ ```bash
118
+ python scripts/fetch_assets.py --manifest configs/assets.lock.json --kind all
119
+ ```
120
+
121
+ 数据来自 [LongMemEval 官方仓库](https://github.com/xiaowu0162/LongMemEval)和[清洗版数据集](https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned)。BGE-M3 与 reranker 来自官方 [BAAI/bge-m3](https://huggingface.co/BAAI/bge-m3) 和 [BAAI/bge-reranker-v2-m3](https://huggingface.co/BAAI/bge-reranker-v2-m3) 页面。
122
+
123
+ 文件类资产在启用前都会按照 `configs/assets.lock.json` 校验。如果当前检出只有 LFS 指针文件,脚本会停止并给出获取权重所需的完整 `git lfs pull` 命令。
124
+
125
+ ### 3. 配置模型服务
126
+
127
+ ```bash
128
+ cp configs/writer.env.example configs/writer.env
129
+ cp configs/answer.env.example configs/answer.env
130
+ cp configs/judge.env.example configs/judge.env
131
+ cp configs/run.env.example configs/run.env
132
+ ```
133
+
134
+ 编辑生成的四个文件:
135
+
136
+ - `writer.env` 配置一个兼容 OpenAI API 格式的 DeepSeek 服务和逗号分隔的 key 池,供 Writer、主体归属判断和证据规划器共同使用。该服务需要接受 `deepseek-v4-flash` 与 `deepseek-v4-pro`,或者在网关中把这两个名称映射到兼容部署。
137
+ - `answer.env` 的兼容 OpenAI API 服务必须以 `gpt-5.4` 这一精确名称提供回答模型。
138
+ - `judge.env` 的服务必须提供 `gpt-4o-2024-08-06`。公开命令传入 `gpt-4o` 别名,内置评测器会将其固定解析为该版本。
139
+ - `run.env` 配置本地路径、并发数与 PyTorch 设备。
140
+
141
+ API key 只保留在本地,仓库的 `.gitignore` 已排除真实配置文件。
142
+
143
+ Judge 默认使用 OpenAI Chat Completions 请求格式。只有当服务端实现了 Responses API 时,才在 `judge.env` 中设置 `OPENAI_WIRE_API=responses`。
144
+
145
+ ### 4. 运行离线 smoke test
146
+
147
+ 该测试只使用合成 fixture,不会调用外部模型服务:
148
+
149
+ ```bash
150
+ bash scripts/reproduce_smoke.sh
151
+ ```
152
+
153
+ ### 5. 运行 LongMemEval-S500
154
+
155
+ ```bash
156
+ bash scripts/reproduce_s500.sh
157
+ ```
158
+
159
+ 脚本会顺序执行并校验:
160
+
161
+ ```text
162
+ 准备数据与有序 QID
163
+ -> 构建 Writer 日志、Source/Fast/Slow 记忆和索引
164
+ -> 召回 Source + Fast + Slow 证据
165
+ -> 整理 operation-bound evidence packet
166
+ -> 使用 GPT-5.4 生成回答
167
+ -> 使用官方 GPT-4o Judge 评测
168
+ -> 导出总成绩与分项成绩
169
+ ```
170
+
171
+ 关键产物位于 `runs/longmemeval_s500/`:
172
+
173
+ ```text
174
+ BUILD_COMPLETE
175
+ retrieval_s500/evidence_windows.jsonl
176
+ retrieval_s500/retrieval_debug.jsonl
177
+ retrieval_s500_compiled/COMPILE_COMPLETE
178
+ retrieval_s500_compiled/evidence_windows.jsonl
179
+ answers_s500/answers.jsonl
180
+ official_judge_gpt4o_20240806.jsonl
181
+ scorecard.json
182
+ ```
183
+
184
+ 进入下一阶段前,脚本会检查完成标记或报告、准确的产物行数,以及与输入一致的有序 QID。Judge 产物还必须包含由 `gpt-4o-2024-08-06` 生成的布尔标签。兼容的召回、证据整理、回答和 Judge 产物会从断点继续;构建恢复采用故障关闭策略,必要时会要求使用构建脚本打印出的显式审查参数。回答协议不兼容时必须使用新的回答输出目录。
185
+
186
+ ## 仓库结构
187
+
188
+ ```text
189
+ configs/ 配置示例与 hash 固定的资产清单
190
+ fixtures/ 合成离线测试数据
191
+ results/ 公开的汇总成绩
192
+ scripts/ 资产下载、校验、smoke 与 S500 运行脚本
193
+ src/ TMCRA V4 benchmark 核心、LongMemEval harness 与最小 adapter
194
+ tests/ 离线 schema、gold 隔离和成绩汇总测试
195
+ ../../models/tmcra_v4_longmemeval_s500_20260715/
196
+ 本 benchmark 使用的 Git LFS checkpoint
197
+ ```
198
+
199
+ 线上 HTTP API、账户系统、计费、租户授权、部署文件、凭据、运行数据库、供应商原始日志和控制面均不在本仓库范围内。
200
+
201
+ ## 校验
202
+
203
+ ```bash
204
+ python -m compileall -q src scripts tests
205
+ python -m pytest
206
+ ```
207
+
208
+ 正式发布前还需要执行 secret scan、依赖漏洞扫描、SBOM,并人工审查暂存文件和 Release 资产。
209
+
210
+ ## 引用
211
+
212
+ 报告本结果时,请注明所使用的仓库版本,并引用 [LongMemEval 项目](https://github.com/xiaowu0162/LongMemEval)及其论文。
213
+
214
+ ## 许可证
215
+
216
+ TMCRA benchmark 代码采用 [Apache License 2.0](../../LICENSE)。LongMemEval 与下载的 BGE 模型保留各自许可证,详见 [`THIRD_PARTY_NOTICES.md`](THIRD_PARTY_NOTICES.md)。
benchmarks/longmemeval/THIRD_PARTY_NOTICES.md ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Third-Party Notices
2
+
3
+ This benchmark runner downloads or adapts third-party material. Those components remain subject to their original licenses; the repository's Apache-2.0 license does not replace them.
4
+
5
+ ## LongMemEval
6
+
7
+ - Project: <https://github.com/xiaowu0162/LongMemEval>
8
+ - Cleaned dataset: <https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned>
9
+ - License source: <https://github.com/xiaowu0162/LongMemEval/blob/main/LICENSE>
10
+ - License: MIT
11
+
12
+ The bundled evaluation wrapper follows the public LongMemEval evaluation prompt and rule. The dataset is downloaded separately by `scripts/fetch_assets.py`.
13
+
14
+ The complete LongMemEval MIT license text follows:
15
+
16
+ > MIT License
17
+ >
18
+ > Copyright (c) 2024 Di Wu
19
+ >
20
+ > Permission is hereby granted, free of charge, to any person obtaining a copy
21
+ > of this software and associated documentation files (the "Software"), to deal
22
+ > in the Software without restriction, including without limitation the rights
23
+ > to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
24
+ > copies of the Software, and to permit persons to whom the Software is
25
+ > furnished to do so, subject to the following conditions:
26
+ >
27
+ > The above copyright notice and this permission notice shall be included in all
28
+ > copies or substantial portions of the Software.
29
+ >
30
+ > THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
31
+ > IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
32
+ > FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
33
+ > AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
34
+ > LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
35
+ > OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
36
+ > SOFTWARE.
37
+
38
+ ## BAAI/bge-m3
39
+
40
+ - Model page: <https://huggingface.co/BAAI/bge-m3>
41
+ - Locked revision: `5617a9f61b028005a4858fdac845db406aefb181`
42
+ - License declared by the model repository: MIT
43
+
44
+ The model files are downloaded from Hugging Face and are not relicensed by TMCRA. Review the model repository and its license before redistribution.
45
+
46
+ ## BAAI/bge-reranker-v2-m3
47
+
48
+ - Model page: <https://huggingface.co/BAAI/bge-reranker-v2-m3>
49
+ - Locked revision: `953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e`
50
+ - License declared by the model repository: Apache License 2.0
51
+
52
+ The model files are downloaded from Hugging Face and are not relicensed by TMCRA. Review the model repository and its license before redistribution.
benchmarks/longmemeval/configs/answer.env.example ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ # The answer runner requires this exact model identifier. The endpoint may be
2
+ # an OpenAI-compatible gateway, but it must expose gpt-5.4 under this name.
3
+ TMCRA_ANSWER_MODEL=gpt-5.4
4
+ TMCRA_ANSWER_BASE_URL=https://your-answer-endpoint.example/v1
5
+ TMCRA_ANSWER_API_KEY=replace_me
6
+ TMCRA_ANSWER_MAX_TOKENS=1536
benchmarks/longmemeval/configs/assets.lock.json ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "schema_version": "tmcra.benchmark-assets.1",
3
+ "files": [
4
+ {
5
+ "id": "longmemeval_s_cleaned",
6
+ "kind": "dataset",
7
+ "destination": "data/longmemeval_s_cleaned.json",
8
+ "url": "https://huggingface.co/datasets/xiaowu0162/longmemeval-cleaned/resolve/main/longmemeval_s_cleaned.json",
9
+ "bytes": 277383467,
10
+ "sha256": "d6f21ea9d60a0d56f34a05b609c79c88a451d2ae03597821ea3d5a9678c3a442"
11
+ },
12
+ {
13
+ "id": "tmcra_node_scorer",
14
+ "kind": "checkpoint",
15
+ "storage": "git_lfs",
16
+ "destination": "../../models/tmcra_v4_longmemeval_s500_20260715/node_scorer.pt",
17
+ "bytes": 207356243,
18
+ "sha256": "d2318aafb07f9a15a1d95e6eac1b3e09afa07c570dd9d15d3f75544c9950f201"
19
+ },
20
+ {
21
+ "id": "tmcra_path_scorer",
22
+ "kind": "checkpoint",
23
+ "storage": "git_lfs",
24
+ "destination": "../../models/tmcra_v4_longmemeval_s500_20260715/path_scorer.pt",
25
+ "bytes": 207356243,
26
+ "sha256": "ecdc2bca51b7646a8a56c7db9cea4aeec581e64c3b61e6e06f7e9fa2a93b3bf5"
27
+ },
28
+ {
29
+ "id": "tmcra_v3_reranker",
30
+ "kind": "checkpoint",
31
+ "storage": "git_lfs",
32
+ "destination": "../../models/tmcra_v4_longmemeval_s500_20260715/tmcra_v3_reranker.pt",
33
+ "bytes": 4232213,
34
+ "sha256": "09a285b484ca857b24b53ad19d5998302f16d224ca9e4fd73a4eb5a52f022942"
35
+ }
36
+ ],
37
+ "huggingface_models": [
38
+ {
39
+ "id": "bge_m3",
40
+ "repo_id": "BAAI/bge-m3",
41
+ "revision": "5617a9f61b028005a4858fdac845db406aefb181",
42
+ "destination": "models/BAAI/bge-m3"
43
+ },
44
+ {
45
+ "id": "bge_reranker_v2_m3",
46
+ "repo_id": "BAAI/bge-reranker-v2-m3",
47
+ "revision": "953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e",
48
+ "destination": "models/BAAI/bge-reranker-v2-m3"
49
+ }
50
+ ]
51
+ }
benchmarks/longmemeval/configs/judge.env.example ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ # OpenAI-compatible Judge endpoint. The public runner selects the gpt-4o alias;
2
+ # the evaluator resolves it to the fixed gpt-4o-2024-08-06 model identifier.
3
+ TMCRA_JUDGE_BASE_URL=https://your-judge-endpoint.example/v1
4
+ TMCRA_JUDGE_API_KEY=replace_me
5
+ TMCRA_OFFICIAL_JUDGE_WORKERS=8
6
+
7
+ # The default wire format is POST /chat/completions. Uncomment this only when
8
+ # the configured endpoint implements POST /responses instead.
9
+ # OPENAI_WIRE_API=responses
benchmarks/longmemeval/configs/run.env.example ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # reproduce_s500.sh changes into benchmarks/longmemeval before loading this
2
+ # file, so $PWD resolves to that directory.
3
+ DATA="$PWD/data/longmemeval_s_cleaned.json"
4
+ QIDS="$PWD/data/longmemeval_s500.qids.txt"
5
+ RUN="$PWD/runs/longmemeval_s500"
6
+ # TMCRA scorer weights are tracked by Git LFS in the repository-level model
7
+ # package. BGE snapshots are downloaded locally under this benchmark directory.
8
+ MODELS="$PWD/../../models/tmcra_v4_longmemeval_s500_20260715"
9
+ BGE_MODELS="$PWD/models/BAAI"
10
+
11
+ # Tune concurrency to provider rate limits and local memory. Device values are
12
+ # passed directly to PyTorch and therefore require a matching local runtime.
13
+ WRITER_CONCURRENCY=20
14
+ SLOW_CONCURRENCY=20
15
+ RETRIEVAL_DEVICE=cuda
16
+ GRAPH_DEVICE=cuda
17
+ ANSWER_WORKERS=12
18
+ JUDGE_WORKERS=8
benchmarks/longmemeval/configs/writer.env.example ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # One OpenAI-compatible DeepSeek endpoint is used by the benchmark Writer,
2
+ # subject-attribution pass, and evidence planner. The gateway must accept the
3
+ # model names deepseek-v4-flash and deepseek-v4-pro, or alias them accordingly.
4
+ TMCRA_WRITER_BASE_URL=https://your-writer-endpoint.example/v1
5
+ TMCRA_WRITER_TIMEOUT_SECONDS=180
6
+ TMCRA_WRITER_MAX_TOKENS=16384
7
+
8
+ # Comma-separated, non-empty, unique API keys. For a single key, keep the
9
+ # count at 1. If you add keys, update the count to the exact number of keys.
10
+ TMCRA_DEEPSEEK_WRITER_BASE_URL=https://your-planner-endpoint.example/v1
11
+ TMCRA_DEEPSEEK_WRITER_KEY_POOL=replace_me
12
+ TMCRA_DEEPSEEK_WRITER_KEY_POOL_COUNT=1
benchmarks/longmemeval/fixtures/tiny_longmemeval.json ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "question_id": "tmcra_fixture_user",
4
+ "question_type": "single-session-user",
5
+ "question": "Which drink did the user choose?",
6
+ "answer": "Tea.",
7
+ "question_date": "2026-01-02",
8
+ "haystack_session_ids": ["fixture-session-1"],
9
+ "haystack_dates": ["2026-01-01"],
10
+ "haystack_sessions": [[{"role": "user", "content": "I chose tea."}]]
11
+ },
12
+ {
13
+ "question_id": "tmcra_fixture_update",
14
+ "question_type": "knowledge-update",
15
+ "question": "Where does the user live now?",
16
+ "answer": "Shanghai.",
17
+ "question_date": "2026-01-03",
18
+ "haystack_session_ids": ["fixture-session-2"],
19
+ "haystack_dates": ["2026-01-02"],
20
+ "haystack_sessions": [[{"role": "user", "content": "I moved from Suzhou to Shanghai."}]]
21
+ }
22
+ ]
benchmarks/longmemeval/pyproject.toml ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [build-system]
2
+ requires = ["setuptools>=80", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "tmcra-agent-memory-longmemeval"
7
+ version = "0.1.0"
8
+ description = "TMCRA LongMemEval benchmark reproduction pipeline"
9
+ readme = "README.md"
10
+ requires-python = ">=3.12,<3.13"
11
+ license = "Apache-2.0"
12
+ authors = [{ name = "Yu Haoxin" }]
13
+ dependencies = [
14
+ "huggingface-hub==1.23.0",
15
+ "numpy==2.4.3",
16
+ "safetensors==0.8.0",
17
+ "tokenizers==0.22.2",
18
+ "transformers==5.10.2",
19
+ ]
20
+
21
+ [project.optional-dependencies]
22
+ test = ["pytest==9.0.2"]
23
+
24
+ [tool.setuptools.packages.find]
25
+ where = ["src"]
26
+ namespaces = true
27
+
28
+ [tool.pytest.ini_options]
29
+ testpaths = ["tests"]
30
+ addopts = "-q"
benchmarks/longmemeval/requirements.txt ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ --extra-index-url https://download.pytorch.org/whl/cu128
2
+ torch==2.11.0+cu128
3
+ huggingface-hub==1.23.0
4
+ numpy==2.4.3
5
+ safetensors==0.8.0
6
+ tokenizers==0.22.2
7
+ transformers==5.10.2
8
+ pytest==9.0.2
benchmarks/longmemeval/results/benchmark.json ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "schema_version": "tmcra.benchmark-result.1",
3
+ "benchmark": "LongMemEval",
4
+ "metric": "accuracy",
5
+ "score_percent": 82.2,
6
+ "correct": 411,
7
+ "total": 500,
8
+ "by_question_type": {
9
+ "knowledge-update": {"correct": 71, "total": 78, "score_percent": 91.0},
10
+ "multi-session": {"correct": 90, "total": 133, "score_percent": 67.7},
11
+ "single-session-assistant": {"correct": 55, "total": 56, "score_percent": 98.2},
12
+ "single-session-preference": {"correct": 27, "total": 30, "score_percent": 90.0},
13
+ "single-session-user": {"correct": 67, "total": 70, "score_percent": 95.7},
14
+ "temporal-reasoning": {"correct": 101, "total": 133, "score_percent": 75.9}
15
+ }
16
+ }
benchmarks/longmemeval/scripts/fetch_assets.py ADDED
@@ -0,0 +1,95 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import json
7
+ import os
8
+ import urllib.request
9
+ from pathlib import Path
10
+
11
+ LFS_INCLUDE = "models/tmcra_v4_longmemeval_s500_20260715/*.pt"
12
+
13
+
14
+ def digest(path: Path) -> str:
15
+ hasher = hashlib.sha256()
16
+ with path.open("rb") as handle:
17
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
18
+ hasher.update(chunk)
19
+ return hasher.hexdigest()
20
+
21
+
22
+ def fetch_file(root: Path, row: dict[str, object]) -> None:
23
+ destination = (root / str(row["destination"])).resolve()
24
+ expected = str(row["sha256"])
25
+ expected_bytes = int(row["bytes"])
26
+ if destination.is_file() and destination.stat().st_size == expected_bytes and digest(destination) == expected:
27
+ print(f"verified {row['id']}: {destination}")
28
+ return
29
+ if row.get("storage") == "git_lfs":
30
+ detail = "file is missing"
31
+ if destination.is_file():
32
+ with destination.open("rb") as handle:
33
+ prefix = handle.read(128)
34
+ if prefix.startswith(b"version https://git-lfs.github.com/spec/v1"):
35
+ detail = "only a Git LFS pointer is present"
36
+ else:
37
+ detail = f"size or SHA-256 mismatch ({destination.stat().st_size} bytes)"
38
+ raise RuntimeError(
39
+ f"Git LFS checkpoint {row['id']} is unavailable: {detail}: {destination}\n"
40
+ f"run from the repository root: git lfs pull --include=\"{LFS_INCLUDE}\""
41
+ )
42
+ if not row.get("url"):
43
+ raise RuntimeError(f"asset {row['id']} has neither a valid local Git LFS file nor a download URL")
44
+ destination.parent.mkdir(parents=True, exist_ok=True)
45
+ temporary = destination.with_suffix(destination.suffix + ".part")
46
+ temporary.unlink(missing_ok=True)
47
+ print(f"downloading {row['id']} -> {destination}")
48
+ urllib.request.urlretrieve(str(row["url"]), temporary)
49
+ if temporary.stat().st_size != expected_bytes:
50
+ raise RuntimeError(f"size mismatch for {row['id']}")
51
+ actual = digest(temporary)
52
+ if actual != expected:
53
+ raise RuntimeError(f"SHA-256 mismatch for {row['id']}: {actual}")
54
+ os.replace(temporary, destination)
55
+
56
+
57
+ def fetch_model(root: Path, row: dict[str, object]) -> None:
58
+ from huggingface_hub import snapshot_download
59
+
60
+ destination = root / str(row["destination"])
61
+ destination.parent.mkdir(parents=True, exist_ok=True)
62
+ print(f"downloading {row['repo_id']}@{row['revision']} -> {destination}")
63
+ snapshot_download(
64
+ repo_id=str(row["repo_id"]),
65
+ revision=str(row["revision"]),
66
+ local_dir=destination,
67
+ )
68
+
69
+
70
+ def main() -> int:
71
+ parser = argparse.ArgumentParser()
72
+ parser.add_argument("--manifest", type=Path, default=Path("configs/assets.lock.json"))
73
+ parser.add_argument("--root", type=Path, default=Path.cwd())
74
+ parser.add_argument("--kind", choices=("all", "dataset", "checkpoint", "model"), default="all")
75
+ args = parser.parse_args()
76
+ payload = json.loads(args.manifest.read_text(encoding="utf-8"))
77
+ if args.kind in {"all", "dataset", "checkpoint"}:
78
+ selected = [
79
+ row
80
+ for row in payload["files"]
81
+ if args.kind == "all" or row["kind"] == args.kind
82
+ ]
83
+ # Validate repository-owned LFS files before starting any large network
84
+ # download, so an incomplete clone fails quickly and predictably.
85
+ selected.sort(key=lambda row: row.get("storage") != "git_lfs")
86
+ for row in selected:
87
+ fetch_file(args.root, row)
88
+ if args.kind in {"all", "model"}:
89
+ for row in payload["huggingface_models"]:
90
+ fetch_model(args.root, row)
91
+ return 0
92
+
93
+
94
+ if __name__ == "__main__":
95
+ raise SystemExit(main())
benchmarks/longmemeval/scripts/prepare_longmemeval.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import json
7
+ from collections import Counter
8
+ from pathlib import Path
9
+
10
+
11
+ QUESTION_TYPES = {
12
+ "knowledge-update",
13
+ "multi-session",
14
+ "single-session-assistant",
15
+ "single-session-preference",
16
+ "single-session-user",
17
+ "temporal-reasoning",
18
+ }
19
+
20
+
21
+ def load_and_validate(path: Path) -> list[dict[str, object]]:
22
+ rows = json.loads(path.read_text(encoding="utf-8"))
23
+ if not isinstance(rows, list) or not rows:
24
+ raise ValueError("LongMemEval input must be a non-empty JSON array")
25
+ seen: set[str] = set()
26
+ for index, row in enumerate(rows):
27
+ if not isinstance(row, dict):
28
+ raise ValueError(f"row {index} is not an object")
29
+ qid = str(row.get("question_id") or "").strip()
30
+ if not qid or qid in seen:
31
+ raise ValueError(f"invalid or duplicate question_id at row {index}: {qid!r}")
32
+ seen.add(qid)
33
+ question_type = str(row.get("question_type") or "").strip()
34
+ if question_type not in QUESTION_TYPES:
35
+ raise ValueError(f"{qid}: unexpected question_type {question_type!r}")
36
+ for field in ("question", "answer", "haystack_sessions", "haystack_dates"):
37
+ if field not in row:
38
+ raise ValueError(f"{qid}: missing {field}")
39
+ return rows
40
+
41
+
42
+ def main() -> int:
43
+ parser = argparse.ArgumentParser()
44
+ parser.add_argument("--data", required=True, type=Path)
45
+ parser.add_argument("--qid-output", required=True, type=Path)
46
+ parser.add_argument("--summary-output", type=Path)
47
+ parser.add_argument("--limit", type=int, default=0)
48
+ args = parser.parse_args()
49
+ rows = load_and_validate(args.data)
50
+ selected = rows[: args.limit] if args.limit else rows
51
+ qids = [str(row["question_id"]) for row in selected]
52
+ args.qid_output.parent.mkdir(parents=True, exist_ok=True)
53
+ args.qid_output.write_text("\n".join(qids) + "\n", encoding="utf-8")
54
+ summary = {
55
+ "schema_version": "tmcra.longmemeval-input.1",
56
+ "dataset_sha256": hashlib.sha256(args.data.read_bytes()).hexdigest(),
57
+ "source_count": len(rows),
58
+ "selected_count": len(selected),
59
+ "question_types": dict(sorted(Counter(str(row["question_type"]) for row in selected).items())),
60
+ "qid_output": str(args.qid_output),
61
+ }
62
+ if args.summary_output:
63
+ args.summary_output.parent.mkdir(parents=True, exist_ok=True)
64
+ args.summary_output.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
65
+ print(json.dumps(summary, sort_keys=True))
66
+ return 0
67
+
68
+
69
+ if __name__ == "__main__":
70
+ raise SystemExit(main())
benchmarks/longmemeval/scripts/reproduce_s500.sh ADDED
@@ -0,0 +1,152 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ set -euo pipefail
3
+
4
+ ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
5
+ cd "$ROOT"
6
+
7
+ for file in configs/writer.env configs/answer.env configs/judge.env configs/run.env; do
8
+ if [[ ! -f "$file" ]]; then
9
+ echo "missing $file; copy the matching .example file and fill in your values" >&2
10
+ exit 2
11
+ fi
12
+ done
13
+
14
+ set -a
15
+ source configs/run.env
16
+ set +a
17
+
18
+ LEGACY="$ROOT/src/tmcra_benchmark/legacy"
19
+ VENDOR="$ROOT/src/tmcra_benchmark/vendor/tmcra_integrated"
20
+ HARNESS="$ROOT/src/tmcra_benchmark/harness/run_lme_s10_native_tmcra.py"
21
+ JUDGE="$ROOT/src/tmcra_benchmark/harness/evaluate_qa_official_judge.py"
22
+ export PYTHONPATH="$LEGACY:$VENDOR:$ROOT/src${PYTHONPATH:+:$PYTHONPATH}"
23
+
24
+ : "${DATA:?DATA is required in configs/run.env}"
25
+ : "${QIDS:?QIDS is required in configs/run.env}"
26
+ : "${RUN:?RUN is required in configs/run.env}"
27
+ : "${MODELS:?MODELS is required in configs/run.env}"
28
+ : "${BGE_MODELS:?BGE_MODELS is required in configs/run.env}"
29
+
30
+ python scripts/prepare_longmemeval.py \
31
+ --data "$DATA" \
32
+ --qid-output "$QIDS" \
33
+ --summary-output "$RUN.input_summary.json"
34
+
35
+ if [[ ! -f "$RUN/BUILD_COMPLETE" ]]; then
36
+ build_resume_args=()
37
+ if [[ -d "$RUN" ]]; then
38
+ build_resume_args+=(--resume)
39
+ fi
40
+ python "$LEGACY/run_tmcra_v4_build.py" \
41
+ "${build_resume_args[@]}" \
42
+ --out-dir "$RUN" \
43
+ --qid-list "$QIDS" \
44
+ --data "$DATA" \
45
+ --repo "$VENDOR" \
46
+ --writer-env configs/writer.env \
47
+ --embedding-model "$BGE_MODELS/bge-m3" \
48
+ --writer-concurrency "${WRITER_CONCURRENCY:-20}" \
49
+ --slow-concurrency "${SLOW_CONCURRENCY:-20}" \
50
+ --index-batch-size 16 \
51
+ --device "${RETRIEVAL_DEVICE:-cuda}"
52
+ fi
53
+ python scripts/validate_stage.py build --run-dir "$RUN" --expected-qids "$QIDS"
54
+
55
+ export TMCRA_NODE_MODEL_PATH="$MODELS/node_scorer.pt"
56
+ export TMCRA_PATH_MODEL_PATH="$MODELS/path_scorer.pt"
57
+ export TMCRA_BGE_M3_PATH="$BGE_MODELS/bge-m3"
58
+ export TMCRA_BGE_RERANKER_PATH="$BGE_MODELS/bge-reranker-v2-m3"
59
+
60
+ if [[ ! -f "$RUN/retrieval_s500/RETRIEVAL_COMPLETE" ]]; then
61
+ retrieval_resume_args=()
62
+ if [[ -d "$RUN/retrieval_s500" ]]; then
63
+ retrieval_resume_args+=(--resume)
64
+ fi
65
+ python "$LEGACY/run_tmcra_v4_retrieve.py" \
66
+ "${retrieval_resume_args[@]}" \
67
+ --run-dir "$RUN" \
68
+ --tag retrieval_s500 \
69
+ --qid-list "$QIDS" \
70
+ --writer-env configs/writer.env \
71
+ --repo "$VENDOR" \
72
+ --harness "$HARNESS" \
73
+ --node-model "$MODELS/node_scorer.pt" \
74
+ --path-model "$MODELS/path_scorer.pt" \
75
+ --checkpoint "$MODELS/tmcra_v3_reranker.pt" \
76
+ --embedding-model "$BGE_MODELS/bge-m3" \
77
+ --cross-model "$BGE_MODELS/bge-reranker-v2-m3" \
78
+ --device "${RETRIEVAL_DEVICE:-cuda}" \
79
+ --graph-device "${GRAPH_DEVICE:-cuda}" \
80
+ --dense-k 32 \
81
+ --slow-dense-k 24 \
82
+ --graph-k 24 \
83
+ --composition-mode layered \
84
+ --packing-budget-mode fixed \
85
+ --top-k 8 \
86
+ --cross-batch-size 24
87
+ fi
88
+ python scripts/validate_stage.py retrieve --run-dir "$RUN" --tag retrieval_s500 --expected-qids "$QIDS"
89
+
90
+ if [[ ! -f "$RUN/retrieval_s500_compiled/COMPILE_COMPLETE" ]]; then
91
+ python "$LEGACY/run_tmcra_v4_compile_evidence.py" \
92
+ --evidence "$RUN/retrieval_s500/evidence_windows.jsonl" \
93
+ --retrieval-debug "$RUN/retrieval_s500/retrieval_debug.jsonl" \
94
+ --require-retrieval-debug \
95
+ --out-dir "$RUN/retrieval_s500_compiled" \
96
+ --qid-list "$QIDS" \
97
+ --writer-env configs/writer.env \
98
+ --planner-provider deepseek \
99
+ --planner-model deepseek-v4-pro \
100
+ --workers 4 \
101
+ --timeout 180
102
+ fi
103
+ python scripts/validate_stage.py compile --run-dir "$RUN/retrieval_s500_compiled" --expected-qids "$QIDS"
104
+
105
+ set -a
106
+ source configs/answer.env
107
+ set +a
108
+ # The answer runner resumes from its validated answers.jsonl. Running it again
109
+ # also repairs an incomplete report without repeating completed provider calls.
110
+ python "$LEGACY/run_tmcra_v4_gpt54_answers.py" \
111
+ --evidence "$RUN/retrieval_s500_compiled/evidence_windows.jsonl" \
112
+ --harness "$HARNESS" \
113
+ --out-dir "$RUN/answers_s500" \
114
+ --qid-list "$QIDS" \
115
+ --workers "${ANSWER_WORKERS:-12}" \
116
+ --attempts 3 \
117
+ --answer-window-limit 8 \
118
+ --lane production
119
+ python scripts/validate_stage.py answer --run-dir "$RUN/answers_s500" --expected-qids "$QIDS"
120
+
121
+ set -a
122
+ source configs/judge.env
123
+ set +a
124
+ : "${TMCRA_JUDGE_BASE_URL:?TMCRA_JUDGE_BASE_URL is required}"
125
+ : "${TMCRA_JUDGE_API_KEY:?TMCRA_JUDGE_API_KEY is required}"
126
+
127
+ python "$LEGACY/run_tmcra_v4_parallel_official_judge.py" \
128
+ --judge "$JUDGE" \
129
+ --metric-model gpt-4o \
130
+ --hyp-file "$RUN/answers_s500/answers.jsonl" \
131
+ --ref-file "$DATA" \
132
+ --result-file "$RUN/official_judge_gpt4o_20240806.jsonl" \
133
+ --work-dir "$RUN/official_judge_work" \
134
+ --base-url "$TMCRA_JUDGE_BASE_URL" \
135
+ --api-key-env TMCRA_JUDGE_API_KEY \
136
+ --workers "${JUDGE_WORKERS:-8}" \
137
+ --timeout 180 \
138
+ --process-timeout 240 \
139
+ --max-retries 0 \
140
+ --resume
141
+
142
+ python scripts/validate_stage.py judge \
143
+ --run-dir "$RUN" \
144
+ --expected-qids "$QIDS" \
145
+ --judge-results "$RUN/official_judge_gpt4o_20240806.jsonl"
146
+
147
+ python scripts/summarize_official_judge.py \
148
+ --judge-results "$RUN/official_judge_gpt4o_20240806.jsonl" \
149
+ --dataset "$DATA" \
150
+ --output "$RUN/scorecard.json"
151
+
152
+ echo "TMCRA LongMemEval run complete: $RUN/scorecard.json"
benchmarks/longmemeval/scripts/reproduce_smoke.sh ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ set -euo pipefail
3
+
4
+ ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
5
+ cd "$ROOT"
6
+
7
+ DATA="${DATA:-$ROOT/fixtures/tiny_longmemeval.json}"
8
+ OUT="${OUT:-$ROOT/runs/offline_fixture_check}"
9
+ mkdir -p "$OUT"
10
+
11
+ python scripts/prepare_longmemeval.py \
12
+ --data "$DATA" \
13
+ --qid-output "$OUT/qids.txt" \
14
+ --summary-output "$OUT/input_summary.json"
15
+ python -m pytest
16
+
17
+ echo "Offline fixture and unit tests passed."
18
+ echo "This smoke test does not call Writer, Answer, or Judge providers."
benchmarks/longmemeval/scripts/summarize_official_judge.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import json
6
+ from collections import defaultdict
7
+ from pathlib import Path
8
+ from typing import Any
9
+
10
+
11
+ def read_jsonl(path: Path) -> list[dict[str, Any]]:
12
+ return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
13
+
14
+
15
+ def label(row: dict[str, Any]) -> bool:
16
+ value = row.get("autoeval_label")
17
+ if isinstance(value, dict):
18
+ value = value.get("label")
19
+ if isinstance(value, bool):
20
+ return value
21
+ if isinstance(value, str):
22
+ normalized = value.strip().lower()
23
+ if normalized in {"true", "1", "yes"}:
24
+ return True
25
+ if normalized in {"false", "0", "no"}:
26
+ return False
27
+ raise ValueError(f"invalid autoeval label: {value!r}")
28
+
29
+
30
+ def summarize(dataset: list[dict[str, Any]], judged: list[dict[str, Any]]) -> dict[str, Any]:
31
+ type_by_qid = {str(row["question_id"]): str(row["question_type"]) for row in dataset}
32
+ seen: set[str] = set()
33
+ buckets: dict[str, list[bool]] = defaultdict(list)
34
+ all_labels: list[bool] = []
35
+ for row in judged:
36
+ qid = str(row.get("question_id") or "")
37
+ if qid not in type_by_qid or qid in seen:
38
+ raise ValueError(f"unknown or duplicate question_id: {qid!r}")
39
+ seen.add(qid)
40
+ result = label(row)
41
+ buckets[type_by_qid[qid]].append(result)
42
+ all_labels.append(result)
43
+ if len(seen) != len(dataset):
44
+ missing = sorted(set(type_by_qid) - seen)
45
+ raise ValueError(f"judge results are incomplete: {len(seen)}/{len(dataset)}; first missing={missing[:5]}")
46
+ by_type = {}
47
+ for name, values in sorted(buckets.items()):
48
+ correct = sum(values)
49
+ by_type[name] = {
50
+ "correct": correct,
51
+ "total": len(values),
52
+ "accuracy": round(correct / len(values), 6),
53
+ "accuracy_percent": round(correct * 100 / len(values), 1),
54
+ }
55
+ correct = sum(all_labels)
56
+ return {
57
+ "schema_version": "tmcra.longmemeval-scorecard.1",
58
+ "benchmark": "LongMemEval",
59
+ "correct": correct,
60
+ "total": len(all_labels),
61
+ "accuracy": round(correct / len(all_labels), 6),
62
+ "accuracy_percent": round(correct * 100 / len(all_labels), 1),
63
+ "by_question_type": by_type,
64
+ }
65
+
66
+
67
+ def main() -> int:
68
+ parser = argparse.ArgumentParser()
69
+ parser.add_argument("--judge-results", required=True, type=Path)
70
+ parser.add_argument("--dataset", required=True, type=Path)
71
+ parser.add_argument("--output", required=True, type=Path)
72
+ args = parser.parse_args()
73
+ dataset = json.loads(args.dataset.read_text(encoding="utf-8"))
74
+ scorecard = summarize(dataset, read_jsonl(args.judge_results))
75
+ args.output.parent.mkdir(parents=True, exist_ok=True)
76
+ args.output.write_text(json.dumps(scorecard, indent=2, sort_keys=True) + "\n", encoding="utf-8")
77
+ print(json.dumps(scorecard, sort_keys=True))
78
+ return 0
79
+
80
+
81
+ if __name__ == "__main__":
82
+ raise SystemExit(main())
benchmarks/longmemeval/scripts/validate_stage.py ADDED
@@ -0,0 +1,237 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import json
6
+ from collections.abc import Mapping
7
+ from pathlib import Path
8
+ from typing import Any
9
+
10
+
11
+ PRODUCTION_ANSWER_PROTOCOL = "evidence_operation_bound_v6"
12
+
13
+
14
+ def fail(message: str) -> None:
15
+ raise SystemExit(message)
16
+
17
+
18
+ def require(path: Path) -> Path:
19
+ if not path.is_file():
20
+ fail(f"missing required artifact: {path}")
21
+ return path
22
+
23
+
24
+ def read_json(path: Path) -> dict[str, Any]:
25
+ require(path)
26
+ try:
27
+ value = json.loads(path.read_text(encoding="utf-8"))
28
+ except (OSError, json.JSONDecodeError) as exc:
29
+ fail(f"invalid JSON artifact {path}: {exc}")
30
+ if not isinstance(value, dict):
31
+ fail(f"JSON artifact must be an object: {path}")
32
+ return value
33
+
34
+
35
+ def read_jsonl(path: Path) -> list[dict[str, Any]]:
36
+ require(path)
37
+ rows: list[dict[str, Any]] = []
38
+ for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
39
+ if not line.strip():
40
+ continue
41
+ try:
42
+ value = json.loads(line)
43
+ except json.JSONDecodeError as exc:
44
+ fail(f"invalid JSONL row {path}:{line_number}: {exc}")
45
+ if not isinstance(value, dict):
46
+ fail(f"JSONL row must be an object: {path}:{line_number}")
47
+ rows.append(value)
48
+ return rows
49
+
50
+
51
+ def qids(path: Path) -> list[str]:
52
+ require(path)
53
+ values = [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
54
+ if not values:
55
+ fail(f"QID list is empty: {path}")
56
+ if len(values) != len(set(values)):
57
+ fail(f"QID list contains duplicates: {path}")
58
+ return values
59
+
60
+
61
+ def row_qids(path: Path) -> list[str]:
62
+ values = [str(row.get("question_id") or row.get("qid") or "").strip() for row in read_jsonl(path)]
63
+ if any(not value for value in values):
64
+ fail(f"artifact contains an empty QID: {path}")
65
+ if len(values) != len(set(values)):
66
+ fail(f"artifact contains duplicate QIDs: {path}")
67
+ return values
68
+
69
+
70
+ def assert_qid_order(path: Path, expected: list[str]) -> None:
71
+ actual = row_qids(path)
72
+ if actual != expected:
73
+ fail(f"QID mismatch in {path}: got {len(actual)} ordered rows, expected {len(expected)}")
74
+
75
+
76
+ def assert_complete_report(
77
+ path: Path,
78
+ *,
79
+ expected_count: int,
80
+ count_field: str,
81
+ ) -> dict[str, Any]:
82
+ report = read_json(path)
83
+ if report.get("status") != "complete":
84
+ fail(f"report is not complete: {path}")
85
+ try:
86
+ actual_count = int(report.get(count_field, -1))
87
+ except (TypeError, ValueError):
88
+ fail(f"report has an invalid {count_field}: {path}")
89
+ if actual_count != expected_count:
90
+ fail(f"report {count_field} mismatch in {path}: {actual_count}/{expected_count}")
91
+ return report
92
+
93
+
94
+ def validate_build(run_dir: Path, expected: list[str]) -> None:
95
+ require(run_dir / "BUILD_COMPLETE")
96
+ require(run_dir / "qids.txt")
97
+ if qids(run_dir / "qids.txt") != expected:
98
+ fail("build qids.txt does not match the requested QID order")
99
+
100
+ manifest = read_json(run_dir / "input_manifest.json")
101
+ if manifest.get("status") != "prepared":
102
+ fail("build input manifest is not in prepared state")
103
+ if list(manifest.get("qids") or []) != expected:
104
+ fail("build input manifest QIDs do not match the requested QID order")
105
+ if int(manifest.get("row_count", -1)) != len(expected):
106
+ fail("build input manifest row_count is inconsistent")
107
+ workers = manifest.get("workers")
108
+ if not isinstance(workers, list):
109
+ fail("build input manifest workers must be a list")
110
+ worker_qids = [str(row.get("question_id") or "").strip() for row in workers if isinstance(row, Mapping)]
111
+ if len(worker_qids) != len(workers) or worker_qids != expected:
112
+ fail("build worker QIDs do not match the requested QID order")
113
+
114
+ assert_qid_order(run_dir / "scope_manifest.jsonl", expected)
115
+ assert_qid_order(run_dir / "query_manifest.jsonl", expected)
116
+ assert_complete_report(
117
+ run_dir / "build_report.json",
118
+ expected_count=len(expected),
119
+ count_field="row_count",
120
+ )
121
+ assert_complete_report(
122
+ run_dir / "index_report.json",
123
+ expected_count=len(expected),
124
+ count_field="row_count",
125
+ )
126
+
127
+
128
+ def validate_retrieve(run_dir: Path, tag: str, expected: list[str]) -> None:
129
+ directory = run_dir / tag
130
+ require(directory / "RETRIEVAL_COMPLETE")
131
+ assert_complete_report(
132
+ directory / "report.json",
133
+ expected_count=len(expected),
134
+ count_field="query_count",
135
+ )
136
+ run_report = assert_complete_report(
137
+ directory / "retrieval_run_report.json",
138
+ expected_count=len(expected),
139
+ count_field="query_count",
140
+ )
141
+ if str(run_report.get("tag") or "") != tag:
142
+ fail(f"retrieval tag mismatch in {directory / 'retrieval_run_report.json'}")
143
+ assert_qid_order(directory / "evidence_windows.jsonl", expected)
144
+ assert_qid_order(directory / "retrieval_debug.jsonl", expected)
145
+
146
+
147
+ def validate_compile(run_dir: Path, expected: list[str]) -> None:
148
+ require(run_dir / "COMPILE_COMPLETE")
149
+ assert_complete_report(
150
+ run_dir / "report.json",
151
+ expected_count=len(expected),
152
+ count_field="row_count",
153
+ )
154
+ assert_qid_order(run_dir / "evidence_windows.jsonl", expected)
155
+ assert_qid_order(run_dir / "retrieval_debug.jsonl", expected)
156
+
157
+
158
+ def validate_answer(run_dir: Path, expected: list[str]) -> None:
159
+ report = assert_complete_report(
160
+ run_dir / "report.json",
161
+ expected_count=len(expected),
162
+ count_field="requested",
163
+ )
164
+ if int(report.get("completed", -1)) != len(expected):
165
+ fail(f"answer report completed count mismatch: {report.get('completed')}/{len(expected)}")
166
+ if int(report.get("failure_count", -1)) != 0:
167
+ fail(f"answer report contains failures: {report.get('failure_count')}")
168
+ if report.get("answer_model") != "gpt-5.4":
169
+ fail(f"answer report used an unexpected model: {report.get('answer_model')!r}")
170
+ if report.get("answer_protocol") != PRODUCTION_ANSWER_PROTOCOL:
171
+ fail(f"answer report used an unexpected protocol: {report.get('answer_protocol')!r}")
172
+ if report.get("lane") != "production" or report.get("promotion_eligible") is not True:
173
+ fail("answer report is not a promotion-eligible production run")
174
+ if report.get("compiled_evidence_packet_required") is not True:
175
+ fail("answer report did not require compiled evidence packets")
176
+ answers_path = run_dir / "answers.jsonl"
177
+ assert_qid_order(answers_path, expected)
178
+ for row in read_jsonl(answers_path):
179
+ if row.get("answer_model") != "gpt-5.4":
180
+ fail(f"answer row used an unexpected model: {row.get('question_id')}")
181
+ if row.get("answer_protocol") != PRODUCTION_ANSWER_PROTOCOL:
182
+ fail(f"answer row used an unexpected protocol: {row.get('question_id')}")
183
+
184
+
185
+ def validate_judge(results: Path, expected: list[str]) -> None:
186
+ assert_qid_order(results, expected)
187
+ rows = read_jsonl(results)
188
+ for row in rows:
189
+ label = row.get("autoeval_label")
190
+ if not isinstance(label, Mapping) or type(label.get("label")) is not bool:
191
+ fail(f"judge row has no valid boolean label: {row.get('question_id')}")
192
+ if label.get("model") != "gpt-4o-2024-08-06":
193
+ fail(f"judge row used an unexpected model: {label.get('model')!r}")
194
+
195
+ report_path = results.with_name(f"{results.name}.isolated_report.json")
196
+ report = read_json(report_path)
197
+ if report.get("status") != "complete":
198
+ fail(f"judge report is not complete: {report_path}")
199
+ if int(report.get("row_count", -1)) != len(expected):
200
+ fail(f"judge row_count mismatch in {report_path}")
201
+ if int(report.get("committed_count", -1)) != len(expected):
202
+ fail(f"judge committed_count mismatch in {report_path}")
203
+ if int(report.get("failure_count", -1)) != 0:
204
+ fail(f"judge report contains failures: {report_path}")
205
+ if report.get("metric_model") not in {"gpt-4o", "gpt-4o-2024-08-06"}:
206
+ fail(f"judge report used an unexpected model alias: {report.get('metric_model')!r}")
207
+
208
+
209
+ def main() -> int:
210
+ parser = argparse.ArgumentParser()
211
+ parser.add_argument("stage", choices=("build", "retrieve", "compile", "answer", "judge"))
212
+ parser.add_argument("--run-dir", required=True, type=Path)
213
+ parser.add_argument("--expected-qids", required=True, type=Path)
214
+ parser.add_argument("--tag", default="retrieval_s500")
215
+ parser.add_argument("--judge-results", type=Path)
216
+ args = parser.parse_args()
217
+ expected = qids(args.expected_qids)
218
+
219
+ if args.stage == "build":
220
+ validate_build(args.run_dir, expected)
221
+ elif args.stage == "retrieve":
222
+ validate_retrieve(args.run_dir, args.tag, expected)
223
+ elif args.stage == "compile":
224
+ validate_compile(args.run_dir, expected)
225
+ elif args.stage == "answer":
226
+ validate_answer(args.run_dir, expected)
227
+ else:
228
+ if args.judge_results is None:
229
+ fail("--judge-results is required for the judge stage")
230
+ validate_judge(args.judge_results, expected)
231
+
232
+ print(f"{args.stage} validation: OK ({len(expected)} QIDs)")
233
+ return 0
234
+
235
+
236
+ if __name__ == "__main__":
237
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/__init__.py ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ """TMCRA LongMemEval benchmark reproduction package."""
2
+
3
+ __version__ = "0.1.0"
benchmarks/longmemeval/src/tmcra_benchmark/harness/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ """LongMemEval adapter and official judge integration."""
benchmarks/longmemeval/src/tmcra_benchmark/harness/evaluate_qa_official_judge.py ADDED
@@ -0,0 +1,295 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """LongMemEval official-prompt QA judge.
3
+
4
+ This is an implementation-compatible wrapper around the public LongMemEval
5
+ `src/evaluation/evaluate_qa.py` logic. The scoring prompt and yes/no rule are
6
+ kept equivalent, while the runner adds resume support and avoids extra Python
7
+ package dependencies by using urllib.
8
+
9
+ For official reporting, use:
10
+ --metric-model gpt-4o
11
+ with OPENAI_API_KEY set. The resolved model is gpt-4o-2024-08-06, matching the
12
+ official evaluator's model alias.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import argparse
18
+ import json
19
+ import os
20
+ import sys
21
+ import time
22
+ import urllib.error
23
+ import urllib.request
24
+ from pathlib import Path
25
+ from statistics import mean
26
+ from typing import Any
27
+
28
+
29
+ MODEL_ZOO = {
30
+ "gpt5.4": ("gpt-5.4", "openai"),
31
+ "gpt-5.4": ("gpt-5.4", "openai"),
32
+ "gpt-4o-mini": ("gpt-4o-mini-2024-07-18", "openai"),
33
+ "gpt-4o": ("gpt-4o-2024-08-06", "openai"),
34
+ "gpt-4o-2024-08-06": ("gpt-4o-2024-08-06", "openai"),
35
+ "llama-3.1-70b-instruct": ("meta-llama/Meta-Llama-3.1-70B-Instruct", "local"),
36
+ }
37
+
38
+
39
+ def get_anscheck_prompt(task: str, question: str, answer: str, response: str, abstention: bool = False) -> str:
40
+ if not abstention:
41
+ if task in ["single-session-user", "single-session-assistant", "multi-session"]:
42
+ template = (
43
+ "I will give you a question, a correct answer, and a response from a model. "
44
+ "Please answer yes if the response contains the correct answer. Otherwise, answer no. "
45
+ "If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also answer yes. "
46
+ "If the response only contains a subset of the information required by the answer, answer no. \n\n"
47
+ "Question: {}\n\nCorrect Answer: {}\n\nModel Response: {}\n\n"
48
+ "Is the model response correct? Answer yes or no only."
49
+ )
50
+ return template.format(question, answer, response)
51
+ if task == "temporal-reasoning":
52
+ template = (
53
+ "I will give you a question, a correct answer, and a response from a model. "
54
+ "Please answer yes if the response contains the correct answer. Otherwise, answer no. "
55
+ "If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also answer yes. "
56
+ "If the response only contains a subset of the information required by the answer, answer no. "
57
+ "In addition, do not penalize off-by-one errors for the number of days. "
58
+ "If the question asks for the number of days/weeks/months, etc., and the model makes off-by-one errors (e.g., predicting 19 days when the answer is 18), the model's response is still correct. \n\n"
59
+ "Question: {}\n\nCorrect Answer: {}\n\nModel Response: {}\n\n"
60
+ "Is the model response correct? Answer yes or no only."
61
+ )
62
+ return template.format(question, answer, response)
63
+ if task == "knowledge-update":
64
+ template = (
65
+ "I will give you a question, a correct answer, and a response from a model. "
66
+ "Please answer yes if the response contains the correct answer. Otherwise, answer no. "
67
+ "If the response contains some previous information along with an updated answer, the response should be considered as correct as long as the updated answer is the required answer.\n\n"
68
+ "Question: {}\n\nCorrect Answer: {}\n\nModel Response: {}\n\n"
69
+ "Is the model response correct? Answer yes or no only."
70
+ )
71
+ return template.format(question, answer, response)
72
+ if task == "single-session-preference":
73
+ template = (
74
+ "I will give you a question, a rubric for desired personalized response, and a response from a model. "
75
+ "Please answer yes if the response satisfies the desired response. Otherwise, answer no. "
76
+ "The model does not need to reflect all the points in the rubric. "
77
+ "The response is correct as long as it recalls and utilizes the user's personal information correctly.\n\n"
78
+ "Question: {}\n\nRubric: {}\n\nModel Response: {}\n\n"
79
+ "Is the model response correct? Answer yes or no only."
80
+ )
81
+ return template.format(question, answer, response)
82
+ raise NotImplementedError(f"Unsupported question_type: {task}")
83
+
84
+ template = (
85
+ "I will give you an unanswerable question, an explanation, and a response from a model. "
86
+ "Please answer yes if the model correctly identifies the question as unanswerable. "
87
+ "The model could say that the information is incomplete, or some other information is given but the asked information is not.\n\n"
88
+ "Question: {}\n\nExplanation: {}\n\nModel Response: {}\n\n"
89
+ "Does the model correctly identify the question as unanswerable? Answer yes or no only."
90
+ )
91
+ return template.format(question, answer, response)
92
+
93
+
94
+ def load_json_or_jsonl(path: Path) -> list[dict[str, Any]]:
95
+ text = path.read_text(encoding="utf-8")
96
+ try:
97
+ data = json.loads(text)
98
+ if isinstance(data, list):
99
+ return data
100
+ if isinstance(data, dict):
101
+ return [data]
102
+ raise ValueError(f"{path} JSON root is not a list")
103
+ except json.JSONDecodeError:
104
+ return [json.loads(line) for line in text.splitlines() if line.strip()]
105
+
106
+
107
+ def call_chat_completion(
108
+ *,
109
+ base_url: str,
110
+ api_key: str,
111
+ model: str,
112
+ prompt: str,
113
+ timeout: int,
114
+ max_retries: int,
115
+ ) -> str:
116
+ wire_api = os.environ.get("OPENAI_WIRE_API", "").strip().lower()
117
+ if wire_api == "responses":
118
+ url = base_url.rstrip("/") + "/responses"
119
+ payload = {
120
+ "model": model,
121
+ "input": [{"role": "user", "content": prompt}],
122
+ "temperature": 0,
123
+ "max_output_tokens": 10,
124
+ }
125
+ else:
126
+ url = base_url.rstrip("/") + "/chat/completions"
127
+ payload = {
128
+ "model": model,
129
+ "messages": [{"role": "user", "content": prompt}],
130
+ "n": 1,
131
+ "temperature": 0,
132
+ "max_tokens": 10,
133
+ }
134
+ body = json.dumps(payload).encode("utf-8")
135
+ headers = {
136
+ "Content-Type": "application/json",
137
+ "Authorization": f"Bearer {api_key}",
138
+ }
139
+ last_error: Exception | None = None
140
+ for attempt in range(max_retries + 1):
141
+ req = urllib.request.Request(url, data=body, headers=headers, method="POST")
142
+ try:
143
+ with urllib.request.urlopen(req, timeout=timeout) as resp:
144
+ data = json.loads(resp.read().decode("utf-8"))
145
+ if wire_api == "responses":
146
+ output_text = str(data.get("output_text", "")).strip()
147
+ if output_text:
148
+ return output_text
149
+ chunks: list[str] = []
150
+ for item in data.get("output", []) or []:
151
+ for content in item.get("content", []) or []:
152
+ text = content.get("text") or content.get("content") or ""
153
+ if text:
154
+ chunks.append(str(text))
155
+ return "\n".join(chunks).strip()
156
+ return data["choices"][0]["message"]["content"].strip()
157
+ except (urllib.error.HTTPError, urllib.error.URLError, TimeoutError) as exc:
158
+ last_error = exc
159
+ if attempt >= max_retries:
160
+ break
161
+ time.sleep(min(60, 2**attempt))
162
+ raise RuntimeError(f"chat completion failed after {max_retries + 1} attempts: {last_error}")
163
+
164
+
165
+ def read_existing_labels(result_file: Path) -> dict[str, dict[str, Any]]:
166
+ done: dict[str, dict[str, Any]] = {}
167
+ if not result_file.exists():
168
+ return done
169
+ with result_file.open("r", encoding="utf-8") as f:
170
+ for line in f:
171
+ if not line.strip():
172
+ continue
173
+ item = json.loads(line)
174
+ qid = str(item.get("question_id", ""))
175
+ if qid:
176
+ done[qid] = item
177
+ return done
178
+
179
+
180
+ def main() -> int:
181
+ parser = argparse.ArgumentParser()
182
+ parser.add_argument("--metric-model", default="gpt-4o", choices=sorted(MODEL_ZOO))
183
+ parser.add_argument("--hyp-file", required=True, type=Path)
184
+ parser.add_argument("--ref-file", required=True, type=Path)
185
+ parser.add_argument("--result-file", type=Path)
186
+ parser.add_argument("--base-url", default=None)
187
+ parser.add_argument("--api-key-env", default="OPENAI_API_KEY")
188
+ parser.add_argument("--limit", type=int, default=0)
189
+ parser.add_argument("--timeout", type=int, default=120)
190
+ parser.add_argument("--max-retries", type=int, default=6)
191
+ parser.add_argument("--resume", action="store_true")
192
+ parser.add_argument("--quiet", action="store_true")
193
+ args = parser.parse_args()
194
+
195
+ model, source = MODEL_ZOO[args.metric_model]
196
+ if source == "openai":
197
+ base_url = args.base_url or "https://api.openai.com/v1"
198
+ else:
199
+ base_url = args.base_url or "http://localhost:8001/v1"
200
+
201
+ api_key = os.getenv(args.api_key_env)
202
+ if not api_key:
203
+ print(f"Missing API key env: {args.api_key_env}", file=sys.stderr)
204
+ return 2
205
+
206
+ hypotheses = load_json_or_jsonl(args.hyp_file)
207
+ references = load_json_or_jsonl(args.ref_file)
208
+ qid2qdata = {str(entry["question_id"]): entry for entry in references}
209
+ qid2qtype = {str(entry["question_id"]): entry["question_type"] for entry in references}
210
+
211
+ if args.limit and args.limit > 0:
212
+ hypotheses = hypotheses[: args.limit]
213
+
214
+ result_file = args.result_file or Path(str(args.hyp_file) + f".eval-results-{args.metric_model}")
215
+ result_file.parent.mkdir(parents=True, exist_ok=True)
216
+ existing = read_existing_labels(result_file) if args.resume else {}
217
+
218
+ logs: list[dict[str, Any]] = list(existing.values())
219
+ qtype2acc: dict[str, list[int]] = {}
220
+ for item in logs:
221
+ qid = str(item.get("question_id"))
222
+ if qid in qid2qtype and item.get("autoeval_label"):
223
+ qtype2acc.setdefault(qid2qtype[qid], []).append(1 if item["autoeval_label"].get("label") else 0)
224
+
225
+ mode = "a" if args.resume and result_file.exists() else "w"
226
+ with result_file.open(mode, encoding="utf-8") as out_f:
227
+ for idx, entry in enumerate(hypotheses, 1):
228
+ qid = str(entry.get("question_id", ""))
229
+ if not qid or qid not in qid2qdata:
230
+ print(f"Warning: skipping unknown question_id={qid}", file=sys.stderr)
231
+ continue
232
+ if qid in existing:
233
+ continue
234
+ qtype = qid2qtype[qid]
235
+ qdata = qid2qdata[qid]
236
+ prompt = get_anscheck_prompt(
237
+ qtype,
238
+ str(qdata["question"]),
239
+ str(qdata["answer"]),
240
+ str(entry.get("hypothesis", "")),
241
+ abstention="_abs" in qid,
242
+ )
243
+ eval_response = call_chat_completion(
244
+ base_url=base_url,
245
+ api_key=api_key,
246
+ model=model,
247
+ prompt=prompt,
248
+ timeout=args.timeout,
249
+ max_retries=args.max_retries,
250
+ )
251
+ label = "yes" in eval_response.lower()
252
+ output = dict(entry)
253
+ output["autoeval_label"] = {"model": model, "label": label, "raw_response": eval_response}
254
+ print(json.dumps(output, ensure_ascii=False), file=out_f, flush=True)
255
+ logs.append(output)
256
+ qtype2acc.setdefault(qtype, []).append(1 if label else 0)
257
+ if not args.quiet:
258
+ print(
259
+ json.dumps(
260
+ {
261
+ "index": idx,
262
+ "question_id": qid,
263
+ "label": label,
264
+ "raw_response": eval_response,
265
+ "question": qdata["question"],
266
+ "answer": qdata["answer"],
267
+ "hypothesis": entry.get("hypothesis", ""),
268
+ },
269
+ ensure_ascii=False,
270
+ ),
271
+ flush=True,
272
+ )
273
+
274
+ labels = [1 if x.get("autoeval_label", {}).get("label") else 0 for x in logs if x.get("autoeval_label")]
275
+ summary = {
276
+ "metric_model_alias": args.metric_model,
277
+ "metric_model": model,
278
+ "base_url": base_url,
279
+ "hyp_file": str(args.hyp_file),
280
+ "ref_file": str(args.ref_file),
281
+ "result_file": str(result_file),
282
+ "evaluated_count": len(labels),
283
+ "accuracy": round(mean(labels), 4) if labels else None,
284
+ "by_question_type": {
285
+ k: {"accuracy": round(mean(v), 4), "count": len(v)} for k, v in sorted(qtype2acc.items()) if v
286
+ },
287
+ }
288
+ summary_file = Path(str(result_file) + ".summary.json")
289
+ summary_file.write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
290
+ print(json.dumps(summary, ensure_ascii=False, indent=2))
291
+ return 0
292
+
293
+
294
+ if __name__ == "__main__":
295
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/harness/run_lme_s10_native_tmcra.py ADDED
The diff for this file is too large to render. See raw diff
 
benchmarks/longmemeval/src/tmcra_benchmark/legacy/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ """Hash-pinned TMCRA V4 benchmark implementation."""
benchmarks/longmemeval/src/tmcra_benchmark/legacy/audit_tmcra_v4_chain.py ADDED
The diff for this file is too large to render. See raw diff
 
benchmarks/longmemeval/src/tmcra_benchmark/legacy/build_v3_runtime_dataset.py ADDED
@@ -0,0 +1,791 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import json
7
+ import re
8
+ import time
9
+ from collections import Counter
10
+ from pathlib import Path
11
+ from typing import Any, Iterable, Mapping, Sequence
12
+
13
+ import torch
14
+
15
+ from tmcra_v2_lme_pipeline import EmbeddingCacheVectorizer, session_text_chunks
16
+ from tmcra_v3_schema import (
17
+ CHANNEL_NAMES,
18
+ SCHEMA_VERSION,
19
+ clean_text,
20
+ read_jsonl,
21
+ validate_sample,
22
+ validate_split_isolation,
23
+ write_jsonl,
24
+ )
25
+
26
+
27
+ EVENT_INDEX_RE = re.compile(
28
+ r"^event::longmemeval:(?P<qid>[^:]+):s(?P<session>\d+)_c(?P<parent>\d+)$"
29
+ )
30
+ TEACHER_EVENT_INDEX_RE = re.compile(r":s(?P<session>\d+)_c(?P<parent>\d+)$")
31
+ RAW_TEACHER_EVENT_RE = re.compile(
32
+ r"^lme:[^:]+:session:(?P<session>[^:]+):turn:(?P<turn>\d+)$"
33
+ )
34
+ TURN_BOUNDARY_RE = re.compile(r"\n\[[^\]]+ turn=\d+ role=")
35
+
36
+
37
+ def stable_int(value: str) -> int:
38
+ return int.from_bytes(hashlib.blake2b(value.encode("utf-8"), digest_size=8).digest(), "big")
39
+
40
+
41
+ def normalized_question(value: Any) -> str:
42
+ return clean_text(value).casefold()
43
+
44
+
45
+ def load_teacher_supervision(path: Path) -> dict[str, dict[str, Any]]:
46
+ output: dict[str, dict[str, Any]] = {}
47
+ for row in read_jsonl(path):
48
+ question_key = normalized_question(row.get("question"))
49
+ if not question_key:
50
+ raise RuntimeError("teacher row is missing question text")
51
+ if question_key in output:
52
+ raise RuntimeError(f"duplicate teacher question after exact normalization: {row.get('question')!r}")
53
+ if not bool(row.get("usable_for_training", False)):
54
+ continue
55
+ if not list(row.get("positive_event_ids") or []):
56
+ raise RuntimeError(f"teacher row has no positive events: {row.get('question_id')}")
57
+ output[question_key] = row
58
+ if not output:
59
+ raise RuntimeError(f"no usable teacher supervision rows: {path}")
60
+ return output
61
+
62
+
63
+ def teacher_parent_locations(event_ids: Iterable[Any], valid_locations: set[tuple[int, int]], qid: str) -> list[tuple[int, int]]:
64
+ output: list[tuple[int, int]] = []
65
+ seen: set[tuple[int, int]] = set()
66
+ for raw in event_ids:
67
+ event_id = clean_text(raw)
68
+ match = TEACHER_EVENT_INDEX_RE.search(event_id)
69
+ if match is None:
70
+ raise RuntimeError(f"{qid}: teacher event is not a runtime parent event: {event_id}")
71
+ location = (int(match.group("session")), int(match.group("parent")))
72
+ if location not in valid_locations:
73
+ raise RuntimeError(f"{qid}: teacher event has no parent in the current inventory: {event_id}")
74
+ if location not in seen:
75
+ seen.add(location)
76
+ output.append(location)
77
+ return output
78
+
79
+
80
+ def reconstruct_parent_text(candidates: Sequence[Mapping[str, Any]], indexes: Sequence[int], qid: str) -> str:
81
+ if not indexes:
82
+ raise RuntimeError(f"{qid}: cannot reconstruct an empty parent")
83
+ total = max(int(candidates[index]["source_char_end"]) for index in indexes)
84
+ characters: list[str | None] = [None] * total
85
+ for index in indexes:
86
+ candidate = candidates[index]
87
+ text = str(candidate["text"])
88
+ if "\n" not in text:
89
+ raise RuntimeError(f"{qid}: inventory candidate is missing its payload separator")
90
+ payload = text.split("\n", 1)[1]
91
+ start = int(candidate["source_char_start"])
92
+ end = int(candidate["source_char_end"])
93
+ if len(payload) != end - start:
94
+ raise RuntimeError(f"{qid}: inventory candidate payload length does not match its source span")
95
+ for offset, character in enumerate(payload, start=start):
96
+ existing = characters[offset]
97
+ if existing is not None and existing != character:
98
+ raise RuntimeError(f"{qid}: overlapping inventory subchunks disagree")
99
+ characters[offset] = character
100
+ if any(character is None for character in characters):
101
+ raise RuntimeError(f"{qid}: reconstructed parent contains a character coverage gap")
102
+ return "".join(character for character in characters if character is not None)
103
+
104
+
105
+ def teacher_turn_candidate_indexes(
106
+ event_ids: Iterable[Any],
107
+ *,
108
+ candidates: Sequence[Mapping[str, Any]],
109
+ parent_candidates: Mapping[tuple[int, int], Sequence[int]],
110
+ qid: str,
111
+ ) -> set[int]:
112
+ output: set[int] = set()
113
+ session_locations: dict[str, list[tuple[int, int]]] = {}
114
+ for location, indexes in parent_candidates.items():
115
+ session_id = clean_text(candidates[indexes[0]].get("session_id"))
116
+ session_locations.setdefault(session_id, []).append(location)
117
+ parent_text_cache: dict[tuple[int, int], str] = {}
118
+ for raw in event_ids:
119
+ event_id = clean_text(raw)
120
+ match = RAW_TEACHER_EVENT_RE.match(event_id)
121
+ if match is None:
122
+ raise RuntimeError(f"{qid}: teacher raw event is malformed: {event_id}")
123
+ session_id = clean_text(match.group("session"))
124
+ turn = int(match.group("turn"))
125
+ marker = f"[{session_id} turn={turn} role="
126
+ matches: list[tuple[tuple[int, int], int, int]] = []
127
+ for location in session_locations.get(session_id, []):
128
+ if location not in parent_text_cache:
129
+ parent_text_cache[location] = reconstruct_parent_text(
130
+ candidates, parent_candidates[location], qid
131
+ )
132
+ parent_text = parent_text_cache[location]
133
+ turn_start = parent_text.find(marker)
134
+ if turn_start < 0:
135
+ continue
136
+ next_turn = TURN_BOUNDARY_RE.search(parent_text, turn_start + len(marker))
137
+ turn_end = next_turn.start() if next_turn else len(parent_text)
138
+ matches.append((location, turn_start, turn_end))
139
+ if len(matches) != 1:
140
+ raise RuntimeError(
141
+ f"{qid}: teacher event must map to exactly one current parent turn, got {len(matches)}: {event_id}"
142
+ )
143
+ location, turn_start, turn_end = matches[0]
144
+ matched_subchunks = {
145
+ index
146
+ for index in parent_candidates[location]
147
+ if int(candidates[index]["source_char_start"]) < turn_end
148
+ and int(candidates[index]["source_char_end"]) > turn_start
149
+ }
150
+ if not matched_subchunks:
151
+ raise RuntimeError(f"{qid}: teacher turn has no overlapping subchunk: {event_id}")
152
+ output.update(matched_subchunks)
153
+ return output
154
+
155
+
156
+ def covered_windows(text: str, max_chars: int, overlap: int) -> list[tuple[int, int]]:
157
+ if max_chars <= 0:
158
+ raise RuntimeError("subchunk payload size must be positive")
159
+ if overlap < 0 or overlap >= max_chars:
160
+ raise RuntimeError("subchunk overlap must be non-negative and smaller than the payload")
161
+ if not text:
162
+ return []
163
+ spans: list[tuple[int, int]] = []
164
+ start = 0
165
+ while start < len(text):
166
+ hard_end = min(len(text), start + max_chars)
167
+ end = hard_end
168
+ if hard_end < len(text):
169
+ floor = max(start + max_chars // 2, hard_end - 160)
170
+ split_at = text.rfind(" ", floor, hard_end)
171
+ if split_at > start:
172
+ end = split_at
173
+ if end <= start:
174
+ end = hard_end
175
+ spans.append((start, end))
176
+ if end >= len(text):
177
+ break
178
+ start = max(start + 1, end - overlap)
179
+ if spans[0][0] != 0 or spans[-1][1] != len(text):
180
+ raise AssertionError("subchunk coverage does not reach both text boundaries")
181
+ if any(right_start > left_end for (_, left_end), (right_start, _) in zip(spans, spans[1:])):
182
+ raise AssertionError("subchunk coverage contains a gap")
183
+ return spans
184
+
185
+
186
+ def load_lme(path: Path) -> list[dict[str, Any]]:
187
+ value = json.loads(path.read_text(encoding="utf-8"))
188
+ if not isinstance(value, list) or not value:
189
+ raise RuntimeError(f"LongMemEval data must be a non-empty list: {path}")
190
+ rows = [row for row in value if isinstance(row, dict)]
191
+ qids = [clean_text(row.get("question_id")) for row in rows]
192
+ if len(qids) != len(set(qids)):
193
+ raise RuntimeError("LongMemEval question_id values are not unique")
194
+ return rows
195
+
196
+
197
+ def choose_holdout(rows: Sequence[Mapping[str, Any]], count: int) -> list[str]:
198
+ if count <= 0 or count >= len(rows):
199
+ raise RuntimeError(f"holdout count must be between 1 and {len(rows) - 1}")
200
+ by_type: dict[str, list[str]] = {}
201
+ for row in rows:
202
+ qid = clean_text(row.get("question_id"))
203
+ qtype = clean_text(row.get("question_type")) or "unknown"
204
+ by_type.setdefault(qtype, []).append(qid)
205
+ selected: list[str] = []
206
+ for qtype, qids in sorted(by_type.items()):
207
+ qids.sort(key=lambda qid: (stable_int(f"tmcra-v3-holdout:{qid}"), qid))
208
+ quota = int(len(qids) * count / len(rows))
209
+ selected.extend(qids[:quota])
210
+ selected_set = set(selected)
211
+ remaining = [clean_text(row.get("question_id")) for row in rows if clean_text(row.get("question_id")) not in selected_set]
212
+ remaining.sort(key=lambda qid: (stable_int(f"tmcra-v3-holdout-fill:{qid}"), qid))
213
+ selected.extend(remaining[: max(0, count - len(selected))])
214
+ selected = sorted(selected[:count])
215
+ if len(selected) != count or len(set(selected)) != count:
216
+ raise RuntimeError("failed to construct an exact unique holdout")
217
+ return selected
218
+
219
+
220
+ def make_inventory_row(
221
+ row: Mapping[str, Any],
222
+ split: str,
223
+ parent_chars: int,
224
+ subchunk_chars: int,
225
+ subchunk_overlap: int,
226
+ ) -> dict[str, Any]:
227
+ qid = clean_text(row.get("question_id"))
228
+ question = clean_text(row.get("question"))
229
+ question_date = clean_text(row.get("question_date"))
230
+ query_text = f"{question}\nQuestion date: {question_date}" if question_date else question
231
+ sessions = list(row.get("haystack_sessions") or [])
232
+ session_ids = [clean_text(value) for value in list(row.get("haystack_session_ids") or [])]
233
+ dates = [clean_text(value) for value in list(row.get("haystack_dates") or [])]
234
+ answer_session_ids = {clean_text(value) for value in list(row.get("answer_session_ids") or []) if clean_text(value)}
235
+ if not sessions or len(sessions) != len(session_ids):
236
+ raise RuntimeError(f"{qid}: haystack session/id mismatch")
237
+ candidates: list[dict[str, Any]] = []
238
+ found_gold: set[str] = set()
239
+ parent_chunk_count = 0
240
+ source_parent_chars = 0
241
+ for index, session in enumerate(sessions):
242
+ session_id = session_ids[index]
243
+ date = dates[index] if index < len(dates) else ""
244
+ if not isinstance(session, Sequence) or isinstance(session, (str, bytes, bytearray)):
245
+ raise RuntimeError(f"{qid}: session {index} is not a turn list")
246
+ turns = list(session)
247
+ if any(not isinstance(turn, Mapping) for turn in turns):
248
+ raise RuntimeError(f"{qid}: session {index} contains a non-object turn")
249
+ parent_chunks = session_text_chunks(
250
+ session_id=session_id,
251
+ date=date,
252
+ turns=turns,
253
+ max_chars=parent_chars,
254
+ max_chunks=0,
255
+ )
256
+ if not parent_chunks:
257
+ parent_chunks = [f"LongMemEval session_id={session_id} date={date}"]
258
+ relevant = session_id in answer_session_ids
259
+ if relevant:
260
+ found_gold.add(session_id)
261
+ for parent_index, parent_text in enumerate(parent_chunks, start=1):
262
+ parent_chunk_count += 1
263
+ source_parent_chars += len(parent_text)
264
+ prefix = (
265
+ f"LongMemEval session_id={session_id} date={date} "
266
+ f"parent_chunk={parent_index:02d}"
267
+ )
268
+ payload_chars = subchunk_chars - len(prefix) - 1
269
+ spans = covered_windows(parent_text, payload_chars, subchunk_overlap)
270
+ for subchunk_index, (char_start, char_end) in enumerate(spans, start=1):
271
+ text = f"{prefix}\n{parent_text[char_start:char_end]}"
272
+ if len(text) > subchunk_chars:
273
+ raise AssertionError(f"{qid}: generated subchunk exceeds configured character cap")
274
+ candidates.append(
275
+ {
276
+ "candidate_id": (
277
+ f"chunk::longmemeval:{qid}:s{index:03d}_c{parent_index:02d}_p{subchunk_index:02d}"
278
+ ),
279
+ "text": text,
280
+ "session_id": session_id,
281
+ "session_index": index,
282
+ "parent_chunk_index": parent_index,
283
+ "subchunk_index": subchunk_index,
284
+ "source_char_start": char_start,
285
+ "source_char_end": char_end,
286
+ "labels": {
287
+ "relevance": relevant,
288
+ "hard_negative": False,
289
+ "evidence_role": "positive" if relevant else "negative",
290
+ "target_scope": "answer_session_bag",
291
+ },
292
+ }
293
+ )
294
+ if found_gold != answer_session_ids:
295
+ raise RuntimeError(f"{qid}: answer sessions missing from haystack: {sorted(answer_session_ids - found_gold)}")
296
+ return {
297
+ "schema_version": SCHEMA_VERSION,
298
+ "question_id": qid,
299
+ "question": question,
300
+ "question_date": question_date,
301
+ "query_text": query_text,
302
+ "question_type": clean_text(row.get("question_type")),
303
+ "split": split,
304
+ "gold_answer": clean_text(row.get("answer")),
305
+ "answer_session_ids": sorted(answer_session_ids),
306
+ "supervision": {
307
+ "target_type": "multi_instance_answer_session_bag",
308
+ "loss": "negative_log_positive_probability_mass",
309
+ },
310
+ "inventory_metadata": {
311
+ "session_count": len(sessions),
312
+ "parent_chunk_count": parent_chunk_count,
313
+ "candidate_subchunk_count": len(candidates),
314
+ "source_parent_chars": source_parent_chars,
315
+ "parent_chars": parent_chars,
316
+ "subchunk_chars": subchunk_chars,
317
+ "subchunk_overlap": subchunk_overlap,
318
+ "truncated": False,
319
+ },
320
+ "candidates": candidates,
321
+ }
322
+
323
+
324
+ def command_inventory(args: argparse.Namespace) -> None:
325
+ out_dir = Path(args.out_dir)
326
+ out_dir.mkdir(parents=True, exist_ok=False)
327
+ rows = load_lme(Path(args.data))
328
+ holdout = set(choose_holdout(rows, args.holdout_count))
329
+ inventory_rows = [
330
+ make_inventory_row(
331
+ row,
332
+ "holdout" if clean_text(row.get("question_id")) in holdout else "train",
333
+ args.parent_chars,
334
+ args.subchunk_chars,
335
+ args.subchunk_overlap,
336
+ )
337
+ for row in rows
338
+ ]
339
+ holdout_rows = [row for row in inventory_rows if row["split"] == "holdout"]
340
+ train_rows = [row for row in inventory_rows if row["split"] == "train"]
341
+ isolation = validate_split_isolation(train_rows, holdout_rows)
342
+ write_jsonl(out_dir / "inventory.jsonl", inventory_rows)
343
+ (out_dir / "holdout_qids.txt").write_text("\n".join(sorted(holdout)) + "\n", encoding="utf-8")
344
+ report = {
345
+ "schema_version": SCHEMA_VERSION,
346
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
347
+ "data": str(Path(args.data).resolve()),
348
+ "row_count": len(inventory_rows),
349
+ "train_count": len(train_rows),
350
+ "holdout_count": len(holdout_rows),
351
+ "candidate_count": sum(len(row["candidates"]) for row in inventory_rows),
352
+ "candidate_count_avg": round(sum(len(row["candidates"]) for row in inventory_rows) / len(inventory_rows), 4),
353
+ "parent_chunk_count": sum(row["inventory_metadata"]["parent_chunk_count"] for row in inventory_rows),
354
+ "session_count": sum(row["inventory_metadata"]["session_count"] for row in inventory_rows),
355
+ "positive_bag_candidate_count": sum(
356
+ int(candidate["labels"]["relevance"])
357
+ for row in inventory_rows
358
+ for candidate in row["candidates"]
359
+ ),
360
+ "max_candidate_chars": max(len(candidate["text"]) for row in inventory_rows for candidate in row["candidates"]),
361
+ "truncated_candidate_count": 0,
362
+ "parent_chars": args.parent_chars,
363
+ "subchunk_chars": args.subchunk_chars,
364
+ "subchunk_overlap": args.subchunk_overlap,
365
+ "question_types": dict(Counter(row["question_type"] for row in inventory_rows)),
366
+ "holdout_question_types": dict(Counter(row["question_type"] for row in holdout_rows)),
367
+ "isolation": isolation,
368
+ }
369
+ (out_dir / "inventory_report.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
370
+ print(json.dumps(report, indent=2, sort_keys=True))
371
+
372
+
373
+ def load_graph_debug(run_dir: Path) -> dict[str, dict[str, Any]]:
374
+ output: dict[str, dict[str, Any]] = {}
375
+ paths = sorted(run_dir.glob("shard_*_out/samples_debug.jsonl"))
376
+ if not paths:
377
+ raise FileNotFoundError(f"no graph debug shards under {run_dir}")
378
+ for path in paths:
379
+ for row in read_jsonl(path):
380
+ qid = clean_text(row.get("question_id"))
381
+ if qid:
382
+ output[qid] = row
383
+ return output
384
+
385
+
386
+ def ordered_parent_locations(
387
+ event_ids: Iterable[Any],
388
+ *,
389
+ qid: str,
390
+ valid_locations: set[tuple[int, int]],
391
+ ) -> list[tuple[int, int]]:
392
+ output: list[tuple[int, int]] = []
393
+ seen: set[tuple[int, int]] = set()
394
+ for raw in event_ids:
395
+ event_id = clean_text(raw)
396
+ match = EVENT_INDEX_RE.match(event_id)
397
+ if match is None:
398
+ continue
399
+ if clean_text(match.group("qid")) != qid:
400
+ raise RuntimeError(f"{qid}: graph debug contains an event for another question: {event_id}")
401
+ location = (int(match.group("session")), int(match.group("parent")))
402
+ if location not in valid_locations:
403
+ raise RuntimeError(f"{qid}: graph event has no exact parent chunk mapping: {event_id}")
404
+ if location in seen:
405
+ continue
406
+ seen.add(location)
407
+ output.append(location)
408
+ return output
409
+
410
+
411
+ def expand_parent_locations(
412
+ locations: Iterable[tuple[int, int]],
413
+ parent_candidates: Mapping[tuple[int, int], Sequence[int]],
414
+ ) -> list[int]:
415
+ output: list[int] = []
416
+ seen: set[int] = set()
417
+ for location in locations:
418
+ for index in parent_candidates[location]:
419
+ if index not in seen:
420
+ seen.add(index)
421
+ output.append(index)
422
+ return output
423
+
424
+
425
+ def rrank(rank: int | None) -> float:
426
+ return 0.0 if rank is None else 1.0 / float(rank + 1)
427
+
428
+
429
+ def canonical_candidate(
430
+ raw: Mapping[str, Any],
431
+ *,
432
+ relevant: bool,
433
+ session_relevant: bool,
434
+ teacher_relevant: bool | None,
435
+ dense_score: float,
436
+ dense_rank: int,
437
+ graph_rank: int | None,
438
+ graph_selected: bool,
439
+ graph_final: bool,
440
+ session_count: int,
441
+ forced_positive: bool,
442
+ hard_negative: bool,
443
+ ) -> dict[str, Any]:
444
+ relevant = bool(relevant)
445
+ hard_negative = bool(hard_negative and not relevant)
446
+ channels = {
447
+ "dense_score": float(dense_score),
448
+ "dense_rank_rr": rrank(dense_rank),
449
+ "graph_rank_rr": rrank(graph_rank),
450
+ "graph_selected": float(bool(graph_selected)),
451
+ "graph_final": float(bool(graph_final)),
452
+ "recency_norm": float(raw["session_index"]) / float(max(1, session_count - 1)),
453
+ }
454
+ if tuple(channels.keys()) != CHANNEL_NAMES:
455
+ raise AssertionError("channel construction order changed")
456
+ return {
457
+ "candidate_id": raw["candidate_id"],
458
+ "text": raw["text"],
459
+ "session_id": raw["session_id"],
460
+ "session_index": int(raw["session_index"]),
461
+ "parent_chunk_index": int(raw["parent_chunk_index"]),
462
+ "subchunk_index": int(raw["subchunk_index"]),
463
+ "source_char_start": int(raw["source_char_start"]),
464
+ "source_char_end": int(raw["source_char_end"]),
465
+ "channels": channels,
466
+ "labels": {
467
+ "relevance": relevant,
468
+ "session_relevance": bool(session_relevant),
469
+ "teacher_evidence_relevance": teacher_relevant,
470
+ "hard_negative": hard_negative,
471
+ "evidence_role": "positive" if relevant else ("hard_negative" if hard_negative else "negative"),
472
+ "target_scope": "answer_session_bag",
473
+ },
474
+ "provenance": {
475
+ "dense_retrieved": dense_rank >= 0,
476
+ "graph_retrieved": graph_rank is not None,
477
+ "graph_selected": bool(graph_selected),
478
+ "graph_final": bool(graph_final),
479
+ "forced_positive_for_training": bool(forced_positive),
480
+ },
481
+ }
482
+
483
+
484
+ def metric_hit(order: Sequence[int], positives: set[int], k: int) -> int:
485
+ return int(any(index in positives for index in order[:k]))
486
+
487
+
488
+ def command_finalize(args: argparse.Namespace) -> None:
489
+ inventory_rows = read_jsonl(Path(args.inventory))
490
+ graph_debug = load_graph_debug(Path(args.graph_debug_run))
491
+ teacher_by_question = load_teacher_supervision(Path(args.teacher_labels))
492
+ vectorizer = EmbeddingCacheVectorizer(
493
+ cache_dir=args.embedding_cache,
494
+ dim=args.text_dim,
495
+ expected_backend="hf",
496
+ expected_model=args.embedding_model,
497
+ )
498
+ out_dir = Path(args.out_dir)
499
+ out_dir.mkdir(parents=True, exist_ok=False)
500
+ train_rows: list[dict[str, Any]] = []
501
+ holdout_rows: list[dict[str, Any]] = []
502
+ full_rows: list[dict[str, Any]] = []
503
+ full_inventory_rows: list[dict[str, Any]] = []
504
+ provider = Counter()
505
+ teacher_provider = Counter()
506
+ supervision_counts = Counter()
507
+ for source in inventory_rows:
508
+ qid = clean_text(source["question_id"])
509
+ candidates = list(source["candidates"])
510
+ session_count = int(source["inventory_metadata"]["session_count"])
511
+ candidate_count = len(candidates)
512
+ parent_candidates: dict[tuple[int, int], list[int]] = {}
513
+ for index, candidate in enumerate(candidates):
514
+ location = (int(candidate["session_index"]), int(candidate["parent_chunk_index"]))
515
+ parent_candidates.setdefault(location, []).append(index)
516
+ valid_parent_locations = set(parent_candidates)
517
+ qvec = vectorizer.encode_one(source["query_text"])
518
+ evec = torch.stack([vectorizer.encode_one(candidate["text"]) for candidate in candidates])
519
+ dense_scores = (evec @ qvec).tolist()
520
+ dense_order = sorted(range(candidate_count), key=lambda index: (-dense_scores[index], index))
521
+ dense_rank = {index: rank for rank, index in enumerate(dense_order)}
522
+ debug = graph_debug.get(qid, {})
523
+ retrieval = dict(debug.get("retrieval") or {})
524
+ selected_parents = ordered_parent_locations(
525
+ retrieval.get("selected_event_ids") or [], qid=qid, valid_locations=valid_parent_locations
526
+ )
527
+ final_parents = ordered_parent_locations(
528
+ retrieval.get("final_hit_event_ids") or [], qid=qid, valid_locations=valid_parent_locations
529
+ )
530
+ recall_parents = ordered_parent_locations(
531
+ retrieval.get("recall_event_ids") or [], qid=qid, valid_locations=valid_parent_locations
532
+ )
533
+ graph_parents: list[tuple[int, int]] = []
534
+ graph_parent_seen: set[tuple[int, int]] = set()
535
+ for location in [*selected_parents, *recall_parents]:
536
+ if location not in graph_parent_seen:
537
+ graph_parent_seen.add(location)
538
+ graph_parents.append(location)
539
+ graph_parent_rank = {location: rank for rank, location in enumerate(graph_parents)}
540
+ graph_order = expand_parent_locations(graph_parents, parent_candidates)
541
+ graph_rank = {
542
+ index: graph_parent_rank[(int(candidates[index]["session_index"]), int(candidates[index]["parent_chunk_index"]))]
543
+ for index in graph_order
544
+ }
545
+ selected_indexes = set(expand_parent_locations(selected_parents, parent_candidates))
546
+ final_indexes = set(expand_parent_locations(final_parents, parent_candidates))
547
+ session_positives = {index for index, candidate in enumerate(candidates) if candidate["labels"]["relevance"]}
548
+ teacher_row = teacher_by_question.get(normalized_question(source["question"]))
549
+ teacher_positives: set[int] | None = None
550
+ teacher_hard: set[int] = set()
551
+ if teacher_row is not None:
552
+ teacher_metadata = dict(teacher_row.get("metadata") or {})
553
+ raw_teacher_positive = list(teacher_metadata.get("raw_teacher_positive_event_ids") or [])
554
+ raw_teacher_hard = list(teacher_metadata.get("raw_teacher_hard_negative_event_ids") or [])
555
+ if not raw_teacher_positive:
556
+ raise RuntimeError(f"{qid}: aligned teacher row is missing raw positive turn ids")
557
+ teacher_positives = teacher_turn_candidate_indexes(
558
+ raw_teacher_positive,
559
+ candidates=candidates,
560
+ parent_candidates=parent_candidates,
561
+ qid=qid,
562
+ )
563
+ teacher_hard = teacher_turn_candidate_indexes(
564
+ raw_teacher_hard,
565
+ candidates=candidates,
566
+ parent_candidates=parent_candidates,
567
+ qid=qid,
568
+ )
569
+ if not teacher_positives:
570
+ raise RuntimeError(f"{qid}: teacher supervision produced no positive subchunks")
571
+ mixed_teacher_subchunks = teacher_positives & teacher_hard
572
+ if mixed_teacher_subchunks:
573
+ supervision_counts["teacher_mixed_positive_hard_subchunks"] += len(mixed_teacher_subchunks)
574
+ teacher_hard -= teacher_positives
575
+ positives = teacher_positives
576
+ supervision = {
577
+ "target_type": "teacher_aligned_turn_bag",
578
+ "loss": "negative_log_positive_probability_mass",
579
+ "training_weight": float(teacher_row.get("training_weight_hint", 1.0) or 1.0),
580
+ "teacher_label_model": clean_text((teacher_row.get("metadata") or {}).get("teacher_label_model")),
581
+ "teacher_verify_model": clean_text((teacher_row.get("metadata") or {}).get("teacher_verify_model")),
582
+ "teacher_confidence": clean_text(teacher_row.get("label_confidence")),
583
+ }
584
+ else:
585
+ positives = session_positives
586
+ supervision = {
587
+ "target_type": "multi_instance_answer_session_bag",
588
+ "loss": "negative_log_positive_probability_mass",
589
+ "training_weight": float(args.weak_supervision_weight),
590
+ }
591
+ supervision_counts[supervision["target_type"]] += 1
592
+ provider["n"] += 1
593
+ for k in (1, 5, 10, 20, 32):
594
+ provider[f"primary_dense_r@{k}"] += metric_hit(dense_order, positives, k)
595
+ provider[f"primary_graph_r@{k}"] += metric_hit(graph_order, positives, k)
596
+ provider[f"session_dense_r@{k}"] += metric_hit(dense_order, session_positives, k)
597
+ provider[f"session_graph_r@{k}"] += metric_hit(graph_order, session_positives, k)
598
+ if teacher_positives is not None:
599
+ teacher_provider[f"dense_r@{k}"] += metric_hit(dense_order, teacher_positives, k)
600
+ teacher_provider[f"graph_r@{k}"] += metric_hit(graph_order, teacher_positives, k)
601
+ if teacher_positives is not None:
602
+ teacher_provider["n"] += 1
603
+ runtime_indexes: list[int] = []
604
+ runtime_seen: set[int] = set()
605
+ graph_runtime_order = expand_parent_locations(graph_parents[: args.graph_k], parent_candidates)
606
+ for index in [*dense_order[: args.dense_k], *graph_runtime_order]:
607
+ if index not in runtime_seen:
608
+ runtime_seen.add(index)
609
+ runtime_indexes.append(index)
610
+ provider["union_gold_hit"] += int(bool(positives & runtime_seen))
611
+ provider["union_session_hit"] += int(bool(session_positives & runtime_seen))
612
+ if teacher_positives is not None:
613
+ teacher_provider["union_hit"] += int(bool(teacher_positives & runtime_seen))
614
+ pre_force_hit = bool(positives & runtime_seen)
615
+ training_indexes = list(runtime_indexes)
616
+ forced_indexes: set[int] = set()
617
+ if source["split"] == "train" and not pre_force_hit:
618
+ best_positive = next(index for index in dense_order if index in positives)
619
+ training_indexes.append(best_positive)
620
+ runtime_seen.add(best_positive)
621
+ forced_indexes.add(best_positive)
622
+ def build_row(indexes: Sequence[int], split: str, allow_forced: bool) -> dict[str, Any]:
623
+ output_candidates = []
624
+ for index in indexes:
625
+ hard = index not in positives and (
626
+ index in teacher_hard
627
+ or dense_rank[index] < args.hard_negative_k
628
+ or graph_rank.get(index, 10**9) < args.hard_negative_k
629
+ )
630
+ output_candidates.append(
631
+ canonical_candidate(
632
+ candidates[index],
633
+ relevant=index in positives,
634
+ session_relevant=index in session_positives,
635
+ teacher_relevant=None if teacher_positives is None else index in teacher_positives,
636
+ dense_score=dense_scores[index],
637
+ dense_rank=dense_rank[index],
638
+ graph_rank=graph_rank.get(index),
639
+ graph_selected=index in selected_indexes,
640
+ graph_final=index in final_indexes,
641
+ session_count=session_count,
642
+ forced_positive=allow_forced and index in forced_indexes,
643
+ hard_negative=hard,
644
+ )
645
+ )
646
+ return {
647
+ "schema_version": SCHEMA_VERSION,
648
+ "question_id": qid,
649
+ "question": source["question"],
650
+ "question_date": source["question_date"],
651
+ "query_text": source["query_text"],
652
+ "question_type": source["question_type"],
653
+ "split": split,
654
+ "gold_answer": source["gold_answer"],
655
+ "answer_session_ids": source["answer_session_ids"],
656
+ "supervision": supervision,
657
+ "candidates": output_candidates,
658
+ "pool_metadata": {
659
+ "session_count": session_count,
660
+ "inventory_count": candidate_count,
661
+ "graph_parent_event_count": len(graph_parents),
662
+ "graph_runtime_candidate_count": len(graph_runtime_order),
663
+ "runtime_pool_count": len(runtime_indexes),
664
+ "pre_force_gold_hit": pre_force_hit,
665
+ "pre_force_session_hit": bool(session_positives & set(runtime_indexes)),
666
+ "pre_force_teacher_hit": None if teacher_positives is None else bool(teacher_positives & set(runtime_indexes)),
667
+ "forced_positive_count": len(forced_indexes) if allow_forced else 0,
668
+ "dense_k": args.dense_k,
669
+ "graph_k": args.graph_k,
670
+ "graph_debug_available": bool(debug),
671
+ "graph_retrieval_mode": retrieval.get("retrieval_mode"),
672
+ },
673
+ }
674
+ runtime_row = build_row(runtime_indexes, "full_eval", False)
675
+ validate_sample(runtime_row, require_positive=False)
676
+ full_rows.append(runtime_row)
677
+ all_row = build_row(list(range(candidate_count)), "full_eval", False)
678
+ validate_sample(all_row, require_positive=True)
679
+ full_inventory_rows.append(all_row)
680
+ if source["split"] == "train":
681
+ train_row = build_row(training_indexes, "train", True)
682
+ validate_sample(train_row, require_positive=True)
683
+ train_rows.append(train_row)
684
+ else:
685
+ holdout_row = build_row(runtime_indexes, "holdout", False)
686
+ validate_sample(holdout_row, require_positive=False)
687
+ holdout_rows.append(holdout_row)
688
+ isolation = validate_split_isolation(train_rows, holdout_rows)
689
+ if len(train_rows) + len(holdout_rows) != len(inventory_rows):
690
+ raise RuntimeError("finalized row count mismatch")
691
+ write_jsonl(out_dir / "train.jsonl", train_rows)
692
+ write_jsonl(out_dir / "holdout.jsonl", holdout_rows)
693
+ write_jsonl(out_dir / "full_eval.jsonl", full_rows)
694
+ write_jsonl(out_dir / "full_inventory_eval.jsonl", full_inventory_rows)
695
+ if args.dev_count <= 0 or args.dev_count >= len(train_rows):
696
+ raise RuntimeError("dev_count must leave non-empty optimization train and dev sets")
697
+ ordered_train_qids = sorted(
698
+ (row["question_id"] for row in train_rows),
699
+ key=lambda qid: (stable_int(f"tmcra-v3-dev:{qid}"), qid),
700
+ )
701
+ dev_qids = set(ordered_train_qids[: args.dev_count])
702
+ model_train_rows = [row for row in train_rows if row["question_id"] not in dev_qids]
703
+ full_by_qid = {row["question_id"]: row for row in full_rows}
704
+ model_dev_rows = [full_by_qid[qid] for qid in ordered_train_qids[: args.dev_count]]
705
+ if any(row["pool_metadata"]["forced_positive_count"] for row in model_dev_rows):
706
+ raise RuntimeError("runtime dev rows must never contain forced positives")
707
+ write_jsonl(out_dir / "model_train.jsonl", model_train_rows)
708
+ write_jsonl(out_dir / "model_dev_runtime.jsonl", model_dev_rows)
709
+ (out_dir / "model_dev_qids.txt").write_text("\n".join(ordered_train_qids[: args.dev_count]) + "\n", encoding="utf-8")
710
+ cross_cache_rows = [*train_rows, *holdout_rows]
711
+ if len({row["question_id"] for row in cross_cache_rows}) != len(inventory_rows):
712
+ raise RuntimeError("cross-cache input must contain exactly one row per question")
713
+ write_jsonl(out_dir / "cross_cache_samples.jsonl", cross_cache_rows)
714
+ n = max(1, int(provider["n"]))
715
+ report = {
716
+ "schema_version": SCHEMA_VERSION,
717
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
718
+ "train_count": len(train_rows),
719
+ "holdout_count": len(holdout_rows),
720
+ "full_eval_count": len(full_rows),
721
+ "full_inventory_eval_count": len(full_inventory_rows),
722
+ "cross_cache_sample_count": len(cross_cache_rows),
723
+ "cross_cache_pair_count": sum(len(row["candidates"]) for row in cross_cache_rows),
724
+ "model_train_count": len(model_train_rows),
725
+ "model_dev_runtime_count": len(model_dev_rows),
726
+ "model_dev_runtime_candidate_miss_count": sum(
727
+ int(not any(candidate["labels"]["relevance"] for candidate in row["candidates"]))
728
+ for row in model_dev_rows
729
+ ),
730
+ "isolation": isolation,
731
+ "candidate_pool": {
732
+ "dense_k": args.dense_k,
733
+ "graph_k": args.graph_k,
734
+ "hard_negative_k": args.hard_negative_k,
735
+ "train_candidate_avg": round(sum(len(row["candidates"]) for row in train_rows) / max(1, len(train_rows)), 4),
736
+ "holdout_candidate_avg": round(sum(len(row["candidates"]) for row in holdout_rows) / max(1, len(holdout_rows)), 4),
737
+ "full_inventory_candidate_avg": round(sum(len(row["candidates"]) for row in full_inventory_rows) / max(1, len(full_inventory_rows)), 4),
738
+ "train_forced_positive_rows": sum(int(row["pool_metadata"]["forced_positive_count"] > 0) for row in train_rows),
739
+ },
740
+ "provider_metrics": {
741
+ key: round(value / n, 6) for key, value in sorted(provider.items()) if key != "n"
742
+ },
743
+ "teacher_provider_metrics": {
744
+ key: round(value / max(1, int(teacher_provider["n"])), 6)
745
+ for key, value in sorted(teacher_provider.items())
746
+ if key != "n"
747
+ },
748
+ "teacher_supervision_count": int(teacher_provider["n"]),
749
+ "supervision_counts": dict(supervision_counts),
750
+ "teacher_labels": str(Path(args.teacher_labels).resolve()),
751
+ "weak_supervision_weight": args.weak_supervision_weight,
752
+ "embedding_cache": str(Path(args.embedding_cache).resolve()),
753
+ "graph_debug_run": str(Path(args.graph_debug_run).resolve()),
754
+ }
755
+ (out_dir / "report.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
756
+ print(json.dumps(report, indent=2, sort_keys=True))
757
+
758
+
759
+ def main() -> int:
760
+ parser = argparse.ArgumentParser()
761
+ sub = parser.add_subparsers(dest="command", required=True)
762
+ inventory = sub.add_parser("inventory")
763
+ inventory.add_argument("--data", required=True)
764
+ inventory.add_argument("--out-dir", required=True)
765
+ inventory.add_argument("--holdout-count", type=int, default=100)
766
+ inventory.add_argument("--parent-chars", type=int, default=7000)
767
+ inventory.add_argument("--subchunk-chars", type=int, default=1800)
768
+ inventory.add_argument("--subchunk-overlap", type=int, default=200)
769
+ finalize = sub.add_parser("finalize")
770
+ finalize.add_argument("--inventory", required=True)
771
+ finalize.add_argument("--embedding-cache", required=True)
772
+ finalize.add_argument("--embedding-model", required=True)
773
+ finalize.add_argument("--graph-debug-run", required=True)
774
+ finalize.add_argument("--teacher-labels", required=True)
775
+ finalize.add_argument("--out-dir", required=True)
776
+ finalize.add_argument("--text-dim", type=int, default=1024)
777
+ finalize.add_argument("--dense-k", type=int, default=32)
778
+ finalize.add_argument("--graph-k", type=int, default=24)
779
+ finalize.add_argument("--hard-negative-k", type=int, default=12)
780
+ finalize.add_argument("--dev-count", type=int, default=40)
781
+ finalize.add_argument("--weak-supervision-weight", type=float, default=0.35)
782
+ args = parser.parse_args()
783
+ if args.command == "inventory":
784
+ command_inventory(args)
785
+ else:
786
+ command_finalize(args)
787
+ return 0
788
+
789
+
790
+ if __name__ == "__main__":
791
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/ops/audit_tmcra_v4_subject_attribution.py ADDED
@@ -0,0 +1,877 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Audit user-memory ownership inside pasted or forwarded documents.
3
+
4
+ The deterministic router only selects document-shaped source messages. DeepSeek
5
+ Pro makes every semantic keep/quarantine decision from exact Source excerpts.
6
+ Original Source records remain immutable; quarantined Fast records stay in the
7
+ append-only graph but are excluded from current retrieval.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import argparse
13
+ import hashlib
14
+ import json
15
+ import os
16
+ import re
17
+ import sqlite3
18
+ import sys
19
+ from contextlib import closing
20
+ from datetime import datetime, timezone
21
+ from pathlib import Path
22
+ from typing import Any, Mapping, Protocol, Sequence
23
+
24
+ REPO_ROOT = Path(__file__).resolve().parents[1]
25
+ if str(REPO_ROOT) not in sys.path:
26
+ sys.path.insert(0, str(REPO_ROOT))
27
+
28
+ from tmcra_v4_batch_writer import DeepSeekBatchClient
29
+
30
+
31
+ PROMPT_VERSION = "tmcra-v4-subject-attribution-2026-07-14.3"
32
+ MODEL = "deepseek-v4-pro"
33
+ CURRENT_STATES = {"active", "parallel_active", "promoted", "challenged"}
34
+ DECISIONS = {"keep_user", "quarantine_third_party", "quarantine_ambiguous"}
35
+
36
+ SYSTEM_PROMPT = """You audit subject ownership for a production personal-memory system.
37
+ Return exactly one JSON object and no prose:
38
+ {"decisions":[{"memory_id":"exact supplied ID","decision":"keep_user|quarantine_third_party|quarantine_ambiguous","actual_subject":"chat_user for keep_user; named local subject for quarantine_third_party; otherwise empty","chat_user_bridge_quote":"exact outside-artifact identity bridge for keep_user; otherwise empty","reason":"concise source-grounded reason"}]}.
39
+ Return every supplied candidate exactly once and never invent an ID.
40
+
41
+ The chat user is the person who submitted the outer message to this memory system. An author, sender,
42
+ signatory, quoted speaker, recipient, company, or document subject inside an embedded artifact is a separate
43
+ identity and must never be equated with the chat user merely because the artifact appears in a user-role
44
+ message. Text inside a pasted email thread, article, resume, log, transcript, or document is not the chat
45
+ user's conversational voice. A sender or signatory inside the artifact is third party by default.
46
+
47
+ keep_user is allowed only when text outside the embedded artifact explicitly bridges the chat user to the
48
+ local author or subject, for example "I wrote the email below" or "this is my resume". A sender name,
49
+ signature, first-person wording inside the artifact, or mailbox label is not an identity bridge. If the
50
+ source starts directly with mailbox UI or document text and contains no outside bridge, do not use
51
+ keep_user. A mailbox line such as "to me" identifies the mailbox owner as a recipient, not as the sender.
52
+ For keep_user, actual_subject must be exactly "chat_user". Use quarantine_third_party when the source
53
+ locally attributes the fact to a named author, sender, signatory, quoted speaker, company, or document
54
+ subject. Use quarantine_ambiguous when no actual subject can be established. Useful document facts remain
55
+ in immutable Source; do not force them into the chat user's Fast profile.
56
+ Every keep_user decision must cite chat_user_bridge_quote as an exact Source substring outside the artifact.
57
+ The candidate evidence quote itself cannot be used as that bridge. For both quarantine decisions,
58
+ chat_user_bridge_quote must be empty. Candidate objects intentionally contain no Writer claim or slot;
59
+ judge ownership from Source only.
60
+ Do not use outside knowledge, model confidence, benchmark labels, or the assertion wording as authority.
61
+ Judge against the exact source excerpts and offsets."""
62
+
63
+ HEADER_RE = re.compile(
64
+ r"(?im)^(?:from|to|subject|sent|de|para|asunto|enviado):\s*\S"
65
+ )
66
+ MAILBOX_RE = re.compile(r"(?im)^\s*to\s+(?:me|[A-Z][^\n,]{0,40})(?:,|\s*$)")
67
+ SIGNATURE_RE = re.compile(r"(?im)^\s*(?:regards|best regards|sincerely),?\s*$")
68
+ LEGAL_RE = re.compile(r"(?i)confidentiality notice|electronic communications privacy act")
69
+ THREAD_RE = re.compile(r"(?i)view entire message|scanned by gmail|\bon .{0,120} wrote:\s*$", re.M)
70
+
71
+
72
+ class AttributionError(RuntimeError):
73
+ pass
74
+
75
+
76
+ class AttributionClient(Protocol):
77
+ def complete(self, payload: Mapping[str, Any]) -> tuple[str, Mapping[str, Any]]:
78
+ ...
79
+
80
+
81
+ class DeepSeekProAttributionClient:
82
+ def __init__(self) -> None:
83
+ base_url = (
84
+ os.getenv("TMCRA_DEEPSEEK_PRO_BASE_URL")
85
+ or os.getenv("TMCRA_WRITER_BASE_URL")
86
+ or "https://api.deepseek.com/v1"
87
+ )
88
+ raw_keys = (
89
+ os.getenv("TMCRA_DEEPSEEK_PRO_KEY_POOL")
90
+ or os.getenv("TMCRA_WRITER_API_KEY_POOL")
91
+ or os.getenv("TMCRA_DEEPSEEK_WRITER_KEY_POOL")
92
+ or ""
93
+ )
94
+ keys = [item.strip() for item in raw_keys.split(",") if item.strip()]
95
+ if not keys:
96
+ raise AttributionError("DeepSeek Pro key pool is empty; no fallback is allowed")
97
+ max_tokens = int(os.getenv("TMCRA_DEEPSEEK_PRO_MAX_TOKENS", "16384"))
98
+ self.client = DeepSeekBatchClient(
99
+ base_url=base_url,
100
+ model=MODEL,
101
+ api_keys=keys,
102
+ timeout=float(os.getenv("TMCRA_DEEPSEEK_PRO_TIMEOUT", "180")),
103
+ max_tokens=max_tokens,
104
+ )
105
+
106
+ def complete(self, payload: Mapping[str, Any]) -> tuple[str, Mapping[str, Any]]:
107
+ return self.client._complete(
108
+ model=MODEL,
109
+ system_prompt=SYSTEM_PROMPT,
110
+ payload=payload,
111
+ stage="subject_attribution_pro",
112
+ )
113
+
114
+
115
+ def _now() -> str:
116
+ return datetime.now(timezone.utc).isoformat(timespec="milliseconds")
117
+
118
+
119
+ def _json(value: Any) -> str:
120
+ return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
121
+
122
+
123
+ def _digest(value: Any) -> str:
124
+ return hashlib.sha256(_json(value).encode("utf-8")).hexdigest()
125
+
126
+
127
+ def _text(value: Any) -> str:
128
+ return str(value).strip() if value is not None else ""
129
+
130
+
131
+ def document_route_reasons(content: str) -> list[str]:
132
+ """Route only; these signals never decide whether an assertion is valid."""
133
+ reasons: list[str] = []
134
+ if HEADER_RE.search(content):
135
+ reasons.append("mail_header")
136
+ if MAILBOX_RE.search(content):
137
+ reasons.append("mailbox_recipient_line")
138
+ if SIGNATURE_RE.search(content):
139
+ reasons.append("signature_block")
140
+ if LEGAL_RE.search(content):
141
+ reasons.append("mail_legal_notice")
142
+ if THREAD_RE.search(content):
143
+ reasons.append("quoted_thread")
144
+ return reasons if len(reasons) >= 2 else []
145
+
146
+
147
+ def _metadata(row: Mapping[str, Any]) -> dict[str, Any]:
148
+ try:
149
+ value = json.loads(str(row["metadata_json"]))
150
+ except (KeyError, TypeError, json.JSONDecodeError) as exc:
151
+ raise AttributionError("record metadata is invalid JSON") from exc
152
+ if not isinstance(value, dict):
153
+ raise AttributionError("record metadata must be an object")
154
+ return value
155
+
156
+
157
+ def _initialize(connection: sqlite3.Connection) -> None:
158
+ connection.execute(
159
+ """
160
+ CREATE TABLE IF NOT EXISTS v4_subject_attribution_audits (
161
+ audit_id TEXT PRIMARY KEY,
162
+ scope_id TEXT NOT NULL,
163
+ message_id TEXT NOT NULL,
164
+ prompt_version TEXT NOT NULL,
165
+ model TEXT NOT NULL,
166
+ request_json TEXT NOT NULL,
167
+ request_sha256 TEXT NOT NULL,
168
+ status TEXT NOT NULL,
169
+ response_json TEXT NOT NULL DEFAULT '',
170
+ response_sha256 TEXT NOT NULL DEFAULT '',
171
+ call_metadata_json TEXT NOT NULL DEFAULT '{}',
172
+ decisions_json TEXT NOT NULL DEFAULT '[]',
173
+ error TEXT NOT NULL DEFAULT '',
174
+ created_at TEXT NOT NULL,
175
+ updated_at TEXT NOT NULL,
176
+ UNIQUE(scope_id,message_id,request_sha256)
177
+ )
178
+ """
179
+ )
180
+
181
+
182
+ def _source_rows(connection: sqlite3.Connection, scope_id: str) -> list[sqlite3.Row]:
183
+ return connection.execute(
184
+ "SELECT scope_id,message_id,session_index,message_index,message_role,"
185
+ "source_turn_index,content,content_sha256 FROM v4_source_journal "
186
+ "WHERE scope_id=? AND message_role='user' AND status='enriched' "
187
+ "ORDER BY session_index,message_index",
188
+ (scope_id,),
189
+ ).fetchall()
190
+
191
+
192
+ def _candidate_records(
193
+ connection: sqlite3.Connection, scope_id: str, source: Mapping[str, Any]
194
+ ) -> list[dict[str, Any]]:
195
+ output: list[dict[str, Any]] = []
196
+ rows = connection.execute(
197
+ "SELECT memory_id,value,slot_key,turn_index,state,metadata_json FROM records "
198
+ "WHERE scope_id=? AND turn_index=? ORDER BY memory_id",
199
+ (scope_id, int(source["source_turn_index"])),
200
+ ).fetchall()
201
+ for row in rows:
202
+ metadata = _metadata(row)
203
+ if (
204
+ metadata.get("content_variant") != "product_semantic_memory"
205
+ or metadata.get("memory_layer") != "fast"
206
+ or metadata.get("node_kind") != "atomic_user_assertion"
207
+ or _text(row["state"]) not in CURRENT_STATES
208
+ or _text(metadata.get("message_id")) != _text(source["message_id"])
209
+ ):
210
+ continue
211
+ start = metadata.get("evidence_char_start")
212
+ end = metadata.get("evidence_char_end")
213
+ if (
214
+ isinstance(start, bool)
215
+ or isinstance(end, bool)
216
+ or not isinstance(start, int)
217
+ or not isinstance(end, int)
218
+ or not 0 <= start < end <= len(str(source["content"]))
219
+ ):
220
+ raise AttributionError(f"{row['memory_id']}: evidence offsets are invalid")
221
+ quote = str(source["content"])[start:end]
222
+ expected_quote = _text(metadata.get("source_span") or metadata.get("raw_content"))
223
+ if quote != expected_quote:
224
+ raise AttributionError(f"{row['memory_id']}: Source quote drift")
225
+ output.append(
226
+ {
227
+ "memory_id": str(row["memory_id"]),
228
+ "claim_text": str(row["value"]),
229
+ "canonical_slot": _text(
230
+ metadata.get("canonical_slot_key") or row["slot_key"]
231
+ ),
232
+ "evidence_quote": quote,
233
+ "evidence_char_start": start,
234
+ "evidence_char_end": end,
235
+ }
236
+ )
237
+ return output
238
+
239
+
240
+ def _source_segments(content: str, candidates: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
241
+ ranges = [(0, min(len(content), 700))]
242
+ for candidate in candidates:
243
+ start = max(0, int(candidate["evidence_char_start"]) - 420)
244
+ end = min(len(content), int(candidate["evidence_char_end"]) + 420)
245
+ ranges.append((start, end))
246
+ ranges.sort()
247
+ merged: list[list[int]] = []
248
+ for start, end in ranges:
249
+ if merged and start <= merged[-1][1] + 80:
250
+ merged[-1][1] = max(merged[-1][1], end)
251
+ else:
252
+ merged.append([start, end])
253
+ return [
254
+ {"char_start": start, "char_end": end, "text": content[start:end]}
255
+ for start, end in merged
256
+ ]
257
+
258
+
259
+ def scan_database(database: Path, scope_id: str) -> list[dict[str, Any]]:
260
+ with closing(sqlite3.connect(database)) as connection:
261
+ connection.row_factory = sqlite3.Row
262
+ jobs: list[dict[str, Any]] = []
263
+ for source in _source_rows(connection, scope_id):
264
+ reasons = document_route_reasons(str(source["content"]))
265
+ if not reasons:
266
+ continue
267
+ review_candidates = _candidate_records(connection, scope_id, source)
268
+ if not review_candidates:
269
+ continue
270
+ candidates = [
271
+ {
272
+ "memory_id": item["memory_id"],
273
+ "evidence_quote": item["evidence_quote"],
274
+ "evidence_char_start": item["evidence_char_start"],
275
+ "evidence_char_end": item["evidence_char_end"],
276
+ }
277
+ for item in review_candidates
278
+ ]
279
+ payload = {
280
+ "scope_id": scope_id,
281
+ "message_id": str(source["message_id"]),
282
+ "message_role": "user",
283
+ "route_reasons": reasons,
284
+ "source_segments": _source_segments(
285
+ str(source["content"]), review_candidates
286
+ ),
287
+ "candidates": candidates,
288
+ }
289
+ jobs.append(
290
+ {
291
+ "database": str(database),
292
+ "scope_id": scope_id,
293
+ "message_id": str(source["message_id"]),
294
+ "session_index": int(source["session_index"]),
295
+ "message_index": int(source["message_index"]),
296
+ "source_turn_index": int(source["source_turn_index"]),
297
+ "route_reasons": reasons,
298
+ "payload": payload,
299
+ "review_candidates": review_candidates,
300
+ "request_sha256": _digest(
301
+ {
302
+ "prompt_version": PROMPT_VERSION,
303
+ "model": MODEL,
304
+ "payload": payload,
305
+ }
306
+ ),
307
+ }
308
+ )
309
+ return jobs
310
+
311
+
312
+ def validate_decisions(
313
+ raw: Any, payload: Mapping[str, Any]
314
+ ) -> list[dict[str, str]]:
315
+ candidates = payload.get("candidates")
316
+ source_segments = payload.get("source_segments")
317
+ if not isinstance(candidates, list) or not isinstance(source_segments, list):
318
+ raise AttributionError("attribution request payload is malformed")
319
+ candidate_ids = [
320
+ _text(item.get("memory_id"))
321
+ for item in candidates
322
+ if isinstance(item, Mapping)
323
+ ]
324
+ if len(candidate_ids) != len(candidates) or not all(candidate_ids):
325
+ raise AttributionError("attribution request candidate identity is malformed")
326
+ candidate_by_id = {
327
+ _text(item["memory_id"]): item
328
+ for item in candidates
329
+ if isinstance(item, Mapping)
330
+ }
331
+ source_texts = [
332
+ _text(item.get("text"))
333
+ for item in source_segments
334
+ if isinstance(item, Mapping) and _text(item.get("text"))
335
+ ]
336
+ if not isinstance(raw, Mapping) or set(raw) != {"decisions"}:
337
+ raise AttributionError("attribution response must contain exactly decisions")
338
+ decisions = raw.get("decisions")
339
+ if not isinstance(decisions, list) or len(decisions) != len(candidate_ids):
340
+ raise AttributionError("attribution response decision count changed")
341
+ normalized: list[dict[str, str]] = []
342
+ seen: set[str] = set()
343
+ for index, item in enumerate(decisions):
344
+ if not isinstance(item, Mapping) or set(item) != {
345
+ "memory_id",
346
+ "decision",
347
+ "actual_subject",
348
+ "chat_user_bridge_quote",
349
+ "reason",
350
+ }:
351
+ raise AttributionError(f"decisions[{index}] has an invalid shape")
352
+ memory_id = _text(item.get("memory_id"))
353
+ decision = _text(item.get("decision"))
354
+ actual_subject = _text(item.get("actual_subject"))
355
+ bridge_quote = _text(item.get("chat_user_bridge_quote"))
356
+ reason = _text(item.get("reason"))
357
+ if memory_id not in candidate_ids or memory_id in seen:
358
+ raise AttributionError(f"decisions[{index}] changed candidate identity")
359
+ if decision not in DECISIONS or not reason:
360
+ raise AttributionError(f"decisions[{index}] has an invalid decision")
361
+ if decision == "keep_user" and actual_subject != "chat_user":
362
+ raise AttributionError(
363
+ f"decisions[{index}] keep_user must bind actual_subject to chat_user"
364
+ )
365
+ if decision == "keep_user":
366
+ evidence_quote = _text(candidate_by_id[memory_id].get("evidence_quote"))
367
+ if (
368
+ not bridge_quote
369
+ or bridge_quote == evidence_quote
370
+ or not any(bridge_quote in text for text in source_texts)
371
+ ):
372
+ raise AttributionError(
373
+ f"decisions[{index}] keep_user lacks an exact outside-artifact bridge"
374
+ )
375
+ elif bridge_quote:
376
+ raise AttributionError(
377
+ f"decisions[{index}] quarantine decision cannot cite a chat-user bridge"
378
+ )
379
+ if decision == "quarantine_third_party" and not actual_subject:
380
+ raise AttributionError(
381
+ f"decisions[{index}] must identify the third-party subject"
382
+ )
383
+ if decision == "quarantine_ambiguous" and actual_subject:
384
+ raise AttributionError(
385
+ f"decisions[{index}] ambiguous subject must remain empty"
386
+ )
387
+ seen.add(memory_id)
388
+ normalized.append(
389
+ {
390
+ "memory_id": memory_id,
391
+ "decision": decision,
392
+ "actual_subject": actual_subject,
393
+ "chat_user_bridge_quote": bridge_quote,
394
+ "reason": reason,
395
+ }
396
+ )
397
+ if seen != set(candidate_ids):
398
+ raise AttributionError("attribution response omitted candidate identities")
399
+ normalized.sort(key=lambda item: candidate_ids.index(item["memory_id"]))
400
+ return normalized
401
+
402
+
403
+ def _replacement_head(
404
+ connection: sqlite3.Connection,
405
+ scope_id: str,
406
+ slot_key: str,
407
+ excluded_ids: set[str],
408
+ ) -> str:
409
+ rows = connection.execute(
410
+ "SELECT h.memory_id,r.state FROM slot_history h JOIN records r "
411
+ "ON r.scope_id=h.scope_id AND r.memory_id=h.memory_id "
412
+ "WHERE h.scope_id=? AND h.slot_key=? ORDER BY h.ordinal DESC",
413
+ (scope_id, slot_key),
414
+ ).fetchall()
415
+ return next(
416
+ (
417
+ str(row["memory_id"])
418
+ for row in rows
419
+ if str(row["memory_id"]) not in excluded_ids
420
+ and _text(row["state"]) in CURRENT_STATES
421
+ ),
422
+ "",
423
+ )
424
+
425
+
426
+ def _repair_slot_head(
427
+ connection: sqlite3.Connection,
428
+ scope_id: str,
429
+ slot_key: str,
430
+ memory_id: str,
431
+ excluded_ids: set[str],
432
+ ) -> None:
433
+ head = connection.execute(
434
+ "SELECT memory_id FROM slot_heads WHERE scope_id=? AND slot_key=?",
435
+ (scope_id, slot_key),
436
+ ).fetchone()
437
+ if head is None or str(head["memory_id"]) != memory_id:
438
+ return
439
+ replacement = _replacement_head(connection, scope_id, slot_key, excluded_ids)
440
+ if replacement:
441
+ connection.execute(
442
+ "UPDATE slot_heads SET memory_id=? WHERE scope_id=? AND slot_key=?",
443
+ (replacement, scope_id, slot_key),
444
+ )
445
+ else:
446
+ connection.execute(
447
+ "DELETE FROM slot_heads WHERE scope_id=? AND slot_key=?",
448
+ (scope_id, slot_key),
449
+ )
450
+
451
+
452
+ def _apply_decisions(
453
+ connection: sqlite3.Connection,
454
+ job: Mapping[str, Any],
455
+ audit_id: str,
456
+ decisions: Sequence[Mapping[str, str]],
457
+ ) -> dict[str, Any]:
458
+ candidate_ids = [item["memory_id"] for item in job["payload"]["candidates"]]
459
+ rows = connection.execute(
460
+ "SELECT memory_id,slot_key,state,metadata_json FROM records WHERE scope_id=? "
461
+ f"AND memory_id IN ({','.join('?' for _ in candidate_ids)})",
462
+ (job["scope_id"], *candidate_ids),
463
+ ).fetchall()
464
+ if {str(row["memory_id"]) for row in rows} != set(candidate_ids):
465
+ raise AttributionError("candidate records changed before attribution commit")
466
+ decision_by_id = {item["memory_id"]: item for item in decisions}
467
+ quarantine_ids = {
468
+ item["memory_id"]
469
+ for item in decisions
470
+ if item["decision"] != "keep_user"
471
+ }
472
+ parent_signature_map: dict[str, tuple[str, Mapping[str, str]]] = {}
473
+ for row in rows:
474
+ memory_id = str(row["memory_id"])
475
+ if memory_id not in quarantine_ids:
476
+ continue
477
+ signature = _text(_metadata(row).get("event_signature"))
478
+ if not signature:
479
+ continue
480
+ if signature in parent_signature_map:
481
+ raise AttributionError("quarantined parent event signature is not unique")
482
+ parent_signature_map[signature] = (memory_id, decision_by_id[memory_id])
483
+
484
+ dependent_rows: list[sqlite3.Row] = []
485
+ if parent_signature_map:
486
+ for row in connection.execute(
487
+ "SELECT memory_id,slot_key,state,metadata_json FROM records WHERE scope_id=?",
488
+ (job["scope_id"],),
489
+ ).fetchall():
490
+ if str(row["memory_id"]) in candidate_ids:
491
+ continue
492
+ parent_signature = _text(
493
+ _metadata(row).get("facet_parent_event_signature")
494
+ )
495
+ if parent_signature in parent_signature_map:
496
+ dependent_rows.append(row)
497
+
498
+ dependent_ids = {str(row["memory_id"]) for row in dependent_rows}
499
+ all_quarantine_ids = quarantine_ids | dependent_ids
500
+ changed_ids: list[str] = []
501
+ for row in rows:
502
+ memory_id = str(row["memory_id"])
503
+ decision = decision_by_id[memory_id]
504
+ metadata = _metadata(row)
505
+ metadata["subject_attribution_audit_id"] = audit_id
506
+ metadata["subject_attribution_prompt_version"] = PROMPT_VERSION
507
+ metadata["subject_attribution_model"] = MODEL
508
+ metadata["subject_attribution_decision"] = decision["decision"]
509
+ metadata["subject_attribution_actual_subject"] = decision["actual_subject"]
510
+ metadata["subject_attribution_chat_user_bridge_quote"] = decision[
511
+ "chat_user_bridge_quote"
512
+ ]
513
+ metadata["subject_attribution_reason"] = decision["reason"]
514
+ state = str(row["state"])
515
+ if memory_id in quarantine_ids:
516
+ state = "quarantined"
517
+ metadata["excluded_from_retrieval"] = True
518
+ metadata["conflict_action"] = "subject_attribution_quarantine"
519
+ connection.execute(
520
+ "UPDATE records SET state=?,metadata_json=? WHERE scope_id=? AND memory_id=?",
521
+ (state, _json(metadata), job["scope_id"], memory_id),
522
+ )
523
+ changed_ids.append(memory_id)
524
+ if memory_id not in quarantine_ids:
525
+ continue
526
+ slot_key = str(row["slot_key"])
527
+ _repair_slot_head(
528
+ connection,
529
+ job["scope_id"],
530
+ slot_key,
531
+ memory_id,
532
+ all_quarantine_ids,
533
+ )
534
+
535
+ for row in dependent_rows:
536
+ memory_id = str(row["memory_id"])
537
+ metadata = _metadata(row)
538
+ parent_signature = _text(metadata.get("facet_parent_event_signature"))
539
+ parent_memory_id, parent_decision = parent_signature_map[parent_signature]
540
+ metadata["subject_attribution_audit_id"] = audit_id
541
+ metadata["subject_attribution_prompt_version"] = PROMPT_VERSION
542
+ metadata["subject_attribution_model"] = MODEL
543
+ metadata["subject_attribution_decision"] = "quarantine_with_parent"
544
+ metadata["subject_attribution_parent_memory_id"] = parent_memory_id
545
+ metadata["subject_attribution_parent_decision"] = parent_decision["decision"]
546
+ metadata["subject_attribution_actual_subject"] = parent_decision[
547
+ "actual_subject"
548
+ ]
549
+ metadata["subject_attribution_chat_user_bridge_quote"] = ""
550
+ metadata["subject_attribution_reason"] = (
551
+ "structural facet of quarantined parent assertion"
552
+ )
553
+ metadata["excluded_from_retrieval"] = True
554
+ metadata["conflict_action"] = "subject_attribution_parent_quarantine"
555
+ connection.execute(
556
+ "UPDATE records SET state='quarantined',metadata_json=? "
557
+ "WHERE scope_id=? AND memory_id=?",
558
+ (_json(metadata), job["scope_id"], memory_id),
559
+ )
560
+ changed_ids.append(memory_id)
561
+ _repair_slot_head(
562
+ connection,
563
+ job["scope_id"],
564
+ str(row["slot_key"]),
565
+ memory_id,
566
+ all_quarantine_ids,
567
+ )
568
+ return {
569
+ "changed_memory_ids": sorted(changed_ids),
570
+ "quarantined_memory_ids": sorted(quarantine_ids),
571
+ "cascaded_quarantined_memory_ids": sorted(dependent_ids),
572
+ "kept_memory_ids": sorted(set(candidate_ids) - quarantine_ids),
573
+ }
574
+
575
+
576
+ def _reused_application_state(
577
+ connection: sqlite3.Connection,
578
+ job: Mapping[str, Any],
579
+ audit_id: str,
580
+ decisions: Sequence[Mapping[str, str]],
581
+ ) -> dict[str, Any]:
582
+ candidate_ids = {item["memory_id"] for item in job["payload"]["candidates"]}
583
+ quarantine_ids = {
584
+ item["memory_id"] for item in decisions if item["decision"] != "keep_user"
585
+ }
586
+ dependent_ids: set[str] = set()
587
+ for row in connection.execute(
588
+ "SELECT memory_id,metadata_json FROM records WHERE scope_id=?",
589
+ (job["scope_id"],),
590
+ ).fetchall():
591
+ metadata = _metadata(row)
592
+ if (
593
+ _text(metadata.get("subject_attribution_audit_id")) == audit_id
594
+ and _text(metadata.get("subject_attribution_decision"))
595
+ == "quarantine_with_parent"
596
+ and _text(metadata.get("subject_attribution_parent_memory_id"))
597
+ in quarantine_ids
598
+ ):
599
+ dependent_ids.add(str(row["memory_id"]))
600
+ return {
601
+ "changed_memory_ids": sorted(candidate_ids | dependent_ids),
602
+ "quarantined_memory_ids": sorted(quarantine_ids),
603
+ "cascaded_quarantined_memory_ids": sorted(dependent_ids),
604
+ "kept_memory_ids": sorted(candidate_ids - quarantine_ids),
605
+ }
606
+
607
+
608
+ def _cost(metadata: Mapping[str, Any]) -> float:
609
+ usage = metadata.get("usage")
610
+ usage = usage if isinstance(usage, Mapping) else metadata
611
+ prompt = int(usage.get("prompt_tokens", 0) or 0)
612
+ completion = int(usage.get("completion_tokens", 0) or 0)
613
+ hit = int(usage.get("prompt_cache_hit_tokens", 0) or 0)
614
+ miss = int(usage.get("prompt_cache_miss_tokens", prompt - hit) or 0)
615
+ prompt_rate = float(os.getenv("TMCRA_DEEPSEEK_PRO_PROMPT_COST_PER_MILLION", "3"))
616
+ completion_rate = float(os.getenv("TMCRA_DEEPSEEK_PRO_COMPLETION_COST_PER_MILLION", "6"))
617
+ cache_rate = float(os.getenv("TMCRA_DEEPSEEK_PRO_CACHE_COST_PER_MILLION", "0.025"))
618
+ return (miss * prompt_rate + hit * cache_rate + completion * completion_rate) / 1_000_000
619
+
620
+
621
+ def execute_job(
622
+ database: Path,
623
+ job: Mapping[str, Any],
624
+ client: AttributionClient,
625
+ ) -> dict[str, Any]:
626
+ audit_id = "saa_" + job["request_sha256"][:24]
627
+ request_json = _json(job["payload"])
628
+ with closing(sqlite3.connect(database, timeout=30)) as connection:
629
+ connection.row_factory = sqlite3.Row
630
+ _initialize(connection)
631
+ existing = connection.execute(
632
+ "SELECT status,decisions_json,call_metadata_json FROM "
633
+ "v4_subject_attribution_audits WHERE audit_id=?",
634
+ (audit_id,),
635
+ ).fetchone()
636
+ if existing is not None:
637
+ if str(existing["status"]) != "completed":
638
+ raise AttributionError(
639
+ f"{audit_id}: prior attribution call is not safely reusable"
640
+ )
641
+ decisions = json.loads(str(existing["decisions_json"]))
642
+ applied = _reused_application_state(
643
+ connection, job, audit_id, decisions
644
+ )
645
+ return {
646
+ "audit_id": audit_id,
647
+ "status": "reused",
648
+ "physical_api_calls": 0,
649
+ "decisions": decisions,
650
+ "call_metadata": json.loads(str(existing["call_metadata_json"])),
651
+ **applied,
652
+ }
653
+ now = _now()
654
+ connection.execute(
655
+ "INSERT INTO v4_subject_attribution_audits VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
656
+ (
657
+ audit_id,
658
+ job["scope_id"],
659
+ job["message_id"],
660
+ PROMPT_VERSION,
661
+ MODEL,
662
+ request_json,
663
+ job["request_sha256"],
664
+ "api_started",
665
+ "",
666
+ "",
667
+ "{}",
668
+ "[]",
669
+ "",
670
+ now,
671
+ now,
672
+ ),
673
+ )
674
+ connection.commit()
675
+
676
+ try:
677
+ content, call_metadata = client.complete(job["payload"])
678
+ except Exception as exc:
679
+ with closing(sqlite3.connect(database, timeout=30)) as connection:
680
+ connection.execute(
681
+ "UPDATE v4_subject_attribution_audits SET status='failed',error=?,"
682
+ "call_metadata_json=?,updated_at=? "
683
+ "WHERE audit_id=? AND status='api_started'",
684
+ (
685
+ str(exc),
686
+ _json(dict(getattr(exc, "metadata", {}) or {})),
687
+ _now(),
688
+ audit_id,
689
+ ),
690
+ )
691
+ connection.commit()
692
+ raise
693
+
694
+ with closing(sqlite3.connect(database, timeout=30)) as connection:
695
+ received = connection.execute(
696
+ "UPDATE v4_subject_attribution_audits SET status='response_received',"
697
+ "response_json=?,response_sha256=?,call_metadata_json=?,updated_at=? "
698
+ "WHERE audit_id=? AND status='api_started'",
699
+ (
700
+ content,
701
+ hashlib.sha256(content.encode("utf-8")).hexdigest(),
702
+ _json(dict(call_metadata)),
703
+ _now(),
704
+ audit_id,
705
+ ),
706
+ )
707
+ if received.rowcount != 1:
708
+ raise AttributionError("attribution journal changed before response save")
709
+ connection.commit()
710
+ try:
711
+ raw = json.loads(content)
712
+ decisions = validate_decisions(raw, job["payload"])
713
+ except (json.JSONDecodeError, AttributionError) as exc:
714
+ error = (
715
+ "attribution response content is not JSON"
716
+ if isinstance(exc, json.JSONDecodeError)
717
+ else str(exc)
718
+ )
719
+ with closing(sqlite3.connect(database, timeout=30)) as connection:
720
+ connection.execute(
721
+ "UPDATE v4_subject_attribution_audits SET status='failed',error=?,"
722
+ "updated_at=? WHERE audit_id=? AND status='response_received'",
723
+ (error, _now(), audit_id),
724
+ )
725
+ connection.commit()
726
+ raise AttributionError(error) from exc
727
+ with closing(sqlite3.connect(database, timeout=30)) as connection:
728
+ connection.row_factory = sqlite3.Row
729
+ connection.execute("BEGIN IMMEDIATE")
730
+ applied = _apply_decisions(connection, job, audit_id, decisions)
731
+ superseded = connection.execute(
732
+ "UPDATE v4_subject_attribution_audits SET status='superseded',error=?,"
733
+ "updated_at=? WHERE scope_id=? AND message_id=? "
734
+ "AND status IN ('completed','superseded') "
735
+ "AND audit_id<>?",
736
+ (
737
+ f"superseded_by:{audit_id}",
738
+ _now(),
739
+ job["scope_id"],
740
+ job["message_id"],
741
+ audit_id,
742
+ ),
743
+ ).rowcount
744
+ updated = connection.execute(
745
+ "UPDATE v4_subject_attribution_audits SET status='completed',response_json=?,"
746
+ "response_sha256=?,call_metadata_json=?,decisions_json=?,updated_at=? "
747
+ "WHERE audit_id=? AND status='response_received'",
748
+ (
749
+ content,
750
+ hashlib.sha256(content.encode("utf-8")).hexdigest(),
751
+ _json(dict(call_metadata)),
752
+ _json(decisions),
753
+ _now(),
754
+ audit_id,
755
+ ),
756
+ )
757
+ if updated.rowcount != 1:
758
+ raise AttributionError("attribution journal changed before commit")
759
+ connection.commit()
760
+ return {
761
+ "audit_id": audit_id,
762
+ "status": "completed",
763
+ "physical_api_calls": 1,
764
+ "estimated_cost_cny": _cost(call_metadata),
765
+ "decisions": decisions,
766
+ "call_metadata": dict(call_metadata),
767
+ "superseded_audit_count": int(superseded),
768
+ **applied,
769
+ }
770
+
771
+
772
+ def _worker_databases(run_dir: Path, workers: Sequence[str]) -> list[tuple[str, Path, str]]:
773
+ selected = set(workers)
774
+ output: list[tuple[str, Path, str]] = []
775
+ for database in sorted((run_dir / "writer").glob("worker_*/native_memory.sqlite3")):
776
+ worker = database.parent.name
777
+ if selected and worker not in selected:
778
+ continue
779
+ with closing(sqlite3.connect(database)) as connection:
780
+ row = connection.execute(
781
+ "SELECT scope_id FROM v4_source_journal LIMIT 1"
782
+ ).fetchone()
783
+ if row is None or not _text(row[0]):
784
+ raise AttributionError(f"{worker}: scope is missing")
785
+ output.append((worker, database, _text(row[0])))
786
+ if selected - {item[0] for item in output}:
787
+ raise AttributionError(
788
+ "unknown workers: " + ",".join(sorted(selected - {item[0] for item in output}))
789
+ )
790
+ return output
791
+
792
+
793
+ def main() -> None:
794
+ parser = argparse.ArgumentParser(description=__doc__)
795
+ parser.add_argument("--run-dir", type=Path, required=True)
796
+ parser.add_argument("--workers", default="")
797
+ parser.add_argument("--apply", action="store_true")
798
+ parser.add_argument("--output", type=Path, required=True)
799
+ args = parser.parse_args()
800
+ run_dir = args.run_dir.resolve()
801
+ workers = [item.strip() for item in args.workers.split(",") if item.strip()]
802
+ scanned: list[dict[str, Any]] = []
803
+ worker_by_database: dict[str, str] = {}
804
+ for worker, database, scope_id in _worker_databases(run_dir, workers):
805
+ worker_by_database[str(database)] = worker
806
+ for job in scan_database(database, scope_id):
807
+ scanned.append({"worker": worker, **job})
808
+ results: list[dict[str, Any]] = []
809
+ if args.apply and scanned:
810
+ client = DeepSeekProAttributionClient()
811
+ for job in scanned:
812
+ results.append(
813
+ {
814
+ "worker": job["worker"],
815
+ "message_id": job["message_id"],
816
+ **execute_job(Path(job["database"]), job, client),
817
+ }
818
+ )
819
+ report = {
820
+ "schema_version": "tmcra.v4.subject-attribution-report.1",
821
+ "status": "complete",
822
+ "mode": "apply" if args.apply else "scan_only",
823
+ "prompt_version": PROMPT_VERSION,
824
+ "model": MODEL,
825
+ "run_dir": str(run_dir),
826
+ "scanned_worker_count": len(_worker_databases(run_dir, workers)),
827
+ "routed_message_count": len(scanned),
828
+ "routed_candidate_count": sum(
829
+ len(job["payload"]["candidates"]) for job in scanned
830
+ ),
831
+ "physical_api_calls": sum(
832
+ int(item.get("physical_api_calls", 0) or 0) for item in results
833
+ ),
834
+ "estimated_cost_cny": round(
835
+ sum(float(item.get("estimated_cost_cny", 0.0) or 0.0) for item in results),
836
+ 8,
837
+ ),
838
+ "decision_quarantined_count": sum(
839
+ len(item.get("quarantined_memory_ids", [])) for item in results
840
+ ),
841
+ "cascaded_quarantined_count": sum(
842
+ len(item.get("cascaded_quarantined_memory_ids", []))
843
+ for item in results
844
+ ),
845
+ "quarantined_count": sum(
846
+ len(item.get("quarantined_memory_ids", []))
847
+ + len(item.get("cascaded_quarantined_memory_ids", []))
848
+ for item in results
849
+ ),
850
+ "routed": [
851
+ {
852
+ "worker": job["worker"],
853
+ "database": job["database"],
854
+ "scope_id": job["scope_id"],
855
+ "message_id": job["message_id"],
856
+ "session_index": job["session_index"],
857
+ "message_index": job["message_index"],
858
+ "source_turn_index": job["source_turn_index"],
859
+ "route_reasons": job["route_reasons"],
860
+ "request_sha256": job["request_sha256"],
861
+ "candidate_count": len(job["payload"]["candidates"]),
862
+ "candidates": job["review_candidates"],
863
+ }
864
+ for job in scanned
865
+ ],
866
+ "results": results,
867
+ }
868
+ args.output.parent.mkdir(parents=True, exist_ok=True)
869
+ args.output.write_text(
870
+ json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
871
+ encoding="utf-8",
872
+ )
873
+ print(json.dumps(report, ensure_ascii=False, sort_keys=True))
874
+
875
+
876
+ if __name__ == "__main__":
877
+ main()
benchmarks/longmemeval/src/tmcra_benchmark/legacy/prepare_tmcra_v4_e2e_data.py ADDED
@@ -0,0 +1,261 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import json
7
+ from collections import Counter
8
+ from pathlib import Path
9
+ from typing import Any, Mapping, Sequence
10
+
11
+
12
+ WRITER_FIELDS = (
13
+ "question_id",
14
+ "haystack_sessions",
15
+ "haystack_session_ids",
16
+ "haystack_dates",
17
+ )
18
+ FORBIDDEN_WRITER_FIELDS = frozenset(
19
+ {
20
+ "question",
21
+ "question_date",
22
+ "question_type",
23
+ "answer",
24
+ "gold_answer",
25
+ "answer_session_ids",
26
+ "labels",
27
+ "supervision",
28
+ }
29
+ )
30
+
31
+
32
+ class PreparationError(RuntimeError):
33
+ pass
34
+
35
+
36
+ def clean_text(value: Any) -> str:
37
+ return " ".join(str(value or "").split())
38
+
39
+
40
+ def sha256_file(path: Path) -> str:
41
+ digest = hashlib.sha256()
42
+ with path.open("rb") as handle:
43
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
44
+ digest.update(chunk)
45
+ return digest.hexdigest()
46
+
47
+
48
+ def write_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
49
+ with path.open("x", encoding="utf-8") as handle:
50
+ for row in rows:
51
+ handle.write(
52
+ json.dumps(dict(row), ensure_ascii=False, sort_keys=True) + "\n"
53
+ )
54
+
55
+
56
+ def validate_source_row(row: Mapping[str, Any]) -> None:
57
+ qid = clean_text(row.get("question_id"))
58
+ sessions = row.get("haystack_sessions")
59
+ session_ids = row.get("haystack_session_ids")
60
+ dates = row.get("haystack_dates")
61
+ if not qid or not isinstance(sessions, list) or not sessions:
62
+ raise PreparationError(f"{qid or '<missing-qid>'}: missing history sessions")
63
+ if not isinstance(session_ids, list) or not isinstance(dates, list):
64
+ raise PreparationError(f"{qid}: session IDs and dates must be arrays")
65
+ if len(sessions) != len(session_ids) or len(sessions) != len(dates):
66
+ raise PreparationError(f"{qid}: history/session/date lengths differ")
67
+ for session_index, session in enumerate(sessions):
68
+ if not isinstance(session, list) or not session:
69
+ raise PreparationError(f"{qid}: session {session_index} is empty")
70
+ for message_index, message in enumerate(session):
71
+ if not isinstance(message, Mapping):
72
+ raise PreparationError(
73
+ f"{qid}: session {session_index} message {message_index} is not an object"
74
+ )
75
+ role = clean_text(message.get("role")).lower()
76
+ if role not in {"user", "assistant", "system", "tool"}:
77
+ raise PreparationError(
78
+ f"{qid}: session {session_index} message {message_index} is invalid"
79
+ )
80
+
81
+
82
+ def prepare(
83
+ *, data_path: Path, qid_path: Path, out_dir: Path
84
+ ) -> dict[str, Any]:
85
+ qids = [
86
+ clean_text(value)
87
+ for value in qid_path.read_text(encoding="utf-8").splitlines()
88
+ if clean_text(value)
89
+ ]
90
+ if not qids or len(qids) != len(set(qids)):
91
+ raise PreparationError("qid list must be non-empty and unique")
92
+ source = json.loads(data_path.read_text(encoding="utf-8"))
93
+ if not isinstance(source, list):
94
+ raise PreparationError("source dataset must be a JSON array")
95
+ source_rows: dict[str, Mapping[str, Any]] = {}
96
+ for raw_row in source:
97
+ if not isinstance(raw_row, Mapping):
98
+ raise PreparationError("source dataset contains a non-object row")
99
+ validate_source_row(raw_row)
100
+ source_qid = clean_text(raw_row.get("question_id"))
101
+ if source_qid in source_rows:
102
+ raise PreparationError(f"duplicate source question_id: {source_qid}")
103
+ source_rows[source_qid] = raw_row
104
+ missing = [qid for qid in qids if qid not in source_rows]
105
+ if missing:
106
+ raise PreparationError(f"missing qids in source data: {missing}")
107
+
108
+ out_dir.mkdir(parents=True, exist_ok=False)
109
+ scope_manifest: list[dict[str, Any]] = []
110
+ query_manifest: list[dict[str, Any]] = []
111
+ evaluation_refs: list[dict[str, Any]] = []
112
+ workers: list[dict[str, Any]] = []
113
+ writer_rows: list[dict[str, Any]] = []
114
+ total_input_messages = 0
115
+ total_nonempty_messages = 0
116
+ total_empty_messages = 0
117
+ total_duplicate_session_id_occurrences = 0
118
+ duplicate_session_id_qids: list[str] = []
119
+ for worker_index, qid in enumerate(qids):
120
+ source_row = source_rows[qid]
121
+ session_id_counts = Counter(
122
+ clean_text(value)
123
+ for value in source_row.get("haystack_session_ids") or []
124
+ if clean_text(value)
125
+ )
126
+ duplicate_session_ids = {
127
+ session_id: count
128
+ for session_id, count in sorted(session_id_counts.items())
129
+ if count > 1
130
+ }
131
+ duplicate_occurrences = sum(
132
+ count - 1 for count in duplicate_session_ids.values()
133
+ )
134
+ if duplicate_occurrences:
135
+ duplicate_session_id_qids.append(qid)
136
+ total_duplicate_session_id_occurrences += duplicate_occurrences
137
+ writer_row = {field: source_row.get(field) for field in WRITER_FIELDS}
138
+ if FORBIDDEN_WRITER_FIELDS.intersection(writer_row):
139
+ raise AssertionError("writer sanitizer retained a forbidden field")
140
+ worker_dir = out_dir / "writer" / f"worker_{worker_index:03d}"
141
+ worker_dir.mkdir(parents=True, exist_ok=False)
142
+ writer_input = worker_dir / "input.json"
143
+ writer_input.write_text(
144
+ json.dumps([writer_row], ensure_ascii=False) + "\n", encoding="utf-8"
145
+ )
146
+ writer_rows.append(writer_row)
147
+ database = worker_dir / "native_memory.sqlite3"
148
+ scope_id = f"tmcra_v4:{qid}"
149
+ index_path = out_dir / "indexes" / f"{qid}.pt"
150
+ scope_manifest.append(
151
+ {
152
+ "question_id": qid,
153
+ "db_path": str(database),
154
+ "scope_id": scope_id,
155
+ "index_path": str(index_path),
156
+ }
157
+ )
158
+ query_manifest.append(
159
+ {
160
+ "question_id": qid,
161
+ "question": clean_text(source_row.get("question")),
162
+ "question_date": clean_text(source_row.get("question_date")),
163
+ "question_type": clean_text(source_row.get("question_type")),
164
+ "db_path": str(database),
165
+ "scope_id": scope_id,
166
+ "index_path": str(index_path),
167
+ }
168
+ )
169
+ evaluation_refs.append(
170
+ {
171
+ "question_id": qid,
172
+ "answer": source_row.get("answer"),
173
+ "answer_session_ids": source_row.get("answer_session_ids"),
174
+ }
175
+ )
176
+ workers.append(
177
+ {
178
+ "worker_index": worker_index,
179
+ "question_id": qid,
180
+ "worker_dir": str(worker_dir),
181
+ "input": str(writer_input),
182
+ "input_sha256": sha256_file(writer_input),
183
+ "scope_id": scope_id,
184
+ "session_count": len(source_row.get("haystack_sessions") or []),
185
+ "message_count": sum(
186
+ len(session)
187
+ for session in source_row.get("haystack_sessions") or []
188
+ ),
189
+ "nonempty_message_count": sum(
190
+ 1
191
+ for session in source_row.get("haystack_sessions") or []
192
+ for message in session
193
+ if str(message.get("content") or "").strip()
194
+ ),
195
+ "empty_message_count": sum(
196
+ 1
197
+ for session in source_row.get("haystack_sessions") or []
198
+ for message in session
199
+ if not str(message.get("content") or "").strip()
200
+ ),
201
+ "duplicate_session_ids": duplicate_session_ids,
202
+ "duplicate_session_id_occurrence_count": duplicate_occurrences,
203
+ }
204
+ )
205
+ total_input_messages += workers[-1]["message_count"]
206
+ total_nonempty_messages += workers[-1]["nonempty_message_count"]
207
+ total_empty_messages += workers[-1]["empty_message_count"]
208
+ (out_dir / "indexes").mkdir(parents=True, exist_ok=False)
209
+ (out_dir / "writer_input.json").write_text(
210
+ json.dumps(writer_rows, ensure_ascii=False) + "\n", encoding="utf-8"
211
+ )
212
+ write_jsonl(out_dir / "scope_manifest.jsonl", scope_manifest)
213
+ write_jsonl(out_dir / "query_manifest.jsonl", query_manifest)
214
+ evaluation_dir = out_dir / "evaluation_only"
215
+ evaluation_dir.mkdir(parents=True, exist_ok=False)
216
+ write_jsonl(evaluation_dir / "references.jsonl", evaluation_refs)
217
+ (out_dir / "qids.txt").write_text("\n".join(qids) + "\n", encoding="utf-8")
218
+ manifest = {
219
+ "schema_version": "tmcra.v4.e2e-input.1",
220
+ "status": "prepared",
221
+ "source_data": str(data_path.resolve()),
222
+ "source_data_sha256": sha256_file(data_path),
223
+ "row_count": len(qids),
224
+ "input_message_count": total_input_messages,
225
+ "nonempty_message_count": total_nonempty_messages,
226
+ "empty_message_count": total_empty_messages,
227
+ "duplicate_session_id_occurrence_count": total_duplicate_session_id_occurrences,
228
+ "duplicate_session_id_qids": duplicate_session_id_qids,
229
+ "qids": qids,
230
+ "writer_fields": list(WRITER_FIELDS),
231
+ "forbidden_writer_fields": sorted(FORBIDDEN_WRITER_FIELDS),
232
+ "writer_inputs_have_query_or_evaluation_fields": False,
233
+ "gold_isolation_dir": str(evaluation_dir),
234
+ "workers": workers,
235
+ "combined_writer_input": str(out_dir / "writer_input.json"),
236
+ "scope_manifest": str(out_dir / "scope_manifest.jsonl"),
237
+ "query_manifest": str(out_dir / "query_manifest.jsonl"),
238
+ }
239
+ (out_dir / "input_manifest.json").write_text(
240
+ json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
241
+ )
242
+ return manifest
243
+
244
+
245
+ def main() -> int:
246
+ parser = argparse.ArgumentParser(description="Prepare gold-isolated TMCRA V4 E2E inputs")
247
+ parser.add_argument("--data", required=True, type=Path)
248
+ parser.add_argument("--qid-list", required=True, type=Path)
249
+ parser.add_argument("--out-dir", required=True, type=Path)
250
+ args = parser.parse_args()
251
+ report = prepare(
252
+ data_path=args.data.resolve(),
253
+ qid_path=args.qid_list.resolve(),
254
+ out_dir=args.out_dir.resolve(),
255
+ )
256
+ print(json.dumps(report, indent=2, sort_keys=True))
257
+ return 0
258
+
259
+
260
+ if __name__ == "__main__":
261
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_build.py ADDED
@@ -0,0 +1,984 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import json
6
+ import os
7
+ import sqlite3
8
+ import subprocess
9
+ import sys
10
+ import time
11
+ from contextlib import closing
12
+ from concurrent.futures import ThreadPoolExecutor, as_completed
13
+ from datetime import datetime, timezone
14
+ from pathlib import Path
15
+ from typing import Any, Mapping, Sequence
16
+
17
+ from prepare_tmcra_v4_e2e_data import prepare
18
+ from tmcra_v4_cost_report import build_report, collect_calls
19
+ from tmcra_v4_slow_graph import (
20
+ PROCESS_LOSS_INTERRUPTION_ERROR,
21
+ SlowGraphStore,
22
+ load_graph_schema,
23
+ )
24
+
25
+
26
+ BASE = Path(__file__).resolve().parent
27
+ PACKAGE_ROOT = BASE.parent
28
+ PROJECT_ROOT = BASE.parents[2]
29
+ DEFAULT_DATA = PROJECT_ROOT / "data" / "longmemeval_s_cleaned.json"
30
+ DEFAULT_REPO = PACKAGE_ROOT / "vendor" / "tmcra_integrated"
31
+ DEFAULT_WRITER_ENV = PROJECT_ROOT / "configs" / "writer.env"
32
+ DEFAULT_EMBEDDING = PROJECT_ROOT / "models" / "BAAI" / "bge-m3"
33
+ SUBJECT_ATTRIBUTION_PROMPT_VERSION = "tmcra-v4-subject-attribution-2026-07-14.3"
34
+
35
+
36
+ class BuildError(RuntimeError):
37
+ pass
38
+
39
+
40
+ def _now() -> str:
41
+ return datetime.now(timezone.utc).isoformat(timespec="seconds")
42
+
43
+
44
+ def _stage(root: Path, name: str, **values: Any) -> None:
45
+ record = {"at": _now(), "stage": name, **values}
46
+ with (root / "build.log.jsonl").open("a", encoding="utf-8") as handle:
47
+ handle.write(json.dumps(record, sort_keys=True) + "\n")
48
+ print(json.dumps(record, sort_keys=True), flush=True)
49
+
50
+
51
+ def _load_shell_environment(path: Path) -> dict[str, str]:
52
+ if not path.is_file():
53
+ raise BuildError(f"writer environment file is missing: {path}")
54
+ command = 'set -a; source "$1"; env -0'
55
+ result = subprocess.run(
56
+ ["bash", "-c", command, "tmcra-v4-env", str(path)],
57
+ check=True,
58
+ stdout=subprocess.PIPE,
59
+ stderr=subprocess.PIPE,
60
+ )
61
+ loaded: dict[str, str] = {}
62
+ for entry in result.stdout.decode("utf-8").split("\0"):
63
+ if "=" in entry:
64
+ key, value = entry.split("=", 1)
65
+ loaded[key] = value
66
+ return loaded
67
+
68
+
69
+ def _key_pool(environment: Mapping[str, str]) -> list[str]:
70
+ raw = environment.get("TMCRA_DEEPSEEK_WRITER_KEY_POOL", "")
71
+ keys = [item.strip() for item in raw.split(",") if item.strip()]
72
+ if not keys or len(keys) != len(set(keys)):
73
+ raise BuildError("DeepSeek writer key pool must be non-empty and unique")
74
+ configured_count = environment.get("TMCRA_DEEPSEEK_WRITER_KEY_POOL_COUNT", "")
75
+ if configured_count and int(configured_count) != len(keys):
76
+ raise BuildError("DeepSeek writer key pool count does not match actual keys")
77
+ return keys
78
+
79
+
80
+ def _rotated(keys: Sequence[str], worker_index: int) -> str:
81
+ offset = worker_index % len(keys)
82
+ return ",".join([*keys[offset:], *keys[:offset]])
83
+
84
+
85
+ def _worker_environment(base: Mapping[str, str], keys: Sequence[str], worker_index: int) -> dict[str, str]:
86
+ environment = dict(base)
87
+ pool = _rotated(keys, worker_index)
88
+ base_url = environment.get("TMCRA_DEEPSEEK_WRITER_BASE_URL") or environment.get("TMCRA_WRITER_BASE_URL") or "https://api.deepseek.com/v1"
89
+ max_tokens = environment.get("TMCRA_WRITER_MAX_TOKENS", "16384")
90
+ environment.update(
91
+ {
92
+ "TMCRA_WRITER_MAX_TOKENS": max_tokens,
93
+ "TMCRA_WRITER_BASE_URL": base_url,
94
+ "TMCRA_WRITER_MODEL": "deepseek-v4-flash",
95
+ "TMCRA_WRITER_API_KEY_POOL": pool,
96
+ "TMCRA_DEEPSEEK_FLASH_BASE_URL": base_url,
97
+ "TMCRA_DEEPSEEK_FLASH_KEY_POOL": pool,
98
+ "TMCRA_DEEPSEEK_FLASH_MAX_TOKENS": max_tokens,
99
+ "TMCRA_DEEPSEEK_FLASH_PROMPT_COST_PER_MILLION": "1",
100
+ "TMCRA_DEEPSEEK_FLASH_COMPLETION_COST_PER_MILLION": "2",
101
+ "TMCRA_DEEPSEEK_FLASH_CACHE_COST_PER_MILLION": "0.02",
102
+ "TMCRA_DEEPSEEK_PRO_BASE_URL": base_url,
103
+ "TMCRA_DEEPSEEK_PRO_KEY_POOL": pool,
104
+ "TMCRA_DEEPSEEK_PRO_MAX_TOKENS": max_tokens,
105
+ "TMCRA_DEEPSEEK_PRO_PROMPT_COST_PER_MILLION": "3",
106
+ "TMCRA_DEEPSEEK_PRO_COMPLETION_COST_PER_MILLION": "6",
107
+ "TMCRA_DEEPSEEK_PRO_CACHE_COST_PER_MILLION": "0.025",
108
+ }
109
+ )
110
+ return environment
111
+
112
+
113
+ def _run(command: Sequence[str], log_path: Path, environment: Mapping[str, str]) -> None:
114
+ log_path.parent.mkdir(parents=True, exist_ok=True)
115
+ with log_path.open("x", encoding="utf-8") as log:
116
+ log.write(json.dumps({"command": list(command)}, sort_keys=True) + "\n")
117
+ log.flush()
118
+ subprocess.run(
119
+ list(command),
120
+ check=True,
121
+ stdout=log,
122
+ stderr=subprocess.STDOUT,
123
+ env=dict(environment),
124
+ )
125
+
126
+
127
+ def _resume_log(directory: Path, stem: str) -> Path:
128
+ return directory / f"{stem}.resume.{time.time_ns()}.log"
129
+
130
+
131
+ def _load_resume_manifest(out_dir: Path) -> dict[str, Any]:
132
+ manifest_path = out_dir / "input_manifest.json"
133
+ if not manifest_path.is_file():
134
+ raise BuildError("resume run has no input_manifest.json")
135
+ try:
136
+ manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
137
+ except (OSError, json.JSONDecodeError) as exc:
138
+ raise BuildError("resume input manifest is unreadable") from exc
139
+ if not isinstance(manifest, Mapping):
140
+ raise BuildError("resume input manifest must be an object")
141
+ workers = manifest.get("workers")
142
+ if (
143
+ manifest.get("status") != "prepared"
144
+ or not isinstance(workers, list)
145
+ or not workers
146
+ or int(manifest.get("row_count", 0) or 0) != len(workers)
147
+ ):
148
+ raise BuildError("resume input manifest is incomplete or stale")
149
+ return dict(manifest)
150
+
151
+
152
+ def _verify_resume_writer(worker: Mapping[str, Any]) -> None:
153
+ worker_dir = Path(str(worker["worker_dir"])).resolve()
154
+ database = worker_dir / "native_memory.sqlite3"
155
+ report_path = worker_dir / "product_writer_report.json"
156
+ audit_path = worker_dir / "writer_chain_audit.json"
157
+ if not all(path.is_file() for path in (database, report_path, audit_path)):
158
+ raise BuildError(f"resume writer artifacts are incomplete: {worker_dir}")
159
+ report = json.loads(report_path.read_text(encoding="utf-8"))
160
+ audit = json.loads(audit_path.read_text(encoding="utf-8"))
161
+ if report.get("completed") is not True or audit.get("passed") is not True:
162
+ raise BuildError(f"resume writer audit is not complete: {worker_dir}")
163
+ with closing(sqlite3.connect(database)) as con:
164
+ statuses = dict(
165
+ con.execute("SELECT status,count(*) FROM v4_batch_journal GROUP BY status")
166
+ )
167
+ if statuses != {"committed": int(report.get("batches", -1))}:
168
+ raise BuildError(
169
+ f"resume writer journal is not fully committed: {worker_dir}: {statuses}"
170
+ )
171
+
172
+
173
+ def _interrupted_writer_calls(worker: Mapping[str, Any]) -> list[dict[str, str]]:
174
+ worker_dir = Path(str(worker["worker_dir"])).resolve()
175
+ database = worker_dir / "native_memory.sqlite3"
176
+ if not database.is_file():
177
+ return []
178
+ output: list[dict[str, str]] = []
179
+ with closing(sqlite3.connect(database)) as con:
180
+ con.row_factory = sqlite3.Row
181
+ for row in con.execute(
182
+ "SELECT batch_id FROM v4_batch_journal WHERE status='api_started' ORDER BY batch_index"
183
+ ):
184
+ output.append(
185
+ {
186
+ "worker": str(worker["question_id"]),
187
+ "stage": "batch_flash",
188
+ "call_key": f"flash:{row['batch_id']}",
189
+ }
190
+ )
191
+ if con.execute(
192
+ "SELECT 1 FROM sqlite_master WHERE type='table' AND name='v4_reconciliation_jobs'"
193
+ ).fetchone():
194
+ for row in con.execute(
195
+ "SELECT job_id FROM v4_reconciliation_jobs WHERE status='pro_started' ORDER BY created_at,job_id"
196
+ ):
197
+ output.append(
198
+ {
199
+ "worker": str(worker["question_id"]),
200
+ "stage": "reconciliation_pro",
201
+ "call_key": f"pro:{row['job_id']}",
202
+ }
203
+ )
204
+ return output
205
+
206
+
207
+ def _interrupted_slow_calls(worker: Mapping[str, Any]) -> list[dict[str, Any]]:
208
+ worker_dir = Path(str(worker["worker_dir"])).resolve()
209
+ database = worker_dir / "native_memory.sqlite3"
210
+ if not database.is_file():
211
+ return []
212
+ output: list[dict[str, Any]] = []
213
+ with closing(sqlite3.connect(database)) as con:
214
+ con.row_factory = sqlite3.Row
215
+ tables = {
216
+ str(row[0])
217
+ for row in con.execute(
218
+ "SELECT name FROM sqlite_master WHERE type='table'"
219
+ )
220
+ }
221
+ if not {"slow_graph_jobs", "slow_graph_attempts"} <= tables:
222
+ return []
223
+ recovered_attempts: set[str] = set()
224
+ if "slow_graph_process_loss_recoveries" in tables:
225
+ recovered_attempts = {
226
+ str(row[0])
227
+ for row in con.execute(
228
+ "SELECT attempt_id FROM slow_graph_process_loss_recoveries"
229
+ )
230
+ }
231
+ rows = con.execute(
232
+ "SELECT j.job_id,j.scope_id,j.region_key,j.status,j.claim_owner,"
233
+ "j.lease_expires_at,a.attempt_id,a.status AS attempt_status,"
234
+ "a.created_at FROM slow_graph_jobs j JOIN slow_graph_attempts a "
235
+ "ON a.job_id=j.job_id WHERE ("
236
+ "(j.status='pending' AND j.claim_token IS NOT NULL "
237
+ "AND a.status='started' AND a.claim_token=j.claim_token "
238
+ "AND a.claim_owner=j.claim_owner) OR "
239
+ "(j.status='failed' AND j.claim_token IS NULL AND j.last_error=? "
240
+ "AND a.status='expired' AND a.error=?)) "
241
+ "ORDER BY a.created_at,a.attempt_id",
242
+ (PROCESS_LOSS_INTERRUPTION_ERROR, PROCESS_LOSS_INTERRUPTION_ERROR),
243
+ ).fetchall()
244
+ for row in rows:
245
+ if str(row["attempt_id"]) in recovered_attempts:
246
+ continue
247
+ output.append(
248
+ {
249
+ "worker": str(worker.get("question_id") or worker_dir.name),
250
+ **dict(row),
251
+ }
252
+ )
253
+ return output
254
+
255
+
256
+ def _recover_interrupted_slow_calls(
257
+ worker: Mapping[str, Any], *, repo: Path
258
+ ) -> list[dict[str, Any]]:
259
+ worker_dir = Path(str(worker["worker_dir"])).resolve()
260
+ database = worker_dir / "native_memory.sqlite3"
261
+ store = SlowGraphStore(database, schema=load_graph_schema(repo))
262
+ reviewed = store.interrupted_process_loss_attempts()
263
+ reports = [
264
+ store.recover_interrupted_process_loss(
265
+ str(item["job_id"]),
266
+ expected_attempt_id=str(item["attempt_id"]),
267
+ )
268
+ for item in reviewed
269
+ ]
270
+ if reports:
271
+ report_path = _resume_log(worker_dir, "slow_process_loss_recovery").with_suffix(
272
+ ".json"
273
+ )
274
+ report_path.write_text(
275
+ json.dumps(
276
+ {
277
+ "schema_version": "tmcra.v4.slow-process-loss-worker-report.1",
278
+ "worker": str(worker.get("question_id") or worker_dir.name),
279
+ "physical_api_calls_during_recovery": 0,
280
+ "recoveries": reports,
281
+ },
282
+ indent=2,
283
+ sort_keys=True,
284
+ )
285
+ + "\n",
286
+ encoding="utf-8",
287
+ )
288
+ return reports
289
+
290
+
291
+ def _failed_slow_jobs(database: Path) -> list[dict[str, str]]:
292
+ with closing(sqlite3.connect(database)) as con:
293
+ con.row_factory = sqlite3.Row
294
+ if con.execute(
295
+ "SELECT 1 FROM sqlite_master WHERE type='table' AND name='slow_graph_jobs'"
296
+ ).fetchone() is None:
297
+ return []
298
+ return [
299
+ {
300
+ "job_id": str(row["job_id"]),
301
+ "region_key": str(row["region_key"]),
302
+ "last_error": str(row["last_error"]),
303
+ }
304
+ for row in con.execute(
305
+ "SELECT job_id,region_key,last_error FROM slow_graph_jobs "
306
+ "WHERE status IN ('failed','retryable') ORDER BY created_at,job_id"
307
+ )
308
+ ]
309
+
310
+
311
+ def _slow_worker_resume(
312
+ worker: Mapping[str, Any],
313
+ *,
314
+ repo: Path,
315
+ environment: Mapping[str, str],
316
+ enqueue: bool = True,
317
+ recover_interrupted_slow_calls: bool = False,
318
+ ) -> None:
319
+ worker_dir = Path(str(worker["worker_dir"])).resolve()
320
+ database = worker_dir / "native_memory.sqlite3"
321
+ scope_id = str(worker["scope_id"])
322
+ prefix = [
323
+ sys.executable,
324
+ str(BASE / "tmcra_v4_slow_graph.py"),
325
+ str(database),
326
+ "--repo",
327
+ str(repo),
328
+ ]
329
+ interrupted = _interrupted_slow_calls(worker)
330
+ if interrupted and not recover_interrupted_slow_calls:
331
+ raise BuildError(
332
+ "resume has Slow calls without durable responses; explicit process-loss "
333
+ "recovery flag is required: " + json.dumps(interrupted, sort_keys=True)
334
+ )
335
+ if interrupted:
336
+ recovered = _recover_interrupted_slow_calls(worker, repo=repo)
337
+ if len(recovered) != len(interrupted):
338
+ raise BuildError(
339
+ "Slow process-loss recovery count drifted after preflight"
340
+ )
341
+ if enqueue:
342
+ _run(
343
+ [*prefix, "enqueue", scope_id],
344
+ _resume_log(worker_dir, "slow_enqueue"),
345
+ environment,
346
+ )
347
+ failed = _failed_slow_jobs(database)
348
+ if failed:
349
+ raise BuildError(
350
+ "resume requires explicit revalidation of failed slow jobs: "
351
+ + json.dumps(failed, sort_keys=True)
352
+ )
353
+ _run([*prefix, "drain"], _resume_log(worker_dir, "slow_drain"), environment)
354
+ _run(
355
+ [*prefix, "audit", scope_id, "--require-promotion-coverage"],
356
+ _resume_log(worker_dir, "slow_audit"),
357
+ environment,
358
+ )
359
+
360
+
361
+ def _subject_attribution_stage(
362
+ out_dir: Path, environment: Mapping[str, str]
363
+ ) -> dict[str, Any]:
364
+ report_path = out_dir / "subject_attribution_report.json"
365
+ if not report_path.exists():
366
+ _run(
367
+ [
368
+ sys.executable,
369
+ str(BASE / "ops" / "audit_tmcra_v4_subject_attribution.py"),
370
+ "--run-dir",
371
+ str(out_dir),
372
+ "--apply",
373
+ "--output",
374
+ str(report_path),
375
+ ],
376
+ out_dir / "subject_attribution.log",
377
+ environment,
378
+ )
379
+ try:
380
+ report = json.loads(report_path.read_text(encoding="utf-8"))
381
+ except (OSError, json.JSONDecodeError) as exc:
382
+ raise BuildError("subject-attribution report is unreadable") from exc
383
+ if (
384
+ not isinstance(report, Mapping)
385
+ or report.get("status") != "complete"
386
+ or report.get("mode") != "apply"
387
+ or report.get("prompt_version") != SUBJECT_ATTRIBUTION_PROMPT_VERSION
388
+ or report.get("model") != "deepseek-v4-pro"
389
+ ):
390
+ raise BuildError("subject-attribution stage is incomplete or drifted")
391
+ _stage(
392
+ out_dir,
393
+ "subject_attribution_complete",
394
+ routed_messages=int(report.get("routed_message_count", 0) or 0),
395
+ quarantined=int(report.get("quarantined_count", 0) or 0),
396
+ physical_api_calls=int(report.get("physical_api_calls", 0) or 0),
397
+ )
398
+ return dict(report)
399
+
400
+
401
+ def _slow_process_loss_cost_uncertainty(
402
+ databases: Sequence[Path],
403
+ ) -> dict[str, int]:
404
+ recoveries = 0
405
+ potential_min = 0
406
+ potential_max = 0
407
+ for database in databases:
408
+ with closing(sqlite3.connect(database)) as con:
409
+ if con.execute(
410
+ "SELECT 1 FROM sqlite_master WHERE type='table' "
411
+ "AND name='slow_graph_process_loss_recoveries'"
412
+ ).fetchone() is None:
413
+ continue
414
+ row = con.execute(
415
+ "SELECT count(*),"
416
+ "coalesce(sum(potential_duplicate_physical_calls_min),0),"
417
+ "coalesce(sum(potential_duplicate_physical_calls_max),0) "
418
+ "FROM slow_graph_process_loss_recoveries"
419
+ ).fetchone()
420
+ recoveries += int(row[0])
421
+ potential_min += int(row[1])
422
+ potential_max += int(row[2])
423
+ return {
424
+ "unknown_external_call_outcomes": recoveries,
425
+ "potential_duplicate_physical_calls_min": potential_min,
426
+ "potential_duplicate_physical_calls_max": potential_max,
427
+ }
428
+
429
+
430
+ def _finalize_build(
431
+ *,
432
+ out_dir: Path,
433
+ workers: Sequence[Mapping[str, Any]],
434
+ writer_concurrency: int,
435
+ slow_concurrency: int,
436
+ recovered: bool,
437
+ ) -> dict[str, Any]:
438
+ databases = [
439
+ Path(str(worker["worker_dir"])) / "native_memory.sqlite3"
440
+ for worker in workers
441
+ ]
442
+ interrupted_call_logs = [
443
+ path
444
+ for worker in workers
445
+ if (
446
+ path := Path(str(worker["worker_dir"]))
447
+ / "product_writer_interrupted_calls.jsonl"
448
+ ).is_file()
449
+ ]
450
+ cost = build_report(collect_calls(interrupted_call_logs, databases))
451
+ process_loss = _slow_process_loss_cost_uncertainty(databases)
452
+ cost["slow_process_loss_uncertainty"] = process_loss
453
+ cost["cost_is_fully_observed"] = (
454
+ process_loss["unknown_external_call_outcomes"] == 0
455
+ )
456
+ (out_dir / "build_cost_report.json").write_text(
457
+ json.dumps(cost, indent=2, sort_keys=True) + "\n", encoding="utf-8"
458
+ )
459
+ quality = _writer_quality_report(workers)
460
+ (out_dir / "writer_quality_report.json").write_text(
461
+ json.dumps(quality, indent=2, sort_keys=True) + "\n", encoding="utf-8"
462
+ )
463
+ report = {
464
+ "schema_version": "tmcra.v4.build.1",
465
+ "status": "complete",
466
+ "row_count": len(workers),
467
+ "writer_concurrency": writer_concurrency,
468
+ "slow_concurrency": slow_concurrency,
469
+ "physical_api_call_count": cost["physical_call_count"],
470
+ "exact_cost_cny": cost["exact_cost_cny"],
471
+ "min_cost_cny": cost["min_cost_cny"],
472
+ "max_cost_cny": cost["max_cost_cny"],
473
+ "cost_is_fully_observed": cost["cost_is_fully_observed"],
474
+ "slow_process_loss_unknown_external_outcomes": process_loss[
475
+ "unknown_external_call_outcomes"
476
+ ],
477
+ "slow_process_loss_potential_duplicate_physical_calls_min": (
478
+ process_loss["potential_duplicate_physical_calls_min"]
479
+ ),
480
+ "slow_process_loss_potential_duplicate_physical_calls_max": (
481
+ process_loss["potential_duplicate_physical_calls_max"]
482
+ ),
483
+ "interrupted_calls_without_usage": sum(
484
+ int(item.get("calls_without_usage") or 0)
485
+ for item in cost["by_stage_model"]
486
+ if str(item.get("stage") or "").endswith("_interrupted")
487
+ ),
488
+ "resumed": recovered,
489
+ "writer_quality_requires_review": quality["requires_review"],
490
+ "writer_quality_warning_count": quality["warning_count"],
491
+ "writer_quality_dropped_count": quality["dropped_count"],
492
+ }
493
+ (out_dir / "build_report.json").write_text(
494
+ json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
495
+ )
496
+ failed = out_dir / "FAILED"
497
+ if failed.is_file():
498
+ history = {
499
+ "resolved_at": _now(),
500
+ "resolved_by_resume": recovered,
501
+ "failure": json.loads(failed.read_text(encoding="utf-8")),
502
+ }
503
+ with (out_dir / "build_failure_history.jsonl").open(
504
+ "a", encoding="utf-8"
505
+ ) as handle:
506
+ handle.write(json.dumps(history, sort_keys=True) + "\n")
507
+ failed.unlink()
508
+ (out_dir / "BUILD_COMPLETE").write_text(_now() + "\n", encoding="utf-8")
509
+ _stage(out_dir, "complete", resumed=recovered)
510
+ return report
511
+
512
+
513
+ def _writer_quality_report(
514
+ workers: Sequence[Mapping[str, Any]],
515
+ ) -> dict[str, Any]:
516
+ warning_counts: dict[str, int] = {}
517
+ dropped_counts: dict[str, int] = {}
518
+ message_count = 0
519
+ messages_with_warnings = 0
520
+ semantic_proposals = 0
521
+ semantic_committed = 0
522
+ for worker in workers:
523
+ path = Path(str(worker["worker_dir"])) / "product_write_messages.jsonl"
524
+ if not path.is_file():
525
+ continue
526
+ for raw_line in path.read_text(encoding="utf-8").splitlines():
527
+ if not raw_line.strip():
528
+ continue
529
+ row = json.loads(raw_line)
530
+ message_count += 1
531
+ semantic_proposals += int(row.get("semantic_proposals") or 0)
532
+ semantic_committed += int(row.get("semantic_committed") or 0)
533
+ warnings = list(row.get("validation_warnings") or [])
534
+ if warnings:
535
+ messages_with_warnings += 1
536
+ for warning in warnings:
537
+ code = str(warning.get("code") or "unknown_warning")
538
+ warning_counts[code] = warning_counts.get(code, 0) + 1
539
+ dropped = int(warning.get("dropped_count") or 0)
540
+ if dropped:
541
+ dropped_counts[code] = dropped_counts.get(code, 0) + dropped
542
+ review_codes = {
543
+ "assistant_assertions_dropped",
544
+ "invalid_assertion_quarantined",
545
+ "invalid_interaction_quarantined",
546
+ "invalid_item_collection_defaulted_empty",
547
+ "durability_defaulted_uncertain",
548
+ "temporal_durability_defaulted_uncertain",
549
+ "conflicting_durability_defaulted_uncertain",
550
+ }
551
+ review_events = sum(warning_counts.get(code, 0) for code in review_codes)
552
+ return {
553
+ "schema_version": "tmcra.v4.writer-quality.1",
554
+ "message_count": message_count,
555
+ "messages_with_warnings": messages_with_warnings,
556
+ "warning_count": sum(warning_counts.values()),
557
+ "warnings_by_code": dict(sorted(warning_counts.items())),
558
+ "dropped_count": sum(dropped_counts.values()),
559
+ "dropped_by_code": dict(sorted(dropped_counts.items())),
560
+ "review_event_count": review_events,
561
+ "review_codes": sorted(review_codes),
562
+ "requires_review": review_events > 0,
563
+ "semantic_proposals": semantic_proposals,
564
+ "semantic_committed": semantic_committed,
565
+ }
566
+
567
+
568
+ def _record_failure(out_dir: Path, exc: Exception) -> None:
569
+ failed = out_dir / "FAILED"
570
+ if failed.is_file():
571
+ with (out_dir / "build_failure_history.jsonl").open(
572
+ "a", encoding="utf-8"
573
+ ) as handle:
574
+ handle.write(
575
+ json.dumps(
576
+ {
577
+ "superseded_at": _now(),
578
+ "resolved": False,
579
+ "failure": json.loads(failed.read_text(encoding="utf-8")),
580
+ },
581
+ sort_keys=True,
582
+ )
583
+ + "\n"
584
+ )
585
+ failed.write_text(
586
+ json.dumps({"at": _now(), "error": f"{exc.__class__.__name__}: {exc}"})
587
+ + "\n",
588
+ encoding="utf-8",
589
+ )
590
+
591
+
592
+ def _writer_worker(worker: Mapping[str, Any], *, repo: Path, environment: Mapping[str, str]) -> None:
593
+ worker_dir = Path(str(worker["worker_dir"]))
594
+ _run(
595
+ [
596
+ sys.executable,
597
+ str(BASE / "tmcra_v4_batch_writer.py"),
598
+ "--input",
599
+ str(worker["input"]),
600
+ "--out-dir",
601
+ str(worker_dir),
602
+ "--repo",
603
+ str(repo),
604
+ ],
605
+ worker_dir / "writer.log",
606
+ environment,
607
+ )
608
+ _run(
609
+ [
610
+ sys.executable,
611
+ str(BASE / "audit_tmcra_v4_chain.py"),
612
+ "--run-dir",
613
+ str(worker_dir),
614
+ "--output",
615
+ str(worker_dir / "writer_chain_audit.json"),
616
+ "--worker-db",
617
+ f"worker={worker_dir / 'native_memory.sqlite3'}",
618
+ ],
619
+ worker_dir / "writer_audit.log",
620
+ environment,
621
+ )
622
+
623
+
624
+ def _writer_worker_resume(
625
+ worker: Mapping[str, Any],
626
+ *,
627
+ repo: Path,
628
+ environment: Mapping[str, str],
629
+ recover_interrupted_api_calls: bool,
630
+ ) -> None:
631
+ worker_dir = Path(str(worker["worker_dir"])).resolve()
632
+ command = [
633
+ sys.executable,
634
+ str(BASE / "tmcra_v4_batch_writer.py"),
635
+ "--input",
636
+ str(worker["input"]),
637
+ "--out-dir",
638
+ str(worker_dir),
639
+ "--repo",
640
+ str(repo),
641
+ "--revalidate-failed-raw-response",
642
+ ]
643
+ if recover_interrupted_api_calls:
644
+ command.append("--recover-interrupted-api-calls")
645
+ _run(
646
+ command,
647
+ _resume_log(worker_dir, "writer"),
648
+ environment,
649
+ )
650
+ _run(
651
+ [
652
+ sys.executable,
653
+ str(BASE / "audit_tmcra_v4_chain.py"),
654
+ "--run-dir",
655
+ str(worker_dir),
656
+ "--output",
657
+ str(worker_dir / "writer_chain_audit.json"),
658
+ "--worker-db",
659
+ f"worker={worker_dir / 'native_memory.sqlite3'}",
660
+ ],
661
+ _resume_log(worker_dir, "writer_audit"),
662
+ environment,
663
+ )
664
+
665
+
666
+ def _slow_worker(worker: Mapping[str, Any], *, repo: Path, environment: Mapping[str, str]) -> None:
667
+ worker_dir = Path(str(worker["worker_dir"]))
668
+ database = worker_dir / "native_memory.sqlite3"
669
+ scope_id = str(worker["scope_id"])
670
+ prefix = [sys.executable, str(BASE / "tmcra_v4_slow_graph.py"), str(database), "--repo", str(repo)]
671
+ _run([*prefix, "enqueue", scope_id], worker_dir / "slow_enqueue.log", environment)
672
+ _run([*prefix, "drain"], worker_dir / "slow_drain.log", environment)
673
+ _run(
674
+ [*prefix, "audit", scope_id, "--require-promotion-coverage"],
675
+ worker_dir / "slow_audit.log",
676
+ environment,
677
+ )
678
+
679
+
680
+ def _parallel(
681
+ workers: Sequence[Mapping[str, Any]],
682
+ concurrency: int,
683
+ task: Any,
684
+ environments: Sequence[Mapping[str, str]],
685
+ ) -> None:
686
+ if concurrency <= 0:
687
+ raise BuildError("concurrency must be positive")
688
+ with ThreadPoolExecutor(max_workers=min(concurrency, len(workers))) as executor:
689
+ futures = {
690
+ executor.submit(task, worker, environment=environments[index]): str(worker["question_id"])
691
+ for index, worker in enumerate(workers)
692
+ }
693
+ for future in as_completed(futures):
694
+ future.result()
695
+
696
+
697
+ def resume_build(args: argparse.Namespace) -> dict[str, Any]:
698
+ out_dir = args.out_dir.resolve()
699
+ if not out_dir.is_dir():
700
+ raise BuildError(f"resume output directory does not exist: {out_dir}")
701
+ if (out_dir / "BUILD_COMPLETE").exists():
702
+ raise BuildError("resume output is already complete")
703
+ manifest = _load_resume_manifest(out_dir)
704
+ expected_qids = [
705
+ line.strip()
706
+ for line in args.qid_list.resolve().read_text(encoding="utf-8").splitlines()
707
+ if line.strip()
708
+ ]
709
+ if expected_qids != list(manifest.get("qids") or []):
710
+ raise BuildError("resume qid list does not match the frozen input manifest")
711
+ shell_environment = _load_shell_environment(args.writer_env.resolve())
712
+ base_environment = {**os.environ, **shell_environment}
713
+ keys = _key_pool(base_environment)
714
+ workers = list(manifest["workers"])
715
+ environments = [
716
+ _worker_environment(base_environment, keys, index)
717
+ for index in range(len(workers))
718
+ ]
719
+ try:
720
+ incomplete_workers: list[Mapping[str, Any]] = []
721
+ for worker in workers:
722
+ try:
723
+ _verify_resume_writer(worker)
724
+ except BuildError:
725
+ incomplete_workers.append(worker)
726
+ if incomplete_workers:
727
+ if not getattr(args, "revalidate_failed_writer_raw_response", False):
728
+ raise BuildError(
729
+ "resume has incomplete writer workers; explicit raw-response "
730
+ "revalidation flag is required: "
731
+ + json.dumps(
732
+ [str(worker["question_id"]) for worker in incomplete_workers]
733
+ )
734
+ )
735
+ interrupted = [
736
+ call
737
+ for worker in incomplete_workers
738
+ for call in _interrupted_writer_calls(worker)
739
+ ]
740
+ if interrupted and not getattr(
741
+ args, "recover_interrupted_api_calls", False
742
+ ):
743
+ raise BuildError(
744
+ "resume has started calls without durable responses; explicit "
745
+ "process-loss recovery flag is required: "
746
+ + json.dumps(interrupted, sort_keys=True)
747
+ )
748
+ _stage(
749
+ out_dir,
750
+ "writer_resume_started",
751
+ workers=len(incomplete_workers),
752
+ interrupted_calls=len(interrupted),
753
+ )
754
+ selected_environments = [
755
+ environments[workers.index(worker)] for worker in incomplete_workers
756
+ ]
757
+ _parallel(
758
+ incomplete_workers,
759
+ args.writer_concurrency,
760
+ lambda worker, environment: _writer_worker_resume(
761
+ worker,
762
+ repo=args.repo.resolve(),
763
+ environment=environment,
764
+ recover_interrupted_api_calls=bool(interrupted),
765
+ ),
766
+ selected_environments,
767
+ )
768
+ for worker in workers:
769
+ _verify_resume_writer(worker)
770
+ _stage(out_dir, "writer_resume_complete")
771
+ _stage(out_dir, "subject_attribution_started", resumed=True)
772
+ _subject_attribution_stage(out_dir, base_environment)
773
+ interrupted_slow = [
774
+ call for worker in workers for call in _interrupted_slow_calls(worker)
775
+ ]
776
+ if interrupted_slow and not getattr(
777
+ args, "recover_interrupted_slow_calls", False
778
+ ):
779
+ raise BuildError(
780
+ "resume has Slow calls without durable responses; explicit "
781
+ "--recover-interrupted-slow-calls is required: "
782
+ + json.dumps(interrupted_slow, sort_keys=True)
783
+ )
784
+ _stage(
785
+ out_dir,
786
+ "resume_started",
787
+ row_count=len(workers),
788
+ interrupted_slow_calls=len(interrupted_slow),
789
+ )
790
+ _parallel(
791
+ workers,
792
+ args.slow_concurrency,
793
+ lambda worker, environment: _slow_worker_resume(
794
+ worker,
795
+ repo=args.repo.resolve(),
796
+ environment=environment,
797
+ recover_interrupted_slow_calls=bool(interrupted_slow),
798
+ ),
799
+ environments,
800
+ )
801
+ _stage(out_dir, "slow_graph_complete", resumed=True)
802
+ runtime_environment = dict(base_environment)
803
+ runtime_environment["TMCRA_NODE_MODEL_DEVICE"] = args.device
804
+ _run(
805
+ [
806
+ sys.executable,
807
+ str(BASE / "tmcra_v4_online_runtime.py"),
808
+ "build-index",
809
+ "--scope-manifest",
810
+ str(out_dir / "scope_manifest.jsonl"),
811
+ "--out-report",
812
+ str(out_dir / "index_report.json"),
813
+ "--embedding-model",
814
+ str(args.embedding_model.resolve()),
815
+ "--device",
816
+ args.device,
817
+ "--batch-size",
818
+ str(args.index_batch_size),
819
+ ],
820
+ _resume_log(out_dir, "index"),
821
+ runtime_environment,
822
+ )
823
+ _stage(out_dir, "index_complete", resumed=True)
824
+ _run(
825
+ [
826
+ sys.executable,
827
+ str(BASE / "audit_tmcra_v4_chain.py"),
828
+ "--run-dir",
829
+ str(out_dir),
830
+ "--output",
831
+ str(out_dir / "build_chain_audit.json"),
832
+ "--build-only",
833
+ ],
834
+ _resume_log(out_dir, "build_chain_audit"),
835
+ runtime_environment,
836
+ )
837
+ _stage(out_dir, "build_audit_complete", resumed=True)
838
+ return _finalize_build(
839
+ out_dir=out_dir,
840
+ workers=workers,
841
+ writer_concurrency=args.writer_concurrency,
842
+ slow_concurrency=args.slow_concurrency,
843
+ recovered=True,
844
+ )
845
+ except Exception as exc:
846
+ _record_failure(out_dir, exc)
847
+ raise
848
+
849
+
850
+ def build(args: argparse.Namespace) -> dict[str, Any]:
851
+ if getattr(args, "resume", False):
852
+ return resume_build(args)
853
+ out_dir = args.out_dir.resolve()
854
+ if out_dir.exists():
855
+ raise BuildError(f"output directory already exists: {out_dir}")
856
+ manifest = prepare(
857
+ data_path=args.data.resolve(),
858
+ qid_path=args.qid_list.resolve(),
859
+ out_dir=out_dir,
860
+ )
861
+ _stage(out_dir, "prepared", rows=manifest["row_count"])
862
+ shell_environment = _load_shell_environment(args.writer_env.resolve())
863
+ base_environment = {**os.environ, **shell_environment}
864
+ keys = _key_pool(base_environment)
865
+ workers = list(manifest["workers"])
866
+ environments = [
867
+ _worker_environment(base_environment, keys, index) for index in range(len(workers))
868
+ ]
869
+ try:
870
+ _stage(out_dir, "writer_started", concurrency=args.writer_concurrency)
871
+ _parallel(
872
+ workers,
873
+ args.writer_concurrency,
874
+ lambda worker, environment: _writer_worker(
875
+ worker, repo=args.repo.resolve(), environment=environment
876
+ ),
877
+ environments,
878
+ )
879
+ _stage(out_dir, "writer_complete")
880
+ _stage(out_dir, "subject_attribution_started")
881
+ _subject_attribution_stage(out_dir, base_environment)
882
+ _stage(out_dir, "slow_graph_started", concurrency=args.slow_concurrency)
883
+ _parallel(
884
+ workers,
885
+ args.slow_concurrency,
886
+ lambda worker, environment: _slow_worker(
887
+ worker, repo=args.repo.resolve(), environment=environment
888
+ ),
889
+ environments,
890
+ )
891
+ _stage(out_dir, "slow_graph_complete")
892
+ runtime_environment = dict(base_environment)
893
+ runtime_environment["TMCRA_NODE_MODEL_DEVICE"] = args.device
894
+ _run(
895
+ [
896
+ sys.executable,
897
+ str(BASE / "tmcra_v4_online_runtime.py"),
898
+ "build-index",
899
+ "--scope-manifest",
900
+ str(out_dir / "scope_manifest.jsonl"),
901
+ "--out-report",
902
+ str(out_dir / "index_report.json"),
903
+ "--embedding-model",
904
+ str(args.embedding_model.resolve()),
905
+ "--device",
906
+ args.device,
907
+ "--batch-size",
908
+ str(args.index_batch_size),
909
+ ],
910
+ out_dir / "index.log",
911
+ runtime_environment,
912
+ )
913
+ _stage(out_dir, "index_complete")
914
+ _run(
915
+ [
916
+ sys.executable,
917
+ str(BASE / "audit_tmcra_v4_chain.py"),
918
+ "--run-dir",
919
+ str(out_dir),
920
+ "--output",
921
+ str(out_dir / "build_chain_audit.json"),
922
+ "--build-only",
923
+ ],
924
+ out_dir / "build_chain_audit.log",
925
+ runtime_environment,
926
+ )
927
+ return _finalize_build(
928
+ out_dir=out_dir,
929
+ workers=workers,
930
+ writer_concurrency=args.writer_concurrency,
931
+ slow_concurrency=args.slow_concurrency,
932
+ recovered=False,
933
+ )
934
+ except Exception as exc:
935
+ _record_failure(out_dir, exc)
936
+ raise
937
+
938
+
939
+ def main() -> int:
940
+ parser = argparse.ArgumentParser(description="Build a frozen TMCRA V4 memory corpus")
941
+ parser.add_argument("--out-dir", type=Path, required=True)
942
+ parser.add_argument("--qid-list", type=Path, required=True)
943
+ parser.add_argument("--data", type=Path, default=DEFAULT_DATA)
944
+ parser.add_argument("--repo", type=Path, default=DEFAULT_REPO)
945
+ parser.add_argument("--writer-env", type=Path, default=DEFAULT_WRITER_ENV)
946
+ parser.add_argument("--embedding-model", type=Path, default=DEFAULT_EMBEDDING)
947
+ parser.add_argument("--writer-concurrency", type=int, default=1)
948
+ parser.add_argument("--slow-concurrency", type=int, default=1)
949
+ parser.add_argument("--index-batch-size", type=int, default=16)
950
+ parser.add_argument("--device", default="cuda")
951
+ parser.add_argument(
952
+ "--resume",
953
+ action="store_true",
954
+ help="resume a frozen incomplete build after explicit failed-job review",
955
+ )
956
+ parser.add_argument(
957
+ "--revalidate-failed-writer-raw-response",
958
+ action="store_true",
959
+ help="explicitly revalidate one saved clean writer response per failed worker",
960
+ )
961
+ parser.add_argument(
962
+ "--recover-interrupted-api-calls",
963
+ action="store_true",
964
+ help=(
965
+ "after explicit process-loss review, replace started writer calls "
966
+ "that have no durable response using the same model"
967
+ ),
968
+ )
969
+ parser.add_argument(
970
+ "--recover-interrupted-slow-calls",
971
+ action="store_true",
972
+ help=(
973
+ "after explicit review, journal expired Slow attempts with unknown "
974
+ "external outcomes and reopen only those jobs"
975
+ ),
976
+ )
977
+ args = parser.parse_args()
978
+ report = build(args)
979
+ print(json.dumps(report, sort_keys=True))
980
+ return 0
981
+
982
+
983
+ if __name__ == "__main__":
984
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_compile_evidence.py ADDED
@@ -0,0 +1,1015 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import json
7
+ import os
8
+ import re
9
+ import threading
10
+ import time
11
+ from collections.abc import Mapping, Sequence
12
+ from concurrent.futures import ThreadPoolExecutor, as_completed
13
+ from pathlib import Path
14
+ from typing import Any
15
+
16
+ from run_tmcra_v4_build import DEFAULT_WRITER_ENV, _key_pool, _load_shell_environment
17
+ from tmcra_v4_evidence_operations import (
18
+ PACKET_COMPILER_VERSION,
19
+ PACKET_SCHEMA,
20
+ PLAN_SCHEMA,
21
+ build_evidence_catalog,
22
+ compile_evidence_packet,
23
+ operation_plan_structural_risks,
24
+ unbound_memory_requirement_ids,
25
+ validate_operation_plan,
26
+ )
27
+ from tmcra_v4_evidence_planner import (
28
+ DeepSeekEvidenceOperationPlanner,
29
+ EvidencePlannerError,
30
+ PROMPT_VERSION as EVIDENCE_PLANNER_PROMPT_VERSION,
31
+ normalize_planner_output,
32
+ )
33
+ from tmcra_v4_route_policy import (
34
+ RoutePolicyError,
35
+ validate_diagnostic_retrieval_rows,
36
+ validate_production_retrieval_rows,
37
+ )
38
+
39
+
40
+ class EvidenceCompileError(RuntimeError):
41
+ pass
42
+
43
+
44
+ ARTIFACT_BINDING_SCHEMA = "tmcra.v4.evidence-compiler-artifact-binding.1"
45
+ REVIEW_POLICY_VERSION = "tmcra.v4.evidence-review-policy.4"
46
+
47
+
48
+ def _text(value: Any) -> str:
49
+ return value.strip() if isinstance(value, str) else ""
50
+
51
+
52
+ def _read_jsonl(path: Path) -> list[dict[str, Any]]:
53
+ return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
54
+
55
+
56
+ def _retrieval_debug_report(
57
+ path: Path,
58
+ expected_qids: Sequence[str],
59
+ ) -> dict[str, Any]:
60
+ if not path.is_file():
61
+ raise EvidenceCompileError(f"retrieval debug is missing: {path}")
62
+ raw = path.read_bytes()
63
+ try:
64
+ rows = [
65
+ json.loads(line)
66
+ for line in raw.decode("utf-8").splitlines()
67
+ if line.strip()
68
+ ]
69
+ except (UnicodeDecodeError, json.JSONDecodeError) as exc:
70
+ raise EvidenceCompileError(f"retrieval debug is invalid JSONL: {path}") from exc
71
+ actual_qids = [
72
+ _text(row.get("question_id") or row.get("qid"))
73
+ if isinstance(row, Mapping)
74
+ else ""
75
+ for row in rows
76
+ ]
77
+ if actual_qids != list(expected_qids):
78
+ raise EvidenceCompileError(
79
+ "retrieval debug rows do not match frozen evidence qid order"
80
+ )
81
+ return {
82
+ "schema_version": "tmcra.v4.compiled-retrieval-debug-binding.1",
83
+ "source_path": str(path),
84
+ "row_count": len(rows),
85
+ "source_sha256": hashlib.sha256(raw).hexdigest(),
86
+ }
87
+
88
+
89
+ def _stage_retrieval_debug(
90
+ source: Path,
91
+ out_dir: Path,
92
+ expected_qids: Sequence[str],
93
+ *,
94
+ expected_sha256: str,
95
+ ) -> dict[str, Any]:
96
+ report = _retrieval_debug_report(source, expected_qids)
97
+ if report["source_sha256"] != expected_sha256:
98
+ raise EvidenceCompileError("retrieval debug changed during evidence compilation")
99
+ raw = source.read_bytes()
100
+ destination = out_dir / "retrieval_debug.jsonl"
101
+ temporary = destination.with_name(
102
+ f".{destination.name}.tmp.{os.getpid()}.{time.time_ns()}"
103
+ )
104
+ temporary.write_bytes(raw)
105
+ os.replace(temporary, destination)
106
+ artifact_sha256 = hashlib.sha256(destination.read_bytes()).hexdigest()
107
+ if artifact_sha256 != expected_sha256:
108
+ raise EvidenceCompileError("staged retrieval debug failed its integrity check")
109
+ return {
110
+ **report,
111
+ "artifact_path": destination.name,
112
+ "artifact_sha256": artifact_sha256,
113
+ "status": "staged",
114
+ }
115
+
116
+
117
+ def _atomic_json(path: Path, value: Mapping[str, Any]) -> None:
118
+ temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}.{time.time_ns()}")
119
+ temporary.write_text(json.dumps(dict(value), ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
120
+ os.replace(temporary, path)
121
+
122
+
123
+ def _atomic_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
124
+ temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}.{time.time_ns()}")
125
+ temporary.write_text("".join(json.dumps(dict(row), ensure_ascii=False, sort_keys=True) + "\n" for row in rows), encoding="utf-8")
126
+ os.replace(temporary, path)
127
+
128
+
129
+ def _append_jsonl_atomic(path: Path, row: Mapping[str, Any]) -> None:
130
+ rows = _read_jsonl(path) if path.is_file() else []
131
+ rows.append(dict(row))
132
+ _atomic_jsonl(path, rows)
133
+
134
+
135
+ def _identity(row: Mapping[str, Any]) -> str:
136
+ payload = {
137
+ "question_id": row.get("question_id"),
138
+ "question": row.get("question"),
139
+ "question_date": row.get("question_date"),
140
+ "evidence_windows": row.get("evidence_windows"),
141
+ # Retrieval is part of the compiler input. In particular, a diagnostic
142
+ # row must never collide with a production row for the same question.
143
+ "retrieval_contract": row.get("retrieval_contract"),
144
+ }
145
+ return hashlib.sha256(json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")).hexdigest()
146
+
147
+
148
+ def _planner_contract(
149
+ *,
150
+ provider: str,
151
+ model: str,
152
+ review_policy_version: str = REVIEW_POLICY_VERSION,
153
+ ) -> dict[str, Any]:
154
+ return {
155
+ "provider": provider,
156
+ "model": model,
157
+ "prompt_version": EVIDENCE_PLANNER_PROMPT_VERSION,
158
+ "review_policy_version": review_policy_version,
159
+ "plan_schema": PLAN_SCHEMA,
160
+ }
161
+
162
+
163
+ def _packet_compiler_contract(*, compiler_version: str = PACKET_COMPILER_VERSION) -> dict[str, Any]:
164
+ return {
165
+ "packet_schema": PACKET_SCHEMA,
166
+ "compiler_version": compiler_version,
167
+ }
168
+
169
+
170
+ def _artifact_binding(
171
+ row: Mapping[str, Any],
172
+ *,
173
+ provider: str,
174
+ model: str,
175
+ compiler_version: str = PACKET_COMPILER_VERSION,
176
+ review_policy_version: str = REVIEW_POLICY_VERSION,
177
+ ) -> dict[str, Any]:
178
+ return {
179
+ "schema_version": ARTIFACT_BINDING_SCHEMA,
180
+ "input_sha256": _identity(row),
181
+ "input": {
182
+ "question_id": row.get("question_id"),
183
+ "question": row.get("question"),
184
+ "question_date": row.get("question_date"),
185
+ "evidence_windows": row.get("evidence_windows"),
186
+ "retrieval_route_contract": row.get("retrieval_contract"),
187
+ },
188
+ "planner": _planner_contract(
189
+ provider=provider,
190
+ model=model,
191
+ review_policy_version=review_policy_version,
192
+ ),
193
+ "packet_compiler": _packet_compiler_contract(compiler_version=compiler_version),
194
+ }
195
+
196
+
197
+ def _binding_sha256(binding: Mapping[str, Any]) -> str:
198
+ return hashlib.sha256(
199
+ json.dumps(dict(binding), ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
200
+ ).hexdigest()
201
+
202
+
203
+ def _saved_binding(
204
+ saved: Mapping[str, Any],
205
+ saved_row: Mapping[str, Any],
206
+ ) -> dict[str, Any]:
207
+ binding = saved.get("artifact_binding")
208
+ if isinstance(binding, Mapping):
209
+ return dict(binding)
210
+ planner = saved.get("planner")
211
+ if not isinstance(planner, Mapping):
212
+ raise EvidenceCompileError("completed artifact lacks its planner contract")
213
+ provider = _text(planner.get("provider"))
214
+ model = _text(planner.get("model"))
215
+ if not provider or not model:
216
+ raise EvidenceCompileError("completed artifact lacks its planner provider/model binding")
217
+ packet = saved_row.get("compiled_evidence_packet")
218
+ persisted_plan_schema = (
219
+ packet.get("operation_plan", {}).get("schema_version")
220
+ if isinstance(packet, Mapping)
221
+ and isinstance(packet.get("operation_plan"), Mapping)
222
+ else None
223
+ )
224
+ if persisted_plan_schema != PLAN_SCHEMA:
225
+ raise EvidenceCompileError("completed artifact planner contract schema is stale")
226
+ # Legacy artifacts did not persist this envelope. Reconstructing it is
227
+ # safe only from the persisted row and planner metadata; both are checked
228
+ # against the current invocation before reuse.
229
+ packet_compiler_version = (
230
+ _text(packet.get("packet_compiler_version"))
231
+ if isinstance(packet, Mapping)
232
+ else ""
233
+ )
234
+ return _artifact_binding(
235
+ saved_row,
236
+ provider=provider,
237
+ model=model,
238
+ compiler_version=packet_compiler_version or "legacy-unknown",
239
+ review_policy_version=(
240
+ _text(planner.get("review_policy_version")) or "legacy-unknown"
241
+ ),
242
+ )
243
+
244
+
245
+ def _assert_artifact_binding(
246
+ saved: Mapping[str, Any],
247
+ row: Mapping[str, Any],
248
+ *,
249
+ input_sha256: str,
250
+ expected_binding: Mapping[str, Any],
251
+ ) -> None:
252
+ saved_row = saved.get("row")
253
+ if not isinstance(saved_row, Mapping):
254
+ raise EvidenceCompileError("completed artifact lacks its persisted input row")
255
+ if _identity(saved_row) != input_sha256:
256
+ raise EvidenceCompileError("persisted compiler identity mismatch")
257
+ current_binding = dict(expected_binding)
258
+ persisted_binding = _saved_binding(saved, saved_row)
259
+ if persisted_binding.get("schema_version") not in {None, ARTIFACT_BINDING_SCHEMA}:
260
+ raise EvidenceCompileError("completed artifact binding schema is stale")
261
+ if persisted_binding.get("input_sha256") != input_sha256:
262
+ raise EvidenceCompileError("persisted compiler identity mismatch")
263
+ if persisted_binding.get("input") != current_binding.get("input"):
264
+ raise EvidenceCompileError("completed artifact input contract mismatch")
265
+ if persisted_binding.get("planner") != current_binding.get("planner"):
266
+ raise EvidenceCompileError("completed artifact planner contract/provider/model mismatch")
267
+ persisted_packet_contract = persisted_binding.get("packet_compiler")
268
+ current_packet_contract = current_binding.get("packet_compiler")
269
+ if not isinstance(persisted_packet_contract, Mapping) or not isinstance(current_packet_contract, Mapping):
270
+ raise EvidenceCompileError("completed artifact packet compiler contract is missing")
271
+ if persisted_packet_contract.get("packet_schema") != current_packet_contract.get("packet_schema"):
272
+ raise EvidenceCompileError("completed artifact packet schema contract mismatch")
273
+ if row.get("retrieval_contract") != saved_row.get("retrieval_contract"):
274
+ raise EvidenceCompileError("completed artifact retrieval route contract mismatch")
275
+
276
+ persisted_packet = saved_row.get("compiled_evidence_packet")
277
+ if isinstance(persisted_packet, Mapping) and persisted_packet.get("schema_version") != PACKET_SCHEMA:
278
+ raise EvidenceCompileError("completed artifact packet schema is stale")
279
+
280
+
281
+ def _cost_cny(calls: Sequence[Mapping[str, Any]]) -> float | None:
282
+ if any(call.get("provider") != "deepseek" for call in calls):
283
+ return None
284
+ total = 0.0
285
+ for call in calls:
286
+ usage = dict(call.get("usage") or {})
287
+ total += (
288
+ int(usage.get("prompt_cache_miss_tokens", 0)) * 3.0
289
+ + int(usage.get("completion_tokens", 0)) * 6.0
290
+ + int(usage.get("prompt_cache_hit_tokens", 0)) * 0.025
291
+ ) / 1_000_000.0
292
+ return round(total, 8)
293
+
294
+
295
+ def _physical_calls(values: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
296
+ calls: dict[str, dict[str, Any]] = {}
297
+
298
+ def visit(value: Any) -> None:
299
+ if not isinstance(value, Mapping):
300
+ return
301
+ for key in ("calls", "prior_calls"):
302
+ nested = value.get(key)
303
+ if isinstance(nested, Sequence) and not isinstance(nested, (str, bytes)):
304
+ for item in nested:
305
+ visit(item)
306
+ call_id = _text(value.get("physical_call_id"))
307
+ if value.get("physical_api_call") is True and call_id:
308
+ calls.setdefault(call_id, dict(value))
309
+
310
+ for value in values:
311
+ visit(value)
312
+ return list(calls.values())
313
+
314
+
315
+ def _aggregate_planner_calls(calls: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
316
+ if not calls:
317
+ raise EvidenceCompileError("planner call list is empty")
318
+ usage_keys = ("prompt_tokens", "completion_tokens", "prompt_cache_hit_tokens", "prompt_cache_miss_tokens", "total_tokens")
319
+ usage = {key: sum(int((call.get("usage") or {}).get(key, 0)) for call in calls) for key in usage_keys}
320
+ return {
321
+ **dict(calls[-1]),
322
+ "physical_api_calls": len(calls),
323
+ "usage": usage,
324
+ "review_applied": len(calls) > 1,
325
+ "calls": [dict(call) for call in calls],
326
+ }
327
+
328
+
329
+ RELATIVE_EVENT_REFERENCE_PATTERN = re.compile(
330
+ r"\b(?:when|before|after)\b", re.IGNORECASE
331
+ )
332
+ RELATIVE_DISTANCE_PATTERN = re.compile(
333
+ r"\b(?:ago|before|after)\b", re.IGNORECASE
334
+ )
335
+
336
+
337
+ def _relative_event_anchor_review_needed(
338
+ plan: Mapping[str, Any], catalog: Mapping[str, Any] | None
339
+ ) -> bool:
340
+ if not isinstance(catalog, Mapping):
341
+ return False
342
+ question = _text(catalog.get("question"))
343
+ if not (
344
+ RELATIVE_EVENT_REFERENCE_PATTERN.search(question)
345
+ and RELATIVE_DISTANCE_PATTERN.search(question)
346
+ ):
347
+ return False
348
+ evidence_by_id = {
349
+ _text(item.get("evidence_id")): item
350
+ for item in catalog.get("evidence") or []
351
+ if isinstance(item, Mapping)
352
+ }
353
+ anchor_ids = [
354
+ _text(value)
355
+ for value in catalog.get("lexical_anchor_ids") or []
356
+ if _text(value) in evidence_by_id
357
+ ]
358
+ anchor_sessions = {
359
+ _text(evidence_by_id[evidence_id].get("session_id"))
360
+ for evidence_id in anchor_ids
361
+ }
362
+ if len(anchor_sessions - {""}) < 2:
363
+ return False
364
+ question_atom_ids = {
365
+ _text(item.get("atom_id"))
366
+ for item in catalog.get("atoms") or []
367
+ if isinstance(item, Mapping) and item.get("evidence_id") == "QUESTION"
368
+ }
369
+ for operation in plan.get("operations") or []:
370
+ if (
371
+ not isinstance(operation, Mapping)
372
+ or operation.get("operation_type") != "date_difference"
373
+ or not question_atom_ids.intersection(operation.get("input_atom_ids") or [])
374
+ ):
375
+ continue
376
+ bound_ids = {
377
+ _text(value) for value in operation.get("input_evidence_ids") or []
378
+ }
379
+ bound_sessions = {
380
+ _text(evidence_by_id[value].get("session_id"))
381
+ for value in bound_ids
382
+ if value in evidence_by_id
383
+ }
384
+ if any(
385
+ evidence_id not in bound_ids
386
+ and _text(evidence_by_id[evidence_id].get("session_id"))
387
+ not in bound_sessions
388
+ for evidence_id in anchor_ids
389
+ ):
390
+ return True
391
+ return False
392
+
393
+
394
+ def _required_memory_premise_binding(
395
+ plan: Mapping[str, Any], catalog: Mapping[str, Any] | None
396
+ ) -> dict[str, Any] | None:
397
+ if not isinstance(catalog, Mapping):
398
+ return None
399
+ anchor_ids = {
400
+ _text(value)
401
+ for value in catalog.get("lexical_anchor_ids") or []
402
+ if _text(value)
403
+ }
404
+ if not anchor_ids:
405
+ return None
406
+ required_memory_ids = {
407
+ _text(premise.get("premise_id"))
408
+ for premise in (plan.get("task_contract") or {}).get("premises") or []
409
+ if isinstance(premise, Mapping)
410
+ and premise.get("source") == "memory"
411
+ and premise.get("necessity") == "required"
412
+ and _text(premise.get("premise_id"))
413
+ }
414
+ if not required_memory_ids:
415
+ return None
416
+ bound_ids = {
417
+ _text(value)
418
+ for requirement in plan.get("requirements") or []
419
+ if isinstance(requirement, Mapping)
420
+ and _text(requirement.get("requirement_id")) in required_memory_ids
421
+ for value in requirement.get("evidence_ids") or []
422
+ if _text(value)
423
+ }
424
+ if not bound_ids or bound_ids.intersection(anchor_ids):
425
+ return None
426
+ return {
427
+ "required_memory_requirement_ids": sorted(required_memory_ids),
428
+ "bound_evidence_ids": sorted(bound_ids),
429
+ "lexical_anchor_ids": sorted(anchor_ids),
430
+ }
431
+
432
+
433
+ def _review_context_for_plan(
434
+ plan: Mapping[str, Any], catalog: Mapping[str, Any] | None = None
435
+ ) -> dict[str, Any] | None:
436
+ missing_requirement_ids = unbound_memory_requirement_ids(plan)
437
+ structural_risks = operation_plan_structural_risks(plan)
438
+ relative_event_anchor_risk = _relative_event_anchor_review_needed(plan, catalog)
439
+ memory_relevance_signal = _required_memory_premise_binding(plan, catalog)
440
+ review_reasons = [
441
+ *structural_risks,
442
+ *(["unbound_required_memory_premises"] if missing_requirement_ids else []),
443
+ *(["relative_event_anchor_not_bound"] if relative_event_anchor_risk else []),
444
+ *(
445
+ ["required_memory_premise_outside_query_anchors"]
446
+ if memory_relevance_signal
447
+ else []
448
+ ),
449
+ ]
450
+ if not review_reasons:
451
+ return None
452
+ relative_anchor_instruction = (
453
+ " This review has already established that another remembered event "
454
+ "is the relative-time endpoint. QUESTION is forbidden as a "
455
+ "date_difference input; bind at least two distinct Source event dates."
456
+ if relative_event_anchor_risk
457
+ else ""
458
+ )
459
+ memory_relevance_instruction = (
460
+ " The required memory premise is currently bound only to evidence that "
461
+ "falls outside the query's lexical anchors. Do not satisfy the memory "
462
+ "contract with an unrelated remembered detail. Reinspect the candidate "
463
+ "evidence, including user-authored Source text and source-group context, "
464
+ "and bind memory that materially constrains the answer target. Lexical "
465
+ "anchors are review candidates rather than a hard allowlist, so an "
466
+ "implicitly related memory may remain only when the rebuilt plan changes "
467
+ "the evidence binding or adds relevant corroborating evidence."
468
+ if memory_relevance_signal
469
+ else ""
470
+ )
471
+ hard_constraints: dict[str, Any] = {}
472
+ if relative_event_anchor_risk:
473
+ hard_constraints.update(
474
+ {
475
+ "forbidden_date_difference_evidence_ids": ["QUESTION"],
476
+ "minimum_distinct_source_event_anchors": 2,
477
+ }
478
+ )
479
+ if memory_relevance_signal:
480
+ hard_constraints["forbidden_required_memory_evidence_ids"] = list(
481
+ memory_relevance_signal["bound_evidence_ids"]
482
+ )
483
+ return {
484
+ "initial_plan": dict(plan),
485
+ "review_reasons": review_reasons,
486
+ "missing_requirement_ids": missing_requirement_ids,
487
+ "memory_premise_relevance_signal": memory_relevance_signal or {},
488
+ "hard_constraints": hard_constraints,
489
+ "instruction": (
490
+ "Rebuild the complete plan. Treat a missing state as a hypothesis, "
491
+ "not proof of absence; for memory-conditioned generation, bind the "
492
+ "remembered constraints instead of searching for a ready-made answer. "
493
+ "When a relative-time question names another remembered event with "
494
+ "when, before, or after, compare the distinct event anchors and do not "
495
+ "default to the question date unless the wording explicitly makes it "
496
+ "the comparison endpoint."
497
+ + relative_anchor_instruction
498
+ + memory_relevance_instruction
499
+ ),
500
+ }
501
+
502
+
503
+ def _review_resolution_failure(
504
+ plan: Mapping[str, Any],
505
+ catalog: Mapping[str, Any] | None,
506
+ review_context: Mapping[str, Any] | None,
507
+ ) -> str:
508
+ reasons = set((review_context or {}).get("review_reasons") or [])
509
+ if (
510
+ "relative_event_anchor_not_bound" in reasons
511
+ and _relative_event_anchor_review_needed(plan, catalog)
512
+ ):
513
+ return (
514
+ "relative event-anchor review still binds QUESTION instead of "
515
+ "two distinct Source event anchors"
516
+ )
517
+ if "required_memory_premise_outside_query_anchors" in reasons:
518
+ current_signal = _required_memory_premise_binding(plan, catalog)
519
+ initial_signal = (review_context or {}).get(
520
+ "memory_premise_relevance_signal"
521
+ ) or {}
522
+ if current_signal and set(current_signal.get("bound_evidence_ids") or []) == set(
523
+ initial_signal.get("bound_evidence_ids") or []
524
+ ):
525
+ return (
526
+ "required memory-premise relevance review kept the same "
527
+ "out-of-anchor evidence binding"
528
+ )
529
+ return ""
530
+
531
+
532
+ def _recover_plan_from_failure(
533
+ failure_path: Path,
534
+ *,
535
+ question_id: str,
536
+ input_sha256: str,
537
+ catalog: Mapping[str, Any],
538
+ expected_binding: Mapping[str, Any] | None = None,
539
+ ) -> tuple[dict[str, Any], dict[str, Any]] | None:
540
+ if not failure_path.is_file():
541
+ return None
542
+ persisted = json.loads(failure_path.read_text(encoding="utf-8"))
543
+ if (
544
+ persisted.get("question_id") != question_id
545
+ or persisted.get("input_sha256") != input_sha256
546
+ ):
547
+ raise EvidenceCompileError(
548
+ f"{question_id}: persisted compiler failure identity mismatch"
549
+ )
550
+ metadata = dict(persisted.get("planner") or {})
551
+ if _text(metadata.get("prompt_version")) != EVIDENCE_PLANNER_PROMPT_VERSION:
552
+ return None
553
+ if _text(metadata.get("review_policy_version")) != REVIEW_POLICY_VERSION:
554
+ return None
555
+ if expected_binding is not None:
556
+ expected_planner = expected_binding.get("planner")
557
+ if not isinstance(expected_planner, Mapping) or any(
558
+ metadata.get(field) != expected_planner.get(field)
559
+ for field in ("provider", "model", "prompt_version")
560
+ ):
561
+ return None
562
+ raw_plan = metadata.get("raw_plan")
563
+ if not isinstance(raw_plan, Mapping) and metadata.get(
564
+ "review_resolution_failed"
565
+ ):
566
+ review_context = metadata.get("review_context")
567
+ if isinstance(review_context, Mapping):
568
+ raw_plan = review_context.get("initial_plan")
569
+ if not isinstance(raw_plan, Mapping):
570
+ return None
571
+ try:
572
+ plan, warnings = normalize_planner_output(raw_plan, catalog)
573
+ except Exception:
574
+ # The journal is an optimization, not a permanent failure latch. A raw
575
+ # response that cannot satisfy the current contract must be replanned.
576
+ return None
577
+ return plan, {
578
+ **metadata,
579
+ "raw_plan": dict(raw_plan),
580
+ "recovered_from_persisted_raw_plan": True,
581
+ "normalization_warnings": warnings,
582
+ }
583
+
584
+
585
+ def _replan_context_from_failure(
586
+ failure_path: Path,
587
+ *,
588
+ question_id: str,
589
+ input_sha256: str,
590
+ expected_binding: Mapping[str, Any] | None = None,
591
+ catalog: Mapping[str, Any] | None = None,
592
+ ) -> dict[str, Any] | None:
593
+ if not failure_path.is_file():
594
+ return None
595
+ persisted = json.loads(failure_path.read_text(encoding="utf-8"))
596
+ if (
597
+ persisted.get("question_id") != question_id
598
+ or persisted.get("input_sha256") != input_sha256
599
+ ):
600
+ raise EvidenceCompileError(
601
+ f"{question_id}: persisted compiler failure identity mismatch"
602
+ )
603
+ metadata = dict(persisted.get("planner") or {})
604
+ if _text(metadata.get("prompt_version")) != EVIDENCE_PLANNER_PROMPT_VERSION:
605
+ return None
606
+ if expected_binding is not None:
607
+ expected_planner = expected_binding.get("planner")
608
+ if not isinstance(expected_planner, Mapping) or any(
609
+ metadata.get(field) != expected_planner.get(field)
610
+ for field in ("provider", "model", "prompt_version")
611
+ ):
612
+ return None
613
+ raw_plan = metadata.get("raw_plan")
614
+ error = _text(persisted.get("error"))
615
+ if not isinstance(raw_plan, Mapping) or not error:
616
+ return None
617
+ candidate_memory_evidence_ids = [
618
+ _text(value)
619
+ for value in (catalog or {}).get("lexical_anchor_ids") or []
620
+ if _text(value)
621
+ ]
622
+ return {
623
+ "initial_plan": dict(raw_plan),
624
+ "review_reasons": ["persisted_plan_validation_failure"],
625
+ "prior_validation_error": error,
626
+ "candidate_memory_evidence_ids": candidate_memory_evidence_ids,
627
+ "instruction": (
628
+ "Rebuild the complete plan and correct the prior runtime contract "
629
+ f"validation failure: {error}. Inspect the supplied evidence again. "
630
+ "For memory-conditioned generation, bind at least one relevant "
631
+ "required memory preference, constraint, owned tools, or resources. "
632
+ "For recommendation or advice, a previously mentioned owned tool or "
633
+ "accessory that can directly address the target problem is a valid "
634
+ "memory premise. Do not choose an unrelated preference merely because "
635
+ "it is more explicit; the premise must share the target domain or "
636
+ "intended action with the current question. Reinspect the listed "
637
+ "candidate_memory_evidence_ids plus their user-authored Source and "
638
+ "source-group context, and bind the relevant parent Evidence ID. "
639
+ "The candidate list is a review hint, not a hard allowlist. If no "
640
+ "relevant memory is supportable, "
641
+ "preserve an unbound memory requirement instead of relabeling query "
642
+ "context as remembered evidence."
643
+ ),
644
+ }
645
+
646
+
647
+ def _refresh_completed_artifact(
648
+ saved: Mapping[str, Any],
649
+ row: Mapping[str, Any],
650
+ *,
651
+ question_id: str,
652
+ input_sha256: str,
653
+ expected_binding: Mapping[str, Any] | None = None,
654
+ ) -> tuple[dict[str, Any], bool]:
655
+ if (
656
+ saved.get("question_id") != question_id
657
+ ):
658
+ raise EvidenceCompileError(
659
+ f"{question_id}: persisted compiler identity mismatch"
660
+ )
661
+ saved_row = saved.get("row")
662
+ if saved.get("input_sha256") != input_sha256 and (
663
+ not isinstance(saved_row, Mapping) or _identity(saved_row) != input_sha256
664
+ ):
665
+ raise EvidenceCompileError(
666
+ f"{question_id}: persisted compiler identity mismatch"
667
+ )
668
+ output = dict(saved)
669
+ if expected_binding is None:
670
+ persisted_binding = saved.get("artifact_binding")
671
+ planner = (
672
+ persisted_binding.get("planner")
673
+ if isinstance(persisted_binding, Mapping)
674
+ else saved.get("planner")
675
+ )
676
+ if not isinstance(planner, Mapping):
677
+ raise EvidenceCompileError(
678
+ f"{question_id}: completed artifact lacks its planner contract"
679
+ )
680
+ expected_binding = _artifact_binding(
681
+ row,
682
+ provider=_text(planner.get("provider")),
683
+ model=_text(planner.get("model")),
684
+ )
685
+ _assert_artifact_binding(
686
+ saved,
687
+ row,
688
+ input_sha256=input_sha256,
689
+ expected_binding=expected_binding,
690
+ )
691
+ planner_metadata = output.get("planner")
692
+ if (
693
+ not isinstance(planner_metadata, Mapping)
694
+ or _text(planner_metadata.get("prompt_version"))
695
+ != EVIDENCE_PLANNER_PROMPT_VERSION
696
+ ):
697
+ raise EvidenceCompileError(
698
+ f"{question_id}: completed artifact uses a stale evidence planner contract"
699
+ )
700
+ saved_row = output.get("row")
701
+ packet = saved_row.get("compiled_evidence_packet") if isinstance(saved_row, Mapping) else None
702
+ if not isinstance(packet, Mapping):
703
+ raise EvidenceCompileError(
704
+ f"{question_id}: completed artifact lacks a compiled evidence packet"
705
+ )
706
+ binding_changed = output.get("artifact_binding") != dict(expected_binding)
707
+ if packet.get("packet_compiler_version") == PACKET_COMPILER_VERSION:
708
+ if binding_changed:
709
+ output["artifact_binding"] = dict(expected_binding)
710
+ output["compiler_binding_sha256"] = _binding_sha256(expected_binding)
711
+ return output, binding_changed
712
+ plan = packet.get("operation_plan")
713
+ if not isinstance(plan, Mapping):
714
+ raise EvidenceCompileError(
715
+ f"{question_id}: completed artifact lacks its persisted operation plan"
716
+ )
717
+ output_row = dict(row)
718
+ output_row["compiled_evidence_packet"] = compile_evidence_packet(row, plan)
719
+ output["row"] = output_row
720
+ output["packet_compiler_version"] = PACKET_COMPILER_VERSION
721
+ output["artifact_binding"] = dict(expected_binding)
722
+ output["compiler_binding_sha256"] = _binding_sha256(expected_binding)
723
+ output["local_packet_recompile_count"] = int(
724
+ output.get("local_packet_recompile_count", 0)
725
+ ) + 1
726
+ return output, True
727
+
728
+
729
+ def main() -> int:
730
+ parser = argparse.ArgumentParser(description="Compile immutable TMCRA evidence into operation-bound answer packets")
731
+ parser.add_argument("--evidence", type=Path, required=True)
732
+ parser.add_argument("--out-dir", type=Path, required=True)
733
+ parser.add_argument("--qid-list", type=Path)
734
+ parser.add_argument(
735
+ "--retrieval-debug",
736
+ type=Path,
737
+ help="retrieval_debug.jsonl to bind into the compiled evaluation artifact",
738
+ )
739
+ parser.add_argument(
740
+ "--require-retrieval-debug",
741
+ action="store_true",
742
+ help="fail before planner calls unless retrieval debug is present and qid-aligned",
743
+ )
744
+ parser.add_argument("--writer-env", type=Path, default=DEFAULT_WRITER_ENV)
745
+ parser.add_argument("--workers", type=int, default=4)
746
+ parser.add_argument("--timeout", type=float, default=180.0)
747
+ parser.add_argument("--planner-provider", choices=("deepseek", "xiaomi_mimo"), default="deepseek")
748
+ parser.add_argument("--planner-model")
749
+ parser.add_argument("--planner-base-url")
750
+ parser.add_argument("--planner-key-file", type=Path)
751
+ parser.add_argument(
752
+ "--diagnostic",
753
+ action="store_true",
754
+ help="explicitly compile a diagnostic retrieval lane that cannot enter production evaluation",
755
+ )
756
+ args = parser.parse_args()
757
+ if args.workers <= 0:
758
+ raise EvidenceCompileError("workers must be positive")
759
+ rows = _read_jsonl(args.evidence.resolve())
760
+ if args.qid_list:
761
+ qids = [line.strip() for line in args.qid_list.read_text(encoding="utf-8").splitlines() if line.strip()]
762
+ by_qid = {_text(row.get("question_id")): row for row in rows}
763
+ if not qids or len(qids) != len(set(qids)) or any(qid not in by_qid for qid in qids):
764
+ raise EvidenceCompileError("qid list is empty, duplicated, or absent from evidence")
765
+ rows = [by_qid[qid] for qid in qids]
766
+ try:
767
+ route_report = (
768
+ validate_diagnostic_retrieval_rows(rows)
769
+ if args.diagnostic
770
+ else validate_production_retrieval_rows(rows)
771
+ )
772
+ except RoutePolicyError as exc:
773
+ raise EvidenceCompileError(
774
+ f"retrieval route failed before evidence planner calls: {exc}"
775
+ ) from exc
776
+ expected_qids = [_text(row.get("question_id")) for row in rows]
777
+ retrieval_debug_source = (
778
+ args.retrieval_debug.resolve()
779
+ if args.retrieval_debug
780
+ else args.evidence.resolve().parent / "retrieval_debug.jsonl"
781
+ )
782
+ retrieval_debug_preflight: dict[str, Any] | None = None
783
+ if retrieval_debug_source.is_file():
784
+ retrieval_debug_preflight = _retrieval_debug_report(
785
+ retrieval_debug_source,
786
+ expected_qids,
787
+ )
788
+ elif args.require_retrieval_debug:
789
+ raise EvidenceCompileError(
790
+ f"retrieval debug is required before planner calls: {retrieval_debug_source}"
791
+ )
792
+ if args.planner_provider == "deepseek":
793
+ environment = _load_shell_environment(args.writer_env.resolve())
794
+ keys = _key_pool(environment)
795
+ base_url = args.planner_base_url or environment.get("TMCRA_DEEPSEEK_WRITER_BASE_URL") or environment.get("TMCRA_WRITER_BASE_URL") or "https://api.deepseek.com/v1"
796
+ planner_model = args.planner_model or "deepseek-v4-pro"
797
+ else:
798
+ if not args.planner_key_file or not args.planner_key_file.is_file():
799
+ raise EvidenceCompileError("xiaomi_mimo requires --planner-key-file")
800
+ key = args.planner_key_file.read_text(encoding="utf-8").strip()
801
+ if not key:
802
+ raise EvidenceCompileError("planner key file is empty")
803
+ keys = [key]
804
+ base_url = args.planner_base_url or "https://api.xiaomimimo.com/v1"
805
+ planner_model = args.planner_model or "mimo-v2.5"
806
+ out_dir = args.out_dir.resolve()
807
+ journal_dir = out_dir / "rows"
808
+ journal_dir.mkdir(parents=True, exist_ok=True)
809
+ failure_history_path = out_dir / "planner_failure_history.jsonl"
810
+ lock = threading.Lock()
811
+
812
+ def compile_one(index: int, row: Mapping[str, Any]) -> dict[str, Any]:
813
+ qid = _text(row.get("question_id"))
814
+ identity = _identity(row)
815
+ binding = _artifact_binding(
816
+ row,
817
+ provider=args.planner_provider,
818
+ model=planner_model,
819
+ )
820
+ artifact = journal_dir / f"{index:06d}_{qid}.json"
821
+ failure_path = journal_dir / f"{index:06d}_{qid}.failure.json"
822
+ if artifact.is_file():
823
+ saved = json.loads(artifact.read_text(encoding="utf-8"))
824
+ refreshed, changed = _refresh_completed_artifact(
825
+ saved,
826
+ row,
827
+ question_id=qid,
828
+ input_sha256=identity,
829
+ expected_binding=binding,
830
+ )
831
+ if changed:
832
+ with lock:
833
+ _atomic_json(artifact, refreshed)
834
+ return refreshed
835
+ catalog = build_evidence_catalog(row)
836
+ planner = DeepSeekEvidenceOperationPlanner(
837
+ base_url=base_url,
838
+ api_keys=[keys[index % len(keys)]],
839
+ timeout=args.timeout,
840
+ model=planner_model,
841
+ provider=args.planner_provider,
842
+ )
843
+ planner_calls: list[dict[str, Any]] = []
844
+ recovered_from_persisted_raw_plan = False
845
+ try:
846
+ persisted_replan_context = _replan_context_from_failure(
847
+ failure_path,
848
+ question_id=qid,
849
+ input_sha256=identity,
850
+ expected_binding=binding,
851
+ catalog=catalog,
852
+ )
853
+ recovered = _recover_plan_from_failure(
854
+ failure_path,
855
+ question_id=qid,
856
+ input_sha256=identity,
857
+ catalog=catalog,
858
+ expected_binding=binding,
859
+ )
860
+ if recovered is not None:
861
+ plan, recovered_metadata = recovered
862
+ recovered_from_persisted_raw_plan = True
863
+ planner_calls.append(recovered_metadata)
864
+ if not recovered_from_persisted_raw_plan:
865
+ plan, metadata = planner.plan(
866
+ catalog,
867
+ review_context=persisted_replan_context,
868
+ )
869
+ planner_calls.append(metadata)
870
+ review_context = _review_context_for_plan(plan, catalog)
871
+ if review_context is not None:
872
+ plan, review_metadata = planner.plan(
873
+ catalog,
874
+ review_context=review_context,
875
+ )
876
+ planner_calls.append(review_metadata)
877
+ review_failure = _review_resolution_failure(
878
+ plan,
879
+ catalog,
880
+ review_context,
881
+ )
882
+ if review_failure:
883
+ raise EvidencePlannerError(
884
+ review_failure,
885
+ metadata={
886
+ **review_metadata,
887
+ "review_resolution_failed": True,
888
+ "review_context": review_context,
889
+ },
890
+ )
891
+ metadata = _aggregate_planner_calls(planner_calls)
892
+ metadata["review_policy_version"] = REVIEW_POLICY_VERSION
893
+ metadata["recovered_from_persisted_raw_plan"] = (
894
+ recovered_from_persisted_raw_plan
895
+ )
896
+ metadata["new_physical_api_calls_this_resume"] = len(planner_calls) - int(
897
+ recovered_from_persisted_raw_plan
898
+ )
899
+ except EvidencePlannerError as exc:
900
+ failure = {
901
+ "failure_id": f"ecf_{time.time_ns()}",
902
+ "observed_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
903
+ "question_id": qid,
904
+ "input_sha256": identity,
905
+ "artifact_binding": binding,
906
+ "compiler_binding_sha256": _binding_sha256(binding),
907
+ "status": "failed",
908
+ "error": str(exc),
909
+ "planner": {
910
+ **exc.metadata,
911
+ "review_policy_version": REVIEW_POLICY_VERSION,
912
+ **({"prior_calls": planner_calls} if planner_calls else {}),
913
+ },
914
+ }
915
+ with lock:
916
+ _atomic_json(failure_path, failure)
917
+ _append_jsonl_atomic(failure_history_path, failure)
918
+ raise
919
+ plan = validate_operation_plan(plan, catalog)
920
+ packet = compile_evidence_packet(row, plan)
921
+ output_row = dict(row)
922
+ output_row["compiled_evidence_packet"] = packet
923
+ result = {
924
+ "question_id": qid,
925
+ "input_sha256": identity,
926
+ "artifact_binding": binding,
927
+ "compiler_binding_sha256": _binding_sha256(binding),
928
+ "status": "completed",
929
+ "planner": metadata,
930
+ "row": output_row,
931
+ }
932
+ with lock:
933
+ _atomic_json(artifact, result)
934
+ failure_path.unlink(missing_ok=True)
935
+ return result
936
+
937
+ results: dict[str, dict[str, Any]] = {}
938
+ failures: list[dict[str, str]] = []
939
+ with ThreadPoolExecutor(max_workers=args.workers) as executor:
940
+ futures = {executor.submit(compile_one, index, row): _text(row.get("question_id")) for index, row in enumerate(rows)}
941
+ for future in as_completed(futures):
942
+ qid = futures[future]
943
+ try:
944
+ result = future.result()
945
+ except Exception as exc:
946
+ failures.append({"question_id": qid, "error": f"{exc.__class__.__name__}: {exc}"})
947
+ else:
948
+ results[qid] = result
949
+ if failures:
950
+ _atomic_json(out_dir / "failures.json", {"failures": failures})
951
+ raise EvidenceCompileError(f"evidence compilation failed for {len(failures)} rows")
952
+ ordered_results = [results[_text(row.get("question_id"))] for row in rows]
953
+ (out_dir / "failures.json").unlink(missing_ok=True)
954
+ ordered_rows = [item["row"] for item in ordered_results]
955
+ calls = [item["planner"] for item in ordered_results]
956
+ failure_history = _read_jsonl(failure_history_path) if failure_history_path.is_file() else []
957
+ completed_physical_calls = _physical_calls(calls)
958
+ all_physical_calls = _physical_calls(
959
+ [*calls, *(dict(item.get("planner") or {}) for item in failure_history)]
960
+ )
961
+ _atomic_jsonl(out_dir / "evidence_windows.jsonl", ordered_rows)
962
+ retrieval_debug_report = (
963
+ _stage_retrieval_debug(
964
+ retrieval_debug_source,
965
+ out_dir,
966
+ expected_qids,
967
+ expected_sha256=retrieval_debug_preflight["source_sha256"],
968
+ )
969
+ if retrieval_debug_preflight is not None
970
+ else {
971
+ "schema_version": "tmcra.v4.compiled-retrieval-debug-binding.1",
972
+ "status": "not_provided",
973
+ }
974
+ )
975
+ _atomic_json(
976
+ out_dir / "report.json",
977
+ {
978
+ "schema_version": "tmcra.v4.evidence-compiler-run.1",
979
+ "status": "complete",
980
+ "row_count": len(ordered_rows),
981
+ "physical_call_count": len(all_physical_calls),
982
+ "completed_artifact_physical_call_count": len(completed_physical_calls),
983
+ "failure_history_record_count": len(failure_history),
984
+ "locally_recompiled_row_count": sum(
985
+ int(item.get("local_packet_recompile_count", 0) > 0)
986
+ for item in ordered_results
987
+ ),
988
+ "packet_compiler_version": PACKET_COMPILER_VERSION,
989
+ "planner_provider": args.planner_provider,
990
+ "planner_model": planner_model,
991
+ "retrieval_route_report": route_report,
992
+ "retrieval_debug": retrieval_debug_report,
993
+ "operation_count": sum(len(row["compiled_evidence_packet"]["operation_results"]) for row in ordered_rows),
994
+ "requirement_count": sum(len(row["compiled_evidence_packet"]["requirement_coverage"]) for row in ordered_rows),
995
+ "structural_risk_count": sum(
996
+ len(row["compiled_evidence_packet"].get("structural_risk_signals") or [])
997
+ for row in ordered_rows
998
+ ),
999
+ "typed_semantics_accepted_count": sum(
1000
+ len((row["compiled_evidence_packet"].get("typed_semantics_report") or {}).get("accepted") or [])
1001
+ for row in ordered_rows
1002
+ ),
1003
+ "typed_semantics_rejected_count": sum(
1004
+ len((row["compiled_evidence_packet"].get("typed_semantics_report") or {}).get("rejected") or [])
1005
+ for row in ordered_rows
1006
+ ),
1007
+ "exact_cost_cny": _cost_cny(all_physical_calls),
1008
+ },
1009
+ )
1010
+ (out_dir / "COMPILE_COMPLETE").write_text(time.strftime("%Y-%m-%dT%H:%M:%S%z") + "\n", encoding="utf-8")
1011
+ return 0
1012
+
1013
+
1014
+ if __name__ == "__main__":
1015
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_evaluate.py ADDED
@@ -0,0 +1,725 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import json
6
+ import os
7
+ import subprocess
8
+ import sys
9
+ from collections.abc import Mapping
10
+ from pathlib import Path
11
+ from typing import Any
12
+
13
+ from run_tmcra_v4_build import BASE, DEFAULT_DATA, _now, _resume_log, _run
14
+ from run_tmcra_v4_retrieve import DEFAULT_HARNESS
15
+ from tmcra_v4_route_policy import (
16
+ POLICY_SCHEMA,
17
+ PRODUCTION_ANSWER_MODEL,
18
+ PRODUCTION_ANSWER_PROTOCOL,
19
+ PRODUCTION_ANSWER_RUNNER,
20
+ PRODUCTION_LANE,
21
+ SOURCE_COVERAGE_TRACE_K,
22
+ RoutePolicyError,
23
+ assert_production_answer_runner,
24
+ validate_production_answers,
25
+ validate_production_evidence,
26
+ )
27
+ from tmcra_v4_regression_gate import (
28
+ DEFAULT_CASES as DEFAULT_REGRESSION_CASES,
29
+ evaluate_gate as evaluate_regression_gate,
30
+ )
31
+
32
+
33
+ PROJECT_ROOT = BASE.parents[2]
34
+ PACKAGE_ROOT = BASE.parent
35
+ DEFAULT_ANSWER_ENV = PROJECT_ROOT / "configs" / "answer.env"
36
+ DEFAULT_JUDGE = PACKAGE_ROOT / "harness" / "evaluate_qa_official_judge.py"
37
+ DEFAULT_ANSWER_RUNNER = BASE / PRODUCTION_ANSWER_RUNNER
38
+
39
+
40
+ class EvaluationError(RuntimeError):
41
+ pass
42
+
43
+
44
+ def _load_environment(path: Path) -> dict[str, str]:
45
+ if not path.is_file():
46
+ raise EvaluationError(f"answer environment file is missing: {path}")
47
+ result = subprocess.run(
48
+ ["bash", "-c", 'set -a; source "$1"; env -0', "tmcra-v4-answer-env", str(path)],
49
+ check=True,
50
+ stdout=subprocess.PIPE,
51
+ stderr=subprocess.PIPE,
52
+ )
53
+ output: dict[str, str] = {}
54
+ for entry in result.stdout.decode("utf-8").split("\0"):
55
+ if "=" in entry:
56
+ key, value = entry.split("=", 1)
57
+ output[key] = value
58
+ return output
59
+
60
+
61
+ def _read_jsonl(path: Path) -> list[dict[str, Any]]:
62
+ return [
63
+ json.loads(line)
64
+ for line in path.read_text(encoding="utf-8").splitlines()
65
+ if line.strip()
66
+ ]
67
+
68
+
69
+ def _source_top24_session_ids(
70
+ debug_row: Mapping[str, Any], qid: str
71
+ ) -> list[str]:
72
+ trace = debug_row.get("source_top24_candidates")
73
+ pool = debug_row.get("source_candidate_pool_trace")
74
+ candidate_count = debug_row.get("source_candidate_count")
75
+ trace_k = debug_row.get("source_coverage_trace_k")
76
+ if type(candidate_count) is not int or candidate_count <= 0:
77
+ raise EvaluationError(f"{qid}: Source candidate count is invalid")
78
+ if trace_k != SOURCE_COVERAGE_TRACE_K:
79
+ raise EvaluationError(
80
+ f"{qid}: Source coverage trace contract is not Top{SOURCE_COVERAGE_TRACE_K}"
81
+ )
82
+ if not isinstance(trace, list) or not trace:
83
+ raise EvaluationError(f"{qid}: Source Top24 coverage trace is missing")
84
+ if not isinstance(pool, list):
85
+ raise EvaluationError(f"{qid}: Source candidate pool trace is missing")
86
+ expected_trace_count = min(SOURCE_COVERAGE_TRACE_K, candidate_count)
87
+ if len(trace) != expected_trace_count:
88
+ raise EvaluationError(
89
+ f"{qid}: Source Top24 coverage trace has {len(trace)} rows; "
90
+ f"expected {expected_trace_count}"
91
+ )
92
+ if len(pool) != candidate_count:
93
+ raise EvaluationError(
94
+ f"{qid}: Source candidate pool trace has {len(pool)} rows; "
95
+ f"expected {candidate_count}"
96
+ )
97
+
98
+ def validate_trace_rows(
99
+ rows: list[Any], *, label: str
100
+ ) -> list[tuple[str, str, int, int, int]]:
101
+ identities: list[tuple[str, str, int, int, int]] = []
102
+ locations: set[tuple[str, int, int, int]] = set()
103
+ for expected_rank, item in enumerate(rows, 1):
104
+ if not isinstance(item, Mapping):
105
+ raise EvaluationError(f"{qid}: {label} row must be an object")
106
+ rank = item.get("rank")
107
+ if type(rank) is not int or rank != expected_rank:
108
+ raise EvaluationError(f"{qid}: {label} is malformed")
109
+ candidate_id = str(item.get("candidate_id") or "").strip()
110
+ session_id = str(item.get("session_id") or "").strip()
111
+ coordinates = (
112
+ item.get("session_index"),
113
+ item.get("parent_chunk_index"),
114
+ item.get("subchunk_index"),
115
+ )
116
+ if (
117
+ not candidate_id
118
+ or not session_id
119
+ or any(type(value) is not int or value < 0 for value in coordinates)
120
+ ):
121
+ raise EvaluationError(f"{qid}: {label} has an invalid identity")
122
+ location = (session_id, *coordinates)
123
+ # Candidate IDs are content-derived and can collide when a dataset
124
+ # repeats the same session/chunk at a different session_index. The
125
+ # physical occurrence coordinates remain the authoritative identity.
126
+ if location in locations:
127
+ raise EvaluationError(f"{qid}: {label} contains duplicate evidence")
128
+ locations.add(location)
129
+ identities.append((candidate_id, session_id, *coordinates))
130
+ return identities
131
+
132
+ top_identities = validate_trace_rows(trace, label="Source Top24 coverage trace")
133
+ pool_identities = validate_trace_rows(pool, label="Source candidate pool trace")
134
+ if top_identities != pool_identities[:expected_trace_count]:
135
+ raise EvaluationError(
136
+ f"{qid}: Source Top24 trace is not the exact candidate-pool prefix"
137
+ )
138
+ session_ids: list[str] = []
139
+ for item in trace:
140
+ session_ids.append(str(item["session_id"]).strip())
141
+ return list(dict.fromkeys(session_ids))
142
+
143
+
144
+ def _validate_retrieval_debug_rows(
145
+ expected_qids: list[str], retrieval_debug: list[dict[str, Any]]
146
+ ) -> None:
147
+ actual_qids = [
148
+ str(row.get("question_id") or row.get("qid") or "").strip()
149
+ if isinstance(row, Mapping)
150
+ else ""
151
+ for row in retrieval_debug
152
+ ]
153
+ if actual_qids != expected_qids:
154
+ raise EvaluationError("retrieval debug rows do not match frozen qid order")
155
+ for qid, row in zip(expected_qids, retrieval_debug, strict=True):
156
+ _source_top24_session_ids(row, qid)
157
+
158
+
159
+ def _reference_rows(value: Any, source: str = "reference") -> list[dict[str, Any]]:
160
+ if isinstance(value, list):
161
+ raw_rows = value
162
+ elif isinstance(value, Mapping):
163
+ raw_rows: Any = None
164
+ for key in ("questions", "data", "rows", "references"):
165
+ if key in value:
166
+ raw_rows = value[key]
167
+ break
168
+ if raw_rows is None and value and all(isinstance(item, Mapping) for item in value.values()):
169
+ raw_rows = [dict(item, qid=key) for key, item in value.items()]
170
+ if raw_rows is None and ("qid" in value or "question_id" in value):
171
+ raw_rows = [value]
172
+ if not isinstance(raw_rows, list):
173
+ raise EvaluationError(f"{source} must contain a list of rows")
174
+ else:
175
+ raise EvaluationError(f"{source} must be a JSON object or list")
176
+
177
+ rows: list[dict[str, Any]] = []
178
+ seen: set[str] = set()
179
+ for index, raw_row in enumerate(raw_rows, 1):
180
+ if not isinstance(raw_row, Mapping):
181
+ raise EvaluationError(f"{source} row {index} must be an object")
182
+ qid_value = raw_row.get("qid")
183
+ question_id = raw_row.get("question_id")
184
+ if qid_value is not None and question_id is not None and str(qid_value).strip() != str(question_id).strip():
185
+ raise EvaluationError(f"{source} row {index} has disagreeing qid and question_id")
186
+ qid = str(qid_value or question_id or "").strip()
187
+ if not qid:
188
+ raise EvaluationError(f"{source} row {index} is missing qid")
189
+ if qid in seen:
190
+ raise EvaluationError(f"{source} contains duplicate qid {qid}")
191
+ session_ids = raw_row.get("answer_session_ids")
192
+ if not isinstance(session_ids, list) or not session_ids:
193
+ raise EvaluationError(f"{source} row {index} ({qid}) lacks answer_session_ids")
194
+ normalized_sessions = [str(item).strip() for item in session_ids]
195
+ if any(not item for item in normalized_sessions) or len(set(normalized_sessions)) != len(normalized_sessions):
196
+ raise EvaluationError(
197
+ f"{source} row {index} ({qid}) has blank or duplicate answer_session_ids"
198
+ )
199
+ gold_key = next(
200
+ (key for key in ("gold_answer", "answer", "expected_answer") if key in raw_row),
201
+ None,
202
+ )
203
+ if gold_key is None:
204
+ raise EvaluationError(f"{source} row {index} ({qid}) lacks a gold answer")
205
+ rows.append(
206
+ {
207
+ "qid": qid,
208
+ "gold_answer": raw_row[gold_key],
209
+ "answer_session_ids": normalized_sessions,
210
+ "question_type": str(
211
+ raw_row.get("question_type") or raw_row.get("type") or "unknown"
212
+ ).strip()
213
+ or "unknown",
214
+ }
215
+ )
216
+ seen.add(qid)
217
+ if not rows:
218
+ raise EvaluationError(f"{source} is empty")
219
+ return rows
220
+
221
+
222
+ def _read_reference(path: Path) -> list[dict[str, Any]]:
223
+ try:
224
+ value = json.loads(path.read_text(encoding="utf-8"))
225
+ except OSError as exc:
226
+ raise EvaluationError(f"cannot read reference data: {exc}") from exc
227
+ except json.JSONDecodeError as exc:
228
+ raise EvaluationError("reference data is not valid JSON") from exc
229
+ return _reference_rows(value)
230
+
231
+
232
+ def _build_evaluation_audit(
233
+ expected_qids: list[str],
234
+ reference: list[dict[str, Any]],
235
+ evidence: list[dict[str, Any]],
236
+ retrieval_debug: list[dict[str, Any]],
237
+ judged: list[dict[str, Any]],
238
+ ) -> list[dict[str, Any]]:
239
+ _validate_retrieval_debug_rows(expected_qids, retrieval_debug)
240
+ reference_by_qid = {row["qid"]: row for row in reference}
241
+ evidence_by_qid = {str(row.get("question_id") or row.get("qid") or "").strip(): row for row in evidence}
242
+ judged_by_qid = {str(row.get("question_id") or row.get("qid") or "").strip(): row for row in judged}
243
+ debug_by_qid = {
244
+ str(row.get("question_id") or row.get("qid") or "").strip(): row
245
+ for row in retrieval_debug
246
+ }
247
+ missing_reference = [qid for qid in expected_qids if qid not in reference_by_qid]
248
+ missing_evidence = [qid for qid in expected_qids if qid not in evidence_by_qid]
249
+ missing_judged = [qid for qid in expected_qids if qid not in judged_by_qid]
250
+ missing_debug = [qid for qid in expected_qids if qid not in debug_by_qid]
251
+ if missing_reference:
252
+ raise EvaluationError(f"reference is missing evaluated qids: {missing_reference[:10]}")
253
+ if missing_evidence:
254
+ raise EvaluationError(f"evidence is missing evaluated qids: {missing_evidence[:10]}")
255
+ if missing_judged:
256
+ raise EvaluationError(f"judge is missing evaluated qids: {missing_judged[:10]}")
257
+ if missing_debug:
258
+ raise EvaluationError(
259
+ f"retrieval debug is missing evaluated qids: {missing_debug[:10]}"
260
+ )
261
+
262
+ audit: list[dict[str, Any]] = []
263
+ for qid in expected_qids:
264
+ reference_row = reference_by_qid[qid]
265
+ selected = evidence_by_qid[qid].get("selected_session_ids")
266
+ if not isinstance(selected, list):
267
+ raise EvaluationError(f"{qid}: selected_session_ids must be a list")
268
+ selected_ids = [str(item).strip() for item in selected]
269
+ if any(not item for item in selected_ids) or len(set(selected_ids)) != len(selected_ids):
270
+ raise EvaluationError(f"{qid}: selected_session_ids has blank or duplicate values")
271
+ answer_ids = list(reference_row["answer_session_ids"])
272
+ source_top24_session_ids = _source_top24_session_ids(debug_by_qid[qid], qid)
273
+
274
+ def coverage(candidate_session_ids: list[str]) -> dict[str, Any]:
275
+ missing_ids = [
276
+ session_id
277
+ for session_id in answer_ids
278
+ if session_id not in set(candidate_session_ids)
279
+ ]
280
+ return {
281
+ "complete": not missing_ids,
282
+ "expected_session_count": len(answer_ids),
283
+ "selected_session_count": len(candidate_session_ids),
284
+ "covered_session_count": len(answer_ids) - len(missing_ids),
285
+ "hit_rate": round(
286
+ (len(answer_ids) - len(missing_ids)) / len(answer_ids), 6
287
+ ),
288
+ "missing_session_ids": missing_ids,
289
+ }
290
+
291
+ source_top24_coverage = coverage(source_top24_session_ids)
292
+ final_evidence_coverage = coverage(selected_ids)
293
+ official_label = official_judge_correct(judged_by_qid[qid])
294
+ audit.append(
295
+ {
296
+ "qid": qid,
297
+ "question_id": qid,
298
+ "question_type": reference_row["question_type"],
299
+ "gold_answer": reference_row["gold_answer"],
300
+ "gold_answer_session_ids": answer_ids,
301
+ "answer_session_ids": answer_ids,
302
+ "source_top24_session_ids": source_top24_session_ids,
303
+ "selected_session_ids": selected_ids,
304
+ "source_top24_coverage": source_top24_coverage,
305
+ "final_evidence_coverage": final_evidence_coverage,
306
+ "session_coverage": final_evidence_coverage,
307
+ "source_coverage_complete": source_top24_coverage["complete"],
308
+ "source_top24_coverage_complete": source_top24_coverage["complete"],
309
+ "final_evidence_coverage_complete": final_evidence_coverage["complete"],
310
+ "missing_source_top24_session_ids": source_top24_coverage[
311
+ "missing_session_ids"
312
+ ],
313
+ "missing_session_ids": final_evidence_coverage["missing_session_ids"],
314
+ "missing_answer_session_ids": final_evidence_coverage[
315
+ "missing_session_ids"
316
+ ],
317
+ "official_label": official_label,
318
+ "official_judge_label": official_label,
319
+ "official_correct": official_label,
320
+ }
321
+ )
322
+ return audit
323
+
324
+
325
+ def _write_jsonl(path: Path, rows: list[Mapping[str, Any]]) -> None:
326
+ temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}")
327
+ temporary.write_text(
328
+ "".join(json.dumps(dict(row), ensure_ascii=False, sort_keys=True) + "\n" for row in rows),
329
+ encoding="utf-8",
330
+ )
331
+ os.replace(temporary, path)
332
+
333
+
334
+ def _validate_answer_facing_evidence(rows: list[dict[str, Any]]) -> None:
335
+ forbidden_fields = {
336
+ "answer",
337
+ "gold_answer",
338
+ "expected_answer",
339
+ "answer_session_ids",
340
+ "labels",
341
+ "supervision",
342
+ }
343
+ for row in rows:
344
+ qid = str(row.get("question_id") or "").strip() or "<unknown>"
345
+ leaked = sorted(forbidden_fields.intersection(row))
346
+ if leaked:
347
+ raise EvaluationError(
348
+ f"{qid}: answer-facing evidence contains benchmark fields: {', '.join(leaked)}"
349
+ )
350
+
351
+
352
+ def official_judge_correct(row: Mapping[str, Any]) -> bool:
353
+ if str(row.get("error") or "").strip():
354
+ raise EvaluationError("official judge row contains an error")
355
+ autoeval = row.get("autoeval_label")
356
+ if not isinstance(autoeval, Mapping) or type(autoeval.get("label")) is not bool:
357
+ raise EvaluationError(
358
+ "official judge row lacks boolean autoeval_label.label"
359
+ )
360
+ return bool(autoeval["label"])
361
+
362
+
363
+ def _validate_judge_resume_rows(
364
+ answers: list[dict[str, Any]], judged: list[dict[str, Any]]
365
+ ) -> None:
366
+ answer_qids = [str(row.get("question_id") or "").strip() for row in answers]
367
+ judged_qids = [str(row.get("question_id") or "").strip() for row in judged]
368
+ if judged_qids != answer_qids[: len(judged_qids)]:
369
+ raise EvaluationError("existing judge rows are not a frozen answer prefix")
370
+ answer_by_qid = {str(row["question_id"]): row for row in answers}
371
+ for row in judged:
372
+ qid = str(row.get("question_id") or "").strip()
373
+ answer = answer_by_qid[qid]
374
+ for key in ("question", "hypothesis", "evidence_sha256", "answer_model"):
375
+ if row.get(key) != answer.get(key):
376
+ raise EvaluationError(
377
+ f"{qid}: existing judge row is bound to different answer evidence"
378
+ )
379
+
380
+
381
+ def _archive_failure_marker(run_dir: Path, retrieval_tag: str) -> None:
382
+ marker = run_dir / f"FAILED.answer_{retrieval_tag}"
383
+ if not marker.is_file():
384
+ return
385
+ history = run_dir / f"answer_{retrieval_tag}.failure_history.jsonl"
386
+ with history.open("a", encoding="utf-8") as handle:
387
+ for line in marker.read_text(encoding="utf-8").splitlines():
388
+ if line.strip():
389
+ handle.write(line.rstrip() + "\n")
390
+ marker.unlink()
391
+
392
+
393
+ def _record_failure(run_dir: Path, retrieval_tag: str, exc: Exception) -> None:
394
+ _archive_failure_marker(run_dir, retrieval_tag)
395
+ marker = run_dir / f"FAILED.answer_{retrieval_tag}"
396
+ temporary = marker.with_name(f".{marker.name}.tmp.{os.getpid()}")
397
+ temporary.write_text(
398
+ json.dumps({"at": _now(), "error": f"{exc.__class__.__name__}: {exc}"})
399
+ + "\n",
400
+ encoding="utf-8",
401
+ )
402
+ os.replace(temporary, marker)
403
+
404
+
405
+ def evaluate(args: argparse.Namespace) -> dict[str, Any]:
406
+ run_dir = args.run_dir.resolve()
407
+ retrieval_dir = run_dir / args.retrieval_tag
408
+ if not any(
409
+ (retrieval_dir / marker).is_file()
410
+ for marker in ("RETRIEVAL_COMPLETE", "COMPILE_COMPLETE")
411
+ ):
412
+ raise EvaluationError("retrieval or evidence compilation run is incomplete")
413
+ answer_dir = run_dir / f"answer_{args.retrieval_tag}"
414
+ resume = bool(getattr(args, "resume", False))
415
+ failure_marker = run_dir / f"FAILED.answer_{args.retrieval_tag}"
416
+ judge_path = run_dir / f"{args.retrieval_tag}.official_judge.jsonl"
417
+ if answer_dir.exists() and not resume:
418
+ raise EvaluationError(f"answer output already exists: {answer_dir}")
419
+ if failure_marker.exists() and not resume:
420
+ raise EvaluationError("prior answer failure exists; use explicit --resume")
421
+ if resume and not (
422
+ answer_dir.exists() or failure_marker.exists() or judge_path.exists()
423
+ ):
424
+ raise EvaluationError("resume requested but no prior answer state exists")
425
+ try:
426
+ qid_list = Path(
427
+ getattr(args, "qid_list", None) or run_dir / "qids.txt"
428
+ ).resolve()
429
+ if not qid_list.is_file():
430
+ raise EvaluationError(f"answer qid list is missing: {qid_list}")
431
+ expected_qids = [
432
+ line.strip()
433
+ for line in qid_list.read_text(encoding="utf-8").splitlines()
434
+ if line.strip()
435
+ ]
436
+ if not expected_qids or len(expected_qids) != len(set(expected_qids)):
437
+ raise EvaluationError("answer qid list is empty or contains duplicates")
438
+ baseline_judge = getattr(args, "baseline_judge", None)
439
+ if baseline_judge is not None:
440
+ baseline_judge = Path(baseline_judge).resolve()
441
+ if not baseline_judge.is_file():
442
+ raise EvaluationError(
443
+ f"frozen baseline judge is missing: {baseline_judge}"
444
+ )
445
+
446
+ try:
447
+ assert_production_answer_runner(args.answer_runner.resolve())
448
+ answer_facing_evidence = _read_jsonl(retrieval_dir / "evidence_windows.jsonl")
449
+ _validate_answer_facing_evidence(answer_facing_evidence)
450
+ route_report = validate_production_evidence(answer_facing_evidence)
451
+ except RoutePolicyError as exc:
452
+ raise EvaluationError(f"production route policy rejected the run: {exc}") from exc
453
+ if [item.get("question_id") for item in answer_facing_evidence] != expected_qids:
454
+ raise EvaluationError("evidence rows do not match frozen qid order")
455
+ retrieval_debug = _read_jsonl(retrieval_dir / "retrieval_debug.jsonl")
456
+ _validate_retrieval_debug_rows(expected_qids, retrieval_debug)
457
+
458
+ loaded = _load_environment(args.answer_env.resolve())
459
+ environment = {**os.environ, **loaded}
460
+ if environment.get("TMCRA_ANSWER_MODEL") != PRODUCTION_ANSWER_MODEL:
461
+ raise EvaluationError(
462
+ f"answer layer must remain fixed to {PRODUCTION_ANSWER_MODEL}"
463
+ )
464
+ _run(
465
+ [
466
+ sys.executable,
467
+ str(args.answer_runner.resolve()),
468
+ "--evidence",
469
+ str(retrieval_dir / "evidence_windows.jsonl"),
470
+ "--harness",
471
+ str(args.harness.resolve()),
472
+ "--out-dir",
473
+ str(answer_dir),
474
+ "--workers",
475
+ str(args.answer_workers),
476
+ "--attempts",
477
+ str(args.answer_attempts),
478
+ "--answer-window-limit",
479
+ str(args.answer_window_limit),
480
+ "--qid-list",
481
+ str(qid_list),
482
+ "--lane",
483
+ "production",
484
+ ],
485
+ (
486
+ _resume_log(run_dir, f"answer_{args.retrieval_tag}")
487
+ if resume
488
+ else run_dir / f"answer_{args.retrieval_tag}.log"
489
+ ),
490
+ environment,
491
+ )
492
+ answers = _read_jsonl(answer_dir / "answers.jsonl")
493
+ try:
494
+ validate_production_answers(answers)
495
+ except RoutePolicyError as exc:
496
+ raise EvaluationError(f"production route policy rejected answers: {exc}") from exc
497
+ if [item.get("question_id") for item in answers] != expected_qids:
498
+ raise EvaluationError("answer rows do not match frozen qid order")
499
+ judge_resume_options: list[str] = []
500
+ if judge_path.is_file():
501
+ if not resume:
502
+ raise EvaluationError("judge output already exists without --resume")
503
+ existing_judged = _read_jsonl(judge_path)
504
+ _validate_judge_resume_rows(answers, existing_judged)
505
+ judge_resume_options = ["--resume"]
506
+ _run(
507
+ [
508
+ sys.executable,
509
+ str(args.judge.resolve()),
510
+ "--metric-model",
511
+ "gpt-5.4",
512
+ "--hyp-file",
513
+ str(answer_dir / "answers.jsonl"),
514
+ "--ref-file",
515
+ str(args.data.resolve()),
516
+ "--result-file",
517
+ str(judge_path),
518
+ "--base-url",
519
+ environment["TMCRA_ANSWER_BASE_URL"],
520
+ "--api-key-env",
521
+ "TMCRA_ANSWER_API_KEY",
522
+ "--timeout",
523
+ "180",
524
+ "--max-retries",
525
+ "0",
526
+ *judge_resume_options,
527
+ "--quiet",
528
+ ],
529
+ (
530
+ _resume_log(run_dir, f"{args.retrieval_tag}.judge")
531
+ if resume
532
+ else run_dir / f"{args.retrieval_tag}.judge.log"
533
+ ),
534
+ environment,
535
+ )
536
+ judged = _read_jsonl(judge_path)
537
+ if [item.get("question_id") for item in answers] != expected_qids:
538
+ raise EvaluationError("answer rows do not match frozen qid order")
539
+ if [item.get("question_id") for item in judged] != expected_qids:
540
+ raise EvaluationError("judge rows do not match frozen qid order")
541
+ if any(item.get("answer_model") != PRODUCTION_ANSWER_MODEL for item in answers):
542
+ raise EvaluationError(f"answer output did not use {PRODUCTION_ANSWER_MODEL}")
543
+ labels = [official_judge_correct(item) for item in judged]
544
+ correct_count = sum(labels)
545
+ reference = _read_reference(args.data.resolve())
546
+ evaluation_audit = _build_evaluation_audit(
547
+ expected_qids,
548
+ reference,
549
+ answer_facing_evidence,
550
+ retrieval_debug,
551
+ judged,
552
+ )
553
+ evaluation_audit_path = answer_dir / "evaluation_audit.jsonl"
554
+ _write_jsonl(evaluation_audit_path, evaluation_audit)
555
+ source_complete_count = sum(
556
+ bool(row["source_top24_coverage_complete"]) for row in evaluation_audit
557
+ )
558
+ final_complete_count = sum(
559
+ bool(row["final_evidence_coverage_complete"]) for row in evaluation_audit
560
+ )
561
+ expected_gold_sessions = sum(
562
+ int(row["source_top24_coverage"]["expected_session_count"])
563
+ for row in evaluation_audit
564
+ )
565
+ source_covered_sessions = sum(
566
+ int(row["source_top24_coverage"]["covered_session_count"])
567
+ for row in evaluation_audit
568
+ )
569
+ final_covered_sessions = sum(
570
+ int(row["final_evidence_coverage"]["covered_session_count"])
571
+ for row in evaluation_audit
572
+ )
573
+ source_top24_hit_rate = (
574
+ source_covered_sessions / expected_gold_sessions
575
+ if expected_gold_sessions
576
+ else 0.0
577
+ )
578
+ final_evidence_hit_rate = (
579
+ final_covered_sessions / expected_gold_sessions
580
+ if expected_gold_sessions
581
+ else 0.0
582
+ )
583
+ accuracy = correct_count / len(judged) if judged else 0.0
584
+ evaluation_quality_gate_passed = (
585
+ source_top24_hit_rate == 1.0
586
+ and final_evidence_hit_rate >= 0.9
587
+ and accuracy >= 0.9
588
+ )
589
+ regression_gate_report: dict[str, Any] = {
590
+ "schema_version": "tmcra.v4.regression-gate-reference.1",
591
+ "status": "not_run",
592
+ "promotion_eligible": False,
593
+ "reason": "frozen baseline judge was not supplied",
594
+ }
595
+ if baseline_judge is not None:
596
+ regression_evidence_path = answer_dir / "source_top24_regression_evidence.jsonl"
597
+ _write_jsonl(
598
+ regression_evidence_path,
599
+ [
600
+ {
601
+ "qid": row["qid"],
602
+ "question_id": row["question_id"],
603
+ "selected_session_ids": row["source_top24_session_ids"],
604
+ }
605
+ for row in evaluation_audit
606
+ ],
607
+ )
608
+ regression_gate_report = evaluate_regression_gate(
609
+ baseline_judge,
610
+ judge_path,
611
+ regression_evidence_path,
612
+ args.data.resolve(),
613
+ cases_path=Path(
614
+ getattr(args, "regression_cases", DEFAULT_REGRESSION_CASES)
615
+ ).resolve(),
616
+ min_candidate_accuracy=None,
617
+ max_regressions=0,
618
+ min_improvements=0,
619
+ )
620
+ (answer_dir / "v4_regression_gate_report.json").write_text(
621
+ json.dumps(regression_gate_report, indent=2, sort_keys=True) + "\n",
622
+ encoding="utf-8",
623
+ )
624
+ regression_gate_passed = bool(
625
+ regression_gate_report.get("promotion_eligible")
626
+ )
627
+ promotion_gate_passed = (
628
+ evaluation_quality_gate_passed and regression_gate_passed
629
+ )
630
+ summary = {
631
+ "schema_version": "tmcra.v4.evaluation.3",
632
+ "status": "complete",
633
+ "route_policy_schema": POLICY_SCHEMA,
634
+ "lane": PRODUCTION_LANE,
635
+ "answer_protocol": PRODUCTION_ANSWER_PROTOCOL,
636
+ "promotion_eligible": promotion_gate_passed,
637
+ "route_report": route_report,
638
+ "retrieval_tag": args.retrieval_tag,
639
+ "evaluated_count": len(judged),
640
+ "answer_model": PRODUCTION_ANSWER_MODEL,
641
+ "answer_attempt_limit": int(args.answer_attempts),
642
+ "judge_retry_limit": 0,
643
+ "judge_label_field": "autoeval_label.label",
644
+ "correct_count": correct_count,
645
+ "accuracy": round(accuracy, 6),
646
+ "evaluation_quality_gate_passed": evaluation_quality_gate_passed,
647
+ "promotion_gate": {
648
+ "passed": promotion_gate_passed,
649
+ "evaluation_quality_gate_passed": evaluation_quality_gate_passed,
650
+ "zero_api_regression_gate_passed": regression_gate_passed,
651
+ "source_top24_gold_session_hit_rate_min": 1.0,
652
+ "final_evidence_gold_session_hit_rate_min": 0.9,
653
+ "official_accuracy_min": 0.9,
654
+ "max_regressions": 0,
655
+ },
656
+ "regression_gate": regression_gate_report,
657
+ "evaluation_audit": {
658
+ "path": str(evaluation_audit_path),
659
+ "row_count": len(evaluation_audit),
660
+ "reference_row_count": len(reference),
661
+ "reference_is_superset": len(reference) >= len(expected_qids)
662
+ and set(expected_qids).issubset({row["qid"] for row in reference}),
663
+ "source_top24_coverage_complete_count": source_complete_count,
664
+ "source_top24_coverage_complete_rate": round(
665
+ source_complete_count / len(evaluation_audit), 6
666
+ )
667
+ if evaluation_audit
668
+ else 0.0,
669
+ "source_top24_gold_session_hit_rate": round(
670
+ source_top24_hit_rate, 6
671
+ ),
672
+ "final_evidence_coverage_complete_count": final_complete_count,
673
+ "final_evidence_coverage_complete_rate": round(
674
+ final_complete_count / len(evaluation_audit), 6
675
+ )
676
+ if evaluation_audit
677
+ else 0.0,
678
+ "final_evidence_gold_session_hit_rate": round(
679
+ final_evidence_hit_rate, 6
680
+ ),
681
+ "session_coverage_complete_count": final_complete_count,
682
+ "session_coverage_complete_rate": round(
683
+ final_complete_count / len(evaluation_audit), 6
684
+ )
685
+ if evaluation_audit
686
+ else 0.0,
687
+ },
688
+ }
689
+ (answer_dir / "v4_evaluation_report.json").write_text(
690
+ json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
691
+ )
692
+ (answer_dir / "ANSWER_COMPLETE").write_text(_now() + "\n", encoding="utf-8")
693
+ _archive_failure_marker(run_dir, args.retrieval_tag)
694
+ return summary
695
+ except Exception as exc:
696
+ _record_failure(run_dir, args.retrieval_tag, exc)
697
+ raise
698
+
699
+
700
+ def main() -> int:
701
+ parser = argparse.ArgumentParser(description="Evaluate a completed TMCRA V4 retrieval run")
702
+ parser.add_argument("--run-dir", type=Path, required=True)
703
+ parser.add_argument("--retrieval-tag", default="retrieval_1")
704
+ parser.add_argument("--resume", action="store_true")
705
+ parser.add_argument("--answer-env", type=Path, default=DEFAULT_ANSWER_ENV)
706
+ parser.add_argument("--answer-runner", type=Path, default=DEFAULT_ANSWER_RUNNER)
707
+ parser.add_argument("--harness", type=Path, default=DEFAULT_HARNESS)
708
+ parser.add_argument("--judge", type=Path, default=DEFAULT_JUDGE)
709
+ parser.add_argument("--data", type=Path, default=DEFAULT_DATA)
710
+ parser.add_argument("--qid-list", type=Path)
711
+ parser.add_argument("--answer-workers", type=int, default=1)
712
+ parser.add_argument("--answer-attempts", type=int, choices=(1, 2, 3), default=1)
713
+ parser.add_argument("--answer-window-limit", type=int, default=8)
714
+ parser.add_argument("--baseline-judge", type=Path)
715
+ parser.add_argument(
716
+ "--regression-cases", type=Path, default=DEFAULT_REGRESSION_CASES
717
+ )
718
+ args = parser.parse_args()
719
+ summary = evaluate(args)
720
+ print(json.dumps(summary, sort_keys=True))
721
+ return 0
722
+
723
+
724
+ if __name__ == "__main__":
725
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_gpt54_answers.py ADDED
@@ -0,0 +1,847 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import importlib.util
7
+ import json
8
+ import os
9
+ import sys
10
+ import threading
11
+ import time
12
+ from concurrent.futures import ThreadPoolExecutor, as_completed
13
+ from pathlib import Path
14
+ from typing import Any, Mapping
15
+
16
+ from tmcra_v4_evidence_operations import (
17
+ ANSWER_SCHEMA,
18
+ PACKET_SCHEMA,
19
+ EvidenceOperationError,
20
+ validate_evidence_bound_answer,
21
+ )
22
+ from tmcra_v4_semantic_evidence import (
23
+ SEMANTIC_ANSWER_SCHEMA,
24
+ SEMANTIC_PACKET_SCHEMA,
25
+ SemanticEvidenceError,
26
+ normalize_semantic_answer_output,
27
+ semantic_packet_to_advisory_packet,
28
+ validate_semantic_answer,
29
+ )
30
+ from tmcra_v4_route_policy import (
31
+ PRODUCTION_ANSWER_PROTOCOL,
32
+ RoutePolicyError,
33
+ SHADOW_ANSWER_PROTOCOL,
34
+ validate_production_answers,
35
+ validate_production_evidence,
36
+ )
37
+
38
+
39
+ FORBIDDEN_EVALUATION_FIELDS = {
40
+ "answer",
41
+ "gold_answer",
42
+ "answer_session_ids",
43
+ "labels",
44
+ "supervision",
45
+ }
46
+ USER_MEMORY_CUE_LIMIT = 32
47
+ USER_MEMORY_CUE_TEXT_LIMIT = 640
48
+
49
+
50
+ def clean_text(value: Any) -> str:
51
+ return " ".join(str(value or "").split())
52
+
53
+
54
+ def fixed_answer_environment(answer_window_limit: int) -> dict[str, str]:
55
+ return {
56
+ "TMCRA_ANSWER_WINDOW_PLANNER_MODE": "off",
57
+ "TMCRA_ANSWER_WINDOW_PLANNER_SEMANTIC_MODE": "off",
58
+ "TMCRA_LLM_EVIDENCE_SELECTOR_MODE": "off",
59
+ "TMCRA_EVIDENCE_UNIT_PLANNER_MODE": "off",
60
+ "TMCRA_UNIFIED_OPERATION_PLANNER_MODE": "off",
61
+ "TMCRA_LLM_CHANNEL_PLANNER_MODE": "off",
62
+ "TMCRA_FINAL_ANSWER_SURFACE_CHANNEL_MODE": "off",
63
+ "TMCRA_ANSWER_REASONING_ORGANIZER": "off",
64
+ "TMCRA_LLM_CACHE_MODE": "off",
65
+ "TMCRA_LME_FORCE_ANSWER": "1",
66
+ "TMCRA_HTTP_JSON_MAX_ATTEMPTS": "1",
67
+ "TMCRA_ANSWER_EVIDENCE_WINDOW_LIMIT": str(max(1, answer_window_limit)),
68
+ }
69
+
70
+
71
+ def read_jsonl(path: Path) -> list[dict[str, Any]]:
72
+ rows: list[dict[str, Any]] = []
73
+ with path.open("r", encoding="utf-8", errors="strict") as handle:
74
+ for line_no, line in enumerate(handle, start=1):
75
+ if not line.strip():
76
+ continue
77
+ value = json.loads(line)
78
+ if not isinstance(value, dict):
79
+ raise RuntimeError(f"row is not an object at {path}:{line_no}")
80
+ rows.append(value)
81
+ if not rows:
82
+ raise RuntimeError(f"no rows: {path}")
83
+ return rows
84
+
85
+
86
+ def evidence_sha256(row: Mapping[str, Any]) -> str:
87
+ return hashlib.sha256(
88
+ json.dumps(
89
+ dict(row), ensure_ascii=False, sort_keys=True, separators=(",", ":")
90
+ ).encode("utf-8")
91
+ ).hexdigest()
92
+
93
+
94
+ def atomic_write_jsonl(path: Path, rows: list[Mapping[str, Any]]) -> None:
95
+ temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}.{time.time_ns()}")
96
+ temporary.write_text(
97
+ "".join(json.dumps(dict(row), ensure_ascii=False, sort_keys=True) + "\n" for row in rows),
98
+ encoding="utf-8",
99
+ )
100
+ os.replace(temporary, path)
101
+
102
+
103
+ def bind_existing_answers_to_evidence(
104
+ existing: dict[str, dict[str, Any]],
105
+ rows: list[dict[str, Any]],
106
+ ) -> tuple[dict[str, str], int]:
107
+ evidence_by_qid = {clean_text(row.get("question_id")): row for row in rows}
108
+ if not all(evidence_by_qid) or len(evidence_by_qid) != len(rows):
109
+ raise RuntimeError("answer evidence contains missing or duplicate question_id")
110
+ hashes = {qid: evidence_sha256(row) for qid, row in evidence_by_qid.items()}
111
+ backfills = 0
112
+ for qid, row in existing.items():
113
+ if qid not in evidence_by_qid:
114
+ continue
115
+ expected = evidence_by_qid[qid]
116
+ expected_hash = hashes[qid]
117
+ observed_hash = clean_text(row.get("evidence_sha256"))
118
+ if observed_hash and observed_hash != expected_hash:
119
+ raise RuntimeError(f"{qid}: existing answer is bound to different retrieval evidence")
120
+ if not observed_hash:
121
+ if (
122
+ clean_text(row.get("question")) != clean_text(expected.get("question"))
123
+ or list(row.get("selected_session_ids") or [])
124
+ != list(expected.get("selected_session_ids") or [])
125
+ or clean_text(row.get("answer_model")) != "gpt-5.4"
126
+ ):
127
+ raise RuntimeError(f"{qid}: legacy answer cannot be bound to current evidence")
128
+ row["evidence_sha256"] = expected_hash
129
+ backfills += 1
130
+ return hashes, backfills
131
+
132
+
133
+ def ordered_completed_answers(
134
+ existing: Mapping[str, Mapping[str, Any]],
135
+ rows: list[Mapping[str, Any]],
136
+ ) -> list[dict[str, Any]]:
137
+ qids = [clean_text(row.get("question_id")) for row in rows]
138
+ if not all(qids) or len(qids) != len(set(qids)):
139
+ raise RuntimeError("answer evidence contains missing or duplicate question_id")
140
+ unexpected = sorted(set(existing) - set(qids))
141
+ if unexpected:
142
+ raise RuntimeError(
143
+ f"existing answer output contains qids outside the frozen request: {unexpected[:10]}"
144
+ )
145
+ return [dict(existing[qid]) for qid in qids if qid in existing]
146
+
147
+
148
+ def load_harness(path: Path):
149
+ module_name = "tmcra_v3_answer_harness"
150
+ spec = importlib.util.spec_from_file_location(module_name, path)
151
+ if spec is None or spec.loader is None:
152
+ raise RuntimeError(f"cannot import answer harness: {path}")
153
+ module = importlib.util.module_from_spec(spec)
154
+ sys.modules[module_name] = module
155
+ spec.loader.exec_module(module)
156
+ return module
157
+
158
+
159
+ def reject_evaluation_fields(rows: list[dict[str, Any]]) -> None:
160
+ for row in rows:
161
+ present = sorted(FORBIDDEN_EVALUATION_FIELDS.intersection(row))
162
+ if present:
163
+ qid = clean_text(row.get("question_id")) or "<missing-qid>"
164
+ raise RuntimeError(
165
+ f"{qid}: answer input contains evaluation-only fields: {present}"
166
+ )
167
+
168
+
169
+ def render_hierarchical_evidence(window: Mapping[str, Any]) -> str:
170
+ source = clean_text(window.get("text"))
171
+ if not source:
172
+ raise RuntimeError("evidence window has empty source text")
173
+ blocks: list[str] = []
174
+ for context in list(window.get("memory_contexts") or []):
175
+ if not isinstance(context, Mapping):
176
+ raise RuntimeError("memory_context must be an object")
177
+ claim = clean_text(context.get("claim_text"))
178
+ slot = clean_text(context.get("canonical_slot"))
179
+ role = clean_text(context.get("role"))
180
+ if not claim or not slot or not role:
181
+ raise RuntimeError("memory_context lacks role, canonical_slot, or claim_text")
182
+ blocks.append(f"[TMCRA slow memory role={role} slot={slot}]\n{claim}")
183
+ for attachment in list(window.get("attachments") or []):
184
+ if not isinstance(attachment, Mapping):
185
+ raise RuntimeError("memory attachment must be an object")
186
+ role = clean_text(attachment.get("role"))
187
+ slot = clean_text(attachment.get("canonical_slot"))
188
+ if role == "context_only":
189
+ summary = clean_text(attachment.get("summary"))
190
+ if not summary or not slot:
191
+ raise RuntimeError("slow context attachment lacks slot or summary")
192
+ blocks.append(
193
+ f"[TMCRA slow memory role=context_only slot={slot}]\n{summary}"
194
+ )
195
+ elif role in {"fast_context", "override"}:
196
+ override_text = clean_text(attachment.get("text"))
197
+ if not override_text or not slot:
198
+ raise RuntimeError("fast memory attachment lacks slot or text")
199
+ if role == "override":
200
+ blocks.append(
201
+ "[TMCRA fast memory role=override "
202
+ f"precedence=newer_fast_evidence slot={slot}]\n"
203
+ "This newer fast-layer evidence overrides conflicting slow-memory "
204
+ f"claims.\n{override_text}"
205
+ )
206
+ else:
207
+ blocks.append(
208
+ f"[TMCRA fast memory role=context slot={slot}]\n{override_text}"
209
+ )
210
+ else:
211
+ raise RuntimeError(f"unsupported memory attachment role: {role!r}")
212
+ blocks.append("[TMCRA immutable source evidence]\n" + source)
213
+ return "\n\n".join(blocks)
214
+
215
+
216
+ def build_user_memory_cues(
217
+ evidence: list[Mapping[str, Any]],
218
+ ) -> list[dict[str, Any]]:
219
+ """Build a compact index of durable claims and user-authored evidence.
220
+
221
+ The index only points into the immutable Source reservoir. Round-robin
222
+ emission prevents an early, verbose source from hiding cues in later
223
+ retrieval windows.
224
+ """
225
+ cues_by_source: list[list[dict[str, Any]]] = []
226
+ for item in evidence:
227
+ evidence_id = clean_text(item.get("evidence_id"))
228
+ source_record_id = clean_text(item.get("source_record_id"))
229
+ source_cues: list[dict[str, Any]] = []
230
+
231
+ def add_cue(
232
+ kind: str,
233
+ text: Any,
234
+ *,
235
+ record_id: str = source_record_id,
236
+ historical_date: Any = None,
237
+ timestamp: Any = None,
238
+ ) -> None:
239
+ normalized = clean_text(text)
240
+ if not normalized:
241
+ return
242
+ source_cues.append(
243
+ {
244
+ "evidence_id": evidence_id,
245
+ "kind": kind,
246
+ "source_record_id": record_id,
247
+ "historical_date": clean_text(
248
+ historical_date
249
+ if historical_date is not None
250
+ else item.get("historical_date")
251
+ ),
252
+ "timestamp": clean_text(
253
+ timestamp if timestamp is not None else item.get("timestamp")
254
+ ),
255
+ "text": normalized[:USER_MEMORY_CUE_TEXT_LIMIT],
256
+ "text_truncated": len(normalized) > USER_MEMORY_CUE_TEXT_LIMIT,
257
+ }
258
+ )
259
+
260
+ for context in list(item.get("memory_contexts") or []):
261
+ if isinstance(context, Mapping):
262
+ add_cue(
263
+ "slow_claim",
264
+ context.get("claim_text") or context.get("capsule_summary"),
265
+ )
266
+ for attachment in list(item.get("attachments") or []):
267
+ if isinstance(attachment, Mapping):
268
+ add_cue(
269
+ "fast_memory",
270
+ attachment.get("text") or attachment.get("summary"),
271
+ )
272
+ if clean_text(item.get("message_role")).lower() == "user":
273
+ add_cue("user_source", item.get("text"))
274
+ for context in list(item.get("source_group_context") or []):
275
+ if not isinstance(context, Mapping):
276
+ continue
277
+ role = clean_text(
278
+ context.get("message_role") or context.get("role")
279
+ ).lower()
280
+ if role == "user":
281
+ add_cue(
282
+ "user_context",
283
+ context.get("text"),
284
+ record_id=clean_text(context.get("source_record_id")),
285
+ historical_date=context.get("historical_date"),
286
+ timestamp=context.get("timestamp"),
287
+ )
288
+ cues_by_source.append(source_cues)
289
+
290
+ output: list[dict[str, Any]] = []
291
+ seen_text: set[str] = set()
292
+ max_source_cues = max((len(items) for items in cues_by_source), default=0)
293
+ for cue_index in range(max_source_cues):
294
+ for source_cues in cues_by_source:
295
+ if cue_index >= len(source_cues):
296
+ continue
297
+ cue = source_cues[cue_index]
298
+ fingerprint = clean_text(cue["text"]).casefold()
299
+ if fingerprint in seen_text:
300
+ continue
301
+ seen_text.add(fingerprint)
302
+ output.append({"cue_order": len(output) + 1, **cue})
303
+ if len(output) >= USER_MEMORY_CUE_LIMIT:
304
+ return output
305
+ return output
306
+
307
+
308
+ EVIDENCE_OPERATION_ANSWER_PROMPT = f"""You answer one memory question from a compiled, immutable evidence packet.
309
+ Return exactly one JSON object and no prose using this schema:
310
+ {{"schema_version":"{ANSWER_SCHEMA}","answerability":"sufficient|insufficient","claims":[{{"claim_id":"C01","text":"one factual claim","origin":"memory_fact|memory_inference|memory_derived|query_context|model_knowledge","support_ids":["E01"],"computation_ids":["O01"]}}],"missing_requirements":["R01"],"answer":"short complete answer"}}
311
+
312
+ Read task_intent.output_origin before answering. task_intent and planner_source_groups are routing hints, not truth or answerability decisions. requirement_catalog lists the response requirement IDs but deliberately carries no planner missing/satisfied judgment. Determine support or absence yourself from the complete Source reservoir. user_memory_cues is a deterministic navigation index of slow-memory claims and user-authored Source or source-group context; it is not separate proof, so inspect the referenced Source before binding a remembered claim. For memory_conditioned_generation, first recover the user's grounded preferences or constraints and then generate the requested recommendation or advice; a ready-made historical answer is not required. Carry forward relevant prior choices, owned tools, resources, genres, constraints, and already-adopted solutions instead of defaulting to generic advice. When several relevant cues differ in specificity, preserve the narrower grounded detail rather than replacing it with a broad category. The final answer itself must explicitly show how at least one specific remembered constraint shaped the response; do not leave personalization only inside the claims array. Do not declare the task insufficient merely because the exact requested recommendation was not historical when grounded memory can personalize a new answer and model knowledge can fulfill the current target. Query context never requires Source evidence IDs. Model knowledge does not require Source evidence IDs when it is allowed for memory-conditioned generation. External-tool premises remain unresolved until a tool result exists.
313
+ Source evidence is ordered by frozen retrieval priority and then by source order within each selected session; source_order is not an event date. Each Source item carries its immutable historical_date and timestamp when available. Resolve words such as today and yesterday against that Source item's timestamp, never against the question date. Inspect source_group_context as neighboring source context and use each neighbor's own timestamp. Before answering, internally build an exact evidence ledger for the requested entity: list every matching event or item, its event date or stated relative time, state (actual, planned, hypothetical, or mentioned), and Source ID. Deduplicate repeated mentions of the same event or item before counting. For latest/earliest/order/duration questions, compare event times grounded in Source text plus its timestamp, not source_order or database session_index. If a when, before, or after clause names another remembered event, treat that event as a candidate temporal endpoint and compare both Source dates before using the question date. Use the question date only when the wording actually asks for distance from the query time.
314
+ Each Source item may include memory_contexts and attachments. A slow_context is the slow graph's durable interpretation and must retain its capsule, claim, and provenance identity. A fast_context is an atomic current-memory interpretation. An attachment marked override with precedence=newer_fast_evidence has controller-verified precedence over a conflicting slow claim in the same canonical slot. These graph views organize meaning but do not replace Source proof: support final remembered claims with the Source evidence ID carrying the context, and inspect its immutable text before relying on it.
315
+ Use only supplied Source evidence IDs and verified computation IDs. memory_fact claims require direct support_ids. memory_inference claims require support_ids for every Source premise and represent source-grounded comparison or synthesis that has no verified operation result. memory_derived claims require computation_ids. query_context claims may have empty binding lists in every task when they state only information supplied by the current question, such as the question date or arithmetic derived solely from it. Only memory_conditioned_generation may use model_knowledge claims with empty binding lists; label their origin explicitly, and include at least one memory_fact, memory_inference, or memory_derived claim showing how remembered constraints shaped the answer. Never label remembered information as query_context or model_knowledge. planner_operation_candidates expose no result because their operand selection has not been semantically certified; independently inspect their Source supports. lexical_anchor_ids are deterministic search hints, not proof. Never transfer values between different entities, people, places, events, or activities. Preserve requested temporal order, distinct values, totals, and update precedence.
316
+ For preferences, use the specific historical experiences in Source evidence. Treat named entities and activities exactly: a shorter, broader, or related term is not the requested multiword entity unless the source explicitly says they are equivalent. For false premises or entity mismatch, explicitly state that the requested target is not established, then optionally state related counterevidence. For count/list questions, enumerate distinct matching items before producing the total. For changing facts, apply latest update precedence using event time and source-group context. A sufficient answer must cover every requirement and have no missing_requirements. An insufficient answer must name supplied requirement IDs that remain unresolved after inspecting all Source evidence. Never use hidden supervision or outside facts for remembered claims."""
317
+
318
+
319
+ SEMANTIC_EVIDENCE_ANSWER_PROMPT = f"""Execute one verified memory task contract.
320
+ Return exactly one JSON object and no prose using this schema:
321
+ {{"schema_version":"{SEMANTIC_ANSWER_SCHEMA}","response_status":"answered|partial|answered_with_context|requires_tool|clarification|not_answerable","claims":[{{"claim_id":"C01","text":"one final claim","origin":"memory_fact|memory_derived|model_knowledge|tool_result","premise_ids":["P01"],"source_claim_ids":["M01"],"computation_ids":["O01"]}}],"unresolved_premise_ids":[],"answer":"concise complete response"}}
322
+
323
+ The task contract defines the requested target and output shape. The answerability certificate is authoritative; do not redo retrieval, invent missing memory, or change its answer policy.
324
+ For memory_fact claims, use only Source claim IDs bound to the cited premises. For memory_derived claims, use only completed computation IDs bound to the cited premises. Preserve entity, time, polarity, modality, cardinality, ordering, and update state.
325
+ When the policy is answer_with_memory_context, treat remembered claims as constraints or background. You may produce model_knowledge claims needed to fulfill the current request, but never describe them as remembered facts and never attach Source claim IDs to them.
326
+ When the policy is answer_partial, answer only the supported portion, explicitly state the unresolved limitation, use response_status partial, and preserve the certificate's unresolved premise IDs. For a memory-conditioned task, the partial answer must explicitly use at least one bound memory claim rather than falling back to a generic response. Never turn a partial count or other incomplete derived aggregate into a final total.
327
+ When the policy is call_tool, clarify, or abstain, use the matching response_status and explain the unresolved need without fabricating an answer. unresolved_premise_ids must exactly match the certificate.
328
+ Do not cite raw evidence IDs directly. Do not use benchmark labels or hidden supervision."""
329
+
330
+
331
+ def render_compiled_packet(packet: Mapping[str, Any]) -> dict[str, Any]:
332
+ if packet.get("schema_version") != PACKET_SCHEMA:
333
+ raise RuntimeError("compiled evidence packet schema is invalid")
334
+ evidence = list(packet.get("raw_evidence_reservoir") or [])
335
+ evidence_ids = {clean_text(item.get("evidence_id")) for item in evidence}
336
+ if not evidence or not all(evidence_ids):
337
+ raise RuntimeError("compiled evidence packet has no source reservoir")
338
+ bundles = []
339
+ for bundle in list(packet.get("evidence_bundles") or []):
340
+ ids = list(bundle.get("evidence_ids") or [])
341
+ if any(item not in evidence_ids for item in ids):
342
+ raise RuntimeError("compiled evidence bundle cites unknown source")
343
+ bundles.append(
344
+ {
345
+ "bundle_id": bundle.get("bundle_id"),
346
+ "role": bundle.get("role"),
347
+ "evidence_ids": ids,
348
+ }
349
+ )
350
+ question_contract = dict(packet.get("question_contract") or {})
351
+ task_contract = packet.get("task_contract")
352
+ if not isinstance(task_contract, Mapping):
353
+ task_contract = question_contract.get("task_contract")
354
+ question_contract.pop("task_contract", None)
355
+ requirements = list(packet.get("requirement_coverage") or [])
356
+ premise_by_id = {
357
+ clean_text(item.get("premise_id")): item
358
+ for item in (
359
+ task_contract.get("premises") or []
360
+ if isinstance(task_contract, Mapping)
361
+ else []
362
+ )
363
+ if isinstance(item, Mapping)
364
+ }
365
+ output_origin = (
366
+ clean_text(task_contract.get("output_origin"))
367
+ if isinstance(task_contract, Mapping)
368
+ else clean_text(question_contract.get("output_origin"))
369
+ )
370
+ requirement_catalog = []
371
+ bound_constraints = []
372
+ for requirement in requirements:
373
+ requirement_id = clean_text(requirement.get("requirement_id"))
374
+ evidence_ids_for_requirement = [
375
+ item
376
+ for item in list(requirement.get("evidence_ids") or [])
377
+ if item in evidence_ids
378
+ ]
379
+ premise = premise_by_id.get(requirement_id) or {}
380
+ description = clean_text(requirement.get("description"))
381
+ source_kind = clean_text(premise.get("source")) or "memory"
382
+ catalog_item = {
383
+ "requirement_id": requirement_id,
384
+ "description": (
385
+ "remembered personalization constraints; inspect all Source evidence"
386
+ if output_origin == "memory_conditioned_generation"
387
+ and source_kind == "memory"
388
+ and not evidence_ids_for_requirement
389
+ else description
390
+ ),
391
+ "source_kind": source_kind,
392
+ }
393
+ if evidence_ids_for_requirement:
394
+ catalog_item["candidate_source_ids"] = evidence_ids_for_requirement
395
+ requirement_catalog.append(catalog_item)
396
+ if evidence_ids_for_requirement and isinstance(premise, Mapping):
397
+ bound_constraints.append(
398
+ {
399
+ "premise_id": requirement_id,
400
+ "description": clean_text(premise.get("description")) or description,
401
+ "source_kind": source_kind,
402
+ "grounded_constraints": list(premise.get("grounded_constraints") or []),
403
+ "candidate_source_ids": evidence_ids_for_requirement,
404
+ }
405
+ )
406
+ operation_results = list(packet.get("operation_results") or [])
407
+ verified_operations = [
408
+ dict(item)
409
+ for item in operation_results
410
+ if item.get("status") == "completed"
411
+ and item.get("answer_authoritative") is True
412
+ ]
413
+ planner_operation_candidates = [
414
+ {
415
+ "operation_id": item.get("operation_id"),
416
+ "operation_type": item.get("operation_type"),
417
+ "support_ids": list(item.get("support_ids") or []),
418
+ "semantic_status": "uncertified_operand_selection",
419
+ }
420
+ for item in operation_results
421
+ if item.get("status") == "completed"
422
+ and item.get("answer_authoritative") is not True
423
+ ]
424
+ task_intent: dict[str, Any] = {"output_origin": output_origin}
425
+ if isinstance(task_contract, Mapping):
426
+ task_intent["target"] = dict(task_contract.get("target") or {})
427
+ task_intent["output"] = dict(task_contract.get("output") or {})
428
+ task_intent["bound_constraints"] = bound_constraints
429
+ rendered = {
430
+ "answer_view_schema": "tmcra.v4.answer-evidence-view.5",
431
+ "question_contract": question_contract,
432
+ "task_intent": task_intent,
433
+ "user_memory_cues": build_user_memory_cues(evidence),
434
+ "requirement_catalog": requirement_catalog,
435
+ "verified_operation_results": verified_operations,
436
+ "planner_operation_candidates": planner_operation_candidates,
437
+ "planner_source_groups": bundles,
438
+ "lexical_anchor_ids": list(packet.get("lexical_anchor_ids") or []),
439
+ "source_evidence": [
440
+ {
441
+ "source_order": index,
442
+ "evidence_id": item.get("evidence_id"),
443
+ "session_id": item.get("session_id"),
444
+ "session_index": item.get("session_index"),
445
+ "parent_chunk_index": item.get("parent_chunk_index"),
446
+ "historical_date": item.get("historical_date"),
447
+ "timestamp": item.get("timestamp"),
448
+ "message_role": item.get("message_role"),
449
+ "question_overlap_terms": item.get("question_overlap_terms") or [],
450
+ "question_overlap_score": int(item.get("question_overlap_score", 0)),
451
+ "text": item.get("text"),
452
+ "source_group_id": item.get("source_group_id"),
453
+ "source_group_context": list(item.get("source_group_context") or []),
454
+ "memory_contexts": list(item.get("memory_contexts") or []),
455
+ "attachments": list(item.get("attachments") or []),
456
+ "provenance": list(item.get("provenance") or []),
457
+ }
458
+ for index, item in enumerate(evidence, start=1)
459
+ ],
460
+ }
461
+ if isinstance(packet.get("semantic_advisory"), Mapping):
462
+ rendered["semantic_advisory"] = dict(packet["semantic_advisory"])
463
+ return rendered
464
+
465
+
466
+ def parse_operation_bound_answer(raw: str, packet: Mapping[str, Any]) -> dict[str, Any]:
467
+ try:
468
+ value = json.loads(raw)
469
+ except json.JSONDecodeError as exc:
470
+ raise RuntimeError(f"operation-bound answer is not strict JSON: {exc}") from exc
471
+ try:
472
+ return validate_evidence_bound_answer(value, packet)
473
+ except EvidenceOperationError as exc:
474
+ raise RuntimeError(f"operation-bound answer failed validation: {exc}") from exc
475
+
476
+
477
+ def render_semantic_packet(packet: Mapping[str, Any]) -> dict[str, Any]:
478
+ if packet.get("schema_version") != SEMANTIC_PACKET_SCHEMA:
479
+ raise RuntimeError("semantic evidence packet schema is invalid")
480
+ contract = packet.get("task_contract")
481
+ certificate = packet.get("answerability_certificate")
482
+ if not isinstance(contract, Mapping) or not isinstance(certificate, Mapping):
483
+ raise RuntimeError("semantic evidence packet lacks contract or certificate")
484
+ premise_by_id = {
485
+ clean_text(item.get("premise_id")): item
486
+ for item in contract.get("premises") or []
487
+ if isinstance(item, Mapping)
488
+ }
489
+ claims_by_id = {
490
+ clean_text(item.get("claim_id")): item
491
+ for item in packet.get("claim_ledger") or []
492
+ if isinstance(item, Mapping)
493
+ }
494
+ evidence_by_id = {
495
+ clean_text(item.get("evidence_id")): item
496
+ for item in packet.get("raw_evidence_reservoir") or []
497
+ if isinstance(item, Mapping)
498
+ }
499
+ operation_by_id = {
500
+ clean_text(item.get("operation_id")): item
501
+ for item in (packet.get("resolution_program") or {}).get("operation_results") or []
502
+ if isinstance(item, Mapping)
503
+ }
504
+ resolved_premises = []
505
+ for binding in packet.get("premise_bindings") or []:
506
+ if not isinstance(binding, Mapping):
507
+ raise RuntimeError("semantic premise binding must be an object")
508
+ premise_id = clean_text(binding.get("premise_id"))
509
+ if premise_id not in premise_by_id:
510
+ raise RuntimeError("semantic premise binding cites an unknown premise")
511
+ claim_ids = list(binding.get("claim_ids") or [])
512
+ evidence_ids = list(binding.get("evidence_ids") or [])
513
+ operation_ids = list(binding.get("operation_ids") or [])
514
+ if any(item not in claims_by_id for item in claim_ids):
515
+ raise RuntimeError("semantic premise binding cites an unknown Source claim")
516
+ if any(item not in evidence_by_id for item in evidence_ids):
517
+ raise RuntimeError("semantic premise binding cites unknown evidence")
518
+ if any(item not in operation_by_id for item in operation_ids):
519
+ raise RuntimeError("semantic premise binding cites an unknown operation")
520
+ resolved_premises.append(
521
+ {
522
+ "premise": premise_by_id[premise_id],
523
+ "coverage": binding.get("coverage"),
524
+ "relation": binding.get("relation"),
525
+ "source_claims": [claims_by_id[item] for item in claim_ids],
526
+ "operation_results": [operation_by_id[item] for item in operation_ids],
527
+ "source_evidence": [
528
+ {
529
+ "evidence_id": item,
530
+ "session_id": evidence_by_id[item].get("session_id"),
531
+ "session_index": evidence_by_id[item].get("session_index"),
532
+ "text": evidence_by_id[item].get("text"),
533
+ }
534
+ for item in evidence_ids
535
+ ],
536
+ }
537
+ )
538
+ return {
539
+ "query_context": dict(packet.get("query_context") or {}),
540
+ "task_contract": dict(contract),
541
+ "answerability_certificate": dict(certificate),
542
+ "resolved_premises": resolved_premises,
543
+ "non_memory_premises": [
544
+ premise
545
+ for premise in contract.get("premises") or []
546
+ if isinstance(premise, Mapping) and premise.get("source") != "memory"
547
+ ],
548
+ "target_operation_results": [
549
+ operation_by_id[item]
550
+ for item in certificate.get("target_operation_ids") or []
551
+ if item in operation_by_id
552
+ ],
553
+ }
554
+
555
+
556
+ def parse_semantic_bound_answer(raw: str, packet: Mapping[str, Any]) -> dict[str, Any]:
557
+ try:
558
+ value = json.loads(raw)
559
+ except json.JSONDecodeError as exc:
560
+ raise RuntimeError(f"semantic-bound answer is not strict JSON: {exc}") from exc
561
+ try:
562
+ return validate_semantic_answer(normalize_semantic_answer_output(value, packet), packet)
563
+ except SemanticEvidenceError as exc:
564
+ raise RuntimeError(f"semantic-bound answer failed validation: {exc}") from exc
565
+
566
+
567
+ def main() -> int:
568
+ parser = argparse.ArgumentParser()
569
+ parser.add_argument("--evidence", required=True)
570
+ parser.add_argument("--harness", required=True)
571
+ parser.add_argument("--out-dir", required=True)
572
+ parser.add_argument("--workers", type=int, default=12)
573
+ parser.add_argument("--attempts", type=int, default=3)
574
+ parser.add_argument("--answer-window-limit", type=int, default=8)
575
+ parser.add_argument("--limit", type=int, default=0)
576
+ parser.add_argument("--qid-list", default="")
577
+ parser.add_argument("--lane", choices=("production", "shadow"), default="production")
578
+ args = parser.parse_args()
579
+ answer_model = clean_text(os.getenv("TMCRA_ANSWER_MODEL"))
580
+ answer_base_url = clean_text(os.getenv("TMCRA_ANSWER_BASE_URL"))
581
+ answer_api_key = clean_text(os.getenv("TMCRA_ANSWER_API_KEY"))
582
+ if answer_model != "gpt-5.4":
583
+ raise RuntimeError(f"fixed answer model must be exactly gpt-5.4, got {answer_model!r}")
584
+ if not answer_base_url or not answer_api_key:
585
+ raise RuntimeError("TMCRA_ANSWER_BASE_URL and TMCRA_ANSWER_API_KEY are required")
586
+ fixed_answer_layer = fixed_answer_environment(args.answer_window_limit)
587
+ os.environ.update(fixed_answer_layer)
588
+ rows = read_jsonl(Path(args.evidence))
589
+ if args.qid_list:
590
+ requested = [
591
+ clean_text(value)
592
+ for value in Path(args.qid_list).read_text(encoding="utf-8").splitlines()
593
+ if clean_text(value)
594
+ ]
595
+ by_qid = {clean_text(row.get("question_id")): row for row in rows}
596
+ missing = [qid for qid in requested if qid not in by_qid]
597
+ if missing:
598
+ raise RuntimeError(f"requested qids are missing from evidence: {missing[:10]}")
599
+ rows = [by_qid[qid] for qid in requested]
600
+ if args.limit > 0:
601
+ rows = rows[: args.limit]
602
+ if not rows:
603
+ raise RuntimeError("no answer rows selected")
604
+ reject_evaluation_fields(rows)
605
+ if args.lane == "production":
606
+ try:
607
+ validate_production_evidence(rows)
608
+ except RoutePolicyError as exc:
609
+ raise RuntimeError(
610
+ f"production route failed before GPT-5.4 calls: {exc}"
611
+ ) from exc
612
+ harness = load_harness(Path(args.harness))
613
+ configured_base_url, configured_model, configured_api_key = harness.answer_llm_config()
614
+ if clean_text(configured_model) != "gpt-5.4":
615
+ raise RuntimeError(f"harness resolved unexpected model: {configured_model!r}")
616
+ if not clean_text(configured_base_url) or not clean_text(configured_api_key):
617
+ raise RuntimeError("harness did not resolve the required GPT5.4 endpoint")
618
+
619
+ # The legacy harness contains optional answer-side repair hooks. Keep its
620
+ # primary prompt, but make one logical GPT call the enforced contract.
621
+ harness.complete_profile_answer_from_evidence = lambda _question, answer, _windows: answer
622
+ call_state = threading.local()
623
+ original_chat_completion = harness.chat_completion
624
+
625
+ def counted_chat_completion(*call_args: Any, **call_kwargs: Any) -> str:
626
+ call_state.count = int(getattr(call_state, "count", 0)) + 1
627
+ return original_chat_completion(*call_args, **call_kwargs)
628
+
629
+ harness.chat_completion = counted_chat_completion
630
+ out_dir = Path(args.out_dir)
631
+ out_dir.mkdir(parents=True, exist_ok=True)
632
+ raw_attempt_dir = out_dir / "raw_attempts"
633
+ raw_attempt_dir.mkdir(exist_ok=True)
634
+ answers_path = out_dir / "answers.jsonl"
635
+ existing: dict[str, dict[str, Any]] = {}
636
+ if answers_path.exists():
637
+ for row in read_jsonl(answers_path):
638
+ qid = clean_text(row.get("question_id"))
639
+ if qid:
640
+ if qid in existing:
641
+ raise RuntimeError(f"duplicate existing answer row: {qid}")
642
+ existing[qid] = row
643
+ if existing and args.lane == "production":
644
+ try:
645
+ validate_production_answers(list(existing.values()))
646
+ except RoutePolicyError as exc:
647
+ raise RuntimeError(
648
+ "existing answers use an incompatible production protocol; "
649
+ "start a new answer output directory"
650
+ ) from exc
651
+ evidence_hashes, legacy_evidence_hash_backfills = bind_existing_answers_to_evidence(
652
+ existing, rows
653
+ )
654
+ lock = threading.Lock()
655
+
656
+ def answer_one(row: Mapping[str, Any]) -> dict[str, Any]:
657
+ qid = clean_text(row.get("question_id"))
658
+ question = clean_text(row.get("question"))
659
+ question_date = clean_text(row.get("question_date"))
660
+ runtime_question = f"{question}\nQuestion date: {question_date}" if question_date else question
661
+ semantic_packet = row.get("semantic_evidence_packet")
662
+ operation_packet = row.get("compiled_evidence_packet")
663
+ if args.lane == "production":
664
+ if isinstance(semantic_packet, Mapping):
665
+ raise RuntimeError(f"{qid}: semantic packets are forbidden in the production lane")
666
+ if not isinstance(operation_packet, Mapping):
667
+ raise RuntimeError(f"{qid}: production lane requires a compiled evidence packet")
668
+ packet = operation_packet
669
+ rendered_packet = render_compiled_packet(packet)
670
+ answer_prompt = EVIDENCE_OPERATION_ANSWER_PROMPT
671
+ answer_parser = parse_operation_bound_answer
672
+ answer_protocol = PRODUCTION_ANSWER_PROTOCOL
673
+ else:
674
+ if not isinstance(semantic_packet, Mapping) or isinstance(operation_packet, Mapping):
675
+ raise RuntimeError(f"{qid}: shadow lane requires only a semantic evidence packet")
676
+ packet = semantic_packet_to_advisory_packet(semantic_packet)
677
+ rendered_packet = render_compiled_packet(packet)
678
+ answer_prompt = EVIDENCE_OPERATION_ANSWER_PROMPT
679
+ answer_parser = parse_operation_bound_answer
680
+ answer_protocol = SHADOW_ANSWER_PROTOCOL
681
+ messages = [
682
+ {"role": "system", "content": answer_prompt},
683
+ {
684
+ "role": "user",
685
+ "content": json.dumps(rendered_packet, ensure_ascii=False, separators=(",", ":")),
686
+ },
687
+ ]
688
+ started = time.time()
689
+ last_error: Exception | None = None
690
+ answer = ""
691
+ evidence_binding: dict[str, Any] = {}
692
+ attempts_used = 0
693
+ logical_answer_calls = 0
694
+ validation_error = ""
695
+ for attempt in range(1, max(1, args.attempts) + 1):
696
+ attempts_used = attempt
697
+ call_state.count = 0
698
+ try:
699
+ attempt_path = raw_attempt_dir / f"{qid}.attempt_{attempt}.json"
700
+ if attempt_path.is_file():
701
+ persisted = json.loads(attempt_path.read_text(encoding="utf-8"))
702
+ if persisted.get("question_id") != qid or persisted.get("evidence_sha256") != evidence_hashes[qid]:
703
+ raise RuntimeError(f"{qid}: persisted raw answer identity mismatch")
704
+ raw_answer = str(persisted["raw_answer"])
705
+ attempt_call_count = 0
706
+ else:
707
+ attempt_messages = list(messages)
708
+ if validation_error:
709
+ attempt_messages.append(
710
+ {
711
+ "role": "user",
712
+ "content": (
713
+ "The previous JSON failed the deterministic validator: "
714
+ f"{validation_error}. Correct only the JSON structure or evidence bindings. "
715
+ "Use the same compiled packet and return the required schema."
716
+ ),
717
+ }
718
+ )
719
+ raw_answer = harness.chat_completion(
720
+ configured_base_url,
721
+ configured_model,
722
+ attempt_messages,
723
+ max_tokens=int(os.getenv("TMCRA_ANSWER_MAX_TOKENS", "1536")),
724
+ temperature=0.0,
725
+ api_key=configured_api_key,
726
+ )
727
+ attempt_call_count = int(getattr(call_state, "count", 0))
728
+ if attempt_call_count != 1:
729
+ raise RuntimeError(
730
+ f"{qid}: answer attempt made {attempt_call_count} logical GPT calls; expected exactly 1"
731
+ )
732
+ temporary = attempt_path.with_name(f".{attempt_path.name}.tmp.{os.getpid()}.{time.time_ns()}")
733
+ temporary.write_text(
734
+ json.dumps(
735
+ {
736
+ "question_id": qid,
737
+ "evidence_sha256": evidence_hashes[qid],
738
+ "attempt": attempt,
739
+ "raw_answer": raw_answer,
740
+ },
741
+ ensure_ascii=False,
742
+ sort_keys=True,
743
+ )
744
+ + "\n",
745
+ encoding="utf-8",
746
+ )
747
+ os.replace(temporary, attempt_path)
748
+ logical_answer_calls += attempt_call_count
749
+ if attempt_call_count not in {0, 1}:
750
+ raise RuntimeError(
751
+ f"{qid}: invalid physical call count for answer attempt"
752
+ )
753
+ evidence_binding = answer_parser(raw_answer, packet)
754
+ answer = evidence_binding["answer"]
755
+ break
756
+ except Exception as exc:
757
+ last_error = exc
758
+ validation_error = f"{exc.__class__.__name__}: {exc}"
759
+ if attempt >= max(1, args.attempts):
760
+ break
761
+ time.sleep(min(8.0, 1.5 * (2 ** (attempt - 1))))
762
+ if not answer:
763
+ assert last_error is not None
764
+ raise RuntimeError(
765
+ f"{qid}: GPT5.4 failed after {max(1, args.attempts)} same-model attempts: "
766
+ f"{last_error.__class__.__name__}:{last_error}"
767
+ )
768
+ return {
769
+ "question_id": qid,
770
+ "question_type": clean_text(row.get("question_type")),
771
+ "question": question,
772
+ "evidence_sha256": evidence_hashes[qid],
773
+ "hypothesis": answer,
774
+ "evidence_binding": evidence_binding,
775
+ "answer_protocol": answer_protocol,
776
+ "answer_model": configured_model,
777
+ "answer_attempt_limit": max(1, args.attempts),
778
+ "answer_attempts_used": attempts_used,
779
+ "logical_answer_calls": logical_answer_calls,
780
+ "selected_session_ids": list(row.get("selected_session_ids") or []),
781
+ "latency_sec": round(time.time() - started, 3),
782
+ "error": "",
783
+ }
784
+
785
+ pending = [row for row in rows if clean_text(row.get("question_id")) not in existing]
786
+ failures: list[dict[str, str]] = []
787
+ completed = 0
788
+ started = time.time()
789
+ with ThreadPoolExecutor(max_workers=max(1, args.workers)) as executor:
790
+ futures = {executor.submit(answer_one, row): clean_text(row.get("question_id")) for row in pending}
791
+ for future in as_completed(futures):
792
+ qid = futures[future]
793
+ try:
794
+ result = future.result()
795
+ except Exception as exc:
796
+ failures.append({"question_id": qid, "error": f"{exc.__class__.__name__}:{str(exc)[:500]}"})
797
+ continue
798
+ with lock:
799
+ with answers_path.open("a", encoding="utf-8") as handle:
800
+ handle.write(json.dumps(result, ensure_ascii=False, sort_keys=True) + "\n")
801
+ existing[qid] = result
802
+ completed += 1
803
+ if completed % 10 == 0 or completed == len(pending):
804
+ print(
805
+ json.dumps(
806
+ {
807
+ "completed_this_run": completed,
808
+ "pending_this_run": len(pending),
809
+ "total_saved": len(existing),
810
+ "failures": len(failures),
811
+ "elapsed_sec": round(time.time() - started, 3),
812
+ }
813
+ ),
814
+ flush=True,
815
+ )
816
+ final_rows = ordered_completed_answers(existing, rows)
817
+ atomic_write_jsonl(answers_path, final_rows)
818
+ protocols = sorted({str(row.get("answer_protocol") or "") for row in final_rows})
819
+ report = {
820
+ "status": "complete" if len(final_rows) == len(rows) and not failures else "incomplete",
821
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
822
+ "answer_model": configured_model,
823
+ "lane": args.lane,
824
+ "promotion_eligible": args.lane == "production",
825
+ "fixed_answer_layer": fixed_answer_layer,
826
+ "disabled_harness_hooks": ["complete_profile_answer_from_evidence"],
827
+ "max_logical_answer_calls_per_success": max(1, args.attempts),
828
+ "hierarchical_memory_context_rendering": True,
829
+ "compiled_evidence_packet_required": True,
830
+ "answer_protocol": protocols[0] if len(protocols) == 1 else "mixed",
831
+ "evidence_binding_persisted": True,
832
+ "answers_bound_to_evidence_sha256": True,
833
+ "legacy_evidence_hash_backfills": legacy_evidence_hash_backfills,
834
+ "requested": len(rows),
835
+ "completed": len(final_rows),
836
+ "failure_count": len(failures),
837
+ "avg_latency_sec": round(sum(float(row["latency_sec"]) for row in final_rows) / max(1, len(final_rows)), 3),
838
+ "answers": str(answers_path),
839
+ }
840
+ (out_dir / "failures.json").write_text(json.dumps(failures, indent=2, sort_keys=True), encoding="utf-8")
841
+ (out_dir / "report.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
842
+ print(json.dumps(report, indent=2, sort_keys=True))
843
+ return 0 if report["status"] == "complete" else 2
844
+
845
+
846
+ if __name__ == "__main__":
847
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_parallel_official_judge.py ADDED
@@ -0,0 +1,273 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Run one isolated official LongMemEval judge process per question."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import argparse
7
+ import hashlib
8
+ import json
9
+ import os
10
+ import subprocess
11
+ import sys
12
+ from concurrent.futures import ThreadPoolExecutor, as_completed
13
+ from pathlib import Path
14
+ from typing import Any, Mapping
15
+
16
+
17
+ DEFAULT_OFFICIAL_JUDGE = (
18
+ Path(__file__).resolve().parent.parent / "harness" / "evaluate_qa_official_judge.py"
19
+ )
20
+ METRIC_MODEL_ALIASES = {"gpt-4o-2024-08-06": "gpt-4o"}
21
+ BOUND_FIELDS = ("question", "hypothesis", "evidence_sha256", "answer_model")
22
+
23
+
24
+ class ParallelJudgeError(RuntimeError):
25
+ pass
26
+
27
+
28
+ def _read_jsonl(path: Path) -> list[dict[str, Any]]:
29
+ if not path.is_file():
30
+ return []
31
+ return [
32
+ json.loads(line)
33
+ for line in path.read_text(encoding="utf-8").splitlines()
34
+ if line.strip()
35
+ ]
36
+
37
+
38
+ def _write_jsonl_atomic(path: Path, rows: list[Mapping[str, Any]]) -> None:
39
+ path.parent.mkdir(parents=True, exist_ok=True)
40
+ temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}")
41
+ with temporary.open("w", encoding="utf-8") as handle:
42
+ for row in rows:
43
+ handle.write(json.dumps(row, ensure_ascii=False) + "\n")
44
+ handle.flush()
45
+ os.fsync(handle.fileno())
46
+ os.replace(temporary, path)
47
+
48
+
49
+ def _validate_result(row: Mapping[str, Any], answer: Mapping[str, Any]) -> None:
50
+ qid = str(answer.get("question_id") or "").strip()
51
+ if str(row.get("question_id") or "").strip() != qid:
52
+ raise ParallelJudgeError(f"{qid}: judge result has a different question ID")
53
+ for key in BOUND_FIELDS:
54
+ if row.get(key) != answer.get(key):
55
+ raise ParallelJudgeError(f"{qid}: judge result is bound to different {key}")
56
+ label = row.get("autoeval_label")
57
+ if (
58
+ not isinstance(label, Mapping)
59
+ or type(label.get("label")) is not bool
60
+ or not str(label.get("model") or "").strip()
61
+ or not str(label.get("raw_response") or "").strip()
62
+ ):
63
+ raise ParallelJudgeError(f"{qid}: judge result has no valid official label")
64
+
65
+
66
+ def _question_dir(work_dir: Path, index: int, qid: str) -> Path:
67
+ digest = hashlib.sha256(qid.encode("utf-8")).hexdigest()[:16]
68
+ return work_dir / f"{index:04d}_{digest}"
69
+
70
+
71
+ def _judge_one(
72
+ *,
73
+ index: int,
74
+ answer: Mapping[str, Any],
75
+ args: argparse.Namespace,
76
+ environment: Mapping[str, str],
77
+ work_dir: Path,
78
+ ) -> tuple[int, dict[str, Any], bool]:
79
+ qid = str(answer["question_id"])
80
+ question_dir = _question_dir(work_dir, index, qid)
81
+ question_dir.mkdir(parents=True, exist_ok=True)
82
+ hypothesis_path = question_dir / "hypothesis.jsonl"
83
+ result_path = question_dir / "official_judge.jsonl"
84
+ log_path = question_dir / "official_judge.log"
85
+ _write_jsonl_atomic(hypothesis_path, [answer])
86
+
87
+ cached = _read_jsonl(result_path)
88
+ if cached:
89
+ if len(cached) != 1:
90
+ raise ParallelJudgeError(f"{qid}: cached judge result is not one row")
91
+ _validate_result(cached[0], answer)
92
+ return index, cached[0], True
93
+
94
+ command = [
95
+ sys.executable,
96
+ str(args.judge.resolve()),
97
+ "--metric-model",
98
+ METRIC_MODEL_ALIASES.get(args.metric_model, args.metric_model),
99
+ "--hyp-file",
100
+ str(hypothesis_path),
101
+ "--ref-file",
102
+ str(args.ref_file.resolve()),
103
+ "--result-file",
104
+ str(result_path),
105
+ "--base-url",
106
+ args.base_url,
107
+ "--api-key-env",
108
+ args.api_key_env,
109
+ "--timeout",
110
+ str(args.timeout),
111
+ "--max-retries",
112
+ str(args.max_retries),
113
+ "--quiet",
114
+ ]
115
+ completed = subprocess.run(
116
+ command,
117
+ env=dict(environment),
118
+ stdout=subprocess.PIPE,
119
+ stderr=subprocess.STDOUT,
120
+ text=True,
121
+ timeout=args.process_timeout,
122
+ check=False,
123
+ )
124
+ log_path.write_text(completed.stdout or "", encoding="utf-8")
125
+ if completed.returncode != 0:
126
+ raise ParallelJudgeError(
127
+ f"{qid}: official judge process exited {completed.returncode}; log={log_path}"
128
+ )
129
+ rows = _read_jsonl(result_path)
130
+ if len(rows) != 1:
131
+ raise ParallelJudgeError(f"{qid}: official judge process did not write one row")
132
+ _validate_result(rows[0], answer)
133
+ return index, rows[0], False
134
+
135
+
136
+ def run_parallel_judge(
137
+ args: argparse.Namespace, *, environment: Mapping[str, str] | None = None
138
+ ) -> dict[str, Any]:
139
+ environment = dict(environment or os.environ)
140
+ if not environment.get(args.api_key_env):
141
+ raise ParallelJudgeError(f"missing API key env: {args.api_key_env}")
142
+ if not args.judge.is_file():
143
+ raise ParallelJudgeError(f"official judge is missing: {args.judge}")
144
+
145
+ hypotheses = _read_jsonl(args.hyp_file)
146
+ if args.limit and args.limit > 0:
147
+ hypotheses = hypotheses[: args.limit]
148
+ expected_qids = [str(item.get("question_id") or "").strip() for item in hypotheses]
149
+ if not expected_qids or any(not qid for qid in expected_qids):
150
+ raise ParallelJudgeError("hypotheses are empty or contain empty question IDs")
151
+ if len(expected_qids) != len(set(expected_qids)):
152
+ raise ParallelJudgeError("hypotheses contain duplicate question IDs")
153
+
154
+ result_file = args.result_file or Path(
155
+ str(args.hyp_file) + f".eval-results-{args.metric_model}"
156
+ )
157
+ if result_file.exists() and not args.resume:
158
+ raise ParallelJudgeError("judge result exists; use --resume")
159
+ existing = _read_jsonl(result_file)
160
+ if len(existing) > len(hypotheses):
161
+ raise ParallelJudgeError("existing judge result is longer than hypotheses")
162
+ for index, row in enumerate(existing):
163
+ _validate_result(row, hypotheses[index])
164
+
165
+ committed: dict[int, dict[str, Any]] = {
166
+ index: row for index, row in enumerate(existing)
167
+ }
168
+ work_dir = args.work_dir or result_file.with_name(f"{result_file.name}.work")
169
+ work_dir.mkdir(parents=True, exist_ok=True)
170
+ pending_indices = list(range(len(existing), len(hypotheses)))
171
+ failures: list[str] = []
172
+ cached_count = 0
173
+ called_count = 0
174
+
175
+ def commit_contiguous_prefix() -> int:
176
+ prefix: list[dict[str, Any]] = []
177
+ for index in range(len(hypotheses)):
178
+ if index not in committed:
179
+ break
180
+ prefix.append(committed[index])
181
+ _write_jsonl_atomic(result_file, prefix)
182
+ return len(prefix)
183
+
184
+ with ThreadPoolExecutor(max_workers=args.workers) as executor:
185
+ futures = {
186
+ executor.submit(
187
+ _judge_one,
188
+ index=index,
189
+ answer=hypotheses[index],
190
+ args=args,
191
+ environment=environment,
192
+ work_dir=work_dir,
193
+ ): index
194
+ for index in pending_indices
195
+ }
196
+ for future in as_completed(futures):
197
+ index = futures[future]
198
+ try:
199
+ completed_index, row, cached = future.result()
200
+ committed[completed_index] = row
201
+ cached_count += int(cached)
202
+ called_count += int(not cached)
203
+ commit_contiguous_prefix()
204
+ except Exception as exc:
205
+ failures.append(f"{expected_qids[index]}: {exc}")
206
+
207
+ committed_count = commit_contiguous_prefix()
208
+ report = {
209
+ "status": "complete" if committed_count == len(hypotheses) else "incomplete",
210
+ "metric_model": args.metric_model,
211
+ "row_count": len(hypotheses),
212
+ "committed_count": committed_count,
213
+ "existing_count": len(existing),
214
+ "cached_count": cached_count,
215
+ "physical_call_count": called_count,
216
+ "failure_count": len(failures),
217
+ "failures": failures,
218
+ "workers": args.workers,
219
+ "work_dir": str(work_dir),
220
+ }
221
+ if failures or committed_count != len(hypotheses):
222
+ report_path = result_file.with_name(f"{result_file.name}.isolated_report.json")
223
+ temporary = report_path.with_name(f".{report_path.name}.tmp.{os.getpid()}")
224
+ temporary.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
225
+ os.replace(temporary, report_path)
226
+ raise ParallelJudgeError(
227
+ f"official judge incomplete: committed={committed_count}/{len(hypotheses)}, "
228
+ f"failures={len(failures)}; report={report_path}"
229
+ )
230
+ labels = [bool((committed[index]["autoeval_label"])["label"]) for index in range(len(hypotheses))]
231
+ report["correct_count"] = sum(labels)
232
+ report["accuracy"] = round(sum(labels) / len(labels), 6)
233
+ report_path = result_file.with_name(f"{result_file.name}.isolated_report.json")
234
+ temporary = report_path.with_name(f".{report_path.name}.tmp.{os.getpid()}")
235
+ temporary.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
236
+ os.replace(temporary, report_path)
237
+ return report
238
+
239
+
240
+ def main() -> int:
241
+ parser = argparse.ArgumentParser()
242
+ parser.add_argument("--judge", type=Path, default=DEFAULT_OFFICIAL_JUDGE)
243
+ parser.add_argument("--metric-model", default="gpt-4o")
244
+ parser.add_argument("--hyp-file", required=True, type=Path)
245
+ parser.add_argument("--ref-file", required=True, type=Path)
246
+ parser.add_argument("--result-file", type=Path)
247
+ parser.add_argument("--work-dir", type=Path)
248
+ parser.add_argument("--base-url", required=True)
249
+ parser.add_argument("--api-key-env", default="OPENAI_API_KEY")
250
+ parser.add_argument("--limit", type=int, default=0)
251
+ parser.add_argument("--timeout", type=int, default=180)
252
+ parser.add_argument("--process-timeout", type=int, default=240)
253
+ parser.add_argument("--max-retries", type=int, default=0)
254
+ parser.add_argument("--resume", action="store_true")
255
+ parser.add_argument("--quiet", action="store_true")
256
+ parser.add_argument(
257
+ "--workers",
258
+ type=int,
259
+ default=int(os.environ.get("TMCRA_OFFICIAL_JUDGE_WORKERS", "8")),
260
+ )
261
+ args = parser.parse_args()
262
+ if args.workers <= 0:
263
+ raise ParallelJudgeError("workers must be positive")
264
+ if args.process_timeout <= args.timeout:
265
+ raise ParallelJudgeError("process timeout must exceed request timeout")
266
+ summary = run_parallel_judge(args)
267
+ if not args.quiet:
268
+ print(json.dumps(summary, sort_keys=True))
269
+ return 0
270
+
271
+
272
+ if __name__ == "__main__":
273
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_retrieve.py ADDED
@@ -0,0 +1,524 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import json
7
+ import os
8
+ import re
9
+ import sys
10
+ from pathlib import Path
11
+ from typing import Any
12
+
13
+ import tmcra_v4_online_runtime as online
14
+ from run_tmcra_v4_build import (
15
+ BASE,
16
+ DEFAULT_REPO,
17
+ DEFAULT_WRITER_ENV,
18
+ _key_pool,
19
+ _load_shell_environment,
20
+ _now,
21
+ _resume_log,
22
+ _run,
23
+ )
24
+ from tmcra_v4_cost_report import build_report, collect_calls
25
+ from tmcra_v4_route_policy import (
26
+ RoutePolicyError,
27
+ validate_diagnostic_retrieval_rows,
28
+ validate_production_packing_budget,
29
+ validate_production_retrieval_mode,
30
+ validate_production_retrieval_rows,
31
+ )
32
+
33
+
34
+ PROJECT_ROOT = BASE.parents[2]
35
+ PACKAGE_ROOT = BASE.parent
36
+ DEFAULT_HARNESS = PACKAGE_ROOT / "harness" / "run_lme_s10_native_tmcra.py"
37
+ DEFAULT_NODE = PROJECT_ROOT / "models" / "tmcra" / "node_scorer.pt"
38
+ DEFAULT_PATH_MODEL = PROJECT_ROOT / "models" / "tmcra" / "path_scorer.pt"
39
+ DEFAULT_CHECKPOINT = PROJECT_ROOT / "models" / "tmcra" / "tmcra_v3_reranker.pt"
40
+ DEFAULT_EMBEDDING = PROJECT_ROOT / "models" / "BAAI" / "bge-m3"
41
+ DEFAULT_CROSS = PROJECT_ROOT / "models" / "BAAI" / "bge-reranker-v2-m3"
42
+
43
+
44
+ class RetrievalError(RuntimeError):
45
+ pass
46
+
47
+
48
+ def _jsonl(path: Path) -> list[dict[str, Any]]:
49
+ return [
50
+ json.loads(line)
51
+ for line in path.read_text(encoding="utf-8").splitlines()
52
+ if line.strip()
53
+ ]
54
+
55
+
56
+ def _fingerprints(rows: list[dict[str, Any]]) -> dict[str, str]:
57
+ return {
58
+ str(row["question_id"]): online.scope_fingerprint(
59
+ Path(str(row["db_path"])).resolve(), str(row["scope_id"])
60
+ )
61
+ for row in rows
62
+ }
63
+
64
+
65
+ def _file_sha256(path: Path) -> str:
66
+ digest = hashlib.sha256()
67
+ with path.open("rb") as handle:
68
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
69
+ digest.update(chunk)
70
+ return digest.hexdigest()
71
+
72
+
73
+ def _atomic_write_text(path: Path, text: str) -> None:
74
+ temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}")
75
+ temporary.write_text(text, encoding="utf-8")
76
+ os.replace(temporary, path)
77
+
78
+
79
+ def _prepare_graph_boundary(
80
+ *,
81
+ run_dir: Path,
82
+ tag: str,
83
+ fingerprints: dict[str, str],
84
+ resume: bool,
85
+ ) -> None:
86
+ before_path = run_dir / f"{tag}.graph_before.json"
87
+ failed_path = run_dir / f"FAILED.{tag}"
88
+ staging_path = run_dir / f".{tag}.staging"
89
+ if resume:
90
+ if not failed_path.is_file():
91
+ raise RetrievalError("resume requires the existing FAILED marker")
92
+ if not before_path.is_file():
93
+ raise RetrievalError("resume requires the original graph_before boundary")
94
+ try:
95
+ original = json.loads(before_path.read_text(encoding="utf-8"))
96
+ except (OSError, json.JSONDecodeError) as exc:
97
+ raise RetrievalError("original graph_before boundary is unreadable") from exc
98
+ if original != fingerprints:
99
+ raise RetrievalError("graph changed since the failed retrieval began")
100
+ return
101
+ if failed_path.exists() or before_path.exists() or staging_path.exists():
102
+ raise RetrievalError(
103
+ "prior retrieval state exists; inspect it and use explicit --resume"
104
+ )
105
+ _atomic_write_text(
106
+ before_path,
107
+ json.dumps(fingerprints, indent=2, sort_keys=True) + "\n",
108
+ )
109
+
110
+
111
+ def _archive_failed_marker(run_dir: Path, tag: str) -> None:
112
+ failed_path = run_dir / f"FAILED.{tag}"
113
+ if not failed_path.is_file():
114
+ return
115
+ history_path = run_dir / f"{tag}.failure_history.jsonl"
116
+ content = failed_path.read_text(encoding="utf-8")
117
+ with history_path.open("a", encoding="utf-8") as handle:
118
+ for line in content.splitlines():
119
+ if line.strip():
120
+ handle.write(line.rstrip() + "\n")
121
+ failed_path.unlink()
122
+
123
+
124
+ def _record_failure(run_dir: Path, tag: str, exc: Exception) -> None:
125
+ _archive_failed_marker(run_dir, tag)
126
+ _atomic_write_text(
127
+ run_dir / f"FAILED.{tag}",
128
+ json.dumps(
129
+ {"at": _now(), "error": f"{exc.__class__.__name__}: {exc}"}
130
+ )
131
+ + "\n",
132
+ )
133
+
134
+
135
+ def _validate_manifest_pair(
136
+ scope_rows: list[dict[str, Any]], query_rows: list[dict[str, Any]]
137
+ ) -> None:
138
+ def inventory(rows: list[dict[str, Any]], label: str) -> dict[str, tuple[str, str, str]]:
139
+ result: dict[str, tuple[str, str, str]] = {}
140
+ for row in rows:
141
+ qid = str(row.get("question_id") or "").strip()
142
+ identity = tuple(
143
+ str(Path(str(row[key])).resolve()) if key != "scope_id" else str(row[key]).strip()
144
+ for key in ("db_path", "scope_id", "index_path")
145
+ )
146
+ if not qid or qid in result or not all(identity):
147
+ raise RetrievalError(f"{label} has a missing or duplicate retrieval identity")
148
+ result[qid] = identity
149
+ return result
150
+
151
+ scope_inventory = inventory(scope_rows, "scope manifest")
152
+ query_inventory = inventory(query_rows, "query manifest")
153
+ if scope_inventory != query_inventory:
154
+ raise RetrievalError("scope/query manifests do not reference the same frozen indexes")
155
+
156
+
157
+ def _select_manifest_rows(
158
+ scope_rows: list[dict[str, Any]],
159
+ query_rows: list[dict[str, Any]],
160
+ qid_path: Path,
161
+ ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
162
+ qids = [
163
+ line.strip()
164
+ for line in qid_path.read_text(encoding="utf-8").splitlines()
165
+ if line.strip()
166
+ ]
167
+ if not qids or len(qids) != len(set(qids)):
168
+ raise RetrievalError("qid list is empty or duplicated")
169
+ scope_by_qid = {str(row.get("question_id") or "").strip(): row for row in scope_rows}
170
+ query_by_qid = {str(row.get("question_id") or "").strip(): row for row in query_rows}
171
+ missing = [
172
+ qid
173
+ for qid in qids
174
+ if qid not in scope_by_qid or qid not in query_by_qid
175
+ ]
176
+ if missing:
177
+ raise RetrievalError(f"qid list contains unknown questions: {missing[:10]}")
178
+ return (
179
+ [scope_by_qid[qid] for qid in qids],
180
+ [query_by_qid[qid] for qid in qids],
181
+ )
182
+
183
+
184
+ def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
185
+ _atomic_write_text(
186
+ path,
187
+ "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in rows),
188
+ )
189
+
190
+
191
+ def _run_production_build_preflight(
192
+ run_dir: Path, tag: str, *, resume: bool
193
+ ) -> dict[str, Any]:
194
+ output = run_dir / f"{tag}.production_build_preflight.json"
195
+ _run(
196
+ [
197
+ sys.executable,
198
+ str(BASE / "audit_tmcra_v4_chain.py"),
199
+ "--run-dir",
200
+ str(run_dir),
201
+ "--output",
202
+ str(output),
203
+ "--build-only",
204
+ ],
205
+ (
206
+ _resume_log(run_dir, f"{tag}.production_build_preflight")
207
+ if resume
208
+ else run_dir / f"{tag}.production_build_preflight.log"
209
+ ),
210
+ dict(os.environ),
211
+ )
212
+ if not output.is_file():
213
+ raise RetrievalError("production build preflight did not emit an audit report")
214
+ try:
215
+ report = json.loads(output.read_text(encoding="utf-8"))
216
+ except (OSError, json.JSONDecodeError) as exc:
217
+ raise RetrievalError("production build preflight report is unreadable") from exc
218
+ if report.get("passed") is not True:
219
+ raise RetrievalError("production build preflight did not pass")
220
+ return report
221
+
222
+
223
+ def _validate_committed_retrieval_output(
224
+ out_dir: Path,
225
+ *,
226
+ expected_qids: list[str],
227
+ execution_lane: str,
228
+ ) -> dict[str, Any]:
229
+ evidence_rows = _jsonl(out_dir / "evidence_windows.jsonl")
230
+ debug_rows = _jsonl(out_dir / "retrieval_debug.jsonl")
231
+ try:
232
+ route_report = (
233
+ validate_production_retrieval_rows(evidence_rows)
234
+ if execution_lane == "production"
235
+ else validate_diagnostic_retrieval_rows(evidence_rows)
236
+ )
237
+ except RoutePolicyError as exc:
238
+ raise RetrievalError(
239
+ f"committed retrieval output does not match the requested lane: {exc}"
240
+ ) from exc
241
+ evidence_qids = [str(row.get("question_id") or "").strip() for row in evidence_rows]
242
+ debug_qids = [
243
+ str(row.get("question_id") or row.get("qid") or "").strip()
244
+ for row in debug_rows
245
+ ]
246
+ if evidence_qids != expected_qids or debug_qids != expected_qids:
247
+ raise RetrievalError(
248
+ "committed retrieval output does not match the selected query manifest"
249
+ )
250
+ try:
251
+ report = json.loads((out_dir / "report.json").read_text(encoding="utf-8"))
252
+ except (OSError, json.JSONDecodeError) as exc:
253
+ raise RetrievalError("committed retrieval report is unreadable") from exc
254
+ if report.get("query_count") != len(expected_qids):
255
+ raise RetrievalError("committed retrieval report query count is inconsistent")
256
+ return route_report
257
+
258
+
259
+ def retrieve(args: argparse.Namespace) -> dict[str, Any]:
260
+ # The logical run path is part of durable retrieval identity. Following a
261
+ # storage symlink here would invalidate resumable checkpoints after a move.
262
+ run_dir = args.run_dir.absolute()
263
+ execution_lane = "diagnostic" if bool(getattr(args, "diagnostic", False)) else "production"
264
+ try:
265
+ validate_production_retrieval_mode(
266
+ getattr(args, "composition_mode", None), execution_lane=execution_lane
267
+ )
268
+ validate_production_packing_budget(
269
+ getattr(args, "packing_budget_mode", "fixed"),
270
+ getattr(args, "top_k", 8),
271
+ execution_lane=execution_lane,
272
+ )
273
+ except RoutePolicyError as exc:
274
+ raise RetrievalError(f"retrieval route policy rejected the run: {exc}") from exc
275
+ if not (run_dir / "BUILD_COMPLETE").is_file():
276
+ raise RetrievalError("frozen V4 build is incomplete")
277
+ if not re.fullmatch(r"[A-Za-z0-9_.-]+", args.tag):
278
+ raise RetrievalError("tag contains unsupported characters")
279
+ out_dir = run_dir / args.tag
280
+ resume = bool(getattr(args, "resume", False))
281
+ reuse_committed_output = False
282
+ if out_dir.exists():
283
+ required = (
284
+ out_dir / "evidence_windows.jsonl",
285
+ out_dir / "retrieval_debug.jsonl",
286
+ out_dir / "report.json",
287
+ )
288
+ if not resume:
289
+ raise RetrievalError(f"retrieval output already exists: {out_dir}")
290
+ if not out_dir.is_dir() or not all(path.is_file() for path in required):
291
+ raise RetrievalError(
292
+ f"resume found an incomplete committed retrieval output: {out_dir}"
293
+ )
294
+ reuse_committed_output = True
295
+ query_manifest = Path(
296
+ getattr(args, "query_manifest", None) or run_dir / "query_manifest.jsonl"
297
+ ).resolve()
298
+ scope_manifest = Path(
299
+ getattr(args, "scope_manifest", None) or run_dir / "scope_manifest.jsonl"
300
+ ).resolve()
301
+ if not query_manifest.is_file() or not scope_manifest.is_file():
302
+ raise RetrievalError("scope/query manifest is missing")
303
+ rows = _jsonl(scope_manifest)
304
+ query_rows = _jsonl(query_manifest)
305
+ if not rows or not query_rows:
306
+ raise RetrievalError("scope/query manifest is empty")
307
+ _validate_manifest_pair(rows, query_rows)
308
+ qid_list = getattr(args, "qid_list", None)
309
+ if qid_list is not None:
310
+ if not Path(qid_list).resolve().is_file():
311
+ raise RetrievalError("qid list is missing")
312
+ rows, query_rows = _select_manifest_rows(
313
+ rows, query_rows, Path(qid_list).resolve()
314
+ )
315
+ scope_manifest = run_dir / f"{args.tag}.scope_manifest.selected.jsonl"
316
+ query_manifest = run_dir / f"{args.tag}.query_manifest.selected.jsonl"
317
+ _write_jsonl(scope_manifest, rows)
318
+ _write_jsonl(query_manifest, query_rows)
319
+ _validate_manifest_pair(rows, query_rows)
320
+
321
+ if reuse_committed_output:
322
+ _validate_committed_retrieval_output(
323
+ out_dir,
324
+ expected_qids=[str(row["question_id"]).strip() for row in query_rows],
325
+ execution_lane=execution_lane,
326
+ )
327
+
328
+ if execution_lane == "production":
329
+ _run_production_build_preflight(run_dir, args.tag, resume=resume)
330
+
331
+ shell_environment = _load_shell_environment(args.writer_env.resolve())
332
+ environment = {**os.environ, **shell_environment}
333
+ keys = _key_pool(environment)
334
+ base_url = environment.get("TMCRA_DEEPSEEK_WRITER_BASE_URL") or environment.get("TMCRA_WRITER_BASE_URL") or "https://api.deepseek.com/v1"
335
+ environment.update(
336
+ {
337
+ "TMCRA_RECALL_PLANNER_BASE_URL": base_url,
338
+ "TMCRA_RECALL_PLANNER_MODEL": "deepseek-v4-flash",
339
+ "TMCRA_RECALL_PLANNER_API_KEY_POOL": ",".join(keys),
340
+ "TMCRA_NODE_MODEL_PATH": str(args.node_model.resolve()),
341
+ "TMCRA_PATH_MODEL_PATH": str(args.path_model.resolve()),
342
+ "TMCRA_NODE_MODEL_DEVICE": args.graph_device,
343
+ }
344
+ )
345
+ before = _fingerprints(rows)
346
+ _prepare_graph_boundary(
347
+ run_dir=run_dir,
348
+ tag=args.tag,
349
+ fingerprints=before,
350
+ resume=resume,
351
+ )
352
+ try:
353
+ if not reuse_committed_output:
354
+ runtime_options = ["--resume"] if resume else []
355
+ if args.planner_replay_dir is not None:
356
+ runtime_options.extend(
357
+ ["--planner-replay-dir", str(args.planner_replay_dir.resolve())]
358
+ )
359
+ _run(
360
+ [
361
+ sys.executable,
362
+ str(BASE / "tmcra_v4_online_runtime.py"),
363
+ "retrieve",
364
+ *runtime_options,
365
+ "--query-manifest",
366
+ str(query_manifest),
367
+ "--out-dir",
368
+ str(out_dir),
369
+ "--checkpoint",
370
+ str(args.checkpoint.resolve()),
371
+ "--cross-model",
372
+ str(args.cross_model.resolve()),
373
+ "--repo",
374
+ str(args.repo.resolve()),
375
+ "--harness",
376
+ str(args.harness.resolve()),
377
+ "--node-model",
378
+ str(args.node_model.resolve()),
379
+ "--path-model",
380
+ str(args.path_model.resolve()),
381
+ "--graph-device",
382
+ args.graph_device,
383
+ "--embedding-model",
384
+ str(args.embedding_model.resolve()),
385
+ "--slow-dense-k",
386
+ str(args.slow_dense_k),
387
+ "--dense-k",
388
+ str(args.dense_k),
389
+ "--graph-k",
390
+ str(args.graph_k),
391
+ "--composition-mode",
392
+ str(args.composition_mode),
393
+ "--execution-lane",
394
+ execution_lane,
395
+ "--packing-budget-mode",
396
+ str(args.packing_budget_mode),
397
+ "--top-k",
398
+ str(args.top_k),
399
+ "--adaptive-simple-k",
400
+ str(args.adaptive_simple_k),
401
+ "--adaptive-standard-k",
402
+ str(args.adaptive_standard_k),
403
+ "--adaptive-complex-k",
404
+ str(args.adaptive_complex_k),
405
+ "--cross-batch-size",
406
+ str(args.cross_batch_size),
407
+ "--device",
408
+ args.device,
409
+ ],
410
+ (
411
+ _resume_log(run_dir, args.tag)
412
+ if resume
413
+ else run_dir / f"{args.tag}.log"
414
+ ),
415
+ environment,
416
+ )
417
+ after = _fingerprints(rows)
418
+ _atomic_write_text(
419
+ run_dir / f"{args.tag}.graph_after.json",
420
+ json.dumps(after, indent=2, sort_keys=True) + "\n",
421
+ )
422
+ if after != before:
423
+ raise RetrievalError("retrieval changed graph structure or graph records")
424
+ _run(
425
+ [
426
+ sys.executable,
427
+ str(BASE / "audit_tmcra_v4_chain.py"),
428
+ "--run-dir",
429
+ str(run_dir),
430
+ "--retrieval-dir",
431
+ str(out_dir),
432
+ "--output",
433
+ str(run_dir / f"{args.tag}.chain_audit.json"),
434
+ ],
435
+ (
436
+ _resume_log(run_dir, f"{args.tag}.chain_audit")
437
+ if resume
438
+ else run_dir / f"{args.tag}.chain_audit.log"
439
+ ),
440
+ environment,
441
+ )
442
+ # Retrieval cost is incremental. Frozen graph databases contain the
443
+ # historical writer ledger and must not be charged to every query run.
444
+ calls = collect_calls([out_dir / "retrieval_debug.jsonl"], [])
445
+ cost = build_report(calls)
446
+ (run_dir / f"{args.tag}.cost_report.json").write_text(
447
+ json.dumps(cost, indent=2, sort_keys=True) + "\n", encoding="utf-8"
448
+ )
449
+ report = json.loads((out_dir / "report.json").read_text(encoding="utf-8"))
450
+ summary = {
451
+ "schema_version": "tmcra.v4.retrieval-run.1",
452
+ "status": "complete",
453
+ "tag": args.tag,
454
+ "execution_lane": execution_lane,
455
+ "composition_mode": str(args.composition_mode),
456
+ "query_count": report["query_count"],
457
+ "selected_output": str(out_dir / "evidence_windows.jsonl"),
458
+ "scope_manifest": str(scope_manifest),
459
+ "scope_manifest_sha256": _file_sha256(scope_manifest),
460
+ "query_manifest": str(query_manifest),
461
+ "query_manifest_sha256": _file_sha256(query_manifest),
462
+ "qid_list": str(Path(qid_list).resolve()) if qid_list is not None else None,
463
+ "qid_list_sha256": _file_sha256(Path(qid_list).resolve()) if qid_list is not None else None,
464
+ "graph_structure_unchanged": True,
465
+ "resumed": resume,
466
+ "reused_committed_retrieval_output": reuse_committed_output,
467
+ "physical_api_call_count_total": cost["physical_call_count"],
468
+ "exact_cost_cny_total": cost["exact_cost_cny"],
469
+ "min_cost_cny_total": cost["min_cost_cny"],
470
+ "max_cost_cny_total": cost["max_cost_cny"],
471
+ }
472
+ (out_dir / "retrieval_run_report.json").write_text(
473
+ json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
474
+ )
475
+ (out_dir / "RETRIEVAL_COMPLETE").write_text(_now() + "\n", encoding="utf-8")
476
+ _archive_failed_marker(run_dir, args.tag)
477
+ return summary
478
+ except Exception as exc:
479
+ _record_failure(run_dir, args.tag, exc)
480
+ raise
481
+
482
+
483
+ def main() -> int:
484
+ parser = argparse.ArgumentParser(description="Run one TMCRA V4 retrieval pass over a frozen build")
485
+ parser.add_argument("--run-dir", type=Path, required=True)
486
+ parser.add_argument("--tag", default="retrieval_1")
487
+ parser.add_argument("--resume", action="store_true")
488
+ parser.add_argument("--scope-manifest", type=Path)
489
+ parser.add_argument("--query-manifest", type=Path)
490
+ parser.add_argument("--qid-list", type=Path)
491
+ parser.add_argument("--planner-replay-dir", type=Path)
492
+ parser.add_argument("--writer-env", type=Path, default=DEFAULT_WRITER_ENV)
493
+ parser.add_argument("--repo", type=Path, default=DEFAULT_REPO)
494
+ parser.add_argument("--harness", type=Path, default=DEFAULT_HARNESS)
495
+ parser.add_argument("--node-model", type=Path, default=DEFAULT_NODE)
496
+ parser.add_argument("--path-model", type=Path, default=DEFAULT_PATH_MODEL)
497
+ parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
498
+ parser.add_argument("--embedding-model", type=Path, default=DEFAULT_EMBEDDING)
499
+ parser.add_argument("--cross-model", type=Path, default=DEFAULT_CROSS)
500
+ parser.add_argument("--device", default="cuda")
501
+ parser.add_argument("--graph-device", default="cuda")
502
+ parser.add_argument("--dense-k", type=int, default=32)
503
+ parser.add_argument("--slow-dense-k", type=int, default=24)
504
+ parser.add_argument("--graph-k", type=int, default=24)
505
+ parser.add_argument("--composition-mode", choices=("layered", "source-only-diagnostic"), default="layered")
506
+ parser.add_argument(
507
+ "--diagnostic",
508
+ action="store_true",
509
+ help="explicitly allow a non-production retrieval composition",
510
+ )
511
+ parser.add_argument("--packing-budget-mode", choices=("fixed", "adaptive"), default="fixed")
512
+ parser.add_argument("--top-k", type=int, default=8)
513
+ parser.add_argument("--adaptive-simple-k", type=int, default=8)
514
+ parser.add_argument("--adaptive-standard-k", type=int, default=12)
515
+ parser.add_argument("--adaptive-complex-k", type=int, default=16)
516
+ parser.add_argument("--cross-batch-size", type=int, default=24)
517
+ args = parser.parse_args()
518
+ summary = retrieve(args)
519
+ print(json.dumps(summary, sort_keys=True))
520
+ return 0
521
+
522
+
523
+ if __name__ == "__main__":
524
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_model_assets.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ from collections.abc import Mapping
5
+ from pathlib import Path
6
+ from typing import Any
7
+
8
+
9
+ def load_optional_hf_model_manifest(model_dir: str | Path) -> dict[str, Any]:
10
+ """Read optional provenance metadata for a locally cached Hugging Face model.
11
+
12
+ ``TMCRA_MODEL_MANIFEST.json`` is a TMCRA sidecar, not part of the Hugging
13
+ Face model layout. Official BGE snapshots therefore remain valid without
14
+ it. When a sidecar is supplied, validate it before using its provenance
15
+ fields in reports.
16
+ """
17
+
18
+ path = Path(model_dir) / "TMCRA_MODEL_MANIFEST.json"
19
+ if not path.exists():
20
+ return {}
21
+ if not path.is_file():
22
+ raise ValueError(f"model manifest must be a file: {path}")
23
+ try:
24
+ payload = json.loads(path.read_text(encoding="utf-8"))
25
+ except (OSError, UnicodeError, json.JSONDecodeError) as exc:
26
+ raise ValueError(f"invalid model manifest: {path}") from exc
27
+ if not isinstance(payload, Mapping):
28
+ raise ValueError(f"model manifest must contain a JSON object: {path}")
29
+ return dict(payload)
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v2_lme_pipeline.py ADDED
@@ -0,0 +1,855 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import json
7
+ import math
8
+ import os
9
+ import random
10
+ import re
11
+ import time
12
+ import urllib.request
13
+ from dataclasses import dataclass
14
+ from pathlib import Path
15
+ from typing import Any, Mapping, Sequence
16
+
17
+ import torch
18
+ import torch.nn as nn
19
+ import torch.nn.functional as F
20
+
21
+ TOKEN_RE = re.compile(r"[a-zA-Z0-9]+")
22
+ EVENT_RE = re.compile(r"^event::longmemeval:(?P<qid>[^:]+):s(?P<session>\d+)_c(?P<chunk>\d+)$")
23
+
24
+
25
+ def clean_text(value: Any) -> str:
26
+ return " ".join(str(value or "").split())
27
+
28
+
29
+ def norm_tokens(text: Any) -> list[str]:
30
+ return [m.group(0).lower() for m in TOKEN_RE.finditer(clean_text(text))]
31
+
32
+
33
+ def stable_hash(value: str) -> int:
34
+ return int.from_bytes(hashlib.blake2b(value.encode("utf-8"), digest_size=8).digest(), "big")
35
+
36
+
37
+ def embedding_text_key(text: Any) -> str:
38
+ cleaned = clean_text(text)
39
+ return hashlib.sha256(cleaned.encode("utf-8")).hexdigest()
40
+
41
+
42
+ def iter_sample_texts(samples: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
43
+ items: list[dict[str, Any]] = []
44
+ seen: set[str] = set()
45
+ for sample in samples:
46
+ qid = clean_text(sample.get("question_id"))
47
+ question = clean_text(sample.get("question", ""))
48
+ if question:
49
+ key = embedding_text_key(question)
50
+ if key not in seen:
51
+ seen.add(key)
52
+ items.append({"key": key, "kind": "question", "question_id": qid, "text": question})
53
+ for cand in list(sample.get("candidates") or []):
54
+ text = clean_text(cand.get("text", ""))
55
+ if not text:
56
+ continue
57
+ key = embedding_text_key(text)
58
+ if key not in seen:
59
+ seen.add(key)
60
+ items.append({
61
+ "key": key,
62
+ "kind": "candidate",
63
+ "question_id": qid,
64
+ "event_id": clean_text(cand.get("event_id", "")),
65
+ "text": text,
66
+ })
67
+ if not items:
68
+ raise RuntimeError("no texts found for embedding precompute")
69
+ return items
70
+
71
+
72
+ def split_name(qid: str, val_ratio: float = 0.2) -> str:
73
+ bucket = stable_hash(qid) % 10000
74
+ return "val" if bucket < int(val_ratio * 10000) else "train"
75
+
76
+
77
+ def read_jsonl(path: Path) -> list[dict[str, Any]]:
78
+ if not path.exists():
79
+ raise FileNotFoundError(f"required jsonl file not found: {path}")
80
+ rows=[]
81
+ with path.open("r", encoding="utf-8", errors="replace") as f:
82
+ for lineno, line in enumerate(f, start=1):
83
+ if line.strip():
84
+ try:
85
+ rows.append(json.loads(line))
86
+ except json.JSONDecodeError as exc:
87
+ raise RuntimeError(f"invalid jsonl at {path}:{lineno}: {exc}") from exc
88
+ if not rows:
89
+ raise RuntimeError(f"required jsonl file is empty: {path}")
90
+ return rows
91
+
92
+
93
+ def write_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
94
+ with path.open("w", encoding="utf-8") as f:
95
+ for row in rows:
96
+ f.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
97
+
98
+
99
+ def load_shard_rows(shard_glob: str) -> dict[str, list[dict[str, Any]]]:
100
+ base = Path("/") if shard_glob.startswith("/") else Path()
101
+ pattern = shard_glob[1:] if shard_glob.startswith("/") else shard_glob
102
+ shard_paths = sorted(base.glob(pattern))
103
+ if not shard_paths:
104
+ raise FileNotFoundError(f"no shard files matched --shard-glob: {shard_glob}")
105
+ rows_by_qid: dict[str, list[dict[str, Any]]] = {}
106
+ for shard in shard_paths:
107
+ data = json.loads(shard.read_text(encoding="utf-8", errors="replace"))
108
+ if not isinstance(data, list):
109
+ raise RuntimeError(f"shard is not a JSON list: {shard}")
110
+ for idx, row in enumerate(data):
111
+ if not isinstance(row, dict):
112
+ raise RuntimeError(f"shard row is not an object: {shard}[{idx}]")
113
+ qid = clean_text(row.get("question_id"))
114
+ if qid.endswith("_abs"):
115
+ qid = qid[:-4]
116
+ if not qid:
117
+ raise RuntimeError(f"shard row missing question_id: {shard}[{idx}]")
118
+ rows_by_qid.setdefault(qid, []).append(row)
119
+ if not rows_by_qid:
120
+ raise RuntimeError(f"no usable rows loaded from --shard-glob: {shard_glob}")
121
+ return rows_by_qid
122
+
123
+
124
+ def session_text_chunks(*, session_id: str, date: str, turns: list[Mapping[str, Any]], max_chars: int = 7000, max_chunks: int = 0) -> list[str]:
125
+ chunks: list[str] = []
126
+ current: list[str] = [f"LongMemEval session_id={session_id} date={date}"]
127
+ current_len = len(current[0])
128
+ for index, turn in enumerate(turns, start=1):
129
+ role = clean_text(turn.get("role", "unknown"))
130
+ content = clean_text(turn.get("content", ""))
131
+ if not content:
132
+ continue
133
+ part = f"[{session_id} turn={index} role={role}] {content}"
134
+ if current_len + len(part) + 2 > max_chars and len(current) > 1:
135
+ chunks.append("\n".join(current))
136
+ if max_chunks > 0 and len(chunks) >= max_chunks:
137
+ return chunks
138
+ current = [f"LongMemEval session_id={session_id} date={date} continued=true"]
139
+ current_len = len(current[0])
140
+ current.append(part)
141
+ current_len += len(part) + 1
142
+ if len(current) > 1 and (max_chunks <= 0 or len(chunks) < max_chunks):
143
+ chunks.append("\n".join(current))
144
+ return chunks
145
+
146
+
147
+ def event_text_map_for_row(row: Mapping[str, Any], qid: str) -> dict[str, dict[str, Any]]:
148
+ sessions = list(row.get("haystack_sessions") or [])
149
+ session_ids = [clean_text(x) for x in list(row.get("haystack_session_ids") or [])]
150
+ dates = [clean_text(x) for x in list(row.get("haystack_dates") or [])]
151
+ out: dict[str, dict[str, Any]] = {}
152
+ for sidx, sid in enumerate(session_ids):
153
+ if sidx >= len(sessions):
154
+ continue
155
+ chunks = session_text_chunks(
156
+ session_id=sid,
157
+ date=dates[sidx] if sidx < len(dates) else "",
158
+ turns=list(sessions[sidx] or []),
159
+ max_chars=7000,
160
+ max_chunks=0,
161
+ )
162
+ for cidx, text in enumerate(chunks, start=1):
163
+ eid = f"event::longmemeval:{qid}:s{sidx:03d}_c{cidx:02d}"
164
+ out[eid] = {
165
+ "event_id": eid,
166
+ "text": text,
167
+ "session_id": sid,
168
+ "session_index": sidx,
169
+ "chunk_index": cidx,
170
+ "date": dates[sidx] if sidx < len(dates) else "",
171
+ }
172
+ return out
173
+
174
+
175
+ def choose_row_for_query(rows: list[dict[str, Any]], candidate_ids: Sequence[str], positive_ids: Sequence[str]) -> dict[str, Any] | None:
176
+ if not rows:
177
+ return None
178
+ wanted = set(candidate_ids or []) | set(positive_ids or [])
179
+ if not wanted:
180
+ return None
181
+ qid = clean_text(rows[0].get("question_id"))
182
+ if qid.endswith("_abs"):
183
+ qid = qid[:-4]
184
+ best = None
185
+ best_hits = -1
186
+ for row in rows:
187
+ rqid = clean_text(row.get("question_id"))
188
+ if rqid.endswith("_abs"):
189
+ rqid = rqid[:-4]
190
+ events = set(event_text_map_for_row(row, rqid))
191
+ hits = len(wanted & events)
192
+ if hits > best_hits:
193
+ best = row
194
+ best_hits = hits
195
+ return best
196
+
197
+
198
+ def token_overlap(a: str, b: str) -> float:
199
+ at, bt = set(norm_tokens(a)), set(norm_tokens(b))
200
+ if not at or not bt:
201
+ return 0.0
202
+ return len(at & bt) / max(1, len(at | bt))
203
+
204
+
205
+ def contains_number_overlap(a: str, b: str) -> float:
206
+ nums_a = {t for t in norm_tokens(a) if any(ch.isdigit() for ch in t)}
207
+ nums_b = {t for t in norm_tokens(b) if any(ch.isdigit() for ch in t)}
208
+ if not nums_a or not nums_b:
209
+ return 0.0
210
+ return len(nums_a & nums_b) / max(1, len(nums_a | nums_b))
211
+
212
+
213
+ def build_dataset(args: argparse.Namespace) -> None:
214
+ label_path = Path(args.aligned_queries)
215
+ out_dir = Path(args.out_dir)
216
+ out_dir.mkdir(parents=True, exist_ok=True)
217
+ labels = read_jsonl(label_path)
218
+ rows_by_qid = load_shard_rows(args.shard_glob)
219
+ samples=[]; rejected=[]
220
+ for row in labels:
221
+ qid = clean_text(row.get("question_id"))
222
+ if qid.endswith("_abs"):
223
+ qid = qid[:-4]
224
+ question = clean_text(row.get("question"))
225
+ positives = list(row.get("positive_event_ids") or [])
226
+ hard_negs = set(row.get("hard_negative_event_ids") or [])
227
+ candidate_ids = list(row.get("candidate_event_ids") or [])
228
+ source_rows = rows_by_qid.get(qid, [])
229
+ source_row = choose_row_for_query(source_rows, candidate_ids, positives)
230
+ if not source_row:
231
+ rejected.append({"question_id": qid, "reason": "missing_source_row"})
232
+ continue
233
+ event_map = event_text_map_for_row(source_row, qid)
234
+ if not candidate_ids:
235
+ rejected.append({"question_id": qid, "reason": "missing_candidate_event_ids"})
236
+ continue
237
+ before_filter_count = len(candidate_ids)
238
+ candidate_ids = [eid for eid in candidate_ids if eid in event_map]
239
+ if len(candidate_ids) != before_filter_count:
240
+ rejected.append({"question_id": qid, "reason": "candidate_ids_not_in_event_map", "candidate_count": before_filter_count, "mapped_count": len(candidate_ids)})
241
+ continue
242
+ if not candidate_ids:
243
+ rejected.append({"question_id": qid, "reason": "no_candidate_events"})
244
+ continue
245
+ if not any(eid in event_map for eid in positives):
246
+ rejected.append({"question_id": qid, "reason": "positive_not_in_event_map", "positive_event_ids": positives[:10]})
247
+ continue
248
+ q_tokens = set(norm_tokens(question))
249
+ candidates=[]
250
+ max_session = max(1, len(source_row.get("haystack_session_ids") or []))
251
+ for eid in candidate_ids:
252
+ info = event_map[eid]
253
+ text = clean_text(info["text"])
254
+ pos = eid in positives
255
+ hard = eid in hard_negs
256
+ candidates.append({
257
+ "event_id": eid,
258
+ "text": text,
259
+ "session_id": info.get("session_id", ""),
260
+ "session_index": int(info.get("session_index", 0)),
261
+ "chunk_index": int(info.get("chunk_index", 0)),
262
+ "features": [
263
+ token_overlap(question, text),
264
+ contains_number_overlap(question, text),
265
+ min(1.0, len(set(norm_tokens(text)) & q_tokens) / max(1, len(q_tokens))),
266
+ float(info.get("session_index", 0)) / float(max_session),
267
+ float(info.get("chunk_index", 0)) / 10.0,
268
+ 1.0 if hard else 0.0,
269
+ ],
270
+ "label": 1 if pos else 0,
271
+ "role_label": 1 if pos else (2 if hard else 0),
272
+ })
273
+ if not any(c["label"] for c in candidates):
274
+ rejected.append({"question_id": qid, "reason": "positive_filtered_out", "positive_event_ids": positives[:10]})
275
+ continue
276
+ metadata = dict(row.get("metadata") or {})
277
+ metadata.setdefault("question_type", clean_text(source_row.get("question_type", "")))
278
+ metadata.setdefault("gold_answer", clean_text(source_row.get("answer", "")))
279
+ metadata.setdefault("answer_session_ids", list(source_row.get("answer_session_ids") or []))
280
+ samples.append({
281
+ "question_id": qid,
282
+ "question": question,
283
+ "question_type": clean_text(source_row.get("question_type", "")),
284
+ "split": split_name(qid, args.val_ratio),
285
+ "positive_event_ids": positives,
286
+ "positive_path_ids": list(row.get("positive_path_ids") or []),
287
+ "metadata": metadata,
288
+ "candidates": candidates,
289
+ })
290
+ write_jsonl(out_dir / "samples.jsonl", samples)
291
+ write_jsonl(out_dir / "rejected.jsonl", rejected)
292
+ report = {
293
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
294
+ "aligned_queries": str(label_path),
295
+ "sample_count": len(samples),
296
+ "rejected_count": len(rejected),
297
+ "train_count": sum(1 for s in samples if s["split"] == "train"),
298
+ "val_count": sum(1 for s in samples if s["split"] == "val"),
299
+ "candidate_count_avg": round(sum(len(s["candidates"]) for s in samples) / max(1, len(samples)), 3),
300
+ "positive_count_avg": round(sum(sum(c["label"] for c in s["candidates"]) for s in samples) / max(1, len(samples)), 3),
301
+ "outputs": {"samples": str(out_dir / "samples.jsonl"), "rejected": str(out_dir / "rejected.jsonl")},
302
+ }
303
+ (out_dir / "dataset_report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
304
+ print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True))
305
+
306
+
307
+ class OpenAIEmbeddingVectorizer:
308
+ def __init__(self, *, dim: int, base_url: str, model: str, api_key: str = ""):
309
+ self.dim = int(dim)
310
+ self.base_url = base_url.rstrip("/")
311
+ self.model = model
312
+ self.api_key = api_key
313
+ self.cache: dict[str, torch.Tensor] = {}
314
+
315
+ def encode_one(self, text: str) -> torch.Tensor:
316
+ text = clean_text(text)
317
+ cached = self.cache.get(text)
318
+ if cached is not None:
319
+ return cached
320
+ payload = json.dumps({"model": self.model, "input": text}, ensure_ascii=False).encode("utf-8")
321
+ headers = {"Content-Type": "application/json"}
322
+ if self.api_key:
323
+ headers["Authorization"] = f"Bearer {self.api_key}"
324
+ req = urllib.request.Request(self.base_url + "/embeddings", data=payload, method="POST", headers=headers)
325
+ with urllib.request.urlopen(req, timeout=120) as resp:
326
+ body = json.loads(resp.read().decode("utf-8", "replace"))
327
+ data = body.get("data") or []
328
+ if not data or "embedding" not in data[0]:
329
+ raise RuntimeError(f"embedding response missing data[0].embedding from {self.base_url}")
330
+ vec = torch.tensor(list(data[0]["embedding"]), dtype=torch.float32)
331
+ if vec.numel() != self.dim:
332
+ raise RuntimeError(f"embedding dim mismatch: got {vec.numel()} expected {self.dim}")
333
+ vec = vec / vec.norm().clamp_min(1e-6)
334
+ self.cache[text] = vec
335
+ return vec
336
+
337
+
338
+ class BgeM3DenseVectorizer:
339
+ def __init__(
340
+ self,
341
+ *,
342
+ dim: int,
343
+ model_path: str,
344
+ device: str = "cpu",
345
+ max_length: int = 8192,
346
+ strict_max_length: bool = False,
347
+ ):
348
+ from transformers import AutoModel, AutoTokenizer
349
+ self.dim = int(dim)
350
+ self.cache: dict[str, torch.Tensor] = {}
351
+ self.device = torch.device(device)
352
+ self.max_length = int(max_length)
353
+ self.strict_max_length = bool(strict_max_length)
354
+ if self.max_length <= 0:
355
+ raise RuntimeError("--embedding-max-length must be positive")
356
+ resolved = Path(model_path).resolve()
357
+ if "bge-m3" not in resolved.name.lower() and "bge-m3" not in str(resolved).lower():
358
+ raise RuntimeError(f"hf backend is currently wired only for local BAAI/bge-m3, got: {model_path}")
359
+ self.model_path = str(resolved)
360
+ self.tokenizer = AutoTokenizer.from_pretrained(
361
+ str(resolved),
362
+ local_files_only=True,
363
+ trust_remote_code=False,
364
+ )
365
+ self.model = AutoModel.from_pretrained(
366
+ str(resolved),
367
+ local_files_only=True,
368
+ trust_remote_code=False,
369
+ ).to(self.device)
370
+ self.model.eval()
371
+ hidden_size = int(getattr(self.model.config, "hidden_size", 0) or 0)
372
+ if hidden_size != self.dim:
373
+ raise RuntimeError(f"bge-m3 hidden_size mismatch: model has {hidden_size}, --text-dim is {self.dim}")
374
+
375
+ def encode_batch(self, texts: Sequence[str], batch_size: int = 4) -> torch.Tensor:
376
+ cleaned = [clean_text(t) for t in texts]
377
+ if not cleaned:
378
+ raise RuntimeError("encode_batch got no texts")
379
+ out_chunks: list[torch.Tensor] = []
380
+ with torch.no_grad():
381
+ for start in range(0, len(cleaned), max(1, int(batch_size))):
382
+ batch_texts = cleaned[start:start + max(1, int(batch_size))]
383
+ encoded = self.tokenizer(
384
+ batch_texts,
385
+ return_tensors="pt",
386
+ padding=True,
387
+ truncation=not self.strict_max_length,
388
+ max_length=self.max_length if not self.strict_max_length else None,
389
+ )
390
+ if self.strict_max_length:
391
+ token_lengths = encoded["attention_mask"].sum(dim=1)
392
+ longest = int(token_lengths.max())
393
+ if longest > self.max_length:
394
+ offending = int(torch.argmax(token_lengths))
395
+ raise RuntimeError(
396
+ f"BGE input has {longest} tokens, exceeding strict max_length={self.max_length}; "
397
+ f"batch_index={offending} chars={len(batch_texts[offending])}"
398
+ )
399
+ encoded = {k: v.to(self.device) for k, v in encoded.items()}
400
+ with torch.autocast(
401
+ device_type=self.device.type,
402
+ dtype=torch.float16,
403
+ enabled=self.device.type == "cuda",
404
+ ):
405
+ model_out = self.model(**encoded)
406
+ if not hasattr(model_out, "last_hidden_state"):
407
+ raise RuntimeError("bge-m3 model output missing last_hidden_state")
408
+ vecs = model_out.last_hidden_state[:, 0].detach().cpu().float()
409
+ if vecs.shape[1] != self.dim:
410
+ raise RuntimeError(f"bge-m3 embedding dim mismatch: got {vecs.shape[1]} expected {self.dim}")
411
+ vecs = vecs / vecs.norm(dim=1, keepdim=True).clamp_min(1e-6)
412
+ out_chunks.append(vecs)
413
+ return torch.cat(out_chunks, dim=0)
414
+
415
+ def encode_one(self, text: str) -> torch.Tensor:
416
+ text = clean_text(text)
417
+ cached = self.cache.get(text)
418
+ if cached is not None:
419
+ return cached
420
+ vec = self.encode_batch([text], batch_size=1)[0]
421
+ self.cache[text] = vec
422
+ return vec
423
+
424
+
425
+ class EmbeddingCacheVectorizer:
426
+ def __init__(self, *, cache_dir: str, dim: int, expected_backend: str = "", expected_model: str = ""):
427
+ root = Path(cache_dir)
428
+ manifest_path = root / "manifest.json"
429
+ tensor_path = root / "embeddings.pt"
430
+ if not manifest_path.exists():
431
+ raise FileNotFoundError(f"embedding cache manifest not found: {manifest_path}")
432
+ if not tensor_path.exists():
433
+ raise FileNotFoundError(f"embedding cache tensor not found: {tensor_path}")
434
+ self.manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
435
+ self.dim = int(dim)
436
+ if int(self.manifest.get("text_dim", 0) or 0) != self.dim:
437
+ raise RuntimeError(f"embedding cache dim mismatch: cache has {self.manifest.get('text_dim')} expected {self.dim}")
438
+ backend = clean_text(self.manifest.get("embedding_backend", ""))
439
+ model = clean_text(self.manifest.get("embedding_model", ""))
440
+ if expected_backend and backend != expected_backend:
441
+ raise RuntimeError(f"embedding cache backend mismatch: cache has {backend} expected {expected_backend}")
442
+ if expected_model and model != expected_model:
443
+ raise RuntimeError(f"embedding cache model mismatch: cache has {model} expected {expected_model}")
444
+ payload = torch.load(tensor_path, map_location="cpu", weights_only=True)
445
+ keys = list(payload.get("keys") or [])
446
+ vectors = payload.get("vectors")
447
+ if not keys or vectors is None:
448
+ raise RuntimeError(f"invalid embedding cache payload: {tensor_path}")
449
+ if len(keys) != int(vectors.shape[0]):
450
+ raise RuntimeError(f"embedding cache row mismatch: {len(keys)} keys vs {vectors.shape[0]} vectors")
451
+ if int(vectors.shape[1]) != self.dim:
452
+ raise RuntimeError(f"embedding cache tensor dim mismatch: {vectors.shape[1]} expected {self.dim}")
453
+ self.keys = keys
454
+ self.index = {k: i for i, k in enumerate(keys)}
455
+ if len(self.index) != len(keys):
456
+ raise RuntimeError("embedding cache contains duplicate keys")
457
+ self.vectors = vectors.float().contiguous()
458
+
459
+ def encode_one(self, text: str) -> torch.Tensor:
460
+ key = embedding_text_key(text)
461
+ idx = self.index.get(key)
462
+ if idx is None:
463
+ preview = clean_text(text)[:120]
464
+ raise RuntimeError(f"embedding cache miss: key={key} preview={preview!r}")
465
+ return self.vectors[idx]
466
+
467
+
468
+ def build_vectorizer(args: argparse.Namespace, device: torch.device):
469
+ backend = clean_text(getattr(args, "embedding_backend", ""))
470
+ cache_dir = clean_text(getattr(args, "embedding_cache", ""))
471
+ if not backend:
472
+ raise RuntimeError("embedding backend is required; pass --embedding-backend openai or --embedding-backend hf")
473
+ if cache_dir:
474
+ return EmbeddingCacheVectorizer(
475
+ cache_dir=cache_dir,
476
+ dim=args.text_dim,
477
+ expected_backend=backend,
478
+ expected_model=clean_text(getattr(args, "embedding_model", "")),
479
+ )
480
+ if backend == "openai":
481
+ base_url = clean_text(getattr(args, "embedding_base_url", ""))
482
+ model = clean_text(getattr(args, "embedding_model", ""))
483
+ api_key = clean_text(getattr(args, "embedding_api_key", ""))
484
+ if not base_url or not model:
485
+ raise RuntimeError("openai embedding backend requires explicit --embedding-base-url and --embedding-model")
486
+ return OpenAIEmbeddingVectorizer(dim=args.text_dim, base_url=base_url, model=model, api_key=api_key)
487
+ if backend == "hf":
488
+ model_path = clean_text(getattr(args, "embedding_model", ""))
489
+ if not model_path:
490
+ raise RuntimeError("hf embedding backend requires explicit --embedding-model local path")
491
+ if not os.path.exists(model_path):
492
+ raise RuntimeError(f"hf embedding backend requires a local model path; not found: {model_path}")
493
+ return BgeM3DenseVectorizer(dim=args.text_dim, model_path=model_path, device=str(device), max_length=args.embedding_max_length)
494
+ raise RuntimeError(f"unsupported embedding backend: {backend}")
495
+
496
+
497
+ class V2EvidenceScorer(nn.Module):
498
+ def __init__(self, text_dim: int, feature_dim: int, hidden_dim: int = 256, layers: int = 2, roles: int = 3):
499
+ super().__init__()
500
+ self.query_proj = nn.Sequential(nn.LayerNorm(text_dim), nn.Linear(text_dim, hidden_dim), nn.SiLU())
501
+ self.event_proj = nn.Sequential(nn.LayerNorm(text_dim), nn.Linear(text_dim, hidden_dim), nn.SiLU())
502
+ self.feature_proj = nn.Sequential(nn.LayerNorm(feature_dim), nn.Linear(feature_dim, hidden_dim // 2), nn.SiLU())
503
+ self.session_embedding = nn.Embedding(256, hidden_dim)
504
+ self.chunk_embedding = nn.Embedding(32, hidden_dim)
505
+ encoder_layer = nn.TransformerEncoderLayer(
506
+ d_model=hidden_dim,
507
+ nhead=8,
508
+ dim_feedforward=hidden_dim * 4,
509
+ dropout=0.1,
510
+ activation="gelu",
511
+ batch_first=True,
512
+ norm_first=True,
513
+ )
514
+ self.graph_encoder = nn.TransformerEncoder(encoder_layer, num_layers=max(1, int(layers)))
515
+ pair_dim = hidden_dim * 4 + hidden_dim // 2
516
+ self.event_head = nn.Sequential(nn.LayerNorm(pair_dim), nn.Linear(pair_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1))
517
+ self.pack_head = nn.Sequential(nn.LayerNorm(pair_dim), nn.Linear(pair_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 1))
518
+ self.role_head = nn.Sequential(nn.LayerNorm(pair_dim), nn.Linear(pair_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, roles))
519
+
520
+ def forward(self, q_vec: torch.Tensor, e_vec: torch.Tensor, features: torch.Tensor, session_idx: torch.Tensor, chunk_idx: torch.Tensor, mask: torch.Tensor) -> dict[str, torch.Tensor]:
521
+ q = self.query_proj(q_vec)
522
+ e = self.event_proj(e_vec)
523
+ f = self.feature_proj(features)
524
+ session_idx = session_idx.clamp_min(0).clamp_max(255)
525
+ chunk_idx = chunk_idx.clamp_min(0).clamp_max(31)
526
+ h = e + self.session_embedding(session_idx) + self.chunk_embedding(chunk_idx)
527
+ h = self.graph_encoder(h, src_key_padding_mask=~mask)
528
+ qx = q.unsqueeze(1).expand_as(h)
529
+ pair = torch.cat([qx, h, qx * h, torch.abs(qx - h), f], dim=-1)
530
+ event_logits = self.event_head(pair).squeeze(-1).masked_fill(~mask, -1e4)
531
+ pack_logits = self.pack_head(pair).squeeze(-1).masked_fill(~mask, -1e4)
532
+ role_logits = self.role_head(pair).masked_fill(~mask.unsqueeze(-1), 0.0)
533
+ return {"event_logits": event_logits, "pack_logits": pack_logits, "role_logits": role_logits}
534
+
535
+
536
+ @dataclass
537
+ class Batch:
538
+ q_vec: torch.Tensor
539
+ e_vec: torch.Tensor
540
+ features: torch.Tensor
541
+ session_idx: torch.Tensor
542
+ chunk_idx: torch.Tensor
543
+ mask: torch.Tensor
544
+ labels: torch.Tensor
545
+ role_labels: torch.Tensor
546
+ qids: list[str]
547
+ event_ids: list[list[str]]
548
+
549
+
550
+ def make_batch(samples: Sequence[Mapping[str, Any]], vectorizer: Any, text_dim: int, feature_dim: int, device: torch.device) -> Batch:
551
+ max_c = max(len(s["candidates"]) for s in samples)
552
+ b = len(samples)
553
+ q_vec = torch.zeros((b, text_dim), dtype=torch.float32)
554
+ e_vec = torch.zeros((b, max_c, text_dim), dtype=torch.float32)
555
+ features = torch.zeros((b, max_c, feature_dim), dtype=torch.float32)
556
+ session_idx = torch.zeros((b, max_c), dtype=torch.long)
557
+ chunk_idx = torch.zeros((b, max_c), dtype=torch.long)
558
+ mask = torch.zeros((b, max_c), dtype=torch.bool)
559
+ labels = torch.zeros((b, max_c), dtype=torch.float32)
560
+ role_labels = torch.zeros((b, max_c), dtype=torch.long)
561
+ qids=[]; event_ids=[]
562
+ for i, sample in enumerate(samples):
563
+ qids.append(str(sample["question_id"])); event_ids.append([])
564
+ q_vec[i] = vectorizer.encode_one(str(sample.get("question", "")))
565
+ for j, cand in enumerate(sample["candidates"]):
566
+ e_vec[i, j] = vectorizer.encode_one(str(cand.get("text", "")))
567
+ feat = list(cand.get("features") or [])
568
+ features[i, j, : min(feature_dim, len(feat))] = torch.tensor(feat[:feature_dim], dtype=torch.float32)
569
+ session_idx[i, j] = int(cand.get("session_index", 0) or 0)
570
+ chunk_idx[i, j] = int(cand.get("chunk_index", 0) or 0)
571
+ mask[i, j] = True
572
+ labels[i, j] = float(cand.get("label", 0) or 0)
573
+ role_labels[i, j] = int(cand.get("role_label", 0) or 0)
574
+ event_ids[-1].append(str(cand.get("event_id", "")))
575
+ return Batch(q_vec.to(device), e_vec.to(device), features.to(device), session_idx.to(device), chunk_idx.to(device), mask.to(device), labels.to(device), role_labels.to(device), qids, event_ids)
576
+
577
+
578
+ def ranking_loss(event_logits: torch.Tensor, labels: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
579
+ losses=[]
580
+ for logits, y, m in zip(event_logits, labels, mask):
581
+ valid_logits = logits[m]
582
+ valid_y = y[m]
583
+ pos = valid_y > 0.5
584
+ if not bool(pos.any()):
585
+ continue
586
+ log_probs = F.log_softmax(valid_logits, dim=0)
587
+ target = valid_y / valid_y.sum().clamp_min(1.0)
588
+ losses.append(-(target * log_probs).sum())
589
+ if not losses:
590
+ return event_logits.sum() * 0.0
591
+ return torch.stack(losses).mean()
592
+
593
+
594
+ def eval_model(model: V2EvidenceScorer, samples: Sequence[Mapping[str, Any]], vectorizer: Any, args: argparse.Namespace, device: torch.device) -> dict[str, Any]:
595
+ model.eval()
596
+ totals={"n":0,"event_recall_at_1":0,"event_recall_at_5":0,"event_recall_at_10":0,"pack_recall_at_10":0,"mrr":0.0}
597
+ with torch.no_grad():
598
+ for start in range(0, len(samples), args.batch_size):
599
+ batch_samples=samples[start:start+args.batch_size]
600
+ batch=make_batch(batch_samples, vectorizer, args.text_dim, args.feature_dim, device)
601
+ out=model(batch.q_vec,batch.e_vec,batch.features,batch.session_idx,batch.chunk_idx,batch.mask)
602
+ event_scores=out["event_logits"].detach().cpu()
603
+ pack_scores=out["pack_logits"].detach().cpu()
604
+ labels=batch.labels.detach().cpu()
605
+ mask=batch.mask.detach().cpu()
606
+ for i in range(len(batch_samples)):
607
+ valid=torch.nonzero(mask[i], as_tuple=False).squeeze(-1)
608
+ if valid.numel()==0: continue
609
+ pos=set(torch.nonzero(labels[i][valid]>0.5, as_tuple=False).squeeze(-1).tolist())
610
+ if not pos: continue
611
+ order=torch.argsort(event_scores[i][valid], descending=True).tolist()
612
+ pack_order=torch.argsort(pack_scores[i][valid], descending=True).tolist()
613
+ totals["n"]+=1
614
+ totals["event_recall_at_1"] += int(any(idx in pos for idx in order[:1]))
615
+ totals["event_recall_at_5"] += int(any(idx in pos for idx in order[:5]))
616
+ totals["event_recall_at_10"] += int(any(idx in pos for idx in order[:10]))
617
+ totals["pack_recall_at_10"] += int(any(idx in pos for idx in pack_order[:10]))
618
+ first_rank=next((r+1 for r,idx in enumerate(order) if idx in pos), None)
619
+ totals["mrr"] += 0.0 if first_rank is None else 1.0/first_rank
620
+ n=max(1, totals["n"])
621
+ return {k:(round(v/n,6) if k!="n" else v) for k,v in totals.items()}
622
+
623
+
624
+ def precompute_embeddings(args: argparse.Namespace) -> None:
625
+ out_dir = Path(args.out_dir)
626
+ out_dir.mkdir(parents=True, exist_ok=True)
627
+ tensor_path = out_dir / "embeddings.pt"
628
+ manifest_path = out_dir / "manifest.json"
629
+ items_path = out_dir / "items.jsonl"
630
+ if (tensor_path.exists() or manifest_path.exists() or items_path.exists()) and not args.overwrite:
631
+ raise RuntimeError(f"embedding cache output already exists; pass --overwrite explicitly: {out_dir}")
632
+ samples = read_jsonl(Path(args.samples))
633
+ if args.limit > 0:
634
+ samples = samples[: args.limit]
635
+ if not samples:
636
+ raise RuntimeError("no samples for embedding precompute")
637
+ if not args.cpu and not torch.cuda.is_available():
638
+ raise RuntimeError("CUDA is unavailable; pass --cpu explicitly for CPU precompute")
639
+ if args.embedding_backend != "hf":
640
+ raise RuntimeError("precompute currently supports only --embedding-backend hf for local bge-m3")
641
+ device = torch.device("cpu" if args.cpu else "cuda")
642
+ vectorizer = BgeM3DenseVectorizer(
643
+ dim=args.text_dim,
644
+ model_path=args.embedding_model,
645
+ device=str(device),
646
+ max_length=args.embedding_max_length,
647
+ )
648
+ items = iter_sample_texts(samples)
649
+ keys = [item["key"] for item in items]
650
+ texts = [item["text"] for item in items]
651
+ vectors: list[torch.Tensor] = []
652
+ total = len(texts)
653
+ started = time.time()
654
+ for start in range(0, total, args.batch_size):
655
+ batch_texts = texts[start:start + args.batch_size]
656
+ vecs = vectorizer.encode_batch(batch_texts, batch_size=args.batch_size)
657
+ vectors.append(vecs.to(dtype=torch.float16 if args.dtype == "float16" else torch.float32))
658
+ done = min(total, start + len(batch_texts))
659
+ if done == total or done % max(1, args.log_every) == 0:
660
+ elapsed = time.time() - started
661
+ print(json.dumps({"done": done, "total": total, "elapsed_sec": round(elapsed, 3)}, ensure_ascii=False), flush=True)
662
+ matrix = torch.cat(vectors, dim=0).contiguous()
663
+ if matrix.shape != (len(keys), args.text_dim):
664
+ raise RuntimeError(f"embedding matrix shape mismatch: got {tuple(matrix.shape)} expected {(len(keys), args.text_dim)}")
665
+ torch.save({"keys": keys, "vectors": matrix}, tensor_path)
666
+ item_rows=[]
667
+ for item in items:
668
+ item_rows.append({
669
+ "key": item["key"],
670
+ "kind": item.get("kind", ""),
671
+ "question_id": item.get("question_id", ""),
672
+ "event_id": item.get("event_id", ""),
673
+ "text_len": len(item.get("text", "")),
674
+ "preview": clean_text(item.get("text", ""))[:200],
675
+ })
676
+ write_jsonl(items_path, item_rows)
677
+ manifest = {
678
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
679
+ "samples": str(Path(args.samples).resolve()),
680
+ "sample_count": len(samples),
681
+ "unique_text_count": len(keys),
682
+ "text_dim": int(args.text_dim),
683
+ "dtype": args.dtype,
684
+ "embedding_backend": args.embedding_backend,
685
+ "embedding_model": clean_text(args.embedding_model),
686
+ "embedding_max_length": int(args.embedding_max_length),
687
+ "device": str(device),
688
+ "outputs": {"tensor": str(tensor_path), "items": str(items_path), "manifest": str(manifest_path)},
689
+ }
690
+ manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
691
+ print(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True))
692
+
693
+
694
+ def smoke(args: argparse.Namespace) -> None:
695
+ out_dir = Path(args.out_dir)
696
+ out_dir.mkdir(parents=True, exist_ok=True)
697
+ samples = read_jsonl(Path(args.samples))
698
+ if args.limit > 0:
699
+ samples = samples[: args.limit]
700
+ if not samples:
701
+ raise RuntimeError("no samples for smoke")
702
+ if not args.cpu and not torch.cuda.is_available():
703
+ raise RuntimeError("CUDA is unavailable; pass --cpu explicitly for CPU smoke")
704
+ device = torch.device("cpu" if args.cpu else "cuda")
705
+ vectorizer = build_vectorizer(args, device)
706
+ model = V2EvidenceScorer(args.text_dim, args.feature_dim, args.hidden_dim, args.layers).to(device)
707
+ model.eval()
708
+ batch = make_batch(samples, vectorizer, args.text_dim, args.feature_dim, device)
709
+ with torch.no_grad():
710
+ out = model(batch.q_vec, batch.e_vec, batch.features, batch.session_idx, batch.chunk_idx, batch.mask)
711
+ loss_rank = ranking_loss(out["event_logits"], batch.labels, batch.mask)
712
+ pack_loss = F.binary_cross_entropy_with_logits(out["pack_logits"][batch.mask], batch.labels[batch.mask])
713
+ role_loss = F.cross_entropy(out["role_logits"][batch.mask], batch.role_labels[batch.mask])
714
+ metrics = eval_model(model, samples, vectorizer, args, device)
715
+ report = {
716
+ "status": "ok",
717
+ "mode": "forward_only_no_optimizer_step",
718
+ "device": str(device),
719
+ "sample_count": len(samples),
720
+ "candidate_count_max": max(len(s.get("candidates", [])) for s in samples),
721
+ "candidate_count_avg": round(sum(len(s.get("candidates", [])) for s in samples) / max(1, len(samples)), 3),
722
+ "losses": {
723
+ "ranking_loss": round(float(loss_rank.detach().cpu()), 6),
724
+ "pack_loss": round(float(pack_loss.detach().cpu()), 6),
725
+ "role_loss": round(float(role_loss.detach().cpu()), 6),
726
+ },
727
+ "metrics": metrics,
728
+ "embedding_backend": clean_text(getattr(args, "embedding_backend", "")),
729
+ "embedding_model": clean_text(getattr(args, "embedding_model", "")),
730
+ "embedding_cache": clean_text(getattr(args, "embedding_cache", "")),
731
+ "embedding_max_length": int(getattr(args, "embedding_max_length", 0) or 0),
732
+ "text_dim": int(args.text_dim),
733
+ "hidden_dim": int(args.hidden_dim),
734
+ "layers": int(args.layers),
735
+ }
736
+ (out_dir / "smoke_report.json").write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
737
+ print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True))
738
+
739
+
740
+ def train(args: argparse.Namespace) -> None:
741
+ out_dir=Path(args.out_dir); out_dir.mkdir(parents=True, exist_ok=True)
742
+ samples=read_jsonl(Path(args.samples))
743
+ if args.limit>0:
744
+ samples=samples[:args.limit]
745
+ train_samples=[s for s in samples if s.get("split") == "train"]
746
+ val_samples=[s for s in samples if s.get("split") == "val"]
747
+ if not train_samples:
748
+ raise RuntimeError("no train samples")
749
+ if not val_samples:
750
+ raise RuntimeError("no val samples; rebuild dataset with a non-empty validation split")
751
+ if not args.cpu and not torch.cuda.is_available():
752
+ raise RuntimeError("CUDA is unavailable; pass --cpu explicitly for CPU train")
753
+ device=torch.device("cpu" if args.cpu else "cuda")
754
+ vectorizer=build_vectorizer(args, device)
755
+ model=V2EvidenceScorer(args.text_dim,args.feature_dim,args.hidden_dim,args.layers).to(device)
756
+ opt=torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
757
+ best=None
758
+ history=[]
759
+ rng=random.Random(args.seed)
760
+ for epoch in range(1,args.epochs+1):
761
+ model.train(); rng.shuffle(train_samples)
762
+ losses=[]
763
+ for start in range(0,len(train_samples),args.batch_size):
764
+ batch_samples=train_samples[start:start+args.batch_size]
765
+ batch=make_batch(batch_samples, vectorizer, args.text_dim, args.feature_dim, device)
766
+ out=model(batch.q_vec,batch.e_vec,batch.features,batch.session_idx,batch.chunk_idx,batch.mask)
767
+ loss_rank=ranking_loss(out["event_logits"],batch.labels,batch.mask)
768
+ bce=F.binary_cross_entropy_with_logits(out["pack_logits"][batch.mask], batch.labels[batch.mask])
769
+ role=F.cross_entropy(out["role_logits"][batch.mask], batch.role_labels[batch.mask])
770
+ loss=loss_rank + args.pack_loss_weight*bce + args.role_loss_weight*role
771
+ opt.zero_grad(set_to_none=True); loss.backward(); nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
772
+ losses.append(float(loss.detach().cpu()))
773
+ val=eval_model(model,val_samples,vectorizer,args,device)
774
+ train_eval=eval_model(model,train_samples[:min(len(train_samples),100)],vectorizer,args,device)
775
+ summary={"epoch":epoch,"loss":round(sum(losses)/max(1,len(losses)),6),"train":train_eval,"val":val}
776
+ history.append(summary)
777
+ print(json.dumps(summary,ensure_ascii=False), flush=True)
778
+ score=val.get("event_recall_at_5",0.0)+0.1*val.get("mrr",0.0)
779
+ if best is None or score > best[0]:
780
+ best=(score,epoch)
781
+ torch.save({"model_state":model.state_dict(),"args":vars(args),"epoch":epoch,"val":val}, out_dir/"tmcra_v2_scorer.pt")
782
+ (out_dir/"train_history.json").write_text(json.dumps(history,ensure_ascii=False,indent=2),encoding="utf-8")
783
+ report={"best": {"score": best[0] if best else None, "epoch": best[1] if best else None}, "final": history[-1] if history else {}, "checkpoint": str(out_dir/"tmcra_v2_scorer.pt")}
784
+ (out_dir/"train_report.json").write_text(json.dumps(report,ensure_ascii=False,indent=2,sort_keys=True),encoding="utf-8")
785
+ print(json.dumps(report,ensure_ascii=False,indent=2,sort_keys=True))
786
+
787
+
788
+ def main() -> int:
789
+ p=argparse.ArgumentParser(description="TMCRA v2 LongMemEval semantic-graph scorer pipeline")
790
+ sub=p.add_subparsers(dest="cmd", required=True)
791
+ b=sub.add_parser("build")
792
+ b.add_argument("--aligned-queries", required=True)
793
+ b.add_argument("--shard-glob", default=os.getenv("TMCRA_SHARD_GLOB", "runs/shard_*.json"))
794
+ b.add_argument("--out-dir", required=True)
795
+ b.add_argument("--val-ratio", type=float, default=0.2)
796
+ pc=sub.add_parser("precompute")
797
+ pc.add_argument("--samples", required=True)
798
+ pc.add_argument("--out-dir", required=True)
799
+ pc.add_argument("--text-dim", type=int, default=1024)
800
+ pc.add_argument("--embedding-backend", choices=["hf"], required=True)
801
+ pc.add_argument("--embedding-model", required=True)
802
+ pc.add_argument("--embedding-max-length", type=int, default=8192)
803
+ pc.add_argument("--batch-size", type=int, default=2)
804
+ pc.add_argument("--log-every", type=int, default=100)
805
+ pc.add_argument("--dtype", choices=["float16", "float32"], default="float16")
806
+ pc.add_argument("--limit", type=int, default=0)
807
+ pc.add_argument("--overwrite", action="store_true")
808
+ pc.add_argument("--cpu", action="store_true")
809
+ t=sub.add_parser("train")
810
+ t.add_argument("--samples", required=True)
811
+ t.add_argument("--out-dir", required=True)
812
+ t.add_argument("--text-dim", type=int, default=1024)
813
+ t.add_argument("--embedding-backend", choices=["openai", "hf"], required=True)
814
+ t.add_argument("--embedding-base-url", default="")
815
+ t.add_argument("--embedding-model", default="")
816
+ t.add_argument("--embedding-cache", default="")
817
+ t.add_argument("--embedding-api-key", default="")
818
+ t.add_argument("--embedding-max-length", type=int, default=8192)
819
+ t.add_argument("--feature-dim", type=int, default=6)
820
+ t.add_argument("--hidden-dim", type=int, default=256)
821
+ t.add_argument("--layers", type=int, default=2)
822
+ t.add_argument("--batch-size", type=int, default=8)
823
+ t.add_argument("--epochs", type=int, default=8)
824
+ t.add_argument("--lr", type=float, default=2e-4)
825
+ t.add_argument("--weight-decay", type=float, default=0.01)
826
+ t.add_argument("--pack-loss-weight", type=float, default=0.3)
827
+ t.add_argument("--role-loss-weight", type=float, default=0.15)
828
+ t.add_argument("--seed", type=int, default=13)
829
+ t.add_argument("--limit", type=int, default=0)
830
+ t.add_argument("--cpu", action="store_true")
831
+ sm=sub.add_parser("smoke")
832
+ sm.add_argument("--samples", required=True)
833
+ sm.add_argument("--out-dir", required=True)
834
+ sm.add_argument("--text-dim", type=int, default=1024)
835
+ sm.add_argument("--embedding-backend", choices=["openai", "hf"], required=True)
836
+ sm.add_argument("--embedding-base-url", default="")
837
+ sm.add_argument("--embedding-model", default="")
838
+ sm.add_argument("--embedding-cache", default="")
839
+ sm.add_argument("--embedding-api-key", default="")
840
+ sm.add_argument("--embedding-max-length", type=int, default=8192)
841
+ sm.add_argument("--feature-dim", type=int, default=6)
842
+ sm.add_argument("--hidden-dim", type=int, default=96)
843
+ sm.add_argument("--layers", type=int, default=1)
844
+ sm.add_argument("--batch-size", type=int, default=8)
845
+ sm.add_argument("--limit", type=int, default=4)
846
+ sm.add_argument("--cpu", action="store_true")
847
+ args=p.parse_args()
848
+ if args.cmd == "build": build_dataset(args)
849
+ elif args.cmd == "precompute": precompute_embeddings(args)
850
+ elif args.cmd == "train": train(args)
851
+ elif args.cmd == "smoke": smoke(args)
852
+ return 0
853
+
854
+ if __name__ == "__main__":
855
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_online_runtime.py ADDED
@@ -0,0 +1,1867 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import importlib.util
7
+ import json
8
+ import os
9
+ import re
10
+ import sqlite3
11
+ import sys
12
+ import time
13
+ from collections import Counter
14
+ from pathlib import Path
15
+ from typing import Any, Iterable, Mapping, Sequence
16
+
17
+ import torch
18
+
19
+ from build_v3_runtime_dataset import covered_windows, rrank
20
+ from tmcra_v2_lme_pipeline import BgeM3DenseVectorizer
21
+ from tmcra_v3_reranker import ChannelAwareMemoryReranker
22
+ from tmcra_model_assets import load_optional_hf_model_manifest
23
+ from tmcra_v3_schema import CHANNEL_NAMES, SCHEMA_VERSION, clean_text, write_jsonl
24
+ from tmcra_v3_recall_planner import (
25
+ DEEPSEEK_FLASH_MODEL,
26
+ DeepSeekFlashRecallPlanner,
27
+ RecallPlannerError,
28
+ apply_recall_plan,
29
+ )
30
+
31
+
32
+ LEGACY_CHUNK_ID_RE = re.compile(r"^s(?P<session>\d+)_c(?P<parent>\d+)$")
33
+ MESSAGE_ID_RE = re.compile(r"^s(?P<session>\d+)_m(?P<parent>\d+)$")
34
+ EVENT_PARENT_RE = re.compile(r":s(?P<session>\d+)_[cm](?P<parent>\d+)$")
35
+ CURRENT_FAST_STATES = frozenset(
36
+ {"active", "parallel_active", "promoted", "challenged"}
37
+ )
38
+ ONLINE_INDEX_SCHEMA_VERSION = "tmcra.v3.online-index.3"
39
+ FAST_SEMANTIC_STATE_POLICY = (
40
+ "current-fast-states-v1:" + ",".join(sorted(CURRENT_FAST_STATES))
41
+ )
42
+ INGEST_PREFIX_RE = re.compile(r"^\[[^\]\n]+\]\s+user:\s*", flags=re.IGNORECASE)
43
+ SESSION_HEADER_RE = re.compile(
44
+ r"^LongMemEval session_id=(?P<session_id>\S+) date=(?P<date>.+?)(?: continued=true)? \[",
45
+ flags=re.DOTALL,
46
+ )
47
+
48
+
49
+ def read_jsonl(path: Path) -> list[dict[str, Any]]:
50
+ rows: list[dict[str, Any]] = []
51
+ with path.open("r", encoding="utf-8", errors="strict") as handle:
52
+ for line_no, line in enumerate(handle, start=1):
53
+ if not line.strip():
54
+ continue
55
+ value = json.loads(line)
56
+ if not isinstance(value, dict):
57
+ raise RuntimeError(f"row is not an object at {path}:{line_no}")
58
+ rows.append(value)
59
+ if not rows:
60
+ raise RuntimeError(f"no rows: {path}")
61
+ return rows
62
+
63
+
64
+ def atomic_torch_save(payload: Mapping[str, Any], target: Path) -> None:
65
+ target.parent.mkdir(parents=True, exist_ok=True)
66
+ temporary = target.with_suffix(target.suffix + ".tmp")
67
+ torch.save(dict(payload), temporary)
68
+ os.replace(temporary, target)
69
+
70
+
71
+ def sha256_file(path: Path) -> str:
72
+ digest = hashlib.sha256()
73
+ with path.open("rb") as handle:
74
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
75
+ digest.update(chunk)
76
+ return digest.hexdigest()
77
+
78
+
79
+ def load_native_harness(path: Path, repo: Path):
80
+ resolved_repo = str(repo.resolve())
81
+ if resolved_repo not in sys.path:
82
+ sys.path.insert(0, resolved_repo)
83
+ module_name = "tmcra_v3_native_runtime_harness"
84
+ spec = importlib.util.spec_from_file_location(module_name, path)
85
+ if spec is None or spec.loader is None:
86
+ raise RuntimeError(f"cannot import native TMCRA harness: {path}")
87
+ module = importlib.util.module_from_spec(spec)
88
+ sys.modules[module_name] = module
89
+ spec.loader.exec_module(module)
90
+ return module
91
+
92
+
93
+ def graph_runtime_env(args: argparse.Namespace) -> dict[str, str]:
94
+ values = {
95
+ "TMCRA_RETRIEVAL_MODE": "hybrid_node_scored",
96
+ "TMCRA_NODE_MODEL_PATH": str(Path(args.node_model).resolve()),
97
+ "TMCRA_PATH_MODEL_PATH": str(Path(args.path_model).resolve()),
98
+ "TMCRA_NODE_MODEL_DEVICE": args.graph_device,
99
+ "TMCRA_SUPPORT_PATH_K": str(args.support_path_k),
100
+ "TMCRA_PATH_TUNNEL_RESCUE_K": str(args.path_tunnel_rescue_k),
101
+ "TMCRA_CANDIDATE_EVENT_K": str(args.candidate_event_k),
102
+ "TMCRA_DEEPSEEK_GRAPH_MODEL_MODE": "off",
103
+ "TMCRA_TOPIC_BUCKET_MODE": "off",
104
+ "TMCRA_LLM_CHANNEL_PLANNER_MODE": "off",
105
+ "TMCRA_LLM_EVIDENCE_SELECTOR_MODE": "off",
106
+ "TMCRA_EVIDENCE_UNIT_PLANNER_MODE": "off",
107
+ "TMCRA_UNIFIED_OPERATION_PLANNER_MODE": "off",
108
+ }
109
+ for name, value in values.items():
110
+ os.environ[name] = value
111
+ return values
112
+
113
+
114
+ def scope_counts(db_path: Path, scope_id: str) -> dict[str, int]:
115
+ with sqlite3.connect(db_path) as connection:
116
+ output = {}
117
+ for table in ("records", "memory_edges", "audit_turn_log", "audit_retrieval_log"):
118
+ output[table] = int(
119
+ connection.execute(f'SELECT COUNT(*) FROM "{table}" WHERE scope_id=?', (scope_id,)).fetchone()[0]
120
+ )
121
+ return output
122
+
123
+
124
+ def scope_fingerprint(db_path: Path, scope_id: str) -> str:
125
+ relevant_tables = ("records", "memory_edges", "slot_heads", "slot_history")
126
+ snapshot: dict[str, Any] = {}
127
+ with sqlite3.connect(db_path) as connection:
128
+ connection.row_factory = sqlite3.Row
129
+ for table in relevant_tables:
130
+ exists = connection.execute(
131
+ "SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,)
132
+ ).fetchone()
133
+ if exists is None:
134
+ snapshot[table] = None
135
+ continue
136
+ rows = []
137
+ for row in connection.execute(
138
+ f'SELECT * FROM "{table}" WHERE scope_id=?', (scope_id,)
139
+ ):
140
+ normalized: dict[str, Any] = {}
141
+ for key in row.keys():
142
+ value = row[key]
143
+ if key.endswith("_json") and isinstance(value, str):
144
+ try:
145
+ value = json.loads(value)
146
+ except json.JSONDecodeError as exc:
147
+ raise RuntimeError(
148
+ f"invalid JSON in {table}.{key} while fingerprinting scope"
149
+ ) from exc
150
+ normalized[key] = value
151
+ rows.append(normalized)
152
+ snapshot[table] = sorted(
153
+ rows,
154
+ key=lambda item: json.dumps(
155
+ item, ensure_ascii=False, sort_keys=True, separators=(",", ":")
156
+ ),
157
+ )
158
+ return hashlib.sha256(
159
+ json.dumps(
160
+ snapshot, ensure_ascii=False, sort_keys=True, separators=(",", ":")
161
+ ).encode("utf-8")
162
+ ).hexdigest()
163
+
164
+
165
+ def load_recent_dialogue_context(
166
+ db_path: Path,
167
+ scope_id: str,
168
+ *,
169
+ current_query: str,
170
+ limit: int = 8,
171
+ ) -> list[dict[str, Any]]:
172
+ """Load a small, metadata-free dialogue tail for planner reference resolution."""
173
+ if limit <= 0 or limit > 8:
174
+ raise RuntimeError("recent dialogue limit must be between 1 and 8")
175
+ with sqlite3.connect(db_path) as connection:
176
+ columns = {
177
+ str(row[1]) for row in connection.execute('PRAGMA table_info("audit_turn_log")')
178
+ }
179
+ required = {"scope_id", "event_index", "payload_json"}
180
+ if not required.issubset(columns):
181
+ raise RuntimeError(
182
+ "audit_turn_log lacks required production dialogue columns: "
183
+ + ",".join(sorted(required - columns))
184
+ )
185
+ rows = connection.execute(
186
+ 'SELECT event_index,payload_json FROM "audit_turn_log" '
187
+ "WHERE scope_id=? ORDER BY event_index DESC LIMIT ?",
188
+ (scope_id, limit + 1),
189
+ ).fetchall()
190
+ dialogue: list[dict[str, Any]] = []
191
+ for event_index, raw_payload in reversed(rows):
192
+ try:
193
+ payload = json.loads(raw_payload)
194
+ except (TypeError, json.JSONDecodeError) as exc:
195
+ raise RuntimeError(f"invalid audit_turn_log payload at event_index={event_index}") from exc
196
+ if not isinstance(payload, Mapping):
197
+ raise RuntimeError(f"audit_turn_log payload is not an object at event_index={event_index}")
198
+ speaker = clean_text(payload.get("speaker")).lower()
199
+ text = clean_text(payload.get("text"))
200
+ if speaker not in {"user", "assistant"} or not text:
201
+ raise RuntimeError(f"audit_turn_log turn is not planner-safe at event_index={event_index}")
202
+ dialogue.append(
203
+ {"turn_index": int(event_index), "speaker": speaker, "text": text}
204
+ )
205
+ current = clean_text(current_query)
206
+ if dialogue and dialogue[-1]["speaker"] == "user" and dialogue[-1]["text"] == current:
207
+ dialogue.pop()
208
+ return dialogue[-limit:]
209
+
210
+
211
+ def append_layered_retrieval_audit(
212
+ *,
213
+ repo: Path,
214
+ db_path: Path,
215
+ scope_id: str,
216
+ operation_id: str,
217
+ evidence: Mapping[str, Any],
218
+ debug: Mapping[str, Any],
219
+ ) -> dict[str, Any]:
220
+ """Persist the final layered retrieval decision without rewriting graph state."""
221
+ resolved_repo = str(repo.resolve())
222
+ if resolved_repo not in sys.path:
223
+ sys.path.insert(0, resolved_repo)
224
+ from experiments.replacement.memory_graph import SQLiteSessionMemoryStore
225
+
226
+ with sqlite3.connect(db_path) as connection:
227
+ row = connection.execute(
228
+ "SELECT value_json FROM meta WHERE scope_id=? AND key='audit_retention'",
229
+ (scope_id,),
230
+ ).fetchone()
231
+ if row is None:
232
+ raise RuntimeError(f"{scope_id}: missing audit_retention metadata")
233
+ audit_retention = int(json.loads(row[0]) or 0)
234
+ if audit_retention <= 0:
235
+ raise RuntimeError(f"{scope_id}: invalid audit_retention metadata: {audit_retention}")
236
+
237
+ windows = list(evidence.get("evidence_windows") or [])
238
+ memory_contexts = [
239
+ context
240
+ for window in windows
241
+ for context in list(window.get("memory_contexts") or [])
242
+ ]
243
+ graph = dict(debug.get("graph") or {})
244
+ payload = {
245
+ "event_kind": "tmcra.v3.layered_retrieval",
246
+ "schema_version": "tmcra.v3.layered-retrieval-audit.1",
247
+ "operation_id": clean_text(operation_id),
248
+ "question_id": clean_text(evidence.get("question_id")),
249
+ "query": clean_text(evidence.get("question")),
250
+ "question_date": clean_text(evidence.get("question_date")),
251
+ "recall_plan": dict(evidence.get("recall_plan") or {}),
252
+ "selected_memory_ids": [
253
+ clean_text(window.get("memory_id"))
254
+ for window in windows
255
+ if clean_text(window.get("memory_id"))
256
+ ],
257
+ "selected_session_ids": list(evidence.get("selected_session_ids") or []),
258
+ "selected_capsule_ids": sorted(
259
+ {
260
+ clean_text(context.get("capsule_id"))
261
+ for context in memory_contexts
262
+ if clean_text(context.get("capsule_id"))
263
+ }
264
+ ),
265
+ "selected_claim_ids": sorted(
266
+ {
267
+ clean_text(dict(context.get("provenance") or {}).get("claim_id"))
268
+ for context in memory_contexts
269
+ if clean_text(dict(context.get("provenance") or {}).get("claim_id"))
270
+ }
271
+ ),
272
+ "selected_graph_event_ids": list(graph.get("selected_event_ids") or []),
273
+ "runtime_input_has_gold": bool(debug.get("runtime_input_has_gold")),
274
+ "graph_fingerprint": clean_text(debug.get("graph_fingerprint")),
275
+ "evidence_sha256": hashlib.sha256(
276
+ json.dumps(
277
+ dict(evidence),
278
+ ensure_ascii=False,
279
+ sort_keys=True,
280
+ separators=(",", ":"),
281
+ ).encode("utf-8")
282
+ ).hexdigest(),
283
+ }
284
+ if payload["runtime_input_has_gold"]:
285
+ raise RuntimeError(f"{scope_id}: refusing to audit retrieval with evaluation labels")
286
+ if not payload["operation_id"]:
287
+ raise RuntimeError(f"{scope_id}: layered retrieval audit lacks operation_id")
288
+ store = SQLiteSessionMemoryStore(str(db_path), audit_retention=audit_retention)
289
+ return dict(
290
+ store.append_audit_event(
291
+ scope_id,
292
+ "retrieval_log",
293
+ payload,
294
+ idempotency_key=payload["operation_id"],
295
+ )
296
+ )
297
+
298
+
299
+ def load_persisted_parent_chunks(db_path: Path, scope_id: str) -> list[dict[str, Any]]:
300
+ if not db_path.exists():
301
+ raise FileNotFoundError(db_path)
302
+ with sqlite3.connect(db_path) as connection:
303
+ record_rows = connection.execute(
304
+ "SELECT memory_id, turn_index, metadata_json FROM records WHERE scope_id=?",
305
+ (scope_id,),
306
+ ).fetchall()
307
+ if not record_rows:
308
+ raise RuntimeError(f"scope has no graph records: {scope_id}")
309
+ product_parents: list[dict[str, Any]] = []
310
+ product_message_ids: set[str] = set()
311
+ for memory_id, turn_index, raw_metadata in record_rows:
312
+ metadata = json.loads(raw_metadata)
313
+ if clean_text(metadata.get("content_variant")) != "source_message":
314
+ continue
315
+ message_id = clean_text(metadata.get("message_id"))
316
+ match = MESSAGE_ID_RE.fullmatch(message_id)
317
+ if match is None:
318
+ raise RuntimeError(f"{scope_id}: malformed persisted product message id: {message_id!r}")
319
+ session_index = int(metadata.get("session_index", -1))
320
+ message_index = int(metadata.get("message_index", -1))
321
+ if session_index != int(match.group("session")) or message_index != int(match.group("parent")):
322
+ raise RuntimeError(f"{scope_id}: product message location metadata disagrees with {message_id}")
323
+ raw_content = metadata.get("raw_content")
324
+ if not isinstance(raw_content, str) or not raw_content:
325
+ raise RuntimeError(f"{scope_id}: product source {message_id} has no persisted raw content")
326
+ if message_id in product_message_ids:
327
+ raise RuntimeError(f"{scope_id}: duplicate immutable product source message: {message_id}")
328
+ product_message_ids.add(message_id)
329
+ role = clean_text(metadata.get("speaker") or dict(metadata.get("sidecar_hint_metadata") or {}).get("role"))
330
+ session_id = clean_text(metadata.get("session_id") or dict(metadata.get("sidecar_hint_metadata") or {}).get("session_id"))
331
+ date = clean_text(metadata.get("historical_date") or dict(metadata.get("sidecar_hint_metadata") or {}).get("historical_date"))
332
+ timestamp = clean_text(metadata.get("timestamp"))
333
+ if role not in {"user", "assistant", "system", "tool"} or not session_id or not date or not timestamp:
334
+ raise RuntimeError(f"{scope_id}: incomplete product source metadata for {message_id}")
335
+ product_parents.append(
336
+ {
337
+ "chunk_id": message_id,
338
+ "parent_kind": "message",
339
+ "session_index": session_index,
340
+ "parent_chunk_index": message_index,
341
+ "message_index": message_index,
342
+ "session_id": session_id,
343
+ "date": date,
344
+ "timestamp": timestamp,
345
+ "role": role,
346
+ "text": raw_content,
347
+ "turn_index": int(turn_index),
348
+ "source_record_id": clean_text(memory_id),
349
+ }
350
+ )
351
+
352
+ audit_rows = connection.execute(
353
+ "SELECT event_index, payload_json FROM audit_turn_log WHERE scope_id=? ORDER BY event_index",
354
+ (scope_id,),
355
+ ).fetchall()
356
+ if product_parents:
357
+ audit_message_ids = {
358
+ clean_text(dict(json.loads(raw_payload).get("metadata") or {}).get("message_id"))
359
+ for _, raw_payload in audit_rows
360
+ if clean_text(dict(json.loads(raw_payload).get("metadata") or {}).get("message_id"))
361
+ }
362
+ if audit_message_ids != product_message_ids:
363
+ raise RuntimeError(
364
+ f"{scope_id}: immutable source/audit message sets differ: "
365
+ f"source={len(product_message_ids)} audit={len(audit_message_ids)}"
366
+ )
367
+ product_parents.sort(key=lambda row: (row["session_index"], row["message_index"]))
368
+ return product_parents
369
+
370
+ chunk_by_record: dict[str, str] = {}
371
+ chunks_by_turn: dict[int, set[str]] = {}
372
+ for memory_id, turn_index, raw_metadata in record_rows:
373
+ metadata = json.loads(raw_metadata)
374
+ sidecar = dict(metadata.get("sidecar_hint_metadata") or {})
375
+ chunk_id = clean_text(sidecar.get("chunk_id"))
376
+ if not chunk_id:
377
+ dia_id = clean_text(metadata.get("dia_id"))
378
+ chunk_id = dia_id.rsplit(":", 1)[-1] if dia_id else ""
379
+ if not LEGACY_CHUNK_ID_RE.match(chunk_id):
380
+ continue
381
+ chunk_by_record[clean_text(memory_id)] = chunk_id
382
+ chunks_by_turn.setdefault(int(turn_index), set()).add(chunk_id)
383
+
384
+ if not audit_rows:
385
+ raise RuntimeError(f"scope has no persisted write audit: {scope_id}")
386
+
387
+ parents: list[dict[str, Any]] = []
388
+ seen_locations: set[tuple[int, int]] = set()
389
+ for event_index, raw_payload in audit_rows:
390
+ payload = json.loads(raw_payload)
391
+ if clean_text(payload.get("kind")) != "memory_write":
392
+ continue
393
+ turn_index = int(payload.get("turn_index", 0) or 0)
394
+ record_ids = [clean_text(value) for value in list(payload.get("record_ids") or [])]
395
+ candidate_chunks = {chunk_by_record[value] for value in record_ids if value in chunk_by_record}
396
+ candidate_chunks.update(chunks_by_turn.get(turn_index, set()))
397
+ if len(candidate_chunks) != 1:
398
+ raise RuntimeError(
399
+ f"{scope_id}: persisted turn {turn_index} maps to {len(candidate_chunks)} chunk ids: {sorted(candidate_chunks)}"
400
+ )
401
+ chunk_id = next(iter(candidate_chunks))
402
+ match = LEGACY_CHUNK_ID_RE.match(chunk_id)
403
+ assert match is not None
404
+ session_index = int(match.group("session"))
405
+ parent_index = int(match.group("parent"))
406
+ location = (session_index, parent_index)
407
+ if location in seen_locations:
408
+ raise RuntimeError(f"{scope_id}: duplicate persisted parent location: {location}")
409
+ seen_locations.add(location)
410
+ persisted_text = str(payload.get("text", ""))
411
+ parent_text = INGEST_PREFIX_RE.sub("", persisted_text, count=1)
412
+ if parent_text == persisted_text or not parent_text:
413
+ raise RuntimeError(f"{scope_id}: cannot recover raw parent text from audit turn {turn_index}")
414
+ header = SESSION_HEADER_RE.match(parent_text)
415
+ if header is None:
416
+ raise RuntimeError(f"{scope_id}: malformed persisted LongMemEval parent header at {chunk_id}")
417
+ parents.append(
418
+ {
419
+ "chunk_id": chunk_id,
420
+ "parent_kind": "legacy_chunk",
421
+ "session_index": session_index,
422
+ "parent_chunk_index": parent_index,
423
+ "session_id": clean_text(header.group("session_id")),
424
+ "date": clean_text(header.group("date")),
425
+ "text": parent_text,
426
+ "turn_index": turn_index,
427
+ "audit_event_index": int(event_index),
428
+ }
429
+ )
430
+ parents.sort(key=lambda row: (row["session_index"], row["parent_chunk_index"]))
431
+ if len(parents) != len(audit_rows):
432
+ raise RuntimeError(
433
+ f"{scope_id}: not every persisted audit turn became a parent chunk: parents={len(parents)} audit={len(audit_rows)}"
434
+ )
435
+ return parents
436
+
437
+
438
+ def parent_subchunks(
439
+ parents: Sequence[Mapping[str, Any]],
440
+ *,
441
+ scope_id: str,
442
+ subchunk_chars: int,
443
+ subchunk_overlap: int,
444
+ ) -> list[dict[str, Any]]:
445
+ candidates: list[dict[str, Any]] = []
446
+ for parent in parents:
447
+ if clean_text(parent.get("parent_kind")) == "message":
448
+ prefix = (
449
+ f"TMCRA conversation_id={parent['session_id']} timestamp={parent['timestamp']} "
450
+ f"message={int(parent['message_index']):03d} role={parent['role']}"
451
+ )
452
+ else:
453
+ prefix = (
454
+ f"LongMemEval session_id={parent['session_id']} date={parent['date']} "
455
+ f"parent_chunk={int(parent['parent_chunk_index']):02d}"
456
+ )
457
+ payload_chars = int(subchunk_chars) - len(prefix) - 1
458
+ spans = covered_windows(str(parent["text"]), payload_chars, int(subchunk_overlap))
459
+ for subchunk_index, (char_start, char_end) in enumerate(spans, start=1):
460
+ text = f"{prefix}\n{str(parent['text'])[char_start:char_end]}"
461
+ if len(text) > subchunk_chars:
462
+ raise RuntimeError("generated online subchunk exceeds the strict character limit")
463
+ candidates.append(
464
+ {
465
+ "candidate_id": f"chunk::{scope_id}:{parent['chunk_id']}_p{subchunk_index:02d}",
466
+ "text": text,
467
+ "session_id": parent["session_id"],
468
+ "session_index": int(parent["session_index"]),
469
+ "parent_chunk_index": int(parent["parent_chunk_index"]),
470
+ "parent_kind": clean_text(parent.get("parent_kind")) or "legacy_chunk",
471
+ "role": clean_text(parent.get("role")),
472
+ "message_role": clean_text(parent.get("role")),
473
+ "historical_date": clean_text(parent.get("date")),
474
+ "timestamp": clean_text(parent.get("timestamp")),
475
+ "subchunk_index": subchunk_index,
476
+ "source_char_start": char_start,
477
+ "source_char_end": char_end,
478
+ "source_record_id": clean_text(parent.get("source_record_id")),
479
+ }
480
+ )
481
+ if not candidates:
482
+ raise RuntimeError(f"no online candidates from persisted scope: {scope_id}")
483
+ return candidates
484
+
485
+
486
+ def _metadata_json(value: Any, *, label: str) -> dict[str, Any]:
487
+ try:
488
+ parsed = json.loads(value)
489
+ except (TypeError, json.JSONDecodeError) as exc:
490
+ raise RuntimeError(f"invalid {label} metadata JSON") from exc
491
+ if not isinstance(parsed, dict):
492
+ raise RuntimeError(f"invalid {label} metadata object")
493
+ return parsed
494
+
495
+
496
+ def _canonical_slot(value: Any, *, label: str) -> str:
497
+ slot = clean_text(value)
498
+ if not slot:
499
+ raise RuntimeError(f"{label} lacks canonical_slot")
500
+ return slot
501
+
502
+
503
+ def _normalize_source_parents(value: Any, *, valid_locations: set[tuple[int, int]], label: str) -> list[dict[str, Any]]:
504
+ if not isinstance(value, Sequence) or isinstance(value, (str, bytes)) or not value:
505
+ raise RuntimeError(f"{label} lacks structured source_parents")
506
+ output: list[dict[str, Any]] = []
507
+ seen: set[tuple[int, int, int, int, str]] = set()
508
+ for raw in value:
509
+ if not isinstance(raw, Mapping):
510
+ raise RuntimeError(f"{label} source_parent is not an object")
511
+ try:
512
+ session = int(raw["session_index"])
513
+ parent = int(raw.get("parent_chunk_index", raw.get("message_index")))
514
+ except (KeyError, TypeError, ValueError) as exc:
515
+ raise RuntimeError(f"{label} source_parent lacks integer session/parent coordinates") from exc
516
+ message_index: int | None = None
517
+ if "message_index" in raw:
518
+ try:
519
+ message_index = int(raw["message_index"])
520
+ except (TypeError, ValueError) as exc:
521
+ raise RuntimeError(
522
+ f"{label} source_parent has a non-integer message_index"
523
+ ) from exc
524
+ if message_index != parent:
525
+ raise RuntimeError(
526
+ f"{label} source_parent message_index differs from parent_chunk_index"
527
+ )
528
+ location = (session, parent)
529
+ if location not in valid_locations:
530
+ raise RuntimeError(f"{label} source_parent cannot map to a persisted parent: {location}")
531
+ try:
532
+ char_start = int(raw["evidence_char_start"])
533
+ char_end = int(raw["evidence_char_end"])
534
+ except (KeyError, TypeError, ValueError) as exc:
535
+ raise RuntimeError(
536
+ f"{label} source_parent lacks integer evidence character span"
537
+ ) from exc
538
+ source_record_id = clean_text(raw.get("source_record_id"))
539
+ if char_start < 0 or char_end <= char_start or not source_record_id:
540
+ raise RuntimeError(f"{label} source_parent has invalid evidence provenance")
541
+ identity = (session, parent, char_start, char_end, source_record_id)
542
+ if identity not in seen:
543
+ seen.add(identity)
544
+ item = dict(raw)
545
+ item["session_index"] = session
546
+ item["parent_chunk_index"] = parent
547
+ item["evidence_char_start"] = char_start
548
+ item["evidence_char_end"] = char_end
549
+ item["source_record_id"] = source_record_id
550
+ if message_index is not None:
551
+ item["message_index"] = message_index
552
+ output.append(item)
553
+ if not output:
554
+ raise RuntimeError(f"{label} has no mapped source_parents")
555
+ return output
556
+
557
+
558
+ def load_layered_inventory(
559
+ db_path: Path, scope_id: str, parents: Sequence[Mapping[str, Any]]
560
+ ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
561
+ """Load current slow heads and map immutable fast semantic leaves to parents."""
562
+ valid_locations = {(int(item["session_index"]), int(item["parent_chunk_index"])) for item in parents}
563
+ slow_rows: dict[str, list[tuple[str, str, dict[str, Any]]]] = {}
564
+ semantic: list[dict[str, Any]] = []
565
+ with sqlite3.connect(db_path) as con:
566
+ rows = con.execute(
567
+ "SELECT memory_id, value, state, metadata_json FROM records WHERE scope_id=?", (scope_id,)
568
+ ).fetchall()
569
+ for memory_id, value, state, raw_metadata in rows:
570
+ metadata = _metadata_json(raw_metadata, label=str(memory_id))
571
+ layer = clean_text(metadata.get("memory_layer"))
572
+ variant = clean_text(metadata.get("content_variant"))
573
+ if layer == "slow" or variant == "slow_memory_capsule":
574
+ if layer != "slow" or variant != "slow_memory_capsule":
575
+ raise RuntimeError(f"{memory_id}: malformed slow capsule layer/variant")
576
+ capsule_id = clean_text(metadata.get("capsule_id"))
577
+ revision = metadata.get("revision")
578
+ if not capsule_id or not isinstance(revision, int) or revision < 1:
579
+ raise RuntimeError(f"{memory_id}: invalid capsule_id or revision")
580
+ slow_rows.setdefault(capsule_id, []).append((str(memory_id), str(state), {**metadata, "value": str(value)}))
581
+ continue
582
+ if layer != "fast" or variant != "product_semantic_memory":
583
+ continue
584
+ if clean_text(state) not in CURRENT_FAST_STATES:
585
+ continue
586
+ if (
587
+ clean_text(metadata.get("node_kind")) != "atomic_user_assertion"
588
+ or metadata.get("atomic_evidence_leaf") is not True
589
+ or clean_text(metadata.get("authority")) != "user_assertion"
590
+ ):
591
+ raise RuntimeError(f"{memory_id}: malformed fast semantic evidence leaf")
592
+ slot = clean_text(metadata.get("canonical_slot") or metadata.get("canonical_slot_key"))
593
+ if not slot:
594
+ # Non-slot fast leaves still remain retrievable through their source parent.
595
+ continue
596
+ try:
597
+ location = (int(metadata["session_index"]), int(metadata.get("parent_chunk_index", metadata.get("message_index"))))
598
+ except (KeyError, TypeError, ValueError) as exc:
599
+ raise RuntimeError(f"{memory_id}: fast semantic record lacks source parent location") from exc
600
+ if location not in valid_locations:
601
+ raise RuntimeError(f"{memory_id}: fast semantic record cannot map to persisted parent {location}")
602
+ source_record_id = clean_text(metadata.get("source_record_id"))
603
+ if not source_record_id:
604
+ raise RuntimeError(f"{memory_id}: fast semantic leaf lacks metadata.source_record_id")
605
+ try:
606
+ evidence_char_start = int(metadata["evidence_char_start"])
607
+ evidence_char_end = int(metadata["evidence_char_end"])
608
+ except (KeyError, TypeError, ValueError) as exc:
609
+ raise RuntimeError(
610
+ f"{memory_id}: fast semantic leaf lacks evidence character span"
611
+ ) from exc
612
+ if evidence_char_start < 0 or evidence_char_end <= evidence_char_start:
613
+ raise RuntimeError(f"{memory_id}: fast semantic leaf has invalid evidence character span")
614
+ semantic.append({
615
+ "memory_id": str(memory_id), "record_state": clean_text(state), "canonical_slot": slot, "source_parent": {
616
+ "session_index": location[0], "parent_chunk_index": location[1],
617
+ "source_record_id": source_record_id,
618
+ "evidence_char_start": evidence_char_start,
619
+ "evidence_char_end": evidence_char_end,
620
+ }, "provenance": {"memory_layer": "fast", "content_variant": variant, "source_record_id": source_record_id, "semantic_memory_id": str(memory_id)},
621
+ })
622
+ capsules: list[dict[str, Any]] = []
623
+ for capsule_id, revisions in slow_rows.items():
624
+ max_revision = max(int(meta["revision"]) for _, _, meta in revisions)
625
+ current = [(memory_id, state, meta) for memory_id, state, meta in revisions if int(meta["revision"]) == max_revision]
626
+ if len(current) != 1:
627
+ raise RuntimeError(f"{scope_id}: capsule {capsule_id} lacks a unique latest active/challenged revision")
628
+ memory_id, state, meta = current[0]
629
+ if state != "active" or clean_text(meta.get("status")) not in {"active", "challenged"}:
630
+ continue
631
+ claims = meta.get("claims")
632
+ if not isinstance(claims, Sequence) or isinstance(claims, (str, bytes)) or not claims:
633
+ raise RuntimeError(f"{memory_id}: capsule lacks claims")
634
+ capsule_parents = meta.get("source_parents")
635
+ for claim_index, claim in enumerate(claims):
636
+ if not isinstance(claim, Mapping):
637
+ raise RuntimeError(f"{memory_id}: claim {claim_index} is not an object")
638
+ slot = _canonical_slot(claim.get("canonical_slot"), label=f"{memory_id}: claim {claim_index}")
639
+ source_parents = _normalize_source_parents(
640
+ claim.get("source_parents", capsule_parents), valid_locations=valid_locations,
641
+ label=f"{memory_id}: claim {claim_index}",
642
+ )
643
+ claim_text = clean_text(claim.get("text"))
644
+ claim_id = clean_text(claim.get("claim_id"))
645
+ if not claim_text or not claim_id:
646
+ raise RuntimeError(f"{memory_id}: claim {claim_index} lacks claim_id or text")
647
+ capsules.append({
648
+ "candidate_id": f"capsule::{capsule_id}:r{meta['revision']}:c{claim_index}",
649
+ "memory_id": memory_id, "capsule_id": capsule_id, "revision": int(meta["revision"]),
650
+ "status": clean_text(meta.get("status")), "canonical_slot": slot, "claims": [dict(claim)],
651
+ "source_parents": source_parents, "text": claim_text,
652
+ "provenance": {"memory_layer": "slow", "content_variant": "slow_memory_capsule", "capsule_id": capsule_id, "revision": int(meta["revision"]), "claim_id": claim_id, "canonical_slot": slot, "patch_id": clean_text(meta.get("patch_id")), "source_parents": source_parents},
653
+ })
654
+ return capsules, semantic
655
+
656
+
657
+ def command_build_index(args: argparse.Namespace) -> None:
658
+ rows = read_jsonl(Path(args.scope_manifest))
659
+ device = torch.device(args.device)
660
+ if device.type == "cuda" and not torch.cuda.is_available():
661
+ raise RuntimeError("CUDA is unavailable")
662
+ vectorizer: BgeM3DenseVectorizer | None = None
663
+ started = time.time()
664
+ report_rows: list[dict[str, Any]] = []
665
+ reused_index_count = 0
666
+ for row_index, row in enumerate(rows, start=1):
667
+ db_path = Path(row["db_path"]).resolve()
668
+ scope_id = clean_text(row.get("scope_id"))
669
+ index_path = Path(row["index_path"]).resolve()
670
+ if index_path.exists():
671
+ payload = torch.load(index_path, map_location="cpu", weights_only=True)
672
+ if not isinstance(payload, Mapping):
673
+ raise RuntimeError(f"existing online index is not an object: {index_path}")
674
+ expected = {
675
+ "schema_version": ONLINE_INDEX_SCHEMA_VERSION,
676
+ "fast_semantic_state_policy": FAST_SEMANTIC_STATE_POLICY,
677
+ "scope_id": scope_id,
678
+ "db_path": str(db_path),
679
+ "subchunk_chars": int(args.subchunk_chars),
680
+ "subchunk_overlap": int(args.subchunk_overlap),
681
+ "embedding_model": str(Path(args.embedding_model).resolve()),
682
+ "embedding_max_length": int(args.embedding_max_length),
683
+ "strict_no_truncation": True,
684
+ }
685
+ mismatches = {
686
+ key: {"expected": value, "actual": payload.get(key)}
687
+ for key, value in expected.items()
688
+ if payload.get(key) != value
689
+ }
690
+ current_fingerprint = scope_fingerprint(db_path, scope_id)
691
+ if clean_text(payload.get("graph_fingerprint")) != current_fingerprint:
692
+ mismatches["graph_fingerprint"] = {
693
+ "expected": current_fingerprint,
694
+ "actual": payload.get("graph_fingerprint"),
695
+ }
696
+ if mismatches:
697
+ raise RuntimeError(
698
+ f"existing online index is incompatible with current scope: "
699
+ f"{index_path} mismatches={json.dumps(mismatches, sort_keys=True)}"
700
+ )
701
+ report_row = {
702
+ "question_id": clean_text(row.get("question_id")),
703
+ "scope_id": scope_id,
704
+ "db_path": str(db_path),
705
+ "index_path": str(index_path),
706
+ "parent_count": int(payload["parent_count"]),
707
+ "candidate_count": int(payload["candidate_count"]),
708
+ "slow_capsule_count": int(payload["slow_capsule_count"]),
709
+ "fast_semantic_record_count": int(payload["fast_semantic_record_count"]),
710
+ "fast_semantic_state_policy": payload["fast_semantic_state_policy"],
711
+ "graph_counts": dict(payload["graph_counts_at_index"]),
712
+ "graph_fingerprint": current_fingerprint,
713
+ "reused_existing_index": True,
714
+ }
715
+ reused_index_count += 1
716
+ report_rows.append(report_row)
717
+ print(
718
+ json.dumps(
719
+ {"status": "index_reused", "row": row_index, "total": len(rows), **report_row}
720
+ ),
721
+ flush=True,
722
+ )
723
+ continue
724
+ if vectorizer is None:
725
+ vectorizer = BgeM3DenseVectorizer(
726
+ dim=args.text_dim,
727
+ model_path=args.embedding_model,
728
+ device=str(device),
729
+ max_length=args.embedding_max_length,
730
+ strict_max_length=True,
731
+ )
732
+ graph_fingerprint = scope_fingerprint(db_path, scope_id)
733
+ parents = load_persisted_parent_chunks(db_path, scope_id)
734
+ candidates = parent_subchunks(
735
+ parents,
736
+ scope_id=scope_id,
737
+ subchunk_chars=args.subchunk_chars,
738
+ subchunk_overlap=args.subchunk_overlap,
739
+ )
740
+ slow_capsules, fast_semantic_records = load_layered_inventory(db_path, scope_id, parents)
741
+ fast_vectors = vectorizer.encode_batch([candidate["text"] for candidate in candidates], batch_size=args.batch_size)
742
+ slow_vectors = (
743
+ vectorizer.encode_batch([capsule["text"] for capsule in slow_capsules], batch_size=args.batch_size)
744
+ if slow_capsules else torch.empty((0, args.text_dim), dtype=torch.float32)
745
+ )
746
+ fast_vectors = fast_vectors.to(torch.float16).contiguous()
747
+ slow_vectors = slow_vectors.to(torch.float16).contiguous()
748
+ counts = scope_counts(db_path, scope_id)
749
+ if scope_fingerprint(db_path, scope_id) != graph_fingerprint:
750
+ raise RuntimeError(
751
+ f"{scope_id}: graph changed while the online index was being built"
752
+ )
753
+ payload = {
754
+ "schema_version": ONLINE_INDEX_SCHEMA_VERSION,
755
+ "fast_semantic_state_policy": FAST_SEMANTIC_STATE_POLICY,
756
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
757
+ "scope_id": scope_id,
758
+ "db_path": str(db_path),
759
+ "graph_counts_at_index": counts,
760
+ "graph_fingerprint": graph_fingerprint,
761
+ "parent_count": len(parents),
762
+ "candidate_count": len(candidates),
763
+ "slow_capsule_count": len(slow_capsules),
764
+ "fast_semantic_record_count": len(fast_semantic_records),
765
+ "subchunk_chars": args.subchunk_chars,
766
+ "subchunk_overlap": args.subchunk_overlap,
767
+ "embedding_model": str(Path(args.embedding_model).resolve()),
768
+ "embedding_max_length": args.embedding_max_length,
769
+ "strict_no_truncation": True,
770
+ "fast_candidates": candidates,
771
+ "fast_vectors": fast_vectors,
772
+ "slow_capsules": slow_capsules,
773
+ "slow_vectors": slow_vectors,
774
+ "fast_semantic_records": fast_semantic_records,
775
+ }
776
+ atomic_torch_save(payload, index_path)
777
+ report_row = {
778
+ "question_id": clean_text(row.get("question_id")),
779
+ "scope_id": scope_id,
780
+ "db_path": str(db_path),
781
+ "index_path": str(index_path),
782
+ "parent_count": len(parents),
783
+ "candidate_count": len(candidates),
784
+ "slow_capsule_count": len(slow_capsules),
785
+ "fast_semantic_record_count": len(fast_semantic_records),
786
+ "fast_semantic_state_policy": FAST_SEMANTIC_STATE_POLICY,
787
+ "graph_counts": counts,
788
+ "graph_fingerprint": graph_fingerprint,
789
+ "reused_existing_index": False,
790
+ }
791
+ report_rows.append(report_row)
792
+ print(json.dumps({"status": "indexed", "row": row_index, "total": len(rows), **report_row}), flush=True)
793
+ report = {
794
+ "status": "complete",
795
+ "schema_version": "tmcra.v3.online-index-report.2",
796
+ "row_count": len(report_rows),
797
+ "parent_count": sum(row["parent_count"] for row in report_rows),
798
+ "candidate_count": sum(row["candidate_count"] for row in report_rows),
799
+ "slow_capsule_count": sum(row["slow_capsule_count"] for row in report_rows),
800
+ "reused_index_count": reused_index_count,
801
+ "elapsed_sec": round(time.time() - started, 3),
802
+ "rows": report_rows,
803
+ }
804
+ out_report = Path(args.out_report)
805
+ out_report.parent.mkdir(parents=True, exist_ok=True)
806
+ _atomic_write_text(out_report, json.dumps(report, indent=2, sort_keys=True) + "\n")
807
+
808
+
809
+ def ordered_graph_parents(
810
+ event_ids: Iterable[Any],
811
+ *,
812
+ valid_locations: set[tuple[int, int]],
813
+ strict_prefix: bool = False,
814
+ ) -> tuple[list[tuple[int, int]], list[str]]:
815
+ output: list[tuple[int, int]] = []
816
+ seen: set[tuple[int, int]] = set()
817
+ unmapped: list[str] = []
818
+ for raw in event_ids:
819
+ event_id = clean_text(raw)
820
+ match = EVENT_PARENT_RE.search(event_id)
821
+ if match is None:
822
+ if strict_prefix and event_id.startswith(("event::longmemeval:", "event::tmcra:")):
823
+ unmapped.append(event_id)
824
+ continue
825
+ location = (int(match.group("session")), int(match.group("parent")))
826
+ if location not in valid_locations:
827
+ unmapped.append(event_id)
828
+ continue
829
+ if location not in seen:
830
+ seen.add(location)
831
+ output.append(location)
832
+ return output, unmapped
833
+
834
+
835
+ def expand_parent_locations(
836
+ locations: Iterable[tuple[int, int]],
837
+ parent_candidates: Mapping[tuple[int, int], Sequence[int]],
838
+ ) -> list[int]:
839
+ output: list[int] = []
840
+ seen: set[int] = set()
841
+ for location in locations:
842
+ for candidate_index in parent_candidates[location]:
843
+ if candidate_index not in seen:
844
+ seen.add(candidate_index)
845
+ output.append(candidate_index)
846
+ return output
847
+
848
+
849
+ class OnlineModels:
850
+ def __init__(self, args: argparse.Namespace):
851
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer
852
+
853
+ self.device = torch.device(args.device)
854
+ if self.device.type == "cuda" and not torch.cuda.is_available():
855
+ raise RuntimeError("CUDA is unavailable")
856
+ self.dense = BgeM3DenseVectorizer(
857
+ dim=args.text_dim,
858
+ model_path=args.embedding_model,
859
+ device=str(self.device),
860
+ max_length=args.embedding_max_length,
861
+ strict_max_length=True,
862
+ )
863
+ cross_path = Path(args.cross_model).resolve()
864
+ self.cross_manifest = load_optional_hf_model_manifest(cross_path)
865
+ self.cross_tokenizer = AutoTokenizer.from_pretrained(
866
+ str(cross_path),
867
+ local_files_only=True,
868
+ trust_remote_code=False,
869
+ )
870
+ self.cross_model = AutoModelForSequenceClassification.from_pretrained(
871
+ str(cross_path),
872
+ local_files_only=True,
873
+ trust_remote_code=False,
874
+ torch_dtype=torch.float16 if self.device.type == "cuda" else torch.float32,
875
+ ).to(self.device)
876
+ self.cross_model.eval()
877
+ hidden_size = int(getattr(self.cross_model.config, "hidden_size", 0) or 0)
878
+ if hidden_size <= 0:
879
+ raise RuntimeError("cross encoder hidden size is unavailable")
880
+ checkpoint_path = Path(args.checkpoint).resolve()
881
+ checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=True)
882
+ if checkpoint.get("schema_version") != SCHEMA_VERSION:
883
+ raise RuntimeError("online checkpoint schema mismatch")
884
+ if tuple(checkpoint.get("channel_names") or ()) != CHANNEL_NAMES:
885
+ raise RuntimeError("online checkpoint channel mismatch")
886
+ config = dict(checkpoint.get("model_config") or {})
887
+ self.fusion = ChannelAwareMemoryReranker(
888
+ representation_dim=hidden_size,
889
+ channel_dim=len(CHANNEL_NAMES),
890
+ hidden_dim=int(config["hidden_dim"]),
891
+ layers=int(config["layers"]),
892
+ ).to(self.device)
893
+ self.fusion.load_state_dict(checkpoint["model_state"])
894
+ self.fusion.eval()
895
+ self.checkpoint = checkpoint
896
+ self.checkpoint_path = checkpoint_path
897
+ self.checkpoint_sha256 = sha256_file(checkpoint_path)
898
+ self.cross_max_length = int(args.cross_max_length)
899
+ self.cross_batch_size = int(args.cross_batch_size)
900
+
901
+ def encode_cross(self, query: str, texts: Sequence[str]) -> tuple[torch.Tensor, torch.Tensor]:
902
+ representations: list[torch.Tensor] = []
903
+ logits: list[torch.Tensor] = []
904
+ with torch.inference_mode():
905
+ for start in range(0, len(texts), self.cross_batch_size):
906
+ batch_texts = list(texts[start : start + self.cross_batch_size])
907
+ encoded = self.cross_tokenizer(
908
+ [query] * len(batch_texts),
909
+ batch_texts,
910
+ padding=True,
911
+ truncation=False,
912
+ return_tensors="pt",
913
+ )
914
+ token_lengths = encoded["attention_mask"].sum(dim=1)
915
+ longest = int(token_lengths.max())
916
+ if longest > self.cross_max_length:
917
+ offending = int(torch.argmax(token_lengths))
918
+ raise RuntimeError(
919
+ f"online cross pair has {longest} tokens, exceeding strict max={self.cross_max_length}; "
920
+ f"candidate_batch_index={offending}"
921
+ )
922
+ encoded = {key: value.to(self.device) for key, value in encoded.items()}
923
+ output = self.cross_model(**encoded, output_hidden_states=True, return_dict=True)
924
+ if output.hidden_states is None:
925
+ raise RuntimeError("online cross encoder returned no hidden states")
926
+ representations.append(output.hidden_states[-1][:, 0].float())
927
+ semantic = output.logits.float()
928
+ if semantic.ndim == 2 and semantic.shape[1] == 1:
929
+ semantic = semantic[:, 0]
930
+ elif semantic.ndim != 1:
931
+ raise RuntimeError(f"online cross logits have invalid shape: {tuple(semantic.shape)}")
932
+ logits.append(semantic)
933
+ return torch.cat(representations, dim=0), torch.cat(logits, dim=0)
934
+
935
+
936
+ def load_online_index(path: Path, expected_db: Path, expected_scope: str) -> tuple[list[dict[str, Any]], torch.Tensor, list[dict[str, Any]], torch.Tensor, list[dict[str, Any]], dict[str, Any]]:
937
+ payload = torch.load(path, map_location="cpu", weights_only=True)
938
+ if payload.get("schema_version") != ONLINE_INDEX_SCHEMA_VERSION:
939
+ raise RuntimeError(f"online index schema mismatch: {path}")
940
+ if payload.get("fast_semantic_state_policy") != FAST_SEMANTIC_STATE_POLICY:
941
+ raise RuntimeError(f"online index fast semantic state policy mismatch: {path}")
942
+ if clean_text(payload.get("scope_id")) != expected_scope:
943
+ raise RuntimeError(f"online index scope mismatch: {path}")
944
+ if Path(payload.get("db_path", "")).resolve() != expected_db.resolve():
945
+ raise RuntimeError(f"online index database mismatch: {path}")
946
+ candidates = list(payload.get("fast_candidates") or [])
947
+ vectors = payload.get("fast_vectors")
948
+ slow_capsules = list(payload.get("slow_capsules") or [])
949
+ slow_vectors = payload.get("slow_vectors")
950
+ semantic_records = list(payload.get("fast_semantic_records") or [])
951
+ if not candidates or vectors is None or tuple(vectors.shape) != (len(candidates), 1024):
952
+ raise RuntimeError(f"online index payload is incomplete: {path}")
953
+ if slow_vectors is None or tuple(slow_vectors.shape) != (len(slow_capsules), 1024):
954
+ raise RuntimeError(f"online slow index payload is incomplete: {path}")
955
+ if any("labels" in candidate for candidate in candidates):
956
+ raise RuntimeError("runtime index must not contain benchmark labels")
957
+ return candidates, vectors.float().contiguous(), slow_capsules, slow_vectors.float().contiguous(), semantic_records, payload
958
+
959
+
960
+ def choose_diverse(
961
+ candidates: Sequence[Mapping[str, Any]],
962
+ scores: torch.Tensor,
963
+ *,
964
+ top_k: int,
965
+ max_per_parent: int,
966
+ max_per_session: int,
967
+ ) -> list[int]:
968
+ order = torch.argsort(scores, descending=True).tolist()
969
+ selected: list[int] = []
970
+ parent_counts: Counter[tuple[int, int]] = Counter()
971
+ session_counts: Counter[str] = Counter()
972
+ for index in order:
973
+ candidate = candidates[index]
974
+ parent = (int(candidate["session_index"]), int(candidate["parent_chunk_index"]))
975
+ session_id = clean_text(candidate.get("session_id"))
976
+ if max_per_parent > 0 and parent_counts[parent] >= max_per_parent:
977
+ continue
978
+ if max_per_session > 0 and session_counts[session_id] >= max_per_session:
979
+ continue
980
+ selected.append(index)
981
+ parent_counts[parent] += 1
982
+ session_counts[session_id] += 1
983
+ if len(selected) >= top_k:
984
+ break
985
+ if len(selected) != top_k:
986
+ raise RuntimeError(f"online diversity policy produced only {len(selected)} of {top_k} required windows")
987
+ return selected
988
+
989
+
990
+ def _deprecated_legacy_retrieve_one_v1(
991
+ row: Mapping[str, Any],
992
+ *,
993
+ args: argparse.Namespace,
994
+ harness: Any,
995
+ models: OnlineModels,
996
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
997
+ """Retained only for forensic comparison; the v2 command never calls this schema .1 path."""
998
+ started = time.time()
999
+ qid = clean_text(row.get("question_id"))
1000
+ question = clean_text(row.get("question"))
1001
+ question_date = clean_text(row.get("question_date"))
1002
+ runtime_question = f"{question}\nQuestion date: {question_date}" if question_date else question
1003
+ db_path = Path(row["db_path"]).resolve()
1004
+ scope_id = clean_text(row.get("scope_id"))
1005
+ index_path = Path(row["index_path"]).resolve()
1006
+ if not qid or not question or not scope_id:
1007
+ raise RuntimeError("online retrieval manifest row lacks qid, question, or scope")
1008
+ candidates, dense_vectors, index_payload = load_online_index(index_path, db_path, scope_id)
1009
+ counts_before = scope_counts(db_path, scope_id)
1010
+ if counts_before["records"] != int(index_payload["graph_counts_at_index"]["records"]):
1011
+ raise RuntimeError(f"{qid}: graph records changed after online index creation")
1012
+
1013
+ adapter = harness.build_adapter(scope_id, db_path)
1014
+ graph_started = time.time()
1015
+ retrieval = adapter.retrieve(runtime_question, top_k=args.graph_top_k)
1016
+ graph_elapsed = time.time() - graph_started
1017
+ metadata = dict(getattr(retrieval, "metadata", {}) or {})
1018
+ if clean_text(metadata.get("retrieval_mode")) != "hybrid_node_scored":
1019
+ raise RuntimeError(f"{qid}: graph retrieval did not use hybrid_node_scored")
1020
+ if not bool(metadata.get("hybrid_enabled")):
1021
+ raise RuntimeError(f"{qid}: graph hybrid model path is not active")
1022
+ selected_event_ids = list(metadata.get("selected_event_ids") or [])
1023
+ recall_event_ids = list(metadata.get("recall_event_ids") or [])
1024
+ final_event_ids = list(metadata.get("final_hit_event_ids") or [])
1025
+ if not selected_event_ids:
1026
+ raise RuntimeError(f"{qid}: graph model selected no events")
1027
+
1028
+ parent_candidates: dict[tuple[int, int], list[int]] = {}
1029
+ for index, candidate in enumerate(candidates):
1030
+ location = (int(candidate["session_index"]), int(candidate["parent_chunk_index"]))
1031
+ parent_candidates.setdefault(location, []).append(index)
1032
+ valid_locations = set(parent_candidates)
1033
+ selected_parents, selected_unmapped = ordered_graph_parents(
1034
+ selected_event_ids, valid_locations=valid_locations, strict_prefix=True
1035
+ )
1036
+ recall_parents, recall_unmapped = ordered_graph_parents(recall_event_ids, valid_locations=valid_locations)
1037
+ final_parents, final_unmapped = ordered_graph_parents(final_event_ids, valid_locations=valid_locations)
1038
+ critical_unmapped = [*selected_unmapped, *final_unmapped]
1039
+ if critical_unmapped:
1040
+ raise RuntimeError(f"{qid}: selected/final graph events cannot map to persisted chunks: {critical_unmapped[:8]}")
1041
+ graph_parents: list[tuple[int, int]] = []
1042
+ graph_parent_seen: set[tuple[int, int]] = set()
1043
+ for location in [*selected_parents, *recall_parents]:
1044
+ if location not in graph_parent_seen:
1045
+ graph_parent_seen.add(location)
1046
+ graph_parents.append(location)
1047
+ graph_parents = graph_parents[: args.graph_k]
1048
+ graph_parent_rank = {location: rank for rank, location in enumerate(graph_parents)}
1049
+ graph_runtime_order = expand_parent_locations(graph_parents, parent_candidates)
1050
+ graph_rank = {
1051
+ index: graph_parent_rank[(int(candidates[index]["session_index"]), int(candidates[index]["parent_chunk_index"]))]
1052
+ for index in graph_runtime_order
1053
+ }
1054
+ selected_indexes = set(expand_parent_locations(selected_parents, parent_candidates))
1055
+ final_indexes = set(expand_parent_locations(final_parents, parent_candidates))
1056
+
1057
+ dense_started = time.time()
1058
+ query_vector = models.dense.encode_one(runtime_question)
1059
+ dense_scores = dense_vectors @ query_vector
1060
+ dense_order = sorted(range(len(candidates)), key=lambda index: (-float(dense_scores[index]), index))
1061
+ dense_rank = {index: rank for rank, index in enumerate(dense_order)}
1062
+ dense_elapsed = time.time() - dense_started
1063
+ runtime_indexes: list[int] = []
1064
+ runtime_seen: set[int] = set()
1065
+ for index in [*dense_order[: args.dense_k], *graph_runtime_order]:
1066
+ if index not in runtime_seen:
1067
+ runtime_seen.add(index)
1068
+ runtime_indexes.append(index)
1069
+ if not runtime_indexes:
1070
+ raise RuntimeError(f"{qid}: graph+dense union is empty")
1071
+
1072
+ session_count = max(int(candidate["session_index"]) for candidate in candidates) + 1
1073
+ runtime_candidates: list[dict[str, Any]] = []
1074
+ for index in runtime_indexes:
1075
+ source = dict(candidates[index])
1076
+ channels = {
1077
+ "dense_score": float(dense_scores[index]),
1078
+ "dense_rank_rr": rrank(dense_rank[index]),
1079
+ "graph_rank_rr": rrank(graph_rank.get(index)),
1080
+ "graph_selected": float(index in selected_indexes),
1081
+ "graph_final": float(index in final_indexes),
1082
+ "recency_norm": float(source["session_index"]) / float(max(1, session_count - 1)),
1083
+ }
1084
+ if tuple(channels) != CHANNEL_NAMES:
1085
+ raise AssertionError("online channel order changed")
1086
+ source["channels"] = channels
1087
+ runtime_candidates.append(source)
1088
+
1089
+ cross_started = time.time()
1090
+ representations, semantic_logits = models.encode_cross(
1091
+ runtime_question, [candidate["text"] for candidate in runtime_candidates]
1092
+ )
1093
+ channel_tensor = torch.tensor(
1094
+ [[candidate["channels"][name] for name in CHANNEL_NAMES] for candidate in runtime_candidates],
1095
+ dtype=torch.float32,
1096
+ device=models.device,
1097
+ )
1098
+ mask = torch.ones((1, len(runtime_candidates)), dtype=torch.bool, device=models.device)
1099
+ with torch.inference_mode():
1100
+ fusion_scores = models.fusion(
1101
+ representations.unsqueeze(0),
1102
+ semantic_logits.unsqueeze(0),
1103
+ channel_tensor.unsqueeze(0),
1104
+ mask,
1105
+ ablation="full",
1106
+ )[0].detach().cpu()
1107
+ semantic_cpu = semantic_logits.detach().cpu()
1108
+ cross_elapsed = time.time() - cross_started
1109
+ selected = choose_diverse(
1110
+ runtime_candidates,
1111
+ fusion_scores,
1112
+ top_k=args.top_k,
1113
+ max_per_parent=args.max_per_parent,
1114
+ max_per_session=args.max_per_session,
1115
+ )
1116
+ evidence_windows = []
1117
+ for rank, index in enumerate(selected, start=1):
1118
+ candidate = runtime_candidates[index]
1119
+ evidence_windows.append(
1120
+ {
1121
+ "memory_id": candidate["candidate_id"],
1122
+ "session_id": candidate["session_id"],
1123
+ "session_index": candidate["session_index"],
1124
+ "parent_chunk_index": candidate["parent_chunk_index"],
1125
+ "subchunk_index": candidate["subchunk_index"],
1126
+ "rank": rank,
1127
+ "score": round(float(fusion_scores[index]), 6),
1128
+ "semantic_logit": round(float(semantic_cpu[index]), 6),
1129
+ "channels": candidate["channels"],
1130
+ "text": candidate["text"],
1131
+ }
1132
+ )
1133
+ evidence_row = {
1134
+ "schema_version": SCHEMA_VERSION,
1135
+ "runtime_schema_version": "tmcra.v3.online-retrieval.1",
1136
+ "question_id": qid,
1137
+ "question": question,
1138
+ "question_date": question_date,
1139
+ "question_type": clean_text(row.get("question_type")),
1140
+ "selected_session_ids": [window["session_id"] for window in evidence_windows],
1141
+ "evidence_windows": evidence_windows,
1142
+ }
1143
+ debug_row = {
1144
+ "question_id": qid,
1145
+ "scope_id": scope_id,
1146
+ "db_path": str(db_path),
1147
+ "index_path": str(index_path),
1148
+ "runtime_input_has_gold": any(
1149
+ key in row for key in ("answer", "gold_answer", "answer_session_ids", "labels", "supervision")
1150
+ ),
1151
+ "graph": {
1152
+ "retrieval_mode": metadata.get("retrieval_mode"),
1153
+ "hybrid_enabled": metadata.get("hybrid_enabled"),
1154
+ "decision_score_source": metadata.get("decision_score_source"),
1155
+ "selected_event_count": len(selected_event_ids),
1156
+ "recall_event_count": len(recall_event_ids),
1157
+ "final_event_count": len(final_event_ids),
1158
+ "selected_parent_count": len(selected_parents),
1159
+ "recall_parent_count": len(recall_parents),
1160
+ "final_parent_count": len(final_parents),
1161
+ "unmapped_recall_event_ids": recall_unmapped,
1162
+ "selected_event_ids": selected_event_ids,
1163
+ "recall_event_ids": recall_event_ids,
1164
+ "final_hit_event_ids": final_event_ids,
1165
+ },
1166
+ "inventory_count": len(candidates),
1167
+ "dense_k": args.dense_k,
1168
+ "graph_k": args.graph_k,
1169
+ "union_count": len(runtime_candidates),
1170
+ "selected_count": len(evidence_windows),
1171
+ "checkpoint": str(models.checkpoint_path),
1172
+ "checkpoint_sha256": models.checkpoint_sha256,
1173
+ "cross_model_revision": models.cross_manifest.get("revision"),
1174
+ "strict_no_truncation": True,
1175
+ "restart_boundary_verified": True,
1176
+ "graph_counts_before_query": counts_before,
1177
+ "latency_sec": {
1178
+ "graph": round(graph_elapsed, 4),
1179
+ "dense": round(dense_elapsed, 4),
1180
+ "cross_and_fusion": round(cross_elapsed, 4),
1181
+ "total": round(time.time() - started, 4),
1182
+ },
1183
+ "ranked_union": [
1184
+ {
1185
+ "candidate_id": runtime_candidates[index]["candidate_id"],
1186
+ "session_id": runtime_candidates[index]["session_id"],
1187
+ "parent_chunk_index": runtime_candidates[index]["parent_chunk_index"],
1188
+ "subchunk_index": runtime_candidates[index]["subchunk_index"],
1189
+ "fusion_score": round(float(fusion_scores[index]), 6),
1190
+ "semantic_logit": round(float(semantic_cpu[index]), 6),
1191
+ "channels": runtime_candidates[index]["channels"],
1192
+ }
1193
+ for index in torch.argsort(fusion_scores, descending=True).tolist()
1194
+ ],
1195
+ }
1196
+ if debug_row["runtime_input_has_gold"]:
1197
+ raise RuntimeError(f"{qid}: runtime retrieval manifest contains forbidden evaluation labels")
1198
+ return evidence_row, debug_row
1199
+
1200
+
1201
+ def planner_from_env() -> DeepSeekFlashRecallPlanner:
1202
+ pool = [part.strip() for part in os.environ.get("TMCRA_RECALL_PLANNER_API_KEY_POOL", "").split(",") if part.strip()]
1203
+ return DeepSeekFlashRecallPlanner(
1204
+ base_url=os.environ.get("TMCRA_RECALL_PLANNER_BASE_URL", ""),
1205
+ model=os.environ.get("TMCRA_RECALL_PLANNER_MODEL", DEEPSEEK_FLASH_MODEL), api_keys=pool,
1206
+ )
1207
+
1208
+
1209
+ def _fast_candidates_with_slots(candidates: Sequence[Mapping[str, Any]], semantic_records: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
1210
+ by_parent: dict[tuple[int, int], list[Mapping[str, Any]]] = {}
1211
+ for record in semantic_records:
1212
+ parent = dict(record["source_parent"])
1213
+ by_parent.setdefault((int(parent["session_index"]), int(parent["parent_chunk_index"])), []).append(record)
1214
+ output: list[dict[str, Any]] = []
1215
+ for candidate in candidates:
1216
+ location = (int(candidate["session_index"]), int(candidate["parent_chunk_index"]))
1217
+ records = [
1218
+ record
1219
+ for record in by_parent.get(location, [])
1220
+ if int(candidate["source_char_end"])
1221
+ > int(record["source_parent"]["evidence_char_start"])
1222
+ and int(candidate["source_char_start"])
1223
+ < int(record["source_parent"]["evidence_char_end"])
1224
+ ]
1225
+ by_slot: dict[str, list[Mapping[str, Any]]] = {}
1226
+ for record in records:
1227
+ by_slot.setdefault(_canonical_slot(record.get("canonical_slot"), label="fast semantic record"), []).append(record)
1228
+ if not by_slot:
1229
+ by_slot[f"fast.parent.s{location[0]}.p{location[1]}"] = []
1230
+ for slot, slot_records in by_slot.items():
1231
+ item = dict(candidate)
1232
+ item["canonical_slot"] = slot
1233
+ item["semantic_record_ids"] = [str(record["memory_id"]) for record in slot_records]
1234
+ item["provenance"] = {
1235
+ "memory_layer": "fast",
1236
+ "content_variant": "source_subchunk",
1237
+ "source_parent": {
1238
+ "session_index": location[0],
1239
+ "parent_chunk_index": location[1],
1240
+ "source_record_ids": sorted(
1241
+ {
1242
+ clean_text(record["source_parent"].get("source_record_id"))
1243
+ for record in slot_records
1244
+ if clean_text(record["source_parent"].get("source_record_id"))
1245
+ }
1246
+ )
1247
+ or ([clean_text(candidate.get("source_record_id"))] if clean_text(candidate.get("source_record_id")) else []),
1248
+ },
1249
+ }
1250
+ output.append(item)
1251
+ return output
1252
+
1253
+
1254
+ def _unit_windows(unit: Mapping[str, Any], fast_candidates: Sequence[Mapping[str, Any]], *, qid: str) -> list[dict[str, Any]]:
1255
+ """Fully descend one composed unit, merging duplicate physical source windows."""
1256
+ by_parent: dict[tuple[int, int], list[Mapping[str, Any]]] = {}
1257
+ for candidate in fast_candidates:
1258
+ by_parent.setdefault((int(candidate["session_index"]), int(candidate["parent_chunk_index"])), []).append(candidate)
1259
+ output: dict[tuple[int, int, int], dict[str, Any]] = {}
1260
+
1261
+ def add(candidate: Mapping[str, Any], *, role: str, capsule: Mapping[str, Any] | None = None) -> None:
1262
+ key = (int(candidate["session_index"]), int(candidate["parent_chunk_index"]), int(candidate["subchunk_index"]))
1263
+ item = output.setdefault(key, {**dict(candidate), "roles": [], "capsules": []})
1264
+ if role not in item["roles"]:
1265
+ item["roles"].append(role)
1266
+ if capsule is not None and capsule not in item["capsules"]:
1267
+ item["capsules"].append(dict(capsule))
1268
+
1269
+ def descend(capsule: Mapping[str, Any], role: str) -> None:
1270
+ for parent in capsule["source_parents"]:
1271
+ location = (int(parent["session_index"]), int(parent["parent_chunk_index"]))
1272
+ char_start = int(parent["evidence_char_start"])
1273
+ char_end = int(parent["evidence_char_end"])
1274
+ matches = [
1275
+ candidate
1276
+ for candidate in by_parent.get(location, [])
1277
+ if int(candidate["source_char_end"]) > char_start
1278
+ and int(candidate["source_char_start"]) < char_end
1279
+ ]
1280
+ if not matches:
1281
+ raise RuntimeError(f"{qid}: slow source_parent is unmapped during descent: {location}")
1282
+ for candidate in matches:
1283
+ add(candidate, role=role, capsule=capsule)
1284
+
1285
+ kind = str(unit["unit_type"])
1286
+ if kind.startswith("fast_primary"):
1287
+ add(unit["fast_candidate"], role="primary")
1288
+ elif kind == "slow_primary_with_fast_override":
1289
+ descend(unit["slow_capsule"], "primary")
1290
+ for candidate in unit.get("fast_overrides", []):
1291
+ add(candidate, role="override")
1292
+ elif kind.startswith("slow_primary"):
1293
+ descend(unit["slow_capsule"], "primary")
1294
+ elif kind == "conflict_group":
1295
+ for capsule in unit["slow_capsules"]:
1296
+ descend(capsule, "slow_conflict_candidate")
1297
+ for candidate in unit["fast_candidates"]:
1298
+ add(candidate, role="fast_conflict_candidate")
1299
+ else:
1300
+ raise RuntimeError(f"{qid}: unsupported planned unit type {kind}")
1301
+ return list(output.values())
1302
+
1303
+
1304
+ def _unit_attachments(unit: Mapping[str, Any]) -> list[dict[str, Any]]:
1305
+ if unit.get("slow_context"):
1306
+ return [{"role": "context_only", "capsule_id": item["capsule_id"], "canonical_slot": item["canonical_slot"], "summary": item["text"], "source_parents": item["source_parents"], "provenance": item["provenance"]} for item in unit["slow_context"]]
1307
+ if unit.get("fast_overrides"):
1308
+ return [{"role": "override", "memory_id": item["candidate_id"], "canonical_slot": item["canonical_slot"], "text": item["text"], "provenance": item["provenance"]} for item in unit["fast_overrides"]]
1309
+ return []
1310
+
1311
+
1312
+ def pack_recall_units(units: Sequence[Mapping[str, Any]], fast_candidates: Sequence[Mapping[str, Any]], *, top_k: int, qid: str) -> list[tuple[Mapping[str, Any], list[dict[str, Any]]]]:
1313
+ if top_k <= 0:
1314
+ raise RuntimeError("top_k must be positive")
1315
+ packed: list[tuple[Mapping[str, Any], list[dict[str, Any]]]] = []
1316
+ used: set[tuple[int, int, int]] = set()
1317
+ for unit in units:
1318
+ windows = _unit_windows(unit, fast_candidates, qid=qid)
1319
+ unique = [item for item in windows if (int(item["session_index"]), int(item["parent_chunk_index"]), int(item["subchunk_index"])) not in used]
1320
+ if not unique:
1321
+ if windows:
1322
+ packed.append((unit, windows))
1323
+ continue
1324
+ if len(unique) > top_k and not packed:
1325
+ raise RuntimeError(f"{qid}: first atomic recall unit requires {len(unique)} windows, exceeding strict packing budget {top_k}")
1326
+ if len(used) + len(unique) > top_k:
1327
+ break
1328
+ packed.append((unit, windows))
1329
+ used.update((int(item["session_index"]), int(item["parent_chunk_index"]), int(item["subchunk_index"])) for item in unique)
1330
+ if not packed:
1331
+ raise RuntimeError(f"{qid}: recall plan produced no packable evidence units")
1332
+ return packed
1333
+
1334
+
1335
+ def retrieve_one(
1336
+ row: Mapping[str, Any], *, args: argparse.Namespace, harness: Any, models: OnlineModels,
1337
+ planner: DeepSeekFlashRecallPlanner,
1338
+ graph_adapter_cache: dict[tuple[str, str], Any] | None = None,
1339
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
1340
+ started = time.time()
1341
+ qid, question, question_date = clean_text(row.get("question_id")), clean_text(row.get("question")), clean_text(row.get("question_date"))
1342
+ if not qid or not question:
1343
+ raise RuntimeError("online retrieval manifest row lacks qid or question")
1344
+ if int(args.slow_dense_k) <= 0:
1345
+ raise RuntimeError("slow_dense_k must be positive")
1346
+ db_path, scope_id, index_path = Path(row["db_path"]).resolve(), clean_text(row.get("scope_id")), Path(row["index_path"]).resolve()
1347
+ fast, fast_vectors, slow, slow_vectors, semantic_records, payload = load_online_index(index_path, db_path, scope_id)
1348
+ counts_before = scope_counts(db_path, scope_id)
1349
+ if counts_before["records"] != int(payload["graph_counts_at_index"]["records"]):
1350
+ raise RuntimeError(f"{qid}: graph records changed after online index creation")
1351
+ graph_fingerprint = scope_fingerprint(db_path, scope_id)
1352
+ if graph_fingerprint != clean_text(payload.get("graph_fingerprint")):
1353
+ raise RuntimeError(f"{qid}: graph fingerprint changed after online index creation")
1354
+ recent_dialogue = load_recent_dialogue_context(
1355
+ db_path, scope_id, current_query=question, limit=8
1356
+ )
1357
+ raw_plan, planner_metadata = planner.plan(
1358
+ query=question, question_date=question_date or "unknown",
1359
+ recent_dialogue=recent_dialogue,
1360
+ available_layers={"fast": {"available": bool(fast), "candidate_count": len(fast)}, "slow": {"available": bool(slow), "capsule_count": len(slow)}},
1361
+ )
1362
+ plan = dict(raw_plan)
1363
+ planner_decision_reason = clean_text(plan.get("decision_reason"))
1364
+ if not planner_decision_reason:
1365
+ raise RuntimeError(f"{qid}: Flash recall plan lacks a decision reason")
1366
+ resolved_query = clean_text(plan.get("resolved_query"))
1367
+ if not resolved_query:
1368
+ raise RuntimeError(f"{qid}: Flash recall plan lacks a resolved query")
1369
+ runtime_question = (
1370
+ f"{resolved_query}\nQuestion date: {question_date}"
1371
+ if question_date
1372
+ else resolved_query
1373
+ )
1374
+ authoritative_plan = dict(plan)
1375
+ authoritative_plan.pop("decision_reason", None)
1376
+ if not slow and plan["mode"] != "FAST_ONLY":
1377
+ raise RuntimeError(f"{qid}: Flash must select FAST_ONLY when no slow capsules are available")
1378
+
1379
+ graph_metadata: dict[str, Any] = {"skipped": plan["mode"] == "SLOW_ONLY"}
1380
+ graph_elapsed = 0.0
1381
+ graph_runtime_order: list[int] = []
1382
+ selected_indexes: set[int] = set()
1383
+ final_indexes: set[int] = set()
1384
+ graph_rank: dict[int, int] = {}
1385
+ parent_candidates: dict[tuple[int, int], list[int]] = {}
1386
+ graph_adapter_reused = False
1387
+ for index, candidate in enumerate(fast):
1388
+ parent_candidates.setdefault((int(candidate["session_index"]), int(candidate["parent_chunk_index"])), []).append(index)
1389
+ if plan["mode"] != "SLOW_ONLY":
1390
+ adapter_key = (scope_id, str(db_path))
1391
+ adapter = graph_adapter_cache.get(adapter_key) if graph_adapter_cache is not None else None
1392
+ if adapter is None:
1393
+ adapter = harness.build_adapter(scope_id, db_path)
1394
+ if graph_adapter_cache is not None:
1395
+ graph_adapter_cache[adapter_key] = adapter
1396
+ else:
1397
+ graph_adapter_reused = True
1398
+ graph_started = time.time()
1399
+ retrieval = adapter.retrieve(runtime_question, top_k=args.graph_top_k)
1400
+ graph_elapsed = time.time() - graph_started
1401
+ metadata = dict(getattr(retrieval, "metadata", {}) or {})
1402
+ if clean_text(metadata.get("retrieval_mode")) != "hybrid_node_scored" or not bool(metadata.get("hybrid_enabled")):
1403
+ raise RuntimeError(f"{qid}: graph hybrid_node_scored path is not active")
1404
+ selected_events, recall_events, final_events = list(metadata.get("selected_event_ids") or []), list(metadata.get("recall_event_ids") or []), list(metadata.get("final_hit_event_ids") or [])
1405
+ slow_graph_events = [str(event_id) for event_id in [*selected_events, *recall_events, *final_events] if str(event_id).startswith("slow.")]
1406
+ if slow_graph_events:
1407
+ raise RuntimeError(f"{qid}: fast graph execution crossed the slow-layer boundary: " + ",".join(dict.fromkeys(slow_graph_events)))
1408
+ if not selected_events:
1409
+ raise RuntimeError(f"{qid}: graph model selected no events")
1410
+ valid_locations = set(parent_candidates)
1411
+ selected_parents, selected_unmapped = ordered_graph_parents(selected_events, valid_locations=valid_locations, strict_prefix=True)
1412
+ recall_parents, recall_unmapped = ordered_graph_parents(recall_events, valid_locations=valid_locations)
1413
+ final_parents, final_unmapped = ordered_graph_parents(final_events, valid_locations=valid_locations)
1414
+ if [*selected_unmapped, *final_unmapped]:
1415
+ raise RuntimeError(f"{qid}: graph events cannot map to persisted chunks")
1416
+ graph_parents = list(dict.fromkeys([*selected_parents, *recall_parents]))[:args.graph_k]
1417
+ graph_runtime_order = expand_parent_locations(graph_parents, parent_candidates)
1418
+ graph_rank = {index: rank for rank, location in enumerate(graph_parents) for index in parent_candidates[location]}
1419
+ selected_indexes, final_indexes = set(expand_parent_locations(selected_parents, parent_candidates)), set(expand_parent_locations(final_parents, parent_candidates))
1420
+ graph_metadata = {"skipped": False, "adapter_reused": graph_adapter_reused, "selected_event_ids": selected_events, "recall_event_ids": recall_events, "final_hit_event_ids": final_events, "unmapped_recall_event_ids": recall_unmapped, "retrieval_mode": metadata.get("retrieval_mode"), "runtime_graph_cache_hit": bool(metadata.get("runtime_graph_cache_hit", False)), "hybrid_candidate_union_rescored": bool(metadata.get("hybrid_candidate_union_rescored", False)), "node_runtime_profile": dict(metadata.get("node_runtime_profile", {}) or {}), "hybrid_runtime_profile": dict(metadata.get("hybrid_runtime_profile", {}) or {}), "adapter_runtime_profile": dict(metadata.get("adapter_runtime_profile", {}) or {})}
1421
+
1422
+ fast_ranked: list[dict[str, Any]] = []
1423
+ fast_scores: dict[str, float] = {}
1424
+ dense_elapsed = 0.0
1425
+ cross_elapsed = 0.0
1426
+ if plan["mode"] != "SLOW_ONLY":
1427
+ dense_started = time.time()
1428
+ dense_scores = fast_vectors @ models.dense.encode_one(runtime_question)
1429
+ dense_order = sorted(range(len(fast)), key=lambda i: (-float(dense_scores[i]), i))
1430
+ dense_rank = {index: rank for rank, index in enumerate(dense_order)}
1431
+ dense_elapsed = time.time() - dense_started
1432
+ indexes = list(dict.fromkeys([*dense_order[:args.dense_k], *graph_runtime_order]))
1433
+ if not indexes:
1434
+ raise RuntimeError(f"{qid}: graph+dense union is empty")
1435
+ session_count = max(int(candidate["session_index"]) for candidate in fast) + 1
1436
+ runtime = []
1437
+ for index in indexes:
1438
+ item = dict(fast[index])
1439
+ item["channels"] = {"dense_score": float(dense_scores[index]), "dense_rank_rr": rrank(dense_rank[index]), "graph_rank_rr": rrank(graph_rank.get(index)), "graph_selected": float(index in selected_indexes), "graph_final": float(index in final_indexes), "recency_norm": float(item["session_index"]) / max(1, session_count - 1)}
1440
+ runtime.append(item)
1441
+ cross_started = time.time()
1442
+ reps, logits = models.encode_cross(runtime_question, [item["text"] for item in runtime])
1443
+ channel_tensor = torch.tensor([[item["channels"][name] for name in CHANNEL_NAMES] for item in runtime], dtype=torch.float32, device=models.device)
1444
+ with torch.inference_mode():
1445
+ scores = models.fusion(reps.unsqueeze(0), logits.unsqueeze(0), channel_tensor.unsqueeze(0), torch.ones((1, len(runtime)), dtype=torch.bool, device=models.device), ablation="full")[0].detach().cpu()
1446
+ cross_elapsed += time.time() - cross_started
1447
+ for item, score, semantic in zip(runtime, scores.tolist(), logits.detach().cpu().tolist()):
1448
+ item["score"], item["semantic_logit"] = float(score), float(semantic)
1449
+ fast_ranked = _fast_candidates_with_slots(sorted(runtime, key=lambda item: -float(item["score"])), semantic_records)
1450
+ fast_scores = {str(item["candidate_id"]): float(item["score"]) for item in fast_ranked}
1451
+
1452
+ slow_ranked: list[dict[str, Any]] = []
1453
+ slow_dense_elapsed = 0.0
1454
+ if slow and plan["mode"] != "FAST_ONLY":
1455
+ slow_dense_started = time.time()
1456
+ slow_dense_scores = slow_vectors @ models.dense.encode_one(runtime_question)
1457
+ slow_dense_order = sorted(range(len(slow)), key=lambda i: (-float(slow_dense_scores[i]), i))
1458
+ slow_dense_elapsed = time.time() - slow_dense_started
1459
+ slow_indexes = slow_dense_order[: min(len(slow_dense_order), int(args.slow_dense_k))]
1460
+ if not slow_indexes:
1461
+ raise RuntimeError(f"{qid}: slow inventory is nonempty but dense shortlist is empty")
1462
+ slow_ranked = [dict(slow[index]) for index in slow_indexes]
1463
+ for index, item in zip(slow_indexes, slow_ranked):
1464
+ item["slow_dense_score"] = float(slow_dense_scores[index])
1465
+ item["slow_dense_rank"] = int(slow_dense_order.index(index))
1466
+ cross_started = time.time()
1467
+ _, logits = models.encode_cross(runtime_question, [item["text"] for item in slow_ranked])
1468
+ cross_elapsed += time.time() - cross_started
1469
+ for item, score in zip(slow_ranked, logits.detach().cpu().tolist()):
1470
+ item["semantic_logit"] = float(score)
1471
+ slow_ranked.sort(key=lambda item: -float(item["semantic_logit"]))
1472
+ try:
1473
+ units = apply_recall_plan(plan, fast_ranked, slow_ranked)
1474
+ except RecallPlannerError as exc:
1475
+ raise RuntimeError(f"{qid}: invalid planned evidence composition: {exc}") from exc
1476
+ if plan["mode"] == "CONFLICT_COMPARE":
1477
+ units.sort(
1478
+ key=lambda item: (
1479
+ -max(
1480
+ (
1481
+ fast_scores.get(str(candidate.get("candidate_id", "")), float("-inf"))
1482
+ for candidate in item.get("fast_candidates", [])
1483
+ ),
1484
+ default=float("-inf"),
1485
+ ),
1486
+ str(item["canonical_slot"]),
1487
+ )
1488
+ )
1489
+ elif plan["primary_layer"] == "fast":
1490
+ units.sort(key=lambda item: -fast_scores.get(str(item.get("fast_candidate", {}).get("candidate_id", "")), float("-inf")))
1491
+ else:
1492
+ units.sort(key=lambda item: -float(item.get("slow_capsule", {}).get("semantic_logit", float("-inf"))))
1493
+
1494
+ packed_units = pack_recall_units(units, fast, top_k=args.top_k, qid=qid)
1495
+ evidence_by_location: dict[tuple[int, int, int], dict[str, Any]] = {}
1496
+ for unit_rank, (unit, entries) in enumerate(packed_units, start=1):
1497
+ attachments = _unit_attachments(unit)
1498
+ for entry in entries:
1499
+ capsules = list(entry.get("capsules") or [])
1500
+ memory_contexts = [
1501
+ {
1502
+ "role": next(
1503
+ (
1504
+ role
1505
+ for role in entry["roles"]
1506
+ if role in {"primary", "slow_conflict_candidate"}
1507
+ ),
1508
+ "slow_context",
1509
+ ),
1510
+ "capsule_id": item["capsule_id"],
1511
+ "canonical_slot": item["canonical_slot"],
1512
+ "claim_text": item["text"],
1513
+ "provenance": item["provenance"],
1514
+ }
1515
+ for item in capsules
1516
+ ]
1517
+ location = (
1518
+ int(entry["session_index"]),
1519
+ int(entry["parent_chunk_index"]),
1520
+ int(entry["subchunk_index"]),
1521
+ )
1522
+ candidate = {
1523
+ "memory_id": entry.get("candidate_id"),
1524
+ "session_id": entry["session_id"],
1525
+ "session_index": entry["session_index"],
1526
+ "parent_chunk_index": entry["parent_chunk_index"],
1527
+ "subchunk_index": entry["subchunk_index"],
1528
+ "score": entry.get("score", entry.get("semantic_logit")),
1529
+ "semantic_logit": entry.get("semantic_logit"),
1530
+ "channels": entry.get("channels"),
1531
+ "text": entry["text"],
1532
+ "unit_type": unit["unit_type"],
1533
+ "unit_types": [unit["unit_type"]],
1534
+ "canonical_slot": unit["canonical_slot"],
1535
+ "canonical_slots": [unit["canonical_slot"]],
1536
+ "capsule_id": [item["capsule_id"] for item in capsules],
1537
+ "role": list(entry["roles"]),
1538
+ "provenance": [item["provenance"] for item in capsules]
1539
+ or [entry.get("provenance")],
1540
+ "memory_contexts": memory_contexts,
1541
+ "attachments": attachments,
1542
+ }
1543
+ existing = evidence_by_location.get(location)
1544
+ if existing is None:
1545
+ evidence_by_location[location] = candidate
1546
+ continue
1547
+ for field in (
1548
+ "unit_types",
1549
+ "canonical_slots",
1550
+ "capsule_id",
1551
+ "role",
1552
+ "provenance",
1553
+ "memory_contexts",
1554
+ "attachments",
1555
+ ):
1556
+ for item in candidate[field]:
1557
+ if item not in existing[field]:
1558
+ existing[field].append(item)
1559
+ evidence_windows = list(evidence_by_location.values())
1560
+ for rank, item in enumerate(evidence_windows, start=1):
1561
+ item["rank"] = rank
1562
+ selected_session_ids = list(
1563
+ dict.fromkeys(str(item["session_id"]) for item in evidence_windows)
1564
+ )
1565
+ evidence = {"schema_version": SCHEMA_VERSION, "runtime_schema_version": "tmcra.v3.online-retrieval.3", "question_id": qid, "question": question, "question_date": question_date, "question_type": clean_text(row.get("question_type")), "selected_session_ids": selected_session_ids, "recall_plan": authoritative_plan, "evidence_windows": evidence_windows}
1566
+ runtime_input_has_gold = any(key in row for key in ("answer", "gold_answer", "answer_session_ids", "labels", "supervision"))
1567
+ debug = {"question_id": qid, "scope_id": scope_id, "db_path": str(db_path), "index_path": str(index_path), "recall_plan": authoritative_plan, "planner_decision_reason": planner_decision_reason, "planner_resolved_query": resolved_query, "recent_dialogue_count": len(recent_dialogue), "planner": planner_metadata, "cross_layer_weighted_fusion": False, "graph": graph_metadata, "runtime_input_has_gold": runtime_input_has_gold, "inventory_count": len(fast) + len(slow), "union_count": len(fast_ranked) + len(slow_ranked), "fast_inventory_count": len(fast), "fast_shortlist_count": len(fast_ranked), "slow_capsule_count": len(slow), "slow_shortlist_count": len(slow_ranked), "slow_dense_k": int(args.slow_dense_k), "fast_semantic_record_count": len(semantic_records), "planned_unit_count": len(units), "packed_unit_count": len(packed_units), "budget_excluded_unit_count": max(0, len(units) - len(packed_units)), "selected_count": len(evidence_windows), "atomic_unit_packing": True, "strict_no_truncation": True, "restart_boundary_verified": True, "graph_counts_before_query": counts_before, "graph_fingerprint": graph_fingerprint, "latency_sec": {"graph": round(graph_elapsed, 4), "dense": round(dense_elapsed, 4), "slow_dense": round(slow_dense_elapsed, 4), "cross": round(cross_elapsed, 4), "total": round(time.time() - started, 4)}}
1568
+ if any(key in row for key in ("answer", "gold_answer", "answer_session_ids", "labels", "supervision")):
1569
+ raise RuntimeError(f"{qid}: runtime retrieval manifest contains forbidden evaluation labels")
1570
+ return evidence, debug
1571
+
1572
+
1573
+ def layered_retrieval_operation_id(out_dir: Path, scope_id: str, question_id: str) -> str:
1574
+ identity = "\n".join(
1575
+ (
1576
+ "tmcra.v3.layered-retrieval",
1577
+ str(out_dir.absolute()),
1578
+ clean_text(scope_id),
1579
+ clean_text(question_id),
1580
+ )
1581
+ )
1582
+ return hashlib.sha256(identity.encode("utf-8")).hexdigest()
1583
+
1584
+
1585
+ def _atomic_write_text(path: Path, text: str) -> None:
1586
+ temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}.{time.time_ns()}")
1587
+ temporary.write_text(text, encoding="utf-8")
1588
+ os.replace(temporary, path)
1589
+
1590
+
1591
+ def _reconcile_completed_retrieval_output(
1592
+ *,
1593
+ out_dir: Path,
1594
+ rows: Sequence[Mapping[str, Any]],
1595
+ repo: Path,
1596
+ ) -> dict[str, Any] | None:
1597
+ """Reuse a committed retrieval directory and repair only missing idempotent audit state."""
1598
+ if not out_dir.exists():
1599
+ return None
1600
+ required = {
1601
+ "evidence": out_dir / "evidence_windows.jsonl",
1602
+ "debug": out_dir / "retrieval_debug.jsonl",
1603
+ "report": out_dir / "report.json",
1604
+ }
1605
+ missing = [name for name, path in required.items() if not path.is_file()]
1606
+ if missing:
1607
+ raise RuntimeError(
1608
+ f"retrieval output directory is incomplete and will not be overwritten: "
1609
+ f"{out_dir} missing={','.join(missing)}"
1610
+ )
1611
+ evidence_rows = read_jsonl(required["evidence"])
1612
+ debug_rows = read_jsonl(required["debug"])
1613
+ report = json.loads(required["report"].read_text(encoding="utf-8"))
1614
+ expected_qids = [clean_text(row.get("question_id")) for row in rows]
1615
+ evidence_qids = [clean_text(row.get("question_id")) for row in evidence_rows]
1616
+ debug_qids = [clean_text(row.get("question_id")) for row in debug_rows]
1617
+ if (
1618
+ report.get("status") != "complete"
1619
+ or evidence_qids != expected_qids
1620
+ or debug_qids != expected_qids
1621
+ ):
1622
+ raise RuntimeError(f"existing retrieval output does not match its query manifest: {out_dir}")
1623
+
1624
+ newly_appended = 0
1625
+ query_ids: list[str] = []
1626
+ for evidence, debug in zip(evidence_rows, debug_rows):
1627
+ plan = dict(evidence.get("recall_plan") or {})
1628
+ legacy_reason = clean_text(plan.pop("decision_reason", ""))
1629
+ if legacy_reason:
1630
+ observed_reason = clean_text(debug.get("planner_decision_reason"))
1631
+ if observed_reason and observed_reason != legacy_reason:
1632
+ raise RuntimeError(
1633
+ f"{evidence['question_id']}: persisted planner reason disagrees with debug output"
1634
+ )
1635
+ debug["planner_decision_reason"] = legacy_reason
1636
+ evidence["recall_plan"] = plan
1637
+ operation_id = layered_retrieval_operation_id(
1638
+ out_dir,
1639
+ str(debug["scope_id"]),
1640
+ str(evidence["question_id"]),
1641
+ )
1642
+ persisted = append_layered_retrieval_audit(
1643
+ repo=repo,
1644
+ db_path=Path(debug["db_path"]),
1645
+ scope_id=str(debug["scope_id"]),
1646
+ operation_id=operation_id,
1647
+ evidence=evidence,
1648
+ debug=debug,
1649
+ )
1650
+ audit_payload = dict(persisted.get("payload") or {})
1651
+ query_id = clean_text(audit_payload.get("query_id"))
1652
+ if not query_id:
1653
+ raise RuntimeError(f"{evidence['question_id']}: persisted retrieval audit lacks query_id")
1654
+ newly_appended += int(bool(persisted.get("appended")))
1655
+ query_ids.append(query_id)
1656
+ debug["layered_retrieval_audit"] = {
1657
+ "operation_id": operation_id,
1658
+ "query_id": query_id,
1659
+ "event_total": int(persisted.get("event_total") or 0),
1660
+ "trimmed_total": int(persisted.get("trimmed_total") or 0),
1661
+ "newly_appended": bool(persisted.get("appended")),
1662
+ }
1663
+
1664
+ report["layered_retrieval_audit_confirmed_count"] = len(query_ids)
1665
+ report["layered_retrieval_audit_newly_appended_count"] = newly_appended
1666
+ report["layered_retrieval_audit_query_ids"] = query_ids
1667
+ report["reused_completed_output"] = True
1668
+ _atomic_write_text(
1669
+ required["evidence"],
1670
+ "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in evidence_rows),
1671
+ )
1672
+ _atomic_write_text(
1673
+ required["debug"],
1674
+ "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in debug_rows),
1675
+ )
1676
+ _atomic_write_text(
1677
+ required["report"], json.dumps(report, indent=2, sort_keys=True) + "\n"
1678
+ )
1679
+ return report
1680
+
1681
+
1682
+ def command_retrieve(args: argparse.Namespace) -> None:
1683
+ query_manifest = Path(args.query_manifest)
1684
+ try:
1685
+ rows = read_jsonl(query_manifest)
1686
+ except RuntimeError as exc:
1687
+ if str(exc) == f"no rows: {query_manifest}":
1688
+ raise RuntimeError(f"query manifest is empty: {query_manifest}") from exc
1689
+ raise
1690
+ out_dir = Path(args.out_dir).absolute()
1691
+ completed_report = _reconcile_completed_retrieval_output(
1692
+ out_dir=out_dir,
1693
+ rows=rows,
1694
+ repo=Path(args.repo),
1695
+ )
1696
+ if completed_report is not None:
1697
+ print(json.dumps(completed_report, indent=2, sort_keys=True))
1698
+ return
1699
+ runtime_env = graph_runtime_env(args)
1700
+ for path in (Path(args.node_model), Path(args.path_model), Path(args.checkpoint)):
1701
+ if not path.exists():
1702
+ raise FileNotFoundError(path)
1703
+ planner = planner_from_env()
1704
+ harness = load_native_harness(Path(args.harness), Path(args.repo))
1705
+ harness.disable_topic_bucket_runtime()
1706
+ models = OnlineModels(args)
1707
+ started = time.time()
1708
+ evidence_rows: list[dict[str, Any]] = []
1709
+ debug_rows: list[dict[str, Any]] = []
1710
+ graph_adapter_cache: dict[tuple[str, str], Any] = {}
1711
+ newly_appended_audits = 0
1712
+ for row_index, row in enumerate(rows, start=1):
1713
+ evidence, debug = retrieve_one(
1714
+ row,
1715
+ args=args,
1716
+ harness=harness,
1717
+ models=models,
1718
+ planner=planner,
1719
+ graph_adapter_cache=graph_adapter_cache,
1720
+ )
1721
+ operation_id = layered_retrieval_operation_id(
1722
+ out_dir,
1723
+ str(debug["scope_id"]),
1724
+ str(evidence["question_id"]),
1725
+ )
1726
+ persisted_audit = append_layered_retrieval_audit(
1727
+ repo=Path(args.repo),
1728
+ db_path=Path(debug["db_path"]),
1729
+ scope_id=str(debug["scope_id"]),
1730
+ operation_id=operation_id,
1731
+ evidence=evidence,
1732
+ debug=debug,
1733
+ )
1734
+ newly_appended_audits += int(bool(persisted_audit.get("appended")))
1735
+ debug["layered_retrieval_audit"] = {
1736
+ "operation_id": operation_id,
1737
+ "query_id": clean_text(dict(persisted_audit.get("payload") or {}).get("query_id")),
1738
+ "event_total": int(persisted_audit.get("event_total") or 0),
1739
+ "trimmed_total": int(persisted_audit.get("trimmed_total") or 0),
1740
+ "newly_appended": bool(persisted_audit.get("appended")),
1741
+ }
1742
+ evidence_rows.append(evidence)
1743
+ debug_rows.append(debug)
1744
+ print(
1745
+ json.dumps(
1746
+ {
1747
+ "status": "retrieved",
1748
+ "row": row_index,
1749
+ "total": len(rows),
1750
+ "question_id": evidence["question_id"],
1751
+ "inventory": debug["inventory_count"],
1752
+ "union": debug["union_count"],
1753
+ "selected": debug["selected_count"],
1754
+ "latency_sec": debug["latency_sec"]["total"],
1755
+ }
1756
+ ),
1757
+ flush=True,
1758
+ )
1759
+ latency = [float(row["latency_sec"]["total"]) for row in debug_rows]
1760
+ report = {
1761
+ "status": "complete",
1762
+ "schema_version": "tmcra.v3.online-retrieval-report.2",
1763
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
1764
+ "query_count": len(rows),
1765
+ "runtime_input_with_gold_count": sum(int(row["runtime_input_has_gold"]) for row in debug_rows),
1766
+ "restart_boundary_verified_count": sum(int(row["restart_boundary_verified"]) for row in debug_rows),
1767
+ "layered_retrieval_audit_confirmed_count": sum(
1768
+ int(bool(row.get("layered_retrieval_audit"))) for row in debug_rows
1769
+ ),
1770
+ "layered_retrieval_audit_newly_appended_count": newly_appended_audits,
1771
+ "layered_retrieval_audit_query_ids": [
1772
+ row["layered_retrieval_audit"]["query_id"] for row in debug_rows
1773
+ ],
1774
+ "checkpoint": str(models.checkpoint_path),
1775
+ "checkpoint_sha256": models.checkpoint_sha256,
1776
+ "cross_model_revision": models.cross_manifest.get("revision"),
1777
+ "graph_runtime_env": runtime_env,
1778
+ "strict_no_truncation": True,
1779
+ "cross_layer_weighted_fusion": False,
1780
+ "graph_adapter_cache_size": len(graph_adapter_cache),
1781
+ "planner": {
1782
+ "base_url_configured": bool(os.environ.get("TMCRA_RECALL_PLANNER_BASE_URL")),
1783
+ "model": os.environ.get("TMCRA_RECALL_PLANNER_MODEL", DEEPSEEK_FLASH_MODEL),
1784
+ "api_key_pool_size": len(planner.api_keys),
1785
+ },
1786
+ "top_k": args.top_k,
1787
+ "dense_k": args.dense_k,
1788
+ "graph_k": args.graph_k,
1789
+ "slow_dense_k": args.slow_dense_k,
1790
+ "atomic_unit_packing": True,
1791
+ "avg_planned_unit_count": round(sum(row["planned_unit_count"] for row in debug_rows) / len(debug_rows), 4),
1792
+ "avg_packed_unit_count": round(sum(row["packed_unit_count"] for row in debug_rows) / len(debug_rows), 4),
1793
+ "avg_budget_excluded_unit_count": round(sum(row["budget_excluded_unit_count"] for row in debug_rows) / len(debug_rows), 4),
1794
+ "avg_inventory_count": round(sum(row["inventory_count"] for row in debug_rows) / len(debug_rows), 4),
1795
+ "avg_union_count": round(sum(row["union_count"] for row in debug_rows) / len(debug_rows), 4),
1796
+ "avg_latency_sec": round(sum(latency) / len(latency), 4),
1797
+ "max_latency_sec": round(max(latency), 4),
1798
+ "elapsed_sec": round(time.time() - started, 3),
1799
+ "evidence": str(out_dir / "evidence_windows.jsonl"),
1800
+ "debug": str(out_dir / "retrieval_debug.jsonl"),
1801
+ "reused_completed_output": False,
1802
+ }
1803
+ out_dir.parent.mkdir(parents=True, exist_ok=True)
1804
+ staging = out_dir.with_name(
1805
+ f".{out_dir.name}.staging.{os.getpid()}.{time.time_ns()}"
1806
+ )
1807
+ staging.mkdir(parents=False, exist_ok=False)
1808
+ write_jsonl(staging / "evidence_windows.jsonl", evidence_rows)
1809
+ write_jsonl(staging / "retrieval_debug.jsonl", debug_rows)
1810
+ (staging / "report.json").write_text(
1811
+ json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
1812
+ )
1813
+ os.replace(staging, out_dir)
1814
+ print(json.dumps(report, indent=2, sort_keys=True))
1815
+
1816
+
1817
+ def add_common_model_args(parser: argparse.ArgumentParser) -> None:
1818
+ parser.add_argument("--embedding-model", default=os.getenv("TMCRA_BGE_M3_PATH", "BAAI/bge-m3"))
1819
+ parser.add_argument("--text-dim", type=int, default=1024)
1820
+ parser.add_argument("--embedding-max-length", type=int, default=8192)
1821
+ parser.add_argument("--device", default="cuda")
1822
+
1823
+
1824
+ def main() -> int:
1825
+ parser = argparse.ArgumentParser(description="TMCRA V3 production online index and retrieval runtime")
1826
+ sub = parser.add_subparsers(dest="command", required=True)
1827
+ index = sub.add_parser("build-index")
1828
+ index.add_argument("--scope-manifest", required=True)
1829
+ index.add_argument("--out-report", required=True)
1830
+ index.add_argument("--subchunk-chars", type=int, default=1800)
1831
+ index.add_argument("--subchunk-overlap", type=int, default=200)
1832
+ index.add_argument("--batch-size", type=int, default=16)
1833
+ add_common_model_args(index)
1834
+
1835
+ retrieve = sub.add_parser("retrieve")
1836
+ retrieve.add_argument("--query-manifest", required=True)
1837
+ retrieve.add_argument("--out-dir", required=True)
1838
+ retrieve.add_argument("--checkpoint", required=True)
1839
+ retrieve.add_argument("--cross-model", default=os.getenv("TMCRA_BGE_RERANKER_PATH", "BAAI/bge-reranker-v2-m3"))
1840
+ retrieve.add_argument("--cross-max-length", type=int, default=1280)
1841
+ retrieve.add_argument("--cross-batch-size", type=int, default=24)
1842
+ retrieve.add_argument("--repo", required=True)
1843
+ retrieve.add_argument("--harness", required=True)
1844
+ retrieve.add_argument("--node-model", required=True)
1845
+ retrieve.add_argument("--path-model", required=True)
1846
+ retrieve.add_argument("--graph-device", default="cuda")
1847
+ retrieve.add_argument("--candidate-event-k", type=int, default=24)
1848
+ retrieve.add_argument("--support-path-k", type=int, default=3)
1849
+ retrieve.add_argument("--path-tunnel-rescue-k", type=int, default=2)
1850
+ retrieve.add_argument("--graph-top-k", type=int, default=12)
1851
+ retrieve.add_argument("--dense-k", type=int, default=32)
1852
+ retrieve.add_argument("--slow-dense-k", type=int, default=24)
1853
+ retrieve.add_argument("--graph-k", type=int, default=24)
1854
+ retrieve.add_argument("--top-k", type=int, default=8)
1855
+ retrieve.add_argument("--max-per-parent", type=int, default=2)
1856
+ retrieve.add_argument("--max-per-session", type=int, default=4)
1857
+ add_common_model_args(retrieve)
1858
+ args = parser.parse_args()
1859
+ if args.command == "build-index":
1860
+ command_build_index(args)
1861
+ else:
1862
+ command_retrieve(args)
1863
+ return 0
1864
+
1865
+
1866
+ if __name__ == "__main__":
1867
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_product_writer.py ADDED
The diff for this file is too large to render. See raw diff
 
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_recall_planner.py ADDED
@@ -0,0 +1,406 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Strict DeepSeek Flash control plane for recall composition."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import hashlib
6
+ import json
7
+ import time
8
+ import urllib.error
9
+ import urllib.request
10
+ from collections.abc import Mapping, Sequence
11
+ from typing import Any
12
+
13
+
14
+ PLANNER_VERSION = "tmcra.recall-plan.v4"
15
+ PLANNER_PROMPT_VERSION = "tmcra-recall-planner-2026-07-11.1"
16
+ DEEPSEEK_FLASH_MODEL = "deepseek-v4-flash"
17
+ MODES = frozenset(
18
+ {
19
+ "FAST_ONLY",
20
+ "SLOW_ONLY",
21
+ "SLOW_WITH_FAST_OVERRIDE",
22
+ "FAST_WITH_SLOW_CONTEXT",
23
+ "CONFLICT_COMPARE",
24
+ }
25
+ )
26
+ PLAN_FIELDS = frozenset(
27
+ {
28
+ "mode",
29
+ "primary_layer",
30
+ "fast_role",
31
+ "slow_role",
32
+ "requires_conflict_pairs",
33
+ "decision_code",
34
+ "resolved_query",
35
+ "decision_reason",
36
+ "planner_version",
37
+ }
38
+ )
39
+ _EXPECTED_POLICIES = {
40
+ "FAST_ONLY": ("fast", "primary", "excluded", False),
41
+ "SLOW_ONLY": ("slow", "excluded", "primary", False),
42
+ "SLOW_WITH_FAST_OVERRIDE": ("slow", "override", "primary", False),
43
+ "FAST_WITH_SLOW_CONTEXT": ("fast", "primary", "context_only", False),
44
+ "CONFLICT_COMPARE": ("fast", "conflict_candidate", "conflict_candidate", True),
45
+ }
46
+ _EXPECTED_DECISION_CODES = {
47
+ "FAST_ONLY": "recent_event",
48
+ "SLOW_ONLY": "stable_or_historical_durable",
49
+ "SLOW_WITH_FAST_OVERRIDE": "current_durable",
50
+ "FAST_WITH_SLOW_CONTEXT": "current_task_with_durable_context",
51
+ "CONFLICT_COMPARE": "explicit_conflict_or_comparison",
52
+ }
53
+
54
+ SYSTEM_PROMPT = f"""You are the retrieval control plane of a production hierarchical memory graph.
55
+ Return exactly one JSON object with these fields and no others:
56
+ {{"mode":"FAST_ONLY|SLOW_ONLY|SLOW_WITH_FAST_OVERRIDE|FAST_WITH_SLOW_CONTEXT|CONFLICT_COMPARE",
57
+ "primary_layer":"fast|slow",
58
+ "fast_role":"primary|excluded|override|conflict_candidate",
59
+ "slow_role":"primary|excluded|context_only|conflict_candidate",
60
+ "requires_conflict_pairs":true|false,
61
+ "decision_code":"recent_event|stable_or_historical_durable|current_durable|current_task_with_durable_context|explicit_conflict_or_comparison",
62
+ "resolved_query":"standalone retrieval query in the user's language",
63
+ "decision_reason":"concise reason grounded only in the query and recent dialogue",
64
+ "planner_version":"{PLANNER_VERSION}"}}
65
+
66
+ Mode contracts are exact:
67
+ - FAST_ONLY: explicitly recent episodic event, current task, or interaction lookup that needs no durable user background.
68
+ - SLOW_ONLY: explicitly historical baseline, long-term background, or stable summary where the current effective value is not requested.
69
+ - SLOW_WITH_FAST_OVERRIDE: the default for a present/current durable identity, preference, relationship, constraint, or state. Slow is primary and fast must check for recent corrections even when the query does not say "latest".
70
+ - FAST_WITH_SLOW_CONTEXT: a current event/task is primary and durable user context is background-only.
71
+ - CONFLICT_COMPARE: the query explicitly asks about change, contradiction, before/after, or competing states.
72
+
73
+ Hard routing order:
74
+ 1. Explicit comparison, contradiction, or before/after change -> CONFLICT_COMPARE.
75
+ 2. A present/current/latest/still standing fact -> SLOW_WITH_FAST_OVERRIDE, never SLOW_ONLY.
76
+ 3. A current task or event that needs durable personalization -> FAST_WITH_SLOW_CONTEXT.
77
+ 4. An isolated recent event, task, or exact recent-turn lookup -> FAST_ONLY.
78
+ 5. SLOW_ONLY is allowed only for an explicitly historical baseline or stable long-term summary.
79
+ 6. An unresolved or uncertain follow-up must never use SLOW_ONLY; resolve it from recent_dialogue first.
80
+
81
+ Resolve references using recent_dialogue only when needed. resolved_query must preserve the user's language and intent,
82
+ be independently searchable without recent_dialogue, and must not answer the query or invent a memory fact. If the query
83
+ is already standalone or the dialogue is unrelated, copy the query unchanged. The current query always overrides stale
84
+ or unrelated dialogue.
85
+
86
+ Never emit scores, weights, candidate IDs, benchmark labels, or an answer to the query. Never use a layer marked
87
+ unavailable. Choose exactly one mode. This planner chooses roles only; it does not rank or fuse evidence."""
88
+
89
+
90
+ class RecallPlannerError(RuntimeError):
91
+ pass
92
+
93
+
94
+ class RecallPlannerResponseError(RecallPlannerError):
95
+ def __init__(self, message: str, *, response_content: str = "", request_metadata: Mapping[str, Any] | None = None) -> None:
96
+ super().__init__(message)
97
+ self.response_content = response_content
98
+ self.request_metadata = dict(request_metadata or {})
99
+
100
+
101
+ def _text(value: Any) -> str:
102
+ return value.strip() if isinstance(value, str) else ""
103
+
104
+
105
+ def _sha256(value: str) -> str:
106
+ return hashlib.sha256(value.encode("utf-8")).hexdigest()
107
+
108
+
109
+ def _metadata(*, key_index: int, started: float, finish_reason: str, content: str = "", usage: Mapping[str, Any] | None = None, http_status: int | None = None, error_type: str | None = None, request_sha256: str = "") -> dict[str, Any]:
110
+ usage = usage or {}
111
+ result = {
112
+ "provider": "deepseek",
113
+ "model": DEEPSEEK_FLASH_MODEL,
114
+ "api_key_index": key_index,
115
+ "physical_call_count": 1,
116
+ "latency_seconds": round(time.time() - started, 3),
117
+ "response_sha256": _sha256(content),
118
+ "finish_reason": finish_reason,
119
+ "planner_version": PLANNER_VERSION,
120
+ "prompt_version": PLANNER_PROMPT_VERSION,
121
+ "request_sha256": request_sha256,
122
+ "prompt_tokens": int(usage.get("prompt_tokens", 0) or 0),
123
+ "completion_tokens": int(usage.get("completion_tokens", 0) or 0),
124
+ "total_tokens": int(usage.get("total_tokens", 0) or 0),
125
+ }
126
+ if http_status is not None:
127
+ result["http_status"] = int(http_status)
128
+ if error_type:
129
+ result["error_type"] = error_type
130
+ return result
131
+
132
+
133
+ def _reject_gold(value: Any, path: str = "input") -> None:
134
+ if isinstance(value, Mapping):
135
+ for key, item in value.items():
136
+ key_text = str(key).lower()
137
+ if "gold" in key_text or "benchmark" in key_text or "expected_answer" in key_text:
138
+ raise RecallPlannerError(f"{path}.{key} is not permitted in planner input")
139
+ _reject_gold(item, f"{path}.{key}")
140
+ elif isinstance(value, list):
141
+ for index, item in enumerate(value):
142
+ _reject_gold(item, f"{path}[{index}]")
143
+
144
+
145
+ def validate_recall_plan(value: Mapping[str, Any]) -> dict[str, Any]:
146
+ if not isinstance(value, Mapping) or set(value) != PLAN_FIELDS:
147
+ raise RecallPlannerError("RecallPlan root must contain exactly the required schema fields")
148
+ mode = _text(value.get("mode"))
149
+ if mode not in MODES:
150
+ raise RecallPlannerError(f"unsupported RecallPlan mode: {mode!r}")
151
+ primary_layer, fast_role, slow_role, requires_pairs = _EXPECTED_POLICIES[mode]
152
+ if (
153
+ value.get("primary_layer") != primary_layer
154
+ or value.get("fast_role") != fast_role
155
+ or value.get("slow_role") != slow_role
156
+ or value.get("requires_conflict_pairs") is not requires_pairs
157
+ ):
158
+ raise RecallPlannerError(f"RecallPlan role policy does not match mode {mode}")
159
+ if value.get("decision_code") != _EXPECTED_DECISION_CODES[mode]:
160
+ raise RecallPlannerError(f"RecallPlan decision_code does not match mode {mode}")
161
+ reason = _text(value.get("decision_reason"))
162
+ if not reason:
163
+ raise RecallPlannerError("RecallPlan decision_reason is required")
164
+ resolved_query = _text(value.get("resolved_query"))
165
+ if not resolved_query:
166
+ raise RecallPlannerError("RecallPlan resolved_query is required")
167
+ if len(resolved_query) > 2000:
168
+ raise RecallPlannerError("RecallPlan resolved_query exceeds 2000 characters")
169
+ if value.get("planner_version") != PLANNER_VERSION:
170
+ raise RecallPlannerError(f"RecallPlan planner_version must be {PLANNER_VERSION!r}")
171
+ return {field: value[field] for field in PLAN_FIELDS}
172
+
173
+
174
+ def _validate_recent_dialogue(value: Any) -> list[dict[str, Any]]:
175
+ if value is None:
176
+ return []
177
+ if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
178
+ raise RecallPlannerError("recent_dialogue must be a sequence of turn objects")
179
+ if len(value) > 8:
180
+ raise RecallPlannerError("recent_dialogue may contain at most 8 turns")
181
+ output: list[dict[str, Any]] = []
182
+ for index, item in enumerate(value):
183
+ if not isinstance(item, Mapping) or set(item) != {"turn_index", "speaker", "text"}:
184
+ raise RecallPlannerError(f"recent_dialogue[{index}] must contain exactly turn_index, speaker, and text")
185
+ try:
186
+ turn_index = int(item["turn_index"])
187
+ except (TypeError, ValueError) as exc:
188
+ raise RecallPlannerError(f"recent_dialogue[{index}].turn_index must be an integer") from exc
189
+ speaker = _text(item.get("speaker")).lower()
190
+ text = _text(item.get("text"))
191
+ if speaker not in {"user", "assistant"} or not text:
192
+ raise RecallPlannerError(f"recent_dialogue[{index}] has an invalid speaker or empty text")
193
+ if len(text) > 4000:
194
+ raise RecallPlannerError(f"recent_dialogue[{index}].text exceeds 4000 characters")
195
+ output.append({"turn_index": turn_index, "speaker": speaker, "text": text})
196
+ if any(left["turn_index"] >= right["turn_index"] for left, right in zip(output, output[1:])):
197
+ raise RecallPlannerError("recent_dialogue must be in strictly increasing turn order")
198
+ _reject_gold(output, "recent_dialogue")
199
+ return output
200
+
201
+
202
+ def _layer_available(value: Any) -> bool:
203
+ if isinstance(value, Mapping) and "available" in value:
204
+ return value.get("available") is True
205
+ if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
206
+ return bool(value)
207
+ return bool(value)
208
+
209
+
210
+ def validate_plan_availability(plan: Mapping[str, Any], available_layers: Mapping[str, Any]) -> None:
211
+ fast_available = _layer_available(available_layers.get("fast"))
212
+ slow_available = _layer_available(available_layers.get("slow"))
213
+ mode = plan["mode"]
214
+ if mode in {"FAST_ONLY", "SLOW_WITH_FAST_OVERRIDE", "FAST_WITH_SLOW_CONTEXT", "CONFLICT_COMPARE"} and not fast_available:
215
+ raise RecallPlannerError(f"RecallPlan mode {mode} requires an unavailable fast layer")
216
+ if mode in {"SLOW_ONLY", "SLOW_WITH_FAST_OVERRIDE", "FAST_WITH_SLOW_CONTEXT", "CONFLICT_COMPARE"} and not slow_available:
217
+ raise RecallPlannerError(f"RecallPlan mode {mode} requires an unavailable slow layer")
218
+
219
+
220
+ class DeepSeekFlashRecallPlanner:
221
+ """One physical Flash call per plan request; errors are never retried or routed."""
222
+
223
+ def __init__(self, *, base_url: str, model: str, api_keys: Sequence[str], timeout: float = 30, max_tokens: int = 512) -> None:
224
+ self.base_url = _text(base_url).rstrip("/")
225
+ self.model = model
226
+ self.api_keys = list(dict.fromkeys(_text(key) for key in api_keys if _text(key)))
227
+ self.timeout = max(1.0, float(timeout))
228
+ self.max_tokens = max(128, int(max_tokens))
229
+ self.request_index = 0
230
+ if not self.base_url or not self.api_keys:
231
+ raise RecallPlannerError("planner base_url and API key pool are required")
232
+ if self.model != DEEPSEEK_FLASH_MODEL:
233
+ raise RecallPlannerError(f"recall planner requires {DEEPSEEK_FLASH_MODEL}, got {self.model!r}")
234
+
235
+ def plan(self, *, query: str, question_date: str, available_layers: Mapping[str, Any], recent_dialogue: Sequence[Mapping[str, Any]] | None = None) -> tuple[dict[str, Any], dict[str, Any]]:
236
+ if not _text(query) or not _text(question_date):
237
+ raise RecallPlannerError("query and question_date are required")
238
+ if not isinstance(available_layers, Mapping) or set(available_layers) - {"fast", "slow"}:
239
+ raise RecallPlannerError("available_layers may contain only fast and slow summaries")
240
+ _reject_gold(available_layers, "available_layers")
241
+ dialogue = _validate_recent_dialogue(recent_dialogue)
242
+ payload = {"query": query, "question_date": question_date, "recent_dialogue": dialogue, "available_layers": dict(available_layers)}
243
+ key_index = self.request_index % len(self.api_keys)
244
+ self.request_index += 1
245
+ request_body = {
246
+ "model": DEEPSEEK_FLASH_MODEL,
247
+ "messages": [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": json.dumps(payload, ensure_ascii=False, separators=(",", ":"))}],
248
+ "temperature": 0,
249
+ "max_tokens": self.max_tokens,
250
+ "response_format": {"type": "json_object"},
251
+ "thinking": {"type": "disabled"},
252
+ }
253
+ request_sha256 = _sha256(json.dumps(request_body, ensure_ascii=False, sort_keys=True))
254
+ request = urllib.request.Request(
255
+ f"{self.base_url}/chat/completions",
256
+ data=json.dumps(request_body, ensure_ascii=False).encode("utf-8"),
257
+ headers={"Content-Type": "application/json", "Authorization": f"Bearer {self.api_keys[key_index]}"},
258
+ method="POST",
259
+ )
260
+ started = time.time()
261
+ try:
262
+ with urllib.request.urlopen(request, timeout=self.timeout) as response:
263
+ raw_http = response.read().decode("utf-8")
264
+ response_payload = json.loads(raw_http)
265
+ except urllib.error.HTTPError as exc:
266
+ detail = exc.read().decode("utf-8", errors="replace")[:1000]
267
+ raise RecallPlannerResponseError(f"recall planner HTTP {exc.code}: {detail}", response_content=detail, request_metadata=_metadata(key_index=key_index, started=started, finish_reason="http_error", content=detail, http_status=exc.code, request_sha256=request_sha256)) from exc
268
+ except (UnicodeDecodeError, json.JSONDecodeError) as exc:
269
+ raw = locals().get("raw_http", "")
270
+ raise RecallPlannerResponseError("recall planner returned invalid HTTP JSON", response_content=raw, request_metadata=_metadata(key_index=key_index, started=started, finish_reason="invalid_http_json", content=raw, error_type=exc.__class__.__name__, request_sha256=request_sha256)) from exc
271
+ except Exception as exc:
272
+ raise RecallPlannerResponseError(f"recall planner request failed: {exc.__class__.__name__}: {exc}", request_metadata=_metadata(key_index=key_index, started=started, finish_reason="request_error", error_type=exc.__class__.__name__, request_sha256=request_sha256)) from exc
273
+ usage = response_payload.get("usage") if isinstance(response_payload, Mapping) else {}
274
+ choices = response_payload.get("choices") if isinstance(response_payload, Mapping) else None
275
+ if not isinstance(choices, list) or len(choices) != 1 or not isinstance(choices[0], Mapping):
276
+ raise RecallPlannerResponseError("recall planner response must contain exactly one choice", response_content=raw_http, request_metadata=_metadata(key_index=key_index, started=started, finish_reason="invalid_response", content=raw_http, usage=usage if isinstance(usage, Mapping) else {}, request_sha256=request_sha256))
277
+ choice = choices[0]
278
+ content = (choice.get("message") or {}).get("content") if isinstance(choice.get("message"), Mapping) else None
279
+ finish_reason = _text(choice.get("finish_reason"))
280
+ metadata = _metadata(key_index=key_index, started=started, finish_reason=finish_reason, content=content if isinstance(content, str) else raw_http, usage=usage if isinstance(usage, Mapping) else {}, request_sha256=request_sha256)
281
+ if finish_reason != "stop" or not isinstance(content, str):
282
+ raise RecallPlannerResponseError("recall planner response did not finish with a JSON string", response_content=raw_http, request_metadata=metadata)
283
+ try:
284
+ plan = validate_recall_plan(json.loads(content))
285
+ except (json.JSONDecodeError, RecallPlannerError) as exc:
286
+ raise RecallPlannerResponseError(f"recall planner returned invalid RecallPlan: {exc}", response_content=content, request_metadata=metadata) from exc
287
+ try:
288
+ validate_plan_availability(plan, available_layers)
289
+ except RecallPlannerError as exc:
290
+ raise RecallPlannerResponseError(
291
+ f"recall planner selected an unavailable layer: {exc}",
292
+ response_content=content,
293
+ request_metadata=metadata,
294
+ ) from exc
295
+ return plan, metadata
296
+
297
+
298
+ def _source_parents(item: Mapping[str, Any]) -> list[dict[str, Any]]:
299
+ raw_values: list[Any] = []
300
+ if isinstance(item.get("source_parent"), Mapping):
301
+ raw_values.append(item["source_parent"])
302
+ if isinstance(item.get("source_parents"), Sequence) and not isinstance(item.get("source_parents"), (str, bytes)):
303
+ raw_values.extend(item["source_parents"])
304
+ parents: list[dict[str, Any]] = []
305
+ seen: set[tuple[int, int]] = set()
306
+ for raw in raw_values:
307
+ if not isinstance(raw, Mapping):
308
+ raise RecallPlannerError("slow capsule source parent must be an object")
309
+ try:
310
+ location = (int(raw["session_index"]), int(raw["parent_chunk_index"]))
311
+ except (KeyError, TypeError, ValueError) as exc:
312
+ raise RecallPlannerError("slow capsule source parent lacks integer session/parent coordinates") from exc
313
+ if location in seen:
314
+ continue
315
+ seen.add(location)
316
+ parents.append(dict(raw))
317
+ return parents
318
+
319
+
320
+ def _group_by_slot(items: Sequence[Mapping[str, Any]], *, layer: str, require_parent: bool) -> dict[str, list[Mapping[str, Any]]]:
321
+ grouped: dict[str, list[Mapping[str, Any]]] = {}
322
+ for item in items:
323
+ if not isinstance(item, Mapping):
324
+ raise RecallPlannerError(f"{layer} candidate must be a mapping")
325
+ slot = _text(item.get("canonical_slot"))
326
+ if not slot:
327
+ raise RecallPlannerError(f"{layer} candidate lacks canonical_slot")
328
+ if require_parent and not _source_parents(item):
329
+ raise RecallPlannerError(f"slow capsule {slot!r} lacks source_parent mapping")
330
+ grouped.setdefault(slot, []).append(item)
331
+ return grouped
332
+
333
+
334
+ def apply_recall_plan(plan: Mapping[str, Any], fast_candidates: Sequence[Mapping[str, Any]], slow_capsules: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]:
335
+ """Compose evidence units by role, without combining or recalculating layer scores."""
336
+ normalized = validate_recall_plan(plan)
337
+ fast = _group_by_slot(fast_candidates, layer="fast", require_parent=False)
338
+ slow = _group_by_slot(slow_capsules, layer="slow", require_parent=True)
339
+ mode = normalized["mode"]
340
+ units: list[dict[str, Any]] = []
341
+ if mode == "FAST_ONLY":
342
+ return [
343
+ {"unit_type": "fast_primary", "canonical_slot": slot, "primary_layer": "fast", "fast_candidate": candidate}
344
+ for slot, candidates in fast.items()
345
+ for candidate in candidates
346
+ ]
347
+ if mode == "SLOW_ONLY":
348
+ return [
349
+ {"unit_type": "slow_primary", "canonical_slot": slot, "primary_layer": "slow", "slow_capsule": capsule}
350
+ for slot, capsules in slow.items()
351
+ for capsule in capsules
352
+ ]
353
+ if mode == "SLOW_WITH_FAST_OVERRIDE":
354
+ for slot, capsules in slow.items():
355
+ for capsule in capsules:
356
+ units.append(
357
+ {
358
+ "unit_type": "slow_primary_with_fast_override",
359
+ "canonical_slot": slot,
360
+ "primary_layer": "slow",
361
+ "slow_capsule": capsule,
362
+ "fast_overrides": list(fast.get(slot, [])),
363
+ }
364
+ )
365
+ return units
366
+ if mode == "FAST_WITH_SLOW_CONTEXT":
367
+ all_slow_context = [capsule for capsules in slow.values() for capsule in capsules]
368
+ for slot, candidates in fast.items():
369
+ for candidate in candidates:
370
+ units.append(
371
+ {
372
+ "unit_type": "fast_primary_with_slow_context",
373
+ "canonical_slot": slot,
374
+ "primary_layer": "fast",
375
+ "fast_candidate": candidate,
376
+ "slow_context": list(all_slow_context),
377
+ }
378
+ )
379
+ return units
380
+ for slot in fast.keys() & slow.keys():
381
+ units.append(
382
+ {
383
+ "unit_type": "conflict_group",
384
+ "canonical_slot": slot,
385
+ "primary_layer": "fast",
386
+ "fast_candidates": list(fast[slot]),
387
+ "slow_capsules": list(slow[slot]),
388
+ }
389
+ )
390
+ if not units and fast and slow:
391
+ ranked_fast = [candidate for candidates in fast.values() for candidate in candidates]
392
+ ranked_slow = [capsule for capsules in slow.values() for capsule in capsules]
393
+ pair_count = max(len(ranked_fast), len(ranked_slow))
394
+ for index in range(pair_count):
395
+ units.append(
396
+ {
397
+ "unit_type": "conflict_group",
398
+ "canonical_slot": f"semantic_conflict.{index}",
399
+ "primary_layer": "fast",
400
+ "fast_candidates": [ranked_fast[min(index, len(ranked_fast) - 1)]],
401
+ "slow_capsules": [ranked_slow[min(index, len(ranked_slow) - 1)]],
402
+ }
403
+ )
404
+ if not units:
405
+ raise RecallPlannerError("CONFLICT_COMPARE produced no semantic conflict pairs")
406
+ return units
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_reranker.py ADDED
@@ -0,0 +1,1014 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import hashlib
6
+ import json
7
+ import os
8
+ import random
9
+ import time
10
+ from collections import Counter, defaultdict
11
+ from pathlib import Path
12
+ from typing import Any, Mapping, Sequence
13
+
14
+ import torch
15
+ import torch.nn as nn
16
+ import torch.nn.functional as F
17
+
18
+ from tmcra_v3_schema import (
19
+ CHANNEL_NAMES,
20
+ SCHEMA_VERSION,
21
+ channel_vector,
22
+ clean_text,
23
+ read_jsonl,
24
+ validate_sample,
25
+ write_jsonl,
26
+ )
27
+ from tmcra_model_assets import load_optional_hf_model_manifest
28
+
29
+
30
+ def pair_key(sample: Mapping[str, Any], candidate: Mapping[str, Any]) -> str:
31
+ payload = "\0".join(
32
+ (
33
+ clean_text(sample.get("question_id")),
34
+ clean_text(candidate.get("candidate_id")),
35
+ clean_text(sample.get("query_text")),
36
+ clean_text(candidate.get("text")),
37
+ )
38
+ )
39
+ return hashlib.sha256(payload.encode("utf-8")).hexdigest()
40
+
41
+
42
+ def stable_int(value: str) -> int:
43
+ return int.from_bytes(hashlib.blake2b(value.encode("utf-8"), digest_size=8).digest(), "big")
44
+
45
+
46
+ def atomic_torch_save(payload: Mapping[str, Any], target: Path) -> None:
47
+ temporary = target.with_suffix(target.suffix + ".tmp")
48
+ torch.save(dict(payload), temporary)
49
+ os.replace(temporary, target)
50
+
51
+
52
+ def command_precompute_cross(args: argparse.Namespace) -> None:
53
+ if not torch.cuda.is_available() and not args.cpu:
54
+ raise RuntimeError("CUDA is unavailable; pass --cpu explicitly")
55
+ from transformers import AutoModelForSequenceClassification, AutoTokenizer
56
+
57
+ sample_paths = [Path(value) for value in args.samples]
58
+ rows = [row for path in sample_paths for row in read_jsonl(path)]
59
+ for row in rows:
60
+ validate_sample(
61
+ row,
62
+ require_positive=False,
63
+ allowed_splits=("full_eval", "train", "holdout", "aux_train", "aux_dev"),
64
+ )
65
+ model_dir = Path(args.model).resolve()
66
+ model_manifest = load_optional_hf_model_manifest(model_dir)
67
+ out_dir = Path(args.out_dir)
68
+ out_dir.mkdir(parents=True, exist_ok=True)
69
+ shard_dir = out_dir / "shards"
70
+ shard_dir.mkdir(parents=True, exist_ok=True)
71
+ items: list[dict[str, Any]] = []
72
+ seen: set[str] = set()
73
+ raw_pair_count = 0
74
+ duplicate_pair_count = 0
75
+ for row in rows:
76
+ for candidate in row["candidates"]:
77
+ raw_pair_count += 1
78
+ key = pair_key(row, candidate)
79
+ if key in seen:
80
+ duplicate_pair_count += 1
81
+ continue
82
+ seen.add(key)
83
+ items.append(
84
+ {
85
+ "key": key,
86
+ "question_id": row["question_id"],
87
+ "candidate_id": candidate["candidate_id"],
88
+ "query_text": row["query_text"],
89
+ "candidate_text": candidate["text"],
90
+ }
91
+ )
92
+ write_jsonl(
93
+ out_dir / "items.jsonl",
94
+ [
95
+ {
96
+ "key": item["key"],
97
+ "question_id": item["question_id"],
98
+ "candidate_id": item["candidate_id"],
99
+ "query_len": len(item["query_text"]),
100
+ "candidate_len": len(item["candidate_text"]),
101
+ }
102
+ for item in items
103
+ ],
104
+ )
105
+ device = torch.device("cpu" if args.cpu else "cuda")
106
+ tokenizer = AutoTokenizer.from_pretrained(
107
+ str(model_dir),
108
+ local_files_only=True,
109
+ trust_remote_code=False,
110
+ )
111
+ model = AutoModelForSequenceClassification.from_pretrained(
112
+ str(model_dir),
113
+ local_files_only=True,
114
+ trust_remote_code=False,
115
+ torch_dtype=torch.float32 if args.cpu else torch.float16,
116
+ ).to(device)
117
+ model.eval()
118
+ hidden_size = int(getattr(model.config, "hidden_size", 0) or 0)
119
+ if hidden_size <= 0:
120
+ raise RuntimeError("reranker config.hidden_size is missing")
121
+ started = time.time()
122
+ shard_paths: list[Path] = []
123
+ for start in range(0, len(items), args.shard_size):
124
+ stop = min(len(items), start + args.shard_size)
125
+ shard_index = start // args.shard_size
126
+ shard_path = shard_dir / f"shard_{shard_index:05d}.pt"
127
+ expected_keys = [item["key"] for item in items[start:stop]]
128
+ if shard_path.exists():
129
+ cached = torch.load(shard_path, map_location="cpu", weights_only=True)
130
+ if list(cached.get("keys") or []) != expected_keys:
131
+ raise RuntimeError(f"cross-cache resume key mismatch: {shard_path}")
132
+ representations = cached.get("representations")
133
+ logits = cached.get("semantic_logits")
134
+ if representations is None or logits is None:
135
+ raise RuntimeError(f"cross-cache resume payload invalid: {shard_path}")
136
+ if tuple(representations.shape) != (len(expected_keys), hidden_size):
137
+ raise RuntimeError(f"cross-cache resume representation shape mismatch: {shard_path}")
138
+ print(json.dumps({"status": "resume", "shard": shard_index, "done": stop, "total": len(items)}), flush=True)
139
+ shard_paths.append(shard_path)
140
+ continue
141
+ rep_chunks: list[torch.Tensor] = []
142
+ logit_chunks: list[torch.Tensor] = []
143
+ shard_items = items[start:stop]
144
+ with torch.inference_mode():
145
+ for batch_start in range(0, len(shard_items), args.batch_size):
146
+ batch = shard_items[batch_start : batch_start + args.batch_size]
147
+ encoded = tokenizer(
148
+ [item["query_text"] for item in batch],
149
+ [item["candidate_text"] for item in batch],
150
+ padding=True,
151
+ truncation=False,
152
+ return_tensors="pt",
153
+ )
154
+ token_lengths = encoded["attention_mask"].sum(dim=1)
155
+ longest = int(token_lengths.max())
156
+ if longest > args.max_length:
157
+ offending = int(torch.argmax(token_lengths))
158
+ raise RuntimeError(
159
+ f"cross-encoder pair has {longest} tokens, exceeding strict max_length={args.max_length}; "
160
+ f"qid={batch[offending]['question_id']} candidate={batch[offending]['candidate_id']}"
161
+ )
162
+ encoded = {key: value.to(device) for key, value in encoded.items()}
163
+ output = model(**encoded, output_hidden_states=True, return_dict=True)
164
+ if output.hidden_states is None:
165
+ raise RuntimeError("reranker did not return hidden states")
166
+ representation = output.hidden_states[-1][:, 0].detach().cpu().float()
167
+ semantic = output.logits.detach().cpu().float()
168
+ if semantic.ndim == 2 and semantic.shape[1] == 1:
169
+ semantic = semantic[:, 0]
170
+ elif semantic.ndim != 1:
171
+ raise RuntimeError(f"expected one reranker logit per pair, got {tuple(semantic.shape)}")
172
+ rep_chunks.append(representation.to(torch.float16))
173
+ logit_chunks.append(semantic)
174
+ representations = torch.cat(rep_chunks, dim=0).contiguous()
175
+ logits = torch.cat(logit_chunks, dim=0).contiguous()
176
+ atomic_torch_save(
177
+ {"keys": expected_keys, "representations": representations, "semantic_logits": logits},
178
+ shard_path,
179
+ )
180
+ print(
181
+ json.dumps(
182
+ {
183
+ "status": "encoded",
184
+ "shard": shard_index,
185
+ "done": stop,
186
+ "total": len(items),
187
+ "elapsed_sec": round(time.time() - started, 3),
188
+ }
189
+ ),
190
+ flush=True,
191
+ )
192
+ shard_paths.append(shard_path)
193
+ all_keys: list[str] = []
194
+ all_representations: list[torch.Tensor] = []
195
+ all_logits: list[torch.Tensor] = []
196
+ for path in shard_paths:
197
+ payload = torch.load(path, map_location="cpu", weights_only=True)
198
+ all_keys.extend(list(payload["keys"]))
199
+ all_representations.append(payload["representations"])
200
+ all_logits.append(payload["semantic_logits"])
201
+ expected_all_keys = [item["key"] for item in items]
202
+ representations = torch.cat(all_representations, dim=0).contiguous()
203
+ logits = torch.cat(all_logits, dim=0).contiguous()
204
+ if all_keys != expected_all_keys:
205
+ raise RuntimeError("cross-cache final key order mismatch")
206
+ if tuple(representations.shape) != (len(items), hidden_size) or tuple(logits.shape) != (len(items),):
207
+ raise RuntimeError("cross-cache final tensor shape mismatch")
208
+ atomic_torch_save(
209
+ {"keys": all_keys, "representations": representations, "semantic_logits": logits},
210
+ out_dir / "cross_cache.pt",
211
+ )
212
+ manifest = {
213
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
214
+ "schema_version": SCHEMA_VERSION,
215
+ "samples": [str(path.resolve()) for path in sample_paths],
216
+ "sample_count": len(rows),
217
+ "pair_count": len(items),
218
+ "raw_pair_count": raw_pair_count,
219
+ "duplicate_pair_count": duplicate_pair_count,
220
+ "hidden_size": hidden_size,
221
+ "max_length": args.max_length,
222
+ "strict_no_truncation": True,
223
+ "dtype": "float16",
224
+ "model": str(model_dir),
225
+ "model_repo_id": model_manifest.get("repo_id"),
226
+ "model_revision": model_manifest.get("revision"),
227
+ "elapsed_sec": round(time.time() - started, 3),
228
+ }
229
+ (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True), encoding="utf-8")
230
+ print(json.dumps(manifest, indent=2, sort_keys=True))
231
+
232
+
233
+ class CrossRepresentationCache:
234
+ def __init__(self, root: str):
235
+ directory = Path(root)
236
+ manifest_path = directory / "manifest.json"
237
+ cache_path = directory / "cross_cache.pt"
238
+ if not manifest_path.exists() or not cache_path.exists():
239
+ raise FileNotFoundError(f"complete cross cache is required: {directory}")
240
+ self.manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
241
+ payload = torch.load(cache_path, map_location="cpu", weights_only=True)
242
+ keys = list(payload.get("keys") or [])
243
+ representations = payload.get("representations")
244
+ semantic_logits = payload.get("semantic_logits")
245
+ if not keys or representations is None or semantic_logits is None:
246
+ raise RuntimeError(f"invalid cross cache: {cache_path}")
247
+ if len(keys) != len(set(keys)) or len(keys) != representations.shape[0] or len(keys) != semantic_logits.shape[0]:
248
+ raise RuntimeError(f"cross cache row mismatch: {cache_path}")
249
+ self.index = {key: index for index, key in enumerate(keys)}
250
+ self.representations = representations.float().contiguous()
251
+ self.semantic_logits = semantic_logits.float().contiguous()
252
+ self.hidden_size = int(representations.shape[1])
253
+
254
+ def get(self, sample: Mapping[str, Any], candidate: Mapping[str, Any]) -> tuple[torch.Tensor, torch.Tensor]:
255
+ key = pair_key(sample, candidate)
256
+ index = self.index.get(key)
257
+ if index is None:
258
+ raise RuntimeError(
259
+ f"cross cache miss: qid={sample.get('question_id')} candidate={candidate.get('candidate_id')} key={key}"
260
+ )
261
+ return self.representations[index], self.semantic_logits[index]
262
+
263
+
264
+ class ChannelAwareMemoryReranker(nn.Module):
265
+ def __init__(
266
+ self,
267
+ representation_dim: int,
268
+ channel_dim: int,
269
+ hidden_dim: int = 256,
270
+ layers: int = 2,
271
+ channel_mean: torch.Tensor | None = None,
272
+ channel_scale: torch.Tensor | None = None,
273
+ ):
274
+ super().__init__()
275
+ mean = torch.zeros(channel_dim, dtype=torch.float32) if channel_mean is None else channel_mean.float()
276
+ scale = torch.ones(channel_dim, dtype=torch.float32) if channel_scale is None else channel_scale.float()
277
+ if tuple(mean.shape) != (channel_dim,) or tuple(scale.shape) != (channel_dim,):
278
+ raise RuntimeError("channel normalization tensors have the wrong shape")
279
+ self.register_buffer("channel_mean", mean.contiguous())
280
+ self.register_buffer("channel_scale", scale.clamp_min(1e-4).contiguous())
281
+ self.representation_proj = nn.Sequential(
282
+ nn.LayerNorm(representation_dim),
283
+ nn.Linear(representation_dim, hidden_dim),
284
+ nn.GELU(),
285
+ )
286
+ self.semantic_proj = nn.Sequential(nn.Linear(1, hidden_dim), nn.GELU())
287
+ self.channel_proj = nn.Sequential(nn.Linear(channel_dim, hidden_dim), nn.GELU())
288
+ self.channel_gate = nn.Sequential(nn.Linear(hidden_dim * 2, hidden_dim), nn.Sigmoid())
289
+ encoder_layer = nn.TransformerEncoderLayer(
290
+ d_model=hidden_dim,
291
+ nhead=8,
292
+ dim_feedforward=hidden_dim * 3,
293
+ dropout=0.1,
294
+ activation="gelu",
295
+ batch_first=True,
296
+ norm_first=True,
297
+ )
298
+ self.set_encoder = nn.TransformerEncoder(encoder_layer, num_layers=max(1, int(layers)))
299
+ self.delta_head = nn.Sequential(
300
+ nn.LayerNorm(hidden_dim),
301
+ nn.Linear(hidden_dim, hidden_dim),
302
+ nn.GELU(),
303
+ nn.Dropout(0.1),
304
+ nn.Linear(hidden_dim, 1),
305
+ )
306
+ nn.init.zeros_(self.delta_head[-1].weight)
307
+ nn.init.zeros_(self.delta_head[-1].bias)
308
+
309
+ def forward(
310
+ self,
311
+ representations: torch.Tensor,
312
+ semantic_logits: torch.Tensor,
313
+ channels: torch.Tensor,
314
+ mask: torch.Tensor,
315
+ *,
316
+ ablation: str = "full",
317
+ ) -> torch.Tensor:
318
+ if ablation not in {"full", "no_language", "no_channels", "no_graph"}:
319
+ raise RuntimeError(f"unknown ablation: {ablation}")
320
+ reps = representations
321
+ semantic = semantic_logits
322
+ feature_values = (channels - self.channel_mean) / self.channel_scale
323
+ if ablation == "no_language":
324
+ reps = torch.zeros_like(reps)
325
+ semantic = torch.zeros_like(semantic)
326
+ if ablation == "no_channels":
327
+ feature_values = torch.zeros_like(feature_values)
328
+ elif ablation == "no_graph":
329
+ feature_values = feature_values.clone()
330
+ feature_values[..., 2:5] = 0.0
331
+ rep_hidden = self.representation_proj(reps)
332
+ semantic_hidden = self.semantic_proj(semantic.unsqueeze(-1))
333
+ channel_hidden = self.channel_proj(feature_values)
334
+ gate = self.channel_gate(torch.cat((rep_hidden, semantic_hidden), dim=-1))
335
+ hidden = rep_hidden + semantic_hidden + gate * channel_hidden
336
+ hidden = self.set_encoder(hidden, src_key_padding_mask=~mask)
337
+ delta = self.delta_head(hidden).squeeze(-1)
338
+ scores = semantic + delta
339
+ return scores.masked_fill(~mask, -1e4)
340
+
341
+
342
+ def select_training_candidates(sample: Mapping[str, Any], max_candidates: int, epoch: int) -> list[dict[str, Any]]:
343
+ candidates = list(sample["candidates"])
344
+ positives = [candidate for candidate in candidates if candidate["labels"]["relevance"]]
345
+ negatives = [candidate for candidate in candidates if not candidate["labels"]["relevance"]]
346
+ if not positives:
347
+ raise RuntimeError(f"training sample has no positive: {sample.get('question_id')}")
348
+ if max_candidates < 2:
349
+ raise RuntimeError("max_train_candidates must be at least 2")
350
+ positive_budget = min(len(positives), max(1, max_candidates // 2))
351
+ positives.sort(
352
+ key=lambda candidate: (
353
+ -float(candidate["channels"]["graph_rank_rr"]),
354
+ -float(candidate["channels"]["dense_score"]),
355
+ clean_text(candidate["candidate_id"]),
356
+ )
357
+ )
358
+ if len(positives) > positive_budget:
359
+ anchor_count = min(2, positive_budget)
360
+ anchors = positives[:anchor_count]
361
+ rotating = positives[anchor_count:]
362
+ rng = random.Random(stable_int(f"positive-bag:{epoch}:{sample['question_id']}"))
363
+ rng.shuffle(rotating)
364
+ positives = anchors + rotating[: positive_budget - anchor_count]
365
+ negatives.sort(
366
+ key=lambda candidate: (
367
+ -int(bool(candidate["labels"]["hard_negative"])),
368
+ -float(candidate["channels"]["graph_rank_rr"]),
369
+ -float(candidate["channels"]["dense_score"]),
370
+ clean_text(candidate["candidate_id"]),
371
+ )
372
+ )
373
+ selected = positives + negatives[: max(0, max_candidates - len(positives))]
374
+ rng = random.Random(stable_int(f"candidate-shuffle:{epoch}:{sample['question_id']}"))
375
+ rng.shuffle(selected)
376
+ return selected
377
+
378
+
379
+ def make_batch(
380
+ samples: Sequence[Mapping[str, Any]],
381
+ cache: CrossRepresentationCache,
382
+ device: torch.device,
383
+ *,
384
+ training: bool,
385
+ max_candidates: int,
386
+ epoch: int,
387
+ ) -> dict[str, Any]:
388
+ selected_by_sample = [
389
+ select_training_candidates(sample, max_candidates, epoch) if training else list(sample["candidates"])
390
+ for sample in samples
391
+ ]
392
+ max_count = max(len(candidates) for candidates in selected_by_sample)
393
+ batch_size = len(samples)
394
+ representations = torch.zeros((batch_size, max_count, cache.hidden_size), dtype=torch.float32)
395
+ semantic = torch.zeros((batch_size, max_count), dtype=torch.float32)
396
+ channels = torch.zeros((batch_size, max_count, len(CHANNEL_NAMES)), dtype=torch.float32)
397
+ labels = torch.zeros((batch_size, max_count), dtype=torch.float32)
398
+ mask = torch.zeros((batch_size, max_count), dtype=torch.bool)
399
+ candidate_ids: list[list[str]] = []
400
+ sample_weights = torch.tensor(
401
+ [float((sample.get("supervision") or {}).get("training_weight", 1.0) or 1.0) for sample in samples],
402
+ dtype=torch.float32,
403
+ )
404
+ for row_index, (sample, candidates) in enumerate(zip(samples, selected_by_sample)):
405
+ ids: list[str] = []
406
+ for candidate_index, candidate in enumerate(candidates):
407
+ representation, semantic_logit = cache.get(sample, candidate)
408
+ representations[row_index, candidate_index] = representation
409
+ semantic[row_index, candidate_index] = semantic_logit
410
+ channels[row_index, candidate_index] = torch.tensor(channel_vector(candidate), dtype=torch.float32)
411
+ labels[row_index, candidate_index] = float(bool(candidate["labels"]["relevance"]))
412
+ mask[row_index, candidate_index] = True
413
+ ids.append(clean_text(candidate["candidate_id"]))
414
+ candidate_ids.append(ids)
415
+ return {
416
+ "representations": representations.to(device),
417
+ "semantic_logits": semantic.to(device),
418
+ "channels": channels.to(device),
419
+ "labels": labels.to(device),
420
+ "mask": mask.to(device),
421
+ "candidate_ids": candidate_ids,
422
+ "sample_weights": sample_weights.to(device),
423
+ }
424
+
425
+
426
+ def listwise_loss(
427
+ scores: torch.Tensor,
428
+ labels: torch.Tensor,
429
+ mask: torch.Tensor,
430
+ sample_weights: torch.Tensor | None = None,
431
+ ) -> torch.Tensor:
432
+ losses: list[torch.Tensor] = []
433
+ weights: list[torch.Tensor] = []
434
+ if sample_weights is None:
435
+ sample_weights = torch.ones(scores.shape[0], dtype=scores.dtype, device=scores.device)
436
+ for row_scores, row_labels, row_mask, row_weight in zip(scores, labels, mask, sample_weights):
437
+ valid_scores = row_scores[row_mask]
438
+ valid_labels = row_labels[row_mask]
439
+ if not bool((valid_labels > 0.5).any()):
440
+ continue
441
+ log_probabilities = F.log_softmax(valid_scores, dim=0)
442
+ positive_mask = valid_labels > 0.5
443
+ losses.append(-torch.logsumexp(log_probabilities[positive_mask], dim=0))
444
+ weights.append(row_weight)
445
+ if not losses:
446
+ return scores.sum() * 0.0
447
+ loss_tensor = torch.stack(losses)
448
+ weight_tensor = torch.stack(weights).to(dtype=loss_tensor.dtype).clamp_min(1e-6)
449
+ return (loss_tensor * weight_tensor).sum() / weight_tensor.sum()
450
+
451
+
452
+ def rank_metrics_update(target: dict[str, float], scores: torch.Tensor, labels: torch.Tensor, valid: torch.Tensor) -> None:
453
+ valid_scores = scores[valid]
454
+ valid_labels = labels[valid]
455
+ positives = set(torch.nonzero(valid_labels > 0.5, as_tuple=False).squeeze(-1).tolist())
456
+ target["n"] += 1
457
+ if not positives:
458
+ target["candidate_miss"] += 1
459
+ return
460
+ order = torch.argsort(valid_scores, descending=True).tolist()
461
+ rank = next((position + 1 for position, index in enumerate(order) if index in positives), None)
462
+ for cutoff in (1, 3, 5, 10, 20):
463
+ target[f"recall@{cutoff}"] += int(rank is not None and rank <= cutoff)
464
+ target["mrr"] += 0.0 if rank is None else 1.0 / rank
465
+
466
+
467
+ def finalize_metrics(values: Mapping[str, float]) -> dict[str, Any]:
468
+ n = max(1.0, float(values.get("n", 0.0)))
469
+ output: dict[str, Any] = {"n": int(values.get("n", 0.0))}
470
+ for key, value in sorted(values.items()):
471
+ if key == "n":
472
+ continue
473
+ output[key] = round(float(value) / n, 6)
474
+ return output
475
+
476
+
477
+ def evaluate_model(
478
+ model: ChannelAwareMemoryReranker,
479
+ rows: Sequence[Mapping[str, Any]],
480
+ cache: CrossRepresentationCache,
481
+ device: torch.device,
482
+ *,
483
+ batch_size: int,
484
+ ablation: str,
485
+ ) -> dict[str, Any]:
486
+ model.eval()
487
+ totals: dict[str, float] = defaultdict(float)
488
+ with torch.no_grad():
489
+ for start in range(0, len(rows), batch_size):
490
+ samples = rows[start : start + batch_size]
491
+ batch = make_batch(samples, cache, device, training=False, max_candidates=0, epoch=0)
492
+ scores = model(
493
+ batch["representations"],
494
+ batch["semantic_logits"],
495
+ batch["channels"],
496
+ batch["mask"],
497
+ ablation=ablation,
498
+ ).cpu()
499
+ labels = batch["labels"].cpu()
500
+ mask = batch["mask"].cpu()
501
+ for index in range(len(samples)):
502
+ rank_metrics_update(totals, scores[index], labels[index], mask[index])
503
+ return finalize_metrics(totals)
504
+
505
+
506
+ def split_train_dev(rows: Sequence[Mapping[str, Any]], dev_count: int) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
507
+ ordered = sorted(rows, key=lambda row: (stable_int(f"tmcra-v3-dev:{row['question_id']}"), row["question_id"]))
508
+ if dev_count <= 0 or dev_count >= len(ordered):
509
+ raise RuntimeError("dev_count must leave non-empty train and dev sets")
510
+ dev = list(ordered[:dev_count])
511
+ train = list(ordered[dev_count:])
512
+ return train, dev
513
+
514
+
515
+ def compute_channel_statistics(rows: Sequence[Mapping[str, Any]]) -> tuple[torch.Tensor, torch.Tensor]:
516
+ values = [channel_vector(candidate) for row in rows for candidate in row["candidates"]]
517
+ if not values:
518
+ raise RuntimeError("cannot compute channel statistics from an empty candidate set")
519
+ matrix = torch.tensor(values, dtype=torch.float32)
520
+ mean = matrix.mean(dim=0)
521
+ scale = matrix.std(dim=0, unbiased=False)
522
+ scale = torch.where(scale < 1e-4, torch.ones_like(scale), scale)
523
+ return mean, scale
524
+
525
+
526
+ def load_model_checkpoint(path: Path, cache: CrossRepresentationCache, device: torch.device) -> tuple[ChannelAwareMemoryReranker, dict[str, Any]]:
527
+ checkpoint = torch.load(path, map_location=device, weights_only=True)
528
+ if checkpoint.get("schema_version") != SCHEMA_VERSION:
529
+ raise RuntimeError("checkpoint schema version mismatch")
530
+ if tuple(checkpoint.get("channel_names") or ()) != CHANNEL_NAMES:
531
+ raise RuntimeError("checkpoint channel contract mismatch")
532
+ config = checkpoint["model_config"]
533
+ model = ChannelAwareMemoryReranker(
534
+ representation_dim=cache.hidden_size,
535
+ channel_dim=len(CHANNEL_NAMES),
536
+ hidden_dim=int(config["hidden_dim"]),
537
+ layers=int(config["layers"]),
538
+ ).to(device)
539
+ model.load_state_dict(checkpoint["model_state"])
540
+ return model, checkpoint
541
+
542
+
543
+ def command_train(args: argparse.Namespace) -> None:
544
+ rows = read_jsonl(Path(args.train_samples))
545
+ for row in rows:
546
+ validate_sample(row, require_positive=True, allowed_splits=("train", "aux_train"))
547
+ if args.dev_samples:
548
+ train_rows = list(rows)
549
+ dev_rows = read_jsonl(Path(args.dev_samples))
550
+ for row in dev_rows:
551
+ validate_sample(row, require_positive=False, allowed_splits=("train", "full_eval", "aux_dev"))
552
+ else:
553
+ train_rows, dev_rows = split_train_dev(rows, args.dev_count)
554
+ cache = CrossRepresentationCache(args.cross_cache)
555
+ device = torch.device("cpu" if args.cpu else "cuda")
556
+ if device.type == "cuda" and not torch.cuda.is_available():
557
+ raise RuntimeError("CUDA is unavailable; pass --cpu explicitly")
558
+ channel_mean, channel_scale = compute_channel_statistics(train_rows)
559
+ model = ChannelAwareMemoryReranker(
560
+ representation_dim=cache.hidden_size,
561
+ channel_dim=len(CHANNEL_NAMES),
562
+ hidden_dim=args.hidden_dim,
563
+ layers=args.layers,
564
+ channel_mean=channel_mean,
565
+ channel_scale=channel_scale,
566
+ ).to(device)
567
+ if args.init_checkpoint:
568
+ initial = torch.load(Path(args.init_checkpoint), map_location="cpu", weights_only=True)
569
+ if initial.get("schema_version") != SCHEMA_VERSION:
570
+ raise RuntimeError("initial checkpoint schema version mismatch")
571
+ if tuple(initial.get("channel_names") or ()) != CHANNEL_NAMES:
572
+ raise RuntimeError("initial checkpoint channel contract mismatch")
573
+ initial_config = initial.get("model_config") or {}
574
+ if int(initial_config.get("hidden_dim", -1)) != args.hidden_dim or int(initial_config.get("layers", -1)) != args.layers:
575
+ raise RuntimeError("initial checkpoint architecture does not match requested hidden_dim/layers")
576
+ transferable = {
577
+ key: value
578
+ for key, value in initial["model_state"].items()
579
+ if key not in {"channel_mean", "channel_scale"}
580
+ }
581
+ missing, unexpected = model.load_state_dict(transferable, strict=False)
582
+ if set(missing) != {"channel_mean", "channel_scale"} or unexpected:
583
+ raise RuntimeError(f"unexpected initial checkpoint keys: missing={missing} unexpected={unexpected}")
584
+ optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
585
+ rng = random.Random(args.seed)
586
+ out_dir = Path(args.out_dir)
587
+ out_dir.mkdir(parents=True, exist_ok=False)
588
+ best_score: float | None = None
589
+ best_epoch = 0
590
+ stale_epochs = 0
591
+ history: list[dict[str, Any]] = []
592
+ for epoch in range(1, args.epochs + 1):
593
+ model.train()
594
+ rng.shuffle(train_rows)
595
+ losses: list[float] = []
596
+ for start in range(0, len(train_rows), args.batch_size):
597
+ samples = train_rows[start : start + args.batch_size]
598
+ batch = make_batch(
599
+ samples,
600
+ cache,
601
+ device,
602
+ training=True,
603
+ max_candidates=args.max_train_candidates,
604
+ epoch=epoch,
605
+ )
606
+ scores = model(
607
+ batch["representations"],
608
+ batch["semantic_logits"],
609
+ batch["channels"],
610
+ batch["mask"],
611
+ )
612
+ ranking_loss = listwise_loss(scores, batch["labels"], batch["mask"], batch["sample_weights"])
613
+ valid_delta = (scores - batch["semantic_logits"])[batch["mask"]]
614
+ regularization = valid_delta.square().mean() if valid_delta.numel() else scores.sum() * 0.0
615
+ loss = ranking_loss + args.delta_l2 * regularization
616
+ optimizer.zero_grad(set_to_none=True)
617
+ loss.backward()
618
+ nn.utils.clip_grad_norm_(model.parameters(), 1.0)
619
+ optimizer.step()
620
+ losses.append(float(loss.detach().cpu()))
621
+ dev_metrics = evaluate_model(
622
+ model,
623
+ dev_rows,
624
+ cache,
625
+ device,
626
+ batch_size=args.batch_size,
627
+ ablation="full",
628
+ )
629
+ teacher_dev_rows = [
630
+ row
631
+ for row in dev_rows
632
+ if clean_text((row.get("supervision") or {}).get("target_type")) == "teacher_aligned_turn_bag"
633
+ ]
634
+ dev_teacher_metrics = (
635
+ evaluate_model(
636
+ model,
637
+ teacher_dev_rows,
638
+ cache,
639
+ device,
640
+ batch_size=args.batch_size,
641
+ ablation="full",
642
+ )
643
+ if teacher_dev_rows
644
+ else {"n": 0}
645
+ )
646
+ summary = {
647
+ "epoch": epoch,
648
+ "loss": round(sum(losses) / max(1, len(losses)), 6),
649
+ "dev": dev_metrics,
650
+ "dev_teacher_aligned": dev_teacher_metrics,
651
+ }
652
+ history.append(summary)
653
+ print(json.dumps(summary), flush=True)
654
+ score_metrics = dev_teacher_metrics if int(dev_teacher_metrics.get("n", 0)) >= 10 else dev_metrics
655
+ score = (
656
+ float(score_metrics.get("recall@1", 0.0))
657
+ + 0.25 * float(score_metrics.get("mrr", 0.0))
658
+ + 0.05 * float(score_metrics.get("recall@5", 0.0))
659
+ )
660
+ if best_score is None or score > best_score:
661
+ best_score = score
662
+ best_epoch = epoch
663
+ stale_epochs = 0
664
+ atomic_torch_save(
665
+ {
666
+ "schema_version": SCHEMA_VERSION,
667
+ "channel_names": CHANNEL_NAMES,
668
+ "model_config": {"hidden_dim": args.hidden_dim, "layers": args.layers},
669
+ "model_state": model.state_dict(),
670
+ "channel_statistics": {
671
+ "mean": channel_mean.tolist(),
672
+ "scale": channel_scale.tolist(),
673
+ },
674
+ "epoch": epoch,
675
+ "dev_metrics": dev_metrics,
676
+ "dev_teacher_aligned_metrics": dev_teacher_metrics,
677
+ "cross_cache_manifest": cache.manifest,
678
+ "train_args": vars(args),
679
+ },
680
+ out_dir / "tmcra_v3_reranker.pt",
681
+ )
682
+ else:
683
+ stale_epochs += 1
684
+ if args.patience > 0 and stale_epochs >= args.patience:
685
+ print(json.dumps({"status": "early_stop", "epoch": epoch, "best_epoch": best_epoch}), flush=True)
686
+ break
687
+ report = {
688
+ "status": "complete",
689
+ "schema_version": SCHEMA_VERSION,
690
+ "train_count": len(train_rows),
691
+ "dev_count": len(dev_rows),
692
+ "best_epoch": best_epoch,
693
+ "best_score": best_score,
694
+ "init_checkpoint": str(Path(args.init_checkpoint).resolve()) if args.init_checkpoint else "",
695
+ "channel_statistics": {"mean": channel_mean.tolist(), "scale": channel_scale.tolist()},
696
+ "selection_metric_scope": "teacher_aligned_dev" if any(
697
+ int(item.get("dev_teacher_aligned", {}).get("n", 0)) >= 10 for item in history
698
+ ) else "all_dev",
699
+ "checkpoint": str(out_dir / "tmcra_v3_reranker.pt"),
700
+ "history": history,
701
+ }
702
+ (out_dir / "train_report.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
703
+ print(json.dumps(report, indent=2, sort_keys=True))
704
+
705
+
706
+ def direct_scores(sample: Mapping[str, Any], method: str) -> list[float]:
707
+ candidates = list(sample["candidates"])
708
+ if method == "dense":
709
+ return [float(candidate["channels"]["dense_score"]) for candidate in candidates]
710
+ if method == "graph":
711
+ return [float(candidate["channels"]["graph_rank_rr"]) for candidate in candidates]
712
+ raise RuntimeError(method)
713
+
714
+
715
+ def command_eval(args: argparse.Namespace) -> None:
716
+ rows = read_jsonl(Path(args.samples))
717
+ if args.qid_list:
718
+ qids = {clean_text(value) for value in Path(args.qid_list).read_text(encoding="utf-8").splitlines() if clean_text(value)}
719
+ rows = [row for row in rows if clean_text(row.get("question_id")) in qids]
720
+ if args.limit > 0:
721
+ rows = rows[: args.limit]
722
+ if not rows:
723
+ raise RuntimeError("no evaluation rows")
724
+ for row in rows:
725
+ validate_sample(row, require_positive=False)
726
+ cache = CrossRepresentationCache(args.cross_cache)
727
+ device = torch.device("cpu" if args.cpu else "cuda")
728
+ if device.type == "cuda" and not torch.cuda.is_available():
729
+ raise RuntimeError("CUDA is unavailable; pass --cpu explicitly")
730
+ model, checkpoint = load_model_checkpoint(Path(args.checkpoint), cache, device)
731
+ methods = ("dense", "graph", "base_cross", "fusion", "no_language", "no_channels", "no_graph")
732
+ totals = {method: defaultdict(float) for method in methods}
733
+ by_type: dict[str, dict[str, dict[str, float]]] = {
734
+ method: defaultdict(lambda: defaultdict(float)) for method in methods
735
+ }
736
+ by_supervision: dict[str, dict[str, dict[str, float]]] = {
737
+ method: defaultdict(lambda: defaultdict(float)) for method in methods
738
+ }
739
+ predictions: list[dict[str, Any]] = []
740
+ model.eval()
741
+ with torch.no_grad():
742
+ for start in range(0, len(rows), args.batch_size):
743
+ samples = rows[start : start + args.batch_size]
744
+ batch = make_batch(samples, cache, device, training=False, max_candidates=0, epoch=0)
745
+ model_scores = {
746
+ "base_cross": batch["semantic_logits"].cpu(),
747
+ "fusion": model(batch["representations"], batch["semantic_logits"], batch["channels"], batch["mask"], ablation="full").cpu(),
748
+ "no_language": model(batch["representations"], batch["semantic_logits"], batch["channels"], batch["mask"], ablation="no_language").cpu(),
749
+ "no_channels": model(batch["representations"], batch["semantic_logits"], batch["channels"], batch["mask"], ablation="no_channels").cpu(),
750
+ "no_graph": model(batch["representations"], batch["semantic_logits"], batch["channels"], batch["mask"], ablation="no_graph").cpu(),
751
+ }
752
+ labels = batch["labels"].cpu()
753
+ mask = batch["mask"].cpu()
754
+ for row_index, sample in enumerate(samples):
755
+ qtype = clean_text(sample.get("question_type")) or "unknown"
756
+ supervision_type = clean_text((sample.get("supervision") or {}).get("target_type")) or "unspecified"
757
+ valid = mask[row_index]
758
+ score_map: dict[str, torch.Tensor] = {
759
+ "dense": torch.tensor(direct_scores(sample, "dense"), dtype=torch.float32),
760
+ "graph": torch.tensor(direct_scores(sample, "graph"), dtype=torch.float32),
761
+ }
762
+ score_map.update({method: values[row_index][valid] for method, values in model_scores.items()})
763
+ ranks: dict[str, int | None] = {}
764
+ top_ids: dict[str, list[str]] = {}
765
+ valid_labels = labels[row_index][valid]
766
+ positives = set(torch.nonzero(valid_labels > 0.5, as_tuple=False).squeeze(-1).tolist())
767
+ for method in methods:
768
+ scores = score_map[method]
769
+ method_valid = torch.ones_like(scores, dtype=torch.bool)
770
+ rank_metrics_update(totals[method], scores, valid_labels, method_valid)
771
+ rank_metrics_update(by_type[method][qtype], scores, valid_labels, method_valid)
772
+ rank_metrics_update(by_supervision[method][supervision_type], scores, valid_labels, method_valid)
773
+ order = torch.argsort(scores, descending=True).tolist()
774
+ ranks[method] = next((position + 1 for position, index in enumerate(order) if index in positives), None)
775
+ top_ids[method] = [batch["candidate_ids"][row_index][index] for index in order[:10]]
776
+ predictions.append(
777
+ {
778
+ "question_id": sample["question_id"],
779
+ "question_type": qtype,
780
+ "supervision_type": supervision_type,
781
+ "candidate_count": int(valid.sum()),
782
+ "gold_candidate_count": len(positives),
783
+ "ranks": ranks,
784
+ "top10": top_ids,
785
+ }
786
+ )
787
+ out_dir = Path(args.out_dir)
788
+ out_dir.mkdir(parents=True, exist_ok=False)
789
+ write_jsonl(out_dir / "predictions.jsonl", predictions)
790
+ report = {
791
+ "status": "complete",
792
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
793
+ "schema_version": SCHEMA_VERSION,
794
+ "samples": str(Path(args.samples).resolve()),
795
+ "sample_count": len(rows),
796
+ "checkpoint": str(Path(args.checkpoint).resolve()),
797
+ "checkpoint_epoch": checkpoint.get("epoch"),
798
+ "cross_cache": str(Path(args.cross_cache).resolve()),
799
+ "metrics": {method: finalize_metrics(totals[method]) for method in methods},
800
+ "metrics_by_question_type": {
801
+ method: {qtype: finalize_metrics(values) for qtype, values in sorted(by_type[method].items())}
802
+ for method in methods
803
+ },
804
+ "metrics_by_supervision": {
805
+ method: {name: finalize_metrics(values) for name, values in sorted(by_supervision[method].items())}
806
+ for method in methods
807
+ },
808
+ }
809
+ (out_dir / "report.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
810
+ print(json.dumps(report, indent=2, sort_keys=True))
811
+
812
+
813
+ def command_export_evidence(args: argparse.Namespace) -> None:
814
+ rows = read_jsonl(Path(args.samples))
815
+ if args.qid_list:
816
+ requested = [
817
+ clean_text(value)
818
+ for value in Path(args.qid_list).read_text(encoding="utf-8").splitlines()
819
+ if clean_text(value)
820
+ ]
821
+ by_qid = {clean_text(row.get("question_id")): row for row in rows}
822
+ missing = [qid for qid in requested if qid not in by_qid]
823
+ if missing:
824
+ raise RuntimeError(f"requested qids are missing from samples: {missing[:10]}")
825
+ rows = [by_qid[qid] for qid in requested]
826
+ if args.limit > 0:
827
+ rows = rows[: args.limit]
828
+ if not rows:
829
+ raise RuntimeError("no rows for evidence export")
830
+ for row in rows:
831
+ validate_sample(row, require_positive=False)
832
+ cache = CrossRepresentationCache(args.cross_cache)
833
+ device = torch.device("cpu" if args.cpu else "cuda")
834
+ if device.type == "cuda" and not torch.cuda.is_available():
835
+ raise RuntimeError("CUDA is unavailable; pass --cpu explicitly")
836
+ model, checkpoint = load_model_checkpoint(Path(args.checkpoint), cache, device)
837
+ model.eval()
838
+ output_rows: list[dict[str, Any]] = []
839
+ recall = Counter()
840
+ with torch.no_grad():
841
+ for start in range(0, len(rows), args.batch_size):
842
+ samples = rows[start : start + args.batch_size]
843
+ batch = make_batch(samples, cache, device, training=False, max_candidates=0, epoch=0)
844
+ if args.method == "fusion":
845
+ scores = model(
846
+ batch["representations"],
847
+ batch["semantic_logits"],
848
+ batch["channels"],
849
+ batch["mask"],
850
+ ablation="full",
851
+ ).cpu()
852
+ elif args.method == "base_cross":
853
+ scores = batch["semantic_logits"].cpu().masked_fill(~batch["mask"].cpu(), -1e4)
854
+ elif args.method in {"dense", "graph"}:
855
+ channel_name = "dense_score" if args.method == "dense" else "graph_rank_rr"
856
+ scores = torch.full(batch["mask"].shape, -1e4, dtype=torch.float32)
857
+ for row_index, sample in enumerate(samples):
858
+ values = [float(candidate["channels"][channel_name]) for candidate in sample["candidates"]]
859
+ scores[row_index, : len(values)] = torch.tensor(values, dtype=torch.float32)
860
+ else:
861
+ raise RuntimeError(f"unsupported evidence method: {args.method}")
862
+ labels = batch["labels"].cpu()
863
+ mask = batch["mask"].cpu()
864
+ for row_index, sample in enumerate(samples):
865
+ candidates = list(sample["candidates"])
866
+ valid_count = int(mask[row_index].sum())
867
+ order = torch.argsort(scores[row_index, :valid_count], descending=True).tolist()
868
+ selected_indexes: list[int] = []
869
+ parent_counts: Counter[tuple[int, int]] = Counter()
870
+ session_counts: Counter[str] = Counter()
871
+ for candidate_index in order:
872
+ candidate = candidates[candidate_index]
873
+ parent_key = (int(candidate["session_index"]), int(candidate.get("parent_chunk_index", 0)))
874
+ session_id = clean_text(candidate.get("session_id"))
875
+ if args.max_per_parent > 0 and parent_counts[parent_key] >= args.max_per_parent:
876
+ continue
877
+ if args.max_per_session > 0 and session_counts[session_id] >= args.max_per_session:
878
+ continue
879
+ selected_indexes.append(candidate_index)
880
+ parent_counts[parent_key] += 1
881
+ session_counts[session_id] += 1
882
+ if len(selected_indexes) >= args.top_k:
883
+ break
884
+ positives = set(
885
+ torch.nonzero(labels[row_index, :valid_count] > 0.5, as_tuple=False).squeeze(-1).tolist()
886
+ )
887
+ recall["n"] += 1
888
+ recall["hit"] += int(any(index in positives for index in selected_indexes))
889
+ supervision_type = clean_text((sample.get("supervision") or {}).get("target_type")) or "unspecified"
890
+ recall[f"n:{supervision_type}"] += 1
891
+ recall[f"hit:{supervision_type}"] += int(any(index in positives for index in selected_indexes))
892
+ evidence_windows = []
893
+ for rank, candidate_index in enumerate(selected_indexes, start=1):
894
+ candidate = candidates[candidate_index]
895
+ evidence_windows.append(
896
+ {
897
+ "memory_id": candidate["candidate_id"],
898
+ "session_id": candidate["session_id"],
899
+ "session_index": candidate["session_index"],
900
+ "parent_chunk_index": candidate.get("parent_chunk_index", 0),
901
+ "subchunk_index": candidate.get("subchunk_index", 0),
902
+ "score": round(float(scores[row_index, candidate_index]), 6),
903
+ "rank": rank,
904
+ "text": candidate["text"],
905
+ "channels": candidate["channels"],
906
+ }
907
+ )
908
+ output_rows.append(
909
+ {
910
+ "schema_version": SCHEMA_VERSION,
911
+ "question_id": sample["question_id"],
912
+ "question": sample["question"],
913
+ "question_date": sample.get("question_date", ""),
914
+ "question_type": sample.get("question_type", ""),
915
+ "gold_answer": sample.get("gold_answer", ""),
916
+ "answer_session_ids": sample.get("answer_session_ids", []),
917
+ "selected_session_ids": [window["session_id"] for window in evidence_windows],
918
+ "evidence_windows": evidence_windows,
919
+ }
920
+ )
921
+ out_dir = Path(args.out_dir)
922
+ out_dir.mkdir(parents=True, exist_ok=False)
923
+ write_jsonl(out_dir / "evidence_windows.jsonl", output_rows)
924
+ report = {
925
+ "status": "complete",
926
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
927
+ "schema_version": SCHEMA_VERSION,
928
+ "samples": str(Path(args.samples).resolve()),
929
+ "sample_count": len(output_rows),
930
+ "top_k": args.top_k,
931
+ "method": args.method,
932
+ "max_per_parent": args.max_per_parent,
933
+ "max_per_session": args.max_per_session,
934
+ "evidence_recall": round(float(recall["hit"]) / max(1, int(recall["n"])), 6),
935
+ "evidence_recall_by_supervision": {
936
+ name.removeprefix("n:"): round(
937
+ float(recall[f"hit:{name.removeprefix('n:')}"]) / max(1, int(count)), 6
938
+ )
939
+ for name, count in sorted(recall.items())
940
+ if name.startswith("n:")
941
+ },
942
+ "checkpoint": str(Path(args.checkpoint).resolve()),
943
+ "checkpoint_epoch": checkpoint.get("epoch"),
944
+ "output": str(out_dir / "evidence_windows.jsonl"),
945
+ }
946
+ (out_dir / "report.json").write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
947
+ print(json.dumps(report, indent=2, sort_keys=True))
948
+
949
+
950
+ def main() -> int:
951
+ parser = argparse.ArgumentParser()
952
+ sub = parser.add_subparsers(dest="command", required=True)
953
+ precompute = sub.add_parser("precompute-cross")
954
+ precompute.add_argument("--samples", required=True, nargs="+")
955
+ precompute.add_argument("--model", required=True)
956
+ precompute.add_argument("--out-dir", required=True)
957
+ precompute.add_argument("--max-length", type=int, default=1280)
958
+ precompute.add_argument("--batch-size", type=int, default=4)
959
+ precompute.add_argument("--shard-size", type=int, default=500)
960
+ precompute.add_argument("--cpu", action="store_true")
961
+ train = sub.add_parser("train")
962
+ train.add_argument("--train-samples", required=True)
963
+ train.add_argument("--dev-samples", default="")
964
+ train.add_argument("--cross-cache", required=True)
965
+ train.add_argument("--out-dir", required=True)
966
+ train.add_argument("--init-checkpoint", default="")
967
+ train.add_argument("--hidden-dim", type=int, default=256)
968
+ train.add_argument("--layers", type=int, default=2)
969
+ train.add_argument("--batch-size", type=int, default=8)
970
+ train.add_argument("--max-train-candidates", type=int, default=32)
971
+ train.add_argument("--dev-count", type=int, default=40)
972
+ train.add_argument("--epochs", type=int, default=16)
973
+ train.add_argument("--lr", type=float, default=2e-4)
974
+ train.add_argument("--weight-decay", type=float, default=0.01)
975
+ train.add_argument("--delta-l2", type=float, default=0.002)
976
+ train.add_argument("--patience", type=int, default=4)
977
+ train.add_argument("--seed", type=int, default=31)
978
+ train.add_argument("--cpu", action="store_true")
979
+ evaluate = sub.add_parser("eval")
980
+ evaluate.add_argument("--samples", required=True)
981
+ evaluate.add_argument("--cross-cache", required=True)
982
+ evaluate.add_argument("--checkpoint", required=True)
983
+ evaluate.add_argument("--out-dir", required=True)
984
+ evaluate.add_argument("--batch-size", type=int, default=8)
985
+ evaluate.add_argument("--limit", type=int, default=0)
986
+ evaluate.add_argument("--qid-list", default="")
987
+ evaluate.add_argument("--cpu", action="store_true")
988
+ export = sub.add_parser("export-evidence")
989
+ export.add_argument("--samples", required=True)
990
+ export.add_argument("--cross-cache", required=True)
991
+ export.add_argument("--checkpoint", required=True)
992
+ export.add_argument("--out-dir", required=True)
993
+ export.add_argument("--top-k", type=int, default=8)
994
+ export.add_argument("--method", choices=("dense", "graph", "base_cross", "fusion"), default="fusion")
995
+ export.add_argument("--max-per-parent", type=int, default=0)
996
+ export.add_argument("--max-per-session", type=int, default=0)
997
+ export.add_argument("--batch-size", type=int, default=8)
998
+ export.add_argument("--limit", type=int, default=0)
999
+ export.add_argument("--qid-list", default="")
1000
+ export.add_argument("--cpu", action="store_true")
1001
+ args = parser.parse_args()
1002
+ if args.command == "precompute-cross":
1003
+ command_precompute_cross(args)
1004
+ elif args.command == "train":
1005
+ command_train(args)
1006
+ elif args.command == "eval":
1007
+ command_eval(args)
1008
+ else:
1009
+ command_export_evidence(args)
1010
+ return 0
1011
+
1012
+
1013
+ if __name__ == "__main__":
1014
+ raise SystemExit(main())
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_schema.py ADDED
@@ -0,0 +1,168 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import json
5
+ import math
6
+ from pathlib import Path
7
+ from typing import Any, Mapping, Sequence
8
+
9
+
10
+ SCHEMA_VERSION = "tmcra.memory_recall.v3.0"
11
+ CHANNEL_NAMES = (
12
+ "dense_score",
13
+ "dense_rank_rr",
14
+ "graph_rank_rr",
15
+ "graph_selected",
16
+ "graph_final",
17
+ "recency_norm",
18
+ )
19
+ FORBIDDEN_CHANNEL_TOKENS = ("label", "positive", "negative", "hard", "gold", "answer")
20
+
21
+
22
+ def clean_text(value: Any) -> str:
23
+ return " ".join(str(value or "").split())
24
+
25
+
26
+ def read_jsonl(path: Path) -> list[dict[str, Any]]:
27
+ if not path.exists():
28
+ raise FileNotFoundError(path)
29
+ rows: list[dict[str, Any]] = []
30
+ with path.open("r", encoding="utf-8", errors="strict") as handle:
31
+ for line_no, line in enumerate(handle, start=1):
32
+ if not line.strip():
33
+ continue
34
+ try:
35
+ value = json.loads(line)
36
+ except json.JSONDecodeError as exc:
37
+ raise RuntimeError(f"invalid jsonl at {path}:{line_no}: {exc}") from exc
38
+ if not isinstance(value, dict):
39
+ raise RuntimeError(f"jsonl row is not an object at {path}:{line_no}")
40
+ rows.append(value)
41
+ if not rows:
42
+ raise RuntimeError(f"jsonl is empty: {path}")
43
+ return rows
44
+
45
+
46
+ def write_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
47
+ path.parent.mkdir(parents=True, exist_ok=True)
48
+ with path.open("w", encoding="utf-8") as handle:
49
+ for row in rows:
50
+ handle.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")
51
+
52
+
53
+ def channel_vector(candidate: Mapping[str, Any]) -> list[float]:
54
+ channels = candidate.get("channels")
55
+ if not isinstance(channels, Mapping):
56
+ raise RuntimeError("candidate.channels must be an object")
57
+ return [float(channels[name]) for name in CHANNEL_NAMES]
58
+
59
+
60
+ def _require_finite(value: Any, context: str) -> float:
61
+ try:
62
+ number = float(value)
63
+ except (TypeError, ValueError) as exc:
64
+ raise RuntimeError(f"{context} is not numeric: {value!r}") from exc
65
+ if not math.isfinite(number):
66
+ raise RuntimeError(f"{context} is not finite: {number!r}")
67
+ return number
68
+
69
+
70
+ def validate_candidate(candidate: Mapping[str, Any], *, context: str) -> None:
71
+ candidate_id = clean_text(candidate.get("candidate_id"))
72
+ text = clean_text(candidate.get("text"))
73
+ if not candidate_id:
74
+ raise RuntimeError(f"{context}: candidate_id is required")
75
+ if not text:
76
+ raise RuntimeError(f"{context}: text is required")
77
+ channels = candidate.get("channels")
78
+ if not isinstance(channels, Mapping):
79
+ raise RuntimeError(f"{context}: channels must be an object")
80
+ if len(channels) != len(CHANNEL_NAMES) or set(channels) != set(CHANNEL_NAMES):
81
+ raise RuntimeError(
82
+ f"{context}: channel keys must exactly equal {CHANNEL_NAMES}, got {tuple(channels.keys())}"
83
+ )
84
+ for name, value in channels.items():
85
+ lowered = name.lower()
86
+ if any(token in lowered for token in FORBIDDEN_CHANNEL_TOKENS):
87
+ raise RuntimeError(f"{context}: label-derived channel name is forbidden: {name}")
88
+ _require_finite(value, f"{context}.channels.{name}")
89
+ labels = candidate.get("labels")
90
+ if not isinstance(labels, Mapping):
91
+ raise RuntimeError(f"{context}: labels must be an object")
92
+ relevance = bool(labels.get("relevance", False))
93
+ hard_negative = bool(labels.get("hard_negative", False))
94
+ if relevance and hard_negative:
95
+ raise RuntimeError(f"{context}: a positive candidate cannot be a hard negative")
96
+ role = clean_text(labels.get("evidence_role"))
97
+ expected_role = "positive" if relevance else ("hard_negative" if hard_negative else "negative")
98
+ if role != expected_role:
99
+ raise RuntimeError(f"{context}: evidence_role={role!r}, expected {expected_role!r}")
100
+ target_scope = clean_text(labels.get("target_scope"))
101
+ if target_scope and target_scope != "answer_session_bag":
102
+ raise RuntimeError(f"{context}: unsupported target_scope={target_scope!r}")
103
+
104
+
105
+ def validate_sample(
106
+ sample: Mapping[str, Any],
107
+ *,
108
+ require_positive: bool,
109
+ allowed_splits: Sequence[str] = ("train", "holdout", "full_eval"),
110
+ ) -> None:
111
+ if clean_text(sample.get("schema_version")) != SCHEMA_VERSION:
112
+ raise RuntimeError(f"unsupported schema_version: {sample.get('schema_version')!r}")
113
+ qid = clean_text(sample.get("question_id"))
114
+ query = clean_text(sample.get("query_text"))
115
+ split = clean_text(sample.get("split"))
116
+ if not qid or not query:
117
+ raise RuntimeError("question_id and query_text are required")
118
+ if split not in set(allowed_splits):
119
+ raise RuntimeError(f"invalid split for {qid}: {split!r}")
120
+ supervision = sample.get("supervision")
121
+ if supervision is not None:
122
+ if not isinstance(supervision, Mapping):
123
+ raise RuntimeError(f"{qid}: supervision must be an object")
124
+ if clean_text(supervision.get("target_type")) not in {
125
+ "multi_instance_answer_session_bag",
126
+ "teacher_aligned_turn_bag",
127
+ }:
128
+ raise RuntimeError(f"{qid}: unsupported supervision target_type")
129
+ if clean_text(supervision.get("loss")) != "negative_log_positive_probability_mass":
130
+ raise RuntimeError(f"{qid}: unsupported supervision loss")
131
+ weight = _require_finite(supervision.get("training_weight", 1.0), f"{qid}.supervision.training_weight")
132
+ if weight <= 0.0 or weight > 1.0:
133
+ raise RuntimeError(f"{qid}: supervision training_weight must be in (0, 1]")
134
+ candidates = sample.get("candidates")
135
+ if not isinstance(candidates, list) or not candidates:
136
+ raise RuntimeError(f"{qid}: candidates must be a non-empty list")
137
+ seen: set[str] = set()
138
+ positive_count = 0
139
+ for index, candidate in enumerate(candidates):
140
+ if not isinstance(candidate, Mapping):
141
+ raise RuntimeError(f"{qid}: candidate {index} is not an object")
142
+ validate_candidate(candidate, context=f"{qid}.candidates[{index}]")
143
+ candidate_id = clean_text(candidate.get("candidate_id"))
144
+ if candidate_id in seen:
145
+ raise RuntimeError(f"{qid}: duplicate candidate_id: {candidate_id}")
146
+ seen.add(candidate_id)
147
+ positive_count += int(bool(candidate["labels"].get("relevance", False)))
148
+ if require_positive and positive_count == 0:
149
+ raise RuntimeError(f"{qid}: no positive candidate")
150
+
151
+
152
+ def validate_split_isolation(train_rows: Sequence[Mapping[str, Any]], holdout_rows: Sequence[Mapping[str, Any]]) -> dict[str, int]:
153
+ train_qids = {clean_text(row.get("question_id")) for row in train_rows}
154
+ holdout_qids = {clean_text(row.get("question_id")) for row in holdout_rows}
155
+ train_queries = {clean_text(row.get("query_text")) for row in train_rows}
156
+ holdout_queries = {clean_text(row.get("query_text")) for row in holdout_rows}
157
+ qid_overlap = train_qids & holdout_qids
158
+ query_overlap = train_queries & holdout_queries
159
+ if qid_overlap:
160
+ raise RuntimeError(f"train/holdout qid overlap: {sorted(qid_overlap)[:10]}")
161
+ if query_overlap:
162
+ raise RuntimeError(f"train/holdout query overlap: {sorted(query_overlap)[:3]}")
163
+ return {
164
+ "train_qids": len(train_qids),
165
+ "holdout_qids": len(holdout_qids),
166
+ "qid_overlap": 0,
167
+ "query_text_overlap": 0,
168
+ }
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_slow_graph.py ADDED
@@ -0,0 +1,2138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Strict, append-only slow-memory graph controller."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import argparse
7
+ import hashlib
8
+ import importlib
9
+ import json
10
+ import os
11
+ import sqlite3
12
+ import sys
13
+ import threading
14
+ import time
15
+ import urllib.error
16
+ import urllib.request
17
+ import uuid
18
+ from contextlib import contextmanager
19
+ from dataclasses import dataclass
20
+ from pathlib import Path
21
+ from typing import Any, Iterable, Mapping, Protocol
22
+
23
+ LEAF_VARIANT = "product_semantic_memory"
24
+ CAPSULE_VARIANT = "slow_memory_capsule"
25
+ PATCH_ACTIONS = frozenset(
26
+ {
27
+ "create",
28
+ "revise",
29
+ "challenge",
30
+ "resolve_challenge",
31
+ "retire",
32
+ "noop",
33
+ }
34
+ )
35
+ EDGE_TYPES = frozenset(
36
+ {
37
+ "supports",
38
+ "contradicts",
39
+ "derived_from",
40
+ "supersedes",
41
+ "challenges",
42
+ "invalidates",
43
+ }
44
+ )
45
+ SCHEMA_VERSION = "slow-graph-patch/v4"
46
+ SLOW_EDGE_SOURCE = "slow_graph_control_plane"
47
+ DEFAULT_CLAIM_LEASE_SECONDS = 15 * 60
48
+
49
+
50
+ class SlowGraphError(RuntimeError):
51
+ pass
52
+
53
+
54
+ class PatchValidationError(SlowGraphError):
55
+ pass
56
+
57
+
58
+ class EvidencePolicyError(SlowGraphError):
59
+ pass
60
+
61
+
62
+ class StaleRevisionError(SlowGraphError):
63
+ pass
64
+
65
+
66
+ class AuditError(SlowGraphError):
67
+ pass
68
+
69
+
70
+ class DeepSeekCallError(SlowGraphError):
71
+ def __init__(self, message: str, *, retryable: bool) -> None:
72
+ super().__init__(message)
73
+ self.retryable = retryable
74
+
75
+
76
+ @dataclass(frozen=True)
77
+ class JobClaim:
78
+ job_id: str
79
+ attempt_id: str
80
+ token: str
81
+ owner: str
82
+
83
+
84
+ def _now() -> int:
85
+ return int(time.time())
86
+
87
+
88
+ def _clean(value: Any) -> str:
89
+ return str(value).strip() if value is not None else ""
90
+
91
+
92
+ def _json(value: Any) -> str:
93
+ return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
94
+
95
+
96
+ def _digest(value: Any) -> str:
97
+ return hashlib.sha256(_json(value).encode("utf-8")).hexdigest()
98
+
99
+
100
+ def _strict_json(value: str, *, label: str, expected: type) -> Any:
101
+ try:
102
+ result = json.loads(value)
103
+ except (TypeError, json.JSONDecodeError) as exc:
104
+ raise SlowGraphError(f"invalid {label} JSON") from exc
105
+ if not isinstance(result, expected):
106
+ raise SlowGraphError(f"{label} JSON must be {expected.__name__}")
107
+ return result
108
+
109
+
110
+ def _required_text(value: Any, label: str) -> str:
111
+ text = _clean(value)
112
+ if not text:
113
+ raise PatchValidationError(f"{label} is required")
114
+ return text
115
+
116
+
117
+ def load_graph_schema(repo: str | Path) -> tuple[type[Any], type[Any]]:
118
+ """Load the real graph record and edge types from an explicit TMCRA repo."""
119
+ root = Path(repo).resolve()
120
+ candidates = [root, root / "tmcra_code"]
121
+ package_roots = [
122
+ candidate
123
+ for candidate in candidates
124
+ if (candidate / "experiments" / "replacement" / "memory_graph.py").is_file()
125
+ ]
126
+ if len(package_roots) != 1:
127
+ raise SlowGraphError(
128
+ "--repo must resolve exactly one experiments/replacement/memory_graph.py "
129
+ f"at the repo root or repo/tmcra_code: {root}"
130
+ )
131
+ package_root = package_roots[0]
132
+ module_path = package_root / "experiments" / "replacement" / "memory_graph.py"
133
+ if str(package_root) not in sys.path:
134
+ sys.path.insert(0, str(package_root))
135
+ importlib.invalidate_caches()
136
+ try:
137
+ module = importlib.import_module("experiments.replacement.memory_graph")
138
+ return module.SessionMemoryRecordV2, module.SessionMemoryEdgeV2
139
+ except (ImportError, AttributeError) as exc:
140
+ raise SlowGraphError(
141
+ "unable to import real SessionMemoryRecordV2/SessionMemoryEdgeV2"
142
+ ) from exc
143
+
144
+
145
+ class PatchManager(Protocol):
146
+ model_config: Mapping[str, Any]
147
+ prompt_hash: str
148
+ last_call_metadata: Mapping[str, Any]
149
+
150
+ def propose(
151
+ self, region: Mapping[str, Any], capsules: list[Mapping[str, Any]]
152
+ ) -> Mapping[str, Any]: ...
153
+
154
+
155
+ @dataclass(frozen=True)
156
+ class DeepSeekProConfig:
157
+ base_url: str
158
+ key_pool: tuple[str, ...]
159
+ max_tokens: int
160
+ model: str = "deepseek-v4-pro"
161
+
162
+ @classmethod
163
+ def from_env(cls) -> "DeepSeekProConfig":
164
+ base_url = _clean(os.getenv("TMCRA_DEEPSEEK_PRO_BASE_URL"))
165
+ keys = tuple(
166
+ item.strip()
167
+ for item in _clean(os.getenv("TMCRA_DEEPSEEK_PRO_KEY_POOL")).split(",")
168
+ if item.strip()
169
+ )
170
+ try:
171
+ max_tokens = int(_clean(os.getenv("TMCRA_DEEPSEEK_PRO_MAX_TOKENS")))
172
+ except ValueError as exc:
173
+ raise SlowGraphError(
174
+ "TMCRA_DEEPSEEK_PRO_MAX_TOKENS must be an integer"
175
+ ) from exc
176
+ if not base_url or not keys or max_tokens <= 0:
177
+ raise SlowGraphError(
178
+ "DeepSeek requires BASE_URL, KEY_POOL, and positive MAX_TOKENS"
179
+ )
180
+ return cls(base_url.rstrip("/"), keys, max_tokens)
181
+
182
+
183
+ class DeepSeekProGraphPatchManager:
184
+ """The only production patch manager: no fallback and no hidden defaults."""
185
+
186
+ def __init__(self, config: DeepSeekProConfig) -> None:
187
+ if config.model != "deepseek-v4-pro":
188
+ raise SlowGraphError("DeepSeek model must be deepseek-v4-pro")
189
+ self.config = config
190
+ self._key_index = 0
191
+ self.model_config = {
192
+ "model": config.model,
193
+ "temperature": 0,
194
+ "thinking": "disabled",
195
+ "max_tokens": config.max_tokens,
196
+ }
197
+ self.prompt_hash = _digest(self._messages("<region>", []))
198
+ self.last_call_metadata: Mapping[str, Any] = {}
199
+
200
+ def _messages(self, region: Any, capsules: Any) -> list[dict[str, str]]:
201
+ schema = {
202
+ "operations": [
203
+ {
204
+ "action": "create|revise|challenge|resolve_challenge|retire|noop",
205
+ "capsule_id": "required for non-create operations; omit for create because controller derives it from region_key",
206
+ "base_revision": "required for mutations of an existing capsule",
207
+ "claims": [
208
+ {
209
+ "canonical_slot": "stable semantic slot shared by corrections of the same durable fact",
210
+ "text": "string",
211
+ "support": ["fast id"],
212
+ "counterevidence": ["fast id"],
213
+ }
214
+ ],
215
+ }
216
+ ]
217
+ }
218
+ return [
219
+ {
220
+ "role": "system",
221
+ "content": (
222
+ "You manage only the slow durable layer above immutable fast evidence leaves. "
223
+ "Return exactly one JSON GraphPatch and no prose. Create or revise a capsule only for "
224
+ "a durable preference, identity fact, routine, relationship, standing constraint, or stable "
225
+ "long-running state. A single explicit durable assertion is sufficient; never require an "
226
+ "arbitrary repetition count. Transient events, one-off tasks, quoted third-party claims, "
227
+ "and evidence that does not establish durable memory must produce an empty operations list. "
228
+ "Use challenge when new fast evidence conflicts but does not resolve the old claim; use "
229
+ "resolve_challenge when the supplied evidence resolves an existing challenge; use retire only "
230
+ "when evidence establishes that a capsule is no longer applicable. Topology merge and split "
231
+ "are not part of this regional V1 controller. canonical_slot must remain identical across corrections "
232
+ "of the same fact. Every claim canonical_slot must exactly copy canonical_slot_key from at "
233
+ "least one fast evidence leaf cited by that claim. Do not invent evidence IDs, capsule IDs "
234
+ "for create, confidence, source "
235
+ "parents, claim_id, or fields outside the supplied schema. The controller assigns claim_id. "
236
+ "Always emit support and counterevidence arrays, using [] for the empty side; the controller "
237
+ "also accepts an omitted empty side as a transport-level normalization. Every non-noop claim "
238
+ "must cite supplied fast leaf evidence. record_state and slow_graph_evidence_role are "
239
+ "controller-owned. Never promote historical_noncurrent evidence as a current claim when "
240
+ "current_authoritative evidence exists for the same canonical_slot; use noncurrent evidence "
241
+ "only as history or counterevidence. Each region has exactly one controller-owned capsule: "
242
+ "emit create only when "
243
+ "capsules is empty; otherwise mutate the supplied capsule_id or emit an empty operations list. "
244
+ "Return at most one operation. Empty operations is the canonical region-level noop."
245
+ ),
246
+ },
247
+ {
248
+ "role": "user",
249
+ "content": _json(
250
+ {
251
+ "schema_version": SCHEMA_VERSION,
252
+ "schema": schema,
253
+ "region": region,
254
+ "capsules": capsules,
255
+ }
256
+ ),
257
+ },
258
+ ]
259
+
260
+ def propose(
261
+ self, region: Mapping[str, Any], capsules: list[Mapping[str, Any]]
262
+ ) -> Mapping[str, Any]:
263
+ key_index = self._key_index % len(self.config.key_pool)
264
+ key = self.config.key_pool[key_index]
265
+ self._key_index += 1
266
+ body = {
267
+ "model": "deepseek-v4-pro",
268
+ "temperature": 0,
269
+ "max_tokens": self.config.max_tokens,
270
+ "thinking": {"type": "disabled"},
271
+ "enable_thinking": False,
272
+ "response_format": {"type": "json_object"},
273
+ "messages": self._messages(region, capsules),
274
+ }
275
+ request = urllib.request.Request(
276
+ self.config.base_url + "/chat/completions",
277
+ data=_json(body).encode("utf-8"),
278
+ headers={
279
+ "Authorization": "Bearer " + key,
280
+ "Content-Type": "application/json",
281
+ },
282
+ method="POST",
283
+ )
284
+ physical_call_id = "dsc_" + uuid.uuid4().hex
285
+ started_at = _now()
286
+ self.last_call_metadata = {
287
+ "request": body,
288
+ "key_index": key_index,
289
+ "started_at": started_at,
290
+ "physical_call_id": physical_call_id,
291
+ "status": "started",
292
+ }
293
+ try:
294
+ with urllib.request.urlopen(request, timeout=90) as response: # nosec B310
295
+ status = response.getcode()
296
+ raw_text = response.read().decode("utf-8")
297
+ except urllib.error.HTTPError as exc:
298
+ try:
299
+ detail = exc.read().decode("utf-8", "replace")
300
+ except (AttributeError, OSError):
301
+ detail = ""
302
+ self.last_call_metadata = {
303
+ **self.last_call_metadata,
304
+ "status": "http_error",
305
+ "http_status": exc.code,
306
+ "error": detail,
307
+ "completed_at": _now(),
308
+ }
309
+ raise DeepSeekCallError(
310
+ f"DeepSeek HTTP {exc.code}: {detail}",
311
+ retryable=exc.code == 429 or exc.code >= 500,
312
+ ) from exc
313
+ except (urllib.error.URLError, TimeoutError, OSError) as exc:
314
+ self.last_call_metadata = {
315
+ **self.last_call_metadata,
316
+ "status": "request_error",
317
+ "error": f"{exc.__class__.__name__}: {exc}",
318
+ "completed_at": _now(),
319
+ }
320
+ raise DeepSeekCallError(
321
+ f"DeepSeek transport failure: {exc}", retryable=True
322
+ ) from exc
323
+ if status < 200 or status >= 300:
324
+ self.last_call_metadata = {
325
+ **self.last_call_metadata,
326
+ "status": "unexpected_http_status",
327
+ "http_status": status,
328
+ "raw_response": raw_text,
329
+ "completed_at": _now(),
330
+ }
331
+ raise DeepSeekCallError(
332
+ f"DeepSeek returned HTTP {status}", retryable=status >= 500
333
+ )
334
+ self.last_call_metadata = {
335
+ **self.last_call_metadata,
336
+ "status": "response_received",
337
+ "http_status": status,
338
+ "raw_response": raw_text,
339
+ "completed_at": _now(),
340
+ }
341
+ raw = _strict_json(raw_text, label="DeepSeek response", expected=dict)
342
+ choices = raw.get("choices")
343
+ usage = raw.get("usage")
344
+ if (
345
+ not isinstance(choices, list)
346
+ or len(choices) != 1
347
+ or not isinstance(usage, Mapping)
348
+ ):
349
+ raise DeepSeekCallError(
350
+ "DeepSeek response missing strict choices/usage", retryable=False
351
+ )
352
+ choice = choices[0]
353
+ if not isinstance(choice, Mapping) or not _clean(choice.get("finish_reason")):
354
+ raise DeepSeekCallError(
355
+ "DeepSeek response missing finish_reason", retryable=False
356
+ )
357
+ message = choice.get("message")
358
+ if not isinstance(message, Mapping) or not _clean(message.get("content")):
359
+ raise DeepSeekCallError(
360
+ "DeepSeek response missing choices[0].message.content", retryable=False
361
+ )
362
+ finish_reason = _clean(choice["finish_reason"])
363
+ self.last_call_metadata = {
364
+ **self.last_call_metadata,
365
+ "finish_reason": finish_reason,
366
+ "content": _clean(message["content"]),
367
+ "usage": dict(usage),
368
+ }
369
+ if finish_reason != "stop":
370
+ self.last_call_metadata = {
371
+ **self.last_call_metadata,
372
+ "status": "incomplete_response",
373
+ }
374
+ raise DeepSeekCallError(
375
+ f"DeepSeek response finish_reason must be stop, got {finish_reason!r}",
376
+ retryable=False,
377
+ )
378
+ self.last_call_metadata = {
379
+ **self.last_call_metadata,
380
+ "request": body,
381
+ "response_id": _required_text(raw.get("id"), "DeepSeek response id"),
382
+ "finish_reason": finish_reason,
383
+ "choices": choices,
384
+ "content": _clean(message["content"]),
385
+ "usage": dict(usage),
386
+ "http_status": status,
387
+ "physical_call_id": physical_call_id,
388
+ "status": "completed",
389
+ "completed_at": _now(),
390
+ }
391
+ patch = _strict_json(
392
+ _clean(message["content"]), label="DeepSeek content", expected=dict
393
+ )
394
+ transport_normalizations = []
395
+ for operation in patch.get("operations", []):
396
+ if (
397
+ isinstance(operation, Mapping)
398
+ and operation.get("action") == "create"
399
+ and "capsule_id" in operation
400
+ and operation.get("capsule_id") is None
401
+ ):
402
+ transport_normalizations.append(
403
+ {
404
+ "code": "create_null_capsule_id_ignored",
405
+ "field": "capsule_id",
406
+ }
407
+ )
408
+ if transport_normalizations:
409
+ self.last_call_metadata = {
410
+ **self.last_call_metadata,
411
+ "transport_normalizations": transport_normalizations,
412
+ }
413
+ validate_patch(patch)
414
+ return patch
415
+
416
+
417
+ def _ids(value: Any, label: str) -> list[str]:
418
+ if not isinstance(value, list) or not value:
419
+ raise PatchValidationError(f"{label} must be a non-empty list")
420
+ result = [_required_text(item, label) for item in value]
421
+ if len(set(result)) != len(result):
422
+ raise PatchValidationError(f"{label} contains duplicates")
423
+ return result
424
+
425
+
426
+ def _validate_source_parents(value: Any) -> list[dict[str, Any]]:
427
+ if not isinstance(value, list) or not value:
428
+ raise PatchValidationError("source_parents must be a non-empty list")
429
+ expected = {
430
+ "session_index",
431
+ "parent_chunk_index",
432
+ "message_index",
433
+ "source_record_id",
434
+ "event_id",
435
+ "evidence_char_start",
436
+ "evidence_char_end",
437
+ }
438
+ result: list[dict[str, Any]] = []
439
+ seen: set[str] = set()
440
+ for parent in value:
441
+ if not isinstance(parent, Mapping) or set(parent) != expected:
442
+ raise PatchValidationError("source_parent has an invalid schema")
443
+ normalized = dict(parent)
444
+ for key in ("session_index", "parent_chunk_index", "message_index"):
445
+ if not isinstance(normalized[key], int) or normalized[key] < 0:
446
+ raise PatchValidationError(f"source_parent {key} must be non-negative integer")
447
+ if normalized["parent_chunk_index"] != normalized["message_index"]:
448
+ raise PatchValidationError("source_parent chunk/message coordinates disagree")
449
+ if (
450
+ not isinstance(normalized["evidence_char_start"], int)
451
+ or not isinstance(normalized["evidence_char_end"], int)
452
+ or normalized["evidence_char_start"] < 0
453
+ or normalized["evidence_char_end"] <= normalized["evidence_char_start"]
454
+ ):
455
+ raise PatchValidationError("source_parent evidence character span is invalid")
456
+ _required_text(normalized["source_record_id"], "source_record_id")
457
+ _required_text(normalized["event_id"], "event_id")
458
+ digest = _json(normalized)
459
+ if digest not in seen:
460
+ seen.add(digest)
461
+ result.append(normalized)
462
+ return result
463
+
464
+
465
+ def _validate_claims(
466
+ value: Any, *, stored: bool = False
467
+ ) -> list[dict[str, Any]]:
468
+ if not isinstance(value, list) or not value:
469
+ raise PatchValidationError("claims must be a non-empty list")
470
+ claims: list[dict[str, Any]] = []
471
+ seen = set()
472
+ for claim in value:
473
+ expected = {"canonical_slot", "text", "support", "counterevidence"}
474
+ if stored:
475
+ expected.add("claim_id")
476
+ expected.add("source_parents")
477
+ if not isinstance(claim, Mapping):
478
+ raise PatchValidationError("each claim must be an object")
479
+ if stored:
480
+ valid_shape = set(claim) == expected
481
+ else:
482
+ valid_shape = (
483
+ {"canonical_slot", "text"}.issubset(claim)
484
+ and set(claim).issubset(expected)
485
+ )
486
+ if not valid_shape:
487
+ raise PatchValidationError(
488
+ "each claim has an invalid schema"
489
+ )
490
+ support = claim.get("support", [])
491
+ counter = claim.get("counterevidence", [])
492
+ if not isinstance(support, list) or not isinstance(counter, list):
493
+ raise PatchValidationError("claim evidence must be lists")
494
+ normalized_support = _ids(support, "claim support") if support else []
495
+ normalized_counter = (
496
+ _ids(counter, "claim counterevidence") if counter else []
497
+ )
498
+ if not normalized_support and not normalized_counter:
499
+ raise PatchValidationError("each claim needs support or counterevidence")
500
+ canonical_slot = _required_text(
501
+ claim.get("canonical_slot"), "canonical_slot"
502
+ )
503
+ claim_text = _required_text(claim.get("text"), "claim text")
504
+ claim_id = (
505
+ _required_text(claim.get("claim_id"), "claim_id")
506
+ if stored
507
+ else "clm_"
508
+ + _digest(
509
+ {
510
+ "canonical_slot": canonical_slot,
511
+ "text": claim_text,
512
+ "support": normalized_support,
513
+ "counterevidence": normalized_counter,
514
+ }
515
+ )[:24]
516
+ )
517
+ if claim_id in seen:
518
+ raise PatchValidationError("claim_id must be unique")
519
+ seen.add(claim_id)
520
+ claims.append(
521
+ {
522
+ "claim_id": claim_id,
523
+ "canonical_slot": canonical_slot,
524
+ "text": claim_text,
525
+ "support": normalized_support,
526
+ "counterevidence": normalized_counter,
527
+ **(
528
+ {
529
+ "source_parents": _validate_source_parents(
530
+ claim.get("source_parents")
531
+ )
532
+ }
533
+ if stored
534
+ else {}
535
+ ),
536
+ }
537
+ )
538
+ return claims
539
+
540
+
541
+ def validate_patch(patch: Mapping[str, Any]) -> None:
542
+ if (
543
+ not isinstance(patch, Mapping)
544
+ or set(patch) != {"operations"}
545
+ or not isinstance(patch["operations"], list)
546
+ ):
547
+ raise PatchValidationError("GraphPatch must contain exactly an operations list")
548
+ if len(patch["operations"]) > 1:
549
+ raise PatchValidationError("GraphPatch may contain at most one region operation")
550
+ for operation in patch["operations"]:
551
+ if not isinstance(operation, Mapping):
552
+ raise PatchValidationError("GraphPatch operation must be an object")
553
+ action = _required_text(operation.get("action"), "action")
554
+ if action not in PATCH_ACTIONS:
555
+ raise PatchValidationError("unknown GraphPatch action")
556
+ if "confidence" in operation:
557
+ raise PatchValidationError(
558
+ "model confidence is not an authoritative graph field"
559
+ )
560
+ allowed = {"action", "summary", "claims"}
561
+ if action == "create":
562
+ if "capsule_key" in operation:
563
+ raise PatchValidationError(
564
+ "create capsule identity is controller-derived from region_key"
565
+ )
566
+ if "capsule_id" in operation:
567
+ if operation.get("capsule_id") is not None:
568
+ raise PatchValidationError(
569
+ "create capsule identity is controller-derived from region_key"
570
+ )
571
+ allowed.add("capsule_id")
572
+ elif action == "noop":
573
+ allowed.add("capsule_id")
574
+ if "capsule_id" in operation:
575
+ _required_text(operation.get("capsule_id"), "capsule_id")
576
+ else:
577
+ allowed.update({"capsule_id", "base_revision"})
578
+ _required_text(operation.get("capsule_id"), "capsule_id")
579
+ if (
580
+ not isinstance(operation.get("base_revision"), int)
581
+ or operation["base_revision"] < 1
582
+ ):
583
+ raise PatchValidationError("base_revision must be a positive integer")
584
+ if set(operation) - allowed:
585
+ raise PatchValidationError("unexpected GraphPatch operation fields")
586
+ if action != "noop":
587
+ _validate_claims(operation.get("claims"))
588
+
589
+
590
+ class SlowGraphStore:
591
+ def __init__(
592
+ self,
593
+ database: str | Path,
594
+ *,
595
+ schema: tuple[type[Any], type[Any]],
596
+ claim_lease_seconds: int = DEFAULT_CLAIM_LEASE_SECONDS,
597
+ ) -> None:
598
+ if claim_lease_seconds <= 0:
599
+ raise SlowGraphError("claim lease duration must be positive")
600
+ self.database = Path(database)
601
+ self.record_type, self.edge_type = schema
602
+ self.claim_lease_seconds = claim_lease_seconds
603
+ self.database.parent.mkdir(parents=True, exist_ok=True)
604
+ self._init_schema()
605
+
606
+ def connect(self) -> sqlite3.Connection:
607
+ con = sqlite3.connect(self.database)
608
+ con.row_factory = sqlite3.Row
609
+ con.execute("PRAGMA foreign_keys=ON")
610
+ return con
611
+
612
+ @contextmanager
613
+ def connection(self):
614
+ con = self.connect()
615
+ try:
616
+ yield con
617
+ con.commit()
618
+ except Exception:
619
+ con.rollback()
620
+ raise
621
+ finally:
622
+ con.close()
623
+
624
+ def _init_schema(self) -> None:
625
+ with self.connection() as con:
626
+ con.executescript("""
627
+ CREATE TABLE IF NOT EXISTS slow_graph_jobs (
628
+ job_id TEXT PRIMARY KEY, idempotency_key TEXT NOT NULL UNIQUE, scope_id TEXT NOT NULL, region_key TEXT NOT NULL,
629
+ evidence_ids_json TEXT NOT NULL, metadata_json TEXT NOT NULL, status TEXT NOT NULL CHECK(status IN ('pending','retryable','failed','completed')),
630
+ attempts INTEGER NOT NULL, last_error TEXT NOT NULL, created_at INTEGER NOT NULL, updated_at INTEGER NOT NULL,
631
+ claim_token TEXT, claim_owner TEXT, lease_expires_at INTEGER);
632
+ CREATE TABLE IF NOT EXISTS slow_graph_attempts (
633
+ attempt_id TEXT PRIMARY KEY, job_id TEXT NOT NULL, scope_id TEXT NOT NULL, status TEXT NOT NULL,
634
+ call_metadata_json TEXT NOT NULL, error TEXT NOT NULL, created_at INTEGER NOT NULL, completed_at INTEGER,
635
+ claim_token TEXT, claim_owner TEXT);
636
+ CREATE TABLE IF NOT EXISTS slow_graph_batches (
637
+ batch_id TEXT PRIMARY KEY, batch_key TEXT NOT NULL UNIQUE, scope_id TEXT NOT NULL,
638
+ evidence_snapshot_hash TEXT NOT NULL, job_ids_json TEXT NOT NULL, manager_metadata_json TEXT NOT NULL,
639
+ created_at INTEGER NOT NULL);
640
+ CREATE TABLE IF NOT EXISTS slow_graph_patches (
641
+ patch_id TEXT PRIMARY KEY, job_id TEXT NOT NULL, scope_id TEXT NOT NULL, region_key TEXT NOT NULL, manager_model TEXT NOT NULL,
642
+ patch_json TEXT NOT NULL, call_metadata_json TEXT NOT NULL, applied_at INTEGER NOT NULL);
643
+ CREATE TABLE IF NOT EXISTS slow_graph_patch_operations (
644
+ operation_id TEXT PRIMARY KEY, patch_id TEXT NOT NULL, ordinal INTEGER NOT NULL, capsule_id TEXT NOT NULL, action TEXT NOT NULL,
645
+ base_revision INTEGER, result_revision INTEGER, operation_json TEXT NOT NULL, created_at INTEGER NOT NULL);
646
+ CREATE TABLE IF NOT EXISTS slow_graph_provenance (
647
+ provenance_id TEXT PRIMARY KEY, patch_id TEXT NOT NULL, scope_id TEXT NOT NULL, capsule_id TEXT NOT NULL, revision INTEGER NOT NULL,
648
+ evidence_memory_id TEXT NOT NULL, claim_id TEXT NOT NULL, polarity TEXT NOT NULL, source_parent_json TEXT NOT NULL, created_at INTEGER NOT NULL);
649
+ CREATE INDEX IF NOT EXISTS idx_slow_graph_jobs_pending ON slow_graph_jobs(status, created_at);
650
+ CREATE INDEX IF NOT EXISTS idx_slow_graph_provenance_capsule ON slow_graph_provenance(scope_id, capsule_id, revision);
651
+ CREATE TABLE IF NOT EXISTS memory_edges (
652
+ scope_id TEXT NOT NULL, edge_id TEXT NOT NULL, source_memory_id TEXT NOT NULL, target_memory_id TEXT NOT NULL, edge_type TEXT NOT NULL,
653
+ score REAL NOT NULL, model_score REAL NOT NULL, evidence_turn INTEGER NOT NULL, evidence TEXT NOT NULL, metadata_json TEXT NOT NULL,
654
+ PRIMARY KEY(scope_id, edge_id));
655
+ CREATE TABLE IF NOT EXISTS slot_heads (
656
+ scope_id TEXT NOT NULL, slot_key TEXT NOT NULL, memory_id TEXT NOT NULL,
657
+ PRIMARY KEY(scope_id,slot_key));
658
+ CREATE TABLE IF NOT EXISTS slot_history (
659
+ scope_id TEXT NOT NULL, slot_key TEXT NOT NULL, ordinal INTEGER NOT NULL, memory_id TEXT NOT NULL,
660
+ PRIMARY KEY(scope_id,slot_key,ordinal));
661
+ """)
662
+ self._add_column_if_missing(con, "slow_graph_jobs", "claim_token TEXT")
663
+ self._add_column_if_missing(con, "slow_graph_jobs", "claim_owner TEXT")
664
+ self._add_column_if_missing(
665
+ con, "slow_graph_jobs", "lease_expires_at INTEGER"
666
+ )
667
+ self._add_column_if_missing(
668
+ con, "slow_graph_attempts", "claim_token TEXT"
669
+ )
670
+ self._add_column_if_missing(
671
+ con, "slow_graph_attempts", "claim_owner TEXT"
672
+ )
673
+ con.execute(
674
+ "CREATE INDEX IF NOT EXISTS idx_slow_graph_jobs_claimable "
675
+ "ON slow_graph_jobs(status, claim_token, created_at)"
676
+ )
677
+
678
+ @staticmethod
679
+ def _add_column_if_missing(
680
+ con: sqlite3.Connection, table: str, declaration: str
681
+ ) -> None:
682
+ column = declaration.split()[0]
683
+ columns = {
684
+ str(row["name"])
685
+ for row in con.execute("PRAGMA table_info(" + table + ")")
686
+ }
687
+ if column not in columns:
688
+ con.execute("ALTER TABLE " + table + " ADD COLUMN " + declaration)
689
+
690
+ def _records_table_exists(self, con: sqlite3.Connection) -> None:
691
+ if (
692
+ con.execute(
693
+ "SELECT 1 FROM sqlite_master WHERE type='table' AND name='records'"
694
+ ).fetchone()
695
+ is None
696
+ ):
697
+ raise SlowGraphError("real graph records table is missing")
698
+
699
+ def _metadata(self, row: sqlite3.Row, label: str) -> dict[str, Any]:
700
+ return _strict_json(row["metadata_json"], label=label, expected=dict)
701
+
702
+ def fast_regions(self, scope_id: str) -> dict[str, list[dict[str, Any]]]:
703
+ regions: dict[str, list[dict[str, Any]]] = {}
704
+ with self.connection() as con:
705
+ self._records_table_exists(con)
706
+ rows = con.execute(
707
+ "SELECT memory_id,value,relation,turn_index,metadata_json FROM records WHERE scope_id=?",
708
+ (scope_id,),
709
+ ).fetchall()
710
+ for row in rows:
711
+ metadata = self._metadata(row, "fast evidence metadata")
712
+ if (
713
+ metadata.get("content_variant") != LEAF_VARIANT
714
+ or metadata.get("memory_layer") != "fast"
715
+ or metadata.get("node_kind") != "atomic_user_assertion"
716
+ or metadata.get("atomic_evidence_leaf") is not True
717
+ or metadata.get("authority") != "user_assertion"
718
+ ):
719
+ continue
720
+ key = _clean(
721
+ metadata.get("graph_entity_key")
722
+ or metadata.get("entity_key")
723
+ or metadata.get("domain")
724
+ )
725
+ if not key:
726
+ raise SlowGraphError("fast evidence is missing region key")
727
+ regions.setdefault(key, []).append(
728
+ {
729
+ "memory_id": row["memory_id"],
730
+ "value": row["value"],
731
+ "relation": row["relation"],
732
+ "turn_index": row["turn_index"],
733
+ "metadata": metadata,
734
+ }
735
+ )
736
+ return regions
737
+
738
+ def _capsules(
739
+ self, con: sqlite3.Connection, scope_id: str, region_key: str
740
+ ) -> list[dict[str, Any]]:
741
+ rows = con.execute(
742
+ "SELECT memory_id,state,value,metadata_json FROM records WHERE scope_id=?",
743
+ (scope_id,),
744
+ ).fetchall()
745
+ revisions: list[dict[str, Any]] = []
746
+ for row in rows:
747
+ meta = self._metadata(row, "capsule metadata")
748
+ if (
749
+ meta.get("content_variant") == CAPSULE_VARIANT
750
+ and meta.get("region_key") == region_key
751
+ ):
752
+ revisions.append(
753
+ {
754
+ "memory_id": row["memory_id"],
755
+ "record_state": row["state"],
756
+ "value": row["value"],
757
+ **meta,
758
+ }
759
+ )
760
+ if not revisions:
761
+ return []
762
+ latest_revision = max(int(item["revision"]) for item in revisions)
763
+ latest = [
764
+ item for item in revisions if int(item["revision"]) == latest_revision
765
+ ]
766
+ if len(latest) != 1:
767
+ raise AuditError("region capsule lacks one latest revision")
768
+ return latest
769
+
770
+ def _job_metadata(
771
+ self,
772
+ con: sqlite3.Connection,
773
+ scope_id: str,
774
+ region_key: str,
775
+ evidence_ids: list[str],
776
+ manager: PatchManager | None,
777
+ ) -> dict[str, Any]:
778
+ evidence = self._evidence(con, scope_id, evidence_ids)
779
+ normalized_region = _required_text(region_key, "region key")
780
+ foreign_evidence = [
781
+ item["memory_id"]
782
+ for item in evidence
783
+ if _clean(item["metadata"].get("graph_entity_key"))
784
+ != normalized_region
785
+ ]
786
+ if foreign_evidence:
787
+ raise EvidencePolicyError(
788
+ "fast evidence graph_entity_key does not match job region: "
789
+ + ",".join(foreign_evidence)
790
+ )
791
+ model = dict(manager.model_config) if manager else {"model": "unbound"}
792
+ prompt_hash = manager.prompt_hash if manager else "unbound"
793
+ capsules = self._capsules(con, scope_id, region_key)
794
+ return {
795
+ "schema_version": SCHEMA_VERSION,
796
+ "schema_hash": _digest(
797
+ {"version": SCHEMA_VERSION, "actions": sorted(PATCH_ACTIONS)}
798
+ ),
799
+ "prompt_hash": prompt_hash,
800
+ "model_config": model,
801
+ "evidence_content_hash": _digest(evidence),
802
+ "capsule_revision_hash": _digest(capsules),
803
+ }
804
+
805
+ def _enqueue_in_connection(
806
+ self,
807
+ con: sqlite3.Connection,
808
+ scope_id: str,
809
+ region_key: str,
810
+ evidence_ids: list[str],
811
+ *,
812
+ manager: PatchManager | None = None,
813
+ ) -> str:
814
+ now = _now()
815
+ metadata = self._job_metadata(
816
+ con, scope_id, region_key, evidence_ids, manager
817
+ )
818
+ idem = _digest(
819
+ {
820
+ "scope_id": scope_id,
821
+ "region_key": region_key,
822
+ "evidence_ids": evidence_ids,
823
+ **{
824
+ key: value
825
+ for key, value in metadata.items()
826
+ if key != "capsule_revision_hash"
827
+ },
828
+ }
829
+ )
830
+ job_id = "sgj_" + uuid.uuid4().hex
831
+ con.execute(
832
+ "INSERT OR IGNORE INTO slow_graph_jobs("
833
+ "job_id,idempotency_key,scope_id,region_key,evidence_ids_json,metadata_json,"
834
+ "status,attempts,last_error,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)",
835
+ (
836
+ job_id,
837
+ idem,
838
+ scope_id,
839
+ region_key,
840
+ _json(evidence_ids),
841
+ _json(metadata),
842
+ "pending",
843
+ 0,
844
+ "",
845
+ now,
846
+ now,
847
+ ),
848
+ )
849
+ row = con.execute(
850
+ "SELECT job_id FROM slow_graph_jobs WHERE idempotency_key=?", (idem,)
851
+ ).fetchone()
852
+ return str(row["job_id"])
853
+
854
+ def enqueue(
855
+ self,
856
+ scope_id: str,
857
+ region_key: str,
858
+ evidence_ids: Iterable[str],
859
+ *,
860
+ manager: PatchManager | None = None,
861
+ ) -> str:
862
+ normalized_ids = sorted(
863
+ set(_required_text(item, "evidence id") for item in evidence_ids)
864
+ )
865
+ if not normalized_ids:
866
+ raise SlowGraphError("cannot enqueue without fast evidence")
867
+ with self.connection() as con:
868
+ self._records_table_exists(con)
869
+ return self._enqueue_in_connection(
870
+ con,
871
+ scope_id,
872
+ region_key,
873
+ normalized_ids,
874
+ manager=manager,
875
+ )
876
+
877
+ def enqueue_regions(
878
+ self, scope_id: str, *, manager: PatchManager | None = None
879
+ ) -> list[str]:
880
+ regions = self.fast_regions(scope_id)
881
+ manager_metadata = {
882
+ "schema_version": SCHEMA_VERSION,
883
+ "prompt_hash": manager.prompt_hash if manager else "unbound",
884
+ "model_config": dict(manager.model_config)
885
+ if manager
886
+ else {"model": "unbound"},
887
+ }
888
+ snapshot = {
889
+ key: [
890
+ {
891
+ "memory_id": item["memory_id"],
892
+ "value": item["value"],
893
+ "metadata": item["metadata"],
894
+ }
895
+ for item in values
896
+ ]
897
+ for key, values in sorted(regions.items())
898
+ }
899
+ evidence_snapshot_hash = _digest(snapshot)
900
+ batch_key = _digest(
901
+ {
902
+ "scope_id": scope_id,
903
+ "evidence_snapshot_hash": evidence_snapshot_hash,
904
+ **manager_metadata,
905
+ }
906
+ )
907
+ with self.connection() as con:
908
+ self._records_table_exists(con)
909
+ existing = con.execute(
910
+ "SELECT job_ids_json FROM slow_graph_batches WHERE batch_key=?",
911
+ (batch_key,),
912
+ ).fetchone()
913
+ if existing is not None:
914
+ return [
915
+ str(item)
916
+ for item in _strict_json(
917
+ existing["job_ids_json"], label="batch job IDs", expected=list
918
+ )
919
+ ]
920
+ job_ids = [
921
+ self._enqueue_in_connection(
922
+ con,
923
+ scope_id,
924
+ key,
925
+ sorted(str(item["memory_id"]) for item in values),
926
+ manager=manager,
927
+ )
928
+ for key, values in sorted(regions.items())
929
+ ]
930
+ con.execute(
931
+ "INSERT INTO slow_graph_batches VALUES(?,?,?,?,?,?,?)",
932
+ (
933
+ "sgb_" + uuid.uuid4().hex,
934
+ batch_key,
935
+ scope_id,
936
+ evidence_snapshot_hash,
937
+ _json(job_ids),
938
+ _json(manager_metadata),
939
+ _now(),
940
+ ),
941
+ )
942
+ return job_ids
943
+
944
+ def _job(self, job_id: str) -> sqlite3.Row:
945
+ with self.connection() as con:
946
+ row = con.execute(
947
+ "SELECT * FROM slow_graph_jobs WHERE job_id=?", (job_id,)
948
+ ).fetchone()
949
+ if row is None:
950
+ raise SlowGraphError("unknown slow graph job: " + job_id)
951
+ self._metadata(row, "job")
952
+ return row
953
+
954
+ def _evidence(
955
+ self, con: sqlite3.Connection, scope_id: str, ids: Iterable[str]
956
+ ) -> list[dict[str, Any]]:
957
+ result = []
958
+ for evidence_id in sorted(set(ids)):
959
+ row = con.execute(
960
+ "SELECT memory_id,value,turn_index,state,metadata_json FROM records WHERE scope_id=? AND memory_id=?",
961
+ (scope_id, evidence_id),
962
+ ).fetchone()
963
+ if row is None:
964
+ raise EvidencePolicyError("unknown evidence: " + evidence_id)
965
+ meta = self._metadata(row, "fast evidence")
966
+ if (
967
+ meta.get("content_variant") != LEAF_VARIANT
968
+ or meta.get("memory_layer") != "fast"
969
+ or meta.get("node_kind") != "atomic_user_assertion"
970
+ or meta.get("atomic_evidence_leaf") is not True
971
+ or meta.get("authority") != "user_assertion"
972
+ ):
973
+ raise EvidencePolicyError(
974
+ "only fast product_semantic_memory leaves may support capsules: "
975
+ + evidence_id
976
+ )
977
+ parent = {
978
+ key: meta.get(key)
979
+ for key in (
980
+ "session_index",
981
+ "message_index",
982
+ "source_record_id",
983
+ "event_id",
984
+ "evidence_char_start",
985
+ "evidence_char_end",
986
+ )
987
+ }
988
+ if any(value is None or _clean(value) == "" for value in parent.values()):
989
+ raise EvidencePolicyError(
990
+ "fast evidence lacks structured source_parent fields: "
991
+ + evidence_id
992
+ )
993
+ parent["parent_chunk_index"] = parent["message_index"]
994
+ record_state = _clean(row["state"])
995
+ evidence_role = (
996
+ "current_authoritative"
997
+ if record_state in {"active", "parallel_active", "promoted"}
998
+ else "historical_noncurrent"
999
+ )
1000
+ result.append(
1001
+ {
1002
+ "memory_id": row["memory_id"],
1003
+ "value": row["value"],
1004
+ "turn_index": row["turn_index"],
1005
+ "record_state": record_state,
1006
+ "slow_graph_evidence_role": evidence_role,
1007
+ "metadata": {
1008
+ **meta,
1009
+ "record_state": record_state,
1010
+ "slow_graph_evidence_role": evidence_role,
1011
+ },
1012
+ "source_parent": parent,
1013
+ }
1014
+ )
1015
+ return result
1016
+
1017
+ def _head(
1018
+ self, con: sqlite3.Connection, scope_id: str, capsule_id: str
1019
+ ) -> tuple[int, str] | None:
1020
+ rows = con.execute(
1021
+ "SELECT memory_id,metadata_json FROM records WHERE scope_id=?", (scope_id,)
1022
+ ).fetchall()
1023
+ candidates = []
1024
+ for row in rows:
1025
+ meta = self._metadata(row, "capsule metadata")
1026
+ if (
1027
+ meta.get("content_variant") == CAPSULE_VARIANT
1028
+ and meta.get("capsule_id") == capsule_id
1029
+ ):
1030
+ candidates.append((int(meta["revision"]), str(row["memory_id"])))
1031
+ return max(candidates) if candidates else None
1032
+
1033
+ def _capsule_id(self, scope_id: str, region_key: str) -> str:
1034
+ return (
1035
+ "cap_"
1036
+ + _digest(
1037
+ {
1038
+ "scope_id": scope_id,
1039
+ "region_key": _required_text(region_key, "region_key"),
1040
+ }
1041
+ )[:24]
1042
+ )
1043
+
1044
+ def _mark_superseded(
1045
+ self, con: sqlite3.Connection, scope_id: str, memory_id: str
1046
+ ) -> None:
1047
+ row = con.execute(
1048
+ "SELECT slot_key,metadata_json FROM records WHERE scope_id=? AND memory_id=?",
1049
+ (scope_id, memory_id),
1050
+ ).fetchone()
1051
+ if row is None:
1052
+ raise SlowGraphError("cannot supersede missing record: " + memory_id)
1053
+ metadata = self._metadata(row, "superseded record")
1054
+ if metadata.get("content_variant") == CAPSULE_VARIANT:
1055
+ metadata["status"] = "superseded"
1056
+ con.execute(
1057
+ "UPDATE records SET state='superseded',metadata_json=? WHERE scope_id=? AND memory_id=?",
1058
+ (_json(metadata), scope_id, memory_id),
1059
+ )
1060
+ con.execute(
1061
+ "DELETE FROM slot_heads WHERE scope_id=? AND slot_key=? AND memory_id=?",
1062
+ (scope_id, row["slot_key"], memory_id),
1063
+ )
1064
+
1065
+ def _write_edge(
1066
+ self,
1067
+ con: sqlite3.Connection,
1068
+ *,
1069
+ scope_id: str,
1070
+ source: str,
1071
+ target: str,
1072
+ edge_type: str,
1073
+ patch_id: str,
1074
+ evidence_refs: list[str],
1075
+ action: str,
1076
+ turn: int,
1077
+ ) -> None:
1078
+ if edge_type not in EDGE_TYPES or source == target:
1079
+ raise AuditError("invalid graph edge")
1080
+ edge = self.edge_type(
1081
+ edge_id="sge_"
1082
+ + _digest(
1083
+ {
1084
+ "patch": patch_id,
1085
+ "source": source,
1086
+ "target": target,
1087
+ "type": edge_type,
1088
+ }
1089
+ )[:24],
1090
+ source_memory_id=source,
1091
+ target_memory_id=target,
1092
+ edge_type=edge_type,
1093
+ score=0.0,
1094
+ model_score=0.0,
1095
+ evidence_turn=turn,
1096
+ evidence="slow graph patch",
1097
+ metadata={
1098
+ "edge_source": SLOW_EDGE_SOURCE,
1099
+ "patch_id": patch_id,
1100
+ "evidence_refs": evidence_refs,
1101
+ "action": action,
1102
+ },
1103
+ )
1104
+ con.execute(
1105
+ "INSERT INTO memory_edges VALUES(?,?,?,?,?,?,?,?,?,?)",
1106
+ (
1107
+ scope_id,
1108
+ edge.edge_id,
1109
+ edge.source_memory_id,
1110
+ edge.target_memory_id,
1111
+ edge.edge_type,
1112
+ edge.score,
1113
+ edge.model_score,
1114
+ edge.evidence_turn,
1115
+ edge.evidence,
1116
+ _json(edge.metadata),
1117
+ ),
1118
+ )
1119
+
1120
+ def _insert_revision(
1121
+ self,
1122
+ con: sqlite3.Connection,
1123
+ *,
1124
+ job: sqlite3.Row,
1125
+ patch_id: str,
1126
+ operation: Mapping[str, Any],
1127
+ capsule_id: str,
1128
+ revision: int,
1129
+ action: str,
1130
+ old_memory_id: str | None = None,
1131
+ ) -> str:
1132
+ claims = _validate_claims(operation["claims"])
1133
+ evidence_ids = sorted(
1134
+ {
1135
+ item
1136
+ for claim in claims
1137
+ for key in ("support", "counterevidence")
1138
+ for item in claim[key]
1139
+ }
1140
+ )
1141
+ allowed_evidence = set(
1142
+ _strict_json(
1143
+ job["evidence_ids_json"], label="job evidence IDs", expected=list
1144
+ )
1145
+ )
1146
+ if not set(evidence_ids).issubset(allowed_evidence):
1147
+ raise EvidencePolicyError(
1148
+ "GraphPatch cited fast evidence that was not supplied to this job"
1149
+ )
1150
+ evidence = self._evidence(con, job["scope_id"], evidence_ids)
1151
+ evidence_by_id = {item["memory_id"]: item for item in evidence}
1152
+ controller_normalizations: list[dict[str, Any]] = []
1153
+ for index, claim in enumerate(claims):
1154
+ cited_slots = {
1155
+ _clean(evidence_by_id[evidence_id]["metadata"].get("canonical_slot_key"))
1156
+ for evidence_id in [*claim["support"], *claim["counterevidence"]]
1157
+ }
1158
+ cited_slots.discard("")
1159
+ if claim["canonical_slot"] not in cited_slots:
1160
+ if len(cited_slots) != 1:
1161
+ raise EvidencePolicyError(
1162
+ "claim canonical_slot is ambiguous across cited fast leaf slots"
1163
+ )
1164
+ authoritative_slot = next(iter(cited_slots))
1165
+ original_slot = claim["canonical_slot"]
1166
+ normalized_claim = {
1167
+ **claim,
1168
+ "canonical_slot": authoritative_slot,
1169
+ }
1170
+ normalized_claim["claim_id"] = "clm_" + _digest(
1171
+ {
1172
+ "canonical_slot": authoritative_slot,
1173
+ "text": normalized_claim["text"],
1174
+ "support": normalized_claim["support"],
1175
+ "counterevidence": normalized_claim["counterevidence"],
1176
+ }
1177
+ )[:24]
1178
+ claims[index] = normalized_claim
1179
+ controller_normalizations.append(
1180
+ {
1181
+ "code": "canonical_slot_bound_to_unique_cited_leaf",
1182
+ "claim_index": index,
1183
+ "model_value": original_slot,
1184
+ "authoritative_value": authoritative_slot,
1185
+ }
1186
+ )
1187
+ claim_ids = [claim["claim_id"] for claim in claims]
1188
+ if len(set(claim_ids)) != len(claim_ids):
1189
+ raise EvidencePolicyError(
1190
+ "canonical slot binding produced duplicate claims"
1191
+ )
1192
+ status = {
1193
+ "challenge": "challenged",
1194
+ "retire": "retired",
1195
+ "resolve_challenge": "active",
1196
+ }.get(action, "active")
1197
+ stored_claims: list[dict[str, Any]] = []
1198
+ for claim in claims:
1199
+ claim_parents = [
1200
+ evidence_by_id[evidence_id]["source_parent"]
1201
+ for evidence_id in [*claim["support"], *claim["counterevidence"]]
1202
+ ]
1203
+ stored_claims.append(
1204
+ {**claim, "source_parents": _validate_source_parents(claim_parents)}
1205
+ )
1206
+ source_parents = _validate_source_parents(
1207
+ [parent for claim in stored_claims for parent in claim["source_parents"]]
1208
+ )
1209
+ record_id = f"slow.{capsule_id}.r{revision}"
1210
+ metadata = {
1211
+ "memory_layer": "slow",
1212
+ "content_variant": CAPSULE_VARIANT,
1213
+ "capsule_id": capsule_id,
1214
+ "revision": revision,
1215
+ "status": status,
1216
+ "claims": stored_claims,
1217
+ "source_parents": source_parents,
1218
+ "canonical_slots": sorted(
1219
+ {claim["canonical_slot"] for claim in stored_claims}
1220
+ ),
1221
+ "patch_id": patch_id,
1222
+ "region_key": job["region_key"],
1223
+ "action": action,
1224
+ "controller_normalizations": controller_normalizations,
1225
+ }
1226
+ record_state = "active" if status in {"active", "challenged"} else status
1227
+ record = self.record_type(
1228
+ memory_id=record_id,
1229
+ category=CAPSULE_VARIANT,
1230
+ slot_key="slow." + capsule_id,
1231
+ value=_clean(operation.get("summary")) or _json(claims),
1232
+ relation="capsule_revision",
1233
+ anchor_concepts=[job["region_key"]],
1234
+ evidence_anchors=evidence_ids,
1235
+ salience=0.7,
1236
+ confidence=0.0,
1237
+ source_kind="slow_graph",
1238
+ turn_index=max(item["turn_index"] for item in evidence),
1239
+ state=record_state,
1240
+ supersedes=[old_memory_id] if old_memory_id else [],
1241
+ metadata=metadata,
1242
+ )
1243
+ payload = record.to_dict()
1244
+ con.execute(
1245
+ "INSERT INTO records(scope_id,memory_id,category,slot_key,value,relation,anchor_concepts_json,evidence_anchors_json,salience,confidence,source_kind,turn_index,state,supersedes_json,metadata_json) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
1246
+ (
1247
+ job["scope_id"],
1248
+ payload["memory_id"],
1249
+ payload["category"],
1250
+ payload["slot_key"],
1251
+ payload["value"],
1252
+ payload["relation"],
1253
+ _json(payload["anchor_concepts"]),
1254
+ _json(payload["evidence_anchors"]),
1255
+ payload["salience"],
1256
+ payload["confidence"],
1257
+ payload["source_kind"],
1258
+ payload["turn_index"],
1259
+ payload["state"],
1260
+ _json(payload["supersedes"]),
1261
+ _json(payload["metadata"]),
1262
+ ),
1263
+ )
1264
+ slot_key = payload["slot_key"]
1265
+ ordinal = int(
1266
+ con.execute(
1267
+ "SELECT COALESCE(MAX(ordinal),-1)+1 FROM slot_history WHERE scope_id=? AND slot_key=?",
1268
+ (job["scope_id"], slot_key),
1269
+ ).fetchone()[0]
1270
+ )
1271
+ con.execute(
1272
+ "INSERT INTO slot_history(scope_id,slot_key,ordinal,memory_id) VALUES(?,?,?,?)",
1273
+ (job["scope_id"], slot_key, ordinal, record_id),
1274
+ )
1275
+ if record_state == "active":
1276
+ con.execute(
1277
+ "INSERT INTO slot_heads(scope_id,slot_key,memory_id) VALUES(?,?,?) "
1278
+ "ON CONFLICT(scope_id,slot_key) DO UPDATE SET memory_id=excluded.memory_id",
1279
+ (job["scope_id"], slot_key, record_id),
1280
+ )
1281
+ else:
1282
+ con.execute(
1283
+ "DELETE FROM slot_heads WHERE scope_id=? AND slot_key=?",
1284
+ (job["scope_id"], slot_key),
1285
+ )
1286
+ for claim in stored_claims:
1287
+ for polarity, edge_kind in (
1288
+ ("support", "supports"),
1289
+ ("counterevidence", "contradicts"),
1290
+ ):
1291
+ for evidence_id in claim[polarity]:
1292
+ item = evidence_by_id[evidence_id]
1293
+ con.execute(
1294
+ "INSERT INTO slow_graph_provenance VALUES(?,?,?,?,?,?,?,?,?,?)",
1295
+ (
1296
+ "sgv_" + uuid.uuid4().hex,
1297
+ patch_id,
1298
+ job["scope_id"],
1299
+ capsule_id,
1300
+ revision,
1301
+ evidence_id,
1302
+ claim["claim_id"],
1303
+ polarity,
1304
+ _json(item["source_parent"]),
1305
+ _now(),
1306
+ ),
1307
+ )
1308
+ self._write_edge(
1309
+ con,
1310
+ scope_id=job["scope_id"],
1311
+ source=evidence_id,
1312
+ target=record_id,
1313
+ edge_type=edge_kind,
1314
+ patch_id=patch_id,
1315
+ evidence_refs=[evidence_id],
1316
+ action=action,
1317
+ turn=item["turn_index"],
1318
+ )
1319
+ if old_memory_id:
1320
+ self._mark_superseded(con, job["scope_id"], old_memory_id)
1321
+ self._write_edge(
1322
+ con,
1323
+ scope_id=job["scope_id"],
1324
+ source=old_memory_id,
1325
+ target=record_id,
1326
+ edge_type="supersedes",
1327
+ patch_id=patch_id,
1328
+ evidence_refs=evidence_ids,
1329
+ action=action,
1330
+ turn=payload["turn_index"],
1331
+ )
1332
+ return record_id
1333
+
1334
+ def _audit_transaction(self, con: sqlite3.Connection, scope_id: str) -> None:
1335
+ """Validate slow-layer invariants against the uncommitted write transaction."""
1336
+ rows = con.execute(
1337
+ "SELECT memory_id,state,metadata_json FROM records WHERE scope_id=?",
1338
+ (scope_id,),
1339
+ ).fetchall()
1340
+ capsules: dict[str, list[tuple[str, str, int, str]]] = {}
1341
+ known_records = {str(row["memory_id"]) for row in rows}
1342
+ for row in rows:
1343
+ meta = self._metadata(row, "transaction record")
1344
+ if meta.get("content_variant") != CAPSULE_VARIANT:
1345
+ continue
1346
+ capsule_id = _required_text(meta.get("capsule_id"), "capsule_id")
1347
+ revision = meta.get("revision")
1348
+ status = _required_text(meta.get("status"), "capsule status")
1349
+ if not isinstance(revision, int) or revision < 1:
1350
+ raise AuditError("capsule revision is invalid")
1351
+ if status not in {"active", "challenged", "retired", "superseded"}:
1352
+ raise AuditError("capsule status is invalid")
1353
+ claims = _validate_claims(meta.get("claims"), stored=True)
1354
+ capsules.setdefault(capsule_id, []).append(
1355
+ (str(row["memory_id"]), str(row["state"]), revision, status)
1356
+ )
1357
+ for claim in claims:
1358
+ expected = set(claim["support"] + claim["counterevidence"])
1359
+ actual = {
1360
+ str(item["evidence_memory_id"])
1361
+ for item in con.execute(
1362
+ "SELECT evidence_memory_id FROM slow_graph_provenance "
1363
+ "WHERE scope_id=? AND capsule_id=? AND revision=? AND claim_id=?",
1364
+ (scope_id, capsule_id, revision, claim["claim_id"]),
1365
+ )
1366
+ }
1367
+ if actual != expected:
1368
+ raise AuditError("claim provenance does not match claim evidence")
1369
+ for capsule_id, revisions in capsules.items():
1370
+ ordered = sorted(revisions, key=lambda item: item[2])
1371
+ if [item[2] for item in ordered] != list(range(1, len(ordered) + 1)):
1372
+ raise AuditError("capsule revisions are not contiguous")
1373
+ memory_id, state, _, status = ordered[-1]
1374
+ slot_key = "slow." + capsule_id
1375
+ head = con.execute(
1376
+ "SELECT memory_id FROM slot_heads WHERE scope_id=? AND slot_key=?",
1377
+ (scope_id, slot_key),
1378
+ ).fetchone()
1379
+ if status in {"active", "challenged"}:
1380
+ if state != "active" or head is None or head["memory_id"] != memory_id:
1381
+ raise AuditError("active/challenged capsule head is inconsistent")
1382
+ elif head is not None:
1383
+ raise AuditError("inactive capsule must not retain a slot head")
1384
+ history = [
1385
+ str(item["memory_id"])
1386
+ for item in con.execute(
1387
+ "SELECT memory_id FROM slot_history WHERE scope_id=? AND slot_key=? ORDER BY ordinal",
1388
+ (scope_id, slot_key),
1389
+ )
1390
+ ]
1391
+ if history != [item[0] for item in ordered]:
1392
+ raise AuditError("capsule slot history is inconsistent")
1393
+ for edge in con.execute(
1394
+ "SELECT source_memory_id,target_memory_id,edge_type,metadata_json "
1395
+ "FROM memory_edges WHERE scope_id=?",
1396
+ (scope_id,),
1397
+ ):
1398
+ meta = _strict_json(
1399
+ edge["metadata_json"], label="edge metadata", expected=dict
1400
+ )
1401
+ if meta.get("edge_source") != SLOW_EDGE_SOURCE:
1402
+ continue
1403
+ if (
1404
+ edge["source_memory_id"] == edge["target_memory_id"]
1405
+ or edge["edge_type"] not in EDGE_TYPES
1406
+ or edge["source_memory_id"] not in known_records
1407
+ or edge["target_memory_id"] not in known_records
1408
+ or not isinstance(meta.get("evidence_refs"), list)
1409
+ or not _clean(meta.get("patch_id"))
1410
+ ):
1411
+ raise AuditError("slow graph edge is inconsistent")
1412
+
1413
+ def apply_patch(
1414
+ self,
1415
+ job_id: str,
1416
+ patch: Mapping[str, Any],
1417
+ *,
1418
+ manager_model: str,
1419
+ call_metadata: Mapping[str, Any] | None = None,
1420
+ claim: JobClaim,
1421
+ ) -> str:
1422
+ validate_patch(patch)
1423
+ if claim.job_id != job_id:
1424
+ raise SlowGraphError("claim does not belong to job")
1425
+ patch_id = "sgp_" + uuid.uuid4().hex
1426
+ with self.connection() as con:
1427
+ con.execute("BEGIN IMMEDIATE")
1428
+ now = _now()
1429
+ job = con.execute(
1430
+ "SELECT * FROM slow_graph_jobs WHERE job_id=?", (job_id,)
1431
+ ).fetchone()
1432
+ if job is None:
1433
+ raise SlowGraphError("unknown slow graph job")
1434
+ if (
1435
+ job["status"] != "pending"
1436
+ or job["claim_token"] != claim.token
1437
+ or job["claim_owner"] != claim.owner
1438
+ or job["lease_expires_at"] is None
1439
+ or int(job["lease_expires_at"]) < now
1440
+ ):
1441
+ raise SlowGraphError("job claim is no longer active")
1442
+ self._records_table_exists(con)
1443
+ metadata = self._metadata(job, "job")
1444
+ if metadata["schema_version"] != SCHEMA_VERSION:
1445
+ raise SlowGraphError("job schema metadata drift")
1446
+ con.execute(
1447
+ "INSERT INTO slow_graph_patches VALUES(?,?,?,?,?,?,?,?)",
1448
+ (
1449
+ patch_id,
1450
+ job_id,
1451
+ job["scope_id"],
1452
+ job["region_key"],
1453
+ manager_model,
1454
+ _json(patch),
1455
+ _json(dict(call_metadata or {})),
1456
+ now,
1457
+ ),
1458
+ )
1459
+ ordinal = 0
1460
+ for operation in patch["operations"]:
1461
+ action = operation["action"]
1462
+ region_capsule_id = self._capsule_id(
1463
+ job["scope_id"], job["region_key"]
1464
+ )
1465
+ capsule_id = (
1466
+ region_capsule_id
1467
+ if action == "create"
1468
+ else operation.get("capsule_id")
1469
+ or "region:" + job["region_key"]
1470
+ )
1471
+ if action != "create" and operation.get("capsule_id"):
1472
+ if capsule_id != region_capsule_id:
1473
+ raise EvidencePolicyError(
1474
+ "GraphPatch capsule_id does not belong to the current region"
1475
+ )
1476
+ head = self._head(con, job["scope_id"], capsule_id)
1477
+ if action == "create":
1478
+ if head is not None:
1479
+ raise StaleRevisionError("capsule already exists")
1480
+ record_id = self._insert_revision(
1481
+ con,
1482
+ job=job,
1483
+ patch_id=patch_id,
1484
+ operation=operation,
1485
+ capsule_id=capsule_id,
1486
+ revision=1,
1487
+ action=action,
1488
+ )
1489
+ base, revision = None, 1
1490
+ elif action == "noop":
1491
+ if operation.get("capsule_id") and head is None:
1492
+ raise StaleRevisionError("noop capsule does not exist")
1493
+ if head is None:
1494
+ record_id, base, revision = "", None, None
1495
+ else:
1496
+ record_id, (revision, _) = head[1], head
1497
+ base = revision
1498
+ else:
1499
+ if head is None or operation["base_revision"] != head[0]:
1500
+ raise StaleRevisionError(
1501
+ "base_revision is stale for " + capsule_id
1502
+ )
1503
+ base, revision = head[0], head[0] + 1
1504
+ record_id = self._insert_revision(
1505
+ con,
1506
+ job=job,
1507
+ patch_id=patch_id,
1508
+ operation=operation,
1509
+ capsule_id=capsule_id,
1510
+ revision=revision,
1511
+ action=action,
1512
+ old_memory_id=head[1],
1513
+ )
1514
+ if action == "challenge":
1515
+ self._write_edge(
1516
+ con,
1517
+ scope_id=job["scope_id"],
1518
+ source=record_id,
1519
+ target=head[1],
1520
+ edge_type="challenges",
1521
+ patch_id=patch_id,
1522
+ evidence_refs=[],
1523
+ action=action,
1524
+ turn=now,
1525
+ )
1526
+ if action == "retire":
1527
+ self._write_edge(
1528
+ con,
1529
+ scope_id=job["scope_id"],
1530
+ source=record_id,
1531
+ target=head[1],
1532
+ edge_type="invalidates",
1533
+ patch_id=patch_id,
1534
+ evidence_refs=[],
1535
+ action=action,
1536
+ turn=now,
1537
+ )
1538
+ con.execute(
1539
+ "INSERT INTO slow_graph_patch_operations VALUES(?,?,?,?,?,?,?,?,?)",
1540
+ (
1541
+ "sgo_" + uuid.uuid4().hex,
1542
+ patch_id,
1543
+ ordinal,
1544
+ capsule_id,
1545
+ action,
1546
+ base,
1547
+ revision,
1548
+ _json(operation),
1549
+ now,
1550
+ ),
1551
+ )
1552
+ ordinal += 1
1553
+ completed_attempt = con.execute(
1554
+ "UPDATE slow_graph_attempts SET status='completed',"
1555
+ "call_metadata_json=?,completed_at=? WHERE attempt_id=? AND job_id=? "
1556
+ "AND claim_token=? AND claim_owner=? AND status='started'",
1557
+ (
1558
+ _json(dict(call_metadata or {})),
1559
+ now,
1560
+ claim.attempt_id,
1561
+ job_id,
1562
+ claim.token,
1563
+ claim.owner,
1564
+ ),
1565
+ )
1566
+ if completed_attempt.rowcount != 1:
1567
+ raise SlowGraphError("claimed attempt is no longer active")
1568
+ completed_job = con.execute(
1569
+ "UPDATE slow_graph_jobs SET status='completed',attempts=attempts+1,"
1570
+ "last_error='',updated_at=?,claim_token=NULL,claim_owner=NULL,"
1571
+ "lease_expires_at=NULL WHERE job_id=? AND status='pending' "
1572
+ "AND claim_token=? AND claim_owner=? AND lease_expires_at>=?",
1573
+ (now, job_id, claim.token, claim.owner, now),
1574
+ )
1575
+ if completed_job.rowcount != 1:
1576
+ raise SlowGraphError("job claim expired before completion")
1577
+ self._audit_transaction(con, job["scope_id"])
1578
+ return patch_id
1579
+
1580
+ def _claim_pending_job(
1581
+ self, job_id: str | None, *, owner: str
1582
+ ) -> JobClaim | None:
1583
+ token = "sgc_" + uuid.uuid4().hex
1584
+ attempt_id = "sga_" + uuid.uuid4().hex
1585
+ now = _now()
1586
+ with self.connection() as con:
1587
+ con.execute("BEGIN IMMEDIATE")
1588
+ if job_id is None:
1589
+ row = con.execute(
1590
+ "SELECT job_id,scope_id FROM slow_graph_jobs WHERE status='pending' "
1591
+ "AND claim_token IS NULL ORDER BY created_at LIMIT 1"
1592
+ ).fetchone()
1593
+ else:
1594
+ row = con.execute(
1595
+ "SELECT job_id,scope_id FROM slow_graph_jobs WHERE job_id=? "
1596
+ "AND status='pending' AND claim_token IS NULL",
1597
+ (job_id,),
1598
+ ).fetchone()
1599
+ if row is None:
1600
+ return None
1601
+ claimed = con.execute(
1602
+ "UPDATE slow_graph_jobs SET claim_token=?,claim_owner=?,"
1603
+ "lease_expires_at=?,updated_at=? WHERE job_id=? AND status='pending' "
1604
+ "AND claim_token IS NULL",
1605
+ (
1606
+ token,
1607
+ owner,
1608
+ now + self.claim_lease_seconds,
1609
+ now,
1610
+ str(row["job_id"]),
1611
+ ),
1612
+ )
1613
+ if claimed.rowcount != 1:
1614
+ return None
1615
+ con.execute(
1616
+ "INSERT INTO slow_graph_attempts("
1617
+ "attempt_id,job_id,scope_id,status,call_metadata_json,error,created_at,"
1618
+ "completed_at,claim_token,claim_owner) VALUES(?,?,?,?,?,?,?,?,?,?)",
1619
+ (
1620
+ attempt_id,
1621
+ str(row["job_id"]),
1622
+ str(row["scope_id"]),
1623
+ "started",
1624
+ _json({}),
1625
+ "",
1626
+ now,
1627
+ None,
1628
+ token,
1629
+ owner,
1630
+ ),
1631
+ )
1632
+ return JobClaim(str(row["job_id"]), attempt_id, token, owner)
1633
+
1634
+ def _claim_owner(self) -> str:
1635
+ return "pid:" + str(os.getpid()) + ":" + uuid.uuid4().hex
1636
+
1637
+ def _claim_context(
1638
+ self, claim: JobClaim
1639
+ ) -> tuple[dict[str, Any], list[dict[str, Any]]]:
1640
+ with self.connection() as con:
1641
+ job = con.execute(
1642
+ "SELECT * FROM slow_graph_jobs WHERE job_id=?", (claim.job_id,)
1643
+ ).fetchone()
1644
+ if job is None:
1645
+ raise SlowGraphError("unknown slow graph job")
1646
+ if (
1647
+ job["status"] != "pending"
1648
+ or job["claim_token"] != claim.token
1649
+ or job["claim_owner"] != claim.owner
1650
+ or job["lease_expires_at"] is None
1651
+ or int(job["lease_expires_at"]) < _now()
1652
+ ):
1653
+ raise SlowGraphError("job claim is no longer active")
1654
+ evidence_ids = _strict_json(
1655
+ job["evidence_ids_json"], label="job evidence IDs", expected=list
1656
+ )
1657
+ return (
1658
+ {
1659
+ "region_key": job["region_key"],
1660
+ "evidence": self._evidence(con, job["scope_id"], evidence_ids),
1661
+ },
1662
+ self._capsules(con, job["scope_id"], job["region_key"]),
1663
+ )
1664
+
1665
+ def _renew_claim(self, claim: JobClaim) -> None:
1666
+ now = _now()
1667
+ with self.connection() as con:
1668
+ con.execute("BEGIN IMMEDIATE")
1669
+ renewed = con.execute(
1670
+ "UPDATE slow_graph_jobs SET lease_expires_at=?,updated_at=? "
1671
+ "WHERE job_id=? AND status='pending' AND claim_token=? "
1672
+ "AND claim_owner=? AND lease_expires_at>=?",
1673
+ (
1674
+ now + self.claim_lease_seconds,
1675
+ now,
1676
+ claim.job_id,
1677
+ claim.token,
1678
+ claim.owner,
1679
+ now,
1680
+ ),
1681
+ )
1682
+ if renewed.rowcount != 1:
1683
+ raise SlowGraphError("job claim could not be renewed")
1684
+
1685
+ def _propose_with_lease_heartbeat(
1686
+ self,
1687
+ claim: JobClaim,
1688
+ manager: PatchManager,
1689
+ region: Mapping[str, Any],
1690
+ capsules: list[dict[str, Any]],
1691
+ ) -> dict[str, Any]:
1692
+ stop = threading.Event()
1693
+ errors: list[Exception] = []
1694
+ interval = max(0.1, min(30.0, self.claim_lease_seconds / 3.0))
1695
+
1696
+ def heartbeat() -> None:
1697
+ while not stop.wait(interval):
1698
+ try:
1699
+ self._renew_claim(claim)
1700
+ except Exception as exc:
1701
+ errors.append(exc)
1702
+ stop.set()
1703
+
1704
+ thread = threading.Thread(
1705
+ target=heartbeat,
1706
+ name=f"slow-graph-lease-{claim.job_id}",
1707
+ daemon=True,
1708
+ )
1709
+ thread.start()
1710
+ try:
1711
+ patch = manager.propose(region, capsules)
1712
+ finally:
1713
+ stop.set()
1714
+ thread.join(timeout=max(1.0, interval + 1.0))
1715
+ if thread.is_alive():
1716
+ raise SlowGraphError("slow graph claim heartbeat did not stop")
1717
+ if errors:
1718
+ raise SlowGraphError(f"slow graph claim heartbeat failed: {errors[0]}")
1719
+ self._renew_claim(claim)
1720
+ return patch
1721
+
1722
+ def _finish_claim_failure(
1723
+ self, claim: JobClaim, manager: PatchManager, exc: Exception
1724
+ ) -> bool:
1725
+ status = (
1726
+ "retryable"
1727
+ if isinstance(exc, DeepSeekCallError) and exc.retryable
1728
+ else "failed"
1729
+ )
1730
+ now = _now()
1731
+ with self.connection() as con:
1732
+ con.execute("BEGIN IMMEDIATE")
1733
+ failed_job = con.execute(
1734
+ "UPDATE slow_graph_jobs SET status=?,attempts=attempts+1,last_error=?,"
1735
+ "updated_at=?,claim_token=NULL,claim_owner=NULL,lease_expires_at=NULL "
1736
+ "WHERE job_id=? AND status='pending' AND claim_token=? "
1737
+ "AND claim_owner=? AND lease_expires_at>=?",
1738
+ (status, str(exc), now, claim.job_id, claim.token, claim.owner, now),
1739
+ )
1740
+ if failed_job.rowcount != 1:
1741
+ return False
1742
+ failed_attempt = con.execute(
1743
+ "UPDATE slow_graph_attempts SET status=?,call_metadata_json=?,error=?,"
1744
+ "completed_at=? WHERE attempt_id=? AND job_id=? AND claim_token=? "
1745
+ "AND claim_owner=? AND status='started'",
1746
+ (
1747
+ status,
1748
+ _json(dict(manager.last_call_metadata)),
1749
+ str(exc),
1750
+ now,
1751
+ claim.attempt_id,
1752
+ claim.job_id,
1753
+ claim.token,
1754
+ claim.owner,
1755
+ ),
1756
+ )
1757
+ if failed_attempt.rowcount != 1:
1758
+ raise SlowGraphError("claimed attempt is no longer active")
1759
+ return True
1760
+
1761
+ def _run_claimed_job(self, claim: JobClaim, manager: PatchManager) -> str:
1762
+ try:
1763
+ region, capsules = self._claim_context(claim)
1764
+ patch = self._propose_with_lease_heartbeat(
1765
+ claim, manager, region, capsules
1766
+ )
1767
+ patch_id = self.apply_patch(
1768
+ claim.job_id,
1769
+ patch,
1770
+ manager_model=_required_text(
1771
+ manager.model_config.get("model"), "manager model"
1772
+ ),
1773
+ call_metadata=manager.last_call_metadata,
1774
+ claim=claim,
1775
+ )
1776
+ return patch_id
1777
+ except Exception as exc:
1778
+ self._finish_claim_failure(claim, manager, exc)
1779
+ raise
1780
+
1781
+ def run_job(self, job_id: str, manager: PatchManager) -> str:
1782
+ self.recover_interrupted_attempts()
1783
+ claim = self._claim_pending_job(job_id, owner=self._claim_owner())
1784
+ if claim is None:
1785
+ raise SlowGraphError("job is not pending or is already claimed")
1786
+ return self._run_claimed_job(claim, manager)
1787
+
1788
+ def resume(self, job_id: str) -> None:
1789
+ with self.connection() as con:
1790
+ con.execute("BEGIN IMMEDIATE")
1791
+ row = con.execute(
1792
+ "SELECT status,claim_token FROM slow_graph_jobs WHERE job_id=?",
1793
+ (job_id,),
1794
+ ).fetchone()
1795
+ if (
1796
+ row is None
1797
+ or row["status"] not in {"failed", "retryable"}
1798
+ or row["claim_token"] is not None
1799
+ ):
1800
+ raise SlowGraphError(
1801
+ "only failed or retryable jobs can be explicitly reopened"
1802
+ )
1803
+ reopened = con.execute(
1804
+ "UPDATE slow_graph_jobs SET status='pending',last_error='',updated_at=?,"
1805
+ "claim_token=NULL,claim_owner=NULL,lease_expires_at=NULL WHERE job_id=? "
1806
+ "AND status IN ('failed','retryable') AND claim_token IS NULL",
1807
+ (_now(), job_id),
1808
+ )
1809
+ if reopened.rowcount != 1:
1810
+ raise SlowGraphError("slow graph job changed while reopening")
1811
+
1812
+ def recover_interrupted_attempts(self) -> int:
1813
+ now = _now()
1814
+ with self.connection() as con:
1815
+ con.execute("BEGIN IMMEDIATE")
1816
+ rows = con.execute(
1817
+ "SELECT job_id,claim_token,claim_owner FROM slow_graph_jobs "
1818
+ "WHERE status='pending' AND claim_token IS NOT NULL "
1819
+ "AND lease_expires_at<?",
1820
+ (now,),
1821
+ ).fetchall()
1822
+ for row in rows:
1823
+ released = con.execute(
1824
+ "UPDATE slow_graph_jobs SET status='failed',attempts=attempts+1,"
1825
+ "last_error='claim lease expired; external call outcome uncertain; explicit resume required',"
1826
+ "updated_at=?,claim_token=NULL,claim_owner=NULL,lease_expires_at=NULL "
1827
+ "WHERE job_id=? AND status='pending' AND claim_token=? "
1828
+ "AND claim_owner=? AND lease_expires_at<?",
1829
+ (
1830
+ now,
1831
+ row["job_id"],
1832
+ row["claim_token"],
1833
+ row["claim_owner"],
1834
+ now,
1835
+ ),
1836
+ )
1837
+ if released.rowcount != 1:
1838
+ raise SlowGraphError("expired job claim changed during recovery")
1839
+ expired_attempt = con.execute(
1840
+ "UPDATE slow_graph_attempts SET status='expired',"
1841
+ "error='claim lease expired; external call outcome uncertain; explicit resume required',completed_at=? "
1842
+ "WHERE job_id=? AND claim_token=? AND claim_owner=? "
1843
+ "AND status='started'",
1844
+ (now, row["job_id"], row["claim_token"], row["claim_owner"]),
1845
+ )
1846
+ if expired_attempt.rowcount != 1:
1847
+ raise SlowGraphError("expired job has no active claimed attempt")
1848
+ return len(rows)
1849
+
1850
+ def drain(
1851
+ self, manager: PatchManager, *, batch_size: int | None = None
1852
+ ) -> list[str]:
1853
+ self.recover_interrupted_attempts()
1854
+ if batch_size is not None:
1855
+ if batch_size <= 0:
1856
+ raise SlowGraphError("batch_size must be positive")
1857
+ owner = self._claim_owner()
1858
+ result: list[str] = []
1859
+ while batch_size is None or len(result) < batch_size:
1860
+ claim = self._claim_pending_job(None, owner=owner)
1861
+ if claim is None:
1862
+ break
1863
+ result.append(self._run_claimed_job(claim, manager))
1864
+ return result
1865
+
1866
+ def audit(self, scope_id: str) -> dict[str, int]:
1867
+ with self.connection() as con:
1868
+ self._records_table_exists(con)
1869
+ unfinished = con.execute(
1870
+ "SELECT job_id,status,last_error FROM slow_graph_jobs WHERE scope_id=? AND status!='completed' ORDER BY created_at",
1871
+ (scope_id,),
1872
+ ).fetchall()
1873
+ if unfinished:
1874
+ raise AuditError(
1875
+ "scope has unfinished slow graph jobs: "
1876
+ + _json([dict(item) for item in unfinished])
1877
+ )
1878
+ rows = con.execute(
1879
+ "SELECT memory_id,state,metadata_json FROM records WHERE scope_id=?",
1880
+ (scope_id,),
1881
+ ).fetchall()
1882
+ fast_state: dict[str, tuple[str, str]] = {}
1883
+ active_fast_by_slot: dict[str, set[str]] = {}
1884
+ for row in rows:
1885
+ row_meta = self._metadata(row, "record")
1886
+ if (
1887
+ row_meta.get("content_variant") == LEAF_VARIANT
1888
+ and row_meta.get("memory_layer") == "fast"
1889
+ ):
1890
+ slot = _clean(row_meta.get("canonical_slot_key"))
1891
+ state = _clean(row["state"])
1892
+ fast_state[str(row["memory_id"])] = (state, slot)
1893
+ if state in {"active", "parallel_active", "promoted"}:
1894
+ active_fast_by_slot.setdefault(slot, set()).add(
1895
+ str(row["memory_id"])
1896
+ )
1897
+ capsules: dict[str, list[tuple[str, str, int, str]]] = {}
1898
+ for row in rows:
1899
+ meta = self._metadata(row, "record")
1900
+ if meta.get("content_variant") != CAPSULE_VARIANT:
1901
+ continue
1902
+ capsule_id = _required_text(meta.get("capsule_id"), "capsule_id")
1903
+ revision = meta.get("revision")
1904
+ if not isinstance(revision, int) or revision < 1:
1905
+ raise AuditError("capsule revision is invalid")
1906
+ status = _required_text(meta.get("status"), "capsule status")
1907
+ if status not in {"active", "challenged", "retired", "superseded"}:
1908
+ raise AuditError("capsule status is invalid")
1909
+ capsules.setdefault(capsule_id, []).append(
1910
+ (str(row["memory_id"]), str(row["state"]), revision, status)
1911
+ )
1912
+ claims = _validate_claims(meta.get("claims"), stored=True)
1913
+ for claim in claims:
1914
+ if status in {"active", "challenged"}:
1915
+ stale_support = {
1916
+ evidence_id
1917
+ for evidence_id in claim["support"]
1918
+ if fast_state.get(evidence_id, ("", ""))[0]
1919
+ not in {"active", "parallel_active", "promoted"}
1920
+ }
1921
+ current_replacements = active_fast_by_slot.get(
1922
+ claim["canonical_slot"], set()
1923
+ )
1924
+ if (
1925
+ stale_support
1926
+ and current_replacements
1927
+ and not current_replacements.intersection(
1928
+ claim["counterevidence"]
1929
+ )
1930
+ ):
1931
+ raise AuditError(
1932
+ "active capsule claim promotes superseded fast evidence"
1933
+ )
1934
+ evidence = claim["support"] + claim["counterevidence"]
1935
+ provenance = con.execute(
1936
+ "SELECT evidence_memory_id,source_parent_json FROM slow_graph_provenance WHERE scope_id=? AND capsule_id=? AND revision=? AND claim_id=?",
1937
+ (scope_id, capsule_id, revision, claim["claim_id"]),
1938
+ ).fetchall()
1939
+ if {item["evidence_memory_id"] for item in provenance} != set(
1940
+ evidence
1941
+ ):
1942
+ raise AuditError(
1943
+ "claim provenance does not match claim evidence"
1944
+ )
1945
+ for item in provenance:
1946
+ parent = _strict_json(
1947
+ item["source_parent_json"],
1948
+ label="provenance source_parent",
1949
+ expected=dict,
1950
+ )
1951
+ if set(parent) != {
1952
+ "session_index",
1953
+ "parent_chunk_index",
1954
+ "message_index",
1955
+ "source_record_id",
1956
+ "event_id",
1957
+ "evidence_char_start",
1958
+ "evidence_char_end",
1959
+ }:
1960
+ raise AuditError("provenance is not a leaf source_parent")
1961
+ leaf = con.execute(
1962
+ "SELECT metadata_json FROM records WHERE scope_id=? AND memory_id=?",
1963
+ (scope_id, item["evidence_memory_id"]),
1964
+ ).fetchone()
1965
+ if leaf is None:
1966
+ raise AuditError("provenance references missing evidence")
1967
+ leaf_meta = self._metadata(leaf, "provenance leaf")
1968
+ expected_parent = {
1969
+ "session_index": leaf_meta.get("session_index"),
1970
+ "parent_chunk_index": leaf_meta.get("message_index"),
1971
+ "message_index": leaf_meta.get("message_index"),
1972
+ "source_record_id": leaf_meta.get("source_record_id"),
1973
+ "event_id": leaf_meta.get("event_id"),
1974
+ "evidence_char_start": leaf_meta.get(
1975
+ "evidence_char_start"
1976
+ ),
1977
+ "evidence_char_end": leaf_meta.get("evidence_char_end"),
1978
+ }
1979
+ if (
1980
+ leaf_meta.get("content_variant") != LEAF_VARIANT
1981
+ or leaf_meta.get("memory_layer") != "fast"
1982
+ or leaf_meta.get("node_kind") != "atomic_user_assertion"
1983
+ or leaf_meta.get("atomic_evidence_leaf") is not True
1984
+ or leaf_meta.get("authority") != "user_assertion"
1985
+ or expected_parent != parent
1986
+ ):
1987
+ raise AuditError(
1988
+ "provenance does not resolve to the cited fast leaf"
1989
+ )
1990
+ for capsule_id, revisions in capsules.items():
1991
+ revision_numbers = sorted(item[2] for item in revisions)
1992
+ if revision_numbers != list(range(1, max(revision_numbers) + 1)):
1993
+ raise AuditError("capsule revisions are not contiguous")
1994
+ latest = [item for item in revisions if item[2] == max(revision_numbers)]
1995
+ if len(latest) != 1:
1996
+ raise AuditError("capsule lacks one latest revision")
1997
+ memory_id, state, _, status = latest[0]
1998
+ slot_key = "slow." + capsule_id
1999
+ head = con.execute(
2000
+ "SELECT memory_id FROM slot_heads WHERE scope_id=? AND slot_key=?",
2001
+ (scope_id, slot_key),
2002
+ ).fetchone()
2003
+ if status in {"active", "challenged"}:
2004
+ if state != "active" or head is None or head["memory_id"] != memory_id:
2005
+ raise AuditError("active/challenged capsule head is inconsistent")
2006
+ elif head is not None:
2007
+ raise AuditError("inactive capsule must not retain a slot head")
2008
+ history = [
2009
+ item["memory_id"]
2010
+ for item in con.execute(
2011
+ "SELECT memory_id FROM slot_history WHERE scope_id=? AND slot_key=? ORDER BY ordinal",
2012
+ (scope_id, slot_key),
2013
+ )
2014
+ ]
2015
+ expected_history = [
2016
+ item[0] for item in sorted(revisions, key=lambda item: item[2])
2017
+ ]
2018
+ if history != expected_history:
2019
+ raise AuditError("capsule slot history is inconsistent")
2020
+ edge_rows = con.execute(
2021
+ "SELECT source_memory_id,target_memory_id,edge_type,metadata_json FROM memory_edges WHERE scope_id=?",
2022
+ (scope_id,),
2023
+ ).fetchall()
2024
+ for edge in edge_rows:
2025
+ meta = _strict_json(
2026
+ edge["metadata_json"], label="edge metadata", expected=dict
2027
+ )
2028
+ if meta.get("edge_source") != SLOW_EDGE_SOURCE:
2029
+ continue
2030
+ if (
2031
+ edge["source_memory_id"] == edge["target_memory_id"]
2032
+ or edge["edge_type"] not in EDGE_TYPES
2033
+ or not isinstance(meta.get("evidence_refs"), list)
2034
+ or not _clean(meta.get("patch_id"))
2035
+ ):
2036
+ raise AuditError("edge is inconsistent")
2037
+ return {
2038
+ "slow_graph_batches": int(
2039
+ con.execute(
2040
+ "SELECT COUNT(*) FROM slow_graph_batches WHERE scope_id=?",
2041
+ (scope_id,),
2042
+ ).fetchone()[0]
2043
+ ),
2044
+ "slow_graph_jobs": int(
2045
+ con.execute(
2046
+ "SELECT COUNT(*) FROM slow_graph_jobs WHERE scope_id=?",
2047
+ (scope_id,),
2048
+ ).fetchone()[0]
2049
+ ),
2050
+ "slow_graph_patches": int(
2051
+ con.execute(
2052
+ "SELECT COUNT(*) FROM slow_graph_patches WHERE scope_id=?",
2053
+ (scope_id,),
2054
+ ).fetchone()[0]
2055
+ ),
2056
+ "slow_graph_attempts": int(
2057
+ con.execute(
2058
+ "SELECT COUNT(*) FROM slow_graph_attempts WHERE scope_id=?",
2059
+ (scope_id,),
2060
+ ).fetchone()[0]
2061
+ ),
2062
+ "slow_graph_patch_operations": int(
2063
+ con.execute(
2064
+ "SELECT COUNT(*) FROM slow_graph_patch_operations o JOIN slow_graph_patches p ON p.patch_id=o.patch_id WHERE p.scope_id=?",
2065
+ (scope_id,),
2066
+ ).fetchone()[0]
2067
+ ),
2068
+ "slow_graph_provenance": int(
2069
+ con.execute(
2070
+ "SELECT COUNT(*) FROM slow_graph_provenance WHERE scope_id=?",
2071
+ (scope_id,),
2072
+ ).fetchone()[0]
2073
+ ),
2074
+ "memory_edges": sum(
2075
+ 1
2076
+ for edge in edge_rows
2077
+ if _strict_json(
2078
+ edge["metadata_json"], label="edge metadata", expected=dict
2079
+ ).get("edge_source")
2080
+ == SLOW_EDGE_SOURCE
2081
+ ),
2082
+ }
2083
+
2084
+
2085
+ def main() -> None:
2086
+ parser = argparse.ArgumentParser(description="TMCRA slow graph controller")
2087
+ parser.add_argument("database", type=Path)
2088
+ parser.add_argument(
2089
+ "--repo",
2090
+ type=Path,
2091
+ required=True,
2092
+ help="TMCRA repository containing the real graph schema",
2093
+ )
2094
+ sub = parser.add_subparsers(dest="command", required=True)
2095
+ enqueue = sub.add_parser("enqueue")
2096
+ enqueue.add_argument("scope_id")
2097
+ enqueue.add_argument("--region")
2098
+ drain = sub.add_parser("drain")
2099
+ drain.add_argument("--batch-size", type=int)
2100
+ run = sub.add_parser("run")
2101
+ run.add_argument("job_id")
2102
+ resume = sub.add_parser("resume")
2103
+ resume.add_argument("job_id")
2104
+ audit = sub.add_parser("audit")
2105
+ audit.add_argument("scope_id")
2106
+ args = parser.parse_args()
2107
+ store = SlowGraphStore(args.database, schema=load_graph_schema(args.repo))
2108
+ if args.command == "enqueue":
2109
+ manager = DeepSeekProGraphPatchManager(DeepSeekProConfig.from_env())
2110
+ if args.region:
2111
+ region = store.fast_regions(args.scope_id).get(args.region, [])
2112
+ result = [
2113
+ store.enqueue(
2114
+ args.scope_id,
2115
+ args.region,
2116
+ (item["memory_id"] for item in region),
2117
+ manager=manager,
2118
+ )
2119
+ ]
2120
+ else:
2121
+ result = store.enqueue_regions(args.scope_id, manager=manager)
2122
+ elif args.command == "resume":
2123
+ store.resume(args.job_id)
2124
+ result = {"job_id": args.job_id, "status": "pending"}
2125
+ elif args.command == "audit":
2126
+ result = store.audit(args.scope_id)
2127
+ else:
2128
+ manager = DeepSeekProGraphPatchManager(DeepSeekProConfig.from_env())
2129
+ result = (
2130
+ store.drain(manager, batch_size=args.batch_size)
2131
+ if args.command == "drain"
2132
+ else store.run_job(args.job_id, manager)
2133
+ )
2134
+ print(_json(result))
2135
+
2136
+
2137
+ if __name__ == "__main__":
2138
+ main()
benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v4_batch_writer.py ADDED
The diff for this file is too large to render. See raw diff