feat(clustering): RetrievalCluster module + multi-fact corpus mount (436) - #496
Merged
Conversation
Reviewer's GuideIntroduces the intentional clustering substrate as a standalone module and supporting infrastructure: a new clustering library with union-find based candidate clustering and a diversity-aware pack loop, a batched edge fetch helper in MemoryStore, a multi_fact corpus module plus schema support and docs for the bench gate, and scaffolding tests/bench hooks so this can ship independently of retrieval wiring. Sequence diagram for MemoryStore.edges_for_beliefs batched edge fetchsequenceDiagram
participant Caller
participant MemoryStore
participant SQLiteConnection
Caller->>MemoryStore: edges_for_beliefs(belief_ids)
alt empty belief_ids
MemoryStore-->>Caller: empty_edge_list
else nonempty belief_ids
MemoryStore->>MemoryStore: build_placeholders_and_params(belief_ids)
MemoryStore->>SQLiteConnection: execute(SELECT_edges_src_or_dst_in_placeholders, params)
SQLiteConnection-->>MemoryStore: edge_rows
MemoryStore->>MemoryStore: edges = map_rows_to_edges(edge_rows)
MemoryStore-->>Caller: edges
end
Class diagram for the new intentional clustering moduleclassDiagram
class Belief
class Edge
class RetrievalCluster {
+int cluster_id
+tuple[str,...] member_ids
+str representative_id
+float seed_score
}
class _UnionFind {
-dict[str,str] _parent
-dict[str,int] _size
+__init__() void
+make(x str) void
+find(x str) str
+union(a str, b str) void
}
class clustering_module {
<<module>>
+float DEFAULT_CLUSTER_EDGE_FLOOR
+int DEFAULT_CLUSTER_DIVERSITY_TARGET
+float _CHARS_PER_TOKEN
+_belief_tokens(b Belief) int
+cluster_candidates(candidates list[Belief], candidate_scores dict[str,float], edges Iterable[Edge], edge_weight_floor float) list[RetrievalCluster]
+pack_with_clusters(clusters list[RetrievalCluster], belief_by_id dict[str,Belief], token_budget int, cluster_diversity_target int, fallback_to_score bool) list[Belief]
}
clustering_module ..> Belief : uses
clustering_module ..> Edge : uses
clustering_module ..> RetrievalCluster : creates
clustering_module ..> _UnionFind : uses
RetrievalCluster o--> Belief : members
Flow diagram for pack_with_clusters diversity-aware greedy fillflowchart TD
A[Start pack_with_clusters] --> B[Initialize out, used_tokens, seen, covered_clusters]
B --> C[Sort clusters by descending seed_score]
C --> D{Stage 1: more clusters and covered_clusters < cluster_diversity_target?}
D -->|No| G[Stage 2 setup]
D -->|Yes| E[Select next cluster and representative_id]
E --> F{Representative exists, unseen, and fits token_budget?}
F -->|No and fallback_to_score| G
F -->|No and not fallback_to_score| D
F -->|Yes| H[Append representative, update seen, used_tokens, covered_clusters]
H --> D
G --> I[Iterate clusters in descending seed_score]
I --> J[Iterate member_ids within each cluster]
J --> K{Member unseen, exists, and fits remaining budget?}
K -->|No| J
K -->|Yes| L[Append belief, update seen and used_tokens]
L --> J
J --> M{More clusters to scan?}
M -->|Yes| I
M -->|No| N[Return out]
File-Level Changes
Possibly linked issues
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
|
Warning Rate limit exceeded
You’ve run out of usage credits. Purchase more in the billing tab. ⌛ 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 (1)
📝 WalkthroughWalkthroughThis PR implements the "intentional clustering" feature for retrieval-time reranking. It adds clustering algorithms via union-find connected components, two-stage greedy packing under token budgets, storage integration for edge batching, comprehensive unit tests, and corpus schema validation for the new ChangesIntentional Clustering Implementation
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related issues
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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. Comment |
robotrocketscience
added
author-Toug
Owner
Author
|
[claim:review:Gylf:2026-05-08T19:32:22Z] |
|
This PR is now behind Auto-rebase was removed because the bot has no signing key; rebasing as the bot strips author signatures and the |
Owner
Author
|
[claim:review:Setr:2026-05-08T19:33:04Z] |
Owner
Author
|
[release:review:Setr:2026-05-08T19:33:09Z] |
![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 1 security issue, 2 other issues, and left some high level feedback:
Security issues:
- Avoiding SQL string concatenation: untrusted input concatenated with raw SQL query can result in SQL Injection. In order to execute raw query safely, prepared statement should be used. SQLAlchemy provides TextualSQL to easily used prepared statement with named parameters. For complex SQL composition, use SQL Expression Language or Schema Definition Language. In most cases, SQLAlchemy ORM will be a better option. (link)
General comments:
- In
edges_for_beliefs, consider guarding against very largebelief_idsinputs (e.g. chunking the query or asserting on size) to avoid hitting SQLite’s parameter limit when constructing theINclause. - Both
_belief_tokensand_UnionFindduplicate logic that already exists elsewhere (retrieval._belief_tokens/dedup._UnionFind); it may be worth consolidating these into a shared utility to prevent future drift between implementations.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `edges_for_beliefs`, consider guarding against very large `belief_ids` inputs (e.g. chunking the query or asserting on size) to avoid hitting SQLite’s parameter limit when constructing the `IN` clause.
- Both `_belief_tokens` and `_UnionFind` duplicate logic that already exists elsewhere (`retrieval._belief_tokens` / `dedup._UnionFind`); it may be worth consolidating these into a shared utility to prevent future drift between implementations.
## Individual Comments
### Comment 1
<location path="src/aelfrice/clustering.py" line_range="92-99" />
<code_context>
+ self._parent[x] = x
+ self._size[x] = 1
+
+ def find(self, x: str) -> str:
+ path: list[str] = []
+ while self._parent[x] != x:
+ path.append(x)
+ x = self._parent[x]
+ for p in path:
+ self._parent[p] = x
+ return x
+
+ def union(self, a: str, b: str) -> None:
</code_context>
<issue_to_address>
**suggestion:** Consider making _UnionFind.find more defensive when called with unknown elements.
`find` currently assumes `x` is in `_parent` and will raise `KeyError` otherwise. Since `_UnionFind` may be used elsewhere, it would be safer to either assert `x in _parent` for a clearer failure, or treat unknown elements as their own singleton (i.e., perform a lazy `make` inside `find`).
```suggestion
def find(self, x: str) -> str:
if x not in self._parent:
self.make(x)
path: list[str] = []
while self._parent[x] != x:
path.append(x)
x = self._parent[x]
for p in path:
self._parent[p] = x
return x
```
</issue_to_address>
### Comment 2
<location path="src/aelfrice/clustering.py" line_range="155-162" />
<code_context>
+ for cid in candidate_ids:
+ groups.setdefault(uf.find(cid), []).append(cid)
+
+ raw_clusters: list[tuple[float, str, tuple[str, ...]]] = []
+ for members in groups.values():
+ ranked = sorted(
+ members,
+ key=lambda mid: (-candidate_scores.get(mid, 0.0), mid),
+ )
+ seed = candidate_scores.get(ranked[0], 0.0)
+ raw_clusters.append((seed, ranked[0], tuple(ranked)))
+
+ raw_clusters.sort(key=lambda t: (-t[0], t[1]))
</code_context>
<issue_to_address>
**suggestion (bug_risk):** Missing scores default to 0.0, which can silently hide data issues.
Using `candidate_scores.get(mid, 0.0)` in the sort key and for `seed` means missing scores are treated as 0 and won’t surface mismatches between `candidate_ids` and `candidate_scores`. If `candidate_scores` is meant to be complete, prefer `candidate_scores[mid]` or add a check/assert that all `candidate_ids` have entries so inconsistencies fail fast.
Suggested implementation:
```python
groups: dict[str, list[str]] = {}
for cid in candidate_ids:
groups.setdefault(uf.find(cid), []).append(cid)
# Ensure all candidate_ids have corresponding scores to avoid silently
# treating missing scores as 0.0, which can hide data issues.
missing_scores = [cid for cid in candidate_ids if cid not in candidate_scores]
if missing_scores:
raise ValueError(
f"Missing scores for candidate_ids: {', '.join(sorted(missing_scores))}"
)
raw_clusters: list[tuple[float, str, tuple[str, ...]]] = []
```
```python
ranked = sorted(
members,
key=lambda mid: (-candidate_scores[mid], mid),
)
seed = candidate_scores[ranked[0]]
```
</issue_to_address>
### Comment 3
<location path="src/aelfrice/store.py" line_range="2450-2453" />
<code_context>
cur = self._conn.execute(
f"SELECT * FROM edges WHERE src IN ({ph}) OR dst IN ({ph})",
params,
)
</code_context>
<issue_to_address>
**security (python.sqlalchemy.security.sqlalchemy-execute-raw-query):** Avoiding SQL string concatenation: untrusted input concatenated with raw SQL query can result in SQL Injection. In order to execute raw query safely, prepared statement should be used. SQLAlchemy provides TextualSQL to easily used prepared statement with named parameters. For complex SQL composition, use SQL Expression Language or Schema Definition Language. In most cases, SQLAlchemy ORM will be a better option.
*Source: opengrep*
</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
+92
to
+99
| def find(self, x: str) -> str: | ||
| path: list[str] = [] | ||
| while self._parent[x] != x: | ||
| path.append(x) | ||
| x = self._parent[x] | ||
| for p in path: | ||
| self._parent[p] = x | ||
| return x |
There was a problem hiding this comment.
suggestion: Consider making _UnionFind.find more defensive when called with unknown elements.
find currently assumes x is in _parent and will raise KeyError otherwise. Since _UnionFind may be used elsewhere, it would be safer to either assert x in _parent for a clearer failure, or treat unknown elements as their own singleton (i.e., perform a lazy make inside find).
Suggested change
| def find(self, x: str) -> str: | |
| path: list[str] = [] | |
| while self._parent[x] != x: | |
| path.append(x) | |
| x = self._parent[x] | |
| for p in path: | |
| self._parent[p] = x | |
| return x | |
| def find(self, x: str) -> str: | |
| if x not in self._parent: | |
| self.make(x) | |
| path: list[str] = [] | |
| while self._parent[x] != x: | |
| path.append(x) | |
| x = self._parent[x] | |
| for p in path: | |
| self._parent[p] = x | |
| return x |
Comment on lines
+155
to
+162
| raw_clusters: list[tuple[float, str, tuple[str, ...]]] = [] | ||
| for members in groups.values(): | ||
| ranked = sorted( | ||
| members, | ||
| key=lambda mid: (-candidate_scores.get(mid, 0.0), mid), | ||
| ) | ||
| seed = candidate_scores.get(ranked[0], 0.0) | ||
| raw_clusters.append((seed, ranked[0], tuple(ranked))) |
There was a problem hiding this comment.
suggestion (bug_risk): Missing scores default to 0.0, which can silently hide data issues.
Using candidate_scores.get(mid, 0.0) in the sort key and for seed means missing scores are treated as 0 and won’t surface mismatches between candidate_ids and candidate_scores. If candidate_scores is meant to be complete, prefer candidate_scores[mid] or add a check/assert that all candidate_ids have entries so inconsistencies fail fast.
Suggested implementation:
groups: dict[str, list[str]] = {}
for cid in candidate_ids:
groups.setdefault(uf.find(cid), []).append(cid)
# Ensure all candidate_ids have corresponding scores to avoid silently
# treating missing scores as 0.0, which can hide data issues.
missing_scores = [cid for cid in candidate_ids if cid not in candidate_scores]
if missing_scores:
raise ValueError(
f"Missing scores for candidate_ids: {', '.join(sorted(missing_scores))}"
)
raw_clusters: list[tuple[float, str, tuple[str, ...]]] = [] ranked = sorted(
members,
key=lambda mid: (-candidate_scores[mid], mid),
)
seed = candidate_scores[ranked[0]]
Comment on lines
+2450
to
+2453
| cur = self._conn.execute( | ||
| f"SELECT * FROM edges WHERE src IN ({ph}) OR dst IN ({ph})", | ||
| params, | ||
| ) |
There was a problem hiding this comment.
security (python.sqlalchemy.security.sqlalchemy-execute-raw-query): Avoiding SQL string concatenation: untrusted input concatenated with raw SQL query can result in SQL Injection. In order to execute raw query safely, prepared statement should be used. SQLAlchemy provides TextualSQL to easily used prepared statement with named parameters. For complex SQL composition, use SQL Expression Language or Schema Definition Language. In most cases, SQLAlchemy ORM will be a better option.
Source: opengrep
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Owner
Author
|
[claim:review:kulili:2026-05-08T19:34:11Z] |
Owner
Author
|
[release:review:kulili:2026-05-08T19:34:16Z] |
…436) Implements the pure-library half of intentional clustering per docs/feature-intentional-clustering.md. New `src/aelfrice/clustering.py` exposes: - `RetrievalCluster` dataclass — dense cluster_id, score-ranked member_ids, representative_id, seed_score. - `cluster_candidates()` — path-compressed union-find pass over the candidate-induced edge subgraph. Edges below `edge_weight_floor` (default 0.4: includes CITES, excludes RELATES_TO) and edges with one endpoint outside the candidate pool are filtered out. - `pack_with_clusters()` — diversity-aware greedy fill. Stage 1 picks one representative per cluster (descending seed_score) until `cluster_diversity_target=3` distinct clusters are covered or budget is exhausted. Stage 2 fills remaining budget from the score-ranked tail. `fallback_to_score=True` (default) bails Stage 1 on the first budget miss. `MemoryStore.edges_for_beliefs(belief_ids)` — batched fetch of every edge whose src OR dst is in `belief_ids`. Single SQL with `IN (...)`, empty input → empty list (no SQL). Bench-gate harness scaffold at `tests/bench_gate/test_intentional_clustering.py` + corpus mount point at `tests/corpus/v2_0/multi_fact/`. Public CI skips via the autouse `bench_gated` marker; corpus content lives lab-side per the directory-of-origin rule. The retrieval-side wiring (use_intentional_clustering flag in retrieve_v2 + flag-resolution helpers) is the next gate; the module ships independently so the substrate lands without a hot-path edit. Spec status updated to "module shipped; retrieval wiring + bench-gate evidence are the next gates". Reuses the path-compressed union-find pattern from src/aelfrice/dedup.py rather than importing it — neither owns the primitive yet, and a future refactor can promote one of the two.
robotrocketscience
force-pushed
the
feat/issue-436-clustering
branch
from
May 8, 2026 19:36
491b04c to
271d630
Compare
CodeQL flagged 'run_clustering_uplift may be uninitialized' on the try/import + pytest.skip pattern — CodeQL flow analysis can't see that pytest.skip raises. pytest.importorskip is the idiomatic equivalent that returns the module on success and skips otherwise, with no post-import name-binding ambiguity.
Owner
Author
|
Reviewed at the rebased + fix-up tip. Two reviewer-side commits:
Both commits signed (G). Full suite locally 2886 passed / 25 skipped. Discretion clean. Substance vs spec — clean. Module half of the contract matches Two stage-2 nuances to address in the wiring PR — non-blocking here:
SQL-parameter-count side-note. Locked-belief pre-include (spec § Open Q4) is correctly deferred per the PR body's "What's deferred" — the wiring PR adds the L0-skip at the top of the algorithm. Merging once CI on the fix commit clears. |
Owner
Author
|
[release:review:Gylf:2026-05-08T19:44:30Z] |
This was referenced May 8, 2026
robotrocketscience
added a commit
that referenced
this pull request
May 8, 2026
CodeQL flags 'run_doc_linker_uplift may be uninitialized' on the try/import + pytest.skip pattern — flow analysis can't see that pytest.skip raises. pytest.importorskip is the idiomatic equivalent that returns the module on success and skips otherwise, with no post-import name-binding ambiguity. Same fix as 5465b3c applied to the clustering bench-gate at #496.
This was referenced May 8, 2026
Closed
robotrocketscience
added a commit
that referenced
this pull request
May 10, 2026
Substrate landed in v2.0/v2.1 (#496 module + corpus, #498 retrieve_v2 wiring, #504 bench-gate scorer). Bench gate cleared on the production multi-store sweep (R6 60/60 PASS at p99 0.328ms, 15-30x margin under the 5ms A4 budget). Per the resolve_use_intentional_clustering() docstring contract ("the bench gate flips the default after lab-side benchmark evidence clears"), the default is unblocked. Precedence (env > kwarg > TOML > default) is unchanged; only the default value flips False -> True. Reversible via [retrieval] use_intentional_clustering = false in .aelfrice.toml or AELFRICE_INTENTIONAL_CLUSTERING=0 for v2.0.x parity. Tests updated: - test_default_is_off -> test_default_is_on - test_env_garbage_falls_through default-arm assertion flipped - test_default_call_byte_identical_to_explicit_off -> test_default_call_byte_identical_to_explicit_on (ON-byte-identity invariant supersedes the earlier OFF-byte-identity invariant) Floor parameter (DEFAULT_CLUSTER_EDGE_FLOOR = 0.4) is unchanged in this commit. Lab finding (raising to 0.6 triples uplift on 8 rows by excluding CITES inter-cluster edges) is queued as a separate follow-up so this PR stays minimal and reviewable.
robotrocketscience
added a commit
that referenced
this pull request
May 10, 2026
Substrate landed in v2.0/v2.1 (#496 module + corpus, #498 retrieve_v2 wiring, #504 bench-gate scorer). Bench gate cleared on the production multi-store sweep (R6 60/60 PASS at p99 0.328ms, 15-30x margin under the 5ms A4 budget). Per the resolve_use_intentional_clustering() docstring contract ("the bench gate flips the default after lab-side benchmark evidence clears"), the default is unblocked. Precedence (env > kwarg > TOML > default) is unchanged; only the default value flips False -> True. Reversible via [retrieval] use_intentional_clustering = false in .aelfrice.toml or AELFRICE_INTENTIONAL_CLUSTERING=0 for v2.0.x parity. Tests updated: - test_default_is_off -> test_default_is_on - test_env_garbage_falls_through default-arm assertion flipped - test_default_call_byte_identical_to_explicit_off -> test_default_call_byte_identical_to_explicit_on (ON-byte-identity invariant supersedes the earlier OFF-byte-identity invariant) Floor parameter (DEFAULT_CLUSTER_EDGE_FLOOR = 0.4) is unchanged in this commit. Lab finding (raising to 0.6 triples uplift on 8 rows by excluding CITES inter-cluster edges) is queued as a separate follow-up so this PR stays minimal and reviewable.
robotrocketscience
added a commit
that referenced
this pull request
May 10, 2026
Substrate landed in v2.0/v2.1 (#496 module + corpus, #498 retrieve_v2 wiring, #504 bench-gate scorer). Bench gate cleared on the production multi-store sweep (R6 60/60 PASS at p99 0.328ms, 15-30x margin under the 5ms A4 budget). Per the resolve_use_intentional_clustering() docstring contract ("the bench gate flips the default after lab-side benchmark evidence clears"), the default is unblocked. Precedence (env > kwarg > TOML > default) is unchanged; only the default value flips False -> True. Reversible via [retrieval] use_intentional_clustering = false in .aelfrice.toml or AELFRICE_INTENTIONAL_CLUSTERING=0 for v2.0.x parity. Tests updated: - test_default_is_off -> test_default_is_on - test_env_garbage_falls_through default-arm assertion flipped - test_default_call_byte_identical_to_explicit_off -> test_default_call_byte_identical_to_explicit_on (ON-byte-identity invariant supersedes the earlier OFF-byte-identity invariant) Floor parameter (DEFAULT_CLUSTER_EDGE_FLOOR = 0.4) is unchanged in this commit. Lab finding (raising to 0.6 triples uplift on 8 rows by excluding CITES inter-cluster edges) is queued as a separate follow-up so this PR stays minimal and reviewable.
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.
Summary
Module-first slice of intentional clustering per
docs/feature-intentional-clustering.md. Ships the pure library(
src/aelfrice/clustering.py), theMemoryStore.edges_for_beliefsbatched fetch, the multi_fact corpus mount, and the bench-gate
harness scaffold. Retrieval-side wiring is the next gate.
Toward 436.
What landed
src/aelfrice/clustering.py—RetrievalClusterdataclass,cluster_candidates()(path-compressed union-find on candidate-induced edge subgraph),
pack_with_clusters()(diversity-awaregreedy fill: Stage 1 representatives, Stage 2 score-ranked tail).
Defaults:
DEFAULT_CLUSTER_EDGE_FLOOR = 0.4(CITES in, RELATES_TOout per
EDGE_VALENCE),DEFAULT_CLUSTER_DIVERSITY_TARGET = 3.MemoryStore.edges_for_beliefs(belief_ids)— batched fetch:one SQL, returns every edge whose src OR dst is in the input.
tests/bench_gate/test_intentional_clustering.pytests/corpus/v2_0/multi_fact/. Public CI skipsvia the autouse
bench_gatedmarker; labelled rows live lab-side."module shipped; retrieval wiring + bench-gate evidence next gates".
What's deferred (next PR)
use_intentional_clusteringkwarg / env / TOML flag-resolution.retrieve_v2integration (replace pack loop behind default-OFF flag).The split exists because the pack-loop edit touches the hot retrieval
path. Shipping the substrate first lets reviewers verify the module
without a hot-path-edit blocker; the wiring PR is the next ship gate.
Test plan
edge filtering, score-ranked member ordering, pack Stage 1+2,
strict-diversity mode, missing-id race, default constants).
MemoryStore.edges_for_beliefsbatched-lookup test.tests/test_corpus_schema.pyaccepts the newmulti_factmodule schema (with
list[list_str]validator forexpected_clusters).tests/bench_gate/test_intentional_clustering.pyskips cleanlywhen
AELFRICE_CORPUS_ROOTis unset.landing first; the harness can't measure cluster_coverage@k
without
use_intentional_clustering=ONflowing throughretrieve_v2.Summary by Sourcery
Introduce an intentional clustering module for retrieval along with supporting store APIs, corpus schema, and bench-gate scaffolding to validate multi-fact recall improvements.
New Features:
Enhancements:
Tests:
Summary by CodeRabbit
New Features
Documentation
Tests