fix(indexer): wire contextualization into hexagonal pool - #524
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
💤 Files with no reviewable changes (1)
📝 WalkthroughWalkthrough
ChangesContextualization wiring, concurrency refactor, and preset-driven configuration
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
cc4f62a to
a0b3df3
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unit/core/indexing/test_contextualize.py (1)
13-20: ⚡ Quick winAdd branch coverage for all normalized chat payload shapes.
The new test validates only
choices[0].message.content. Adding parametrized cases forchoices[0].text, top-levelcontent, and unsupported payload fallback would lock in the new helper behavior and prevent silent regressions.Proposed test extension
`@pytest.mark.asyncio` async def test_contextualizer_accepts_openai_style_chat_response(): @@ assert result[0].context == "document-level context" assert result[0].content == "chunk body" + + +@pytest.mark.parametrize( + ("response", "expected_context"), + [ + ({"choices": [{"message": {"content": "document-level context"}}]}, "document-level context"), + ({"choices": [{"text": "legacy completion context"}]}, "legacy completion context"), + ({"content": "top-level context"}, "top-level context"), + ({"unexpected": True}, ""), + ], +) +@pytest.mark.asyncio +async def test_contextualizer_normalizes_chat_response_shapes(response, expected_context): + class ShapeLLM: + async def chat(self, messages, **kwargs): + return response + + contextualizer = ChunkContextualizer(ShapeLLM(), "System prompt") + result = await contextualizer.contextualize([Chunk(id="c1", text="chunk body", partition="p")]) + assert result[0].context == expected_context🤖 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/core/indexing/test_contextualize.py` around lines 13 - 20, The test_contextualizer_accepts_openai_style_chat_response function currently only covers one chat response payload format (choices[0].message.content). Parametrize this test using pytest.mark.parametrize to add test cases for alternative normalized payload shapes including choices[0].text, top-level content, and an unsupported payload fallback case. Each parametrized case should test the ChunkContextualizer with DictLLM configured to return the respective payload shape and verify that contextualize correctly extracts the context from each supported format while handling the unsupported format appropriately.
🤖 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.
Nitpick comments:
In `@tests/unit/core/indexing/test_contextualize.py`:
- Around line 13-20: The test_contextualizer_accepts_openai_style_chat_response
function currently only covers one chat response payload format
(choices[0].message.content). Parametrize this test using
pytest.mark.parametrize to add test cases for alternative normalized payload
shapes including choices[0].text, top-level content, and an unsupported payload
fallback case. Each parametrized case should test the ChunkContextualizer with
DictLLM configured to return the respective payload shape and verify that
contextualize correctly extracts the context from each supported format while
handling the unsupported format appropriately.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9a9e0eef-fc3e-460b-8cf6-c44f5c4ec0fb
📒 Files selected for processing (4)
openrag/core/indexing/contextualize.pyopenrag/services/workers/indexer_pool.pytests/unit/core/indexing/test_contextualize.pytests/unit/services/workers/test_indexer_pool.py
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unit/services/workers/test_indexer_pool.py (1)
141-185: ⚡ Quick winConsider adding test coverage for unknown LLM name.
The
_build_contextualizer_factoryshould raise aKeyErrorwhen an unknown LLM name is requested (and no fallback is available), but this error path is not currently tested.📝 Suggested test case
def test_build_contextualizer_factory_raises_on_unknown_name(tmp_path) -> None: from core.config.model_endpoints import ModelEndpointConfig from services.workers.indexer_pool import _build_contextualizer_factory (tmp_path / "chunk_contextualizer_tmpl.txt").write_text("Context prompt", encoding="utf-8") cfg = SimpleNamespace( models=SimpleNamespace( llm={ "known": ModelEndpointConfig( endpoint="http://llm.example/v1", model_name="model", timeout=30, extra={"implementation": "vllm", "api_key": "key"}, ) } ), llm=SimpleNamespace(base_url="", model="", api_key=""), # No fallback chunker=SimpleNamespace(contextualization_timeout=12, max_concurrent_contextualization=3), paths=SimpleNamespace(prompts_dir=str(tmp_path)), prompts=SimpleNamespace(chunk_contextualizer="chunk_contextualizer_tmpl.txt"), ) factory = _build_contextualizer_factory(cfg) with pytest.raises(KeyError, match="Unknown llm 'unknown'"): factory("unknown")🤖 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/services/workers/test_indexer_pool.py` around lines 141 - 185, Add a new test function to cover the error case for _build_contextualizer_factory when an unknown LLM name is requested. Create a test function (e.g., test_build_contextualizer_factory_raises_on_unknown_name) that sets up a configuration with a known LLM endpoint and no fallback settings, then verifies that calling the factory with an unknown LLM name raises a KeyError with an appropriate error message. Use pytest.raises to assert the exception is raised with the expected error message pattern.
🤖 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 `@tests/unit/services/workers/test_indexer_pool.py`:
- Line 156: The llm_registry registration of FakeLLM with key
"test-contextualizer-llm" persists after the test completes, causing potential
test pollution in the shared process. Add cleanup logic to remove the
"test-contextualizer-llm" registration from the llm_registry after the test
finishes, either by using a pytest fixture with an appropriate teardown/cleanup
method, or by wrapping the registration in a try/finally block that unregisters
the entry. Ensure the cleanup mechanism properly removes the registration to
prevent interference with other tests.
---
Nitpick comments:
In `@tests/unit/services/workers/test_indexer_pool.py`:
- Around line 141-185: Add a new test function to cover the error case for
_build_contextualizer_factory when an unknown LLM name is requested. Create a
test function (e.g., test_build_contextualizer_factory_raises_on_unknown_name)
that sets up a configuration with a known LLM endpoint and no fallback settings,
then verifies that calling the factory with an unknown LLM name raises a
KeyError with an appropriate error message. Use pytest.raises to assert the
exception is raised with the expected error message pattern.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 40cc105d-dba2-4d7e-a9f4-6c025e2ed682
📒 Files selected for processing (4)
openrag/core/indexing/contextualize.pyopenrag/services/workers/indexer_pool.pytests/unit/core/indexing/test_contextualize.pytests/unit/services/workers/test_indexer_pool.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/unit/core/indexing/test_contextualize.py
- openrag/core/indexing/contextualize.py
- openrag/services/workers/indexer_pool.py
… LLM semaphore The hexagonal contextualizer called the LLM without sharing the cluster-wide "llmSemaphore", so a large indexing job could flood the vLLM endpoint and run uncapped against query-time calls. Inject the distributed semaphore into ChunkContextualizer through an optional llm_semaphore gate that wraps _llm.chat, built in the services-layer factory so core stays free of any Ray/services import. Also harden the contextualizer factory test: unregister the FakeLLM from the shared llm_registry in a finally block and drop the unused `instances` attribute.
In the refactored pipeline, contextualization is gated only by the per-partition preset's enable_contextualization (default false); the global CONTEXTUAL_RETRIEVAL (chunker.contextual_retrieval) flag was orphaned and never took effect. Wire it back via two complementary paths: - seed the default indexation preset's enable_contextualization from the flag, mirroring the existing reranker.enabled kill-switch in _finalize_seed; and - re-sync the default preset from the flag on every boot (sync_env_toggles), so an existing deployment honours a changed env value after a restart. Named presets (legal/finance) keep their explicit values.
|
Pushed two follow-up commits on top of the contextualization wiring:
Verified end-to-end on a live stack: forcing the preset to |
hedhoud
left a comment
hedhoud
left a comment
There was a problem hiding this comment.
Left one concern about the default preset ownership model.
Ahmath-Gadji
left a comment
Ahmath-Gadji
left a comment
There was a problem hiding this comment.
LGTM. Tested end to end
2 participants
Context
Contextual retrieval was implemented in the refactored pipeline, but the production Ray indexer never provided the contextualizer dependency. That made contextualization silently no-op on refactor/hexagonal even when a partition preset enabled it.
Problem
This restores the missing composition link so the indexer can build contextualizers from named LLM endpoints, with a fallback to the legacy/global LLM config. It also handles OpenAI-compatible chat response payloads correctly before prepending context to chunks.
Validation
Closes #522
Summary by CodeRabbit
Release Notes
New Features
Configuration