youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-04 09:28:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(api): move model_runtime into dify_graph (#32858)

Browse Source
This commit is contained in:
-LAN-
2026-03-02 20:15:32 +08:00
committed by GitHub
parent e985e73bdc
commit 4fd6b52808
253 changed files with 557 additions and 589 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
api/tests/integration_tests/model_runtime/__mock/plugin_model.py
View File

@ -4,20 +4,27 @@ from collections.abc import Generator, Sequence
from decimal import Decimal
from json import dumps
from core.plugin.entities.plugin_daemon import PluginModelProviderEntity
from core.plugin.impl.model import PluginModelClient
# import monkeypatch
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.llm_entities import LLMMode, LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
from core.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage, PromptMessageTool
from core.model_runtime.entities.model_entities import (
from dify_graph.model_runtime.entities.common_entities import I18nObject
from dify_graph.model_runtime.entities.llm_entities import (
LLMMode,
LLMResult,
LLMResultChunk,
LLMResultChunkDelta,
LLMUsage,
)
from dify_graph.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage, PromptMessageTool
from dify_graph.model_runtime.entities.model_entities import (
AIModelEntity,
FetchFrom,
ModelFeature,
ModelPropertyKey,
ModelType,
)
from core.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity
from core.plugin.entities.plugin_daemon import PluginModelProviderEntity
from core.plugin.impl.model import PluginModelClient
from dify_graph.model_runtime.entities.provider_entities import ConfigurateMethod, ProviderEntity
class MockModelClass(PluginModelClient):

4
api/tests/integration_tests/workflow/nodes/__mock/model.py
View File

@ -4,8 +4,8 @@ from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEnti
from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
from core.entities.provider_entities import CustomConfiguration, CustomProviderConfiguration, SystemConfiguration
from core.model_manager import ModelInstance
from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from dify_graph.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from models.provider import ProviderType

12
api/tests/integration_tests/workflow/nodes/test_llm.py
View File

@ -113,8 +113,8 @@ def test_execute_llm():
from decimal import Decimal
from unittest.mock import MagicMock
from core.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from core.model_runtime.entities.message_entities import AssistantPromptMessage
from dify_graph.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from dify_graph.model_runtime.entities.message_entities import AssistantPromptMessage
# Create mock model instance
mock_model_instance = MagicMock(spec=ModelInstance)
@ -158,7 +158,7 @@ def test_execute_llm():
# Mock fetch_prompt_messages to avoid database calls
def mock_fetch_prompt_messages_1(**_kwargs):
from core.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage
from dify_graph.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage
return [
SystemPromptMessage(content="you are a helpful assistant. today's weather is sunny."),
@ -229,8 +229,8 @@ def test_execute_llm_with_jinja2():
from decimal import Decimal
from unittest.mock import MagicMock
from core.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from core.model_runtime.entities.message_entities import AssistantPromptMessage
from dify_graph.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from dify_graph.model_runtime.entities.message_entities import AssistantPromptMessage
# Create mock model instance
mock_model_instance = MagicMock(spec=ModelInstance)
@ -274,7 +274,7 @@ def test_execute_llm_with_jinja2():
# Mock fetch_prompt_messages to avoid database calls
def mock_fetch_prompt_messages_2(**_kwargs):
from core.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage
from dify_graph.model_runtime.entities.message_entities import SystemPromptMessage, UserPromptMessage
return [
SystemPromptMessage(content="you are a helpful assistant. today's weather is sunny."),

2
api/tests/integration_tests/workflow/nodes/test_parameter_extractor.py
View File

@ -5,9 +5,9 @@ from unittest.mock import MagicMock
from core.app.entities.app_invoke_entities import InvokeFrom
from core.model_manager import ModelInstance
from core.model_runtime.entities import AssistantPromptMessage, UserPromptMessage
from dify_graph.entities import GraphInitParams
from dify_graph.enums import WorkflowNodeExecutionStatus
from dify_graph.model_runtime.entities import AssistantPromptMessage, UserPromptMessage
from dify_graph.nodes.llm.protocols import CredentialsProvider, ModelFactory
from dify_graph.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode
from dify_graph.runtime import GraphRuntimeState, VariablePool

2
api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py
View File

@ -31,12 +31,12 @@ from core.app.layers.pause_state_persist_layer import (
PauseStatePersistenceLayer,
WorkflowResumptionContext,
)
from core.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.entities.pause_reason import SchedulingPause
from dify_graph.enums import WorkflowExecutionStatus
from dify_graph.graph_engine.entities.commands import GraphEngineCommand
from dify_graph.graph_engine.layers.base import GraphEngineLayerNotInitializedError
from dify_graph.graph_events.graph import GraphRunPausedEvent
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.runtime.graph_runtime_state import GraphRuntimeState
from dify_graph.runtime.graph_runtime_state_protocol import ReadOnlyGraphRuntimeState
from dify_graph.runtime.read_only_wrappers import ReadOnlyGraphRuntimeStateWrapper

2
api/tests/test_containers_integration_tests/services/test_dataset_service.py
View File

@ -10,8 +10,8 @@ from uuid import uuid4
import pytest
from core.model_runtime.entities.model_entities import ModelType
from core.rag.retrieval.retrieval_methods import RetrievalMethod
from dify_graph.model_runtime.entities.model_entities import ModelType
from extensions.ext_database import db
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, DatasetPermissionEnum, Document, ExternalKnowledgeBindings

2
api/tests/test_containers_integration_tests/services/test_dataset_service_update_dataset.py
View File

@ -3,7 +3,7 @@ from uuid import uuid4
import pytest
from core.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.entities.model_entities import ModelType
from extensions.ext_database import db
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, ExternalKnowledgeBindings

8
api/tests/test_containers_integration_tests/services/test_model_provider_service.py
View File

@ -4,7 +4,7 @@ import pytest
from faker import Faker
from core.entities.model_entities import ModelStatus
from core.model_runtime.entities.model_entities import FetchFrom, ModelType
from dify_graph.model_runtime.entities.model_entities import FetchFrom, ModelType
from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.provider import Provider, ProviderModel, ProviderModelSetting, ProviderType
from services.model_provider_service import ModelProviderService
@ -407,8 +407,8 @@ class TestModelProviderService:
# Create mock models
from core.entities.model_entities import ModelWithProviderEntity, SimpleModelProviderEntity
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.provider_entities import ProviderEntity
from dify_graph.model_runtime.entities.common_entities import I18nObject
from dify_graph.model_runtime.entities.provider_entities import ProviderEntity
# Create real model objects instead of mocks
provider_entity_1 = SimpleModelProviderEntity(
@ -643,7 +643,7 @@ class TestModelProviderService:
# Create mock default model response
from core.entities.model_entities import DefaultModelEntity, DefaultModelProviderEntity
from core.model_runtime.entities.common_entities import I18nObject
from dify_graph.model_runtime.entities.common_entities import I18nObject
mock_default_model = DefaultModelEntity(
model="gpt-3.5-turbo",

2
api/tests/test_containers_integration_tests/services/workflow/test_workflow_converter.py
View File

@ -11,8 +11,8 @@ from core.app.app_config.entities import (
ModelConfigEntity,
PromptTemplateEntity,
)
from core.model_runtime.entities.llm_entities import LLMMode
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from dify_graph.model_runtime.entities.llm_entities import LLMMode
from dify_graph.variables.input_entities import VariableEntity, VariableEntityType
from models import Account, Tenant
from models.api_based_extension import APIBasedExtension

4
api/tests/unit_tests/controllers/console/workspace/test_load_balancing_config.py
View File

@ -13,8 +13,8 @@ from flask import Flask
from flask.views import MethodView
from werkzeug.exceptions import Forbidden
from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from dify_graph.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.errors.validate import CredentialsValidateFailedError
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]

2
api/tests/unit_tests/controllers/service_api/app/test_audio.py
View File

@ -29,7 +29,7 @@ from controllers.service_api.app.error import (
UnsupportedAudioTypeError,
)
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
from core.model_runtime.errors.invoke import InvokeError
from dify_graph.model_runtime.errors.invoke import InvokeError
from services.audio_service import AudioService
from services.errors.app_model_config import AppModelConfigBrokenError
from services.errors.audio import (

2
api/tests/unit_tests/controllers/service_api/app/test_completion.py
View File

@ -34,7 +34,7 @@ from controllers.service_api.app.error import (
NotChatAppError,
)
from core.errors.error import QuotaExceededError
from core.model_runtime.errors.invoke import InvokeError
from dify_graph.model_runtime.errors.invoke import InvokeError
from models.model import App, AppMode, EndUser
from services.app_generate_service import AppGenerateService
from services.app_task_service import AppTaskService

2
api/tests/unit_tests/core/agent/output_parser/test_cot_output_parser.py
View File

@ -3,7 +3,7 @@ from collections.abc import Generator
from core.agent.entities import AgentScratchpadUnit
from core.agent.output_parser.cot_output_parser import CotAgentOutputParser
from core.model_runtime.entities.llm_entities import AssistantPromptMessage, LLMResultChunk, LLMResultChunkDelta
from dify_graph.model_runtime.entities.llm_entities import AssistantPromptMessage, LLMResultChunk, LLMResultChunkDelta
def mock_llm_response(text) -> Generator[LLMResultChunk, None, None]:

2
api/tests/unit_tests/core/app/app_config/features/file_upload/test_manager.py
View File

@ -1,6 +1,6 @@
from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
from core.model_runtime.entities.message_entities import ImagePromptMessageContent
from dify_graph.file.models import FileTransferMethod, FileUploadConfig, ImageConfig
from dify_graph.model_runtime.entities.message_entities import ImagePromptMessageContent
def test_convert_with_vision():

2
api/tests/unit_tests/core/app/apps/chat/test_base_app_runner_multimodal.py
View File

@ -9,8 +9,8 @@ from core.app.apps.base_app_queue_manager import PublishFrom
from core.app.apps.base_app_runner import AppRunner
from core.app.entities.app_invoke_entities import InvokeFrom
from core.app.entities.queue_entities import QueueMessageFileEvent
from core.model_runtime.entities.message_entities import ImagePromptMessageContent
from dify_graph.file.enums import FileTransferMethod, FileType
from dify_graph.model_runtime.entities.message_entities import ImagePromptMessageContent
from models.enums import CreatorUserRole

4
api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_based_generate_task_pipeline.py
View File

@ -25,9 +25,9 @@ from core.app.entities.task_entities import (
)
from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline
from core.base.tts import AppGeneratorTTSPublisher
from core.model_runtime.entities.llm_entities import LLMResult as RuntimeLLMResult
from core.model_runtime.entities.message_entities import TextPromptMessageContent
from core.ops.ops_trace_manager import TraceQueueManager
from dify_graph.model_runtime.entities.llm_entities import LLMResult as RuntimeLLMResult
from dify_graph.model_runtime.entities.message_entities import TextPromptMessageContent
from models.model import AppMode

10
api/tests/unit_tests/core/model_runtime/__base/test_increase_tool_call.py
View File

@ -2,8 +2,8 @@ from unittest.mock import MagicMock, patch
import pytest
from core.model_runtime.entities.message_entities import AssistantPromptMessage
from core.model_runtime.model_providers.__base.large_language_model import _increase_tool_call
from dify_graph.model_runtime.entities.message_entities import AssistantPromptMessage
from dify_graph.model_runtime.model_providers.__base.large_language_model import _increase_tool_call
ToolCall = AssistantPromptMessage.ToolCall
@ -97,7 +97,9 @@ def test__increase_tool_call():
# case 4:
mock_id_generator = MagicMock()
mock_id_generator.side_effect = [_exp_case.id for _exp_case in EXPECTED_CASE_4]
with patch("core.model_runtime.model_providers.__base.large_language_model._gen_tool_call_id", mock_id_generator):
with patch(
"dify_graph.model_runtime.model_providers.__base.large_language_model._gen_tool_call_id", mock_id_generator
):
_run_case(INPUTS_CASE_4, EXPECTED_CASE_4)
@ -107,6 +109,6 @@ def test__increase_tool_call__no_id_no_name_first_delta_should_raise():
ToolCall(id="", type="function", function=ToolCall.ToolCallFunction(name="func_foo", arguments='"value"}')),
]
actual: list[ToolCall] = []
with patch("core.model_runtime.model_providers.__base.large_language_model._gen_tool_call_id", MagicMock()):
with patch("dify_graph.model_runtime.model_providers.__base.large_language_model._gen_tool_call_id", MagicMock()):
with pytest.raises(ValueError):
_increase_tool_call(inputs, actual)

6
api/tests/unit_tests/core/model_runtime/__base/test_large_language_model_non_stream_parsing.py
View File

@ -1,10 +1,10 @@
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
from core.model_runtime.entities.message_entities import (
from dify_graph.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta, LLMUsage
from dify_graph.model_runtime.entities.message_entities import (
AssistantPromptMessage,
TextPromptMessageContent,
UserPromptMessage,
)
from core.model_runtime.model_providers.__base.large_language_model import _normalize_non_stream_plugin_result
from dify_graph.model_runtime.model_providers.__base.large_language_model import _normalize_non_stream_plugin_result
def _make_chunk(

2
api/tests/unit_tests/core/model_runtime/entities/test_llm_entities.py
View File

@ -2,7 +2,7 @@
from decimal import Decimal
from core.model_runtime.entities.llm_entities import LLMUsage, LLMUsageMetadata
from dify_graph.model_runtime.entities.llm_entities import LLMUsage, LLMUsageMetadata
class TestLLMUsage:

16
api/tests/unit_tests/core/plugin/test_plugin_runtime.py
View File

@ -19,14 +19,6 @@ import httpx
import pytest
from pydantic import BaseModel
from core.model_runtime.errors.invoke import (
InvokeAuthorizationError,
InvokeBadRequestError,
InvokeConnectionError,
InvokeRateLimitError,
InvokeServerUnavailableError,
)
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.plugin.entities.plugin_daemon import (
CredentialType,
PluginDaemonInnerError,
@ -44,6 +36,14 @@ from core.plugin.impl.exc import (
)
from core.plugin.impl.plugin import PluginInstaller
from core.plugin.impl.tool import PluginToolManager
from dify_graph.model_runtime.errors.invoke import (
InvokeAuthorizationError,
InvokeBadRequestError,
InvokeConnectionError,
InvokeRateLimitError,
InvokeServerUnavailableError,
)
from dify_graph.model_runtime.errors.validate import CredentialsValidateFailedError
class TestPluginRuntimeExecution:

10
api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py
View File

@ -5,16 +5,16 @@ import pytest
from configs import dify_config
from core.app.app_config.entities import ModelConfigEntity
from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_runtime.entities.message_entities import (
from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from dify_graph.file import File, FileTransferMethod, FileType
from dify_graph.model_runtime.entities.message_entities import (
AssistantPromptMessage,
ImagePromptMessageContent,
PromptMessageRole,
UserPromptMessage,
)
from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
from core.prompt.entities.advanced_prompt_entities import ChatModelMessage, CompletionModelPromptTemplate, MemoryConfig
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from dify_graph.file import File, FileTransferMethod, FileType
from models.model import Conversation

6
api/tests/unit_tests/core/prompt/test_agent_history_prompt_transform.py
View File

@ -5,14 +5,14 @@ from core.app.entities.app_invoke_entities import (
)
from core.entities.provider_configuration import ProviderModelBundle
from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_runtime.entities.message_entities import (
from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform
from dify_graph.model_runtime.entities.message_entities import (
AssistantPromptMessage,
SystemPromptMessage,
ToolPromptMessage,
UserPromptMessage,
)
from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from core.prompt.agent_history_prompt_transform import AgentHistoryPromptTransform
from dify_graph.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
from models.model import Conversation

2
api/tests/unit_tests/core/prompt/test_prompt_message.py
View File

@ -1,4 +1,4 @@
from core.model_runtime.entities.message_entities import (
from dify_graph.model_runtime.entities.message_entities import (
ImagePromptMessageContent,
TextPromptMessageContent,
UserPromptMessage,

8
api/tests/unit_tests/core/prompt/test_prompt_transform.py
View File

@ -2,10 +2,10 @@
# from core.app.app_config.entities import ModelConfigEntity
# from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
# from core.model_runtime.entities.message_entities import UserPromptMessage
# from core.model_runtime.entities.model_entities import AIModelEntity, ModelPropertyKey, ParameterRule
# from core.model_runtime.entities.provider_entities import ProviderEntity
# from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
# from dify_graph.model_runtime.entities.message_entities import UserPromptMessage
# from dify_graph.model_runtime.entities.model_entities import AIModelEntity, ModelPropertyKey, ParameterRule
# from dify_graph.model_runtime.entities.provider_entities import ProviderEntity
# from dify_graph.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
# from core.prompt.prompt_transform import PromptTransform

2
api/tests/unit_tests/core/prompt/test_simple_prompt_transform.py
View File

@ -2,8 +2,8 @@ from unittest.mock import MagicMock
from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
from core.memory.token_buffer_memory import TokenBufferMemory
from core.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage
from core.prompt.simple_prompt_transform import SimplePromptTransform
from dify_graph.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage
from models.model import AppMode, Conversation

8
api/tests/unit_tests/core/rag/embedding/test_embedding_service.py
View File

@ -52,14 +52,14 @@ import pytest
from sqlalchemy.exc import IntegrityError
from core.entities.embedding_type import EmbeddingInputType
from core.model_runtime.entities.model_entities import ModelPropertyKey
from core.model_runtime.entities.text_embedding_entities import EmbeddingResult, EmbeddingUsage
from core.model_runtime.errors.invoke import (
from core.rag.embedding.cached_embedding import CacheEmbedding
from dify_graph.model_runtime.entities.model_entities import ModelPropertyKey
from dify_graph.model_runtime.entities.text_embedding_entities import EmbeddingResult, EmbeddingUsage
from dify_graph.model_runtime.errors.invoke import (
InvokeAuthorizationError,
InvokeConnectionError,
InvokeRateLimitError,
)
from core.rag.embedding.cached_embedding import CacheEmbedding
from models.dataset import Embedding

2
api/tests/unit_tests/core/rag/indexing/test_indexing_runner.py
View File

@ -61,9 +61,9 @@ from core.indexing_runner import (
DocumentIsPausedError,
IndexingRunner,
)
from core.model_runtime.entities.model_entities import ModelType
from core.rag.index_processor.constant.index_type import IndexStructureType
from core.rag.models.document import ChildDocument, Document
from dify_graph.model_runtime.entities.model_entities import ModelType
from libs.datetime_utils import naive_utc_now
from models.dataset import Dataset, DatasetProcessRule
from models.dataset import Document as DatasetDocument

2
api/tests/unit_tests/core/rag/rerank/test_reranker.py
View File

@ -17,13 +17,13 @@ from unittest.mock import MagicMock, Mock, patch
import pytest
from core.model_manager import ModelInstance
from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
from core.rag.models.document import Document
from core.rag.rerank.entity.weight import KeywordSetting, VectorSetting, Weights
from core.rag.rerank.rerank_factory import RerankRunnerFactory
from core.rag.rerank.rerank_model import RerankModelRunner
from core.rag.rerank.rerank_type import RerankMode
from core.rag.rerank.weight_rerank import WeightRerankRunner
from dify_graph.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
def create_mock_model_instance():

2
api/tests/unit_tests/core/test_model_manager.py
View File

@ -6,7 +6,7 @@ from pytest_mock import MockerFixture
from core.entities.provider_entities import ModelLoadBalancingConfiguration
from core.model_manager import LBModelManager
from core.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.entities.model_entities import ModelType
from extensions.ext_redis import redis_client

6
api/tests/unit_tests/core/test_provider_configuration.py
View File

@ -12,9 +12,9 @@ from core.entities.provider_entities import (
RestrictModel,
SystemConfiguration,
)
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.entities.provider_entities import (
from dify_graph.model_runtime.entities.common_entities import I18nObject
from dify_graph.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.entities.provider_entities import (
ConfigurateMethod,
CredentialFormSchema,
FormOption,

2
api/tests/unit_tests/core/test_provider_manager.py
View File

@ -2,8 +2,8 @@ import pytest
from pytest_mock import MockerFixture
from core.entities.provider_entities import ModelSettings
from core.model_runtime.entities.model_entities import ModelType
from core.provider_manager import ProviderManager
from dify_graph.model_runtime.entities.model_entities import ModelType
from models.provider import LoadBalancingModelConfig, ProviderModelSetting

2
api/tests/unit_tests/core/workflow/entities/test_graph_runtime_state.py
View File

@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch
import pytest
from core.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.runtime import GraphRuntimeState, ReadOnlyGraphRuntimeStateWrapper, VariablePool

2
api/tests/unit_tests/core/workflow/graph_engine/layers/test_llm_quota.py
View File

@ -4,10 +4,10 @@ from unittest.mock import MagicMock, patch
from core.app.workflow.layers.llm_quota import LLMQuotaLayer
from core.errors.error import QuotaExceededError
from core.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.enums import NodeType, WorkflowNodeExecutionStatus
from dify_graph.graph_engine.entities.commands import CommandType
from dify_graph.graph_events.node import NodeRunSucceededEvent
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.node_events import NodeRunResult

4
api/tests/unit_tests/core/workflow/graph_engine/test_graph_state_snapshot.py
View File

@ -1,13 +1,13 @@
import time
from collections.abc import Mapping
from core.model_runtime.entities.llm_entities import LLMMode
from core.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.entities import GraphInitParams
from dify_graph.enums import NodeState
from dify_graph.graph import Graph
from dify_graph.graph_engine.graph_state_manager import GraphStateManager
from dify_graph.graph_engine.ready_queue import InMemoryReadyQueue
from dify_graph.model_runtime.entities.llm_entities import LLMMode
from dify_graph.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.nodes.end.end_node import EndNode
from dify_graph.nodes.end.entities import EndNodeData
from dify_graph.nodes.llm.entities import (

2
api/tests/unit_tests/core/workflow/graph_engine/test_human_input_pause_multi_branch.py
View File

@ -4,7 +4,6 @@ from collections.abc import Iterable
from unittest import mock
from unittest.mock import MagicMock
from core.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.entities import GraphInitParams
from dify_graph.graph import Graph
from dify_graph.graph_events import (
@ -17,6 +16,7 @@ from dify_graph.graph_events import (
NodeRunSucceededEvent,
)
from dify_graph.graph_events.node import NodeRunHumanInputFormFilledEvent
from dify_graph.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.nodes.base.entities import OutputVariableEntity, OutputVariableType
from dify_graph.nodes.end.end_node import EndNode
from dify_graph.nodes.end.entities import EndNodeData

2
api/tests/unit_tests/core/workflow/graph_engine/test_human_input_pause_single_branch.py
View File

@ -3,7 +3,6 @@ import time
from unittest import mock
from unittest.mock import MagicMock
from core.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.entities import GraphInitParams
from dify_graph.graph import Graph
from dify_graph.graph_events import (
@ -16,6 +15,7 @@ from dify_graph.graph_events import (
NodeRunSucceededEvent,
)
from dify_graph.graph_events.node import NodeRunHumanInputFormFilledEvent
from dify_graph.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.nodes.base.entities import OutputVariableEntity, OutputVariableType
from dify_graph.nodes.end.end_node import EndNode
from dify_graph.nodes.end.entities import EndNodeData

4
api/tests/unit_tests/core/workflow/graph_engine/test_if_else_streaming.py
View File

@ -1,8 +1,6 @@
import time
from unittest import mock
from core.model_runtime.entities.llm_entities import LLMMode
from core.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.entities import GraphInitParams
from dify_graph.graph import Graph
from dify_graph.graph_events import (
@ -12,6 +10,8 @@ from dify_graph.graph_events import (
NodeRunStreamChunkEvent,
NodeRunSucceededEvent,
)
from dify_graph.model_runtime.entities.llm_entities import LLMMode
from dify_graph.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.nodes.base.entities import OutputVariableEntity, OutputVariableType
from dify_graph.nodes.end.end_node import EndNode
from dify_graph.nodes.end.entities import EndNodeData

2
api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py
View File

@ -11,8 +11,8 @@ from typing import TYPE_CHECKING, Any, Optional
from unittest.mock import MagicMock
from core.model_manager import ModelInstance
from core.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.enums import WorkflowNodeExecutionMetadataKey, WorkflowNodeExecutionStatus
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.node_events import NodeRunResult, StreamChunkEvent, StreamCompletedEvent
from dify_graph.nodes.agent import AgentNode
from dify_graph.nodes.code import CodeNode

4
api/tests/unit_tests/core/workflow/graph_engine/test_parallel_human_input_pause_missing_finish.py
View File

@ -4,8 +4,6 @@ from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Any
from core.model_runtime.entities.llm_entities import LLMMode
from core.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.entities import GraphInitParams
from dify_graph.entities.workflow_start_reason import WorkflowStartReason
from dify_graph.graph import Graph
@ -19,6 +17,8 @@ from dify_graph.graph_events import (
NodeRunStartedEvent,
NodeRunSucceededEvent,
)
from dify_graph.model_runtime.entities.llm_entities import LLMMode
from dify_graph.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.nodes.human_input.entities import HumanInputNodeData, UserAction
from dify_graph.nodes.human_input.enums import HumanInputFormStatus
from dify_graph.nodes.human_input.human_input_node import HumanInputNode

4
api/tests/unit_tests/core/workflow/graph_engine/test_pause_deferred_ready_nodes.py
View File

@ -4,8 +4,6 @@ from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Any
from core.model_runtime.entities.llm_entities import LLMMode
from core.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.entities import GraphInitParams
from dify_graph.entities.workflow_start_reason import WorkflowStartReason
from dify_graph.graph import Graph
@ -18,6 +16,8 @@ from dify_graph.graph_events import (
NodeRunStartedEvent,
NodeRunSucceededEvent,
)
from dify_graph.model_runtime.entities.llm_entities import LLMMode
from dify_graph.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.nodes.end.end_node import EndNode
from dify_graph.nodes.end.entities import EndNodeData
from dify_graph.nodes.human_input.entities import HumanInputNodeData, UserAction

2
api/tests/unit_tests/core/workflow/nodes/knowledge_retrieval/test_knowledge_retrieval_node.py
View File

@ -5,9 +5,9 @@ from unittest.mock import Mock
import pytest
from core.app.entities.app_invoke_entities import InvokeFrom
from core.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.entities import GraphInitParams
from dify_graph.enums import WorkflowNodeExecutionStatus
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.nodes.knowledge_retrieval.entities import (
KnowledgeRetrievalNodeData,
MultipleRetrievalConfig,

14
api/tests/unit_tests/core/workflow/nodes/llm/test_node.py
View File

@ -10,8 +10,11 @@ from core.app.llm.model_access import DifyCredentialsProvider, DifyModelFactory,
from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
from core.entities.provider_entities import CustomConfiguration, SystemConfiguration
from core.model_manager import ModelInstance
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.message_entities import (
from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from dify_graph.entities import GraphInitParams
from dify_graph.file import File, FileTransferMethod, FileType
from dify_graph.model_runtime.entities.common_entities import I18nObject
from dify_graph.model_runtime.entities.message_entities import (
AssistantPromptMessage,
ImagePromptMessageContent,
PromptMessage,
@ -19,11 +22,8 @@ from core.model_runtime.entities.message_entities import (
TextPromptMessageContent,
UserPromptMessage,
)
from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from dify_graph.entities import GraphInitParams
from dify_graph.file import File, FileTransferMethod, FileType
from dify_graph.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
from dify_graph.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from dify_graph.nodes.llm import llm_utils
from dify_graph.nodes.llm.entities import (
ContextConfig,

4
api/tests/unit_tests/core/workflow/nodes/llm/test_scenarios.py
View File

@ -2,9 +2,9 @@ from collections.abc import Mapping, Sequence
from pydantic import BaseModel, Field
from core.model_runtime.entities.message_entities import PromptMessage
from core.model_runtime.entities.model_entities import ModelFeature
from dify_graph.file import File
from dify_graph.model_runtime.entities.message_entities import PromptMessage
from dify_graph.model_runtime.entities.model_entities import ModelFeature
from dify_graph.nodes.llm.entities import LLMNodeChatModelMessage

2
api/tests/unit_tests/core/workflow/nodes/parameter_extractor/test_parameter_extractor_node.py
View File

@ -7,7 +7,7 @@ from typing import Any
import pytest
from core.model_runtime.entities import LLMMode
from dify_graph.model_runtime.entities import LLMMode
from dify_graph.nodes.llm import ModelConfig, VisionConfig
from dify_graph.nodes.parameter_extractor.entities import ParameterConfig, ParameterExtractorNodeData
from dify_graph.nodes.parameter_extractor.exc import (

2
api/tests/unit_tests/core/workflow/nodes/test_question_classifier_node.py
View File

@ -1,4 +1,4 @@
from core.model_runtime.entities import ImagePromptMessageContent
from dify_graph.model_runtime.entities import ImagePromptMessageContent
from dify_graph.nodes.question_classifier import QuestionClassifierNodeData

2
api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node.py
View File

@ -8,11 +8,11 @@ from unittest.mock import MagicMock, patch
import pytest
from core.model_runtime.entities.llm_entities import LLMUsage
from core.tools.entities.tool_entities import ToolInvokeMessage
from core.tools.utils.message_transformer import ToolFileMessageTransformer
from dify_graph.entities import GraphInitParams
from dify_graph.file import File, FileTransferMethod, FileType
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
from dify_graph.node_events import StreamChunkEvent, StreamCompletedEvent
from dify_graph.runtime import GraphRuntimeState, VariablePool
from dify_graph.system_variable import SystemVariable

2
api/tests/unit_tests/repositories/workflow_node_execution/test_sqlalchemy_repository.py
View File

@ -12,7 +12,6 @@ import pytest
from pytest_mock import MockerFixture
from sqlalchemy.orm import Session, sessionmaker
from core.model_runtime.utils.encoders import jsonable_encoder
from core.repositories import SQLAlchemyWorkflowNodeExecutionRepository
from dify_graph.entities import (
WorkflowNodeExecution,
@ -22,6 +21,7 @@ from dify_graph.enums import (
WorkflowNodeExecutionMetadataKey,
WorkflowNodeExecutionStatus,
)
from dify_graph.model_runtime.utils.encoders import jsonable_encoder
from dify_graph.repositories.workflow_node_execution_repository import OrderConfig
from models.account import Account, Tenant
from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionTriggeredFrom

2
api/tests/unit_tests/services/document_service_validation.py
View File

@ -111,7 +111,7 @@ from unittest.mock import Mock, patch
import pytest
from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
from core.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.entities.model_entities import ModelType
from models.dataset import Dataset, DatasetProcessRule, Document
from services.dataset_service import DatasetService, DocumentService
from services.entities.knowledge_entities.knowledge_entities import (

2
api/tests/unit_tests/services/test_dataset_service_create_dataset.py
View File

@ -13,7 +13,7 @@ from uuid import uuid4
import pytest
from core.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.entities.model_entities import ModelType
from models.account import Account
from models.dataset import Dataset, Pipeline
from services.dataset_service import DatasetService

6
api/tests/unit_tests/services/test_model_provider_service_sanitization.py
View File

@ -3,9 +3,9 @@ import types
import pytest
from core.entities.provider_entities import CredentialConfiguration, CustomModelConfiguration
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.entities.provider_entities import ConfigurateMethod
from dify_graph.model_runtime.entities.common_entities import I18nObject
from dify_graph.model_runtime.entities.model_entities import ModelType
from dify_graph.model_runtime.entities.provider_entities import ConfigurateMethod
from models.provider import ProviderType
from services.model_provider_service import ModelProviderService

4
api/tests/unit_tests/services/workflow/test_workflow_converter.py
View File

@ -15,8 +15,8 @@ from core.app.app_config.entities import (
PromptTemplateEntity,
)
from core.helper import encrypter
from core.model_runtime.entities.llm_entities import LLMMode
from core.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.model_runtime.entities.llm_entities import LLMMode
from dify_graph.model_runtime.entities.message_entities import PromptMessageRole
from dify_graph.variables.input_entities import VariableEntity, VariableEntityType
from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint
from models.model import AppMode

2
api/tests/unit_tests/tools/test_mcp_tool.py
View File

@ -13,11 +13,11 @@ from core.mcp.types import (
TextContent,
TextResourceContents,
)
from core.model_runtime.entities.llm_entities import LLMUsage
from core.tools.__base.tool_runtime import ToolRuntime
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_entities import ToolEntity, ToolIdentity, ToolInvokeMessage
from core.tools.mcp_tool.tool import MCPTool
from dify_graph.model_runtime.entities.llm_entities import LLMUsage
def _make_mcp_tool(output_schema: dict | None = None) -> MCPTool:

6
api/tests/unit_tests/utils/structured_output_parser/test_structured_output_parser.py
View File

@ -5,7 +5,7 @@ import pytest
from core.llm_generator.output_parser.errors import OutputParserError
from core.llm_generator.output_parser.structured_output import invoke_llm_with_structured_output
from core.model_runtime.entities.llm_entities import (
from dify_graph.model_runtime.entities.llm_entities import (
LLMResult,
LLMResultChunk,
LLMResultChunkDelta,
@ -13,13 +13,13 @@ from core.model_runtime.entities.llm_entities import (
LLMResultWithStructuredOutput,
LLMUsage,
)
from core.model_runtime.entities.message_entities import (
from dify_graph.model_runtime.entities.message_entities import (
AssistantPromptMessage,
SystemPromptMessage,
TextPromptMessageContent,
UserPromptMessage,
)
from core.model_runtime.entities.model_entities import AIModelEntity, ModelType
from dify_graph.model_runtime.entities.model_entities import AIModelEntity, ModelType
def create_mock_usage(prompt_tokens: int = 10, completion_tokens: int = 5) -> LLMUsage: