fix(chat): honor the partition chat_llm preset when chatting - #634
Conversation
QueryService was built with a single LLM from the global settings and used
it for every chat/completion answer, so a partition's chat_llm
model-endpoint preset (stored, exposed via PATCH /partition/{p}, resolved
into PartitionConfig) was never consulted at answer time — its only
runtime consumer was the multiQuery/hyde retriever fallback.
QueryService now receives the shared named llm_factory and resolves the
LLM per request, mirroring _resolve_chat_history_depth semantics: no
partition / the 'all' sentinel use the default; a multi-partition request
uses the preset only when every partition that sets one names the same
endpoint. chat_llm is not validated on assignment, so an unknown name
(e.g. endpoint deleted after assignment) falls back to the default LLM
with a warning instead of failing the request.
Debug-level breadcrumbs so operators can trace preset resolution end-to-end without extra volume at the default INFO level: - IndexingPipeline.run() emits one line per file with the resolved captioning VLM, contextualization LLM, topic-tagging LLM and embedder names (None = stage disabled); the _select_* helpers now return the resolved endpoint name alongside the instance - VLLMClient (and VLLMVision via inheritance) plus OllamaClient log a '<Class> ready' construction line with model/endpoint/timeout, mirroring the existing VLLMEmbedder one; the component factories cache instances per endpoint name, so this fires once per configured endpoint and maps a preset name to its base URL and model
|
Warning Review limit reached
Next review available in: 25 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughPartition ChangesPartition chat LLM plumbing
Pipeline logging and client readiness
Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unit/services/workers/test_pipeline_builder.py (1)
560-608: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winOnly the "nothing configured" breadcrumb case is covered.
The new debug-log assertions only exercise the path where VLM/contextualizer/topic-tagger are all unconfigured. Given the KeyError-fallback/skip semantics added in
_select_vlm/_select_contextualizer/_select_topic_tagger, a test covering a resolved named endpoint (and ideally the KeyError-skip warning path) would give better confidence, especially since_select_vlm's named-path currently lacks the same guard as the other two selectors (see companion comment inpipeline_builder.py).🤖 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_pipeline_builder.py` around lines 560 - 608, The current test only covers the “nothing configured” debug breadcrumb, so expand the assertions in test_pipeline_builder to also cover a resolved named endpoint path and, if practical, the KeyError fallback/skip warning behavior for _select_vlm, _select_contextualizer, and _select_topic_tagger. Use the existing _RecordingLogger plus build_indexing_pipeline and pipeline.run to verify the debug breadcrumb reflects a selected endpoint rather than all None values, and ensure the warning path is exercised where the selector falls back instead of failing.
🤖 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 208-218: The _select_vlm selector is missing the same KeyError
protection and warning path used by _select_contextualizer and
_select_topic_tagger. Update _select_vlm so the config.vlm branch is wrapped in
try/except KeyError, logs a logger.warning when the named VLM cannot be
resolved, and falls back to self.vlm instead of letting
self.vlm_factory(config.vlm) fail out of run().
---
Nitpick comments:
In `@tests/unit/services/workers/test_pipeline_builder.py`:
- Around line 560-608: The current test only covers the “nothing configured”
debug breadcrumb, so expand the assertions in test_pipeline_builder to also
cover a resolved named endpoint path and, if practical, the KeyError
fallback/skip warning behavior for _select_vlm, _select_contextualizer, and
_select_topic_tagger. Use the existing _RecordingLogger plus
build_indexing_pipeline and pipeline.run to verify the debug breadcrumb reflects
a selected endpoint rather than all None values, and ensure the warning path is
exercised where the selector falls back instead of failing.
🪄 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: a259ebfa-ec15-448d-b4fe-703792e0b403
📒 Files selected for processing (7)
openrag/di/container.pyopenrag/services/inference/ollama_client.pyopenrag/services/inference/vllm_client.pyopenrag/services/orchestrators/query_service.pyopenrag/services/workers/pipeline_builder.pytests/unit/services/orchestrators/test_query_service.pytests/unit/services/workers/test_pipeline_builder.py
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f7dfe5c733
ℹ️ 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".
|
|
||
| payload["messages"] = self._sanitize_messages(payload["messages"]) | ||
| chunk = await self._llm.chat(payload["messages"], **_sampling(payload)) | ||
| chunk = await self._resolve_llm(partitions).chat(payload["messages"], **_sampling(payload)) |
There was a problem hiding this comment.
Run preflight checks against the selected LLM
When a partition has chat_llm, this line can send the request to a different endpoint/model only after the router has already run preflight checks against config.llm (openrag/api/routers/user/chat.py:278 calls check_llm_model_availability, which validates only config.llm in api/dependencies/llm.py:18-39, and token checks use the default model budget before reaching the service). In that scenario a valid partition request is still rejected if the default endpoint is unavailable, and requests can also be admitted or rejected using the wrong context limit for the actual chat_llm model.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Agreed — same finding as hedhoud's note on line 471. Tracking this as a follow-up: the fix belongs in the HTTP preflight layer (check_llm_model_availability + the startup-primed get_max_model_tokens global), which must resolve the partition's effective chat_llm before the Depends and size/validate against that endpoint per-partition. Out of scope for this PR (which only moves LLM resolution into the service); I'll open an issue and link it.
The chat_llm preset now covers the query-contextualization call (ChatBotRag) as well as the final answer: chat/chat_stream/complete resolve the LLM once and hand it down to generate_query, so a request scoped to a partition never leaks its chat history to the default endpoint during query rewriting. Map-reduce (_infer_relevancy) deliberately stays on the default LLM: the feature is slated for a full post-release refactor and preset routing lands with it.
|
Scope update (66008bf): the Deliberately out of scope: map-reduce. |
POST/PATCH now fail with 422 MODEL_ENDPOINT_NOT_FOUND when chat_llm names an LLM endpoint missing from the catalog, mirroring the existing PRESET_NOT_FOUND check for indexation/retrieval presets. Only the incoming value is checked: a stored name that went stale (endpoint deleted after assignment) keeps its fail-open runtime fallback and never blocks unrelated PATCHes. An explicit chat_llm=null in a PATCH now really clears the preset — the service's None-filter used to swallow it, so the admin UI's 'reset to default' was a silent no-op. Blank strings normalize to null at the schema layer.
|
Assignment-time validation (1529ee3): Two deliberate boundaries:
11 new unit tests (6 service-level: reject/accept/clear/stale-doesn't-block on update + reject/accept on create; 5 schema-level for the null/blank normalization) — |
hedhoud
left a comment
There was a problem hiding this comment.
I rechecked the latest commit. The chat_llm assignment validation helps, but I still left a few notes around runtime behavior and one small doc drift.
| ) -> dict: | ||
| """Non-streaming chat completion → finalized OpenAI dict.""" | ||
| metadata = payload.get("metadata") or {} | ||
| llm = self._resolve_llm(partitions) |
There was a problem hiding this comment.
One thing still worries me here: the HTTP layer runs the availability and token checks against the default LLM before we reach this selected chat_llm. So a partition with a valid custom chat model can still be blocked if the default LLM is down, and the request can be checked against the wrong context size. I think those preflight checks should use the same LLM that will actually answer the request.
There was a problem hiding this comment.
Confirmed and agreed this is a real gap: the check_llm_model_availability dependency validates only config.llm, and check_tokens_limit uses get_max_model_tokens(), a single global primed once at startup from the default model. So a partition with a healthy custom chat_llm can be rejected when the default endpoint is down, and its request is sized against the wrong context window.
I'm scoping this as a follow-up rather than folding it into this PR, for three reasons:
- It's a pre-existing limitation of the HTTP preflight layer, independent of this PR's diff — this PR resolves the answer LLM in the service, which is what makes the gap reachable, but the fix lives entirely in
api/routers/user/chat.py+api/dependencies/llm.py. - A correct fix has to resolve the partition→
chat_llmmapping before theDepends, then run availability against that endpoint and cache/fetchmax_model_lenper endpoint (today it's one startup-primed global) — a non-trivial change to a dependency that runs on every chat request, deserving its own tests. - Keeping it separate avoids expanding this PR's blast radius right before merge.
Behaviorally, this PR is a strict improvement over main (the preset now actually answers); the preflight mismatch only degrades the default-down / non-default-context edge case, which existed implicitly before. I'll open a tracking issue for "run chat preflight checks against the resolved per-partition LLM" and link it here. Shout if you'd rather I pull it into this PR instead.
…olvable _select_vlm now guards the named-endpoint lookup with the same try/except KeyError + warning the contextualizer and topic-tagger selectors use. A VLM endpoint deleted or renamed after assignment no longer raises out of run() and fails the whole indexing job — captioning is enrichment, so it falls back to the legacy default VLM with a warning. Also refresh the _resolve_llm docstring: chat_llm IS validated on assignment now (previous commit), but a stored name can still go stale, which is what the runtime fallback covers. Tests: two former fail-fast assertions become fallback assertions, plus a breadcrumb test covering resolved named endpoints (vlm/ctx/topic/ embedder) alongside the existing all-disabled case.
hedhoud
left a comment
There was a problem hiding this comment.
Approved. The issues raised in review were handled well: the VLM fallback is now defensive, the chat_llm docs are aligned with the validation behavior, and the remaining preflight mismatch is tracked separately in #639. CI is green.
Summary
A partition's
chat_llmmodel-endpoint preset was stored and exposed through the admin API (PATCH /partition/{p}) but never used when chatting:QueryServicewas constructed with a single LLM built from the global settings and answered every chat/completion with it. The preset's only runtime consumer was the multiQuery/hyde retriever fallback inRetrievalService.Fix
QueryServicenow receives the shared namedllm_factoryand resolves the answer-generation LLM per request (chat,chat_stream,complete) via a new_resolve_llm(), mirroring the_resolve_chat_history_depthpartition semantics:openrag-allsentinel → default LLMchat_llm, falling back to the default when unsetchat_llmis not validated on assignment (the endpoint can be renamed/deleted afterwards), so an unknown name falls back to the default LLM with a warning instead of failing the request.Observability
Debug-level breadcrumbs to trace preset resolution end-to-end (zero added volume at the default INFO level):
None= stage disabled)chat_llmpreset (or why the default applies)VLLMClient/VLLMVision/OllamaClientlog a<Class> ready [model=… | endpoint=…]construction line, mirroring the existingVLLMEmbedderone — fires once per configured endpoint and maps a preset name to its base URL/modelUnresolvable preset names keep surfacing as warnings (
Skipping contextualization: cannot resolve LLM '…',Partition chat_llm preset not found …).Test plan
_resolve_llm(preset used, default paths, multi-partition unanimity, deleted-endpoint fallback, end-to-endchat()answering with the preset LLM); pipeline test extended for the new resolution line —tests/unit: 1638 passed (the 2test_rate_limit.pyfailures are pre-existing on the base branch)chat_llmendpoint (http://fake-chat-endpoint.invalid/v1) fails withCannot reach LLM at http://fake-chat-endpoint.invalid/v1on both streaming and non-streaming chat, while a preset-less partition answers from the default endpoint; an unknown preset name logs the fallback warning and still answersmodel endpoints resolved for indexing [… vlm=default | contextualization_llm=MistralOR | topic_tagging_llm=default | embedder=default]with contextualization running through the named endpointSummary by CodeRabbit
chat_llm, with automatic fallback to the default model.chat_llmto the default when set tonull, and reject unknown LLM endpoint names.chat_llmselections now safely revert to defaults.