feat: heat-kernel composition in retrieval (#151 slice 2) - #309
Closed
robotrocketscience wants to merge 2 commits into
Closed
feat: heat-kernel composition in retrieval (#151 slice 2)#309robotrocketscience wants to merge 2 commits into
robotrocketscience wants to merge 2 commits into
Conversation
Extend the v1.3 ranking score with an optional log-additive heat-kernel authority term. Default heat_kernel=1.0 contributes log(1)=0, preserving byte-identical output for callers that have no eigenbasis available (the common path until graph-spectral caches land in retrieval). Defence-in-depth floor matches PARTIAL_BAYESIAN_BM25_FLOOR; the primary guard remains graph_spectral.heat_kernel_safe at the eigenbasis layer. Five new tests pin the neutral default, log-additivity at unit weight, weight-zero collapse, floor behaviour, and monotonicity in authority.
…part 2) Adds heat_kernel_for_candidates() helper to graph_spectral: given the eigenbasis cache and an L1 candidate set with positive BM25 relevance magnitudes, returns per-belief heat-kernel authority scores (floored). Threads an optional eigenbasis_cache kwarg through retrieve() and retrieve_with_tiers() into _l1_hits. When the kwarg is supplied AND is_heat_kernel_enabled() resolves True, _l1_hits computes per-candidate heat-kernel scores and passes them as heat_kernel= into partial_bayesian_score. Default heat_kernel=1.0 keeps non-supplied / disabled callers byte-identical. Cache absence, staleness, empty candidates, and unknown belief ids all fail soft to a neutral no-contribution. Tests cover the helper's per-belief output, unknown-id passthrough, stale-cache behaviour, empty inputs, retrieve()-level wiring, and the flag-off / no-cache byte-identity guarantee.
|
Warning Rate limit exceeded
To keep reviews running without waiting, you can enable usage-based add-on for your organization. This allows additional reviews beyond the hourly cap. Account admins can enable it under billing. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
✨ Finishing Touches🧪 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. Review rate limit: 0/1 reviews remaining, refill in 36 minutes and 22 seconds.Comment |
Reviewer's GuideImplements optional heat-kernel authority composition in the v1.3 retrieval score, threading an eigenbasis-backed heat-kernel lane through L1 retrieval while preserving byte-identical behavior for existing callers when the feature is disabled or unavailable. Sequence diagram for L1 retrieval with optional heat-kernel authoritysequenceDiagram
actor Client
participant Retrieval as retrieve
participant RetrievalTiers as retrieve_with_tiers
participant L1 as _l1_hits
participant Store as MemoryStore
participant HKLookup as _heat_kernel_lookup
participant Spectral as heat_kernel_for_candidates
participant Cache as GraphEigenbasisCache
Client->>Retrieval: retrieve(query, eigenbasis_cache)
Retrieval->>L1: _l1_hits(..., eigenbasis_cache)
alt use_bm25f_anchors is True
L1->>Store: bm25f_search_beliefs_scored
Store-->>L1: (Belief, bm25_pos) list
else use_bm25f_anchors is False
L1->>Store: search_beliefs_scored
Store-->>L1: (Belief, bm25_raw) list
end
L1->>HKLookup: _heat_kernel_lookup(eigenbasis_cache, candidates)
alt cache is None or candidates empty
HKLookup-->>L1: {}
else is_heat_kernel_enabled() is False
HKLookup-->>L1: {}
else heat-kernel lane enabled
HKLookup->>Spectral: heat_kernel_for_candidates(Cache, candidates)
Spectral->>Cache: check is_stale(), eigvals, eigvecs, belief_ids
alt cache stale or eigenbasis missing
Spectral-->>HKLookup: {}
else seeds available
Spectral-->>HKLookup: {belief_id: heat_score}
end
HKLookup-->>L1: {belief_id: heat_score}
end
alt posterior_weight == 0.0 and no heat_scores
L1->>Store: search_beliefs(query)
Store-->>L1: Belief list
L1-->>Retrieval: Belief list
else scoring path
L1->>L1: partial_bayesian_score(bm25_raw, alpha, beta, posterior_weight, heat_kernel)
L1-->>Retrieval: ranked Belief list
end
Retrieval-->>Client: Belief list
Client->>RetrievalTiers: retrieve_with_tiers(..., eigenbasis_cache)
RetrievalTiers->>L1: _l1_hits(..., eigenbasis_cache)
L1-->>RetrievalTiers: ranked L1 Belief list
RetrievalTiers-->>Client: tiered results
Flow diagram for partial_bayesian_score with heat-kernel compositionflowchart TD
A["Start partial_bayesian_score"] --> B["relevance_pos = max(-bm25_raw, PARTIAL_BAYESIAN_BM25_FLOOR)"]
B --> C["score = log(relevance_pos)"]
C --> D{"heat_kernel != 1.0 and<br/>heat_kernel_weight != 0.0?"}
D -->|No| F
D -->|Yes| E["heat_safe = max(heat_kernel, PARTIAL_BAYESIAN_BM25_FLOOR)<br/>score += heat_kernel_weight * log(heat_safe)"]
E --> F{posterior_weight == 0.0?}
F -->|Yes| G[Return score]
F -->|No| H["p = posterior_mean(alpha, beta)"]
H --> I["p_safe = p if p > 0.0<br/>else PARTIAL_BAYESIAN_BM25_FLOOR"]
I --> J["score += posterior_weight * log(p_safe)"]
J --> K[Return score]
G --> L[End]
K --> L
File-Level Changes
Possibly linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![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 4 issues, and left some high level feedback:
- In
heat_kernel_for_candidates, rebuildingid_to_rowon every call is O(N) and may dominate cost for frequent L1 calls; consider storing abelief_id → rowindex onGraphEigenbasisCache(built alongsidebelief_ids) and reusing it here. - In
_l1_hitsFTS5 path, whenposterior_weight == 0.0andheat_scoresis empty you re-runsearch_beliefsafter already callingsearch_beliefs_scored; if identical ordering is guaranteed between the two, you could avoid the second query by reusing thescoredresult (e.g., by dropping scores) to save one DB hit.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `heat_kernel_for_candidates`, rebuilding `id_to_row` on every call is O(N) and may dominate cost for frequent L1 calls; consider storing a `belief_id → row` index on `GraphEigenbasisCache` (built alongside `belief_ids`) and reusing it here.
- In `_l1_hits` FTS5 path, when `posterior_weight == 0.0` and `heat_scores` is empty you re-run `search_beliefs` after already calling `search_beliefs_scored`; if identical ordering is guaranteed between the two, you could avoid the second query by reusing the `scored` result (e.g., by dropping scores) to save one DB hit.
## Individual Comments
### Comment 1
<location path="src/aelfrice/graph_spectral.py" line_range="389-398" />
<code_context>
+ id_to_row: dict[str, int] = {bid: i for i, bid in enumerate(cache.belief_ids)}
+ n = len(cache.belief_ids)
+ bm25_full = np.zeros(n, dtype=np.float64)
+ present: list[str] = []
+ for bid, bm25_pos in candidates:
+ row = id_to_row.get(bid)
+ if row is None:
+ continue
+ if bm25_pos > 0.0:
+ bm25_full[row] = bm25_pos
+ present.append(bid)
+ if not present:
+ return {}
+ seeds = seeds_from_bm25(bm25_full, top_k=seed_top_k)
+ if not np.any(seeds):
+ return {}
+ raw = heat_kernel_score(cache.eigvals, cache.eigvecs, seeds, t=t)
+ safe = heat_kernel_safe(raw, floor=floor)
+ return {bid: float(safe[id_to_row[bid]]) for bid in present}
+
+
</code_context>
<issue_to_address>
**suggestion:** Function only returns scores for positively seeded candidates, which slightly contradicts the docstring.
The docstring promises `dict[belief_id, heat_kernel_score]` for every candidate with a row in the eigenbasis, but the code only records candidates when `bm25_pos > 0.0`. So any candidate with an eigenbasis row but non‑positive `bm25_pos` is omitted, even though it would get a finite score from diffusion. Please decide whether these should be included with their propagated scores or intentionally excluded, and update the docstring and/or logic to match that decision.
Suggested implementation:
```python
id_to_row: dict[str, int] = {bid: i for i, bid in enumerate(cache.belief_ids)}
n = len(cache.belief_ids)
bm25_full = np.zeros(n, dtype=np.float64)
```
```python
def heat_kernel_for_candidates(
cache: "GraphEigenbasisCache",
candidates: list[tuple[str, float]],
*,
t: float = DEFAULT_HEAT_BANDWIDTH,
seed_top_k: int = DEFAULT_BM25_SEED_TOP_K,
floor: float = HEAT_SCORE_FLOOR,
) -> dict[str, float]:
"""Compute per-belief heat-kernel authority scores for the
given `candidates`, returning `dict[belief_id, heat_kernel_score]`
for every candidate that has a corresponding row in the eigenbasis.
The heat kernel is seeded from candidates with positive BM25 scores;
candidates with non-positive BM25 scores can still receive non-zero
authority via diffusion, and their scores are included in the result.
"""
if cache.is_stale() or cache.eigvals is None or cache.eigvecs is None:
return {}
if cache.belief_ids is None or not candidates:
return {}
# Map belief IDs to their row indices in the eigenbasis.
id_to_row: dict[str, int] = {bid: i for i, bid in enumerate(cache.belief_ids)}
n = len(cache.belief_ids)
# Dense BM25 vector over the full eigenbasis rows.
bm25_full = np.zeros(n, dtype=np.float64)
# Track which candidates actually have a row in the eigenbasis.
rows_for_bids: dict[str, int] = {}
for bid, bm25_pos in candidates:
row = id_to_row.get(bid)
if row is None:
# Candidate not present in eigenbasis; cannot score it.
continue
rows_for_bids[bid] = row
# Only positive BM25 scores are used as seeds.
if bm25_pos > 0.0:
bm25_full[row] = bm25_pos
# If there are no positive seeds, there is nothing to diffuse.
if not np.any(bm25_full):
return {}
seeds = seeds_from_bm25(bm25_full, top_k=seed_top_k)
if not np.any(seeds):
return {}
raw = heat_kernel_score(cache.eigvals, cache.eigvecs, seeds, t=t)
safe = heat_kernel_safe(raw, floor=floor)
# Return scores for all candidates that had a row in the eigenbasis,
# regardless of whether they were used as positive seeds.
return {bid: float(safe[row]) for bid, row in rows_for_bids.items()}
```
If the existing docstring for `heat_kernel_for_candidates` spans more lines than shown in the snippet, you should ensure the updated description above is consistent with the rest of the documentation. Specifically:
1. Confirm that any references to "positively seeded candidates" are updated to explain that *seeding* uses positive BM25 scores but *outputs* are provided for all candidates with eigenbasis rows.
2. If there are unit tests asserting that only positively seeded candidates are returned, update them to expect scores for all candidates with rows in the eigenbasis, while still expecting an empty dict when there are no positive BM25 seeds.
</issue_to_address>
### Comment 2
<location path="src/aelfrice/scoring.py" line_range="171-180" />
<code_context>
alpha: float,
beta: float,
posterior_weight: float = DEFAULT_POSTERIOR_WEIGHT,
+ heat_kernel: float = 1.0,
+ heat_kernel_weight: float = 1.0,
) -> float:
</code_context>
<issue_to_address>
**issue:** The docstring claims the heat-kernel floor is identical to `heat_kernel_safe`, but the code uses `PARTIAL_BAYESIAN_BM25_FLOOR`.
The docstring states the heat-kernel input must be passed through `heat_kernel_safe` and that the floor is identical, but the code uses `PARTIAL_BAYESIAN_BM25_FLOOR`. If `HEAT_SCORE_FLOOR` and `PARTIAL_BAYESIAN_BM25_FLOOR` diverge, this will no longer hold and double-flooring could bias scores away from the 1.0 baseline. Consider either using the same constant as `heat_kernel_safe` or updating the docstring to say the floor is only similar, not identical.
</issue_to_address>
### Comment 3
<location path="tests/test_bayesian_ranking.py" line_range="511-520" />
<code_context>
+ assert with_default == base
+
+
+def test_heat_kernel_log_additive_at_unit_weight() -> None:
+ """At heat_kernel_weight=1.0 the heat term adds log(heat_kernel)
+ on top of the bm25+posterior baseline."""
+ bm25_raw, alpha, beta, pw = -1.0, 2.0, 1.0, 0.5
+ heat = 0.4
+ base = partial_bayesian_score(bm25_raw, alpha, beta, pw)
+ composed = partial_bayesian_score(
+ bm25_raw, alpha, beta, pw, heat_kernel=heat,
+ )
+ assert abs(composed - (base + math.log(heat))) < 1e-12
+
+
</code_context>
<issue_to_address>
**suggestion (testing):** Add coverage for non-unit `heat_kernel_weight` values.
Current tests cover only `heat_kernel_weight` = 1.0 and 0.0. Please add a parametrized test over a few `(heat_kernel, heat_kernel_weight)` pairs (e.g., `w in {0.5, 2.0}`) that checks `partial_bayesian_score(..., heat_kernel=h, heat_kernel_weight=w) == base + w * log(h)`, and also verifies that `heat_kernel=1.0` is neutral even when `heat_kernel_weight != 1.0` to lock in the weighting semantics.
```suggestion
def test_heat_kernel_log_additive_at_unit_weight() -> None:
"""At heat_kernel_weight=1.0 the heat term adds log(heat_kernel)
on top of the bm25+posterior baseline."""
bm25_raw, alpha, beta, pw = -1.0, 2.0, 1.0, 0.5
heat = 0.4
base = partial_bayesian_score(bm25_raw, alpha, beta, pw)
composed = partial_bayesian_score(
bm25_raw, alpha, beta, pw, heat_kernel=heat,
)
assert abs(composed - (base + math.log(heat))) < 1e-12
@pytest.mark.parametrize(
"heat_kernel, heat_kernel_weight",
[
(0.4, 0.5),
(0.4, 2.0),
(0.1, 0.5),
(0.9, 2.0),
],
)
def test_heat_kernel_log_additive_with_weight(
heat_kernel: float,
heat_kernel_weight: float,
) -> None:
"""For non-unit heat_kernel_weight the heat term scales log(heat_kernel),
and heat_kernel=1.0 remains neutral."""
bm25_raw, alpha, beta, pw = -1.0, 2.0, 1.0, 0.5
base = partial_bayesian_score(bm25_raw, alpha, beta, pw)
composed = partial_bayesian_score(
bm25_raw,
alpha,
beta,
pw,
heat_kernel=heat_kernel,
heat_kernel_weight=heat_kernel_weight,
)
assert abs(
composed - (base + heat_kernel_weight * math.log(heat_kernel))
) < 1e-12
neutral = partial_bayesian_score(
bm25_raw,
alpha,
beta,
pw,
heat_kernel=1.0,
heat_kernel_weight=heat_kernel_weight,
)
assert abs(neutral - base) < 1e-12
```
</issue_to_address>
### Comment 4
<location path="tests/test_bayesian_ranking.py" line_range="596-605" />
<code_context>
+def test_retrieve_eigenbasis_cache_threaded_through(
</code_context>
<issue_to_address>
**suggestion (testing):** Consider adding a BM25F-path smoke test with `eigenbasis_cache` to cover the `_l1_hits` BM25F branch.
The current tests only cover the FTS5 path with/without `eigenbasis_cache` and the flag toggled. Since `_l1_hits` now also passes `eigenbasis_cache` through the BM25F path (`use_bm25f_anchors=True`), that branch is untested. Please add a small smoke test that:
- builds a `GraphEigenbasisCache`,
- calls `retrieve(..., use_bm25f_anchors=True, bm25f_cache=..., eigenbasis_cache=..., posterior_weight=0.0 or >0, AELFRICE_HEAT_KERNEL=1)`,
- asserts the call succeeds and returns non-empty results (optionally checking stable ordering when `posterior_weight=0` with the flag off vs on).
This will mirror the FTS5 coverage and verify the heat-kernel wiring for BM25F as well.
Suggested implementation:
```python
from aelfrice.graph_spectral import GraphEigenbasisCache
s = _equal_bm25_store()
cache = GraphEigenbasisCache(store=s, path=tmp_path / "eb.npz", k=3)
# Smoke-test the BM25F path with eigenbasis_cache threaded through.
# This mirrors the FTS5 coverage and ensures the heat-kernel wiring works
# when `_l1_hits` routes through BM25F with use_bm25f_anchors=True.
bm25f_cache = _equal_bm25f_cache(store=s)
prior = os.environ.get("AELFRICE_HEAT_KERNEL")
os.environ["AELFRICE_HEAT_KERNEL"] = "1"
try:
l1 = retrieve(
s,
"widget",
token_budget=10_000,
posterior_weight=0.0,
use_bm25f_anchors=True,
bm25f_cache=bm25f_cache,
eigenbasis_cache=cache,
)
finally:
if prior is None:
os.environ.pop("AELFRICE_HEAT_KERNEL", None)
else:
os.environ["AELFRICE_HEAT_KERNEL"] = prior
assert l1, "BM25F + eigenbasis_cache should return a non-empty L1"
```
1. Ensure there is a helper `_equal_bm25f_cache(store: Store)` (or equivalent) that builds a BM25F cache compatible with `retrieve(..., bm25f_cache=...)`. If your test module already has a BM25F cache helper, reuse that name instead of `_equal_bm25f_cache`.
2. If `retrieve` is not already imported into this module, add the appropriate import at the top of `tests/test_bayesian_ranking.py` (e.g., `from aelfrice.ranking import retrieve`) consistent with the existing imports and conventions in the file.
3. If your existing FTS5 eigenbasis tests use a different query string, token budget, or fixture setup for `_equal_bm25_store()`, align the parameters here to match those tests for more stable behavior.
</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
+389
to
+398
| present: list[str] = [] | ||
| for bid, bm25_pos in candidates: | ||
| row = id_to_row.get(bid) | ||
| if row is None: | ||
| continue | ||
| if bm25_pos > 0.0: | ||
| bm25_full[row] = bm25_pos | ||
| present.append(bid) | ||
| if not present: | ||
| return {} |
There was a problem hiding this comment.
suggestion: Function only returns scores for positively seeded candidates, which slightly contradicts the docstring.
The docstring promises dict[belief_id, heat_kernel_score] for every candidate with a row in the eigenbasis, but the code only records candidates when bm25_pos > 0.0. So any candidate with an eigenbasis row but non‑positive bm25_pos is omitted, even though it would get a finite score from diffusion. Please decide whether these should be included with their propagated scores or intentionally excluded, and update the docstring and/or logic to match that decision.
Suggested implementation:
id_to_row: dict[str, int] = {bid: i for i, bid in enumerate(cache.belief_ids)}
n = len(cache.belief_ids)
bm25_full = np.zeros(n, dtype=np.float64)def heat_kernel_for_candidates(
cache: "GraphEigenbasisCache",
candidates: list[tuple[str, float]],
*,
t: float = DEFAULT_HEAT_BANDWIDTH,
seed_top_k: int = DEFAULT_BM25_SEED_TOP_K,
floor: float = HEAT_SCORE_FLOOR,
) -> dict[str, float]:
"""Compute per-belief heat-kernel authority scores for the
given `candidates`, returning `dict[belief_id, heat_kernel_score]`
for every candidate that has a corresponding row in the eigenbasis.
The heat kernel is seeded from candidates with positive BM25 scores;
candidates with non-positive BM25 scores can still receive non-zero
authority via diffusion, and their scores are included in the result.
"""
if cache.is_stale() or cache.eigvals is None or cache.eigvecs is None:
return {}
if cache.belief_ids is None or not candidates:
return {}
# Map belief IDs to their row indices in the eigenbasis.
id_to_row: dict[str, int] = {bid: i for i, bid in enumerate(cache.belief_ids)}
n = len(cache.belief_ids)
# Dense BM25 vector over the full eigenbasis rows.
bm25_full = np.zeros(n, dtype=np.float64)
# Track which candidates actually have a row in the eigenbasis.
rows_for_bids: dict[str, int] = {}
for bid, bm25_pos in candidates:
row = id_to_row.get(bid)
if row is None:
# Candidate not present in eigenbasis; cannot score it.
continue
rows_for_bids[bid] = row
# Only positive BM25 scores are used as seeds.
if bm25_pos > 0.0:
bm25_full[row] = bm25_pos
# If there are no positive seeds, there is nothing to diffuse.
if not np.any(bm25_full):
return {}
seeds = seeds_from_bm25(bm25_full, top_k=seed_top_k)
if not np.any(seeds):
return {}
raw = heat_kernel_score(cache.eigvals, cache.eigvecs, seeds, t=t)
safe = heat_kernel_safe(raw, floor=floor)
# Return scores for all candidates that had a row in the eigenbasis,
# regardless of whether they were used as positive seeds.
return {bid: float(safe[row]) for bid, row in rows_for_bids.items()}If the existing docstring for heat_kernel_for_candidates spans more lines than shown in the snippet, you should ensure the updated description above is consistent with the rest of the documentation. Specifically:
- Confirm that any references to "positively seeded candidates" are updated to explain that seeding uses positive BM25 scores but outputs are provided for all candidates with eigenbasis rows.
- If there are unit tests asserting that only positively seeded candidates are returned, update them to expect scores for all candidates with rows in the eigenbasis, while still expecting an empty dict when there are no positive BM25 seeds.
Comment on lines
+171
to
+180
| heat_kernel: float = 1.0, | ||
| heat_kernel_weight: float = 1.0, | ||
| ) -> float: | ||
| """v1.3 partial Bayesian-weighted retrieval score. | ||
| """v1.3 partial Bayesian-weighted retrieval score, optionally | ||
| composed with a heat-kernel authority term (#151 slice 2). | ||
|
|
||
| `score = log(max(-bm25_raw, EPS)) + posterior_weight * log(posterior_mean)` | ||
| Full formula:: | ||
|
|
||
| score = log(max(-bm25_raw, EPS)) | ||
| + heat_kernel_weight * log(max(heat_kernel, EPS)) |
There was a problem hiding this comment.
issue: The docstring claims the heat-kernel floor is identical to heat_kernel_safe, but the code uses PARTIAL_BAYESIAN_BM25_FLOOR.
The docstring states the heat-kernel input must be passed through heat_kernel_safe and that the floor is identical, but the code uses PARTIAL_BAYESIAN_BM25_FLOOR. If HEAT_SCORE_FLOOR and PARTIAL_BAYESIAN_BM25_FLOOR diverge, this will no longer hold and double-flooring could bias scores away from the 1.0 baseline. Consider either using the same constant as heat_kernel_safe or updating the docstring to say the floor is only similar, not identical.
Comment on lines
+511
to
+520
| def test_heat_kernel_log_additive_at_unit_weight() -> None: | ||
| """At heat_kernel_weight=1.0 the heat term adds log(heat_kernel) | ||
| on top of the bm25+posterior baseline.""" | ||
| bm25_raw, alpha, beta, pw = -1.0, 2.0, 1.0, 0.5 | ||
| heat = 0.4 | ||
| base = partial_bayesian_score(bm25_raw, alpha, beta, pw) | ||
| composed = partial_bayesian_score( | ||
| bm25_raw, alpha, beta, pw, heat_kernel=heat, | ||
| ) | ||
| assert abs(composed - (base + math.log(heat))) < 1e-12 |
There was a problem hiding this comment.
suggestion (testing): Add coverage for non-unit heat_kernel_weight values.
Current tests cover only heat_kernel_weight = 1.0 and 0.0. Please add a parametrized test over a few (heat_kernel, heat_kernel_weight) pairs (e.g., w in {0.5, 2.0}) that checks partial_bayesian_score(..., heat_kernel=h, heat_kernel_weight=w) == base + w * log(h), and also verifies that heat_kernel=1.0 is neutral even when heat_kernel_weight != 1.0 to lock in the weighting semantics.
Suggested change
| def test_heat_kernel_log_additive_at_unit_weight() -> None: | |
| """At heat_kernel_weight=1.0 the heat term adds log(heat_kernel) | |
| on top of the bm25+posterior baseline.""" | |
| bm25_raw, alpha, beta, pw = -1.0, 2.0, 1.0, 0.5 | |
| heat = 0.4 | |
| base = partial_bayesian_score(bm25_raw, alpha, beta, pw) | |
| composed = partial_bayesian_score( | |
| bm25_raw, alpha, beta, pw, heat_kernel=heat, | |
| ) | |
| assert abs(composed - (base + math.log(heat))) < 1e-12 | |
| def test_heat_kernel_log_additive_at_unit_weight() -> None: | |
| """At heat_kernel_weight=1.0 the heat term adds log(heat_kernel) | |
| on top of the bm25+posterior baseline.""" | |
| bm25_raw, alpha, beta, pw = -1.0, 2.0, 1.0, 0.5 | |
| heat = 0.4 | |
| base = partial_bayesian_score(bm25_raw, alpha, beta, pw) | |
| composed = partial_bayesian_score( | |
| bm25_raw, alpha, beta, pw, heat_kernel=heat, | |
| ) | |
| assert abs(composed - (base + math.log(heat))) < 1e-12 | |
| @pytest.mark.parametrize( | |
| "heat_kernel, heat_kernel_weight", | |
| [ | |
| (0.4, 0.5), | |
| (0.4, 2.0), | |
| (0.1, 0.5), | |
| (0.9, 2.0), | |
| ], | |
| ) | |
| def test_heat_kernel_log_additive_with_weight( | |
| heat_kernel: float, | |
| heat_kernel_weight: float, | |
| ) -> None: | |
| """For non-unit heat_kernel_weight the heat term scales log(heat_kernel), | |
| and heat_kernel=1.0 remains neutral.""" | |
| bm25_raw, alpha, beta, pw = -1.0, 2.0, 1.0, 0.5 | |
| base = partial_bayesian_score(bm25_raw, alpha, beta, pw) | |
| composed = partial_bayesian_score( | |
| bm25_raw, | |
| alpha, | |
| beta, | |
| pw, | |
| heat_kernel=heat_kernel, | |
| heat_kernel_weight=heat_kernel_weight, | |
| ) | |
| assert abs( | |
| composed - (base + heat_kernel_weight * math.log(heat_kernel)) | |
| ) < 1e-12 | |
| neutral = partial_bayesian_score( | |
| bm25_raw, | |
| alpha, | |
| beta, | |
| pw, | |
| heat_kernel=1.0, | |
| heat_kernel_weight=heat_kernel_weight, | |
| ) | |
| assert abs(neutral - base) < 1e-12 |
Comment on lines
+596
to
+605
| def test_retrieve_eigenbasis_cache_threaded_through( | ||
| tmp_path: Path, | ||
| ) -> None: | ||
| """Smoke test: retrieve() accepts eigenbasis_cache, builds it, | ||
| enables the heat-kernel flag, and returns a non-empty L1 with | ||
| no exceptions. The actual reranking is exercised by the | ||
| scoring-level tests above; this confirms wiring.""" | ||
| import os | ||
| from aelfrice.graph_spectral import GraphEigenbasisCache | ||
| s = _equal_bm25_store() |
There was a problem hiding this comment.
suggestion (testing): Consider adding a BM25F-path smoke test with eigenbasis_cache to cover the _l1_hits BM25F branch.
The current tests only cover the FTS5 path with/without eigenbasis_cache and the flag toggled. Since _l1_hits now also passes eigenbasis_cache through the BM25F path (use_bm25f_anchors=True), that branch is untested. Please add a small smoke test that:
- builds a
GraphEigenbasisCache, - calls
retrieve(..., use_bm25f_anchors=True, bm25f_cache=..., eigenbasis_cache=..., posterior_weight=0.0 or >0, AELFRICE_HEAT_KERNEL=1), - asserts the call succeeds and returns non-empty results (optionally checking stable ordering when
posterior_weight=0with the flag off vs on).
This will mirror the FTS5 coverage and verify the heat-kernel wiring for BM25F as well.
Suggested implementation:
from aelfrice.graph_spectral import GraphEigenbasisCache
s = _equal_bm25_store()
cache = GraphEigenbasisCache(store=s, path=tmp_path / "eb.npz", k=3)
# Smoke-test the BM25F path with eigenbasis_cache threaded through.
# This mirrors the FTS5 coverage and ensures the heat-kernel wiring works
# when `_l1_hits` routes through BM25F with use_bm25f_anchors=True.
bm25f_cache = _equal_bm25f_cache(store=s)
prior = os.environ.get("AELFRICE_HEAT_KERNEL")
os.environ["AELFRICE_HEAT_KERNEL"] = "1"
try:
l1 = retrieve(
s,
"widget",
token_budget=10_000,
posterior_weight=0.0,
use_bm25f_anchors=True,
bm25f_cache=bm25f_cache,
eigenbasis_cache=cache,
)
finally:
if prior is None:
os.environ.pop("AELFRICE_HEAT_KERNEL", None)
else:
os.environ["AELFRICE_HEAT_KERNEL"] = prior
assert l1, "BM25F + eigenbasis_cache should return a non-empty L1"- Ensure there is a helper
_equal_bm25f_cache(store: Store)(or equivalent) that builds a BM25F cache compatible withretrieve(..., bm25f_cache=...). If your test module already has a BM25F cache helper, reuse that name instead of_equal_bm25f_cache. - If
retrieveis not already imported into this module, add the appropriate import at the top oftests/test_bayesian_ranking.py(e.g.,from aelfrice.ranking import retrieve) consistent with the existing imports and conventions in the file. - If your existing FTS5 eigenbasis tests use a different query string, token budget, or fixture setup for
_equal_bm25_store(), align the parameters here to match those tests for more stable behavior.
Owner
Author
|
[claim:review:Setr:2026-04-29T05:54:22Z] |
Owner
Author
|
[release:review:Setr:2026-04-29T05:55:26Z] |
robotrocketscience
enabled auto-merge (squash)
April 29, 2026 05:55
Owner
Author
|
[claim:review:Gylf:2026-04-29T15:48:25Z] |
Owner
Author
|
[release:review:Gylf:2026-04-29T15:49:21Z] |
Owner
Author
|
Superseded by #310 (merged). Closing. |
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
1 participant
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 part of #151 (slice 2 — heat-kernel composition).
Summary
Adds the optional heat-kernel authority term to the v1.3 ranking
score and wires it through retrieval. Default behaviour is
unchanged — every existing caller produces byte-identical output
because the heat-kernel lane requires both an
eigenbasis_cachekwarg AND
is_heat_kernel_enabled()to resolve True.Composition implemented (per #151):
Changes
scoring.partial_bayesian_score()gainsheat_kernel: float = 1.0and
heat_kernel_weight: float = 1.0. Default1.0makeslog(1) == 0, so callers unchanged. Defence-in-depth floor atPARTIAL_BAYESIAN_BM25_FLOOR.graph_spectral.heat_kernel_for_candidates(cache, candidates):new helper that builds a sparse seed vector from the L1 hit set,
runs
heat_kernel_scoreagainst the cached eigenbasis, appliesheat_kernel_safe, and returnsdict[belief_id, score]forcandidates present in the eigenbasis row order. Returns
{}whencache is stale, eigenbasis unbuilt, or no positive seed exists.
retrieval._l1_hitsaccepts aneigenbasis_cachekwarg; boththe BM25F and FTS5 paths thread per-belief heat scores into
partial_bayesian_score. Backward-compat preserved by thedefault
1.0neutral.retrieve()andretrieve_with_tiers()accept and forwardeigenbasis_cacheinto_l1_hits. TheRetrievalCacheandretrieve_v2wrappers stay untouched in this slice; benchmarkcallers using the lower-level entry points get the new lane
immediately.
Tests
tests/test_bayesian_ranking.py(+8):heat_kernel=1.0tests/test_graph_spectral.py(+4):{}{}Full suite: 1897 passed, 8 skipped.
Out of scope (future slice)
RetrievalCache.retrieve()/retrieve_v2()cache-key extensionfor the eigenbasis-cache identity (only matters once the lane is
default-on; current callers wanting it pass through
retrieve()directly).
Discretion check
Discretion grep against the canonical pattern set — CLEAN.
Summary by Sourcery
Add an optional heat-kernel authority term to the v1.3 retrieval ranking score and wire it through L1 retrieval while preserving default, byte-identical behaviour when the feature is disabled or no cache is available.
New Features:
Enhancements:
Tests: