Skip to content

fix(chat): honor the partition chat_llm preset when chatting - #634

Merged
Ahmath-Gadji merged 5 commits into
refactor/hexagonalfrom
fix/partition-chat-llm-preset
Jul 8, 2026
Merged

fix(chat): honor the partition chat_llm preset when chatting#634
Ahmath-Gadji merged 5 commits into
refactor/hexagonalfrom
fix/partition-chat-llm-preset

Conversation

@Ahmath-Gadji

@Ahmath-Gadji Ahmath-Gadji commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Summary

A partition's chat_llm model-endpoint preset was stored and exposed through the admin API (PATCH /partition/{p}) but never used when chatting: QueryService was 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 in RetrievalService.

Fix

  • QueryService now receives the shared named llm_factory and resolves the answer-generation LLM per request (chat, chat_stream, complete) via a new _resolve_llm(), mirroring the _resolve_chat_history_depth partition semantics:
    • no partition / openrag-all sentinel → default LLM
    • single partition → its chat_llm, falling back to the default when unset
    • multi-partition → the preset only when every partition that sets one names the same endpoint
  • chat_llm is 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.
  • The factory caches clients per endpoint name and already gets evicted on endpoint update/rename, so there is no per-request construction cost.

Observability

Debug-level breadcrumbs to trace preset resolution end-to-end (zero added volume at the default INFO level):

  • one line per indexed file with the resolved captioning VLM, contextualization LLM, topic-tagging LLM and embedder names (None = stage disabled)
  • one line per chat request naming the resolved chat_llm preset (or why the default applies)
  • VLLMClient / VLLMVision / OllamaClient log a <Class> ready [model=… | endpoint=…] construction line, mirroring the existing VLLMEmbedder one — fires once per configured endpoint and maps a preset name to its base URL/model

Unresolvable preset names keep surfacing as warnings (Skipping contextualization: cannot resolve LLM '…', Partition chat_llm preset not found …).

