feat: implement cached prompt entity detection#213
Conversation
📝 WalkthroughWalkthroughThe pull request refactors entity detection in Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
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 |
![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.
|
@codex review |
|
@coderabbitai review |
|
You need to increase your spend limit or enable usage-based billing to run background agents. Go to Cursor |
✅ Actions performedReview triggered.
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b8a0987f73
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if matched: | ||
| return matched |
There was a problem hiding this comment.
Returning immediately when any exact-name/alias match is found skips the phonetic branch entirely, so mixed prompts (for example, one directly named entity plus a Hebrew token that only resolves via alias_type='phonetic') now lose the phonetic entity. Before this change, phonetic fallback was evaluated per candidate and could return both entities; this regression reduces entity-context coverage and can misclassify bilingual prompts.
Useful? React with 👍 / 👎.
![macroscopeapp[bot]](https://avatars.githubusercontent.com/in/900172?s=60&v=4){kind=small}
| except sqlite3.Error: | ||
| pass | ||
| _ENTITY_CACHE = { | ||
| "entities_by_name": {}, | ||
| "aliases_by_name": {}, | ||
| "max_name_tokens": 1, | ||
| "max_alias_tokens": 1, | ||
| } | ||
| finally: |
There was a problem hiding this comment.
🟡 Medium hooks/brainlayer-prompt-search.py:453
When _load_entity_cache() fails to load from database "B", it sets _ENTITY_CACHE = {} but leaves _ENTITY_CACHE_DB_PATH unchanged. If database "A" was previously cached, the stale _ENTITY_CACHE_DB_PATH still holds "A", causing a subsequent call for "A" to pass the cache check at lines 393-394 and return the empty dict instead of the original cached data. Consider clearing _ENTITY_CACHE_DB_PATH in the error handler to invalidate the cache.
except sqlite3.Error:
_ENTITY_CACHE = {
"entities_by_name": {},
"aliases_by_name": {},
"max_name_tokens": 1,
"max_alias_tokens": 1,
}
+ _ENTITY_CACHE_DB_PATH = None🚀 Reply "fix it for me" or copy this AI Prompt for your agent:
In file hooks/brainlayer-prompt-search.py around lines 453-460:
When `_load_entity_cache()` fails to load from database "B", it sets `_ENTITY_CACHE = {}` but leaves `_ENTITY_CACHE_DB_PATH` unchanged. If database "A" was previously cached, the stale `_ENTITY_CACHE_DB_PATH` still holds "A", causing a subsequent call for "A" to pass the cache check at lines 393-394 and return the empty dict instead of the original cached data. Consider clearing `_ENTITY_CACHE_DB_PATH` in the error handler to invalidate the cache.
Evidence trail:
hooks/brainlayer-prompt-search.py lines 387-465 (REVIEWED_COMMIT): The `_load_entity_cache` function at line 387 shows: (1) Cache check at lines 391-392 uses both `_ENTITY_CACHE is not None` and `_ENTITY_CACHE_DB_PATH == cache_key`. (2) Success path at lines 447-448 sets both `_ENTITY_CACHE` and `_ENTITY_CACHE_DB_PATH = cache_key`. (3) Error handler at lines 449-456 only sets `_ENTITY_CACHE` to empty structure, does NOT update `_ENTITY_CACHE_DB_PATH`.
![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 `@hooks/brainlayer-prompt-search.py`:
- Around line 530-531: Update the function's docstring to explicitly state that
phonetic fallback matching is skipped when the conn parameter is None and that
callers must pass a valid connection to enable phonetic matching; reference the
early-return check "if conn is None: return matched" and describe that the
docstring currently only documents cache/span matching and must also explain the
opt-out behavior for phonetic matching when conn is not provided.
- Around line 40-41: The module-level globals _ENTITY_CACHE and
_ENTITY_CACHE_DB_PATH are mutated without synchronization; either document the
single-threaded assumption or protect access with a threading.Lock to avoid race
conditions. Update the module to declare a lock (e.g., _ENTITY_CACHE_LOCK) and
wrap all reads/writes of _ENTITY_CACHE and _ENTITY_CACHE_DB_PATH inside that
lock within the _load_entity_cache function and any other helpers that touch
these globals; alternatively, add a clear top-of-file comment stating the module
is not thread-safe and must only be used in single-threaded contexts. Ensure you
reference and use the existing _load_entity_cache function and the
_ENTITY_CACHE/_ENTITY_CACHE_DB_PATH symbols when making the change.
In `@tests/test_hebrew_alias.py`:
- Around line 151-163: The test
test_hook_runs_phonetic_fallback_even_with_exact_match relies on a fresh entity
cache but doesn't explicitly reset module cache state; before creating the
VectorStore or calling prompt_search.detect_entities_in_prompt, clear
prompt_search._ENTITY_CACHE and prompt_search._ENTITY_CACHE_DB_PATH by setting
them to None so the test does not depend on tmp_path uniqueness and matches the
reset behavior used in test_conditional_hooks.py.
🪄 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: 1b2931d9-47fd-47bd-957b-99c14be1a04b
📒 Files selected for processing (3)
hooks/brainlayer-prompt-search.pytests/test_conditional_hooks.pytests/test_hebrew_alias.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.11)
- GitHub Check: test (3.13)
- 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
Files:
tests/test_hebrew_alias.pytests/test_conditional_hooks.pyhooks/brainlayer-prompt-search.py
tests/**/*.py
📄 CodeRabbit inference engine (CLAUDE.md)
Use
pytestfor testing
Files:
tests/test_hebrew_alias.pytests/test_conditional_hooks.py
🧠 Learnings (3)
📓 Common learnings
Learnt from: EtanHey
Repo: EtanHey/brainlayer PR: 198
File: hooks/brainlayer-prompt-search.py:241-259
Timestamp: 2026-04-04T15:21:59.853Z
Learning: In `hooks/brainlayer-prompt-search.py` (Python), `record_injection_event()` is explicitly best-effort telemetry: silent `except sqlite3.Error: pass` is intentional — table non-existence or lock failures are acceptable silent failures. `sqlite3.connect(timeout=2)` is the file-open timeout; `PRAGMA busy_timeout` governs per-statement lock-wait. The `DEADLINE_MS` (450ms) guard applies only to the FTS search phase, not to this side-channel write.
Learnt from: EtanHey
Repo: EtanHey/brainlayer PR: 198
File: hooks/brainlayer-prompt-search.py:169-169
Timestamp: 2026-04-04T15:21:36.497Z
Learning: In EtanHey/brainlayer, `hooks/brainlayer-prompt-search.py` reads `entity_type` directly from existing rows in `kg_entities` (read-only). `contracts/entity-types.yaml` defines the write-side schema only and is not authoritative for what `entity_type` values exist in the DB. The DB already stores `technology` (72 entities), `project` (24), and `tool` (1) as valid `entity_type` values, so `INJECT_TYPES` in the hook should match these DB values, not the contract file.
📚 Learning: 2026-04-04T15:21:36.497Z
Learnt from: EtanHey
Repo: EtanHey/brainlayer PR: 198
File: hooks/brainlayer-prompt-search.py:169-169
Timestamp: 2026-04-04T15:21:36.497Z
Learning: In EtanHey/brainlayer, `hooks/brainlayer-prompt-search.py` reads `entity_type` directly from existing rows in `kg_entities` (read-only). `contracts/entity-types.yaml` defines the write-side schema only and is not authoritative for what `entity_type` values exist in the DB. The DB already stores `technology` (72 entities), `project` (24), and `tool` (1) as valid `entity_type` values, so `INJECT_TYPES` in the hook should match these DB values, not the contract file.
Applied to files:
tests/test_hebrew_alias.pytests/test_conditional_hooks.pyhooks/brainlayer-prompt-search.py
📚 Learning: 2026-04-04T15:21:59.853Z
Learnt from: EtanHey
Repo: EtanHey/brainlayer PR: 198
File: hooks/brainlayer-prompt-search.py:241-259
Timestamp: 2026-04-04T15:21:59.853Z
Learning: In `hooks/brainlayer-prompt-search.py` (Python), `record_injection_event()` is explicitly best-effort telemetry: silent `except sqlite3.Error: pass` is intentional — table non-existence or lock failures are acceptable silent failures. `sqlite3.connect(timeout=2)` is the file-open timeout; `PRAGMA busy_timeout` governs per-statement lock-wait. The `DEADLINE_MS` (450ms) guard applies only to the FTS search phase, not to this side-channel write.
Applied to files:
tests/test_hebrew_alias.pytests/test_conditional_hooks.pyhooks/brainlayer-prompt-search.py
🔇 Additional comments (10)
hooks/brainlayer-prompt-search.py (4)
372-385: LGTM!The cache key resolution logic correctly handles all three cases: no connection (uses configured path), file-based connection (extracts path from PRAGMA), and in-memory connection (uses object id as fallback). Error handling is appropriate.
388-464: LGTM — error recovery design is sound.The cache loading correctly:
- Opens its own connection when
conn=Noneand properly closes it infinally- On
sqlite3.Error, sets an empty cache but leaves_ENTITY_CACHE_DB_PATHasNone, allowing subsequent calls to retry loading (confirmed bytest_load_entity_cache_retries_after_transient_sqlite_error)- Filters entities by supported
inject_types
467-468: LGTM!Clean implementation that tokenizes the prompt in a single regex pass, preserving both original and lowercase forms for case-insensitive matching.
471-496: LGTM!The longest-first span matching algorithm correctly prioritizes multi-word entity names (e.g., "BrainLayer MCP" over "BrainLayer"). The deduplication via
seen_idsprevents the same entity from being returned multiple times even if matched via different spans.tests/test_conditional_hooks.py (5)
205-211: LGTM!Clean test for the no-match case, verifying that an empty result is returned when no entities match the prompt.
213-247: LGTM!Comprehensive test that validates:
- Cache persistence (second call returns same results after DB is emptied)
- Case-insensitive matching ("EtAn HeYmAn" matches "Etan Heyman")
- Longest-first matching ("BrainLayer MCP" preferred over "BrainLayer")
The explicit cache state reset before testing ensures isolation.
249-266: LGTM!Good test for the
conn=Nonecode path, verifying that the cache loading can open the database viapaths.get_db_path()when no explicit connection is provided.
268-291: LGTM!Effective test for transient error recovery. Verifies that:
- On
sqlite3.OperationalError, the cache is set to empty but_ENTITY_CACHE_DB_PATHremainsNone- Subsequent calls can successfully rebuild the cache
The
ErrorConnstub correctly doesn't need aclose()method sinceclose_connis onlyTruewhen_load_entity_cacheopens its own connection.
293-329: LGTM!Thorough test verifying that unsupported entity types (like
"library") are filtered out during cache loading, ensuring only supported inject types (project,person, etc.) are cached for both entities and aliases.tests/test_hebrew_alias.py (1)
11-11: LGTM!Import is needed for the new test at line 156.
| _ENTITY_CACHE = None | ||
| _ENTITY_CACHE_DB_PATH = None |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Module-level cache lacks synchronization primitives.
The global _ENTITY_CACHE and _ENTITY_CACHE_DB_PATH are accessed and mutated without locking. This is acceptable for the current single-threaded CLI hook execution model, but if this module is ever imported into a multi-threaded context (e.g., a web server or worker pool), concurrent calls to _load_entity_cache could cause race conditions or cache corruption.
Consider adding a comment documenting the single-threaded assumption, or using threading.Lock if broader reuse is anticipated.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@hooks/brainlayer-prompt-search.py` around lines 40 - 41, The module-level
globals _ENTITY_CACHE and _ENTITY_CACHE_DB_PATH are mutated without
synchronization; either document the single-threaded assumption or protect
access with a threading.Lock to avoid race conditions. Update the module to
declare a lock (e.g., _ENTITY_CACHE_LOCK) and wrap all reads/writes of
_ENTITY_CACHE and _ENTITY_CACHE_DB_PATH inside that lock within the
_load_entity_cache function and any other helpers that touch these globals;
alternatively, add a clear top-of-file comment stating the module is not
thread-safe and must only be used in single-threaded contexts. Ensure you
reference and use the existing _load_entity_cache function and the
_ENTITY_CACHE/_ENTITY_CACHE_DB_PATH symbols when making the change.
| if conn is None: | ||
| return matched |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Clarify the phonetic fallback behavior in docstring.
The early return when conn is None skips the phonetic fallback. This is intentional per the PR (allows callers to opt out of phonetic matching), but the docstring only describes the cache-based span matching. Consider documenting that phonetic fallback requires an explicit connection.
📝 Suggested docstring update
def detect_entities_in_prompt(prompt, conn=None):
"""Detect known KG entity names in the prompt.
Loads KG entity names into a module-level cache and performs
case-insensitive hash-table lookups over contiguous prompt token spans.
Multi-word names are matched longest-first.
+
+ When conn is provided, additionally runs phonetic fallback matching
+ for Hebrew tokens against phonetic aliases in the database.
"""🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@hooks/brainlayer-prompt-search.py` around lines 530 - 531, Update the
function's docstring to explicitly state that phonetic fallback matching is
skipped when the conn parameter is None and that callers must pass a valid
connection to enable phonetic matching; reference the early-return check "if
conn is None: return matched" and describe that the docstring currently only
documents cache/span matching and must also explain the opt-out behavior for
phonetic matching when conn is not provided.
| def test_hook_runs_phonetic_fallback_even_with_exact_match(prompt_search, tmp_path): | ||
| db_path = tmp_path / "hook-mixed.db" | ||
| store = VectorStore(db_path) | ||
| person_id = _upsert_person(store) | ||
| project_id = store.upsert_entity("project-brainlayer", "project", "BrainLayer", metadata={}) | ||
| store.add_entity_alias(phonetic_key("Etan"), person_id, alias_type="phonetic") | ||
| store.close() | ||
|
|
||
| conn = sqlite3.connect(db_path) | ||
| matches = prompt_search.detect_entities_in_prompt("Tell me about BrainLayer and what does איתן think?", conn) | ||
| conn.close() | ||
|
|
||
| assert {match["id"] for match in matches} == {person_id, project_id} |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider resetting module cache state for test consistency.
Other tests in test_conditional_hooks.py explicitly reset prompt_search._ENTITY_CACHE = None and prompt_search._ENTITY_CACHE_DB_PATH = None before running. While this test works correctly due to tmp_path uniqueness (different cache key per test), adding the explicit reset would improve consistency and make the test behavior more explicit.
📝 Suggested addition for consistency
def test_hook_runs_phonetic_fallback_even_with_exact_match(prompt_search, tmp_path):
db_path = tmp_path / "hook-mixed.db"
store = VectorStore(db_path)
person_id = _upsert_person(store)
project_id = store.upsert_entity("project-brainlayer", "project", "BrainLayer", metadata={})
store.add_entity_alias(phonetic_key("Etan"), person_id, alias_type="phonetic")
store.close()
+ prompt_search._ENTITY_CACHE = None
+ prompt_search._ENTITY_CACHE_DB_PATH = None
+
conn = sqlite3.connect(db_path)
matches = prompt_search.detect_entities_in_prompt("Tell me about BrainLayer and what does איתן think?", conn)
conn.close()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_hebrew_alias.py` around lines 151 - 163, The test
test_hook_runs_phonetic_fallback_even_with_exact_match relies on a fresh entity
cache but doesn't explicitly reset module cache state; before creating the
VectorStore or calling prompt_search.detect_entities_in_prompt, clear
prompt_search._ENTITY_CACHE and prompt_search._ENTITY_CACHE_DB_PATH by setting
them to None so the test does not depend on tmp_path uniqueness and matches the
reset behavior used in test_conditional_hooks.py.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 participant
Summary
Test plan
Notes
Note
Add cached entity detection with longest-span and phonetic fallback to prompt search
detect_entities_in_promptin brainlayer-prompt-search.py to use an in-memory per-DB cache of entities and aliases, built once via_load_entity_cacheand keyed by DB file path._iter_prompt_tokensand_match_entity_spans, so multi-word entity names are preferred over shorter overlapping matches.paths.get_db_path().📊 Macroscope summarized aac9bba. 3 files reviewed, 1 issue evaluated, 0 issues filtered, 1 comment posted
🗂️ Filtered Issues
Summary by CodeRabbit
Bug Fixes
Tests