mirror of
https://github.com/langgenius/dify.git
synced 2026-10-16 15:35:32 +08:00
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:
co-authored by
autofix-ci[bot]
WH-2099
parent
19838972dc
commit
f0b34bdeb4
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user