Phase 11: add parser registration and DI provider bridge - #431
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds parser registration and Ollama registration via import side-effects, a thread-safe cached component factory, ServiceContainer lifecycle and tracked shutdown, a process-level DI container bridge with initialization gating (HTTP 503), broad runtime DI migration to Depends(get_config), worker/bootstrap refactors, and tests for these behaviors. ChangesDI Provider Bridge with Ollama and Parser Support
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
🚥 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 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 review |
✅ Actions performedReview triggered.
|
Add make_component_factory: a generic, thread-safe (double-checked locking) cached factory that builds inference components by name from a registry and config section. Returns a (factory, cache) tuple so the container can close cached httpx clients on shutdown and invalidate entries on endpoint changes. ModelEndpointConfig documents the unified per-endpoint config shape the factory consumes.
Expose `config` and `is_initialized` on ServiceContainer so the provider bridge's get_config and initialization guard resolve against the real container, not just the test fakes. Track inference clients built via create_embedder/llm/reranker/vlm and aclose() them in shutdown() so httpx connection pools are released on teardown.
The lifespan only set request.app.state.container, so callers that resolve the container without a request (Chainlit, background tasks) fell back to an empty ServiceContainer with no settings. Mirror the resolved boot container into di.providers via set_container after startup and clear it on shutdown, matching the degraded-boot 503 path.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@openrag/di/container.py`:
- Around line 130-139: The shutdown method currently aborts cleanup if any
client's aclose() raises; update Container.shutdown to make close operations
best-effort by iterating over self._inference_clients and awaiting each aclose()
inside a try/except that logs the error (but continues), then clear
self._inference_clients regardless; likewise wrap await
self._catalog_store.shutdown() in a try/except if _catalog_store is not None so
store shutdown failures don't stop final state reset, and ensure
self._initialized = False always executes (e.g., via finally or by placing it
after the guarded operations).
- Around line 147-149: The config() method is annotated to return Settings but
can return None via self._settings; fix by either (A) making the return type
Optional[Settings] and updating callers to handle None (change def config(self)
-> Optional[Settings]) or (B) enforce settings by asserting/raising when absent
(in container.config check if self._settings is None and raise a
RuntimeError/ValueError with a clear message), so update the config method and
references to use the chosen approach and adjust any callers accordingly; refer
to the config method and the _settings attribute on the container class 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 1db2d8e7-27db-436e-aed6-449a1d67722f
📒 Files selected for processing (5)
openrag/api/main.pyopenrag/di/container.pyopenrag/di/factories.pyopenrag/di/test_container.pyopenrag/di/test_factories.py
Address review feedback on the phase 11 container: close each inference client inside its own guard so one aclose() failure no longer skips the remaining clients, the store shutdown, or the state reset (now in a finally block). Make the config property raise the no-settings error instead of returning None so its Settings return type holds.
|
Actionable comments posted: 0 |
Replace module-level load_config() calls with explicit settings passing and FastAPI Depends(get_config) so all components source config from the DI container rather than import-time globals. Bootstrap and parser workers now receive settings via initialize_worker_bootstrap(settings) instead of capturing a module-level singleton.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
openrag/api/routers/user/chat.py (1)
176-181:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPropagate injected settings into max-token ceiling calculation.
check_tokens_limit(..., settings=...)currently ignoressettingsformax_tokens_allowedbecauseget_max_model_tokens()is called without it. This can apply limits from a different config than the request’s injected config.💡 Proposed fix
-def get_max_model_tokens() -> int: +def get_max_model_tokens(settings: "Settings | None" = None) -> int: """Return the cached max model token limit (populated at startup).""" if _max_model_tokens is not None: return _max_model_tokens - config = _runtime_config() + config = _runtime_config(settings) return int(config.llm_context.max_llm_context_size) def check_tokens_limit( request: OpenAIChatCompletionRequest | OpenAICompletionRequest, log, settings: "Settings | None" = None, ): """Validate token limit and raise HTTPException(413) if exceeded.""" is_valid, error_message = validate_tokens_limit( request, - max_tokens_allowed=get_max_model_tokens(), + max_tokens_allowed=get_max_model_tokens(settings), settings=settings, )Also applies to: 228-238, 294-296, 389-390
🤖 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 `@openrag/api/routers/user/chat.py` around lines 176 - 181, get_max_model_tokens() currently always reads the global runtime config, causing check_tokens_limit(..., settings=...) to ignore injected per-request settings; update get_max_model_tokens to accept an optional settings parameter (e.g., get_max_model_tokens(settings: Optional[RuntimeSettings] = None)) and have it prefer values from the passed settings.llm_context.max_llm_context_size when provided, then update all callers (notably where check_tokens_limit is invoked, and the sites around the symbols mentioned) to pass the request's injected settings through so the max-token ceiling is computed from the correct config source.openrag/services/workers/parsers/whisper_workers.py (1)
199-204:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMissing
call_ray_actor_with_timeoutwrapper for pool dispatch.
LocalWhisperLoader.parsecallsself.whisper_actor.transcribe.remote(...)directly without a timeout wrapper. WhileWhisperPool.transcribehas internal retry/backoff, the outer Ray call from the loader has no timeout protection. This is inconsistent withMarkerLoader._convert_pdfandDoclingLoader._dispatch, which both wrap their pool calls withcall_ray_actor_with_timeout.🛠️ Suggested fix
async def parse(self, document: Document) -> ProcessedDocument: if not document.raw_bytes: return ProcessedDocument( document_id=document.id, metadata=dict(document.metadata), ) async with document.as_temporary_file() as path: try: - text = await self.whisper_actor.transcribe.remote(str(path)) + text = await call_ray_actor_with_timeout( + self.whisper_actor.transcribe.remote(str(path)), + timeout=self.config.loader.local_whisper.whisper_timeout, + task_description=f"LocalWhisperLoader transcribe ({path})", + ) except Exception as e: logger.error("Error transcribing audio", error=str(e)) raise🤖 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 `@openrag/services/workers/parsers/whisper_workers.py` around lines 199 - 204, LocalWhisperLoader.parse is calling self.whisper_actor.transcribe.remote(...) directly without the call_ray_actor_with_timeout wrapper, leaving the outer Ray call without a timeout; update LocalWhisperLoader.parse to dispatch to the whisper actor using call_ray_actor_with_timeout(self.whisper_actor, "transcribe", args=(str(path),), ...) (matching the pattern used in MarkerLoader._convert_pdf and DoclingLoader._dispatch) so the outer call has the same timeout/retry protection as WhisperPool.transcribe; ensure you preserve the existing exception handling and pass the same timeout value used elsewhere in the loaders.
🤖 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 `@openrag/components/prompts/prompts.py`:
- Around line 25-27: The current assignment uses falsy checks (prompts_dir =
prompts_dir or config.paths.prompts_dir and prompt_mapping = prompt_mapping or
config.prompts) which overrides explicit empty values like {} or ""; change
these to explicit None checks so only None triggers defaults: after calling
load_config(), set prompts_dir = config.paths.prompts_dir only if prompts_dir is
None, and set prompt_mapping = config.prompts only if prompt_mapping is None
(referencing the variables prompts_dir, prompt_mapping and
config.paths.prompts_dir / config.prompts to locate the assignments).
In `@openrag/di/test_workers.py`:
- Around line 14-23: The test replaces sys.modules["services.workers.bootstrap"]
but the finally block only pops it, which loses any pre-existing module and can
make tests order-dependent; save the original entry before inserting (e.g., prev
= sys.modules.get("services.workers.bootstrap")), set the test ModuleType and
run ensure_worker_bootstrap(settings), then in the finally restore the original:
if prev is None remove the test entry, otherwise reassign
sys.modules["services.workers.bootstrap"] = prev so the original module is
preserved; apply this change around the ModuleType/Mock setup used with
ensure_worker_bootstrap.
---
Outside diff comments:
In `@openrag/api/routers/user/chat.py`:
- Around line 176-181: get_max_model_tokens() currently always reads the global
runtime config, causing check_tokens_limit(..., settings=...) to ignore injected
per-request settings; update get_max_model_tokens to accept an optional settings
parameter (e.g., get_max_model_tokens(settings: Optional[RuntimeSettings] =
None)) and have it prefer values from the passed
settings.llm_context.max_llm_context_size when provided, then update all callers
(notably where check_tokens_limit is invoked, and the sites around the symbols
mentioned) to pass the request's injected settings through so the max-token
ceiling is computed from the correct config source.
In `@openrag/services/workers/parsers/whisper_workers.py`:
- Around line 199-204: LocalWhisperLoader.parse is calling
self.whisper_actor.transcribe.remote(...) directly without the
call_ray_actor_with_timeout wrapper, leaving the outer Ray call without a
timeout; update LocalWhisperLoader.parse to dispatch to the whisper actor using
call_ray_actor_with_timeout(self.whisper_actor, "transcribe", args=(str(path),),
...) (matching the pattern used in MarkerLoader._convert_pdf and
DoclingLoader._dispatch) so the outer call has the same timeout/retry protection
as WhisperPool.transcribe; ensure you preserve the existing exception handling
and pass the same timeout value used elsewhere in the loaders.
🪄 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: 220e4d93-8737-4493-8a87-bb46ebf0393d
📒 Files selected for processing (21)
openrag/api/dependencies/auth.pyopenrag/api/dependencies/files.pyopenrag/api/dependencies/llm.pyopenrag/api/dependencies/test_auth.pyopenrag/api/main.pyopenrag/api/routers/admin/indexing.pyopenrag/api/routers/admin/tools.pyopenrag/api/routers/user/chat.pyopenrag/api/schemas/user/chat.pyopenrag/components/indexer/chunker/chunker.pyopenrag/components/indexer/loaders/base.pyopenrag/components/prompts/prompts.pyopenrag/components/utils.pyopenrag/di/test_workers.pyopenrag/di/workers.pyopenrag/services/workers/bootstrap.pyopenrag/services/workers/parsers/doc_serializer.pyopenrag/services/workers/parsers/docling_workers.pyopenrag/services/workers/parsers/marker_workers.pyopenrag/services/workers/parsers/whisper_workers.pyopenrag/utils/logger.py
| config = load_config() | ||
| prompts_dir = prompts_dir or config.paths.prompts_dir | ||
| prompt_mapping = prompt_mapping or config.prompts |
There was a problem hiding this comment.
Preserve explicit empty prompt_mapping overrides.
Using or here treats {} as “not provided” and silently replaces it with config defaults. That breaks intentional caller overrides.
Suggested fix
def load_prompt(
prompt_name: str,
prompts_dir: Path | None = None,
prompt_mapping=None,
) -> str:
config = load_config()
- prompts_dir = prompts_dir or config.paths.prompts_dir
- prompt_mapping = prompt_mapping or config.prompts
+ if prompts_dir is None:
+ prompts_dir = config.paths.prompts_dir
+ if prompt_mapping is None:
+ prompt_mapping = config.prompts
return load_template_by_key(prompts_dir, prompt_mapping, prompt_name)🤖 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 `@openrag/components/prompts/prompts.py` around lines 25 - 27, The current
assignment uses falsy checks (prompts_dir = prompts_dir or
config.paths.prompts_dir and prompt_mapping = prompt_mapping or config.prompts)
which overrides explicit empty values like {} or ""; change these to explicit
None checks so only None triggers defaults: after calling load_config(), set
prompts_dir = config.paths.prompts_dir only if prompts_dir is None, and set
prompt_mapping = config.prompts only if prompt_mapping is None (referencing the
variables prompts_dir, prompt_mapping and config.paths.prompts_dir /
config.prompts to locate the assignments).
| module = ModuleType("services.workers.bootstrap") | ||
| module.initialize_worker_bootstrap = Mock() | ||
| sys.modules["services.workers.bootstrap"] = module | ||
| settings = Settings() | ||
|
|
||
| try: | ||
| ensure_worker_bootstrap(settings) | ||
| finally: | ||
| sys.modules.pop("services.workers.bootstrap", None) | ||
|
|
There was a problem hiding this comment.
Restore original sys.modules entry to keep test isolation.
finally: pop(...) removes the module but doesn’t restore a pre-existing entry. That can make later tests order-dependent.
Suggested fix
def test_ensure_worker_bootstrap_initializes_explicitly() -> None:
"""Startup calls the worker bootstrap function instead of relying on import side effects."""
+ original = sys.modules.get("services.workers.bootstrap")
module = ModuleType("services.workers.bootstrap")
module.initialize_worker_bootstrap = Mock()
sys.modules["services.workers.bootstrap"] = module
settings = Settings()
try:
ensure_worker_bootstrap(settings)
finally:
- sys.modules.pop("services.workers.bootstrap", None)
+ if original is None:
+ sys.modules.pop("services.workers.bootstrap", None)
+ else:
+ sys.modules["services.workers.bootstrap"] = original🤖 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 `@openrag/di/test_workers.py` around lines 14 - 23, The test replaces
sys.modules["services.workers.bootstrap"] but the finally block only pops it,
which loses any pre-existing module and can make tests order-dependent; save the
original entry before inserting (e.g., prev =
sys.modules.get("services.workers.bootstrap")), set the test ModuleType and run
ensure_worker_bootstrap(settings), then in the finally restore the original: if
prev is None remove the test entry, otherwise reassign
sys.modules["services.workers.bootstrap"] = prev so the original module is
preserved; apply this change around the ModuleType/Mock setup used with
ensure_worker_bootstrap.
Context: Phase 11 moves OpenRAG toward a single composition root. This PR covers the low-conflict Person B side first so Person A can continue container work with stable registration and provider boundaries.\n\nProblem: parser implementations were not registered through a DI module, Ollama registration was implicit, and the provider bridge still only worked through request state. That made the next container lifecycle step harder to test and wire cleanly.\n\nExpected behavior: DI registration is explicit, provider access works both from FastAPI request state and controlled test/lifespan overrides, and current Phase 10 request behavior remains unchanged.
Summary by CodeRabbit
New Features
Improvements
Tests