From 92e19ac177ef5bf170b8e63ecea5f6e0e4174f1d Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 17 Sep 2026 15:40:20 +0100 Subject: [PATCH] fix(graphrag): dispatch extraction to the provider that serves the model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _build_extraction_llm passed settings.LLM_PROVIDER with the resolved extraction model id. Those two disagree in any deployment that leaves the provider at its default: the model id comes from GRAPHRAG_EXTRACTION_MODEL or LLM_NAME, while the provider stays "docsgpt" — the hosted public endpoint, which does not serve it. The request is rejected, the shared fallback answers instead, and the graph gets built by a different model than the one configured, with nothing in the summary saying so. Resolve the provider from the model registry (owner-scoped, so per-user BYOM ids resolve too), fall back to settings.LLM_PROVIDER only when the model is unknown, and take the API key for the provider actually dispatched to rather than the generic settings.API_KEY. The effective provider is logged, since a silent swap was the whole failure mode. --- docsgpt/graphrag/extraction.py | 31 +++++++- tests/graphrag/test_extraction.py | 118 ++++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 2 deletions(-) diff --git a/docsgpt/graphrag/extraction.py b/docsgpt/graphrag/extraction.py index beaf3d66..aa76189f 100644 --- a/docsgpt/graphrag/extraction.py +++ b/docsgpt/graphrag/extraction.py @@ -23,6 +23,10 @@ import logging import re from typing import Any, Callable, Dict, List, Optional +from docsgpt.core.model_utils import ( + get_api_key_for_provider, + get_provider_from_model_id, +) from docsgpt.core.settings import settings from docsgpt.llm.llm_creator import LLMCreator from docsgpt.storage.db.source_config import SourceConfig @@ -63,14 +67,37 @@ def _resolve_max_chunks(config: SourceConfig) -> int: return config.graph.max_chunks or settings.GRAPHRAG_MAX_CHUNKS_FOR_EXTRACTION +def _resolve_extraction_provider( + model_id: Optional[str], user: Optional[str] +) -> str: + """The provider that serves ``model_id``, else the deployment default. + + ``settings.LLM_PROVIDER`` is only a default (``docsgpt``, the hosted public + endpoint, out of the box). Dispatching the resolved extraction model + through it sends the request to a provider that does not serve that model: + the call is rejected, the shared fallback answers instead, and the graph is + built by a different model than the one configured — with nothing in the + summary to say so. ``user`` scopes the lookup so a per-user (BYOM) model id + resolves as well. + """ + provider = ( + get_provider_from_model_id(model_id, user_id=user) if model_id else None + ) + return provider or settings.LLM_PROVIDER + + def _build_extraction_llm( model_id: Optional[str], user: Optional[str], request_id: Optional[str] ): """Build the extraction LLM tagged for token-usage attribution to the owner.""" decoded_token = {"sub": user} if user else None + provider = _resolve_extraction_provider(model_id, user) + logger.info( + "Graph extraction dispatching model=%s via provider=%s", model_id, provider + ) llm = LLMCreator.create_llm( - settings.LLM_PROVIDER, - api_key=settings.API_KEY, + provider, + api_key=get_api_key_for_provider(provider), user_api_key=None, decoded_token=decoded_token, model_id=model_id, diff --git a/tests/graphrag/test_extraction.py b/tests/graphrag/test_extraction.py index ad2d670f..eaaf881d 100644 --- a/tests/graphrag/test_extraction.py +++ b/tests/graphrag/test_extraction.py @@ -403,6 +403,124 @@ class TestModelResolution: assert extraction_module._resolve_max_chunks(config) == 5 +@pytest.mark.unit +class TestExtractionProviderResolution: + """The extraction model decides the provider, not ``LLM_PROVIDER``. + + ``settings.LLM_PROVIDER`` is the deployment default (``docsgpt`` out of the + box, i.e. the hosted public endpoint). Dispatching the resolved extraction + model through it sends the call to a provider that never serves that model: + the request is rejected, the shared fallback answers instead, and the graph + is quietly built by a different model than the one configured. + """ + + def _capture_create_llm(self, monkeypatch, llm=None): + captured = {} + + def _create(provider, *args, **kwargs): + captured["provider"] = provider + captured["args"] = args + captured["kwargs"] = kwargs + return llm or _StubLLM([]) + + monkeypatch.setattr( + extraction_module.LLMCreator, "create_llm", staticmethod(_create) + ) + return captured + + def test_provider_comes_from_the_model_registry(self, monkeypatch): + monkeypatch.setattr(extraction_module.settings, "LLM_PROVIDER", "docsgpt") + monkeypatch.setattr( + extraction_module, "get_provider_from_model_id", lambda *a, **k: "openai" + ) + monkeypatch.setattr( + extraction_module, "get_api_key_for_provider", lambda provider: "sk-openai" + ) + captured = self._capture_create_llm(monkeypatch) + + extraction_module._build_extraction_llm("gpt-4o-mini", "owner-1", "req-1") + + assert captured["provider"] == "openai" + assert captured["kwargs"]["api_key"] == "sk-openai" + assert captured["kwargs"]["model_id"] == "gpt-4o-mini" + + def test_owner_scopes_the_registry_lookup(self, monkeypatch): + """A per-user (BYOM) model only resolves when the owner is passed.""" + seen = {} + + def _resolve(model_id, user_id=None): + seen["model_id"] = model_id + seen["user_id"] = user_id + return "anthropic" + + monkeypatch.setattr( + extraction_module, "get_provider_from_model_id", _resolve + ) + monkeypatch.setattr( + extraction_module, "get_api_key_for_provider", lambda provider: "k" + ) + self._capture_create_llm(monkeypatch) + + extraction_module._build_extraction_llm("byom-uuid", "owner-7", "req-1") + + assert seen == {"model_id": "byom-uuid", "user_id": "owner-7"} + + def test_unknown_model_falls_back_to_the_configured_provider(self, monkeypatch): + monkeypatch.setattr(extraction_module.settings, "LLM_PROVIDER", "docsgpt") + monkeypatch.setattr( + extraction_module, "get_provider_from_model_id", lambda *a, **k: None + ) + monkeypatch.setattr( + extraction_module, "get_api_key_for_provider", lambda provider: "fallback-key" + ) + captured = self._capture_create_llm(monkeypatch) + + extraction_module._build_extraction_llm("mystery-model", "owner-1", "req-1") + + assert captured["provider"] == "docsgpt" + assert captured["kwargs"]["api_key"] == "fallback-key" + + def test_no_model_id_skips_the_lookup(self, monkeypatch): + monkeypatch.setattr(extraction_module.settings, "LLM_PROVIDER", "openai") + calls = [] + monkeypatch.setattr( + extraction_module, + "get_provider_from_model_id", + lambda *a, **k: calls.append(a) or "anthropic", + ) + monkeypatch.setattr( + extraction_module, "get_api_key_for_provider", lambda provider: "k" + ) + captured = self._capture_create_llm(monkeypatch) + + extraction_module._build_extraction_llm(None, "owner-1", "req-1") + + assert calls == [] + assert captured["provider"] == "openai" + + def test_api_key_follows_the_resolved_provider(self, monkeypatch): + """The key must match the provider actually dispatched to.""" + monkeypatch.setattr(extraction_module.settings, "LLM_PROVIDER", "docsgpt") + monkeypatch.setattr(extraction_module.settings, "API_KEY", "generic-key") + monkeypatch.setattr( + extraction_module, "get_provider_from_model_id", lambda *a, **k: "anthropic" + ) + keyed_for = {} + + def _key(provider): + keyed_for["provider"] = provider + return "sk-anthropic" + + monkeypatch.setattr(extraction_module, "get_api_key_for_provider", _key) + captured = self._capture_create_llm(monkeypatch) + + extraction_module._build_extraction_llm("claude-x", "owner-1", "req-1") + + assert keyed_for["provider"] == "anthropic" + assert captured["kwargs"]["api_key"] == "sk-anthropic" + assert captured["kwargs"]["api_key"] != "generic-key" + + @pytest.mark.unit class TestParsing: def test_parses_embedded_json(self):