diff --git a/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_utils.py b/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_utils.py index 6a2fc81fef5..3d6f14fbb11 100644 --- a/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_utils.py +++ b/api/tests/unit_tests/core/workflow/nodes/llm/test_llm_utils.py @@ -1,6 +1,7 @@ from unittest import mock import pytest +from pytest_mock import MockerFixture from core.model_manager import ModelInstance from graphon.file import File, FileTransferMethod, FileType @@ -34,6 +35,7 @@ from graphon.nodes.llm.exc import ( ) from graphon.runtime import VariablePool from graphon.variables import ArrayAnySegment, ArrayFileSegment, NoneSegment +from tests.unit_tests.core.model_fixtures import make_model_instance def _build_model_schema( @@ -53,12 +55,14 @@ def _build_model_schema( ) -def _build_model_instance(*, model_schema: AIModelEntity | None = None) -> mock.MagicMock: - model_instance = mock.MagicMock(spec=ModelInstance) +def _build_model_instance(mocker: MockerFixture, *, model_schema: AIModelEntity | None = None) -> ModelInstance: + model_instance = make_model_instance(provider="openai", model="gpt-3.5-turbo") model_instance.model_name = "gpt-3.5-turbo" model_instance.parameters = {} - model_instance.get_model_schema.return_value = model_schema or _build_model_schema(features=[]) - model_instance.get_llm_num_tokens.return_value = 0 + mocker.patch.object( + model_instance, "get_model_schema", return_value=model_schema or _build_model_schema(features=[]) + ) + mocker.patch.object(model_instance, "get_llm_num_tokens", return_value=0) return model_instance @@ -94,9 +98,8 @@ def variable_pool() -> VariablePool: def _fetch_prompt_messages_with_mocked_content(content): variable_pool = VariablePool.empty() - model_instance = mock.MagicMock(spec=ModelInstance) - model_schema = mock.MagicMock() - model_schema.supports_prompt_content_type.side_effect = lambda content_type: content_type == "text" + model_instance = make_model_instance(provider="openai", model="gpt-3.5-turbo") + model_schema = _build_model_schema(features=[]) prompt_template = [ LLMNodeChatModelMessage( text="You are a classifier.", @@ -342,9 +345,9 @@ def test_fetch_prompt_messages_keeps_list_content_when_multiple_supported_items_ ] -def test_fetch_model_schema_raises_when_model_schema_is_missing(): - model_instance = _build_model_instance() - model_instance.get_model_schema.return_value = None +def test_fetch_model_schema_raises_when_model_schema_is_missing(mocker: MockerFixture): + model_instance = _build_model_instance(mocker) + mocker.patch.object(model_instance, "get_model_schema", return_value=None) with pytest.raises(ValueError, match="Model schema not found for gpt-3.5-turbo"): llm_utils.fetch_model_schema(model_instance=model_instance) @@ -617,8 +620,9 @@ def test_combine_message_content_with_role_handles_all_supported_roles(): llm_utils.combine_message_content_with_role(contents=contents, role="custom") # type: ignore[arg-type] -def test_calculate_rest_token_uses_context_size_and_template_alias(): +def test_calculate_rest_token_uses_context_size_and_template_alias(mocker: MockerFixture): model_instance = _build_model_instance( + mocker, model_schema=_build_model_schema( model_properties={ModelPropertyKey.CONTEXT_SIZE: 4096}, parameter_rules=[ @@ -629,10 +633,10 @@ def test_calculate_rest_token_uses_context_size_and_template_alias(): type=ParameterType.INT, ) ], - ) + ), ) model_instance.parameters = {"max_tokens": 512} - model_instance.get_llm_num_tokens.return_value = 256 + mocker.patch.object(model_instance, "get_llm_num_tokens", return_value=256) assert ( llm_utils.calculate_rest_token( @@ -643,8 +647,8 @@ def test_calculate_rest_token_uses_context_size_and_template_alias(): ) -def test_handle_memory_chat_mode_returns_empty_without_memory_and_uses_window_when_present(): - model_instance = _build_model_instance() +def test_handle_memory_chat_mode_returns_empty_without_memory_and_uses_window_when_present(mocker: MockerFixture): + model_instance = _build_model_instance(mocker) memory = mock.MagicMock() memory.get_history_prompt_messages.return_value = [UserPromptMessage(content="Question")] @@ -669,8 +673,8 @@ def test_handle_memory_chat_mode_returns_empty_without_memory_and_uses_window_wh memory.get_history_prompt_messages.assert_called_once_with(max_token_limit=123, message_limit=2) -def test_handle_memory_completion_mode_validates_role_prefix_and_formats_history(): - model_instance = _build_model_instance() +def test_handle_memory_completion_mode_validates_role_prefix_and_formats_history(mocker: MockerFixture): + model_instance = _build_model_instance(mocker) memory = mock.MagicMock() memory.get_history_prompt_messages.return_value = [ UserPromptMessage(content="Question"), @@ -750,8 +754,8 @@ def test_append_file_prompts_merges_with_existing_user_content_or_appends_new_me assert prompt_messages[-1] == UserPromptMessage(content=[file_prompt]) -def test_fetch_prompt_messages_chat_mode_includes_query_memory_and_supported_files(): - model_instance = _build_model_instance(model_schema=_build_model_schema(features=[ModelFeature.VISION])) +def test_fetch_prompt_messages_chat_mode_includes_query_memory_and_supported_files(mocker: MockerFixture): + model_instance = _build_model_instance(mocker, model_schema=_build_model_schema(features=[ModelFeature.VISION])) memory = mock.MagicMock() memory.get_history_prompt_messages.return_value = [AssistantPromptMessage(content="history")] sys_file = _build_image_file(file_id="sys", related_id="sys-related", remote_url="https://example.com/sys.png") @@ -813,8 +817,8 @@ def test_fetch_prompt_messages_chat_mode_includes_query_memory_and_supported_fil ) -def test_fetch_prompt_messages_completion_mode_updates_list_content_with_histories_and_query(): - model_instance = _build_model_instance(model_schema=_build_model_schema(features=[])) +def test_fetch_prompt_messages_completion_mode_updates_list_content_with_histories_and_query(mocker: MockerFixture): + model_instance = _build_model_instance(mocker, model_schema=_build_model_schema(features=[])) memory = mock.MagicMock() memory.get_history_prompt_messages.return_value = [ UserPromptMessage(content="previous question"), @@ -880,8 +884,8 @@ def test_fetch_prompt_messages_completion_mode_updates_list_content_with_histori ] -def test_fetch_prompt_messages_filters_content_unsupported_by_model_features(): - model_instance = _build_model_instance(model_schema=_build_model_schema(features=[ModelFeature.DOCUMENT])) +def test_fetch_prompt_messages_filters_content_unsupported_by_model_features(mocker: MockerFixture): + model_instance = _build_model_instance(mocker, model_schema=_build_model_schema(features=[ModelFeature.DOCUMENT])) prompt_template = [ LLMNodeChatModelMessage( text="You are a classifier.", @@ -927,8 +931,8 @@ def test_fetch_prompt_messages_filters_content_unsupported_by_model_features(): assert prompt_messages == [SystemPromptMessage(content="You are a classifier.")] -def test_fetch_prompt_messages_completion_mode_supports_string_content_and_invalid_template_type(): - model_instance = _build_model_instance(model_schema=_build_model_schema(features=[])) +def test_fetch_prompt_messages_completion_mode_supports_string_content_and_invalid_template_type(mocker: MockerFixture): + model_instance = _build_model_instance(mocker, model_schema=_build_model_schema(features=[])) with ( mock.patch( diff --git a/api/tests/unit_tests/services/test_agent_llm_inner_service.py b/api/tests/unit_tests/services/test_agent_llm_inner_service.py index 37eb8f851ad..8eff7b79fbb 100644 --- a/api/tests/unit_tests/services/test_agent_llm_inner_service.py +++ b/api/tests/unit_tests/services/test_agent_llm_inner_service.py @@ -8,6 +8,7 @@ from uuid import uuid4 import pytest from pydantic import JsonValue +from pytest_mock import MockerFixture from sqlalchemy.orm import Session, sessionmaker from configs import dify_config @@ -21,6 +22,7 @@ from models.model import App, AppMode from models.provider import ProviderType from services.agent_llm_inner_service import AgentLLMInnerService, AgentLLMInnerServiceError, PreparedAgentLLMInvocation from services.entities.agent_llm_inner import AgentLLMInvokeCaller, AgentLLMInvokeRequest, AgentLLMInvokeTarget +from tests.unit_tests.core.model_fixtures import make_model_instance def _request() -> AgentLLMInvokeRequest: @@ -204,7 +206,7 @@ def test_gateway_uses_quota_managed_instance_as_single_credit_owner( provider_invoke.assert_called_once() -def test_gateway_forwards_prompt_messages_without_revalidation() -> None: +def test_gateway_forwards_prompt_messages_without_revalidation(mocker: MockerFixture) -> None: request = _request() prompt_messages: list[dict[str, JsonValue]] = [ { @@ -227,18 +229,18 @@ def test_gateway_forwards_prompt_messages_without_revalidation() -> None: }, ] request.target.prompt_messages = prompt_messages - model_instance = MagicMock(spec=ModelInstance) - model_instance.invoke_llm.return_value = iter([]) + model_instance = make_model_instance(provider=request.target.provider, model=request.target.model) + invoke = mocker.patch.object(model_instance, "invoke_llm", return_value=iter([])) prepared = PreparedAgentLLMInvocation(request=request, model_instance=model_instance) assert list(AgentLLMInnerService().invoke(prepared)) == [] - assert model_instance.invoke_llm.call_args.kwargs["prompt_messages"] is prompt_messages + assert invoke.call_args.kwargs["prompt_messages"] is prompt_messages -def test_standalone_agent_app_gateway_is_attributed_to_app() -> None: +def test_standalone_agent_app_gateway_is_attributed_to_app(mocker: MockerFixture) -> None: request = _request() - model_instance = MagicMock(spec=ModelInstance) - model_instance.invoke_llm.return_value = iter([]) + model_instance = make_model_instance(provider=request.target.provider, model=request.target.model) + invoke = mocker.patch.object(model_instance, "invoke_llm", return_value=iter([])) prepared = PreparedAgentLLMInvocation( request=request, model_instance=model_instance, @@ -246,16 +248,16 @@ def test_standalone_agent_app_gateway_is_attributed_to_app() -> None: ) assert list(AgentLLMInnerService().invoke(prepared)) == [] - request_metadata = model_instance.invoke_llm.call_args.kwargs["request_metadata"] + request_metadata = invoke.call_args.kwargs["request_metadata"] assert request_metadata["created_by"] is CreditUsageCreatedBy.APP assert request_metadata["agent_config_version_kind"] == "draft" -def test_agent_build_draft_gateway_is_attributed_to_build_draft() -> None: +def test_agent_build_draft_gateway_is_attributed_to_build_draft(mocker: MockerFixture) -> None: request = _request() request.caller.agent_config_version_kind = "build_draft" - model_instance = MagicMock(spec=ModelInstance) - model_instance.invoke_llm.return_value = iter([]) + model_instance = make_model_instance(provider=request.target.provider, model=request.target.model) + invoke = mocker.patch.object(model_instance, "invoke_llm", return_value=iter([])) prepared = PreparedAgentLLMInvocation( request=request, model_instance=model_instance, @@ -263,7 +265,7 @@ def test_agent_build_draft_gateway_is_attributed_to_build_draft() -> None: ) assert list(AgentLLMInnerService().invoke(prepared)) == [] - request_metadata = model_instance.invoke_llm.call_args.kwargs["request_metadata"] + request_metadata = invoke.call_args.kwargs["request_metadata"] assert request_metadata["created_by"] is CreditUsageCreatedBy.BUILD_DRAFT assert request_metadata["agent_config_version_kind"] == "build_draft"