feat: auto-enrich on brain_store via Gemini Flash Lite#115
Conversation
…ass) Wire async Gemini enrichment into brain_store's background thread. When a new chunk is stored, embedding happens first (pass 1), then Gemini 2.5 Flash Lite enriches with faceted tags/summary (pass 2, ~600ms). - Add enrich_single() to enrichment_controller.py — enriches one chunk by ID, low retry count (2), short backoff (0.3s base) - Wire into _background_embed_and_flush in store_handler.py - Opt-out via BRAINLAYER_AUTO_ENRICH=0 env var - Never blocks store response — fire-and-forget in daemon thread - 19 new tests with mock Gemini client Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
Your free trial has ended. If you'd like to continue receiving code reviews, you can add a payment method here.
📝 WalkthroughWalkthroughThis PR adds an optional automatic enrichment pass using Gemini: a new Changes
Sequence DiagramsequenceDiagram
participant StoreThread as Store Handler (Background)
participant EnrichCtrl as Enrichment Controller
participant GeminiAPI as Gemini API Client
participant Store as Vector Store (DB)
participant Parser as Enrichment Parser
StoreThread->>Store: Store chunk & embed
StoreThread->>EnrichCtrl: enrich_single(store, chunk_id)
alt AUTO_ENRICH_ENABLED
EnrichCtrl->>Store: get_chunk(chunk_id)
alt Chunk exists
EnrichCtrl->>EnrichCtrl: build_external_prompt(chunk + Sanitizer)
EnrichCtrl->>GeminiAPI: generate_content(prompt) with retry/backoff
alt API Success
GeminiAPI-->>EnrichCtrl: response
EnrichCtrl->>Parser: parse_enrichment(response)
Parser-->>EnrichCtrl: {summary, tags, ...}
EnrichCtrl->>Store: _apply_enrichment(chunk_id, enrichment)
EnrichCtrl-->>StoreThread: enrichment dict
else API Failure
GeminiAPI-->>EnrichCtrl: error (after retries)
EnrichCtrl-->>StoreThread: None (logged)
end
else Chunk Missing
EnrichCtrl-->>StoreThread: None (early return)
end
else Feature Disabled
EnrichCtrl-->>StoreThread: None (early return)
end
Estimated Code Review Effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 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 |
![macroscopeapp[bot]](https://avatars.githubusercontent.com/in/900172?s=60&v=4){kind=small}
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
EtanHey
left a comment
EtanHey
left a comment
There was a problem hiding this comment.
Code Review: PR #115 -- Auto-enrich on brain_store via Gemini Flash Lite
What was done well
- Clean two-pass architecture. The R47 pattern is well-designed: pass 1 (sync embedding) keeps the store response fast, pass 2 (async Gemini enrichment) fires in the existing background thread. No new threads, no new coordination primitives.
- Thorough error isolation. Every failure path in
enrich_singlereturnsNoneinstead of propagating. The store_handler wraps the call in its own try/except, making it truly fire-and-forget. The store result is never affected by enrichment failures. - Good reuse of existing infrastructure.
enrich_singlereuses_get_gemini_client,_retry_with_backoff,build_external_prompt,parse_enrichment, and_apply_enrichment-- no reinvention. - Sensible retry parameters.
max_retries=2withbase_delay=0.3andmax_delay=5.0is appropriate for inline enrichment (vs the batchmax_retries=12withmax_delay=120).
Critical -- Must Fix
1. CI lint failure: tests/test_auto_enrich.py fails ruff format --check
Three formatting issues flagged by ruff: _fake_gemini_response dict, the cursor.execute() call in test_writes_enrichment_to_db, and the @pytest.mark.parametrize list. Run ruff format tests/test_auto_enrich.py to fix.
Important -- Should Fix
2. TestAutoEnrichEnvVar.test_auto_enrich_flag_parsing tests itself, not the production code
(Macroscope already flagged this -- I agree with their assessment.)
Line 327 manually computes value.lower() not in ("0", "false", "no") rather than importing and checking the actual AUTO_ENRICH_ENABLED module attribute. This test will never catch regressions in the production parsing logic. The fix is:
def test_auto_enrich_flag_parsing(self, value, expected, monkeypatch):
import importlib
from brainlayer import enrichment_controller as ctrl
monkeypatch.setenv("BRAINLAYER_AUTO_ENRICH", value)
importlib.reload(ctrl)
assert ctrl.AUTO_ENRICH_ENABLED == expectedNote that importlib.reload is needed because AUTO_ENRICH_ENABLED is computed at module import time. After the test, monkeypatch will restore the env var, but you may want to reload again in a finalizer if other tests in the same process depend on the original value.
3. _FakeGeminiClient has a confusing __init__ / __new__ split
__init__ is defined but never actually runs because __new__ does all the initialization. This is misleading -- a reader sees __init__ setting self._response_text and self.call_count and assumes that is the constructor. Since __new__ already handles everything, just remove the __init__ method entirely to avoid confusion.
4. Integration test thread-join pattern is fragile
for t in threading.enumerate():
if t.daemon and t.name != "MainThread":
t.join(timeout=5.0)This joins all daemon threads in the process, not just the one spawned by _store. In a parallel pytest run or if other fixtures spawn daemon threads, this could join unrelated threads or miss the target thread entirely. A more robust approach: capture the thread reference directly. The background thread is spawned as t = threading.Thread(...) in _store -- you could monkeypatch threading.Thread to capture the thread instance, or use an Event to signal completion from the mock enrich_single.
Suggestions -- Nice to Have
5. Logger instantiation inside function body
enrich_single creates _logger = logging.getLogger(__name__) on every call. This is cheap but unconventional -- the module-level pattern _logger = logging.getLogger(__name__) (once, at import time) is standard in this codebase. Consider moving it to module level for consistency.
6. Consider testing the enriched_at timestamp
test_writes_enrichment_to_db verifies summary and tags but does not verify that enriched_at was set. Since _apply_enrichment calls update_enrichment, which sets enriched_at, asserting on it would confirm the full write path. This also matters for the enrichment pipeline which uses enriched_at to skip already-enriched chunks.
7. Missing test: build_external_prompt failure path
enrich_single handles prompt-build exceptions (lines 141-145), but no test covers this path. A quick test with build_external_prompt monkeypatched to raise would complete the error-path coverage.
8. Missing test: _apply_enrichment failure path
Similarly, the _apply_enrichment exception handler (lines 173-177) has no test coverage. Monkeypatch store.update_enrichment to raise and verify enrich_single returns None.
Summary
The implementation is architecturally sound and follows the existing patterns well. The critical fix is the ruff formatting (CI blocker). The important fixes are the env-var test that doesn't test production code, the confusing __new__/__init__ split, and the fragile thread-join pattern. Everything else is polish.
19 tests all pass locally. The CI test failure (test_digest_connect.py) is pre-existing and unrelated to this PR.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/brainlayer/enrichment_controller.py`:
- Around line 110-121: enrich_single currently may race with enrich_realtime and
cause duplicate Gemini calls; add a pre-flight guard at the start of
enrich_single that queries the chunk's enrichment state (check the enriched_at
field for the chunk_id) and return early (None) if enriched_at is not null, so
only one writer proceeds; update the function's docstring to note this
pre-flight check and reference enrich_realtime and the enriched_at field so
reviewers can see the concurrency safeguard.
In `@tests/test_auto_enrich.py`:
- Around line 323-328: The test test_auto_enrich_flag_parsing is duplicating the
parsing logic instead of asserting the production flag; change it to set the env
via monkeypatch.setenv("BRAINLAYER_AUTO_ENRICH", value) and then import the
module under test and assert its AUTO_ENRICH_ENABLED equals expected (use
importlib.reload(module) after setting the env to pick up module-level
evaluation), or better: extract the parsing into a helper function (e.g.,
parse_auto_enrich_env) in the production module and call that helper from the
test to verify parsing behavior; ensure you reference AUTO_ENRICH_ENABLED or the
new parse_auto_enrich_env helper and remove the duplicated parsing logic from
the test.
- Around line 61-81: The _FakeGeminiClient currently duplicates initialization
in both __init__ and __new__ (setting _response_text and call_count), so remove
the redundant __init__ method and rely on the __new__ implementation to set
instance._response_text, instance.call_count, and instance.models (which uses
the nested _Models class and increments call_count via generate_content); ensure
generate_content still references self._parent and that models remains attached
to the returned instance.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: c84a0766-a663-43a0-becc-0a4bb67a5aa0
📒 Files selected for processing (3)
src/brainlayer/enrichment_controller.pysrc/brainlayer/mcp/store_handler.pytests/test_auto_enrich.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Macroscope - Correctness Check
- GitHub Check: test (3.13)
- GitHub Check: test (3.11)
- GitHub Check: test (3.12)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Flag risky DB or concurrency changes explicitly and do not hand-wave lock behavior
Enforce one-write-at-a-time concurrency constraint; reads are safe but brain_digest is write-heavy and must not run in parallel with other MCP work
Run pytest before claiming behavior changed safely; current test suite has 929 tests
**/*.py: All scripts and CLI must usepaths.py:get_db_path()for database path resolution instead of hardcoding paths
Preserve verbatim content for message types:ai_code,stack_trace,user_messageduring classification and chunking
Skipnoisecontent entirely; summarizebuild_log; extract structure only fromdir_listingduring chunking
Use AST-aware chunking via tree-sitter; never split stack traces; mask large tool output
Implement retry logic onSQLITE_BUSYerrors; each worker must use its own database connection
Override enrichment backend viaBRAINLAYER_ENRICH_BACKENDenvironment variable (valid values:ollama,mlx,groq); default to Groq
Configure enrichment rate viaBRAINLAYER_ENRICH_RATEenvironment variable (default: 0.2 = 12 RPM)
Checkpoint WAL before and after bulk database operations:PRAGMA wal_checkpoint(FULL)
Drop FTS triggers before bulk deletes fromchunkstable; recreate after operation to avoid performance degradation
Batch deletes in 5-10K chunk sizes with WAL checkpoint every 3 batches
Default search queries must exclude lifecycle-managed chunks; useinclude_archived=Trueparameter to show history
Lint and format code withruff check src/ && ruff format src/
Session dedup coordination: SessionStart writes injected chunk_ids to/tmp/brainlayer_session_{id}.json; UserPromptSubmit skips already-injected chunks
Skip auto-search for prompts containing 'handoff' or 'session-handoff' keywords
Files:
src/brainlayer/mcp/store_handler.pysrc/brainlayer/enrichment_controller.pytests/test_auto_enrich.py
**/*test*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use
pytestfor testing
Files:
tests/test_auto_enrich.py
🧠 Learnings (3)
📚 Learning: 2026-03-22T15:55:22.017Z
Learnt from: EtanHey
Repo: EtanHey/brainlayer PR: 100
File: src/brainlayer/enrichment_controller.py:175-199
Timestamp: 2026-03-22T15:55:22.017Z
Learning: In `src/brainlayer/enrichment_controller.py`, the `parallel` parameter in `enrich_local()` is intentionally kept in the function signature (currently unused, suppressed with `# noqa: ARG001`) for API stability. Parallel local enrichment via a thread pool or process pool is planned for a future iteration. Do not flag this as dead code requiring removal.
Applied to files:
src/brainlayer/mcp/store_handler.pysrc/brainlayer/enrichment_controller.py
📚 Learning: 2026-03-26T15:46:16.139Z
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-26T15:46:16.139Z
Learning: Applies to **/*.py : Configure enrichment rate via `BRAINLAYER_ENRICH_RATE` environment variable (default: 0.2 = 12 RPM)
Applied to files:
src/brainlayer/enrichment_controller.py
📚 Learning: 2026-03-26T15:46:16.139Z
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-26T15:46:16.139Z
Learning: Applies to **/*.py : Override enrichment backend via `BRAINLAYER_ENRICH_BACKEND` environment variable (valid values: `ollama`, `mlx`, `groq`); default to Groq
Applied to files:
src/brainlayer/enrichment_controller.pytests/test_auto_enrich.py
🔇 Additional comments (7)
src/brainlayer/enrichment_controller.py (3)
22-23: LGTM!The environment variable parsing correctly handles multiple falsey values (
0,false,no) with case-insensitive comparison and defaults to enabled.
134-138: LGTM!Graceful degradation when Gemini credentials are unavailable — logs at debug level (appropriate since this is expected in dev/test environments) and returns
Nonewithout disrupting the store flow.
157-166: LGTM!The retry parameters (
base_delay=0.3,max_delay=5.0,max_retries=2) are well-tuned for the ~600ms target latency goal stated in the PR objectives. Short backoffs prevent blocking the background thread while still providing resilience against transient failures.src/brainlayer/mcp/store_handler.py (1)
534-542: LGTM!The fire-and-forget enrichment integration is well-structured:
- Exception isolation ensures store success is never affected by enrichment failures
- Debug-level logging is appropriate for expected scenarios (disabled enrichment, missing API key)
- Uses the background thread's dedicated
bg_storeconnection, respecting APSW's same-thread requirementtests/test_auto_enrich.py (3)
19-42: LGTM!Clean fixture setup with proper resource cleanup. The
stored_chunkfixture correctly chains fromstoreand provides a realistic test chunk for enrichment scenarios.
150-150: Verifytime.sleeppatch target is correct.The patch
ctrl.time.sleepassumestimeis imported as a module inenrichment_controller. Iffrom time import sleepwere used instead, this patch would fail silently. Confirm thatenrichment_controller.pyusesimport time(line 8 shows it does).
260-266: Consider using a more robust thread synchronization pattern.The daemon thread enumeration works but may be fragile if other daemon threads exist. A more targeted approach would capture the thread returned by
_storeor use an event:# Alternative: could mock threading.Thread to capture the specific threadThat said, the current approach works for these isolated test cases.
| def enrich_single(store, chunk_id: str, max_retries: int = 2) -> dict[str, Any] | None: | ||
| """Enrich a single chunk by ID via Gemini 2.5 Flash Lite. | ||
|
|
||
| Designed for post-store auto-enrichment (R47 two-pass pattern): | ||
| Pass 1 = sync embedding (immediate, searchable) | ||
| Pass 2 = this function (async Gemini tagging, ~600ms target) | ||
|
|
||
| Bypasses get_enrichment_candidates — works on any chunk regardless | ||
| of enriched_at status. Overwrites stub enrichment with Gemini output. | ||
|
|
||
| Returns the parsed enrichment dict on success, None on failure. | ||
| """ |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if there's any rate limiting logic or deduplication for concurrent enrichment
rg -nC3 "enriched_at" src/brainlayer/Repository: EtanHey/brainlayer
Length of output: 6281
🏁 Script executed:
# Find where enrich_single is called
rg -n "enrich_single" src/brainlayer/Repository: EtanHey/brainlayer
Length of output: 1319
🏁 Script executed:
# Check for any locking/synchronization mechanisms around enrichment
rg -n "Lock|Event|Thread|await|asyncio" src/brainlayer/enrichment_controller.py | head -20Repository: EtanHey/brainlayer
Length of output: 44
🏁 Script executed:
# Look at the _background_embed_and_flush function mentioned in the docstring
rg -n "_background_embed_and_flush" src/brainlayer/ -A 20Repository: EtanHey/brainlayer
Length of output: 3642
🏁 Script executed:
# Check what enrich_realtime does and how it's scheduled
rg -n "enrich_realtime" src/brainlayer/ -B 3 -A 8Repository: EtanHey/brainlayer
Length of output: 4111
🏁 Script executed:
# Check _apply_enrichment implementation to see when enriched_at is written
rg -n "_apply_enrichment" src/brainlayer/enrichment_controller.py -A 25Repository: EtanHey/brainlayer
Length of output: 3608
🏁 Script executed:
# Check store.update_enrichment to see if it's atomic
rg -n "def update_enrichment" src/brainlayer/ -A 30Repository: EtanHey/brainlayer
Length of output: 2307
Concurrent enrichment is possible but unlikely in practice; consider explicit synchronization per concurrency guidelines.
The race condition scenario is technically possible: if enrich_realtime (manual trigger) runs while enrich_single (daemon auto-enrichment) is executing, both may fetch and enrich the same chunk, resulting in duplicate Gemini API calls and unnecessary quota consumption.
However, in typical usage this overlap is unlikely — enrich_realtime is invoked manually via CLI or MCP, not automatically, so it rarely runs while auto-enrichment threads are active. Additionally, the enrich_single docstring explicitly documents that it overwrites existing enrichment by design (two-pass pattern: Pass 1 = embedding, Pass 2 = async Gemini).
That said, the codebase guidelines require explicitly flagging risky concurrent database changes and enforcing one-write-at-a-time constraints. This function currently has no guard against concurrent writes from enrich_realtime. Consider either:
- Adding a pre-flight check to skip chunks where
enriched_at IS NOT NULL(if deliberate overlap is not desired), or - Documenting that concurrent enrichment is acceptable and updating the guidelines to reflect this exception.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/brainlayer/enrichment_controller.py` around lines 110 - 121,
enrich_single currently may race with enrich_realtime and cause duplicate Gemini
calls; add a pre-flight guard at the start of enrich_single that queries the
chunk's enrichment state (check the enriched_at field for the chunk_id) and
return early (None) if enriched_at is not null, so only one writer proceeds;
update the function's docstring to note this pre-flight check and reference
enrich_realtime and the enriched_at field so reviewers can see the concurrency
safeguard.
| class _FakeGeminiClient: | ||
| """Mock Gemini client that returns a valid enrichment response.""" | ||
|
|
||
| def __init__(self, response_text=None): | ||
| self._response_text = response_text or _fake_gemini_response() | ||
| self.call_count = 0 | ||
|
|
||
| class _Models: | ||
| def __init__(self, parent): | ||
| self._parent = parent | ||
|
|
||
| def generate_content(self, **kwargs): | ||
| self._parent.call_count += 1 | ||
| return SimpleNamespace(text=self._parent._response_text) | ||
|
|
||
| def __new__(cls, response_text=None): | ||
| instance = super().__new__(cls) | ||
| instance._response_text = response_text or _fake_gemini_response() | ||
| instance.call_count = 0 | ||
| instance.models = cls._Models(instance) | ||
| return instance |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Redundant initialization in __init__ and __new__.
Both methods initialize _response_text and call_count. Since __init__ runs after __new__, the __new__ initialization is immediately overwritten. Consider removing __init__ entirely since __new__ already handles setup:
♻️ Simplified fake client
class _FakeGeminiClient:
"""Mock Gemini client that returns a valid enrichment response."""
- def __init__(self, response_text=None):
- self._response_text = response_text or _fake_gemini_response()
- self.call_count = 0
-
class _Models:
def __init__(self, parent):
self._parent = parent🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_auto_enrich.py` around lines 61 - 81, The _FakeGeminiClient
currently duplicates initialization in both __init__ and __new__ (setting
_response_text and call_count), so remove the redundant __init__ method and rely
on the __new__ implementation to set instance._response_text,
instance.call_count, and instance.models (which uses the nested _Models class
and increments call_count via generate_content); ensure generate_content still
references self._parent and that models remains attached to the returned
instance.
- Fix ruff format on test file - Remove __new__/__init__ split in _FakeGeminiClient (use __init__ only) - Add test_returns_none_on_prompt_build_failure - Add test_returns_none_on_apply_failure - 21 tests total (was 19) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
| test_store.close() | ||
|
|
||
|
|
||
| # ── Environment variable opt-out ──────────────────────────────── |
There was a problem hiding this comment.
🟡 Medium tests/test_auto_enrich.py:335
test_auto_enrich_flag_parsing uses importlib.reload(ctrl) to test different environment variable values, but never reloads the module again after the test completes. This leaves AUTO_ENRICH_ENABLED set to the value from the last parametrized case (empty string → True), polluting global state for any subsequent tests that don't explicitly monkeypatch the flag.
Consider using a fixture or try/finally to reload the module with a clean environment (or delete ctrl from sys.modules) after each parametrized run.
🚀 Reply "fix it for me" or copy this AI Prompt for your agent:
In file tests/test_auto_enrich.py around line 335:
`test_auto_enrich_flag_parsing` uses `importlib.reload(ctrl)` to test different environment variable values, but never reloads the module again after the test completes. This leaves `AUTO_ENRICH_ENABLED` set to the value from the last parametrized case (empty string → `True`), polluting global state for any subsequent tests that don't explicitly monkeypatch the flag.
Consider using a fixture or `try/finally` to reload the module with a clean environment (or delete `ctrl` from `sys.modules`) after each parametrized run.
Evidence trail:
tests/test_auto_enrich.py lines 343-365 (test code with importlib.reload), src/brainlayer/enrichment_controller.py line 26 (AUTO_ENRICH_ENABLED module-level definition). No cleanup code is present in the test after the reload.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tests/test_auto_enrich.py`:
- Around line 288-295: The test currently joins all daemon threads via
threading.enumerate(), which may catch unrelated threads; modify the code to
wait only for the specific background worker started by _store_new (or signal
completion via an Event). Change _store_new (or the code that spawns the worker)
to accept/return a threading.Thread or accept a threading.Event that it will set
when enrichment is done, then in the test wait on that thread.join(timeout=5.0)
or event.wait(timeout=5.0) before asserting enriched_ids and closing test_store;
reference the _store_new function and the enriched_ids check so the test only
waits for that specific operation.
- Around line 355-364: The test reloads the brainlayer.enrichment_controller
module which mutates sys.modules and leaves AUTO_ENRICH_ENABLED changed for
later tests; to fix, in test_auto_enrich_flag_parsing save the existing module
entry (orig = sys.modules.get('brainlayer.enrichment_controller')) before you
call importlib.reload(ctrl), then after the assertion restore the original
module state by putting orig back into sys.modules (or deleting the entry if
orig was None) so that the reload has no lasting effect on AUTO_ENRICH_ENABLED;
reference the test function test_auto_enrich_flag_parsing, the ctrl variable,
and the AUTO_ENRICH_ENABLED symbol when making the change.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 7d493b29-a4a7-4b40-ac74-82eb303f32f3
📒 Files selected for processing (2)
src/brainlayer/enrichment_controller.pytests/test_auto_enrich.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: test (3.11)
- GitHub Check: test (3.12)
- GitHub Check: test (3.13)
- GitHub Check: Macroscope - Correctness Check
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Flag risky DB or concurrency changes explicitly and do not hand-wave lock behavior
Enforce one-write-at-a-time concurrency constraint; reads are safe but brain_digest is write-heavy and must not run in parallel with other MCP work
Run pytest before claiming behavior changed safely; current test suite has 929 tests
**/*.py: All scripts and CLI must usepaths.py:get_db_path()for database path resolution instead of hardcoding paths
Preserve verbatim content for message types:ai_code,stack_trace,user_messageduring classification and chunking
Skipnoisecontent entirely; summarizebuild_log; extract structure only fromdir_listingduring chunking
Use AST-aware chunking via tree-sitter; never split stack traces; mask large tool output
Implement retry logic onSQLITE_BUSYerrors; each worker must use its own database connection
Override enrichment backend viaBRAINLAYER_ENRICH_BACKENDenvironment variable (valid values:ollama,mlx,groq); default to Groq
Configure enrichment rate viaBRAINLAYER_ENRICH_RATEenvironment variable (default: 0.2 = 12 RPM)
Checkpoint WAL before and after bulk database operations:PRAGMA wal_checkpoint(FULL)
Drop FTS triggers before bulk deletes fromchunkstable; recreate after operation to avoid performance degradation
Batch deletes in 5-10K chunk sizes with WAL checkpoint every 3 batches
Default search queries must exclude lifecycle-managed chunks; useinclude_archived=Trueparameter to show history
Lint and format code withruff check src/ && ruff format src/
Session dedup coordination: SessionStart writes injected chunk_ids to/tmp/brainlayer_session_{id}.json; UserPromptSubmit skips already-injected chunks
Skip auto-search for prompts containing 'handoff' or 'session-handoff' keywords
Files:
tests/test_auto_enrich.pysrc/brainlayer/enrichment_controller.py
**/*test*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use
pytestfor testing
Files:
tests/test_auto_enrich.py
🧠 Learnings (8)
📓 Common learnings
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-26T15:46:16.139Z
Learning: Use Groq as primary enrichment backend (configured in launchd plist); fall back to Gemini, then Ollama as offline last-resort
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-26T15:46:16.139Z
Learning: Applies to **/*.py : Override enrichment backend via `BRAINLAYER_ENRICH_BACKEND` environment variable (valid values: `ollama`, `mlx`, `groq`); default to Groq
Learnt from: EtanHey
Repo: EtanHey/brainlayer PR: 100
File: src/brainlayer/enrichment_controller.py:175-199
Timestamp: 2026-03-22T15:55:22.017Z
Learning: In `src/brainlayer/enrichment_controller.py`, the `parallel` parameter in `enrich_local()` is intentionally kept in the function signature (currently unused, suppressed with `# noqa: ARG001`) for API stability. Parallel local enrichment via a thread pool or process pool is planned for a future iteration. Do not flag this as dead code requiring removal.
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-26T15:46:16.139Z
Learning: Applies to **/*.py : Configure enrichment rate via `BRAINLAYER_ENRICH_RATE` environment variable (default: 0.2 = 12 RPM)
📚 Learning: 2026-03-22T15:55:22.017Z
Learnt from: EtanHey
Repo: EtanHey/brainlayer PR: 100
File: src/brainlayer/enrichment_controller.py:175-199
Timestamp: 2026-03-22T15:55:22.017Z
Learning: In `src/brainlayer/enrichment_controller.py`, the `parallel` parameter in `enrich_local()` is intentionally kept in the function signature (currently unused, suppressed with `# noqa: ARG001`) for API stability. Parallel local enrichment via a thread pool or process pool is planned for a future iteration. Do not flag this as dead code requiring removal.
Applied to files:
tests/test_auto_enrich.pysrc/brainlayer/enrichment_controller.py
📚 Learning: 2026-03-26T15:46:16.139Z
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-26T15:46:16.139Z
Learning: Applies to **/*.py : Override enrichment backend via `BRAINLAYER_ENRICH_BACKEND` environment variable (valid values: `ollama`, `mlx`, `groq`); default to Groq
Applied to files:
tests/test_auto_enrich.pysrc/brainlayer/enrichment_controller.py
📚 Learning: 2026-03-26T15:46:16.139Z
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-26T15:46:16.139Z
Learning: Applies to **/*.py : Configure enrichment rate via `BRAINLAYER_ENRICH_RATE` environment variable (default: 0.2 = 12 RPM)
Applied to files:
src/brainlayer/enrichment_controller.py
📚 Learning: 2026-03-26T15:46:16.139Z
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: CLAUDE.md:0-0
Timestamp: 2026-03-26T15:46:16.139Z
Learning: Never run bulk database operations while enrichment workers are writing; stop enrichment workers first
Applied to files:
src/brainlayer/enrichment_controller.py
📚 Learning: 2026-03-14T02:20:54.656Z
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: AGENTS.md:0-0
Timestamp: 2026-03-14T02:20:54.656Z
Learning: Applies to **/*.py : Enforce one-write-at-a-time concurrency constraint; reads are safe but brain_digest is write-heavy and must not run in parallel with other MCP work
Applied to files:
src/brainlayer/enrichment_controller.py
📚 Learning: 2026-03-14T02:20:54.656Z
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: AGENTS.md:0-0
Timestamp: 2026-03-14T02:20:54.656Z
Learning: Applies to **/*.py : Flag risky DB or concurrency changes explicitly and do not hand-wave lock behavior
Applied to files:
src/brainlayer/enrichment_controller.py
📚 Learning: 2026-03-14T02:20:54.656Z
Learnt from: CR
Repo: EtanHey/brainlayer PR: 0
File: AGENTS.md:0-0
Timestamp: 2026-03-14T02:20:54.656Z
Learning: Be aware of known BrainLayer issues: DB locking during enrichment and WAL growth up to 4.7GB
Applied to files:
src/brainlayer/enrichment_controller.py
🔇 Additional comments (8)
src/brainlayer/enrichment_controller.py (5)
113-124: Concurrent enrichment overlap: unchanged from prior review.The past review flagged that
enrich_singleandenrich_realtimemay concurrently enrich the same chunk since there's no pre-flightenriched_at IS NOT NULLguard. The docstring documents this is by design ("Overwrites stub enrichment"). No code change here since last review.
6-6: LGTM!Module-level logger follows best practices with
__name__scoping.Also applies to: 14-15
25-27: LGTM!Environment variable parsing correctly handles common truthy/falsy values with case-insensitive comparison.
133-137: Verify intentional Gemini-only backend for auto-enrichment.The function uses Gemini directly via
_get_gemini_client(), while learnings indicate "Use Groq as primary enrichment backend; fall back to Gemini." The PR title explicitly mentions "Gemini Flash Lite," suggesting this is intentional for the fast two-pass pattern. Please confirm this is the desired behavior or ifBRAINLAYER_ENRICH_BACKENDshould be respected here as well.Based on learnings: "Use Groq as primary enrichment backend (configured in launchd plist); fall back to Gemini, then Ollama as offline last-resort."
156-162: LGTM!Retry parameters (
base_delay=0.3,max_delay=5.0,max_retries=2) are appropriately tuned for the ~600ms target latency of pass-2 enrichment.tests/test_auto_enrich.py (3)
19-43: LGTM!Fixtures properly manage resource lifecycle with
yieldandclose(). Thestored_chunkfixture correctly depends onstorefor test isolation.
63-78: LGTM!Clean mock implementation. The previous concern about redundant
__init__/__new__has been addressed—only__init__remains.
83-104: LGTM!Good test structure with direct monkeypatching of
AUTO_ENRICH_ENABLEDto avoid module reload complications. Validates both return value and API call count.
| # Wait for background thread to complete | ||
| for t in threading.enumerate(): | ||
| if t.daemon and t.name != "MainThread": | ||
| t.join(timeout=5.0) | ||
|
|
||
| assert len(enriched_ids) == 1 | ||
| assert enriched_ids[0] == chunk_id | ||
| test_store.close() |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Thread-joining approach works but is fragile.
Joining all daemon threads by iterating threading.enumerate() could catch unrelated daemon threads from the test harness. Consider tracking the specific thread spawned by _store_new or using an Event to signal completion.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_auto_enrich.py` around lines 288 - 295, The test currently joins
all daemon threads via threading.enumerate(), which may catch unrelated threads;
modify the code to wait only for the specific background worker started by
_store_new (or signal completion via an Event). Change _store_new (or the code
that spawns the worker) to accept/return a threading.Thread or accept a
threading.Event that it will set when enrichment is done, then in the test wait
on that thread.join(timeout=5.0) or event.wait(timeout=5.0) before asserting
enriched_ids and closing test_store; reference the _store_new function and the
enriched_ids check so the test only waits for that specific operation.
| def test_auto_enrich_flag_parsing(self, value, expected, monkeypatch): | ||
| """AUTO_ENRICH_ENABLED respects environment variable values.""" | ||
| import importlib | ||
|
|
||
| from brainlayer import enrichment_controller as ctrl | ||
|
|
||
| monkeypatch.setenv("BRAINLAYER_AUTO_ENRICH", value) | ||
| importlib.reload(ctrl) | ||
|
|
||
| assert ctrl.AUTO_ENRICH_ENABLED == expected |
There was a problem hiding this comment.
Module reload without cleanup may pollute other tests.
After importlib.reload(ctrl), the AUTO_ENRICH_ENABLED constant retains the last-tested value for the remainder of the test session. While TestEnrichSingle tests explicitly monkeypatch AUTO_ENRICH_ENABLED, any future tests that don't could inherit stale state.
Consider restoring the original module state:
🧪 Proposed fix to restore module state after reload
def test_auto_enrich_flag_parsing(self, value, expected, monkeypatch):
"""AUTO_ENRICH_ENABLED respects environment variable values."""
import importlib
from brainlayer import enrichment_controller as ctrl
+ original_value = ctrl.AUTO_ENRICH_ENABLED
monkeypatch.setenv("BRAINLAYER_AUTO_ENRICH", value)
importlib.reload(ctrl)
- assert ctrl.AUTO_ENRICH_ENABLED == expected
+ try:
+ assert ctrl.AUTO_ENRICH_ENABLED == expected
+ finally:
+ # Restore default state for subsequent tests
+ monkeypatch.delenv("BRAINLAYER_AUTO_ENRICH", raising=False)
+ importlib.reload(ctrl)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_auto_enrich.py` around lines 355 - 364, The test reloads the
brainlayer.enrichment_controller module which mutates sys.modules and leaves
AUTO_ENRICH_ENABLED changed for later tests; to fix, in
test_auto_enrich_flag_parsing save the existing module entry (orig =
sys.modules.get('brainlayer.enrichment_controller')) before you call
importlib.reload(ctrl), then after the assertion restore the original module
state by putting orig back into sys.modules (or deleting the entry if orig was
None) so that the reload has no lasting effect on AUTO_ENRICH_ENABLED; reference
the test function test_auto_enrich_flag_parsing, the ctrl variable, and the
AUTO_ENRICH_ENABLED symbol when making the change.
1 participant
Summary
brain_store: sync embedding (pass 1, immediate) + async Gemini 2.5 Flash Lite enrichment (pass 2, ~600ms target)enrich_single()toenrichment_controller.py— enriches a single chunk by ID with low retry count (2) and short backoffBRAINLAYER_AUTO_ENRICH=0env varChanged files
src/brainlayer/enrichment_controller.pyenrich_single(),AUTO_ENRICH_ENABLEDflagsrc/brainlayer/mcp/store_handler.pyenrich_singlecall into_background_embed_and_flushtests/test_auto_enrich.pyDesign
Test plan
enrich_singleenriches chunk successfully with mock Gemini_storebackground thread triggersenrich_single🤖 Generated with Claude Code
Note
Add auto-enrichment via Gemini Flash Lite on
brain_storeenrich_singleinenrichment_controller.pythat fetches a chunk, calls Gemini to generate enrichment, and persists the result to the store; all failures returnNonewithout raising.store_handler.py,enrich_singleis called for the newly stored chunk; exceptions are caught and logged at debug level, leaving store behavior unchanged.BRAINLAYER_AUTO_ENRICHenvironment variable (disabled when set to0,false, orno).brain_storecall now triggers a background Gemini API request whenBRAINLAYER_AUTO_ENRICHis enabled, which adds latency and external API usage to the background path.Macroscope summarized e6ea5de.
Summary by CodeRabbit
New Features
Tests