feat(api): LLM polling support (#37462)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: WH-2099 <wh2099@pm.me>
This commit is contained in:
QuantumGhost
2026-06-17 23:34:33 +00:00
committed by GitHub
co-authored by autofix-ci[bot] WH-2099
parent 19838972dc
commit f0b34bdeb4
17 changed files with 704 additions and 46 deletions
@@ -7,6 +7,7 @@ from pytest_mock import MockerFixture
from core.plugin.endpoint.exc import EndpointSetupFailedError
from core.plugin.entities.plugin_daemon import PluginDaemonInnerError
from core.plugin.impl.base import PLUGIN_DAEMON_MAX_PATH_LENGTH, BasePluginClient
from core.plugin.impl.exc import PluginLLMPollingUnsupportedError
from core.trigger.errors import (
EventIgnoreError,
TriggerInvokeError,
@@ -167,3 +168,10 @@ class TestBasePluginClientImpl:
with pytest.raises(expected):
client._handle_plugin_daemon_error("PluginInvokeError", message)
def test_handle_plugin_daemon_error_maps_unsupported_polling_to_typed_exception(self):
client = BasePluginClient()
message = json.dumps({"error_type": PluginLLMPollingUnsupportedError.__name__, "message": "m"})
with pytest.raises(PluginLLMPollingUnsupportedError):
client._handle_plugin_daemon_error("PluginInvokeError", message)
@@ -1,13 +1,17 @@
from __future__ import annotations
import io
import json
from types import SimpleNamespace
import pytest
from pytest_mock import MockerFixture
from core.plugin.entities.plugin_daemon import PluginDaemonInnerError
from core.plugin.impl.exc import PluginInvokeError, PluginLLMPollingUnsupportedError
from core.plugin.impl.model import PluginModelClient
from graphon.model_runtime.entities.llm_entities import LLMPollingResult, LLMPollingStatus, LLMResult, LLMUsage
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage
class TestPluginModelClient:
@@ -183,6 +187,113 @@ class TestPluginModelClient:
)
)
def test_start_llm_polling(self, mocker: MockerFixture):
client = PluginModelClient()
polling_result = LLMPollingResult(
status=LLMPollingStatus.RUNNING,
plugin_state={"task_id": "poll-1"},
next_check_after_seconds=3,
)
request_mock = mocker.patch.object(
client,
"_request_with_plugin_daemon_response",
return_value=polling_result,
)
result = client.start_llm_polling(
tenant_id="tenant-1",
user_id="user-1",
plugin_id="org/plugin:1",
provider="provider-a",
model="gpt-test",
credentials={"api_key": "key"},
prompt_messages=[],
model_parameters={"temperature": 0.1},
tools=[],
stop=["STOP"],
json_schema={"type": "object"},
)
assert result == polling_result
call_kwargs = request_mock.call_args.kwargs
assert call_kwargs["path"] == "plugin/tenant-1/dispatch/model/polling/start"
assert call_kwargs["data"]["data"] == {
"provider": "provider-a",
"model_type": "llm",
"model": "gpt-test",
"credentials": {"api_key": "key"},
"prompt_messages": [],
"model_parameters": {"temperature": 0.1},
"tools": [],
"stop": ["STOP"],
"stream": False,
"json_schema": {"type": "object"},
}
def test_check_llm_polling(self, mocker: MockerFixture):
client = PluginModelClient()
polling_result = LLMPollingResult(
status=LLMPollingStatus.SUCCEEDED,
result=LLMResult(
model="gpt-test",
prompt_messages=[],
message=AssistantPromptMessage(content="done"),
usage=LLMUsage.empty_usage(),
),
)
request_mock = mocker.patch.object(
client,
"_request_with_plugin_daemon_response",
return_value=polling_result,
)
result = client.check_llm_polling(
tenant_id="tenant-1",
user_id="user-1",
plugin_id="org/plugin:1",
provider="provider-a",
model="gpt-test",
credentials={"api_key": "key"},
plugin_state={"task_id": "poll-1"},
)
assert result == polling_result
call_kwargs = request_mock.call_args.kwargs
assert call_kwargs["path"] == "plugin/tenant-1/dispatch/model/polling/check"
assert call_kwargs["data"]["data"] == {
"provider": "provider-a",
"model_type": "llm",
"model": "gpt-test",
"credentials": {"api_key": "key"},
"plugin_state": {"task_id": "poll-1"},
}
def test_start_llm_polling_maps_unsupported_polling_invoke_error(self, mocker: MockerFixture):
client = PluginModelClient()
mocker.patch.object(
client,
"_request_with_plugin_daemon_response",
side_effect=PluginInvokeError(
json.dumps(
{
"error_type": PluginLLMPollingUnsupportedError.__name__,
"message": "Model `gpt-test` does not support polling.",
}
)
),
)
with pytest.raises(PluginLLMPollingUnsupportedError):
client.start_llm_polling(
tenant_id="tenant-1",
user_id="user-1",
plugin_id="org/plugin:1",
provider="provider-a",
model="gpt-test",
credentials={"api_key": "key"},
prompt_messages=[],
)
def test_get_llm_num_tokens(self, mocker: MockerFixture):
client = PluginModelClient()
mocker.patch.object(
@@ -14,7 +14,13 @@ from core.plugin.impl.model_runtime import TENANT_SCOPE_SCHEMA_CACHE_USER_ID, Pl
from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime
from core.plugin.plugin_service import PluginService
from graphon.model_runtime.entities.common_entities import I18nObject
from graphon.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta, LLMUsage
from graphon.model_runtime.entities.llm_entities import (
LLMPollingResult,
LLMPollingStatus,
LLMResultChunk,
LLMResultChunkDelta,
LLMUsage,
)
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage
from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
from graphon.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity
@@ -282,6 +288,74 @@ class TestPluginModelRuntime:
stream=True,
)
def test_start_llm_polling_resolves_plugin_fields(self) -> None:
client = Mock(spec=PluginModelClient)
polling_result = LLMPollingResult(
status=LLMPollingStatus.RUNNING,
plugin_state={"task_id": "poll-1"},
next_check_after_seconds=2,
)
client.start_llm_polling.return_value = polling_result
runtime = PluginModelRuntime(tenant_id="tenant", user_id="user", client=client, plugin_service=PluginService)
result = runtime.start_llm_polling(
provider="langgenius/openai/openai",
model="gpt-4o-mini",
credentials={"api_key": "secret"},
model_parameters={"temperature": 0.2},
prompt_messages=[],
tools=None,
stop=("END",),
json_schema={"type": "object"},
)
assert result == polling_result
client.start_llm_polling.assert_called_once_with(
tenant_id="tenant",
user_id="user",
plugin_id="langgenius/openai",
provider="openai",
model="gpt-4o-mini",
credentials={"api_key": "secret"},
prompt_messages=[],
model_parameters={"temperature": 0.2},
tools=None,
stop=["END"],
json_schema={"type": "object"},
)
def test_check_llm_polling_resolves_plugin_fields(self) -> None:
client = Mock(spec=PluginModelClient)
polling_result = LLMPollingResult(
status=LLMPollingStatus.SUCCEEDED,
result=model_runtime_module.LLMResult(
model="gpt-4o-mini",
prompt_messages=[],
message=AssistantPromptMessage(content="done"),
usage=LLMUsage.empty_usage(),
),
)
client.check_llm_polling.return_value = polling_result
runtime = PluginModelRuntime(tenant_id="tenant", user_id="user", client=client, plugin_service=PluginService)
result = runtime.check_llm_polling(
provider="langgenius/openai/openai",
model="gpt-4o-mini",
credentials={"api_key": "secret"},
plugin_state={"task_id": "poll-1"},
)
assert result == polling_result
client.check_llm_polling.assert_called_once_with(
tenant_id="tenant",
user_id="user",
plugin_id="langgenius/openai",
provider="openai",
model="gpt-4o-mini",
credentials={"api_key": "secret"},
plugin_state={"task_id": "poll-1"},
)
def test_invoke_llm_rejects_per_call_user_override(self) -> None:
client = Mock(spec=PluginModelClient)
client.invoke_llm.return_value = sentinel.result
@@ -703,9 +703,12 @@ class TestSchemaResolverClass:
# For schemas without refs, hybrid should be competitive or better
if not expected: # No refs case
# Hybrid might be slightly slower due to JSON serialization overhead,
# but should not be dramatically worse
assert avg_hybrid < avg_recursive * 5 # At most 5x slower
relative_slowdown_limit = 5.0
absolute_noise_budget_seconds = 2e-4
# JSON serialization has a fixed overhead that dominates tiny schemas,
# so allow a small absolute noise budget on top of the relative limit.
assert avg_hybrid < (avg_recursive * relative_slowdown_limit) + absolute_noise_budget_seconds
def test_string_matching_edge_cases(self):
"""Test edge cases for string-based detection"""
@@ -1,18 +1,26 @@
from collections.abc import Mapping
from types import SimpleNamespace
from unittest.mock import MagicMock, patch, sentinel
from unittest.mock import MagicMock, Mock, patch, sentinel
import pytest
from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext, InvokeFrom, UserFrom
from core.plugin.impl.model import PluginModelClient
from core.plugin.impl.model_runtime import PluginModelRuntime
from core.plugin.plugin_service import PluginService
from core.workflow import node_factory
from core.workflow import template_rendering as workflow_template_rendering
from core.workflow.node_runtime import DifyPreparedLLM
from core.workflow.nodes.knowledge_index import KNOWLEDGE_INDEX_NODE_TYPE
from graphon.entities.base_node_data import BaseNodeData
from graphon.enums import BuiltinNodeTypes, NodeType
from graphon.model_runtime.entities.common_entities import I18nObject
from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelFeature, ModelType
from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
from graphon.nodes.code.entities import CodeLanguage
from graphon.nodes.llm.entities import LLMNodeData
from graphon.nodes.llm.node import LLMNode
from graphon.nodes.llm.runtime_protocols import LLMPollingCapableProtocol
from graphon.nodes.parameter_extractor.entities import ParameterExtractorNodeData
from graphon.variables.segments import ArrayObjectSegment, StringSegment
@@ -35,6 +43,41 @@ def _node_constructor(*, return_value):
return constructor
def _build_llm_model_schema(*, features: list[ModelFeature] | None = None) -> AIModelEntity:
return AIModelEntity(
model="model",
label=I18nObject(en_US="Model"),
model_type=ModelType.LLM,
fetch_from=FetchFrom.PREDEFINED_MODEL,
model_properties={},
features=features,
)
class _ModelTypeInstanceStub(LargeLanguageModel):
def __init__(self, *, model_runtime: object) -> None:
self.model_runtime = model_runtime
class _ModelInstanceStub:
def __init__(
self,
*,
model_runtime: object,
model_schema: AIModelEntity,
) -> None:
self.provider = "langgenius/openai/openai"
self.model_name = "model"
self.credentials = {"api_key": "secret"}
self.parameters = {}
self.stop = ()
self.model_type_instance = _ModelTypeInstanceStub(model_runtime=model_runtime)
self._model_schema = model_schema
def get_model_schema(self) -> AIModelEntity:
return self._model_schema
class TestResolveWorkflowNodeClass:
def test_matching_version_uses_registry_mapping(self, monkeypatch) -> None:
document_extractor_class = sentinel.document_extractor_class
@@ -667,7 +710,7 @@ class TestDifyNodeFactoryCreateNode:
memory = sentinel.memory
factory._build_model_instance_for_llm_node = MagicMock(return_value=sentinel.model_instance)
factory._build_memory_for_llm_node = MagicMock(return_value=memory)
with patch.object(node_factory, "DifyPreparedLLM", return_value=wrapped_model_instance) as prepared_llm:
with patch.object(factory, "_wrap_model_instance_for_node", return_value=wrapped_model_instance) as wrap_model:
kwargs = factory._build_llm_compatible_node_init_kwargs(
node_class=sentinel.node_class,
node_data=node_data,
@@ -686,9 +729,70 @@ class TestDifyNodeFactoryCreateNode:
node_data=node_data,
model_instance=sentinel.model_instance,
)
prepared_llm.assert_called_once_with(sentinel.model_instance)
wrap_model.assert_called_once_with(
node_data=node_data,
model_instance=sentinel.model_instance,
)
assert kwargs["model_instance"] is wrapped_model_instance
def test_build_llm_compatible_node_init_kwargs_uses_polling_wrapper_for_polling_llm_node(self, factory):
node_data = LLMNodeData.model_validate(
{
"type": BuiltinNodeTypes.LLM,
"title": "LLM",
"model": {"provider": "provider", "name": "model", "mode": "chat", "completion_params": {}},
"prompt_template": [{"role": "system", "text": "x"}],
"context": {"enabled": False, "variable_selector": []},
"vision": {"enabled": False},
}
)
plugin_runtime = PluginModelRuntime(
tenant_id="tenant-id",
user_id="user-id",
client=Mock(spec=PluginModelClient),
plugin_service=PluginService,
)
model_instance = _ModelInstanceStub(
model_runtime=plugin_runtime,
model_schema=_build_llm_model_schema(features=[ModelFeature.POLLING]),
)
factory._build_model_instance_for_llm_node = MagicMock(return_value=model_instance)
factory._build_memory_for_llm_node = MagicMock(return_value=sentinel.memory)
kwargs = factory._build_llm_compatible_node_init_kwargs(
node_class=sentinel.node_class,
node_data=node_data,
wrap_model_instance=True,
include_http_client=False,
include_llm_file_saver=False,
include_prompt_message_serializer=False,
include_retriever_attachment_loader=False,
include_jinja2_template_renderer=False,
)
assert isinstance(kwargs["model_instance"], LLMPollingCapableProtocol)
@pytest.mark.parametrize("node_type", [BuiltinNodeTypes.QUESTION_CLASSIFIER, BuiltinNodeTypes.PARAMETER_EXTRACTOR])
def test_wrap_model_instance_keeps_non_llm_graph_nodes_on_plain_wrapper(self, node_type):
plugin_runtime = PluginModelRuntime(
tenant_id="tenant-id",
user_id="user-id",
client=Mock(spec=PluginModelClient),
plugin_service=PluginService,
)
model_instance = _ModelInstanceStub(
model_runtime=plugin_runtime,
model_schema=_build_llm_model_schema(features=[ModelFeature.POLLING]),
)
wrapped = node_factory.DifyNodeFactory._wrap_model_instance_for_node(
node_data=SimpleNamespace(type=node_type),
model_instance=model_instance,
)
assert type(wrapped) is DifyPreparedLLM
assert not isinstance(wrapped, LLMPollingCapableProtocol)
def test_create_node_passes_alias_preserving_llm_data_to_constructor(self, monkeypatch, factory):
created_node = object()
constructor = _node_constructor(return_value=created_node)
@@ -7,6 +7,10 @@ import pytest
from core.app.entities.app_invoke_entities import DIFY_RUN_CONTEXT_KEY, DifyRunContext, InvokeFrom, UserFrom
from core.app.file_access import FileAccessScope, bind_file_access_scope, grant_retriever_segment_access
from core.llm_generator.output_parser.errors import OutputParserError
from core.plugin.impl.exc import PluginLLMPollingUnsupportedError
from core.plugin.impl.model import PluginModelClient
from core.plugin.impl.model_runtime import PluginModelRuntime
from core.plugin.plugin_service import PluginService
from core.workflow import node_runtime
from core.workflow.file_reference import parse_file_reference
from core.workflow.human_input_adapter import (
@@ -21,6 +25,7 @@ from core.workflow.node_runtime import (
DifyFileReferenceFactory,
DifyHumanInputNodeRuntime,
DifyPreparedLLM,
DifyPreparedPollingLLM,
DifyPromptMessageSerializer,
DifyRetrieverAttachmentLoader,
DifyToolFileManager,
@@ -31,23 +36,61 @@ from core.workflow.node_runtime import (
)
from graphon.file import File, FileTransferMethod, FileType
from graphon.model_runtime.entities.common_entities import I18nObject
from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
from graphon.model_runtime.entities.llm_entities import LLMPollingResult, LLMPollingStatus
from graphon.model_runtime.entities.message_entities import AssistantPromptMessage
from graphon.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelFeature, ModelType
from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
from graphon.nodes.human_input.entities import FileInputConfig, FileListInputConfig, HumanInputNodeData
from graphon.nodes.llm.runtime_protocols import LLMPollingCapableProtocol
from graphon.nodes.tool.entities import ToolNodeData, ToolProviderType
from graphon.variables.segments import ArrayFileSegment, FileSegment
from tests.workflow_test_utils import build_test_run_context
def _build_model_schema() -> AIModelEntity:
def _build_model_schema(*, features: list[ModelFeature] | None = None) -> AIModelEntity:
return AIModelEntity(
model="gpt-4o-mini",
label=I18nObject(en_US="GPT-4o mini"),
model_type=ModelType.LLM,
fetch_from=FetchFrom.PREDEFINED_MODEL,
model_properties={},
features=features,
)
class _ModelTypeInstanceStub(LargeLanguageModel):
def __init__(
self,
*,
model_schema: AIModelEntity | None,
model_runtime: object | None = None,
) -> None:
self.model_runtime = model_runtime
self.get_model_schema = Mock(return_value=model_schema)
class _ModelInstanceStub:
def __init__(
self,
*,
model_schema: AIModelEntity | None,
model_runtime: object | None = None,
invoke_llm_result: object = sentinel.result,
get_llm_num_tokens_result: int = 32,
) -> None:
self.provider = "langgenius/openai/openai"
self.model_name = "gpt-4o-mini"
self.parameters = {"temperature": 0.2}
self.stop = ("stop",)
self.credentials = {"api_key": "secret"}
self.model_type_instance = _ModelTypeInstanceStub(
model_schema=model_schema,
model_runtime=model_runtime,
)
self.get_llm_num_tokens = Mock(return_value=get_llm_num_tokens_result)
self.invoke_llm = Mock(return_value=invoke_llm_result)
def _build_run_context(*, invoke_from: InvokeFrom | str = InvokeFrom.DEBUGGER) -> dict[str, object]:
return build_test_run_context(
tenant_id="tenant-id",
@@ -126,17 +169,8 @@ def test_dify_file_reference_factory_passes_tenant_id(monkeypatch: pytest.Monkey
def test_dify_prepared_llm_wraps_model_instance_calls() -> None:
model_schema = _build_model_schema()
model_type_instance = SimpleNamespace(get_model_schema=Mock(return_value=model_schema))
model_instance = SimpleNamespace(
provider="langgenius/openai/openai",
model_name="gpt-4o-mini",
parameters={"temperature": 0.2},
stop=("stop",),
credentials={"api_key": "secret"},
model_type_instance=model_type_instance,
get_llm_num_tokens=Mock(return_value=32),
invoke_llm=Mock(return_value=sentinel.result),
)
model_instance = _ModelInstanceStub(model_schema=model_schema)
model_type_instance = model_instance.model_type_instance
prepared = DifyPreparedLLM(model_instance)
assert prepared.provider == "langgenius/openai/openai"
@@ -167,11 +201,8 @@ def test_dify_prepared_llm_wraps_model_instance_calls() -> None:
def test_dify_prepared_llm_requires_model_schema() -> None:
model_instance = SimpleNamespace(
model_name="gpt-4o-mini",
credentials={},
model_type_instance=SimpleNamespace(get_model_schema=Mock(return_value=None)),
)
model_instance = _ModelInstanceStub(model_schema=None)
model_instance.credentials = {}
prepared = DifyPreparedLLM(model_instance)
with pytest.raises(ValueError, match="Model schema not found"):
@@ -179,12 +210,7 @@ def test_dify_prepared_llm_requires_model_schema() -> None:
def test_dify_prepared_llm_delegates_structured_output_helper(monkeypatch: pytest.MonkeyPatch) -> None:
model_instance = SimpleNamespace(
provider="langgenius/openai/openai",
model_name="gpt-4o-mini",
credentials={"api_key": "secret"},
model_type_instance=SimpleNamespace(get_model_schema=Mock(return_value=_build_model_schema())),
)
model_instance = _ModelInstanceStub(model_schema=_build_model_schema())
prepared = DifyPreparedLLM(model_instance)
invoke_structured = MagicMock(return_value=sentinel.structured)
monkeypatch.setattr(node_runtime, "invoke_llm_with_structured_output", invoke_structured)
@@ -217,6 +243,94 @@ def test_dify_prepared_llm_identifies_structured_output_errors() -> None:
assert prepared.is_structured_output_parse_error(ValueError("other")) is False
def test_dify_prepared_polling_llm_delegates_to_plugin_runtime() -> None:
polling_result = LLMPollingResult(
status=LLMPollingStatus.RUNNING,
plugin_state={"task_id": "poll-1"},
next_check_after_seconds=2,
)
plugin_runtime = PluginModelRuntime(
tenant_id="tenant-id",
user_id="user-id",
client=Mock(spec=PluginModelClient),
plugin_service=PluginService,
)
plugin_runtime.start_llm_polling = Mock(return_value=polling_result) # type: ignore[method-assign]
plugin_runtime.check_llm_polling = Mock(return_value=polling_result) # type: ignore[method-assign]
model_instance = _ModelInstanceStub(
model_schema=_build_model_schema(features=[ModelFeature.POLLING]),
model_runtime=plugin_runtime,
)
prepared = DifyPreparedPollingLLM(model_instance)
assert isinstance(prepared, LLMPollingCapableProtocol)
assert (
prepared.start_llm_polling(
prompt_messages=[],
model_parameters={"temperature": 0.1},
tools=[],
stop=("END",),
json_schema={"type": "object"},
)
== polling_result
)
assert (
prepared.check_llm_polling(
plugin_state={"task_id": "poll-1"},
)
== polling_result
)
plugin_runtime.start_llm_polling.assert_called_once_with(
provider="langgenius/openai/openai",
model="gpt-4o-mini",
credentials={"api_key": "secret"},
prompt_messages=[],
model_parameters={"temperature": 0.1},
tools=[],
stop=("END",),
json_schema={"type": "object"},
)
plugin_runtime.check_llm_polling.assert_called_once_with(
provider="langgenius/openai/openai",
model="gpt-4o-mini",
credentials={"api_key": "secret"},
plugin_state={"task_id": "poll-1"},
)
def test_dify_prepared_polling_llm_raise_exception_when_polling_is_unsupported() -> None:
llm_result = node_runtime.LLMResult(
model="gpt-4o-mini",
prompt_messages=[],
message=AssistantPromptMessage(content="sync-result"),
usage=node_runtime.LLMUsage.empty_usage(),
)
plugin_runtime = PluginModelRuntime(
tenant_id="tenant-id",
user_id="user-id",
client=Mock(),
plugin_service=Mock(),
)
plugin_runtime.start_llm_polling = Mock(side_effect=PluginLLMPollingUnsupportedError("Polling unsupported")) # type: ignore[method-assign]
model_instance = _ModelInstanceStub(
model_schema=_build_model_schema(features=[ModelFeature.POLLING]),
model_runtime=plugin_runtime,
invoke_llm_result=llm_result,
)
prepared = DifyPreparedPollingLLM(model_instance)
with pytest.raises(PluginLLMPollingUnsupportedError):
prepared.start_llm_polling(
prompt_messages=[],
model_parameters={"temperature": 0.1},
tools=None,
stop=None,
json_schema=None,
)
def test_dify_prompt_message_serializer_delegates(monkeypatch: pytest.MonkeyPatch) -> None:
serialize = MagicMock(return_value={"prompt": "value"})
monkeypatch.setattr(node_runtime.PromptMessageUtil, "prompt_messages_to_prompt_for_saving", serialize)