diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 2d87a8648..8b4c9c2dc 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -180,50 +180,116 @@ def _resolve_chat_history_depth(self, partition: list[str] | None) -> int: def _resolve_llm(self, partition: list[str] | None) -> LLM: """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. 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. + Resolution order: + + 1. A partition's configured ``chat_llm`` model-endpoint preset (set + via the admin API) wins when the request scopes to one or more + named partitions that agree on a single preset. Resolved once per + request (``chat`` / ``chat_stream`` / ``complete``) and used for + both the query-contextualization call and the final answer. + Map-reduce is the exception: its relevancy/summarisation passes + stay pinned to the static ``self._llm`` (see ``_infer_relevancy``), + so with a non-env default endpoint a map-reduce request's sub-calls + run on a different model than its answer, until that post-release + refactor lands. + 2. Otherwise the **catalog default** endpoint — the ``is_default=True`` + row, exposed by ``llm_factory`` under the ``"default"`` alias and + resolved fresh per request so promoting a new default endpoint at + runtime takes effect immediately. This covers a direct/web-only + request (no partition), the cross-partition ``"all"`` sentinel, + partitions that set no preset, and partitions whose presets + conflict (no single owning partition). Same partition semantics as + ``_resolve_chat_history_depth``. + 3. The static ``self._llm`` (built from ``settings.llm`` at startup) + only as a last resort — no endpoint factory is wired (unit tests) + or the catalog has no default endpoint yet. ``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. + unresolvable name here must not fail the chat request; it falls + through to the catalog default with a warning. + + The resolved preset name is always logged (at debug), including for + the default, so "which model answered?" is answerable from the logs. + """ + chat_llm = self._agreed_partition_chat_llm(partition) + if chat_llm is not None: + try: + llm = self._llm_factory(chat_llm) # factory is not None when chat_llm is set + except KeyError: + logger.warning( + "Partition chat_llm preset not found in the model-endpoint catalog — " + "falling back to the default LLM", + chat_llm=chat_llm, + partitions=partition, + ) + else: + logger.bind(chat_llm=chat_llm, partitions=partition).debug( + "Answering with the partition's chat_llm preset" + ) + return llm + return self._default_llm(partition) + + def _agreed_partition_chat_llm(self, partition: list[str] | None) -> str | None: + """The single ``chat_llm`` preset the request's partitions agree on, else None. + + Returns None — meaning "use the catalog default" — when no endpoint + factory is wired, the request has no partition or uses the ``"all"`` + sentinel, no named partition sets a preset, or the named partitions + name more than one preset (a conflict with no single owning partition). """ 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 + return None 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 + return next(iter(names)) if len(names) == 1 else None + + def _default_llm(self, partition: list[str] | None) -> LLM: + """The catalog default LLM endpoint (``is_default=True``), resolved fresh. + + Bypassing this and returning the static ``self._llm`` was the bug + behind "the default chat model is still the one in .env" reports: + promoting a new default endpoint in the catalog had no effect on the + default chat path, which stayed pinned to the ``settings.llm`` (env) + client built at startup. Going through the factory's ``"default"`` + alias — kept in sync with the ``is_default`` row and cache-invalidated + on every default change — makes the promotion take effect. + + Falls back to the static ``self._llm`` only when no factory is wired + (unit tests) or the catalog has no default endpoint yet (KeyError). + """ + if self._llm_factory is not None: + try: + llm = self._llm_factory("default") + except KeyError: + pass + else: + logger.bind(chat_llm=self._default_llm_name(), partitions=partition).debug( + "Answering with the default chat_llm preset" + ) + return llm + logger.bind(partitions=partition).debug("Answering with the static default LLM (no catalog default endpoint)") + return self._llm + + def _default_llm_name(self) -> str: + """Real endpoint name behind the catalog ``"default"`` alias, for logging. + + ``ModelEndpointService.load_all`` stores the ``is_default`` row's + config under both its own name and the ``"default"`` alias (the *same* + object), so the name is recovered by identity. Returns ``"default"`` + when it can't be resolved (e.g. the alias isn't populated yet).""" + llms = self._config.models.llm + default_cfg = llms.get("default") + if default_cfg is not None: + for name, cfg in llms.items(): + if name != "default" and cfg is default_cfg: + return name + return "default" # ------------------------------------------------------------------ # Query generation (was RagPipeline.generate_query — no LangChain) @@ -268,9 +334,11 @@ async def generate_query(self, messages: list[dict], llm: LLM | None = None) -> # ------------------------------------------------------------------ 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. + # Deliberately pinned to the static ``self._llm`` (the settings.llm env + # client) — NOT the resolved catalog default the answer uses, so a + # map-reduce request's sub-calls may run on a different model than its + # answer. Map-reduce is slated for a full post-release refactor; routing + # it through the resolved chat_llm is part of that work. async with get_llm_semaphore(): try: resp = await self._llm.chat( diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index 6094f93f8..62fa1b27d 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -108,6 +108,7 @@ def _config(mode="SimpleRag"): paths=_PROMPT_CFG.paths, prompts=_PROMPT_CFG.prompts, partitions={}, + models=SimpleNamespace(llm={}), ) @@ -216,25 +217,42 @@ def test_resolve_llm_uses_partition_preset(): 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) +def test_resolve_llm_default_paths_use_the_catalog_default_endpoint(): + # No partition preset applies → resolve the catalog default endpoint + # (is_default=True, exposed by the factory's "default" alias), NOT the + # static env-built self._llm. Promoting a new default endpoint in the + # catalog must take effect on the default chat path. + static_llm, catalog_default = FakeLLM(), FakeLLM() + factory = RecordingFactory({"default": catalog_default, "mistral": FakeLLM()}) + svc = _svc(llm=static_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) + assert svc._resolve_llm(None) is catalog_default # direct/web-only mode + assert svc._resolve_llm(["all"]) is catalog_default # cross-partition sentinel + assert svc._resolve_llm(["p"]) is catalog_default # partition without a preset + assert svc._resolve_llm(["missing"]) is catalog_default # unknown partition + assert factory.calls == ["default", "default", "default", "default"] + + +def test_resolve_llm_falls_back_to_static_llm_without_factory_or_catalog_default(): + # Last-resort static self._llm: no factory wired (unit tests), or the + # factory has no "default" alias yet (catalog not seeded). + static_llm = FakeLLM() + no_factory = _svc(llm=static_llm) no_factory._config.partitions = {"q": _partition(chat_llm="mistral")} - assert no_factory._resolve_llm(["q"]) is default_llm + assert no_factory._resolve_llm(["q"]) is static_llm + assert no_factory._resolve_llm(None) is static_llm + + empty_catalog = RecordingFactory() # raises KeyError for "default" + svc = _svc(llm=static_llm, llm_factory=empty_catalog) + svc._config.partitions = {"p": _partition(chat_llm=None)} + assert svc._resolve_llm(["p"]) is static_llm + assert empty_catalog.calls == ["default"] 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) + catalog_default, preset_llm = FakeLLM(), FakeLLM() + factory = RecordingFactory({"default": catalog_default, "mistral": preset_llm}) + svc = _svc(llm=FakeLLM(), llm_factory=factory) svc._config.partitions = { "a": _partition(chat_llm="mistral"), "b": _partition(chat_llm="mistral"), @@ -242,18 +260,36 @@ def test_resolve_llm_multi_partition_uses_preset_only_when_unanimous(): "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 + assert svc._resolve_llm(["a", "d"]) is catalog_default # conflicting presets → catalog default + +def test_default_llm_name_recovers_the_is_default_endpoint_name(): + # load_all() stores the is_default row's config under both its own name + # and the "default" alias (same object) — _default_llm_name recovers the + # real name by identity, so the logs name the endpoint that answered. + svc = _svc(llm_factory=RecordingFactory()) + toy_cfg = SimpleNamespace(endpoint="http://toy-llm") + svc._config.models.llm = {"base-llm": SimpleNamespace(endpoint="http://base"), "toy-llm": toy_cfg} + svc._config.models.llm["default"] = toy_cfg # alias points at the same object + assert svc._default_llm_name() == "toy-llm" -def test_resolve_llm_unknown_preset_falls_back_to_default(): + +def test_default_llm_name_returns_default_when_alias_absent(): + svc = _svc(llm_factory=RecordingFactory()) + svc._config.models.llm = {"base-llm": SimpleNamespace(endpoint="http://base")} + assert svc._default_llm_name() == "default" + + +def test_resolve_llm_unknown_preset_falls_through_to_catalog_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) + # afterwards) — an unknown name must not fail the request; it falls through + # to the catalog default endpoint, not the static env llm. + static_llm, catalog_default = FakeLLM(), FakeLLM() + factory = RecordingFactory({"default": catalog_default}) # "deleted" raises KeyError + svc = _svc(llm=static_llm, llm_factory=factory) svc._config.partitions = {"p": _partition(chat_llm="deleted")} - assert svc._resolve_llm(["p"]) is default_llm - assert factory.calls == ["deleted"] + assert svc._resolve_llm(["p"]) is catalog_default + assert factory.calls == ["deleted", "default"] @pytest.mark.asyncio