From 05dcd5c1acf80cb5ff6356854a57e10c309329c5 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Tue, 7 Jul 2026 15:35:25 +0000 Subject: [PATCH 1/5] fix(chat): honor the partition chat_llm preset for answer generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- openrag/di/container.py | 7 +- .../services/orchestrators/query_service.py | 50 ++++++++++- .../orchestrators/test_query_service.py | 90 ++++++++++++++++++- 3 files changed, 141 insertions(+), 6 deletions(-) diff --git a/openrag/di/container.py b/openrag/di/container.py index 25a8210e3..06d778ff8 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -531,8 +531,10 @@ def query_service(self) -> QueryService: """QueryService — lazily built, cached for the container's lifetime. Shares the same core LLM construction as ``retrieval_service`` - (built from ``settings.llm``); the web-search service comes from - the ``WebSearchFactory`` (provider is ``None`` when + (built from ``settings.llm``); the named ``llm_factory`` lets the + service honor a partition's ``chat_llm`` model-endpoint preset for + answer generation. The web-search service comes from the + ``WebSearchFactory`` (provider is ``None`` when ``WEBSEARCH_API_TOKEN`` is unset — web search silently disabled). """ if self._query_service is None: @@ -554,6 +556,7 @@ def query_service(self) -> QueryService: config=settings, web_search_service=WebSearchFactory.create_service(settings), workspace_service=self.workspace_service, + llm_factory=self.llm_factory, ) return self._query_service diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 5eff227e7..e02454a04 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -107,9 +107,11 @@ def __init__( config: Settings, web_search_service: Any | None, workspace_service: WorkspaceService, + llm_factory: Callable[[str], LLM] | None = None, ) -> None: self._retrieval = retrieval_service self._llm = llm + self._llm_factory = llm_factory self._web = web_search_service self._workspace = workspace_service @@ -156,6 +158,48 @@ def _resolve_chat_history_depth(self, partition: list[str] | None) -> int: return max(explicit) return self._default_chat_history_depth + def _resolve_llm(self, partition: list[str] | None) -> LLM: + """Effective answer-generation LLM for this request. + + Honors a partition's configured ``chat_llm`` model-endpoint preset + (set via the admin API) over the default LLM. Same partition + semantics as ``_resolve_chat_history_depth``: no partition and the + ``"all"`` sentinel use the default; a multi-partition request uses + the preset only when every partition that sets one names the same + endpoint — with conflicting presets there is no single owning + partition, so the default applies. + + ``chat_llm`` is not validated against the endpoint catalog when it + is assigned, so an unknown name (e.g. an endpoint deleted after + assignment) must not fail the chat request — it falls back to the + default LLM with a warning. + """ + if self._llm_factory is None or not partition or "all" in partition: + logger.bind(partitions=partition).debug("Answering with the default LLM (no partition-scoped preset)") + return self._llm + names = { + cfg.chat_llm + for name in partition + if (cfg := self._config.partitions.get(name)) is not None and cfg.chat_llm + } + if len(names) != 1: + logger.bind(partitions=partition, chat_llm_presets=sorted(names)).debug( + "Answering with the default LLM (no single chat_llm preset among the partitions)" + ) + return self._llm + (chat_llm,) = names + try: + llm = self._llm_factory(chat_llm) + except KeyError: + logger.warning( + "Partition chat_llm preset not found in the model-endpoint catalog — using the default LLM", + chat_llm=chat_llm, + partitions=partition, + ) + return self._llm + logger.bind(chat_llm=chat_llm, partitions=partition).debug("Answering with the partition's chat_llm preset") + return llm + # ------------------------------------------------------------------ # Query generation (was RagPipeline.generate_query — no LangChain) # ------------------------------------------------------------------ @@ -423,7 +467,7 @@ async def chat( sources = prepare_sources(docs, web_results) 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)) chunk["model"] = model_name content = chunk.get("choices", [{}])[0].get("message", {}).get("content", "") or "" clean, citations = extract_and_strip_sources_block(content) @@ -448,7 +492,7 @@ async def chat_stream( sources = prepare_sources(docs, web_results) payload["messages"] = self._sanitize_messages(payload["messages"]) - llm_stream = self._llm.stream_chat(payload["messages"], **_sampling(payload)) + llm_stream = self._resolve_llm(partitions).stream_chat(payload["messages"], **_sampling(payload)) async for sse_line in stream_with_source_filtering(llm_stream, sources, model_name): yield sse_line @@ -466,7 +510,7 @@ async def complete( payload, docs = await self._prepare_completions(partitions, payload) sources = prepare_sources(docs, []) - resp = await self._llm.generate(payload["prompt"], **_sampling(payload, key="prompt")) + resp = await self._resolve_llm(partitions).generate(payload["prompt"], **_sampling(payload, key="prompt")) text = resp.get("choices", [{}])[0].get("text", "") or "" clean, citations = extract_and_strip_sources_block(text) resp["choices"][0]["text"] = clean diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 692375e85..53cc1cbfb 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -99,13 +99,14 @@ def _config(mode="SimpleRag"): ) -def _svc(*, llm=None, retrieval=None, web=None, mode="SimpleRag") -> QueryService: +def _svc(*, llm=None, retrieval=None, web=None, mode="SimpleRag", llm_factory=None) -> QueryService: return QueryService( retrieval_service=retrieval or FakeRetrieval(), llm=llm or FakeLLM(), config=_config(mode), web_search_service=web or FakeWeb(), workspace_service=FakeWorkspace(), + llm_factory=llm_factory, ) @@ -154,6 +155,93 @@ def test_resolve_chat_history_depth_multi_partition_takes_max_explicit(): assert svc._resolve_chat_history_depth(["a", "b", "c"]) == 6 +# --------------------------------------------------------------------------- # +# chat LLM resolution (per-partition chat_llm model-endpoint preset) +# --------------------------------------------------------------------------- # + + +def _partition(chat_llm=None, chat_history_depth=0): + return SimpleNamespace(chat_llm=chat_llm, chat_history_depth=chat_history_depth) + + +class RecordingFactory: + def __init__(self, llms=None): + self._llms = llms or {} + self.calls: list[str] = [] + + def __call__(self, name: str): + self.calls.append(name) + if name not in self._llms: + raise KeyError(name) + return self._llms[name] + + +def test_resolve_llm_uses_partition_preset(): + preset_llm = FakeLLM() + factory = RecordingFactory({"mistral": preset_llm}) + svc = _svc(llm_factory=factory) + svc._config.partitions = {"p": _partition(chat_llm="mistral")} + assert svc._resolve_llm(["p"]) is preset_llm + assert factory.calls == ["mistral"] + + +def test_resolve_llm_defaults_without_factory_partition_or_preset(): + default_llm = FakeLLM() + factory = RecordingFactory() + svc = _svc(llm=default_llm, llm_factory=factory) + svc._config.partitions = {"p": _partition(chat_llm=None), "q": _partition(chat_llm="mistral")} + assert svc._resolve_llm(None) is default_llm # direct/web-only mode + assert svc._resolve_llm(["all"]) is default_llm # cross-partition sentinel + assert svc._resolve_llm(["p"]) is default_llm # partition without a preset + assert svc._resolve_llm(["missing"]) is default_llm # unknown partition + assert factory.calls == [] # default paths never hit the factory + no_factory = _svc(llm=default_llm) + no_factory._config.partitions = {"q": _partition(chat_llm="mistral")} + assert no_factory._resolve_llm(["q"]) is default_llm + + +def test_resolve_llm_multi_partition_uses_preset_only_when_unanimous(): + default_llm, preset_llm = FakeLLM(), FakeLLM() + factory = RecordingFactory({"mistral": preset_llm}) + svc = _svc(llm=default_llm, llm_factory=factory) + svc._config.partitions = { + "a": _partition(chat_llm="mistral"), + "b": _partition(chat_llm="mistral"), + "c": _partition(chat_llm=None), # unset → doesn't veto + "d": _partition(chat_llm="llama"), + } + assert svc._resolve_llm(["a", "b", "c"]) is preset_llm # unanimous among setters + assert svc._resolve_llm(["a", "d"]) is default_llm # conflicting presets → default + + +def test_resolve_llm_unknown_preset_falls_back_to_default(): + # chat_llm is not validated on assignment (the endpoint may be deleted + # afterwards) — an unknown name must not fail the request. + default_llm = FakeLLM() + factory = RecordingFactory() # raises KeyError for every name + svc = _svc(llm=default_llm, llm_factory=factory) + svc._config.partitions = {"p": _partition(chat_llm="deleted")} + assert svc._resolve_llm(["p"]) is default_llm + assert factory.calls == ["deleted"] + + +@pytest.mark.asyncio +async def test_chat_answers_with_partition_chat_llm(): + default_llm = FakeLLM() + preset_llm = FakeLLM(chat_responses=["preset answer [Sources: none]"]) + svc = _svc(llm=default_llm, llm_factory=RecordingFactory({"mistral": preset_llm})) + svc._config.partitions = {"p": _partition(chat_llm="mistral")} + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + prepare_sources=lambda d, w: [], + model_name="openrag-p", + ) + assert out["choices"][0]["message"]["content"] == "preset answer" + assert len(preset_llm.chat_calls) == 1 + assert default_llm.chat_calls == [] # SimpleRag: no query-gen call either + + # --------------------------------------------------------------------------- # # generate_query # --------------------------------------------------------------------------- # From f7dfe5c7336939e7e5143028858c93166aec5e09 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Tue, 7 Jul 2026 15:35:35 +0000 Subject: [PATCH 2/5] feat(logging): surface which model-endpoint preset each stage uses 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 ' 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 --- openrag/services/inference/ollama_client.py | 1 + openrag/services/inference/vllm_client.py | 9 +++ openrag/services/workers/pipeline_builder.py | 77 +++++++++++-------- .../services/workers/test_pipeline_builder.py | 15 ++++ 4 files changed, 71 insertions(+), 31 deletions(-) diff --git a/openrag/services/inference/ollama_client.py b/openrag/services/inference/ollama_client.py index 19259f760..4ea702ecc 100644 --- a/openrag/services/inference/ollama_client.py +++ b/openrag/services/inference/ollama_client.py @@ -69,6 +69,7 @@ def __init__( timeout=timeout, headers={"Content-Type": "application/json"}, ) + logger.bind(model=self._model, endpoint=self._endpoint, timeout=timeout).debug("OllamaClient ready") @with_circuit_breaker("llm") @with_retry(max_attempts=3) diff --git a/openrag/services/inference/vllm_client.py b/openrag/services/inference/vllm_client.py index e132e781e..bdc8dcea1 100644 --- a/openrag/services/inference/vllm_client.py +++ b/openrag/services/inference/vllm_client.py @@ -76,6 +76,15 @@ def __init__( if api_key: headers["Authorization"] = f"Bearer {api_key}" self._client = httpx.AsyncClient(timeout=timeout, headers=headers) + # Same construction breadcrumb as VLLMEmbedder: the component factories + # cache instances per endpoint name, so this fires once per configured + # endpoint and shows which base URL/model a preset name resolved to. + logger.bind( + model=self._model, + endpoint=self._endpoint, + timeout=timeout, + enable_thinking=self._enable_thinking, + ).debug(f"{type(self).__name__} ready") def _resolve_overrides(self, kwargs: dict) -> tuple[str, str, dict[str, str] | None]: """Read ``metadata.llm_override`` from *kwargs* without mutating caller data. diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py index d90cce1c0..1d830c823 100644 --- a/openrag/services/workers/pipeline_builder.py +++ b/openrag/services/workers/pipeline_builder.py @@ -90,8 +90,8 @@ async def run(self, row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: parser = self._select_parser(config) chunker = self._select_chunker(config) embedder = self._select_embedder(row) - contextualizer = self._select_contextualizer(config) - topic_tagger = self._select_topic_tagger(config) + contextualizer, contextualization_llm = self._select_contextualizer(config) + topic_tagger, topic_tagging_llm = self._select_topic_tagger(config) timings: dict[str, float] = {} @@ -104,18 +104,27 @@ async def _timed(name: str, coro: Any) -> None: try: await _timed("parse", parse_stage(row, parser, timeout=self.timeouts.parse)) - if self._should_caption(row, config): - vlm = self._select_vlm(config) - if vlm is not None: - await _timed( - "caption", - caption_stage( - row, - vlm, - timeout=self.timeouts.caption, - per_image_timeout=self.timeouts.caption_per_image, - ), - ) + # The caption decision needs the parsed document (standalone images + # always caption), so the VLM is resolved after parse. + vlm, vlm_name = self._select_vlm(config) if self._should_caption(row, config) else (None, None) + logger.bind( + task_id=row.get("task_id"), + filename=row.get("filename", ""), + vlm=vlm_name if vlm is not None else None, + contextualization_llm=contextualization_llm, + topic_tagging_llm=topic_tagging_llm, + embedder=str(row.get("embedder_name") or "default"), + ).debug("model endpoints resolved for indexing (None = stage disabled)") + if vlm is not None: + await _timed( + "caption", + caption_stage( + row, + vlm, + timeout=self.timeouts.caption, + per_image_timeout=self.timeouts.caption_per_image, + ), + ) await _timed("chunk", chunk_stage(row, chunker, timeout=self.timeouts.chunk)) if contextualizer is not None: await _timed( @@ -196,17 +205,17 @@ def _select_embedder(self, row: MutableMapping[str, Any]) -> Embedder: return self.embedder_factory(str(embedder_name)) return self.embedder - def _select_vlm(self, config: IndexationPipelineConfig | None) -> VLM | None: - """Pick the captioning VLM instance (availability only — policy is in - ``_should_caption``).""" + def _select_vlm(self, config: IndexationPipelineConfig | None) -> tuple[VLM | None, str | None]: + """Pick the captioning VLM instance and its endpoint name for logging + (availability only — policy is in ``_should_caption``).""" if config is not None and self.vlm_factory is not None: if config.vlm: - return self.vlm_factory(config.vlm) + return self.vlm_factory(config.vlm), config.vlm try: - return self.vlm_factory("default") + return self.vlm_factory("default"), "default" except KeyError: - return self.vlm - return self.vlm + return self.vlm, "default" + return self.vlm, "default" def _should_caption(self, row: MutableMapping[str, Any], config: IndexationPipelineConfig | None) -> bool: """Decide whether to caption this document's images. @@ -222,37 +231,43 @@ def _should_caption(self, row: MutableMapping[str, Any], config: IndexationPipel per_partition = config.enable_image_captioning if config is not None else True return self.image_captioning and per_partition - def _select_contextualizer(self, config: IndexationPipelineConfig | None) -> ChunkContextualizer | None: + def _select_contextualizer( + self, config: IndexationPipelineConfig | None + ) -> tuple[ChunkContextualizer | None, str | None]: if config is not None: if not config.enable_contextualization: - return None + return None, None if self.contextualizer_factory is not None: name = config.contextualization_llm or "default" try: - return self.contextualizer_factory(name) + return self.contextualizer_factory(name), name except KeyError as exc: # Only an unresolvable endpoint name (the factory raises KeyError) # is skipped — contextualization is an enhancement and must not fail # the file over a missing/typo'd LLM. Any other factory error (prompt # load, bad client config, ...) is a real fault and is left to surface. logger.warning(f"Skipping contextualization: cannot resolve LLM '{name}' ({exc})") - return None - return self.contextualizer + return None, None + if self.contextualizer is None: + return None, None + return self.contextualizer, "default" - def _select_topic_tagger(self, config: IndexationPipelineConfig | None) -> TopicTagger | None: + def _select_topic_tagger(self, config: IndexationPipelineConfig | None) -> tuple[TopicTagger | None, str | None]: if config is not None: if not config.enable_topic_tagging: - return None + return None, None if self.topic_tagger_factory is not None: name = config.topic_tagging_llm or "default" try: - return self.topic_tagger_factory(name) + return self.topic_tagger_factory(name), name except KeyError as exc: # Only an unresolvable endpoint name (KeyError) is skipped; any other # factory error is a real fault and is left to surface, not masked. logger.warning(f"Skipping topic tagging: cannot resolve LLM '{name}' ({exc})") - return None - return self.topic_tagger + return None, None + if self.topic_tagger is None: + return None, None + return self.topic_tagger, "default" def build_indexing_pipeline( diff --git a/tests/unit/services/workers/test_pipeline_builder.py b/tests/unit/services/workers/test_pipeline_builder.py index efecd395b..1354f0b0f 100644 --- a/tests/unit/services/workers/test_pipeline_builder.py +++ b/tests/unit/services/workers/test_pipeline_builder.py @@ -560,6 +560,8 @@ class _RecordingLogger: def __init__(self) -> None: self.bound: dict | None = None self.message: str | None = None + self.debug_bound: dict | None = None + self.debug_message: str | None = None def bind(self, **kwargs): self.bound = kwargs @@ -568,6 +570,10 @@ def bind(self, **kwargs): def info(self, message: str) -> None: self.message = message + def debug(self, message: str) -> None: + self.debug_message = message + self.debug_bound = self.bound + recorder = _RecordingLogger() monkeypatch.setattr(pb, "logger", recorder) @@ -591,3 +597,12 @@ def info(self, message: str) -> None: assert recorder.bound["task_id"] == "t1" assert recorder.bound["n_chunks"] == 1 assert recorder.bound["filename"] == "note.txt" + + # The endpoint-resolution breadcrumb fired too: no VLM/contextualizer/topic + # tagger were configured, and the row carried no embedder_name. + assert recorder.debug_message == "model endpoints resolved for indexing (None = stage disabled)" + assert recorder.debug_bound is not None + assert recorder.debug_bound["vlm"] is None + assert recorder.debug_bound["contextualization_llm"] is None + assert recorder.debug_bound["topic_tagging_llm"] is None + assert recorder.debug_bound["embedder"] == "default" From 66008bfc07376fb44083ed535c0378638c61966e Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Wed, 8 Jul 2026 07:37:26 +0000 Subject: [PATCH 3/5] fix(chat): route query generation through the partition chat_llm preset 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. --- openrag/di/container.py | 3 +- .../services/orchestrators/query_service.py | 45 ++++++++++++------- .../orchestrators/test_query_service.py | 20 +++++++++ 3 files changed, 50 insertions(+), 18 deletions(-) diff --git a/openrag/di/container.py b/openrag/di/container.py index 06d778ff8..ad3f85fb3 100644 --- a/openrag/di/container.py +++ b/openrag/di/container.py @@ -533,7 +533,8 @@ def query_service(self) -> QueryService: Shares the same core LLM construction as ``retrieval_service`` (built from ``settings.llm``); the named ``llm_factory`` lets the service honor a partition's ``chat_llm`` model-endpoint preset for - answer generation. The web-search service comes from the + query generation and answer generation. The web-search service + comes from the ``WebSearchFactory`` (provider is ``None`` when ``WEBSEARCH_API_TOKEN`` is unset — web search silently disabled). """ diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index e02454a04..5ea9ddfbf 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -159,13 +159,17 @@ def _resolve_chat_history_depth(self, partition: list[str] | None) -> int: return self._default_chat_history_depth def _resolve_llm(self, partition: list[str] | None) -> LLM: - """Effective answer-generation LLM for this request. + """Effective LLM for this request — query generation and answering. Honors a partition's configured ``chat_llm`` model-endpoint preset - (set via the admin API) over the default LLM. Same partition - semantics as ``_resolve_chat_history_depth``: no partition and the - ``"all"`` sentinel use the default; a multi-partition request uses - the preset only when every partition that sets one names the same + (set via the admin API) over the default LLM. Resolved once per + request (``chat`` / ``chat_stream`` / ``complete``) and used for + both the query-contextualization call and the final answer; + map-reduce stays on the default LLM until its post-release + refactor. Same partition semantics as + ``_resolve_chat_history_depth``: no partition and the ``"all"`` + sentinel use the default; a multi-partition request uses the + preset only when every partition that sets one names the same endpoint — with conflicting presets there is no single owning partition, so the default applies. @@ -204,7 +208,8 @@ def _resolve_llm(self, partition: list[str] | None) -> LLM: # Query generation (was RagPipeline.generate_query — no LangChain) # ------------------------------------------------------------------ - async def generate_query(self, messages: list[dict]) -> SearchQueries: + async def generate_query(self, messages: list[dict], llm: LLM | None = None) -> SearchQueries: + llm = llm or self._llm last_user = messages[-1]["content"] if RAGMODE(self._rag_mode) is RAGMODE.SIMPLERAG: return SearchQueries(query_list=[Query(query=last_user)]) @@ -224,7 +229,7 @@ async def generate_query(self, messages: list[dict]) -> SearchQueries: } for attempt in (1, 2): try: - resp = await self._llm.chat(llm_messages, **params) + resp = await llm.chat(llm_messages, **params) content = resp["choices"][0]["message"]["content"] return SearchQueries.model_validate_json(_json_slice(content)) except Exception as exc: @@ -242,6 +247,9 @@ async def generate_query(self, messages: list[dict]) -> SearchQueries: # ------------------------------------------------------------------ async def _infer_relevancy(self, query: str, doc) -> tuple[bool, str]: + # Deliberately pinned to the default LLM: map-reduce is slated for a + # full post-release refactor, and routing it through the partition + # chat_llm preset is part of that work. async with get_llm_semaphore(): try: resp = await self._llm.chat( @@ -290,9 +298,9 @@ async def _batch(chunks: list, summaries: list) -> bool: # Preparation (was RagPipeline._prepare_for_chat_completion) # ------------------------------------------------------------------ - async def _prepare_chat(self, partition: list[str] | None, payload: dict): + async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: LLM | None = None): messages = payload["messages"][-self._resolve_chat_history_depth(partition) :] - queries = await self.generate_query(messages) + queries = await self.generate_query(messages, llm=llm) metadata = payload.get("metadata") or {} use_map_reduce = metadata.get("use_map_reduce", False) @@ -378,9 +386,9 @@ async def _gather_rag_and_web(self, query_list, partition, top_k, filter_params) doc_lists, web_lists = await asyncio.gather(rag, web) return doc_lists, web_lists - async def _prepare_completions(self, partition: list[str], payload: dict): + async def _prepare_completions(self, partition: list[str], payload: dict, llm: LLM | None = None): prompt = payload["prompt"] - queries = await self.generate_query([{"role": "user", "content": prompt}]) + queries = await self.generate_query([{"role": "user", "content": prompt}], llm=llm) chunks = await self._retrieval.retrieve_multi(partitions=partition, search_queries=queries) docs = [c.to_langchain() for c in chunks] context, included = format_context( @@ -460,14 +468,15 @@ async def chat( ) -> dict: """Non-streaming chat completion → finalized OpenAI dict.""" metadata = payload.get("metadata") or {} + llm = self._resolve_llm(partitions) if partitions is None and not metadata.get("websearch", False): docs, web_results = [], [] else: - payload, docs, web_results = await self._prepare_chat(partitions, payload) + payload, docs, web_results = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) payload["messages"] = self._sanitize_messages(payload["messages"]) - chunk = await self._resolve_llm(partitions).chat(payload["messages"], **_sampling(payload)) + chunk = await llm.chat(payload["messages"], **_sampling(payload)) chunk["model"] = model_name content = chunk.get("choices", [{}])[0].get("message", {}).get("content", "") or "" clean, citations = extract_and_strip_sources_block(content) @@ -485,14 +494,15 @@ async def chat_stream( ) -> AsyncIterator[str]: """Streaming chat completion → SSE strings with filtered sources.""" metadata = payload.get("metadata") or {} + llm = self._resolve_llm(partitions) if partitions is None and not metadata.get("websearch", False): docs, web_results = [], [] else: - payload, docs, web_results = await self._prepare_chat(partitions, payload) + payload, docs, web_results = await self._prepare_chat(partitions, payload, llm) sources = prepare_sources(docs, web_results) payload["messages"] = self._sanitize_messages(payload["messages"]) - llm_stream = self._resolve_llm(partitions).stream_chat(payload["messages"], **_sampling(payload)) + llm_stream = llm.stream_chat(payload["messages"], **_sampling(payload)) async for sse_line in stream_with_source_filtering(llm_stream, sources, model_name): yield sse_line @@ -504,13 +514,14 @@ async def complete( prepare_sources: PrepareSources, ) -> dict: """Non-streaming text completion → finalized OpenAI dict.""" + llm = self._resolve_llm(partitions) if partitions is None: docs = [] else: - payload, docs = await self._prepare_completions(partitions, payload) + payload, docs = await self._prepare_completions(partitions, payload, llm) sources = prepare_sources(docs, []) - resp = await self._resolve_llm(partitions).generate(payload["prompt"], **_sampling(payload, key="prompt")) + resp = await llm.generate(payload["prompt"], **_sampling(payload, key="prompt")) text = resp.get("choices", [{}])[0].get("text", "") or "" clean, citations = extract_and_strip_sources_block(text) resp["choices"][0]["text"] = clean diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 53cc1cbfb..9f981cc67 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -242,6 +242,26 @@ async def test_chat_answers_with_partition_chat_llm(): assert default_llm.chat_calls == [] # SimpleRag: no query-gen call either +@pytest.mark.asyncio +async def test_chat_query_generation_uses_partition_chat_llm(): + # ChatBotRag contextualizes the query with an LLM call — that call must + # go through the partition preset too, not just the final answer. + default_llm = FakeLLM() + query_json = json.dumps({"query_list": [{"query": "rewritten", "temporal_filters": None}]}) + preset_llm = FakeLLM(chat_responses=[query_json, "preset answer [Sources: none]"]) + svc = _svc(mode="ChatBotRag", llm=default_llm, llm_factory=RecordingFactory({"mistral": preset_llm})) + svc._config.partitions = {"p": _partition(chat_llm="mistral")} + out = await svc.chat( + partitions=["p"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {}}, + prepare_sources=lambda d, w: [], + model_name="openrag-p", + ) + assert out["choices"][0]["message"]["content"] == "preset answer" + assert len(preset_llm.chat_calls) == 2 # query generation + answer + assert default_llm.chat_calls == [] + + # --------------------------------------------------------------------------- # # generate_query # --------------------------------------------------------------------------- # From 1529ee348599eb2f691a88bdbd321d10034bcb36 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Wed, 8 Jul 2026 07:52:10 +0000 Subject: [PATCH 4/5] fix(partitions): validate chat_llm on assignment and let null clear it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- openrag/api/routers/admin/partitions.py | 2 +- .../api/schemas/admin/partition_schemas.py | 17 ++++ .../orchestrators/partition_service.py | 34 ++++++- .../api/schemas/admin/test_phase14_schemas.py | 21 +++++ .../test_partition_preset_resolution.py | 93 +++++++++++++++++++ 5 files changed, 163 insertions(+), 4 deletions(-) diff --git a/openrag/api/routers/admin/partitions.py b/openrag/api/routers/admin/partitions.py index 8509f43ca..1f098b78d 100644 --- a/openrag/api/routers/admin/partitions.py +++ b/openrag/api/routers/admin/partitions.py @@ -325,7 +325,7 @@ async def create_partition( - `indexation_preset` - `retrieval_preset` - `chat_history_depth` -- `chat_llm` +- `chat_llm` (must name a registered LLM endpoint — 422 otherwise; explicit `null` resets to the default LLM) **Permissions:** - Requires partition owner role diff --git a/openrag/api/schemas/admin/partition_schemas.py b/openrag/api/schemas/admin/partition_schemas.py index e78e4fc46..281a37ab8 100644 --- a/openrag/api/schemas/admin/partition_schemas.py +++ b/openrag/api/schemas/admin/partition_schemas.py @@ -42,6 +42,14 @@ def validate_non_empty_name(cls, value: str) -> str: """Normalize non-null partition and preset names.""" return _normalize_name(value) + @field_validator("chat_llm") + @classmethod + def normalize_chat_llm(cls, value: str | None) -> str | None: + """Blank and null both mean "use the default LLM" (stored as NULL).""" + if value is None: + return None + return value.strip() or None + class UpdatePartitionRequest(BaseModel): """Request body for updating partition preset assignments.""" @@ -73,6 +81,15 @@ def reject_null_chat_history_depth(cls, value: int | None) -> int: """Reject explicit null for chat history depth updates.""" return _reject_explicit_null("chat_history_depth", value) + @field_validator("chat_llm") + @classmethod + def normalize_chat_llm(cls, value: str | None) -> str | None: + """Normalize ``chat_llm`` — unlike the other fields, explicit null is + allowed: it (or a blank string) resets the partition to the default LLM.""" + if value is None: + return None + return value.strip() or None + @model_validator(mode="after") def require_at_least_one_update(self) -> UpdatePartitionRequest: """Reject empty update payloads.""" diff --git a/openrag/services/orchestrators/partition_service.py b/openrag/services/orchestrators/partition_service.py index fa016c5fb..5cfcabc7c 100644 --- a/openrag/services/orchestrators/partition_service.py +++ b/openrag/services/orchestrators/partition_service.py @@ -57,6 +57,11 @@ # ``list_existant_partitions``. Matched case-insensitively. _RESERVED_PARTITION_NAMES = frozenset({"all"}) +# Columns where an explicit ``None`` in a PATCH is a real value (SQL NULL = +# "reset to default"), not the omitted-field sentinel that the None-filter +# in ``update_partition`` gives every other column. +_NULLABLE_COLUMNS = frozenset({"chat_llm"}) + def _validate_limit(limit: int | None) -> None: """Reject negative ``limit`` values before they reach ``rows[:limit]``. @@ -223,6 +228,8 @@ async def create_partition( # Validate the preset references before touching the DB. if self._config is not None: self._validate_preset_refs({"partition": partition, **config_fields}) + if chat_llm: + self._validate_chat_llm_ref(chat_llm) await self._partition_repo.create_partition(name=partition, user_id=user_id, max_owned=max_owned) @@ -271,18 +278,26 @@ async def delete_partition(self, partition: str) -> None: async def update_partition(self, partition: str, **fields: object) -> dict | None: """Update a partition's config columns and re-resolve the cache. - ``None`` values are ignored (so partial PATCH semantics work). When a + ``None`` values are ignored (so partial PATCH semantics work) — + except for the ``_NULLABLE_COLUMNS`` (``chat_llm``), where an + explicit ``None`` clears the column back to its default. When a preset reference changes, the merged row is validated *before* the - write so an unknown preset name fails fast. + write so an unknown preset name fails fast; likewise an assigned + ``chat_llm`` must name a catalogued LLM endpoint. """ await self._ensure_partition(partition) - updates = {k: v for k, v in fields.items() if v is not None} + updates = {k: v for k, v in fields.items() if v is not None or k in _NULLABLE_COLUMNS} if self._config is not None and updates: current = await self._partition_repo.get_partition_row(partition) if current is None: raise PartitionNotFoundError(f"Partition '{partition}' does not exist.") self._validate_preset_refs({**current, **updates}) + # Only the incoming value is checked — a *stored* name that went + # stale (endpoint deleted later) must not block unrelated PATCHes; + # QueryService falls back to the default LLM for those at runtime. + if updates.get("chat_llm"): + self._validate_chat_llm_ref(updates["chat_llm"]) result = await self._partition_repo.update_partition(partition, **updates) @@ -324,6 +339,19 @@ def _validate_preset_refs(self, row: dict) -> None: except ConfigError as exc: raise ValidationError(exc.message, code="PRESET_NOT_FOUND") from exc + def _validate_chat_llm_ref(self, chat_llm: str) -> None: + """Assignment-time check: ``chat_llm`` must name a catalogued LLM endpoint. + + Only guards assignment — a stored name can go stale afterwards (the + endpoint may be renamed or deleted), which QueryService tolerates by + falling back to the default LLM at request time. + """ + if chat_llm not in self._require_config().models.llm: + raise ValidationError( + f"LLM endpoint '{chat_llm}' referenced by chat_llm not found.", + code="MODEL_ENDPOINT_NOT_FOUND", + ) + def _partition_detail(self, row: dict, cfg: PartitionConfig) -> dict: """Shape a resolved row into the ``PartitionDetailResponse`` payload.""" return { diff --git a/tests/unit/api/schemas/admin/test_phase14_schemas.py b/tests/unit/api/schemas/admin/test_phase14_schemas.py index 105071731..c6e10d0a7 100644 --- a/tests/unit/api/schemas/admin/test_phase14_schemas.py +++ b/tests/unit/api/schemas/admin/test_phase14_schemas.py @@ -122,3 +122,24 @@ def test_update_partition_rejects_explicit_nulls(field): """Partition updates reject null for fields that cannot be cleared.""" with pytest.raises(ValidationError, match=rf"{field} cannot be null"): UpdatePartitionRequest(**{field: None}) + + +def test_update_partition_chat_llm_allows_explicit_null_as_reset(): + """Unlike the other fields, chat_llm=null is a real value: reset to default.""" + request = UpdatePartitionRequest(chat_llm=None) + assert request.chat_llm is None + assert "chat_llm" in request.model_fields_set + + +@pytest.mark.parametrize("blank", ["", " "]) +def test_update_partition_chat_llm_normalizes_blank_to_null(blank): + """A blank chat_llm means the same as null: reset to the default LLM.""" + assert UpdatePartitionRequest(chat_llm=blank).chat_llm is None + + +def test_update_partition_chat_llm_strips_whitespace(): + assert UpdatePartitionRequest(chat_llm=" mistral ").chat_llm == "mistral" + + +def test_create_partition_chat_llm_normalizes_blank_to_null(): + assert CreatePartitionRequest(name="legal", chat_llm=" ").chat_llm is None diff --git a/tests/unit/services/orchestrators/test_partition_preset_resolution.py b/tests/unit/services/orchestrators/test_partition_preset_resolution.py index 64b2219d7..391bade25 100644 --- a/tests/unit/services/orchestrators/test_partition_preset_resolution.py +++ b/tests/unit/services/orchestrators/test_partition_preset_resolution.py @@ -274,6 +274,99 @@ async def test_update_partition_missing_raises_404(): await svc.update_partition("ghost", description="x") +# ------------------------------------------------------------------ +# chat_llm assignment (model-endpoint reference) +# ------------------------------------------------------------------ + + +def _settings_with_llm(*names: str): + from core.config.model_endpoints import ModelEndpointConfig + + s = _settings() + s.models.llm.update({n: ModelEndpointConfig(endpoint="http://llm:8000/v1") for n in names}) + return s + + +@pytest.mark.asyncio +async def test_update_partition_rejects_unknown_chat_llm(): + from core.utils.exceptions import ValidationError + + repo = _FakePartitionRepo(rows=[_full_row("p1")]) + svc = _make_service(repo, settings=_settings_with_llm("mistral")) + + with pytest.raises(ValidationError, match="LLM endpoint 'ghost'") as exc: + await svc.update_partition("p1", chat_llm="ghost") + assert exc.value.status_code == 422 + assert exc.value.code == "MODEL_ENDPOINT_NOT_FOUND" + assert repo._store["p1"]["chat_llm"] is None # nothing was written + + +@pytest.mark.asyncio +async def test_update_partition_accepts_catalogued_chat_llm(): + settings = _settings_with_llm("mistral") + repo = _FakePartitionRepo(rows=[_full_row("p1")]) + svc = _make_service(repo, settings=settings) + + await svc.update_partition("p1", chat_llm="mistral") + + assert repo._store["p1"]["chat_llm"] == "mistral" + assert settings.partitions["p1"].chat_llm == "mistral" + + +@pytest.mark.asyncio +async def test_update_partition_explicit_none_clears_chat_llm(): + # The UI resets to the default LLM by PATCHing chat_llm=null — the + # None-filter that gives other columns partial-PATCH semantics must + # not swallow it. + settings = _settings_with_llm("mistral") + repo = _FakePartitionRepo(rows=[_full_row("p1", chat_llm="mistral")]) + svc = _make_service(repo, settings=settings) + + await svc.update_partition("p1", chat_llm=None) + + assert repo._store["p1"]["chat_llm"] is None + assert settings.partitions["p1"].chat_llm is None + + +@pytest.mark.asyncio +async def test_update_partition_stale_stored_chat_llm_does_not_block_other_updates(): + # Endpoint deleted after assignment: the stored name is stale, but a + # PATCH that doesn't touch chat_llm must still succeed (runtime falls + # back to the default LLM for the stale name). + repo = _FakePartitionRepo(rows=[_full_row("p1", chat_llm="deleted-endpoint")]) + svc = _make_service(repo, settings=_settings_with_llm("mistral")) + + await svc.update_partition("p1", description="new") + + assert repo._store["p1"]["description"] == "new" + assert repo._store["p1"]["chat_llm"] == "deleted-endpoint" + + +@pytest.mark.asyncio +async def test_create_partition_rejects_unknown_chat_llm(): + from core.utils.exceptions import ValidationError + + repo = _FakePartitionRepo() + svc = _make_service(repo, settings=_settings_with_llm("mistral")) + + with pytest.raises(ValidationError, match="LLM endpoint 'ghost'") as exc: + await svc.create_partition("p1", user_id=1, chat_llm="ghost") + assert exc.value.code == "MODEL_ENDPOINT_NOT_FOUND" + assert not await repo.partition_exists("p1") + + +@pytest.mark.asyncio +async def test_create_partition_accepts_catalogued_chat_llm(): + settings = _settings_with_llm("mistral") + repo = _FakePartitionRepo() + svc = _make_service(repo, settings=settings) + + await svc.create_partition("p1", user_id=1, chat_llm="mistral") + + assert repo._store["p1"]["chat_llm"] == "mistral" + assert settings.partitions["p1"].chat_llm == "mistral" + + # ------------------------------------------------------------------ # get_partition_config / update_partition_config (PartitionDetailResponse) # ------------------------------------------------------------------ From abad29d8eae17dec070f3fea683297d2e1025786 Mon Sep 17 00:00:00 2001 From: Ahmath-Gadji Date: Wed, 8 Jul 2026 08:12:22 +0000 Subject: [PATCH 5/5] fix(indexing): fall back instead of failing when a named VLM is unresolvable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _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. --- .../services/orchestrators/query_service.py | 10 ++- openrag/services/workers/pipeline_builder.py | 18 +++-- .../test_model_endpoint_registry_e2e.py | 25 +++--- .../services/workers/test_pipeline_builder.py | 80 ++++++++++++++++++- 4 files changed, 109 insertions(+), 24 deletions(-) diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 5ea9ddfbf..173c71ddb 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -173,10 +173,12 @@ def _resolve_llm(self, partition: list[str] | None) -> LLM: endpoint — with conflicting presets there is no single owning partition, so the default applies. - ``chat_llm`` is not validated against the endpoint catalog when it - is assigned, so an unknown name (e.g. an endpoint deleted after - assignment) must not fail the chat request — it falls back to the - default LLM with a warning. + ``chat_llm`` is validated against the endpoint catalog when it is + assigned (``PartitionService`` rejects an unknown name at create / + PATCH time), but a stored name can still go stale afterwards — the + endpoint may be renamed or deleted after assignment — so an + unresolvable name here must not fail the chat request; it falls back + to the default LLM with a warning. """ if self._llm_factory is None or not partition or "all" in partition: logger.bind(partitions=partition).debug("Answering with the default LLM (no partition-scoped preset)") diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py index 1d830c823..9721c4ec5 100644 --- a/openrag/services/workers/pipeline_builder.py +++ b/openrag/services/workers/pipeline_builder.py @@ -207,13 +207,21 @@ def _select_embedder(self, row: MutableMapping[str, Any]) -> Embedder: def _select_vlm(self, config: IndexationPipelineConfig | None) -> tuple[VLM | None, str | None]: """Pick the captioning VLM instance and its endpoint name for logging - (availability only — policy is in ``_should_caption``).""" + (availability only — policy is in ``_should_caption``). + + Captioning is an enrichment step, so an unresolvable endpoint name must + not fail the file — same rationale as ``_select_contextualizer`` / + ``_select_topic_tagger``. A named VLM whose endpoint was deleted or + renamed after assignment (the factory raises ``KeyError``) falls back to + the legacy VLM with a warning instead of breaking indexing for the whole + partition. + """ if config is not None and self.vlm_factory is not None: - if config.vlm: - return self.vlm_factory(config.vlm), config.vlm + name = config.vlm or "default" try: - return self.vlm_factory("default"), "default" - except KeyError: + return self.vlm_factory(name), name + except KeyError as exc: + logger.warning(f"Skipping named VLM: cannot resolve '{name}' ({exc}) — falling back to the default VLM") return self.vlm, "default" return self.vlm, "default" diff --git a/tests/unit/services/workers/test_model_endpoint_registry_e2e.py b/tests/unit/services/workers/test_model_endpoint_registry_e2e.py index 2bffbcb4d..653dac9a4 100644 --- a/tests/unit/services/workers/test_model_endpoint_registry_e2e.py +++ b/tests/unit/services/workers/test_model_endpoint_registry_e2e.py @@ -320,23 +320,26 @@ async def test_indexing_uses_named_vlm_loaded_from_model_endpoint_registry(): @pytest.mark.asyncio -async def test_indexing_fails_for_missing_named_vlm_when_image_captioning_is_enabled(): +async def test_indexing_falls_back_to_default_vlm_when_named_vlm_is_missing(): + # A named VLM endpoint deleted/renamed after assignment must not fail the + # indexing job — captioning falls back to the legacy default VLM with a + # warning (parity with the contextualizer/topic-tagger selectors). settings = _settings() vlm_factory = _build_vlm_factory(settings) await _hydrate(settings, _FakeEndpointRepo()) default_vlm = _DefaultVLM() - with pytest.raises(KeyError, match="Unknown vlm 'missing-vlm'"): - await _run_pipeline( - settings=settings, - vlm_name="missing-vlm", - image_captioning=True, - default_vlm=default_vlm, - vlm_factory=vlm_factory, - ) + row, _, returned_vlm = await _run_pipeline( + settings=settings, + vlm_name="missing-vlm", + image_captioning=True, + default_vlm=default_vlm, + vlm_factory=vlm_factory, + ) - assert default_vlm.calls == [] - assert _RecordingVLM.instances == [] + assert returned_vlm is default_vlm + assert default_vlm.calls == [b"png"] # captioned via the legacy fallback + assert _RecordingVLM.instances == [] # the named endpoint never resolved @pytest.mark.asyncio diff --git a/tests/unit/services/workers/test_pipeline_builder.py b/tests/unit/services/workers/test_pipeline_builder.py index 1354f0b0f..d87d1b341 100644 --- a/tests/unit/services/workers/test_pipeline_builder.py +++ b/tests/unit/services/workers/test_pipeline_builder.py @@ -211,7 +211,10 @@ def _missing_default(_name: str): @pytest.mark.asyncio -async def test_pipeline_fails_when_explicit_named_vlm_is_missing(): +async def test_pipeline_falls_back_when_explicit_named_vlm_is_missing(): + # A named VLM endpoint deleted/renamed after assignment must not fail the + # whole indexing job — captioning is enrichment, so it falls back to the + # legacy VLM with a warning (mirrors the contextualizer/topic-tagger paths). def _missing_named(name: str): raise KeyError(f"Unknown vlm '{name}'") @@ -221,12 +224,13 @@ def _missing_named(name: str): text_blocks=[TextBlock(text="hello")], images=[ImageBlock(image_bytes=b"png")], ) + fallback_vlm = FakeVLM() pipeline = build_indexing_pipeline( parser=FakeParser(processed), chunker=FakeChunker([Chunk(id="c1", text="hello", partition="tenant-a")]), embedder=FakeEmbedder([[1.0]]), vector_store=FakeVectorStore(), - vlm=FakeVLM(), + vlm=fallback_vlm, vlm_factory=_missing_named, ) row = { @@ -239,8 +243,9 @@ def _missing_named(name: str): }, } - with pytest.raises(KeyError, match="missing-vlm"): - await pipeline.run(row) + await pipeline.run(row) # no raise + + assert fallback_vlm.calls == [b"png"] # captioned via the legacy fallback VLM @pytest.mark.asyncio @@ -606,3 +611,70 @@ def debug(self, message: str) -> None: assert recorder.debug_bound["contextualization_llm"] is None assert recorder.debug_bound["topic_tagging_llm"] is None assert recorder.debug_bound["embedder"] == "default" + + +@pytest.mark.asyncio +async def test_pipeline_breadcrumb_reflects_resolved_named_endpoints(monkeypatch): + """The breadcrumb names the actually-resolved endpoint per stage, not just + ``None`` — the named-endpoint counterpart to the all-disabled case above.""" + import services.workers.pipeline_builder as pb + + class _RecordingLogger: + def __init__(self) -> None: + self.bound: dict | None = None + self.debug_bound: dict | None = None + + def bind(self, **kwargs): + self.bound = kwargs + return self + + def info(self, message: str) -> None: # unused here, kept for parity + pass + + def debug(self, message: str) -> None: + self.debug_bound = self.bound + + def warning(self, message: str) -> None: + pass + + recorder = _RecordingLogger() + monkeypatch.setattr(pb, "logger", recorder) + + document = Document(filename="note.txt", text="hello", partition="tenant-a") + processed = ProcessedDocument( + document_id=document.id, + text_blocks=[TextBlock(text="hello")], + images=[ImageBlock(image_bytes=b"png")], + ) + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=FakeChunker([Chunk(id="c1", text="hello", partition="tenant-a")]), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + vlm_factory=lambda name: FakeVLM(), + contextualizer_factory=lambda name: FakeContextualizer(), + topic_tagger_factory=lambda name: FakeTopicTagger(), + ) + row = { + "task_id": "t1", + "document": document, + "partition": "tenant-a", + "filename": "note.txt", + "embedder_name": "embed-fast", + "indexation_config": { + "enable_image_captioning": True, + "vlm": "vlm-fast", + "enable_contextualization": True, + "contextualization_llm": "ctx-llm", + "enable_topic_tagging": True, + "topic_tagging_llm": "topic-llm", + }, + } + + await pipeline.run(row) + + assert recorder.debug_bound is not None + assert recorder.debug_bound["vlm"] == "vlm-fast" + assert recorder.debug_bound["contextualization_llm"] == "ctx-llm" + assert recorder.debug_bound["topic_tagging_llm"] == "topic-llm" + assert recorder.debug_bound["embedder"] == "embed-fast"