Release TMCRA Agent Memory v2026.07.18
Browse filesPublish 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
- AUTHORS.md +25 -0
- CHANGELOG.md +16 -0
- CITATION.cff +27 -0
- LICENSE +201 -21
- NOTICE +9 -0
- README.md +81 -173
- README.zh-CN.md +66 -173
- SECURITY.md +12 -0
- benchmarks/longmemeval/README.md +216 -0
- benchmarks/longmemeval/README.zh-CN.md +216 -0
- benchmarks/longmemeval/THIRD_PARTY_NOTICES.md +52 -0
- benchmarks/longmemeval/configs/answer.env.example +6 -0
- benchmarks/longmemeval/configs/assets.lock.json +51 -0
- benchmarks/longmemeval/configs/judge.env.example +9 -0
- benchmarks/longmemeval/configs/run.env.example +18 -0
- benchmarks/longmemeval/configs/writer.env.example +12 -0
- benchmarks/longmemeval/fixtures/tiny_longmemeval.json +22 -0
- benchmarks/longmemeval/pyproject.toml +30 -0
- benchmarks/longmemeval/requirements.txt +8 -0
- benchmarks/longmemeval/results/benchmark.json +16 -0
- benchmarks/longmemeval/scripts/fetch_assets.py +95 -0
- benchmarks/longmemeval/scripts/prepare_longmemeval.py +70 -0
- benchmarks/longmemeval/scripts/reproduce_s500.sh +152 -0
- benchmarks/longmemeval/scripts/reproduce_smoke.sh +18 -0
- benchmarks/longmemeval/scripts/summarize_official_judge.py +82 -0
- benchmarks/longmemeval/scripts/validate_stage.py +237 -0
- benchmarks/longmemeval/src/tmcra_benchmark/__init__.py +3 -0
- benchmarks/longmemeval/src/tmcra_benchmark/harness/__init__.py +1 -0
- benchmarks/longmemeval/src/tmcra_benchmark/harness/evaluate_qa_official_judge.py +295 -0
- benchmarks/longmemeval/src/tmcra_benchmark/harness/run_lme_s10_native_tmcra.py +0 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/__init__.py +1 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/audit_tmcra_v4_chain.py +0 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/build_v3_runtime_dataset.py +791 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/ops/audit_tmcra_v4_subject_attribution.py +877 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/prepare_tmcra_v4_e2e_data.py +261 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_build.py +984 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_compile_evidence.py +1015 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_evaluate.py +725 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_gpt54_answers.py +847 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_parallel_official_judge.py +273 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/run_tmcra_v4_retrieve.py +524 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_model_assets.py +29 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v2_lme_pipeline.py +855 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_online_runtime.py +1867 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_product_writer.py +0 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_recall_planner.py +406 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_reranker.py +1014 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_schema.py +168 -0
- benchmarks/longmemeval/src/tmcra_benchmark/legacy/tmcra_v3_slow_graph.py +2138 -0
- 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 |
-
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2 |
|
| 3 |
<p align="center">
|
| 4 |
<img src="assets/tmcra-logo.png" alt="TMCRA logo" width="420">
|
| 5 |
</p>
|
| 6 |
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 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 |
-
|
| 22 |
|
| 23 |
-
|
| 24 |
|
| 25 |
-
|
| 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 |
-
|
| 34 |
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
|
|
|
|
|
|
| 42 |
|
| 43 |
-
|
| 44 |
-
D --> G
|
| 45 |
-
E --> G
|
| 46 |
-
F --> G
|
| 47 |
|
| 48 |
-
|
| 49 |
-
G --> I["Learned Path Scorer"]
|
| 50 |
-
H --> J["Evidence Selection"]
|
| 51 |
-
I --> J
|
| 52 |
|
| 53 |
-
|
| 54 |
-
K --> L["Answer Model"]
|
| 55 |
-
L --> M["Agent Response"]
|
| 56 |
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 61 |
```
|
| 62 |
|
| 63 |
-
The
|
| 64 |
|
| 65 |
-
##
|
| 66 |
|
| 67 |
-
|
| 68 |
|
| 69 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 70 |
|
| 71 |
-
|
| 72 |
|
| 73 |
-
|
| 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 |
-
A typical runtime configuration points TMCRA to these weights:
|
| 88 |
|
| 89 |
```bash
|
| 90 |
-
|
| 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 |
-
|
| 99 |
-
|
| 100 |
-
|
| 101 |
-
|
| 102 |
-
The core adapter snapshot is:
|
| 103 |
-
|
| 104 |
-
```text
|
| 105 |
-
code/memory_adapters.py
|
| 106 |
```
|
| 107 |
|
| 108 |
-
For
|
| 109 |
-
|
| 110 |
-
## Dependency Environment
|
| 111 |
|
| 112 |
-
|
| 113 |
|
| 114 |
-
-
|
| 115 |
-
-
|
| 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 |
-
|
| 121 |
|
| 122 |
-
|
|
|
|
| 123 |
|
| 124 |
-
|
| 125 |
|
| 126 |
-
|
| 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 |
-
|
| 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 |
-
|
| 8 |
-
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
许可证:MIT。
|
| 14 |
-
|
| 15 |
-
## TMCRA 的作用
|
| 16 |
-
|
| 17 |
-
TMCRA 在 Agent 应用和回答模型之间增加一层专门的记忆运行时。
|
| 18 |
-
|
| 19 |
-
在写入阶段,TMCRA 会把对话转成记忆节点、事件单元、profile 信号和图路径。这让系统不仅能保存孤立事实,也能保存不同轮次、不同会话之间的事实关系。
|
| 20 |
|
| 21 |
-
|
| 22 |
|
| 23 |
-
|
| 24 |
|
| 25 |
-
|
| 26 |
-
- 助手回答细节记忆
|
| 27 |
-
- profile 与偏好记忆
|
| 28 |
-
- 时间记忆
|
| 29 |
-
- 跨会话图隧穿
|
| 30 |
-
- 学习式节点/路径打分
|
| 31 |
-
- 面向下游 LLM 的紧凑证据选择
|
| 32 |
|
| 33 |
-
|
| 34 |
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
|
|
|
|
|
|
|
| 42 |
|
| 43 |
-
|
| 44 |
-
D --> G
|
| 45 |
-
E --> G
|
| 46 |
-
F --> G
|
| 47 |
|
| 48 |
-
|
| 49 |
-
G --> I["学习式路径 Scorer"]
|
| 50 |
-
H --> J["证据选择"]
|
| 51 |
-
I --> J
|
| 52 |
|
| 53 |
-
|
| 54 |
-
K --> L["回答模型"]
|
| 55 |
-
L --> M["Agent 回复"]
|
| 56 |
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 61 |
```
|
| 62 |
|
| 63 |
-
|
| 64 |
|
| 65 |
-
##
|
| 66 |
|
| 67 |
-
|
| 68 |
|
| 69 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 70 |
|
| 71 |
-
|
| 72 |
|
| 73 |
-
|
| 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 |
-
|
| 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 |
-
|
| 99 |
-
|
| 100 |
-
|
| 101 |
-
|
| 102 |
-
核心适配器代码快照:
|
| 103 |
-
|
| 104 |
-
```text
|
| 105 |
-
code/memory_adapters.py
|
| 106 |
```
|
| 107 |
|
| 108 |
-
|
| 109 |
-
|
| 110 |
-
## 依赖环境
|
| 111 |
|
| 112 |
-
|
| 113 |
|
| 114 |
-
-
|
| 115 |
-
-
|
| 116 |
-
- NumPy 及常见 Python 数据处理库
|
| 117 |
-
- 用于回答层和写入层的 OpenAI 兼容接口或本地 LLM endpoint
|
| 118 |
-
- 从 GitHub 拉取完整模型包时建议支持 Git LFS
|
| 119 |
|
| 120 |
-
|
| 121 |
|
| 122 |
-
|
|
|
|
| 123 |
|
| 124 |
-
|
| 125 |
|
| 126 |
-
|
| 127 |
-
- **LLM planner 接入口**:可选的规划 hook,可以调用外部 LLM 来整理证据、扩展 query intent,或在最终回答前生成 answer plan。它适合实验和更高成本的部署场景;默认基线仍保持核心图 scorer 链路可以独立测量。
|
| 128 |
|
| 129 |
-
|
| 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
|
|
|