diff --git a/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_node_execution_behavior.py b/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_node_execution_behavior.py index f8224772207..32334e561fd 100644 --- a/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_node_execution_behavior.py +++ b/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_node_execution_behavior.py @@ -9,20 +9,24 @@ from unittest.mock import MagicMock, patch import pytest from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom +from core.app.llm.model_access import DifyCredentialsProvider, DifyModelFactory +from core.helper.ssrf_proxy import graphon_ssrf_proxy from core.llm_generator.output_parser.structured_output import _parse_structured_output -from core.model_manager import ModelInstance +from core.workflow.node_runtime import ( + DifyPreparedLLM, + DifyPromptMessageSerializer, + build_dify_llm_file_saver, + resolve_dify_run_context, +) from core.workflow.system_variables import build_system_variables from extensions.ext_database import db from graphon.enums import WorkflowNodeExecutionStatus from graphon.model_runtime.entities.message_entities import PromptMessage from graphon.node_events import StreamCompletedEvent from graphon.nodes.llm.entities import LLMNodeData -from graphon.nodes.llm.file_saver import LLMFileSaver from graphon.nodes.llm.node import LLMNode -from graphon.nodes.llm.protocols import CredentialsProvider, ModelFactory -from graphon.nodes.llm.runtime_protocols import PromptMessageSerializerProtocol -from graphon.nodes.protocols import HttpClientProtocol from graphon.runtime import GraphRuntimeState, VariablePool +from tests.unit_tests.core.model_fixtures import make_model_config, make_model_instance from tests.workflow_test_utils import build_test_graph_init_params """FOR MOCK FIXTURES, DO NOT REMOVE""" @@ -74,23 +78,19 @@ def init_llm_node(config: dict[str, object]) -> LLMNode: variable_pool.add(["abc", "output"], "sunny") graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()) - prompt_message_serializer = MagicMock(spec=PromptMessageSerializerProtocol) - prompt_message_serializer.serialize.side_effect = lambda *, prompt_messages, **_kwargs: [ - message.model_dump(mode="json") for message in prompt_messages - ] - llm_file_saver = MagicMock(spec=LLMFileSaver) + run_context = resolve_dify_run_context(init_params.run_context) node = LLMNode( node_id=str(uuid.uuid4()), data=LLMNodeData.model_validate(config["data"]), graph_init_params=init_params, graph_runtime_state=graph_runtime_state, - credentials_provider=MagicMock(spec=CredentialsProvider), - model_factory=MagicMock(spec=ModelFactory), - model_instance=MagicMock(spec=ModelInstance), - llm_file_saver=llm_file_saver, - prompt_message_serializer=prompt_message_serializer, - http_client=MagicMock(spec=HttpClientProtocol), + credentials_provider=DifyCredentialsProvider(run_context=run_context), + model_factory=DifyModelFactory(run_context=run_context), + model_instance=DifyPreparedLLM(make_model_instance(provider="openai", model="gpt-3.5-turbo")), + llm_file_saver=build_dify_llm_file_saver(run_context=run_context, http_client=graphon_ssrf_proxy), + prompt_message_serializer=DifyPromptMessageSerializer(), + http_client=graphon_ssrf_proxy, ) return node @@ -129,28 +129,16 @@ def test_execute_llm(monkeypatch: pytest.MonkeyPatch) -> None: _mock_db_session_close(monkeypatch) - def build_mock_model_instance() -> MagicMock: + def build_model_instance() -> DifyPreparedLLM: from decimal import Decimal from unittest.mock import MagicMock from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage - # Create mock model instance - mock_model_instance = MagicMock(spec=ModelInstance) - mock_model_instance.provider = "openai" - mock_model_instance.model_name = "gpt-3.5-turbo" - mock_model_instance.credentials = dict[str, object]() - mock_model_instance.parameters = dict[str, object]() - mock_model_instance.stop = list[str]() - mock_model_instance.model_type_instance = MagicMock() - mock_model_instance.model_type_instance.get_model_schema.return_value = MagicMock( - model_properties={}, - parameter_rules=[], - features=[], - ) - mock_model_instance.provider_model_bundle = MagicMock() - mock_model_instance.provider_model_bundle.configuration.using_provider_type = "custom" + model_instance = make_model_instance(provider="openai", model="gpt-3.5-turbo") + schema = make_model_config(provider="openai", model="gpt-3.5-turbo", mode="chat").model_schema + monkeypatch.setattr(model_instance, "get_model_schema", MagicMock(return_value=schema)) mock_usage = LLMUsage( prompt_tokens=30, prompt_unit_price=Decimal("0.001"), @@ -172,9 +160,8 @@ def test_execute_llm(monkeypatch: pytest.MonkeyPatch) -> None: message=mock_message, usage=mock_usage, ) - mock_model_instance.invoke_llm.return_value = mock_llm_result - - return mock_model_instance + monkeypatch.setattr(model_instance, "invoke_llm", MagicMock(return_value=mock_llm_result)) + return DifyPreparedLLM(model_instance) # Mock fetch_prompt_messages to avoid database calls def mock_fetch_prompt_messages_1(**_kwargs: object) -> tuple[list[PromptMessage], list[str]]: @@ -185,7 +172,7 @@ def test_execute_llm(monkeypatch: pytest.MonkeyPatch) -> None: UserPromptMessage(content="what's the weather today?"), ], [] - node._model_instance = build_mock_model_instance() + node._model_instance = build_model_instance() with patch.object(LLMNode, "fetch_prompt_messages", mock_fetch_prompt_messages_1): # execute node @@ -243,28 +230,16 @@ def test_execute_llm_with_jinja2(monkeypatch: pytest.MonkeyPatch) -> None: _mock_db_session_close(monkeypatch) - def build_mock_model_instance() -> MagicMock: + def build_model_instance() -> DifyPreparedLLM: from decimal import Decimal from unittest.mock import MagicMock from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage from graphon.model_runtime.entities.message_entities import AssistantPromptMessage - # Create mock model instance - mock_model_instance = MagicMock(spec=ModelInstance) - mock_model_instance.provider = "openai" - mock_model_instance.model_name = "gpt-3.5-turbo" - mock_model_instance.credentials = dict[str, object]() - mock_model_instance.parameters = dict[str, object]() - mock_model_instance.stop = list[str]() - mock_model_instance.model_type_instance = MagicMock() - mock_model_instance.model_type_instance.get_model_schema.return_value = MagicMock( - model_properties={}, - parameter_rules=[], - features=[], - ) - mock_model_instance.provider_model_bundle = MagicMock() - mock_model_instance.provider_model_bundle.configuration.using_provider_type = "custom" + model_instance = make_model_instance(provider="openai", model="gpt-3.5-turbo") + schema = make_model_config(provider="openai", model="gpt-3.5-turbo", mode="chat").model_schema + monkeypatch.setattr(model_instance, "get_model_schema", MagicMock(return_value=schema)) mock_usage = LLMUsage( prompt_tokens=30, prompt_unit_price=Decimal("0.001"), @@ -286,9 +261,8 @@ def test_execute_llm_with_jinja2(monkeypatch: pytest.MonkeyPatch) -> None: message=mock_message, usage=mock_usage, ) - mock_model_instance.invoke_llm.return_value = mock_llm_result - - return mock_model_instance + monkeypatch.setattr(model_instance, "invoke_llm", MagicMock(return_value=mock_llm_result)) + return DifyPreparedLLM(model_instance) # Mock fetch_prompt_messages to avoid database calls def mock_fetch_prompt_messages_2(**_kwargs: object) -> tuple[list[PromptMessage], list[str]]: @@ -299,7 +273,7 @@ def test_execute_llm_with_jinja2(monkeypatch: pytest.MonkeyPatch) -> None: UserPromptMessage(content="what's the weather today?"), ], [] - node._model_instance = build_mock_model_instance() + node._model_instance = build_model_instance() with patch.object(LLMNode, "fetch_prompt_messages", mock_fetch_prompt_messages_2): # execute node diff --git a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_execution_behavior.py b/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_execution_behavior.py index 42950c563e4..14fe53758a4 100644 --- a/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_execution_behavior.py +++ b/api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_execution_behavior.py @@ -9,18 +9,18 @@ from unittest.mock import MagicMock import pytest from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom -from core.model_manager import ModelInstance -from core.workflow.node_runtime import DifyPromptMessageSerializer +from core.app.llm.model_access import DifyCredentialsProvider, DifyModelFactory +from core.workflow.node_runtime import DifyPreparedLLM, DifyPromptMessageSerializer, resolve_dify_run_context from core.workflow.system_variables import build_system_variables from extensions.ext_database import db from graphon.enums import WorkflowNodeExecutionStatus from graphon.model_runtime.entities import AssistantPromptMessage, UserPromptMessage from graphon.model_runtime.entities.message_entities import PromptMessage from graphon.model_runtime.memory.prompt_message_memory import PromptMessageMemory -from graphon.nodes.llm.protocols import CredentialsProvider, ModelFactory from graphon.nodes.parameter_extractor.entities import ParameterExtractorNodeData from graphon.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode from graphon.runtime import GraphRuntimeState, VariablePool +from tests.unit_tests.core.model_fixtures import make_model_instance from tests.unit_tests.core.workflow.nodes.parameter_extractor.fixtures.model import get_mocked_fetch_model_instance from tests.workflow_test_utils import build_test_graph_init_params @@ -79,14 +79,15 @@ def init_parameter_extractor_node( graph_runtime_state = GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()) + run_context = resolve_dify_run_context(init_params.run_context) node = ParameterExtractorNode( node_id=str(uuid.uuid4()), data=ParameterExtractorNodeData.model_validate(config["data"]), graph_init_params=init_params, graph_runtime_state=graph_runtime_state, - credentials_provider=MagicMock(spec=CredentialsProvider), - model_factory=MagicMock(spec=ModelFactory), - model_instance=MagicMock(spec=ModelInstance), + credentials_provider=DifyCredentialsProvider(run_context=run_context), + model_factory=DifyModelFactory(run_context=run_context), + model_instance=DifyPreparedLLM(make_model_instance(provider="openai", model="gpt-3.5-turbo")), memory=memory, prompt_message_serializer=DifyPromptMessageSerializer(), )