Cache IDF corpus statistics to cut DB round-trips on repeated recall - #55
Cache IDF corpus statistics to cut DB round-trips on repeated recall#55lkxdsb wants to merge 3 commits into
Conversation
…am#54) Add two short-TTL (60s), instance-level caches inside MemorySearcher._build_idf so repeated or overlapping search() over an unchanged corpus reuses corpus_size and per-token document frequencies instead of re-hitting the store: - _corpus_size_cache: keyed by the partition scope. - _df_cache: keyed by (token, partition_scope) — token is in the key so one query's DF never overwrites another's. On a cache miss, the whole token set is fetched in one batched store call (real keyword_doc_freqs API takes a term list) and back-filled per token, including 0/empty results so an empty corpus or unseen term does not re-hit the store. No storage-layer changes; the TTL bounds staleness instead of write-driven invalidation, per the issue. Extra hardening beyond the issue (separable, see PR for details): - asyncio.Lock around the miss→fetch→back-fill section so concurrent same-key misses share one fetch. - scripts/ for real-store / end-to-end verification (ruff-excluded dir). Tests: tests/unit/test_idf_cache.py (8 cases: Hit, Miss, TTL Expiry, Partition Isolation, DF-per-token no-leak, order-independence, two concurrency cases). Co-Authored-By: Claude <noreply@anthropic.com>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
📝 WalkthroughWalkthroughAdds short-TTL, instance-level caching for IDF corpus sizes and token document frequencies in ChangesIDF cache
Estimated code review effort: 4 (Complex) | ~50 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Client
participant HebbMind
participant MemorySearcher
participant SQLiteMemoryStore
Client->>HebbMind: search(query)
HebbMind->>MemorySearcher: _build_idf(partition_ids, tokens)
MemorySearcher->>MemorySearcher: check corpus_size cache
MemorySearcher->>SQLiteMemoryStore: corpus_size (if cache miss)
SQLiteMemoryStore-->>MemorySearcher: corpus count
MemorySearcher->>MemorySearcher: check token DF cache
MemorySearcher->>SQLiteMemoryStore: keyword_doc_freqs (if cache miss)
SQLiteMemoryStore-->>MemorySearcher: token frequencies
MemorySearcher->>MemorySearcher: populate and expire cache entries
MemorySearcher-->>HebbMind: IDF calibration function
HebbMind-->>Client: search result
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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.
Inline comments:
In `@scripts/verify_idf_cache_e2e.py`:
- Line 1: Update the temporary-directory lifecycle in verify_idf_cache_e2e.py by
replacing the manual tempfile.mkdtemp() home allocation with a
tempfile.TemporaryDirectory() context that encloses the verification work,
ensuring cleanup on success and failure. Also update feel_the_issue.py callers
of _new_home() to use TemporaryDirectory contexts for home_a and home_b, while
preserving the existing workload and search behavior.
In `@scripts/verify_idf_cache_real_store.py`:
- Around line 144-162: Strengthen the Check 2 assertion in the three-query
RecallAgent pass so it verifies DF reuse across overlapping tokens, rather than
only requiring one call. Use an upper bound consistent with the documented “far
fewer than 3 queries × terms” expectation, while retaining the corpus_size == 1
and existing lower-bound checks; update the failure message if needed to report
the tightened DF range.
In `@src/hebb/retrieval/searcher.py`:
- Around line 107-133: The IDF caches can grow indefinitely because expired
entries are only removed when their exact keys are accessed. Replace the plain
dicts initialized in the searcher constructor, especially _corpus_size_cache and
_df_cache, with bounded TTL-aware caches configured from _idf_cache_ttl and an
appropriate maximum size, while preserving their existing key/value formats and
cache lookup behavior.
- Around line 593-644: Update _build_idf so cache reads for corpus size and
token document frequencies occur before acquiring _idf_cache_lock, allowing
complete cache hits to proceed without blocking. Enter the lock only when a
required value is missing, then re-check all relevant cache entries under the
lock before performing corpus_size or keyword_doc_freqs and back-filling,
preserving deduplication for identical concurrent misses.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 457ffcda-bddf-4620-966e-a175fb039d0c
📒 Files selected for processing (5)
scripts/feel_the_issue.pyscripts/verify_idf_cache_e2e.pyscripts/verify_idf_cache_real_store.pysrc/hebb/retrieval/searcher.pytests/unit/test_idf_cache.py
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
tests/unit/test_idf_cache.py (1)
70-74: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the required API docstring.
_RecordingStore.keyword_doc_freqshas type hints but no docstring withArgs,Returns, andRaisessections. Apply the same documentation to the other public methods in_RecordingStore.
As per coding guidelines,**/*.pyrequires type hints on all public functions and docstrings withArgs,Returns, andRaisessections for all public APIs.🤖 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 `@tests/unit/test_idf_cache.py` around lines 70 - 74, Add complete API docstrings to _RecordingStore.keyword_doc_freqs and every other public method in _RecordingStore, including Args, Returns, and Raises sections. Preserve the existing type hints and behavior, including defaulting missing terms to document frequency 1.Source: Coding guidelines
scripts/feel_the_issue.py (2)
84-123: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the required public API docstring to
main.
mainis a public function. It has a return annotation, but it has noArgs,Returns, andRaisessections. Add all three sections.As per coding guidelines: public Python APIs require type hints and a docstring with
Args,Returns, andRaisessections.🤖 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 `@scripts/feel_the_issue.py` around lines 84 - 123, Add a docstring to the public main function documenting its arguments (none), integer return value, and any exceptions it may raise, using Args, Returns, and Raises sections. Keep the existing demo behavior and return logic unchanged.Source: Coding guidelines
116-123: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winRequire actual cache savings before reporting success.
The condition at Line 123 succeeds when
saved_c == 0andsaved_d == 0. A disabled cache, or a workload that bypasses IDF, can therefore pass verification. Require positive savings for both counters or assert the expected absolute call counts.Proposed fix
- return 0 if (saved_c >= 0 and saved_d >= 0) else 1 + return 0 if (saved_c > 0 and saved_d > 0) else 1As per PR objectives: the verification must demonstrate cache hits and reused IDF reads.
🤖 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 `@scripts/feel_the_issue.py` around lines 116 - 123, Update the final success condition in the verification flow to require positive savings for both saved_c and saved_d, so zero-savings runs return failure. Keep the existing nonzero success behavior and reporting unchanged.
🧹 Nitpick comments (1)
tests/unit/test_idf_cache.py (1)
70-74: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winTest omitted document-frequency results.
_RecordingStore.keyword_doc_freqsreturns every requested term and defaults unknown terms to1. Therefore,test_zero_statistics_are_cachedcovers an explicit0, but not an omitted term or an empty response.MemorySearcher._build_idfhandles omitted terms withfetched.get(token, 0). Add a fake-store mode that omits the token, then verify that the second call does not incrementdoc_freq_calls. Confirm the realMemoryStore.keyword_doc_freqscontract if omission is not supported.Also applies to: 241-257
🤖 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 `@tests/unit/test_idf_cache.py` around lines 70 - 74, Extend _RecordingStore.keyword_doc_freqs with a mode that omits requested terms, then update test_zero_statistics_are_cached to exercise an omitted-token response and verify the second lookup does not increase doc_freq_calls. Also cover an empty response if appropriate, and confirm the test matches the real MemoryStore.keyword_doc_freqs omission contract.
🤖 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.
Inline comments:
In `@scripts/feel_the_issue.py`:
- Around line 91-108: Update both Settings constructions in the workload setup,
identified by the `sa` and `sb` variables, to explicitly configure `graph=None`,
`vector_search_enabled=False`, `graph_search_enabled=False`, and
`graph_expansion_enabled=False`, matching the lexical-only setup used by
`verify_idf_cache_real_store.py` so `_run_workload` measures only FTS5/keyword
IDF reads.
---
Outside diff comments:
In `@scripts/feel_the_issue.py`:
- Around line 84-123: Add a docstring to the public main function documenting
its arguments (none), integer return value, and any exceptions it may raise,
using Args, Returns, and Raises sections. Keep the existing demo behavior and
return logic unchanged.
- Around line 116-123: Update the final success condition in the verification
flow to require positive savings for both saved_c and saved_d, so zero-savings
runs return failure. Keep the existing nonzero success behavior and reporting
unchanged.
In `@tests/unit/test_idf_cache.py`:
- Around line 70-74: Add complete API docstrings to
_RecordingStore.keyword_doc_freqs and every other public method in
_RecordingStore, including Args, Returns, and Raises sections. Preserve the
existing type hints and behavior, including defaulting missing terms to document
frequency 1.
---
Nitpick comments:
In `@tests/unit/test_idf_cache.py`:
- Around line 70-74: Extend _RecordingStore.keyword_doc_freqs with a mode that
omits requested terms, then update test_zero_statistics_are_cached to exercise
an omitted-token response and verify the second lookup does not increase
doc_freq_calls. Also cover an empty response if appropriate, and confirm the
test matches the real MemoryStore.keyword_doc_freqs omission contract.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 76a4b7e3-723d-4428-9b0d-1c8454700581
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (5)
pyproject.tomlscripts/feel_the_issue.pyscripts/verify_idf_cache_real_store.pysrc/hebb/retrieval/searcher.pytests/unit/test_idf_cache.py
🚧 Files skipped from review as they are similar to previous changes (2)
- src/hebb/retrieval/searcher.py
- scripts/verify_idf_cache_real_store.py
| with _temporary_home() as home_a: | ||
| sa = Settings(home_dir=home_a, llm_model="openai/gpt-4o-mini", | ||
| embedding_provider="noop", embedding_dim=3) | ||
| with HebbMind(config=sa) as hc: | ||
| _seed(hc) | ||
| corpus_a, df_a = _run_workload(hc) | ||
| print(" 3 recall queries + 4 retried same-query searches") | ||
| print(f" → {corpus_a} corpus_size SQL, {df_a} DF SQL") | ||
|
|
||
| # ---- B) Without cache (TTL=0 → every call re-fetches) --------------- # | ||
| print("\n[B] WITHOUT cache (TTL=0 → re-fetch every time):") | ||
| with _temporary_home() as home_b: | ||
| sb = Settings(home_dir=home_b, llm_model="openai/gpt-4o-mini", | ||
| embedding_provider="noop", embedding_dim=3) | ||
| with HebbMind(config=sb) as hc: | ||
| _seed(hc) | ||
| hc._searcher._idf_cache_ttl = 0.0 # neutralise: all entries expire instantly | ||
| corpus_b, df_b = _run_workload(hc) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 5 'load_vec|vector_search_enabled|graph_search_enabled|graph_expansion_enabled|sqlite-vec' scripts src pyproject.toml || trueRepository: afx-team/hebb-mind
Length of output: 26260
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '## feel_the_issue.py outline/section\n'
wc -l scripts/feel_the_issue.py
sed -n '1,180p' scripts/feel_the_issue.py
printf '\n## hebbmind init/call sites\n'
rg -n -C 4 'class HebbMind|def __init__|HebbMind\(|create_stores|Searcher\(' src scripts | head -n 220
printf '\n## settings relevant fields\n'
rg -n -C 3 'embedding_enabled|embedder|skip_search|storage_type|keyword_search_enabled|graph_search_enabled|vector_search_enabled' src/hebb/config src/hebb | head -n 260Repository: afx-team/hebb-mind
Length of output: 36221
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("src/hebb/config/settings.py")
text = p.read_text()
print("settings.py available:", p.exists())
for field in [
"embedding_enabled",
"embedder",
"skip_search",
"storage_type",
"vector_search_enabled",
"keyword_search_enabled",
"graph_search_enabled",
"graph_expansion_enabled",
]:
print(f"\n[{field}] found:", field in text)
idx = text.find(field)
if idx != -1:
line = text[:idx].count("\n") + 1
line_end = text.find("\n", idx + len(field)) + 1
start = max(0, text.rfind("\n", 0, idx - 300))
print(f"line {line}:", text[max(line_idx, start):line_end.strip() or line_end])
PYRepository: afx-team/hebb-mind
Length of output: 333
Make the IDF workload lexical-only by construction.
HebbMind(config=Settings(..., embedding_provider="noop")) still enables the default SQLite-backed search path: embedding_enabled=True can load sqlite-vec, and MemorySearcher defaults enable vector, keyword, graph, and expansion paths. Add the equivalent graph=None, vector_search_enabled=False, graph_search_enabled=False, and graph_expansion_enabled=False setup seen in scripts/verify_idf_cache_real_store.py so this counts only the FTS5/keyword IDF reads.
🤖 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 `@scripts/feel_the_issue.py` around lines 91 - 108, Update both Settings
constructions in the workload setup, identified by the `sa` and `sb` variables,
to explicitly configure `graph=None`, `vector_search_enabled=False`,
`graph_search_enabled=False`, and `graph_expansion_enabled=False`, matching the
lexical-only setup used by `verify_idf_cache_real_store.py` so `_run_workload`
measures only FTS5/keyword IDF reads.
1 participant
Cache IDF corpus statistics to cut DB round-trips on repeated recall
Closes #54.
What
MemorySearcher._build_idfnow keeps two short-TTL (60s), instance-levelstatistic caches so a repeated or overlapping
search()over an unchangedcorpus reuses
corpus_sizeand per-token document frequencies instead ofre-hitting the store:
_corpus_size_cache— keyed by the partition scope (corpus size dependsonly on which partitions are searched, not the query).
_df_cache— keyed by(token, partition_scope); DF depends on both, sothe token is in the key to stop one query overwriting another's DFs.
Miss → store round-trip → back-fill (including
0/empty results, so an emptycorpus or an unseen token does not re-hit the store). No storage-layer
changes, no write-driven invalidation — the short TTL bounds staleness, per
the issue.
Implementation note vs. the issue text
The issue describes
keyword_doc_freqsas a per-term loop. The real store APItakes a
terms: list[str]and returns the wholedict[str, int]in one call.The cache is therefore per-token: on a miss the whole token set is fetched
in one batch (as the API already does), then each token's df is back-filled
into
_df_cacheindividually. This matches the real API and still deliversthe issue's "reuses cached DF for overlapping tokens" outcome — overlapping
tokens hit the per-token cache on the next query rather than being re-fetched.
Flagging this so the divergence from the issue's literal wording is explicit.
Extra hardening (not in the issue — feel free to drop)
Two things I added beyond the issue's scope; both are clearly separable and
I'm happy to remove either if you'd rather keep this PR tightly scoped:
asyncio.Lockguards themiss → fetch → back-fill section of
_build_idf, so concurrent coroutinesthat miss the same key share one fetch (a waiter re-checks the cache under
the lock and finds it populated by the leader). Correctness is unchanged
without it — the result is right either way — but the lock removes a wasted
round-trip under concurrency (e.g.
RecallAgent's parallel recall paths).Covered by
test_concurrent_miss_deduplicates_store_calls.scripts/(verify_idf_cache_real_store.py,verify_idf_cache_e2e.py,feel_the_issue.py). Not part of the test suite(ruff-excluded dir); they exist to demonstrate the cache against a real
HebbMind+ SQLite stack. Drop them if you don't want extra files in tree.Tests
tests/unit/test_idf_cache.py(8 cases, all passing):Verification
ruff check src/— clean.pytest tests/unit/test_idf_cache.py tests/unit/test_audit_retrieval.py— 15 passed, no regressions.The
scripts/checks run the cache against a realSQLiteMemoryStoreand thereal
HebbMindfacade, confirming a repeated search issues zero IDF DBround-trips and the
RecallAgent-style 3-query pass reuses cached stats.mypy src/hebb/ --strict:retrieval/searcher.pyis clean. (The repo reportssome pre-existing errors in
server/and embedding modules from missingthird-party stubs in my local env — unrelated to this change; verified via
git stashthat the count is identical without these edits.)Summary by CodeRabbit
Performance
Bug Fixes
Tests