Test plan

  • 5 new unit tests for _resolve_llm (preset used, default paths, multi-partition unanimity, deleted-endpoint fallback, end-to-end chat() answering with the preset LLM); pipeline test extended for the new resolution line — tests/unit: 1638 passed (the 2 test_rate_limit.py failures are pre-existing on the base branch)
  • Verified live on a compose stack: a partition wired to a fake chat_llm endpoint (http://fake-chat-endpoint.invalid/v1) fails with Cannot reach LLM at http://fake-chat-endpoint.invalid/v1 on 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 answers
  • Verified the indexing line live: model endpoints resolved for indexing [… vlm=default | contextualization_llm=MistralOR | topic_tagging_llm=default | embedder=default] with contextualization running through the named endpoint

Summary by CodeRabbit

  • New Features
    • Partitioned LLM selection for both query generation and final answering using per-partition chat_llm, with automatic fallback to the default model.
  • Bug Fixes
    • Partition updates now correctly reset chat_llm to the default when set to null, and reject unknown LLM endpoint names.
    • Conflicting per-partition chat_llm selections now safely revert to defaults.
  • Logging Improvements
    • Enhanced debug context during inference client initialization and indexing stage selection.
  • Tests
    • Expanded unit coverage for partition-scoped model selection, schema validation, and logging.

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
@coderabbitai

coderabbitai Bot commented Jul 7, 2026

Copy link
Copy Markdown

Review Change Stack

Warning

Review limit reached

@Ahmath-Gadji, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 25 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: abae1205-7beb-48d8-97d4-084a4c30e8c9

📥 Commits

Reviewing files that changed from the base of the PR and between 1529ee3 and abad29d.

📒 Files selected for processing (4)
  • openrag/services/orchestrators/query_service.py
  • openrag/services/workers/pipeline_builder.py
  • tests/unit/services/workers/test_model_endpoint_registry_e2e.py
  • tests/unit/services/workers/test_pipeline_builder.py
📝 Walkthrough

Walkthrough

Partition chat_llm now flows from API validation through partition persistence into QueryService LLM resolution, and indexing pipeline selection now returns endpoint names that are included in debug logging. Inference clients also emit startup debug logs.

Changes

Partition chat LLM plumbing

Layer / File(s) Summary
API contract and normalization
openrag/api/schemas/admin/partition_schemas.py, openrag/api/routers/admin/partitions.py, tests/unit/api/schemas/admin/test_phase14_schemas.py
chat_llm is normalized on create/update, the admin PATCH description documents reset and invalid-name behavior, and schema tests cover null, blank, and whitespace handling.
Partition persistence validation
openrag/services/orchestrators/partition_service.py, tests/unit/services/orchestrators/test_partition_preset_resolution.py
chat_llm is treated as a nullable assignment, validated against configured LLM endpoints on create and update, and covered by tests for valid, invalid, reset, and stale values.
QueryService partition LLM resolution
openrag/di/container.py, openrag/services/orchestrators/query_service.py, tests/unit/services/orchestrators/test_query_service.py
QueryService receives llm_factory, resolves a per-partition LLM for chat and completion flows, keeps relevancy inference on the default LLM, and tests cover preset selection and routing.

Pipeline logging and client readiness

Layer / File(s) Summary
Selection and run logging
openrag/services/workers/pipeline_builder.py, tests/unit/services/workers/test_pipeline_builder.py
IndexingPipeline selection methods return instances plus endpoint names, run() binds them into logger context, and the pipeline logging test asserts the new debug breadcrumb.
Inference client debug logs
openrag/services/inference/ollama_client.py, openrag/services/inference/vllm_client.py
OllamaClient and VLLMClient emit construction-time debug logs with model, endpoint, timeout, and VLLM-specific settings.

Estimated code review effort: 3 (Moderate) | ~25 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title is concise and accurately summarizes the main change: honoring partition chat_llm presets during chat requests.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/partition-chat-llm-preset

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot added the fix Fix issue label Jul 7, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (1)
tests/unit/services/workers/test_pipeline_builder.py (1)

560-608: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Only 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 in pipeline_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

📥 Commits

Reviewing files that changed from the base of the PR and between 55c19eb and f7dfe5c.

📒 Files selected for processing (7)
  • openrag/di/container.py
  • openrag/services/inference/ollama_client.py
  • openrag/services/inference/vllm_client.py
  • openrag/services/orchestrators/query_service.py
  • openrag/services/workers/pipeline_builder.py
  • tests/unit/services/orchestrators/test_query_service.py
  • tests/unit/services/workers/test_pipeline_builder.py

Comment thread openrag/services/workers/pipeline_builder.py
@hedhoud

hedhoud commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge 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 👍 / 👎.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tracking issue opened: #639.

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.
@Ahmath-Gadji

Copy link
Copy Markdown
Collaborator Author

Scope update (66008bf): the chat_llm preset now also covers query generationchat / chat_stream / complete resolve the LLM once per request and hand it down to generate_query, so a partition-scoped request no longer sends its chat history to the default endpoint for query contextualization in ChatBotRag mode. Covered by a new unit test asserting both the query-gen call and the answer go through the preset LLM.

Deliberately out of scope: map-reduce. _infer_relevancy (the map-reduce relevancy/summarization calls) intentionally stays on the default LLM for this release — map-reduce is an entire feature due for a full refactor post-release, and preset routing will land as part of that work. There's a comment in the code marking this, so please don't flag it as an oversight in review.

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.
@Ahmath-Gadji

Copy link
Copy Markdown
Collaborator Author

Assignment-time validation (1529ee3): chat_llm is now validated when it is assignedPOST /partition and PATCH /partition/{p} return 422 MODEL_ENDPOINT_NOT_FOUND when the name isn't in the LLM endpoint catalog, mirroring the existing PRESET_NOT_FOUND check for indexation/retrieval presets. Typos now fail at the API instead of silently answering with the default LLM.

Two deliberate boundaries:

  • Only the incoming value is checked. A stored name that goes stale later (endpoint renamed/deleted after assignment) never blocks unrelated PATCHes — the fail-open runtime fallback from this PR still covers that case.
  • Explicit null now really clears the preset. The service's None-filter used to swallow chat_llm: null, so the admin UI's "reset to default" was a silent no-op — fixed by treating chat_llm as a nullable column; blank strings normalize to null at the schema layer.

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) — tests/unit fully green (1652 passed, incl. the rate-limit pair).

@coderabbitai coderabbitai Bot removed the fix Fix issue label Jul 8, 2026

@hedhoud hedhoud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread openrag/services/orchestrators/query_service.py Outdated
) -> dict:
"""Non-streaming chat completion → finalized OpenAI dict."""
metadata = payload.get("metadata") or {}
llm = self._resolve_llm(partitions)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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.
  2. A correct fix has to resolve the partition→chat_llm mapping before the Depends, then run availability against that endpoint and cache/fetch max_model_len per 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.
  3. 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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Tracking issue opened: #639.

Comment thread openrag/services/workers/pipeline_builder.py Outdated
…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 hedhoud left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fix Fix issue

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants