feat(compression): type-aware compression module + retrieve_v2 integration (#434) - #493
Merged
Merged
Conversation
Adds src/aelfrice/compression.py with the spec-pinned strategy table: fact + unknown render verbatim, snapshot → headline (when not locked), transient → stub (when not locked), locks always render verbatim. Headline extraction splits on the first '. ' or '.\n' outside a balanced ``` fence; falls back to whitespace-aligned hard truncation at MAX_HEADLINE_CHARS (240). Content that is wholly a single code fence renders verbatim (impl-PR resolution of spec open-question 3). Stub falls back to verbatim when the marker would cost more than the source. Pure deterministic — no store, clock, env, or random reads. Token estimator duplicated from retrieval._estimate_tokens to avoid a circular import (retrieval will consume this module in a follow-up). 18 unit tests cover the strategy table, headline edge cases (short content, no boundary, code-fence-only, in-fence period), the locked-override rule, the token-monotone-non-increasing invariant, and deterministic re-evaluation.
Adds the use_type_aware_compression flag with the established 4-stage precedence (env > kwarg > TOML > default-OFF) and the `AELFRICE_TYPE_AWARE_COMPRESSION` env override. Default-OFF preserves byte-identical v1.x adapter behavior at v2.0.0; the bench gate (A2 + A4 in docs/feature-type-aware-compression.md) flips the default after lab-side benchmark evidence clears. RetrievalResult gains a parallel `compressed_beliefs` field (list[CompressedBelief]). Same length and order as `beliefs` when the flag resolves True; empty list when OFF. Consumers that want raw beliefs keep reading `.beliefs`; consumers that want the type-aware-compressed render read `.compressed_beliefs[i].rendered`. Locked-state classification: `belief.lock_level == LOCK_USER`. This mirrors the L0-never-trimmed rule the rest of retrieval already applies (retrieval.py:950). Compression is applied after the existing pack and the optional temporal_sort pass. This commit does not rewrite the budget-accounting pack loops — the parallel-field shape is the minimum viable change called out in the spec at the "RetrievalResult shape" section. The actual recall@k uplift gated by A2 requires the budget rewrite plus the lab corpus bench cut; both are follow-up work. 10 integration tests cover flag precedence (default OFF, kwarg override, env override, env-garbage fall-through, TOML resolution), the parallel-field invariant (same length / order as beliefs), and the strategy-dispatch path through retrieve_v2 (fact verbatim, snapshot headline, transient stub, locked verbatim). Full suite: 2840 passed, 22 skipped, no regressions.
Adds the lab-mounted bench gate for #434. Skips on public CI (autouse `bench_gated` marker) and skips again when the `tests/corpus/v2_0/compression_uplift/` directory is empty. Measures the upstream invariant A2 depends on: on a mixed-retention-class corpus, compressed total tokens must be strictly less than uncompressed total tokens. If compression doesn't reduce cost on the lab corpus, no pack-loop rewrite can deliver A2's recall@k uplift no matter how it is wired. A2's strict recall@k comparison requires the budget-rewrite follow-up that consumes `compressed_beliefs[*].rendered_tokens` in the pack loops. This gate is the precursor measurement. Row schema documented in the module docstring. Mirrors the corpus contract for #197 (dedup) and other bench-gated modules.
Adds the v2.1 type-aware-compression flag to the [retrieval] section header, the .aelfrice.toml example, and the dedicated Keys subsection. Mirrors the documentation pattern established for use_bm25f_anchors and use_hrr_structural. Includes the strategy table for operator reference (fact/snapshot/ transient/unknown × locked/unlocked) and the precedence chain for the flag (env > kwarg > TOML > default-OFF). Notes that the default-on flip is gated on both the lab-side compression_uplift bench and the pack-loop budget rewrite that turns the parallel-field shape into recall@k uplift.
Reviewer's GuideImplements a deterministic, type-aware belief compression module and wires it into retrieve_v2 behind a default-OFF, multi-source flag, adding tests, a bench gate, and configuration docs while preserving existing RetrievalResult behavior when the feature is disabled. Sequence diagram for retrieve_v2 with type-aware compression flagsequenceDiagram
actor Caller
participant Retrieval as retrieve_v2
participant Env as _env_type_aware_compression_override
participant Config as _read_toml_flag_for
participant Compress as compress_for_retrieval
participant Result as RetrievalResult
Caller->>Retrieval: retrieve_v2(query, ..., use_type_aware_compression)
activate Retrieval
note over Retrieval: Resolve use_type_aware_compression
Retrieval->>Env: _env_type_aware_compression_override()
activate Env
Env-->>Retrieval: env_value or None
deactivate Env
alt env_value is not None
Retrieval-->>Retrieval: flag = env_value
else env_value is None
alt explicit kwarg is not None
Retrieval-->>Retrieval: flag = explicit
else explicit is None
Retrieval->>Config: _read_toml_flag_for(TYPE_AWARE_COMPRESSION_FLAG, start)
activate Config
Config-->>Retrieval: toml_value or None
deactivate Config
alt toml_value is not None
Retrieval-->>Retrieval: flag = toml_value
else no decisive source
Retrieval-->>Retrieval: flag = False
end
end
end
note over Retrieval: run retrieval pipeline to produce beliefs
Retrieval-->>Retrieval: beliefs = [...]
alt flag is True
loop for each belief in beliefs
Retrieval->>Compress: compress_for_retrieval(belief, locked=(belief.lock_level == LOCK_USER))
activate Compress
Compress-->>Retrieval: CompressedBelief
deactivate Compress
end
Retrieval-->>Result: construct RetrievalResult(beliefs, compressed_beliefs=[...], ...)
else flag is False
Retrieval-->>Result: construct RetrievalResult(beliefs, compressed_beliefs=[], ...)
end
Result-->>Caller: RetrievalResult
deactivate Retrieval
Class diagram for type-aware compression and RetrievalResult integrationclassDiagram
class Belief {
+str id
+str retention_class
+str content
+str lock_level
}
class CompressedBelief {
+Belief belief
+str rendered
+int rendered_tokens
+str strategy
}
class RetrievalResult {
+list~Belief~ beliefs
+list~CompressedBelief~ compressed_beliefs
+list~str~ entity_hits
+list~str~ locked_ids
+list~str~ l1_ids
+list bfs_chains
}
class compression_module {
+CompressedBelief compress_for_retrieval(belief, locked)
+int _estimate_tokens(text)
+str _headline(content)
+str _stub(belief)
}
class retrieval_module {
+RetrievalResult retrieve_v2(query, bm25f_cache, temporal_sort, temporal_half_life_seconds, use_type_aware_compression)
+bool resolve_use_type_aware_compression(explicit, start)
+bool _env_type_aware_compression_override()
+bool _read_toml_flag_for(key, start)
}
Belief "1" <-- "1" CompressedBelief : wraps
RetrievalResult "*" o-- "*" Belief : beliefs
RetrievalResult "*" o-- "*" CompressedBelief : compressed_beliefs
compression_module ..> Belief : uses
retrieval_module ..> Belief : uses
retrieval_module ..> CompressedBelief : populates
retrieval_module ..> compression_module : calls
Flow diagram for resolving use_type_aware_compression flagflowchart TD
Start([Start resolve_use_type_aware_compression])
Env["Read env AELFRICE_TYPE_AWARE_COMPRESSION"]
EnvDec{Env value<br/>truthy/falsy?}
ExplicitDec{explicit kwarg<br/>is not None?}
Toml["Read TOML [retrieval] use_type_aware_compression"]
TomlDec{TOML value<br/>is not None?}
EnvTrue[Set flag = True]
EnvFalse[Set flag = False]
UseExplicit[Set flag = explicit]
UseToml[Set flag = TOML value]
UseDefault[Set flag = False<br/>default OFF]
End([Return flag])
Start --> Env
Env --> EnvDec
EnvDec -- Yes, truthy --> EnvTrue --> End
EnvDec -- Yes, falsy --> EnvFalse --> End
EnvDec -- No, unset/invalid --> ExplicitDec
ExplicitDec -- Yes --> UseExplicit --> End
ExplicitDec -- No --> Toml
Toml --> TomlDec
TomlDec -- Yes --> UseToml --> End
TomlDec -- No --> UseDefault --> End
File-Level Changes
Possibly linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
robotrocketscience
added
author-Setr
📝 WalkthroughWalkthroughThis PR implements a complete type-aware belief compression feature for the aelfrice retrieval system. It introduces a new ChangesType-Aware Compression Feature
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related issues
Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've found 2 issues, and left some high level feedback:
- The token estimator constants and logic (
_CHARS_PER_TOKEN,_estimate_tokens) are duplicated betweenretrieval.pyandcompression.py; consider centralising them in a small shared module to avoid silent drift if you ever tune the heuristic. - For
ENV_TYPE_AWARE_COMPRESSION, garbage values currently fail open and silently fall through; you might want to log or warn on unrecognised values so misconfigured environments are easier to diagnose.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The token estimator constants and logic (`_CHARS_PER_TOKEN`, `_estimate_tokens`) are duplicated between `retrieval.py` and `compression.py`; consider centralising them in a small shared module to avoid silent drift if you ever tune the heuristic.
- For `ENV_TYPE_AWARE_COMPRESSION`, garbage values currently fail open and silently fall through; you might want to log or warn on unrecognised values so misconfigured environments are easier to diagnose.
## Individual Comments
### Comment 1
<location path="src/aelfrice/compression.py" line_range="49-58" />
<code_context>
+
+
+@dataclass(frozen=True)
+class CompressedBelief:
+ """A `Belief` plus its packed-render form.
+
+ `rendered_tokens` is monotone-non-increasing in `_estimate_tokens(belief.content)`:
+ the compressor never produces a render that costs more than the source.
+ """
+
+ belief: Belief
+ rendered: str
+ rendered_tokens: int
+ strategy: str
+
+
</code_context>
<issue_to_address>
**suggestion:** Tighten `strategy` typing to avoid misuse and ease refactors.
`strategy` only permits three values (`STRATEGY_VERBATIM`, `STRATEGY_HEADLINE`, `STRATEGY_STUB`) but is typed as `str`. Using a `Literal[...]` or small `Enum` would prevent invalid values and give stronger type-checker/editor support as more strategies are added.
Suggested implementation:
```python
StrategyName = Literal["verbatim", "headline", "stub"]
@dataclass(frozen=True)
class CompressedBelief:
"""A `Belief` plus its packed-render form.
`rendered_tokens` is monotone-non-increasing in `_estimate_tokens(belief.content)`:
the compressor never produces a render that costs more than the source.
"""
belief: Belief
rendered: str
rendered_tokens: int
strategy: StrategyName
```
1. Ensure `Literal` is imported from `typing`. For example, if you currently have `from typing import Final`, update it to `from typing import Final, Literal`. If the import line differs, adjust accordingly.
2. If there are other places in this file (or callers) that construct `CompressedBelief` with `strategy` values, verify they only use `"verbatim"`, `"headline"`, or `"stub"` (or the corresponding constants `STRATEGY_VERBATIM`, `STRATEGY_HEADLINE`, `STRATEGY_STUB` whose values must be these strings) to keep mypy and other type-checkers happy.
</issue_to_address>
### Comment 2
<location path="tests/test_compression.py" line_range="90" />
<code_context>
+ assert cb.rendered == b.content
+
+
+def test_snapshot_unlocked_takes_headline() -> None:
+ content = (
+ "the first sentence is the headline. "
+ "the second sentence and rest of body should be dropped. "
+ "and a third sentence."
+ )
+ b = _mk(content, retention_class=RETENTION_SNAPSHOT)
+ cb = compress_for_retrieval(b, locked=False)
+ assert cb.strategy == STRATEGY_HEADLINE
+ assert cb.rendered.startswith("the first sentence is the headline")
+ assert cb.rendered.endswith("…")
+ assert cb.rendered_tokens < _estimate_tokens(content)
+
+
</code_context>
<issue_to_address>
**suggestion (testing):** Add a test that snapshot headline with a terminal period but no trailing body does not get an ellipsis appended
The `_headline` branch that returns the prefix without an ellipsis when `sentence_end >= len(content)` isn’t covered. Please add a test for a snapshot where the headline is a complete single sentence with a trailing period and no further text (e.g. `"single sentence only."`), and assert `cb.rendered == content` so this behavior is locked in and protected against regressions in the splitting logic.
```suggestion
assert cb.rendered_tokens < _estimate_tokens(content)
def test_snapshot_headline_single_sentence_no_ellipsis() -> None:
content = "single sentence only."
b = _mk(content, retention_class=RETENTION_SNAPSHOT)
cb = compress_for_retrieval(b, locked=False)
assert cb.strategy == STRATEGY_HEADLINE
assert cb.rendered == content
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+49
to
+58
| class CompressedBelief: | ||
| """A `Belief` plus its packed-render form. | ||
|
|
||
| `rendered_tokens` is monotone-non-increasing in `_estimate_tokens(belief.content)`: | ||
| the compressor never produces a render that costs more than the source. | ||
| """ | ||
|
|
||
| belief: Belief | ||
| rendered: str | ||
| rendered_tokens: int |
There was a problem hiding this comment.
suggestion: Tighten strategy typing to avoid misuse and ease refactors.
strategy only permits three values (STRATEGY_VERBATIM, STRATEGY_HEADLINE, STRATEGY_STUB) but is typed as str. Using a Literal[...] or small Enum would prevent invalid values and give stronger type-checker/editor support as more strategies are added.
Suggested implementation:
StrategyName = Literal["verbatim", "headline", "stub"]
@dataclass(frozen=True)
class CompressedBelief:
"""A `Belief` plus its packed-render form.
`rendered_tokens` is monotone-non-increasing in `_estimate_tokens(belief.content)`:
the compressor never produces a render that costs more than the source.
"""
belief: Belief
rendered: str
rendered_tokens: int
strategy: StrategyName- Ensure
Literalis imported fromtyping. For example, if you currently havefrom typing import Final, update it tofrom typing import Final, Literal. If the import line differs, adjust accordingly. - If there are other places in this file (or callers) that construct
CompressedBeliefwithstrategyvalues, verify they only use"verbatim","headline", or"stub"(or the corresponding constantsSTRATEGY_VERBATIM,STRATEGY_HEADLINE,STRATEGY_STUBwhose values must be these strings) to keep mypy and other type-checkers happy.
| assert cb.strategy == STRATEGY_HEADLINE | ||
| assert cb.rendered.startswith("the first sentence is the headline") | ||
| assert cb.rendered.endswith("…") | ||
| assert cb.rendered_tokens < _estimate_tokens(content) |
There was a problem hiding this comment.
suggestion (testing): Add a test that snapshot headline with a terminal period but no trailing body does not get an ellipsis appended
The _headline branch that returns the prefix without an ellipsis when sentence_end >= len(content) isn’t covered. Please add a test for a snapshot where the headline is a complete single sentence with a trailing period and no further text (e.g. "single sentence only."), and assert cb.rendered == content so this behavior is locked in and protected against regressions in the splitting logic.
Suggested change
| assert cb.rendered_tokens < _estimate_tokens(content) | |
| assert cb.rendered_tokens < _estimate_tokens(content) | |
| def test_snapshot_headline_single_sentence_no_ellipsis() -> None: | |
| content = "single sentence only." | |
| b = _mk(content, retention_class=RETENTION_SNAPSHOT) | |
| cb = compress_for_retrieval(b, locked=False) | |
| assert cb.strategy == STRATEGY_HEADLINE | |
| assert cb.rendered == content |
Owner
Author
|
[claim:review:Gylf:2026-05-08T18:36:48Z] |
Owner
Author
|
[claim:review:Toug:2026-05-08T18:36:56Z] |
Owner
Author
|
[release:review:Toug:2026-05-08T18:37:01Z] |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| """ | ||
| from __future__ import annotations | ||
|
|
||
| import pytest |
| """ | ||
| from __future__ import annotations | ||
|
|
||
| import os |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/aelfrice/retrieval.py (1)
1451-1451:⚠️ Potential issue | 🟠 Major | ⚡ Quick winFix merge-blocking deadcode on
use_hrr.
use_hrris intentionally forward-compat, but currently unused and is tripping vulture in CI.Proposed minimal fix
def retrieve_v2( @@ ) -> RetrievalResult: @@ + _ = use_hrr # reserved for forward-compat; keeps deadcode checks green ( out, locked_ids_list,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/aelfrice/retrieval.py` at line 1451, The parameter use_hrr is declared but unused and causing CI deadcode warnings; mark it as used without changing behavior by adding a no-op reference (for example, add a single line like "_ = use_hrr" or "del use_hrr # forward-compat" at the top of the function that contains the use_hrr parameter) so tools like vulture no longer flag it; locate the function whose signature includes "use_hrr: bool = False" and add the no-op reference there.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@src/aelfrice/retrieval.py`:
- Line 1451: The parameter use_hrr is declared but unused and causing CI
deadcode warnings; mark it as used without changing behavior by adding a no-op
reference (for example, add a single line like "_ = use_hrr" or "del use_hrr #
forward-compat" at the top of the function that contains the use_hrr parameter)
so tools like vulture no longer flag it; locate the function whose signature
includes "use_hrr: bool = False" and add the no-op reference there.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: ef145d85-16da-40f1-901e-9c47bbfdaed8
📒 Files selected for processing (6)
docs/CONFIG.mdsrc/aelfrice/compression.pysrc/aelfrice/retrieval.pytests/bench_gate/test_compression_uplift.pytests/test_compression.pytests/test_compression_integration.py
Owner
Author
|
[release:review:Gylf:2026-05-08T18:40:38Z] |
This was referenced May 8, 2026
This was referenced May 14, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes-progress on #434 (type-aware compression). Lands the mechanically-correct, default-OFF form of the feature per the spec's bench-gate / ship-or-defer policy (
docs/feature-type-aware-compression.md§ "Bench-gate / ship-or-defer policy").What landed
1.
src/aelfrice/compression.py(commit 1)Pure, deterministic
compress_for_retrieval(belief, *, locked) → CompressedBelief. Strategy table per spec:factsnapshottransientunknown./.\nboundary, with…. Code-fence-aware (does not split inside ``` blocks). Content that is wholly a single code fence renders verbatim (impl-PR resolution of spec open-question 3). Hard-truncates at last whitespace ≤MAX_HEADLINE_CHARS=240if no boundary found.[stub: belief={id} class=transient]— falls back to verbatim when the marker would cost more than the source content.2.
retrieve_v2()integration (commit 2)use_type_aware_compressionwith the established 4-stage precedence (env > kwarg > TOML > default-OFF). Env var:AELFRICE_TYPE_AWARE_COMPRESSION.RetrievalResultgains a parallelcompressed_beliefs: list[CompressedBelief]field. Same length and order asbeliefswhen the flag resolves True; empty when OFF — preserves byte-identical v1.x adapter behavior.3.
tests/bench_gate/test_compression_uplift.py(commit 3)Lab-mounted bench gate. Skips on public CI (autouse
bench_gatedmarker). WhenAELFRICE_CORPUS_ROOTpoints at a populatedcompression_uplift/directory, asserts the upstream invariant A2 depends on: compressed total tokens < uncompressed total tokens. The strict A2 recall@k comparison waits on the budget rewrite.4.
docs/CONFIG.md(commit 4)Adds
use_type_aware_compressionto the[retrieval]section header, the.aelfrice.tomlexample, and a dedicated Keys subsection with the strategy table.Acceptance status
tests/corpus/v2_0/compression_uplift/row schema documented in the gate's module docstring.compressed_beliefs[*].rendered_tokens. Bench-gate harness exists for the precondition; full A2 measurement is the follow-up.test_compression.py::test_compress_is_deterministicand thetest_token_monotone_non_increasingproperty test (16 inputs × 4 retention classes × 2 lock states).docs/RETRIEVAL_COMPOSITION.mddoes not exist yet (per [v2.0] HRR vocabulary bridge — close vocabulary-gap-recovery claim #433's spec note: "or wherever the [retrieval] Pipeline composition tracker — unified retrieve() with feature-flag gate #154 tracker doc lands by ship-time"). When the tracker doc lands, a one-line row foruse_type_aware_compressionshould be added.Test plan
uv run pytest tests/test_compression.py— 18 passeduv run pytest tests/test_compression_integration.py— 10 passeduv run pytest tests/bench_gate/test_compression_uplift.py— skipped (public CI; corpus absent)uv run pytest tests/— 2864 passed, 40 skipped, no regressions vsmainAELFRICE_CORPUS_ROOT=~/projects/aelfrice-lab/tests/corpus/v2_0 uv run pytest tests/bench_gate/test_compression_uplift.py— operator/lab-side; ratifies A2 preconditionFollow-up issues
retrieve_v2/retrieve_with_tierspack loops to usecb.rendered_tokensfor budget accounting whenuse_type_aware_compression=True. Spec text is atdocs/feature-type-aware-compression.md§ "Where compression sits".[rebuilder] token_budget.docs/RETRIEVAL_COMPOSITION.md(or whichever location [retrieval] Pipeline composition tracker — unified retrieve() with feature-flag gate #154 picks).Refs #434, #154 (composition tracker), #437 (canonical bench harness).
Summary by Sourcery
Add a type-aware compression module and integrate it with retrieve_v2 behind a configurable, default-off flag.
New Features:
Enhancements:
Tests:
Summary by CodeRabbit
New Features
Documentation