fix(indexer): hydrate model-endpoint registry so admin-UI-registered LLM endpoints resolve - #560
Conversation
📝 WalkthroughWalkthrough
ChangesLazy Registry Hydration and Pipeline Skip
Sequence Diagram(s)sequenceDiagram
participant process_file
participant _ensure_registry_fresh
participant _registry_reload_decision
participant _reload_registry
participant ModelEndpointService
process_file->>_ensure_registry_fresh: required LLM names
_ensure_registry_fresh->>_registry_reload_decision: loaded_at, miss_at, required_names
_registry_reload_decision-->>_ensure_registry_fresh: initial | ttl | miss | none
alt initial or miss
_ensure_registry_fresh->>_reload_registry: await load
_reload_registry->>ModelEndpointService: load_all()
else ttl
_ensure_registry_fresh->>_reload_registry: create_task
end
sequenceDiagram
participant pipeline_builder
participant contextualizer_factory
participant topic_tagger_factory
pipeline_builder->>contextualizer_factory: llm_name
alt contextualizer resolves
contextualizer_factory-->>pipeline_builder: contextualizer
else KeyError
contextualizer_factory-->>pipeline_builder: KeyError
pipeline_builder->>pipeline_builder: log warning, return None
end
pipeline_builder->>topic_tagger_factory: llm_name
alt topic tagger resolves
topic_tagger_factory-->>pipeline_builder: topic tagger
else KeyError
topic_tagger_factory-->>pipeline_builder: KeyError
pipeline_builder->>pipeline_builder: log warning, return None
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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 (1)
openrag/services/workers/indexer_pool.py (1)
347-380: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winBroaden the cache identity beyond
endpoint#model_name.Both factories rebuild only when the URL or model name changes, but the constructed clients also depend on
implementation,timeout, andmodel_cfg.extra. If an admin rotates an API key or changes implementation/timeout in the DB, this actor keeps reusing the stale cached client until restart.Suggested fix
+def _llm_endpoint_identity(model_cfg: Any) -> tuple[Any, ...]: + impl_kwargs = {key: value for key, value in model_cfg.extra.items() if key != "implementation"} + return ( + model_cfg.extra.get("implementation", "vllm"), + model_cfg.endpoint, + model_cfg.model_name, + model_cfg.timeout, + tuple(sorted(impl_kwargs.items())), + ) + def factory(name: str = "default") -> ChunkContextualizer: model_cfg = named_llms.get(name) if model_cfg is None: if name == "default" and fallback_cfg is not None: model_cfg = fallback_cfg else: raise KeyError(f"Unknown llm '{name}'. Available: {list(named_llms)}") - identity = f"{model_cfg.endpoint}#{model_cfg.model_name}" + identity = _llm_endpoint_identity(model_cfg) entry = cache.get(name) if entry is not None and entry[0] == identity: return entry[1]def factory(name: str = "default") -> TopicTagger: model_cfg = named_llms.get(name) if model_cfg is None: if name == "default" and fallback_cfg is not None: model_cfg = fallback_cfg else: raise KeyError(f"Unknown llm '{name}'. Available: {list(named_llms)}") - identity = f"{model_cfg.endpoint}#{model_cfg.model_name}" + identity = _llm_endpoint_identity(model_cfg) entry = cache.get(name) if entry is not None and entry[0] == identity: return entry[1]Also applies to: 416-436
🤖 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/indexer_pool.py` around lines 347 - 380, The cache key in the worker factory is too narrow because `identity` only uses `model_cfg.endpoint` and `model_cfg.model_name`, so stale clients can be reused when `implementation`, `timeout`, or `model_cfg.extra` changes. Update the cache identity in the `indexer_pool` factory path that builds `llm` and `ChunkContextualizer` to include all configuration inputs that affect client construction, and make the same change in the other referenced factory block so both cache checks invalidate correctly when those settings change.
🤖 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/services/workers/pipeline_builder.py`:
- Around line 171-178: The exception handling in the contextualization factory
path is too broad and is hiding real failures. Update the logic in the
contextualizer creation block used by the relevant pipeline builder methods to
only suppress the unresolved-endpoint case (the expected unknown-LLM lookup,
such as KeyError from the factory), and let other exceptions propagate so
prompt-loading, config, or client-construction bugs are not silently skipped.
Keep the warning/logging only for the handled missing-LLM case, and apply the
same narrower handling in both affected code paths.
---
Outside diff comments:
In `@openrag/services/workers/indexer_pool.py`:
- Around line 347-380: The cache key in the worker factory is too narrow because
`identity` only uses `model_cfg.endpoint` and `model_cfg.model_name`, so stale
clients can be reused when `implementation`, `timeout`, or `model_cfg.extra`
changes. Update the cache identity in the `indexer_pool` factory path that
builds `llm` and `ChunkContextualizer` to include all configuration inputs that
affect client construction, and make the same change in the other referenced
factory block so both cache checks invalidate correctly when those settings
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: 893645af-f981-4d71-92df-9bed2b937b66
📒 Files selected for processing (4)
openrag/services/workers/indexer_pool.pyopenrag/services/workers/pipeline_builder.pytests/unit/services/workers/test_indexer_pool.pytests/unit/services/workers/test_pipeline_builder.py
1c12fdd to
f8e1e39
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tests/unit/services/workers/test_indexer_pool.py (2)
304-401: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMirror cache-invalidation tests for topic tagging.
The production change applies the same identity-based rebuild logic to
_build_topic_tagger_factory(), but these new edit/API-key rotation assertions exercise only contextualization.🤖 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 304 - 401, Add equivalent cache-invalidation coverage for the topic tagging path, since `_build_topic_tagger_factory()` now uses the same identity-based rebuild behavior as `_build_contextualizer_factory()`. Mirror the existing edit and api-key-rotation assertions with topic-tagger-specific tests so a changed `ModelEndpointConfig` causes a fresh client and updated kwargs, using the same identity helper pattern exercised by `_endpoint_identity`.
226-229: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for stale registry plus missing endpoint.
The current cases don’t cover
missing=Truewhenloaded_atis already TTL-expired; that case should return"miss"so the triggering file blocks on one reload instead of skipping the stage.# TTL expiry refreshes (catches edits to existing endpoints). assert _registry_reload_decision(loaded_at=100.0, last_miss_at=None, now=161.0, ttl=60.0, missing=False) == "ttl" + assert _registry_reload_decision(loaded_at=100.0, last_miss_at=None, now=161.0, ttl=60.0, missing=True) == "miss"🤖 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 226 - 229, Add a test case in test_indexer_pool.py for _registry_reload_decision where loaded_at is already past ttl and missing=True, and assert it returns "miss". The current coverage only checks ttl-only and fresh missing cases, so extend the existing _registry_reload_decision assertions to verify the stale-registry-plus-missing-path. This ensures the decision logic still favors a single reload when an endpoint is missing even after registry expiry.
🤖 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/services/workers/indexer_pool.py`:
- Around line 249-254: The refresh decision logic currently checks TTL before
handling a missing endpoint, which can cause a stale registry with a cache miss
to return "ttl" instead of forcing a miss reload. Update the condition ordering
in the refresh-reason logic that uses loaded_at, ttl, missing, and last_miss_at
so that the miss case is evaluated first and returns "miss" whenever a required
endpoint is missing and the miss is eligible to retry, before falling back to
"ttl" or "initial". This should ensure the path used by _ensure_registry_fresh()
blocks once on a miss instead of immediately taking the skip path.
---
Nitpick comments:
In `@tests/unit/services/workers/test_indexer_pool.py`:
- Around line 304-401: Add equivalent cache-invalidation coverage for the topic
tagging path, since `_build_topic_tagger_factory()` now uses the same
identity-based rebuild behavior as `_build_contextualizer_factory()`. Mirror the
existing edit and api-key-rotation assertions with topic-tagger-specific tests
so a changed `ModelEndpointConfig` causes a fresh client and updated kwargs,
using the same identity helper pattern exercised by `_endpoint_identity`.
- Around line 226-229: Add a test case in test_indexer_pool.py for
_registry_reload_decision where loaded_at is already past ttl and missing=True,
and assert it returns "miss". The current coverage only checks ttl-only and
fresh missing cases, so extend the existing _registry_reload_decision assertions
to verify the stale-registry-plus-missing-path. This ensures the decision logic
still favors a single reload when an endpoint is missing even after registry
expiry.
🪄 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: 69bd50b7-7df7-45f2-a44c-d891fb5f1ee6
📒 Files selected for processing (4)
openrag/services/workers/indexer_pool.pyopenrag/services/workers/pipeline_builder.pytests/unit/services/workers/test_indexer_pool.pytests/unit/services/workers/test_pipeline_builder.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unit/services/workers/test_pipeline_builder.py
- openrag/services/workers/pipeline_builder.py
f8e1e39 to
6438187
Compare
|
Addressed CodeRabbit's note: narrowed both enhancement-skip blocks from |
|
Addressed: |
6438187 to
aeba65a
Compare
…ic tagging The indexer Ray actor never hydrated the DB-backed model-endpoint registry (Settings.models) — only the API process did at startup — so contextualization or topic tagging configured with a named LLM endpoint failed with 'KeyError: Unknown llm'. The indexer now hydrates cfg.models from the DB lazily: once on first use, on a 60s TTL refresh (run in the background, off the file's critical path), and on a rate-limited cache miss so an unknown name can't storm the DB; reloads are single-flight. The contextualizer/topic-tagger factories read the live registry and cache one client per endpoint identity (rebuilt on change). An unresolvable enhancement LLM degrades gracefully (warn + skip) instead of failing the file. The actor logger is an instance attribute rather than a module global, so Ray's by-value pickling of the actor class doesn't pull in loguru's enqueue SimpleQueue. Refs #554
aeba65a to
56446fd
Compare
|
Ran an independent multi-lens review pass over the diff. Fixed three items it surfaced:
Added regression tests for the fallback-resolvability case. 28 unit tests pass; actor still cloud-pickles with the enqueue sink active. Noted but intentionally not in scope here: a dedicated |
hedhoud
left a comment
hedhoud
left a comment
There was a problem hiding this comment.
I did an impact pass beyond the happy path. This looks safe for the issue it targets: the indexer can now hydrate DB-backed LLM endpoints, so contextualization and topic tagging can resolve endpoints created from the admin UI. The graceful-degrade behavior also makes sense here: a missing enhancement LLM should not fail the whole file.
I also checked the adjacent paths around API startup, the model-endpoint service, the DI container, dispatcher, and indexer worker. I did not see a regression there, and the focused local tests plus the GitHub checks are green.
One scope note: this fixes the LLM-backed enhancement path. Named embedders and VLMs registered through the admin UI still look like a separate follow-up if we expect per-preset embedder/VLM selection to work fully inside the indexer. I would not block this PR on that, but I would track it separately.
There was a problem hiding this comment.
🧹 Nitpick comments (2)
openrag/services/workers/indexer_pool.py (2)
149-149: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep reload-failure logs structured and avoid stringifying the exception.
str(exc)can include DB/config details from endpoint loading, and the f-string bypasses structured Loguru context. Bind the reload reason and error type instead.Proposed change
- self._logger.warning(f"Model endpoint registry reload failed ({decision}): {exc}") + self._logger.bind( + reload_reason=decision, + error_type=type(exc).__name__, + ).warning("Model endpoint registry reload failed")As per coding guidelines, “Use Loguru for logging with structured logging via
get_logger()”.🤖 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/indexer_pool.py` at line 149, In the reload failure logging inside the indexer pool’s registry reload path, avoid interpolating the exception into the message string and keep the log structured. Update the warning in the reload-failure handling around the registry reload logic to use the logger’s bound context from get_logger()/Loguru, attaching fields like the reload decision and exception type instead of str(exc). Keep the message generic and rely on structured attributes for the reason and error metadata.Source: Coding guidelines
345-364: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winResolve
cfg.models.llminside each factory call.Both factories document call-time lookup, but they capture the initial registry dict. If hydration ever replaces
cfg.models.llminstead of mutating it in place, newly loaded endpoints stay invisible. Moving the lookup insidefactory()makes the contract robust.Proposed change
- models = getattr(cfg, "models", None) - named_llms = models.llm if models is not None else {} fallback_cfg = _global_llm_endpoint_config(cfg) @@ def factory(name: str = "default") -> ChunkContextualizer: + models = getattr(cfg, "models", None) + named_llms = models.llm if models is not None else {} model_cfg = named_llms.get(name)- models = getattr(cfg, "models", None) - named_llms = models.llm if models is not None else {} fallback_cfg = _global_llm_endpoint_config(cfg) @@ def factory(name: str = "default") -> TopicTagger: + models = getattr(cfg, "models", None) + named_llms = models.llm if models is not None else {} model_cfg = named_llms.get(name)Also applies to: 422-436
🤖 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/indexer_pool.py` around lines 345 - 364, The `factory()` closures in `indexer_pool.py` are capturing the initial `cfg.models.llm` mapping instead of resolving it per call, so updated endpoints can be missed if the registry is replaced. Move the `named_llms = models.llm` lookup inside each `factory()` (and keep using `_global_llm_endpoint_config(cfg)` for the default fallback) so `factory()` always reads the current `cfg.models.llm` state before resolving `model_cfg`.
🤖 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 `@openrag/services/workers/indexer_pool.py`:
- Line 149: In the reload failure logging inside the indexer pool’s registry
reload path, avoid interpolating the exception into the message string and keep
the log structured. Update the warning in the reload-failure handling around the
registry reload logic to use the logger’s bound context from
get_logger()/Loguru, attaching fields like the reload decision and exception
type instead of str(exc). Keep the message generic and rely on structured
attributes for the reason and error metadata.
- Around line 345-364: The `factory()` closures in `indexer_pool.py` are
capturing the initial `cfg.models.llm` mapping instead of resolving it per call,
so updated endpoints can be missed if the registry is replaced. Move the
`named_llms = models.llm` lookup inside each `factory()` (and keep using
`_global_llm_endpoint_config(cfg)` for the default fallback) so `factory()`
always reads the current `cfg.models.llm` state before resolving `model_cfg`.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9422e630-47ae-4a8e-9974-8420ad38375c
📒 Files selected for processing (4)
openrag/services/workers/indexer_pool.pyopenrag/services/workers/pipeline_builder.pytests/unit/services/workers/test_indexer_pool.pytests/unit/services/workers/test_pipeline_builder.py
🚧 Files skipped from review as they are similar to previous changes (1)
- openrag/services/workers/pipeline_builder.py
hedhoud
left a comment
hedhoud
left a comment
There was a problem hiding this comment.
Re-reviewed the updated head 56446fda after the extra fallback/reload hardening.
The new default fallback handling looks correct: the reload decision now matches the factory resolution path, so a deployment that only has the global LLM fallback should not keep doing a blocking reload every TTL window. The lazy prompt/semaphore build is also safer now because it avoids publishing half-built shared state.
I reran the focused worker/indexer suites and the adjacent model-endpoint/container/dispatcher tests locally, plus ruff, diff whitespace, and a Ray cloudpickle smoke for the actor class. I do not see a blocker from this update.
Same scope note as before: embedder/VLM endpoint hydration for indexation presets still looks like a separate follow-up, not something I would block this LLM-contextualization fix on.
hedhoud
left a comment
hedhoud
left a comment
There was a problem hiding this comment.
Approving this from my side.
I tested the current head with an isolated stack and exercised the Admin API flow that the UI depends on: create a model endpoint, create a preset with image captioning + contextual retrieval + topic tagging enabled, assign it to a partition, then upload PDFs through the indexer.
The E2E check passed with 10 PDFs submitted concurrently:
- all uploads returned successfully
- all 10 tasks reached COMPLETED
- 0 failed tasks
- indexer-ui showed COMPLETED (10), FAILED (0)
- contextualization was visible in the indexed chunks through the stored context and [CONTEXT] prefix
- topic tags were persisted for all 10 files
GitHub checks are green as well. The remaining comments look like follow-up/nitpick-level hardening, not blockers for this fix.
2 participants
Problem
Contextualization or topic tagging configured with an LLM endpoint registered via the admin UI (i.e. stored in the DB) made indexing fail with
KeyError: Unknown llm '<name>'. The global/default LLM path worked; only DB-registered named endpoints failed.Root cause: the indexer Ray actor builds its enhancement-LLM factories from
Settings.models, but that registry is hydrated from the DB only in the API process at startup — never in the indexer process. So a named endpoint never resolved there, and a restart didn't help.Fix
cfg.modelsfrom the DB inside the indexer actor, lazily: once on first use, on a 60s TTL refresh (run in the background, off the file's critical path), and on a cache miss (rate-limited so an unknown/typo'd name can't storm the DB). Reloads are single-flight.extrakey (api_key,temperature, …) — so any edit, including an api-key rotation that keeps the same URL/model, yields a new identity and rebuilds the client on the next reload (≤60s), without leaking the old one. This matches the API process's invalidate-on-any-change behaviour.enqueue=TrueSimpleQueue.Freshness semantics
Testing
Refs #554
Summary by CodeRabbit
New Features
Bug Fixes