test: use concrete model instances in LLM utility tests (#43202)

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(

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"