mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 12:14:00 +00:00
fix(workflows): dispatch agent nodes by provider, not display label
/api/models reports `display_provider` when a catalog YAML sets one (`foundry`, `azure_foundry`, `cloudflare`), and the builder persists that string as a node's `llm_name`. The engine passed it straight to LLMCreator, which only knows the names in PROVIDERS_BY_NAME and raises `No LLM class found for type <label>`. That fails the node before any LLM call, so the turn ends in ~60ms with an empty answer and no tokens generated. It hits the *default* path: the platform-default model's label is stamped into every newly dragged agent node and validateWorkflow requires an agent node, so a new user's untouched workflow could not produce a token regardless of what they typed. Ordinary chat was unaffected because it resolves the provider from the model registry and never reads the stored name. Adds `resolve_dispatch_provider`, which prefers a stored name that is a real dispatch provider, then the registry lookup, then the parent agent. Nodes already saved with a label are repaired at run time, so no migration is needed. The api_key now resolves from the normalized name too — `get_api_key_for_provider` falls back to settings.API_KEY for names it does not recognize, which would have sent the deployment key to whatever endpoint the label happened to select.
This commit is contained in:
1 parent
fcef3e68e5
commit
f5129a5656
3 files changed
+155
-7
No files matched your search
@@ -291,7 +291,7 @@ class WorkflowEngine:
|
||||
from application.core.model_utils import (
|
||||
get_api_key_for_provider,
|
||||
get_model_capabilities,
|
||||
get_provider_from_model_id,
|
||||
resolve_dispatch_provider,
|
||||
)
|
||||
|
||||
node_config = AgentNodeConfig(**node.config.get("config", node.config))
|
||||
@@ -314,13 +314,22 @@ class WorkflowEngine:
|
||||
if isinstance(self.agent.decoded_token, dict)
|
||||
else None
|
||||
)
|
||||
node_llm_name = (
|
||||
node_config.llm_name
|
||||
or get_provider_from_model_id(
|
||||
node_model_id or "", user_id=node_user_id
|
||||
)
|
||||
or self.agent.llm_name
|
||||
# ``node_config.llm_name`` is whatever the builder stored, which for
|
||||
# catalogs with a ``display_provider`` is a presentation label
|
||||
# ("foundry", "azure_foundry", "cloudflare") that no LLM class is
|
||||
# registered under. Trusting it verbatim fails the node before any LLM
|
||||
# call and hands the user a blank answer, so resolve it to a real
|
||||
# dispatch provider first — this also repairs nodes already saved with
|
||||
# a label, without a migration.
|
||||
node_llm_name = resolve_dispatch_provider(
|
||||
node_config.llm_name,
|
||||
node_model_id,
|
||||
user_id=node_user_id,
|
||||
fallback=self.agent.llm_name,
|
||||
)
|
||||
# Resolve the key from the normalized name: get_api_key_for_provider
|
||||
# silently returns settings.API_KEY for names it does not know, so a
|
||||
# label here would leak the deployment key to the wrong endpoint.
|
||||
node_api_key = get_api_key_for_provider(node_llm_name) or self.agent.api_key
|
||||
|
||||
# Structured output gates on the model's registry capability flags;
|
||||
|
||||
@@ -20,6 +20,52 @@ def get_api_key_for_provider(provider: str) -> Optional[str]:
|
||||
return settings.API_KEY
|
||||
|
||||
|
||||
def resolve_dispatch_provider(
|
||||
stored_llm_name: Optional[str],
|
||||
model_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
fallback: Optional[str] = None,
|
||||
) -> Optional[str]:
|
||||
"""Resolve a stored provider string to one ``LLMCreator`` can dispatch.
|
||||
|
||||
``/api/models`` reports ``display_provider`` when a catalog YAML sets one
|
||||
(``foundry``, ``azure_foundry``, ``cloudflare``, …), and clients persist
|
||||
that label as ``llm_name`` on agents and workflow nodes. Those labels are
|
||||
presentation-only: they are absent from ``PROVIDERS_BY_NAME``, so passing
|
||||
one to ``LLMCreator.create_llm`` raises ``No LLM class found for type
|
||||
<label>``. Normalizing here keeps already-stored records working without a
|
||||
migration.
|
||||
|
||||
Resolution order: a stored name that is a real dispatch provider wins;
|
||||
otherwise the model registry decides; otherwise ``fallback``.
|
||||
|
||||
Args:
|
||||
stored_llm_name: Provider string persisted on the record, possibly a
|
||||
display label.
|
||||
model_id: Model whose registry entry knows the true provider.
|
||||
user_id: BYOM-resolution scope for per-user model records.
|
||||
fallback: Used when neither the stored name nor the registry resolves.
|
||||
|
||||
Returns:
|
||||
A dispatchable provider name, or ``fallback`` when nothing resolves.
|
||||
"""
|
||||
from application.llm.providers import PROVIDERS_BY_NAME
|
||||
|
||||
# Return the *canonical* lowercase name. ``LLMCreator`` lowercases before
|
||||
# its own lookup, but ``get_api_key_for_provider`` matches exactly — so a
|
||||
# stored "OpenAI" would pass this guard and then silently fall through to
|
||||
# ``settings.API_KEY``, which is the key leak this function exists to stop.
|
||||
if stored_llm_name and stored_llm_name.lower() in PROVIDERS_BY_NAME:
|
||||
return stored_llm_name.lower()
|
||||
if model_id:
|
||||
resolved = get_provider_from_model_id(model_id, user_id=user_id)
|
||||
if resolved:
|
||||
return resolved
|
||||
if fallback and fallback.lower() in PROVIDERS_BY_NAME:
|
||||
return fallback.lower()
|
||||
return fallback
|
||||
|
||||
|
||||
def get_all_available_models(
|
||||
user_id: Optional[str] = None,
|
||||
) -> Dict[str, Dict[str, Any]]:
|
||||
|
||||
@@ -507,3 +507,96 @@ class TestWorkflowEngineAdditionalCoverage:
|
||||
# A non-dict schema triggers JsonSchemaValidationError
|
||||
with pytest.raises(ValueError, match="Invalid JSON schema"):
|
||||
engine._normalize_node_json_schema("not_a_dict", "TestNode")
|
||||
|
||||
|
||||
class TestAgentNodeProviderResolution:
|
||||
"""``llm_name`` stored on a node is a *display* label, not a dispatch name.
|
||||
|
||||
``/api/models`` reports ``display_provider`` (e.g. ``foundry``,
|
||||
``azure_foundry``, ``cloudflare``) and the builder stores that string on
|
||||
the node. Handing it to ``LLMCreator`` raises ``No LLM class found for
|
||||
type <label>``, which fails the node before any LLM call and returns a
|
||||
blank answer to the user. The engine must resolve the real dispatch
|
||||
provider from the model registry instead.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _run(monkeypatch, node, *, registry_provider="openai_compatible"):
|
||||
"""Execute one agent node, returning the kwargs the factory saw."""
|
||||
engine = create_engine()
|
||||
engine.state["query"] = "test"
|
||||
captured = {}
|
||||
|
||||
def _capture(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return StubNodeAgent([{"answer": "ok"}])
|
||||
|
||||
monkeypatch.setattr(
|
||||
WorkflowNodeAgentFactory, "create", staticmethod(_capture)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_api_key_for_provider",
|
||||
lambda name: f"key-for-{name}",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_provider_from_model_id",
|
||||
lambda _, **_kwargs: registry_provider,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"application.core.model_utils.get_model_capabilities",
|
||||
lambda _, **_kwargs: None,
|
||||
)
|
||||
list(engine._execute_agent_node(node))
|
||||
return captured
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"display_label", ["azure_foundry", "cloudflare", "foundry"]
|
||||
)
|
||||
def test_display_provider_label_resolves_to_dispatch_provider(
|
||||
self, monkeypatch, display_label
|
||||
):
|
||||
"""A display label must not reach LLMCreator (err#33)."""
|
||||
node = create_agent_node(node_id="n1")
|
||||
node.config["model_id"] = "Kimi-K2.6"
|
||||
node.config["llm_name"] = display_label
|
||||
|
||||
captured = self._run(monkeypatch, node)
|
||||
|
||||
assert captured["llm_name"] == "openai_compatible"
|
||||
# The api_key must follow the *normalized* name: resolving against the
|
||||
# display label silently falls through to settings.API_KEY.
|
||||
assert captured["api_key"] == "key-for-openai_compatible"
|
||||
|
||||
def test_real_provider_name_is_preserved(self, monkeypatch):
|
||||
"""A node storing a genuine dispatch name keeps it."""
|
||||
node = create_agent_node(node_id="n2")
|
||||
node.config["model_id"] = "gpt-4o"
|
||||
node.config["llm_name"] = "openai"
|
||||
|
||||
captured = self._run(monkeypatch, node, registry_provider="openai")
|
||||
|
||||
assert captured["llm_name"] == "openai"
|
||||
|
||||
def test_unresolvable_label_falls_back_to_parent_agent(self, monkeypatch):
|
||||
"""No registry hit: inherit the parent agent rather than dispatching junk."""
|
||||
node = create_agent_node(node_id="n3")
|
||||
node.config["model_id"] = "mystery-model"
|
||||
node.config["llm_name"] = "some_unknown_label"
|
||||
|
||||
captured = self._run(monkeypatch, node, registry_provider=None)
|
||||
|
||||
# create_engine()'s parent agent is llm_name="openai"
|
||||
assert captured["llm_name"] == "openai"
|
||||
|
||||
def test_retriever_config_uses_normalized_provider(self, monkeypatch):
|
||||
"""The agentic retriever path builds its own kwargs — normalize there too."""
|
||||
node = create_agent_node(node_id="n4")
|
||||
node.config["model_id"] = "Kimi-K2.6"
|
||||
node.config["llm_name"] = "azure_foundry"
|
||||
node.config["agent_type"] = "agentic"
|
||||
node.config["sources"] = ["src-1"]
|
||||
|
||||
captured = self._run(monkeypatch, node)
|
||||
|
||||
assert captured["retriever_config"]["llm_name"] == "openai_compatible"
|
||||
assert captured["retriever_config"]["api_key"] == "key-for-openai_compatible"
|
||||
Reference in new issue
Block a user