diff --git a/api/tests/unit_tests/core/workflow/nodes/llm/test_node.py b/api/tests/unit_tests/core/workflow/nodes/llm/test_node.py index 1aff29bf477..090155bc2da 100644 --- a/api/tests/unit_tests/core/workflow/nodes/llm/test_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/llm/test_node.py @@ -16,9 +16,10 @@ from core.app.llm.model_access import ( ) from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle from core.entities.provider_entities import CustomConfiguration, SystemConfiguration +from core.model_manager import ModelInstance from core.plugin.impl.model_runtime_factory import create_plugin_model_assembly from core.prompt.entities.advanced_prompt_entities import MemoryConfig -from core.workflow.node_runtime import DifyFileReferenceFactory, DifyToolFileManager +from core.workflow.node_runtime import DifyFileReferenceFactory, DifyPromptMessageSerializer, DifyToolFileManager from core.workflow.system_variables import default_system_variables from graphon.entities import GraphInitParams from graphon.file import File, FileTransferMethod, FileType @@ -82,9 +83,7 @@ from graphon.nodes.llm.node import ( _handle_memory_completion_mode, _render_jinja2_message, ) -from graphon.nodes.llm.protocols import CredentialsProvider, ModelFactory from graphon.nodes.llm.reasoning import split_reasoning -from graphon.nodes.llm.runtime_protocols import PromptMessageSerializerProtocol from graphon.runtime import GraphRuntimeState, VariablePool from graphon.template_rendering import TemplateRenderError from graphon.variables import ArrayAnySegment, ArrayFileSegment, NoneSegment @@ -104,20 +103,46 @@ class MockTokenBufferMemory: return self.history_messages -def _build_file_saver() -> FileSaverImpl: - run_context = DifyRunContext( +def _run_context() -> DifyRunContext: + return DifyRunContext( tenant_id="tenant", app_id="app", user_id="user", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, ) + + +def _build_file_saver() -> FileSaverImpl: + run_context = _run_context() return FileSaverImpl( tool_file_manager=DifyToolFileManager(run_context), file_reference_factory=DifyFileReferenceFactory(run_context), ) +class _RecordingFileSaver(LLMFileSaver): + """Return a configured file and record the node's save requests without I/O.""" + + def __init__(self, saved_file: File | None) -> None: + # None means no save request is expected in this scenario. + self.saved_file = saved_file + self.binary_calls: list[tuple[bytes, str, FileType, str | None]] = [] + self.remote_calls: list[tuple[str, FileType]] = [] + + def save_binary_string( + self, data: bytes, mime_type: str, file_type: FileType, extension_override: str | None = None + ) -> File: + self.binary_calls.append((data, mime_type, file_type, extension_override)) + assert self.saved_file is not None, "Unexpected binary file save" + return self.saved_file + + def save_remote_url(self, url: str, file_type: FileType) -> File: + self.remote_calls.append((url, file_type)) + assert self.saved_file is not None, "Unexpected remote file save" + return self.saved_file + + def _build_prepared_llm_mock() -> mock.MagicMock: model_instance = mock.MagicMock() model_instance.provider = "openai" @@ -224,20 +249,20 @@ def llm_node( llm_node_data: LLMNodeData, graph_init_params: GraphInitParams, graph_runtime_state: GraphRuntimeState ) -> LLMNode: file_saver = _build_file_saver() - mock_credentials_provider = mock.MagicMock(spec=CredentialsProvider) - mock_model_factory = mock.MagicMock(spec=ModelFactory) - mock_prompt_message_serializer = mock.MagicMock(spec=PromptMessageSerializerProtocol) + credentials_provider = DifyCredentialsProvider(run_context=_run_context()) + model_factory = DifyModelFactory(run_context=_run_context()) + prompt_message_serializer = DifyPromptMessageSerializer() http_client = mock.MagicMock() node = LLMNode( node_id="1", data=llm_node_data, graph_init_params=graph_init_params, graph_runtime_state=graph_runtime_state, - credentials_provider=mock_credentials_provider, - model_factory=mock_model_factory, + credentials_provider=credentials_provider, + model_factory=model_factory, model_instance=_build_prepared_llm_mock(), llm_file_saver=file_saver, - prompt_message_serializer=mock_prompt_message_serializer, + prompt_message_serializer=prompt_message_serializer, http_client=http_client, ) return node @@ -300,9 +325,11 @@ def model_config(monkeypatch: pytest.MonkeyPatch): ) -def test_fetch_model_config_hydrates_model_instance_runtime_settings(model_config: ModelConfigWithCredentialsEntity): - mock_credentials_provider = mock.MagicMock(spec=CredentialsProvider) - mock_model_factory = mock.MagicMock(spec=DifyModelFactory) +def test_fetch_model_config_hydrates_model_instance_runtime_settings( + model_config: ModelConfigWithCredentialsEntity, mocker: MockerFixture +): + credentials_provider = DifyCredentialsProvider(run_context=_run_context()) + model_factory = DifyModelFactory(run_context=_run_context()) provider_model_bundle = model_config.provider_model_bundle model_type_instance = provider_model_bundle.model_type_instance @@ -313,13 +340,10 @@ def test_fetch_model_config_hydrates_model_instance_runtime_settings(model_confi "stop": ["Observation:", "Human:"], } - model_instance = mock.MagicMock( - model_type_instance=model_type_instance, - provider_model_bundle=provider_model_bundle, - ) + model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo", credentials={}) - mock_credentials_provider.fetch.return_value = {"api_key": "test"} - mock_model_factory.init_model_instance.return_value = model_instance + fetch_credentials = mocker.patch.object(credentials_provider, "fetch", return_value={"api_key": "test"}) + initialize_model = mocker.patch.object(model_factory, "init_model_instance", return_value=model_instance) with ( mock.patch.object( @@ -339,8 +363,8 @@ def test_fetch_model_config_hydrates_model_instance_runtime_settings(model_confi mode="chat", completion_params=completion_params, ), - credentials_provider=mock_credentials_provider, - model_factory=mock_model_factory, + credentials_provider=credentials_provider, + model_factory=model_factory, ) assert hydrated_model_instance is model_instance @@ -362,8 +386,8 @@ def test_fetch_model_config_hydrates_model_instance_runtime_settings(model_confi "max_tokens": 256, "stop": ["Observation:", "Human:"], } - mock_credentials_provider.fetch.assert_called_once_with("openai", "gpt-3.5-turbo") - mock_model_factory.init_model_instance.assert_called_once_with("openai", "gpt-3.5-turbo") + fetch_credentials.assert_called_once_with("openai", "gpt-3.5-turbo") + initialize_model.assert_called_once_with("openai", "gpt-3.5-turbo") provider_model.raise_for_status.assert_called_once() @@ -374,9 +398,11 @@ def test_fetch_model_config_hydrates_model_instance_runtime_settings(model_confi ("openai", ""), ], ) -def test_fetch_model_config_rejects_unconfigured_model(provider: str, model_name: str): - credentials_provider = mock.MagicMock(spec=CredentialsProvider) - model_factory = mock.MagicMock(spec=DifyModelFactory) +def test_fetch_model_config_rejects_unconfigured_model(provider: str, model_name: str, mocker: MockerFixture): + credentials_provider = DifyCredentialsProvider(run_context=_run_context()) + model_factory = DifyModelFactory(run_context=_run_context()) + fetch_credentials = mocker.spy(credentials_provider, "fetch") + initialize_model = mocker.spy(model_factory, "init_model_instance") with pytest.raises(ValueError, match="LLM provider and model are required"): fetch_model_config( @@ -390,18 +416,19 @@ def test_fetch_model_config_rejects_unconfigured_model(provider: str, model_name model_factory=model_factory, ) - credentials_provider.fetch.assert_not_called() - model_factory.init_model_instance.assert_not_called() + fetch_credentials.assert_not_called() + initialize_model.assert_not_called() def test_fetch_model_config_reuses_validated_provider_model_from_dify_credentials_provider( model_config: ModelConfigWithCredentialsEntity, + mocker: MockerFixture, ): mock_provider_manager = mock.MagicMock() mock_configurations = mock.MagicMock() mock_provider_configuration = mock.MagicMock() mock_provider_model = mock.MagicMock() - mock_model_factory = mock.MagicMock(spec=DifyModelFactory) + model_factory = DifyModelFactory(run_context=_run_context()) mock_configurations.get.return_value = mock_provider_configuration mock_provider_configuration.get_provider_model.return_value = mock_provider_model @@ -420,11 +447,10 @@ def test_fetch_model_config_reuses_validated_provider_model_from_dify_credential provider_manager=mock_provider_manager, ) - model_instance = mock.MagicMock( - model_type_instance=model_config.provider_model_bundle.model_type_instance, - provider_model_bundle=model_config.provider_model_bundle, + model_instance = ModelInstance( + provider_model_bundle=model_config.provider_model_bundle, model="gpt-3.5-turbo", credentials={} ) - mock_model_factory.init_model_instance.return_value = model_instance + mocker.patch.object(model_factory, "init_model_instance", return_value=model_instance) with mock.patch.object( model_instance.model_type_instance.__class__, @@ -440,7 +466,7 @@ def test_fetch_model_config_reuses_validated_provider_model_from_dify_credential completion_params={}, ), credentials_provider=credentials_provider, - model_factory=mock_model_factory, + model_factory=model_factory, ) mock_provider_configuration.get_provider_model.assert_called_once_with( @@ -1274,30 +1300,32 @@ def test_handle_memory_completion_mode_uses_prompt_message_interface(mocker: Moc @pytest.fixture -def llm_node_for_multimodal(llm_node_data, graph_init_params, graph_runtime_state) -> tuple[LLMNode, LLMFileSaver]: - mock_file_saver: LLMFileSaver = mock.MagicMock(spec=LLMFileSaver) - mock_credentials_provider = mock.MagicMock(spec=CredentialsProvider) - mock_model_factory = mock.MagicMock(spec=ModelFactory) - mock_prompt_message_serializer = mock.MagicMock(spec=PromptMessageSerializerProtocol) +def llm_node_for_multimodal( + llm_node_data, graph_init_params, graph_runtime_state +) -> tuple[LLMNode, _RecordingFileSaver]: + file_saver = _RecordingFileSaver(saved_file=None) + credentials_provider = DifyCredentialsProvider(run_context=_run_context()) + model_factory = DifyModelFactory(run_context=_run_context()) + prompt_message_serializer = DifyPromptMessageSerializer() http_client = mock.MagicMock() node = LLMNode( node_id="1", data=llm_node_data, graph_init_params=graph_init_params, graph_runtime_state=graph_runtime_state, - credentials_provider=mock_credentials_provider, - model_factory=mock_model_factory, + credentials_provider=credentials_provider, + model_factory=model_factory, model_instance=_build_prepared_llm_mock(), - llm_file_saver=mock_file_saver, - prompt_message_serializer=mock_prompt_message_serializer, + llm_file_saver=file_saver, + prompt_message_serializer=prompt_message_serializer, http_client=http_client, ) - return node, mock_file_saver + return node, file_saver class TestLLMNodeSaveMultiModalImageOutput: - def test_llm_node_save_inline_output(self, llm_node_for_multimodal: tuple[LLMNode, LLMFileSaver]): - llm_node, mock_file_saver = llm_node_for_multimodal + def test_llm_node_save_inline_output(self, llm_node_for_multimodal: tuple[LLMNode, _RecordingFileSaver]): + llm_node, file_saver = llm_node_for_multimodal content = ImagePromptMessageContent( format="png", base64_data=base64.b64encode(b"test-data").decode(), @@ -1313,24 +1341,19 @@ class TestLLMNodeSaveMultiModalImageOutput: mime_type="image/png", size=9, ) - mock_file_saver.save_binary_string.return_value = mock_file + file_saver.saved_file = mock_file file = llm_node.save_multimodal_image_output( content=content, - file_saver=mock_file_saver, + file_saver=file_saver, ) # Manually append to _file_outputs since the static method doesn't do it llm_node._file_outputs.append(file) assert llm_node._file_outputs == [mock_file] assert file == mock_file - mock_file_saver.save_binary_string.assert_called_once_with( - data=b"test-data", - mime_type="image/png", - file_type=FileType.IMAGE, - extension_override=".png", - ) + assert file_saver.binary_calls == [(b"test-data", "image/png", FileType.IMAGE, ".png")] - def test_llm_node_save_url_output(self, llm_node_for_multimodal: tuple[LLMNode, LLMFileSaver]): - llm_node, mock_file_saver = llm_node_for_multimodal + def test_llm_node_save_url_output(self, llm_node_for_multimodal: tuple[LLMNode, _RecordingFileSaver]): + llm_node, file_saver = llm_node_for_multimodal content = ImagePromptMessageContent( format="png", url="https://example.com/image.png", @@ -1346,16 +1369,16 @@ class TestLLMNodeSaveMultiModalImageOutput: mime_type="image/png", size=9, ) - mock_file_saver.save_remote_url.return_value = mock_file + file_saver.saved_file = mock_file file = llm_node.save_multimodal_image_output( content=content, - file_saver=mock_file_saver, + file_saver=file_saver, ) # Manually append to _file_outputs since the static method doesn't do it llm_node._file_outputs.append(file) assert llm_node._file_outputs == [mock_file] assert file == mock_file - mock_file_saver.save_remote_url.assert_called_once_with(content.url, FileType.IMAGE) + assert file_saver.remote_calls == [(content.url, FileType.IMAGE)] def test_llm_node_image_file_to_markdown(llm_node: LLMNode): @@ -1372,25 +1395,25 @@ class TestSaveMultimodalOutputAndConvertResultToMarkdown: return "" def test_str_content(self, llm_node_for_multimodal): - llm_node, mock_file_saver = llm_node_for_multimodal + llm_node, file_saver = llm_node_for_multimodal gen = llm_node._save_multimodal_output_and_convert_result_to_markdown( - contents="hello world", file_saver=mock_file_saver, file_outputs=[] + contents="hello world", file_saver=file_saver, file_outputs=[] ) assert list(gen) == ["hello world"] - mock_file_saver.save_binary_string.assert_not_called() - mock_file_saver.save_remote_url.assert_not_called() + assert file_saver.binary_calls == [] + assert file_saver.remote_calls == [] def test_text_prompt_message_content(self, llm_node_for_multimodal): - llm_node, mock_file_saver = llm_node_for_multimodal + llm_node, file_saver = llm_node_for_multimodal gen = llm_node._save_multimodal_output_and_convert_result_to_markdown( - contents=[TextPromptMessageContent(data="hello world")], file_saver=mock_file_saver, file_outputs=[] + contents=[TextPromptMessageContent(data="hello world")], file_saver=file_saver, file_outputs=[] ) assert list(gen) == ["hello world"] - mock_file_saver.save_binary_string.assert_not_called() - mock_file_saver.save_remote_url.assert_not_called() + assert file_saver.binary_calls == [] + assert file_saver.remote_calls == [] def test_image_content_with_inline_data(self, llm_node_for_multimodal, monkeypatch: pytest.MonkeyPatch): - llm_node, mock_file_saver = llm_node_for_multimodal + llm_node, file_saver = llm_node_for_multimodal image_raw_data = b"PNG_DATA" image_b64_data = base64.b64encode(image_raw_data).decode() @@ -1406,7 +1429,7 @@ class TestSaveMultimodalOutputAndConvertResultToMarkdown: url="https://example.com/test.png", storage_key="test_storage_key", ) - mock_file_saver.save_binary_string.return_value = mock_saved_file + file_saver.saved_file = mock_saved_file gen = llm_node._save_multimodal_output_and_convert_result_to_markdown( contents=[ ImagePromptMessageContent( @@ -1415,7 +1438,7 @@ class TestSaveMultimodalOutputAndConvertResultToMarkdown: mime_type="image/png", ) ], - file_saver=mock_file_saver, + file_saver=file_saver, file_outputs=llm_node._file_outputs, ) yielded_strs = list(gen) @@ -1430,45 +1453,40 @@ class TestSaveMultimodalOutputAndConvertResultToMarkdown: assert yielded_strs[0].startswith("![](") assert expected_file_url_path in yielded_strs[0] assert yielded_strs[0].endswith(")") - mock_file_saver.save_binary_string.assert_called_once_with( - data=image_raw_data, - mime_type="image/png", - file_type=FileType.IMAGE, - extension_override=".png", - ) + assert file_saver.binary_calls == [(image_raw_data, "image/png", FileType.IMAGE, ".png")] assert mock_saved_file in llm_node._file_outputs def test_unknown_content_type(self, llm_node_for_multimodal): - llm_node, mock_file_saver = llm_node_for_multimodal + llm_node, file_saver = llm_node_for_multimodal gen = llm_node._save_multimodal_output_and_convert_result_to_markdown( - contents=frozenset(("hello world",)), file_saver=mock_file_saver, file_outputs=[] + contents=frozenset(("hello world",)), file_saver=file_saver, file_outputs=[] ) assert list(gen) == ["hello world"] - mock_file_saver.save_binary_string.assert_not_called() - mock_file_saver.save_remote_url.assert_not_called() + assert file_saver.binary_calls == [] + assert file_saver.remote_calls == [] def test_unknown_item_type(self, llm_node_for_multimodal, caplog: pytest.LogCaptureFixture): - llm_node, mock_file_saver = llm_node_for_multimodal + llm_node, file_saver = llm_node_for_multimodal unknown_item = self._UnknownItem() with caplog.at_level(logging.WARNING, logger="graphon.nodes.llm.node"): gen = llm_node._save_multimodal_output_and_convert_result_to_markdown( - contents=[unknown_item], file_saver=mock_file_saver, file_outputs=[] + contents=[unknown_item], file_saver=file_saver, file_outputs=[] ) assert list(gen) == [str(unknown_item)] assert "unknown item type encountered" in caplog.text - mock_file_saver.save_binary_string.assert_not_called() - mock_file_saver.save_remote_url.assert_not_called() + assert file_saver.binary_calls == [] + assert file_saver.remote_calls == [] def test_none_content(self, llm_node_for_multimodal): - llm_node, mock_file_saver = llm_node_for_multimodal + llm_node, file_saver = llm_node_for_multimodal gen = llm_node._save_multimodal_output_and_convert_result_to_markdown( - contents=None, file_saver=mock_file_saver, file_outputs=[] + contents=None, file_saver=file_saver, file_outputs=[] ) assert list(gen) == [] - mock_file_saver.save_binary_string.assert_not_called() - mock_file_saver.save_remote_url.assert_not_called() + assert file_saver.binary_calls == [] + assert file_saver.remote_calls == [] class TestReasoningFormat: