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

merge

Browse Source
This commit is contained in:
Joel
2025-01-14 11:35:29 +08:00
parent 8865bdca05 0a6b4d01d7
commit d6b9648332
192 changed files with 3813 additions and 1041 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
api/tests/integration_tests/model_runtime/__mock/plugin_daemon.py Normal file
View File

@ -0,0 +1,44 @@
import os
from collections.abc import Callable
import pytest
# import monkeypatch
from _pytest.monkeypatch import MonkeyPatch
from core.plugin.manager.model import PluginModelManager
from tests.integration_tests.model_runtime.__mock.plugin_model import MockModelClass
def mock_plugin_daemon(
monkeypatch: MonkeyPatch,
) -> Callable[[], None]:
"""
mock openai module
:param monkeypatch: pytest monkeypatch fixture
:return: unpatch function
"""
def unpatch() -> None:
monkeypatch.undo()
monkeypatch.setattr(PluginModelManager, "invoke_llm", MockModelClass.invoke_llm)
monkeypatch.setattr(PluginModelManager, "fetch_model_providers", MockModelClass.fetch_model_providers)
monkeypatch.setattr(PluginModelManager, "get_model_schema", MockModelClass.get_model_schema)
return unpatch
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture
def setup_model_mock(monkeypatch):
if MOCK:
unpatch = mock_plugin_daemon(monkeypatch)
yield
if MOCK:
unpatch()

249
api/tests/integration_tests/model_runtime/__mock/plugin_model.py Normal file
View File

@ -0,0 +1,249 @@
import datetime
import uuid
from collections.abc import Generator, Sequence
from decimal import Decimal
from json import dumps
# import monkeypatch
from typing import Optional
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 (
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.manager.model import PluginModelManager
class MockModelClass(PluginModelManager):
def fetch_model_providers(self, tenant_id: str) -> Sequence[PluginModelProviderEntity]:
"""
Fetch model providers for the given tenant.
"""
return [
PluginModelProviderEntity(
id=uuid.uuid4().hex,
created_at=datetime.datetime.now(),
updated_at=datetime.datetime.now(),
provider="openai",
tenant_id=tenant_id,
plugin_unique_identifier="langgenius/openai/openai",
plugin_id="langgenius/openai",
declaration=ProviderEntity(
provider="openai",
label=I18nObject(
en_US="OpenAI",
zh_Hans="OpenAI",
),
description=I18nObject(
en_US="OpenAI",
zh_Hans="OpenAI",
),
icon_small=I18nObject(
en_US="https://example.com/icon_small.png",
zh_Hans="https://example.com/icon_small.png",
),
icon_large=I18nObject(
en_US="https://example.com/icon_large.png",
zh_Hans="https://example.com/icon_large.png",
),
supported_model_types=[ModelType.LLM],
configurate_methods=[ConfigurateMethod.PREDEFINED_MODEL],
models=[
AIModelEntity(
model="gpt-3.5-turbo",
label=I18nObject(
en_US="gpt-3.5-turbo",
zh_Hans="gpt-3.5-turbo",
),
model_type=ModelType.LLM,
fetch_from=FetchFrom.PREDEFINED_MODEL,
model_properties={},
features=[ModelFeature.TOOL_CALL, ModelFeature.MULTI_TOOL_CALL],
),
AIModelEntity(
model="gpt-3.5-turbo-instruct",
label=I18nObject(
en_US="gpt-3.5-turbo-instruct",
zh_Hans="gpt-3.5-turbo-instruct",
),
model_type=ModelType.LLM,
fetch_from=FetchFrom.PREDEFINED_MODEL,
model_properties={
ModelPropertyKey.MODE: LLMMode.COMPLETION,
},
features=[],
),
],
),
)
]
def get_model_schema(
self,
tenant_id: str,
user_id: str,
plugin_id: str,
provider: str,
model_type: str,
model: str,
credentials: dict,
) -> AIModelEntity | None:
"""
Get model schema
"""
return AIModelEntity(
model=model,
label=I18nObject(
en_US="OpenAI",
zh_Hans="OpenAI",
),
model_type=ModelType(model_type),
fetch_from=FetchFrom.PREDEFINED_MODEL,
model_properties={},
features=[ModelFeature.TOOL_CALL, ModelFeature.MULTI_TOOL_CALL] if model == "gpt-3.5-turbo" else [],
)
@staticmethod
def generate_function_call(
tools: Optional[list[PromptMessageTool]],
) -> Optional[AssistantPromptMessage.ToolCall]:
if not tools or len(tools) == 0:
return None
function: PromptMessageTool = tools[0]
function_name = function.name
function_parameters = function.parameters
function_parameters_type = function_parameters["type"]
if function_parameters_type != "object":
return None
function_parameters_properties = function_parameters["properties"]
function_parameters_required = function_parameters["required"]
parameters = {}
for parameter_name, parameter in function_parameters_properties.items():
if parameter_name not in function_parameters_required:
continue
parameter_type = parameter["type"]
if parameter_type == "string":
if "enum" in parameter:
if len(parameter["enum"]) == 0:
continue
parameters[parameter_name] = parameter["enum"][0]
else:
parameters[parameter_name] = "kawaii"
elif parameter_type == "integer":
parameters[parameter_name] = 114514
elif parameter_type == "number":
parameters[parameter_name] = 1919810.0
elif parameter_type == "boolean":
parameters[parameter_name] = True
return AssistantPromptMessage.ToolCall(
id=str(uuid.uuid4()),
type="function",
function=AssistantPromptMessage.ToolCall.ToolCallFunction(
name=function_name,
arguments=dumps(parameters),
),
)
@staticmethod
def mocked_chat_create_sync(
model: str,
prompt_messages: list[PromptMessage],
tools: Optional[list[PromptMessageTool]] = None,
) -> LLMResult:
tool_call = MockModelClass.generate_function_call(tools=tools)
return LLMResult(
id=str(uuid.uuid4()),
model=model,
prompt_messages=prompt_messages,
message=AssistantPromptMessage(content="elaina", tool_calls=[tool_call] if tool_call else []),
usage=LLMUsage(
prompt_tokens=2,
completion_tokens=1,
total_tokens=3,
prompt_unit_price=Decimal(0.0001),
completion_unit_price=Decimal(0.0002),
prompt_price_unit=Decimal(1),
prompt_price=Decimal(0.0001),
completion_price_unit=Decimal(1),
completion_price=Decimal(0.0002),
total_price=Decimal(0.0003),
currency="USD",
latency=0.001,
),
)
@staticmethod
def mocked_chat_create_stream(
model: str,
prompt_messages: list[PromptMessage],
tools: Optional[list[PromptMessageTool]] = None,
) -> Generator[LLMResultChunk, None, None]:
tool_call = MockModelClass.generate_function_call(tools=tools)
full_text = "Hello, world!\n\n```python\nprint('Hello, world!')\n```"
for i in range(0, len(full_text) + 1):
if i == len(full_text):
yield LLMResultChunk(
model=model,
prompt_messages=prompt_messages,
delta=LLMResultChunkDelta(
index=0,
message=AssistantPromptMessage(
content="",
tool_calls=[tool_call] if tool_call else [],
),
),
)
else:
yield LLMResultChunk(
model=model,
prompt_messages=prompt_messages,
delta=LLMResultChunkDelta(
index=0,
message=AssistantPromptMessage(
content=full_text[i],
tool_calls=[tool_call] if tool_call else [],
),
usage=LLMUsage(
prompt_tokens=2,
completion_tokens=17,
total_tokens=19,
prompt_unit_price=Decimal(0.0001),
completion_unit_price=Decimal(0.0002),
prompt_price_unit=Decimal(1),
prompt_price=Decimal(0.0001),
completion_price_unit=Decimal(1),
completion_price=Decimal(0.0002),
total_price=Decimal(0.0003),
currency="USD",
latency=0.001,
),
),
)
def invoke_llm(
self: PluginModelManager,
*,
tenant_id: str,
user_id: str,
plugin_id: str,
provider: str,
model: str,
credentials: dict,
prompt_messages: list[PromptMessage],
model_parameters: Optional[dict] = None,
tools: Optional[list[PromptMessageTool]] = None,
stop: Optional[list[str]] = None,
stream: bool = True,
):
return MockModelClass.mocked_chat_create_stream(model=model, prompt_messages=prompt_messages, tools=tools)

55
api/tests/integration_tests/model_runtime/gpustack/test_speech2text.py
View File

@ -1,55 +0,0 @@
import os
from pathlib import Path
import pytest
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gpustack.speech2text.speech2text import GPUStackSpeech2TextModel
def test_validate_credentials():
model = GPUStackSpeech2TextModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="faster-whisper-medium",
credentials={
"endpoint_url": "invalid_url",
"api_key": "invalid_api_key",
},
)
model.validate_credentials(
model="faster-whisper-medium",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
},
)
def test_invoke_model():
model = GPUStackSpeech2TextModel()
# Get the directory of the current file
current_dir = os.path.dirname(os.path.abspath(__file__))
# Get assets directory
assets_dir = os.path.join(os.path.dirname(current_dir), "assets")
# Construct the path to the audio file
audio_file_path = os.path.join(assets_dir, "audio.mp3")
file = Path(audio_file_path).read_bytes()
result = model.invoke(
model="faster-whisper-medium",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
},
file=file,
)
assert isinstance(result, str)
assert result == "1, 2, 3, 4, 5, 6, 7, 8, 9, 10"

24
api/tests/integration_tests/model_runtime/gpustack/test_tts.py
View File

@ -1,24 +0,0 @@
import os
from core.model_runtime.model_providers.gpustack.tts.tts import GPUStackText2SpeechModel
def test_invoke_model():
model = GPUStackText2SpeechModel()
result = model.invoke(
model="cosyvoice-300m-sft",
tenant_id="test",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
},
content_text="Hello world",
voice="Chinese Female",
)
content = b""
for chunk in result:
content += chunk
assert content != b""

50
api/tests/integration_tests/workflow/nodes/__mock/model.py Normal file
View File

@ -0,0 +1,50 @@
from unittest.mock import MagicMock
from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
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 models.provider import ProviderType
def get_mocked_fetch_model_config(
provider: str,
model: str,
mode: str,
credentials: dict,
):
model_provider_factory = ModelProviderFactory(tenant_id="test_tenant")
model_type_instance = model_provider_factory.get_model_type_instance(provider, ModelType.LLM)
provider_model_bundle = ProviderModelBundle(
configuration=ProviderConfiguration(
tenant_id="1",
provider=model_provider_factory.get_provider_schema(provider),
preferred_provider_type=ProviderType.CUSTOM,
using_provider_type=ProviderType.CUSTOM,
system_configuration=SystemConfiguration(enabled=False),
custom_configuration=CustomConfiguration(provider=CustomProviderConfiguration(credentials=credentials)),
model_settings=[],
),
model_type_instance=model_type_instance,
)
model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model=model)
model_schema = model_provider_factory.get_model_schema(
provider=provider,
model_type=model_type_instance.model_type,
model=model,
credentials=credentials,
)
assert model_schema is not None
model_config = ModelConfigWithCredentialsEntity(
model=model,
provider=provider,
mode=mode,
credentials=credentials,
parameters={},
model_schema=model_schema,
provider_model_bundle=provider_model_bundle,
)
return MagicMock(return_value=(model_instance, model_config))

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

@ -7,12 +7,7 @@ from unittest.mock import MagicMock
import pytest
from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity
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 import ModelProviderFactory
from core.app.entities.app_invoke_entities import InvokeFrom
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.enums import SystemVariableKey
from core.workflow.graph_engine.entities.graph import Graph
@ -22,11 +17,11 @@ from core.workflow.nodes.event import RunCompletedEvent
from core.workflow.nodes.llm.node import LLMNode
from extensions.ext_database import db
from models.enums import UserFrom
from models.provider import ProviderType
from models.workflow import WorkflowNodeExecutionStatus, WorkflowType
from tests.integration_tests.workflow.nodes.__mock.model import get_mocked_fetch_model_config
"""FOR MOCK FIXTURES, DO NOT REMOVE"""
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
from tests.integration_tests.model_runtime.__mock.plugin_daemon import setup_model_mock
from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock
@ -81,15 +76,19 @@ def init_llm_node(config: dict) -> LLMNode:
return node
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_execute_llm(setup_openai_mock):
def test_execute_llm(setup_model_mock):
node = init_llm_node(
config={
"id": "llm",
"data": {
"title": "123",
"type": "llm",
"model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}},
"model": {
"provider": "langgenius/openai/openai",
"name": "gpt-3.5-turbo",
"mode": "chat",
"completion_params": {},
},
"prompt_template": [
{"role": "system", "text": "you are a helpful assistant.\ntoday's weather is {{#abc.output#}}."},
{"role": "user", "text": "{{#sys.query#}}"},
@ -103,37 +102,15 @@ def test_execute_llm(setup_openai_mock):
credentials = {"openai_api_key": os.environ.get("OPENAI_API_KEY")}
provider_instance = ModelProviderFactory().get_provider_instance("openai")
model_type_instance = provider_instance.get_model_instance(ModelType.LLM)
provider_model_bundle = ProviderModelBundle(
configuration=ProviderConfiguration(
tenant_id="1",
provider=provider_instance.get_provider_schema(),
preferred_provider_type=ProviderType.CUSTOM,
using_provider_type=ProviderType.CUSTOM,
system_configuration=SystemConfiguration(enabled=False),
custom_configuration=CustomConfiguration(provider=CustomProviderConfiguration(credentials=credentials)),
model_settings=[],
),
model_type_instance=model_type_instance,
)
model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo")
model_schema = model_type_instance.get_model_schema("gpt-3.5-turbo")
assert model_schema is not None
model_config = ModelConfigWithCredentialsEntity(
model="gpt-3.5-turbo",
provider="openai",
mode="chat",
credentials=credentials,
parameters={},
model_schema=model_schema,
provider_model_bundle=provider_model_bundle,
)
# Mock db.session.close()
db.session.close = MagicMock()
node._fetch_model_config = MagicMock(return_value=(model_instance, model_config))
node._fetch_model_config = get_mocked_fetch_model_config(
provider="langgenius/openai/openai",
model="gpt-3.5-turbo",
mode="chat",
credentials=credentials,
)
# execute node
result = node._run()
@ -149,8 +126,7 @@ def test_execute_llm(setup_openai_mock):
@pytest.mark.parametrize("setup_code_executor_mock", [["none"]], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_openai_mock):
def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_model_mock):
"""
Test execute LLM node with jinja2
"""
@ -190,38 +166,15 @@ def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_openai_mock):
credentials = {"openai_api_key": os.environ.get("OPENAI_API_KEY")}
provider_instance = ModelProviderFactory().get_provider_instance("openai")
model_type_instance = provider_instance.get_model_instance(ModelType.LLM)
provider_model_bundle = ProviderModelBundle(
configuration=ProviderConfiguration(
tenant_id="1",
provider=provider_instance.get_provider_schema(),
preferred_provider_type=ProviderType.CUSTOM,
using_provider_type=ProviderType.CUSTOM,
system_configuration=SystemConfiguration(enabled=False),
custom_configuration=CustomConfiguration(provider=CustomProviderConfiguration(credentials=credentials)),
model_settings=[],
),
model_type_instance=model_type_instance,
)
model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo")
model_schema = model_type_instance.get_model_schema("gpt-3.5-turbo")
assert model_schema is not None
model_config = ModelConfigWithCredentialsEntity(
model="gpt-3.5-turbo",
provider="openai",
mode="chat",
credentials=credentials,
parameters={},
model_schema=model_schema,
provider_model_bundle=provider_model_bundle,
)
# Mock db.session.close()
db.session.close = MagicMock()
node._fetch_model_config = MagicMock(return_value=(model_instance, model_config))
node._fetch_model_config = get_mocked_fetch_model_config(
provider="langgenius/openai/openai",
model="gpt-3.5-turbo",
mode="chat",
credentials=credentials,
)
# execute node
result = node._run()

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

@ -4,14 +4,7 @@ import uuid
from typing import Optional
from unittest.mock import MagicMock
import pytest
from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity
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 core.app.entities.app_invoke_entities import InvokeFrom
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.enums import SystemVariableKey
from core.workflow.graph_engine.entities.graph import Graph
@ -20,53 +13,11 @@ from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntime
from core.workflow.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode
from extensions.ext_database import db
from models.enums import UserFrom
from models.provider import ProviderType
from tests.integration_tests.workflow.nodes.__mock.model import get_mocked_fetch_model_config
"""FOR MOCK FIXTURES, DO NOT REMOVE"""
from models.workflow import WorkflowNodeExecutionStatus, WorkflowType
from tests.integration_tests.model_runtime.__mock.anthropic import setup_anthropic_mock
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
def get_mocked_fetch_model_config(
provider: str,
model: str,
mode: str,
credentials: dict,
):
model_provider_factory = ModelProviderFactory(tenant_id="test_tenant")
model_type_instance = model_provider_factory.get_model_type_instance(provider, ModelType.LLM)
provider_model_bundle = ProviderModelBundle(
configuration=ProviderConfiguration(
tenant_id="1",
provider=model_provider_factory.get_provider_schema(provider),
preferred_provider_type=ProviderType.CUSTOM,
using_provider_type=ProviderType.CUSTOM,
system_configuration=SystemConfiguration(enabled=False),
custom_configuration=CustomConfiguration(provider=CustomProviderConfiguration(credentials=credentials)),
model_settings=[],
),
model_type_instance=model_type_instance,
)
model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model=model)
model_schema = model_provider_factory.get_model_schema(
provider=provider,
model_type=model_type_instance.model_type,
model=model,
credentials=credentials,
)
assert model_schema is not None
model_config = ModelConfigWithCredentialsEntity(
model=model,
provider=provider,
mode=mode,
credentials=credentials,
parameters={},
model_schema=model_schema,
provider_model_bundle=provider_model_bundle,
)
return MagicMock(return_value=(model_instance, model_config))
from tests.integration_tests.model_runtime.__mock.plugin_daemon import setup_model_mock
def get_mocked_fetch_memory(memory_text: str):
@ -133,8 +84,7 @@ def init_parameter_extractor_node(config: dict):
)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_function_calling_parameter_extractor(setup_openai_mock):
def test_function_calling_parameter_extractor(setup_model_mock):
"""
Test function calling for parameter extractor.
"""
@ -144,7 +94,12 @@ def test_function_calling_parameter_extractor(setup_openai_mock):
"data": {
"title": "123",
"type": "parameter-extractor",
"model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}},
"model": {
"provider": "langgenius/openai/openai",
"name": "gpt-3.5-turbo",
"mode": "chat",
"completion_params": {},
},
"query": ["sys", "query"],
"parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
"instruction": "",
@ -155,25 +110,13 @@ def test_function_calling_parameter_extractor(setup_openai_mock):
)
node._fetch_model_config = get_mocked_fetch_model_config(
provider="openai",
provider="langgenius/openai/openai",
model="gpt-3.5-turbo",
mode="chat",
credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
)
db.session.close = MagicMock()
# construct variable pool
pool = VariablePool(
system_variables={
SystemVariableKey.QUERY: "what's the weather in SF",
SystemVariableKey.FILES: [],
SystemVariableKey.CONVERSATION_ID: "abababa",
SystemVariableKey.USER_ID: "aaa",
},
user_inputs={},
environment_variables=[],
)
result = node._run()
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
@ -182,8 +125,7 @@ def test_function_calling_parameter_extractor(setup_openai_mock):
assert result.outputs.get("__reason") == None
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_instructions(setup_openai_mock):
def test_instructions(setup_model_mock):
"""
Test chat parameter extractor.
"""
@ -193,7 +135,12 @@ def test_instructions(setup_openai_mock):
"data": {
"title": "123",
"type": "parameter-extractor",
"model": {"provider": "openai", "name": "gpt-3.5-turbo", "mode": "chat", "completion_params": {}},
"model": {
"provider": "langgenius/openai/openai",
"name": "gpt-3.5-turbo",
"mode": "chat",
"completion_params": {},
},
"query": ["sys", "query"],
"parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
"reasoning_mode": "function_call",
@ -204,7 +151,7 @@ def test_instructions(setup_openai_mock):
)
node._fetch_model_config = get_mocked_fetch_model_config(
provider="openai",
provider="langgenius/openai/openai",
model="gpt-3.5-turbo",
mode="chat",
credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
@ -228,8 +175,7 @@ def test_instructions(setup_openai_mock):
assert "what's the weather in SF" in prompt.get("text")
@pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)
def test_chat_parameter_extractor(setup_anthropic_mock):
def test_chat_parameter_extractor(setup_model_mock):
"""
Test chat parameter extractor.
"""
@ -239,7 +185,12 @@ def test_chat_parameter_extractor(setup_anthropic_mock):
"data": {
"title": "123",
"type": "parameter-extractor",
"model": {"provider": "anthropic", "name": "claude-2", "mode": "chat", "completion_params": {}},
"model": {
"provider": "langgenius/openai/openai",
"name": "gpt-3.5-turbo",
"mode": "chat",
"completion_params": {},
},
"query": ["sys", "query"],
"parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
"reasoning_mode": "prompt",
@ -250,10 +201,10 @@ def test_chat_parameter_extractor(setup_anthropic_mock):
)
node._fetch_model_config = get_mocked_fetch_model_config(
provider="anthropic",
model="claude-2",
provider="langgenius/openai/openai",
model="gpt-3.5-turbo",
mode="chat",
credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")},
credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
)
db.session.close = MagicMock()
@ -275,8 +226,7 @@ def test_chat_parameter_extractor(setup_anthropic_mock):
assert '<structure>\n{"type": "object"' in prompt.get("text")
@pytest.mark.parametrize("setup_openai_mock", [["completion"]], indirect=True)
def test_completion_parameter_extractor(setup_openai_mock):
def test_completion_parameter_extractor(setup_model_mock):
"""
Test completion parameter extractor.
"""
@ -287,7 +237,7 @@ def test_completion_parameter_extractor(setup_openai_mock):
"title": "123",
"type": "parameter-extractor",
"model": {
"provider": "openai",
"provider": "langgenius/openai/openai",
"name": "gpt-3.5-turbo-instruct",
"mode": "completion",
"completion_params": {},
@ -302,7 +252,7 @@ def test_completion_parameter_extractor(setup_openai_mock):
)
node._fetch_model_config = get_mocked_fetch_model_config(
provider="openai",
provider="langgenius/openai/openai",
model="gpt-3.5-turbo-instruct",
mode="completion",
credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
@ -335,7 +285,7 @@ def test_extract_json_response():
"title": "123",
"type": "parameter-extractor",
"model": {
"provider": "openai",
"provider": "langgenius/openai/openai",
"name": "gpt-3.5-turbo-instruct",
"mode": "completion",
"completion_params": {},
@ -361,8 +311,7 @@ def test_extract_json_response():
assert result["location"] == "kawaii"
@pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)
def test_chat_parameter_extractor_with_memory(setup_anthropic_mock):
def test_chat_parameter_extractor_with_memory(setup_model_mock):
"""
Test chat parameter extractor with memory.
"""
@ -372,7 +321,12 @@ def test_chat_parameter_extractor_with_memory(setup_anthropic_mock):
"data": {
"title": "123",
"type": "parameter-extractor",
"model": {"provider": "anthropic", "name": "claude-2", "mode": "chat", "completion_params": {}},
"model": {
"provider": "langgenius/openai/openai",
"name": "gpt-3.5-turbo",
"mode": "chat",
"completion_params": {},
},
"query": ["sys", "query"],
"parameters": [{"name": "location", "type": "string", "description": "location", "required": True}],
"reasoning_mode": "prompt",
@ -383,10 +337,10 @@ def test_chat_parameter_extractor_with_memory(setup_anthropic_mock):
)
node._fetch_model_config = get_mocked_fetch_model_config(
provider="anthropic",
model="claude-2",
provider="langgenius/openai/openai",
model="gpt-3.5-turbo",
mode="chat",
credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")},
credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
)
node._fetch_memory = get_mocked_fetch_memory("customized memory")
db.session.close = MagicMock()

60
api/tests/integration_tests/workflow/nodes/test_tool.py
View File

@ -1,13 +1,15 @@
import time
import uuid
from unittest.mock import MagicMock
from core.app.entities.app_invoke_entities import InvokeFrom
from core.workflow.entities.node_entities import NodeRunResult
from core.tools.utils.configuration import ToolParameterConfigurationManager
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.enums import SystemVariableKey
from core.workflow.graph_engine.entities.graph import Graph
from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams
from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState
from core.workflow.nodes.event.event import RunCompletedEvent
from core.workflow.nodes.tool.tool_node import ToolNode
from models.enums import UserFrom
from models.workflow import WorkflowNodeExecutionStatus, WorkflowType
@ -63,31 +65,28 @@ def test_tool_variable_invoke():
"data": {
"title": "a",
"desc": "a",
"provider_id": "maths",
"provider_id": "time",
"provider_type": "builtin",
"provider_name": "maths",
"tool_name": "eval_expression",
"tool_label": "eval_expression",
"provider_name": "time",
"tool_name": "current_time",
"tool_label": "current_time",
"tool_configurations": {},
"tool_parameters": {
"expression": {
"type": "variable",
"value": ["1", "123", "args1"],
}
},
"tool_parameters": {},
},
}
)
ToolParameterConfigurationManager.decrypt_tool_parameters = MagicMock(return_value={"format": "%Y-%m-%d %H:%M:%S"})
node.graph_runtime_state.variable_pool.add(["1", "123", "args1"], "1+1")
# execute node
result = node._run()
assert isinstance(result, NodeRunResult)
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
assert result.outputs is not None
assert "2" in result.outputs["text"]
assert result.outputs["files"] == []
for item in result:
if isinstance(item, RunCompletedEvent):
assert item.run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED
assert item.run_result.outputs is not None
assert item.run_result.outputs.get("text") is not None
def test_tool_mixed_invoke():
@ -97,28 +96,25 @@ def test_tool_mixed_invoke():
"data": {
"title": "a",
"desc": "a",
"provider_id": "maths",
"provider_id": "time",
"provider_type": "builtin",
"provider_name": "maths",
"tool_name": "eval_expression",
"tool_label": "eval_expression",
"tool_configurations": {},
"tool_parameters": {
"expression": {
"type": "mixed",
"value": "{{#1.args1#}}",
}
"provider_name": "time",
"tool_name": "current_time",
"tool_label": "current_time",
"tool_configurations": {
"format": "%Y-%m-%d %H:%M:%S",
},
"tool_parameters": {},
},
}
)
node.graph_runtime_state.variable_pool.add(["1", "args1"], "1+1")
ToolParameterConfigurationManager.decrypt_tool_parameters = MagicMock(return_value={"format": "%Y-%m-%d %H:%M:%S"})
# execute node
result = node._run()
assert isinstance(result, NodeRunResult)
assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
assert result.outputs is not None
assert "2" in result.outputs["text"]
assert result.outputs["files"] == []
for item in result:
if isinstance(item, RunCompletedEvent):
assert item.run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED
assert item.run_result.outputs is not None
assert item.run_result.outputs.get("text") is not None

2
api/tests/unit_tests/core/helper/test_marketplace.py
View File

@ -2,6 +2,6 @@ from core.helper.marketplace import download_plugin_pkg
def test_download_plugin_pkg():
pkg = download_plugin_pkg("yeuoly/google:0.0.1@4ff79ee644987e5b744d9c5b7a735d459fe66f26b28724326a7834d7e459e708")
pkg = download_plugin_pkg("langgenius/bing:0.0.1@e58735424d2104f208c2bd683c5142e0332045b425927067acf432b26f3d970b")
assert pkg is not None
assert len(pkg) > 0

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

@ -1,52 +1,52 @@
from unittest.mock import MagicMock
# from unittest.mock import MagicMock
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 core.prompt.prompt_transform import PromptTransform
# 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 core.prompt.prompt_transform import PromptTransform
def test__calculate_rest_token():
model_schema_mock = MagicMock(spec=AIModelEntity)
parameter_rule_mock = MagicMock(spec=ParameterRule)
parameter_rule_mock.name = "max_tokens"
model_schema_mock.parameter_rules = [parameter_rule_mock]
model_schema_mock.model_properties = {ModelPropertyKey.CONTEXT_SIZE: 62}
# def test__calculate_rest_token():
# model_schema_mock = MagicMock(spec=AIModelEntity)
# parameter_rule_mock = MagicMock(spec=ParameterRule)
# parameter_rule_mock.name = "max_tokens"
# model_schema_mock.parameter_rules = [parameter_rule_mock]
# model_schema_mock.model_properties = {ModelPropertyKey.CONTEXT_SIZE: 62}
large_language_model_mock = MagicMock(spec=LargeLanguageModel)
large_language_model_mock.get_num_tokens.return_value = 6
# large_language_model_mock = MagicMock(spec=LargeLanguageModel)
# large_language_model_mock.get_num_tokens.return_value = 6
provider_mock = MagicMock(spec=ProviderEntity)
provider_mock.provider = "openai"
# provider_mock = MagicMock(spec=ProviderEntity)
# provider_mock.provider = "openai"
provider_configuration_mock = MagicMock(spec=ProviderConfiguration)
provider_configuration_mock.provider = provider_mock
provider_configuration_mock.model_settings = None
# provider_configuration_mock = MagicMock(spec=ProviderConfiguration)
# provider_configuration_mock.provider = provider_mock
# provider_configuration_mock.model_settings = None
provider_model_bundle_mock = MagicMock(spec=ProviderModelBundle)
provider_model_bundle_mock.model_type_instance = large_language_model_mock
provider_model_bundle_mock.configuration = provider_configuration_mock
# provider_model_bundle_mock = MagicMock(spec=ProviderModelBundle)
# provider_model_bundle_mock.model_type_instance = large_language_model_mock
# provider_model_bundle_mock.configuration = provider_configuration_mock
model_config_mock = MagicMock(spec=ModelConfigEntity)
model_config_mock.model = "gpt-4"
model_config_mock.credentials = {}
model_config_mock.parameters = {"max_tokens": 50}
model_config_mock.model_schema = model_schema_mock
model_config_mock.provider_model_bundle = provider_model_bundle_mock
# model_config_mock = MagicMock(spec=ModelConfigEntity)
# model_config_mock.model = "gpt-4"
# model_config_mock.credentials = {}
# model_config_mock.parameters = {"max_tokens": 50}
# model_config_mock.model_schema = model_schema_mock
# model_config_mock.provider_model_bundle = provider_model_bundle_mock
prompt_transform = PromptTransform()
# prompt_transform = PromptTransform()
prompt_messages = [UserPromptMessage(content="Hello, how are you?")]
rest_tokens = prompt_transform._calculate_rest_token(prompt_messages, model_config_mock)
# prompt_messages = [UserPromptMessage(content="Hello, how are you?")]
# rest_tokens = prompt_transform._calculate_rest_token(prompt_messages, model_config_mock)
# Validate based on the mock configuration and expected logic
expected_rest_tokens = (
model_schema_mock.model_properties[ModelPropertyKey.CONTEXT_SIZE]
- model_config_mock.parameters["max_tokens"]
- large_language_model_mock.get_num_tokens.return_value
)
assert rest_tokens == expected_rest_tokens
assert rest_tokens == 6
# # Validate based on the mock configuration and expected logic
# expected_rest_tokens = (
# model_schema_mock.model_properties[ModelPropertyKey.CONTEXT_SIZE]
# - model_config_mock.parameters["max_tokens"]
# - large_language_model_mock.get_num_tokens.return_value
# )
# assert rest_tokens == expected_rest_tokens
# assert rest_tokens == 6

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

@ -1,186 +1,190 @@
from core.entities.provider_entities import ModelSettings
from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
from core.provider_manager import ProviderManager
from models.provider import LoadBalancingModelConfig, ProviderModelSetting
# from core.entities.provider_entities import ModelSettings
# from core.model_runtime.entities.model_entities import ModelType
# from core.model_runtime.model_providers.model_provider_factory import ModelProviderFactory
# from core.provider_manager import ProviderManager
# from models.provider import LoadBalancingModelConfig, ProviderModelSetting
def test__to_model_settings(mocker):
# Get all provider entities
model_provider_factory = ModelProviderFactory("test_tenant")
provider_entities = model_provider_factory.get_providers()
# def test__to_model_settings(mocker):
# # Get all provider entities
# model_provider_factory = ModelProviderFactory("test_tenant")
# provider_entities = model_provider_factory.get_providers()
provider_entity = None
for provider in provider_entities:
if provider.provider == "openai":
provider_entity = provider
# provider_entity = None
# for provider in provider_entities:
# if provider.provider == "openai":
# provider_entity = provider
# Mocking the inputs
provider_model_settings = [
ProviderModelSetting(
id="id",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=True,
)
]
load_balancing_model_configs = [
LoadBalancingModelConfig(
id="id1",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
name="__inherit__",
encrypted_config=None,
enabled=True,
),
LoadBalancingModelConfig(
id="id2",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
name="first",
encrypted_config='{"openai_api_key": "fake_key"}',
enabled=True,
),
]
# # Mocking the inputs
# provider_model_settings = [
# ProviderModelSetting(
# id="id",
# tenant_id="tenant_id",
# provider_name="openai",
# model_name="gpt-4",
# model_type="text-generation",
# enabled=True,
# load_balancing_enabled=True,
# )
# ]
# load_balancing_model_configs = [
# LoadBalancingModelConfig(
# id="id1",
# tenant_id="tenant_id",
# provider_name="openai",
# model_name="gpt-4",
# model_type="text-generation",
# name="__inherit__",
# encrypted_config=None,
# enabled=True,
# ),
# LoadBalancingModelConfig(
# id="id2",
# tenant_id="tenant_id",
# provider_name="openai",
# model_name="gpt-4",
# model_type="text-generation",
# name="first",
# encrypted_config='{"openai_api_key": "fake_key"}',
# enabled=True,
# ),
# ]
mocker.patch(
"core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
)
# mocker.patch(
# "core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
# )
provider_manager = ProviderManager()
# provider_manager = ProviderManager()
# Running the method
result = provider_manager._to_model_settings(provider_entity, provider_model_settings, load_balancing_model_configs)
# # Running the method
# result = provider_manager._to_model_settings(provider_entity,
# provider_model_settings, load_balancing_model_configs)
# Asserting that the result is as expected
assert len(result) == 1
assert isinstance(result[0], ModelSettings)
assert result[0].model == "gpt-4"
assert result[0].model_type == ModelType.LLM
assert result[0].enabled is True
assert len(result[0].load_balancing_configs) == 2
assert result[0].load_balancing_configs[0].name == "__inherit__"
assert result[0].load_balancing_configs[1].name == "first"
# # Asserting that the result is as expected
# assert len(result) == 1
# assert isinstance(result[0], ModelSettings)
# assert result[0].model == "gpt-4"
# assert result[0].model_type == ModelType.LLM
# assert result[0].enabled is True
# assert len(result[0].load_balancing_configs) == 2
# assert result[0].load_balancing_configs[0].name == "__inherit__"
# assert result[0].load_balancing_configs[1].name == "first"
def test__to_model_settings_only_one_lb(mocker):
# Get all provider entities
model_provider_factory = ModelProviderFactory("test_tenant")
provider_entities = model_provider_factory.get_providers()
# def test__to_model_settings_only_one_lb(mocker):
# # Get all provider entities
# model_provider_factory = ModelProviderFactory("test_tenant")
# provider_entities = model_provider_factory.get_providers()
provider_entity = None
for provider in provider_entities:
if provider.provider == "openai":
provider_entity = provider
# provider_entity = None
# for provider in provider_entities:
# if provider.provider == "openai":
# provider_entity = provider
# Mocking the inputs
provider_model_settings = [
ProviderModelSetting(
id="id",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=True,
)
]
load_balancing_model_configs = [
LoadBalancingModelConfig(
id="id1",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
name="__inherit__",
encrypted_config=None,
enabled=True,
)
]
# # Mocking the inputs
# provider_model_settings = [
# ProviderModelSetting(
# id="id",
# tenant_id="tenant_id",
# provider_name="openai",
# model_name="gpt-4",
# model_type="text-generation",
# enabled=True,
# load_balancing_enabled=True,
# )
# ]
# load_balancing_model_configs = [
# LoadBalancingModelConfig(
# id="id1",
# tenant_id="tenant_id",
# provider_name="openai",
# model_name="gpt-4",
# model_type="text-generation",
# name="__inherit__",
# encrypted_config=None,
# enabled=True,
# )
# ]
mocker.patch(
"core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
)
# mocker.patch(
# "core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
# )
provider_manager = ProviderManager()
# provider_manager = ProviderManager()
# Running the method
result = provider_manager._to_model_settings(provider_entity, provider_model_settings, load_balancing_model_configs)
# # Running the method
# result = provider_manager._to_model_settings(
# provider_entity, provider_model_settings, load_balancing_model_configs)
# Asserting that the result is as expected
assert len(result) == 1
assert isinstance(result[0], ModelSettings)
assert result[0].model == "gpt-4"
assert result[0].model_type == ModelType.LLM
assert result[0].enabled is True
assert len(result[0].load_balancing_configs) == 0
# # Asserting that the result is as expected
# assert len(result) == 1
# assert isinstance(result[0], ModelSettings)
# assert result[0].model == "gpt-4"
# assert result[0].model_type == ModelType.LLM
# assert result[0].enabled is True
# assert len(result[0].load_balancing_configs) == 0
def test__to_model_settings_lb_disabled(mocker):
# Get all provider entities
model_provider_factory = ModelProviderFactory("test_tenant")
provider_entities = model_provider_factory.get_providers()
# def test__to_model_settings_lb_disabled(mocker):
# # Get all provider entities
# model_provider_factory = ModelProviderFactory("test_tenant")
# provider_entities = model_provider_factory.get_providers()
provider_entity = None
for provider in provider_entities:
if provider.provider == "openai":
provider_entity = provider
# provider_entity = None
# for provider in provider_entities:
# if provider.provider == "openai":
# provider_entity = provider
# Mocking the inputs
provider_model_settings = [
ProviderModelSetting(
id="id",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
enabled=True,
load_balancing_enabled=False,
)
]
load_balancing_model_configs = [
LoadBalancingModelConfig(
id="id1",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
name="__inherit__",
encrypted_config=None,
enabled=True,
),
LoadBalancingModelConfig(
id="id2",
tenant_id="tenant_id",
provider_name="openai",
model_name="gpt-4",
model_type="text-generation",
name="first",
encrypted_config='{"openai_api_key": "fake_key"}',
enabled=True,
),
]
# # Mocking the inputs
# provider_model_settings = [
# ProviderModelSetting(
# id="id",
# tenant_id="tenant_id",
# provider_name="openai",
# model_name="gpt-4",
# model_type="text-generation",
# enabled=True,
# load_balancing_enabled=False,
# )
# ]
# load_balancing_model_configs = [
# LoadBalancingModelConfig(
# id="id1",
# tenant_id="tenant_id",
# provider_name="openai",
# model_name="gpt-4",
# model_type="text-generation",
# name="__inherit__",
# encrypted_config=None,
# enabled=True,
# ),
# LoadBalancingModelConfig(
# id="id2",
# tenant_id="tenant_id",
# provider_name="openai",
# model_name="gpt-4",
# model_type="text-generation",
# name="first",
# encrypted_config='{"openai_api_key": "fake_key"}',
# enabled=True,
# ),
# ]
mocker.patch(
"core.helper.model_provider_cache.ProviderCredentialsCache.get", return_value={"openai_api_key": "fake_key"}
)
# mocker.patch(
# "core.helper.model_provider_cache.ProviderCredentialsCache.get",
# return_value={"openai_api_key": "fake_key"}
# )
provider_manager = ProviderManager()
# provider_manager = ProviderManager()
# Running the method
result = provider_manager._to_model_settings(provider_entity, provider_model_settings, load_balancing_model_configs)
# # Running the method
# result = provider_manager._to_model_settings(provider_entity,
# provider_model_settings, load_balancing_model_configs)
# Asserting that the result is as expected
assert len(result) == 1
assert isinstance(result[0], ModelSettings)
assert result[0].model == "gpt-4"
assert result[0].model_type == ModelType.LLM
assert result[0].enabled is True
assert len(result[0].load_balancing_configs) == 0
# # Asserting that the result is as expected
# assert len(result) == 1
# assert isinstance(result[0], ModelSettings)
# assert result[0].model == "gpt-4"
# assert result[0].model_type == ModelType.LLM
# assert result[0].enabled is True
# assert len(result[0].load_balancing_configs) == 0

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

@ -3,24 +3,20 @@ from typing import Optional
import pytest
from configs import dify_config
from core.app.entities.app_invoke_entities import InvokeFrom, ModelConfigWithCredentialsEntity
from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
from core.entities.provider_entities import CustomConfiguration, SystemConfiguration
from core.file import File, FileTransferMethod, FileType
from core.model_runtime.entities.common_entities import I18nObject
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
ImagePromptMessageContent,
PromptMessage,
PromptMessageRole,
SystemPromptMessage,
TextPromptMessageContent,
UserPromptMessage,
)
from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelFeature, ModelType
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 core.variables import ArrayAnySegment, ArrayFileSegment, NoneSegment
from core.workflow.entities.variable_pool import VariablePool
from core.workflow.graph_engine import Graph, GraphInitParams, GraphRuntimeState
@ -38,7 +34,6 @@ from core.workflow.nodes.llm.node import LLMNode
from models.enums import UserFrom
from models.provider import ProviderType
from models.workflow import WorkflowType
from tests.unit_tests.core.workflow.nodes.llm.test_scenarios import LLMNodeTestScenario
class MockTokenBufferMemory:
@ -112,22 +107,21 @@ def llm_node():
@pytest.fixture
def model_config():
# Create actual provider and model type instances
model_provider_factory = ModelProviderFactory()
provider_instance = model_provider_factory.get_provider_instance("openai")
model_type_instance = provider_instance.get_model_instance(ModelType.LLM)
model_provider_factory = ModelProviderFactory(tenant_id="test")
provider_instance = model_provider_factory.get_plugin_model_provider("openai")
model_type_instance = model_provider_factory.get_model_type_instance("openai", ModelType.LLM)
# Create a ProviderModelBundle
provider_model_bundle = ProviderModelBundle(
configuration=ProviderConfiguration(
tenant_id="1",
provider=provider_instance.get_provider_schema(),
provider=provider_instance,
preferred_provider_type=ProviderType.CUSTOM,
using_provider_type=ProviderType.CUSTOM,
system_configuration=SystemConfiguration(enabled=False),
custom_configuration=CustomConfiguration(provider=None),
model_settings=[],
),
provider_instance=provider_instance,
model_type_instance=model_type_instance,
)
@ -211,236 +205,240 @@ def test_fetch_files_with_non_existent_variable(llm_node):
assert result == []
def test_fetch_prompt_messages__vison_disabled(faker, llm_node, model_config):
prompt_template = []
llm_node.node_data.prompt_template = prompt_template
# def test_fetch_prompt_messages__vison_disabled(faker, llm_node, model_config):
# TODO: Add test
# pass
# prompt_template = []
# llm_node.node_data.prompt_template = prompt_template
fake_vision_detail = faker.random_element(
[ImagePromptMessageContent.DETAIL.HIGH, ImagePromptMessageContent.DETAIL.LOW]
)
fake_remote_url = faker.url()
files = [
File(
id="1",
tenant_id="test",
type=FileType.IMAGE,
filename="test1.jpg",
transfer_method=FileTransferMethod.REMOTE_URL,
remote_url=fake_remote_url,
storage_key="",
)
]
# fake_vision_detail = faker.random_element(
# [ImagePromptMessageContent.DETAIL.HIGH, ImagePromptMessageContent.DETAIL.LOW]
# )
# fake_remote_url = faker.url()
# files = [
# File(
# id="1",
# tenant_id="test",
# type=FileType.IMAGE,
# filename="test1.jpg",
# transfer_method=FileTransferMethod.REMOTE_URL,
# remote_url=fake_remote_url,
# storage_key="",
# )
# ]
fake_query = faker.sentence()
# fake_query = faker.sentence()
prompt_messages, _ = llm_node._fetch_prompt_messages(
sys_query=fake_query,
sys_files=files,
context=None,
memory=None,
model_config=model_config,
prompt_template=prompt_template,
memory_config=None,
vision_enabled=False,
vision_detail=fake_vision_detail,
variable_pool=llm_node.graph_runtime_state.variable_pool,
jinja2_variables=[],
)
# prompt_messages, _ = llm_node._fetch_prompt_messages(
# sys_query=fake_query,
# sys_files=files,
# context=None,
# memory=None,
# model_config=model_config,
# prompt_template=prompt_template,
# memory_config=None,
# vision_enabled=False,
# vision_detail=fake_vision_detail,
# variable_pool=llm_node.graph_runtime_state.variable_pool,
# jinja2_variables=[],
# )
assert prompt_messages == [UserPromptMessage(content=fake_query)]
# assert prompt_messages == [UserPromptMessage(content=fake_query)]
def test_fetch_prompt_messages__basic(faker, llm_node, model_config):
# Setup dify config
dify_config.MULTIMODAL_SEND_FORMAT = "url"
# def test_fetch_prompt_messages__basic(faker, llm_node, model_config):
# TODO: Add test
# pass
# Setup dify config
# dify_config.MULTIMODAL_SEND_FORMAT = "url"
# Generate fake values for prompt template
fake_assistant_prompt = faker.sentence()
fake_query = faker.sentence()
fake_context = faker.sentence()
fake_window_size = faker.random_int(min=1, max=3)
fake_vision_detail = faker.random_element(
[ImagePromptMessageContent.DETAIL.HIGH, ImagePromptMessageContent.DETAIL.LOW]
)
fake_remote_url = faker.url()
# # Generate fake values for prompt template
# fake_assistant_prompt = faker.sentence()
# fake_query = faker.sentence()
# fake_context = faker.sentence()
# fake_window_size = faker.random_int(min=1, max=3)
# fake_vision_detail = faker.random_element(
# [ImagePromptMessageContent.DETAIL.HIGH, ImagePromptMessageContent.DETAIL.LOW]
# )
# fake_remote_url = faker.url()
# Setup mock memory with history messages
mock_history = [
UserPromptMessage(content=faker.sentence()),
AssistantPromptMessage(content=faker.sentence()),
UserPromptMessage(content=faker.sentence()),
AssistantPromptMessage(content=faker.sentence()),
UserPromptMessage(content=faker.sentence()),
AssistantPromptMessage(content=faker.sentence()),
]
# # Setup mock memory with history messages
# mock_history = [
# UserPromptMessage(content=faker.sentence()),
# AssistantPromptMessage(content=faker.sentence()),
# UserPromptMessage(content=faker.sentence()),
# AssistantPromptMessage(content=faker.sentence()),
# UserPromptMessage(content=faker.sentence()),
# AssistantPromptMessage(content=faker.sentence()),
# ]
# Setup memory configuration
memory_config = MemoryConfig(
role_prefix=MemoryConfig.RolePrefix(user="Human", assistant="Assistant"),
window=MemoryConfig.WindowConfig(enabled=True, size=fake_window_size),
query_prompt_template=None,
)
# # Setup memory configuration
# memory_config = MemoryConfig(
# role_prefix=MemoryConfig.RolePrefix(user="Human", assistant="Assistant"),
# window=MemoryConfig.WindowConfig(enabled=True, size=fake_window_size),
# query_prompt_template=None,
# )
memory = MockTokenBufferMemory(history_messages=mock_history)
# memory = MockTokenBufferMemory(history_messages=mock_history)
# Test scenarios covering different file input combinations
test_scenarios = [
LLMNodeTestScenario(
description="No files",
sys_query=fake_query,
sys_files=[],
features=[],
vision_enabled=False,
vision_detail=None,
window_size=fake_window_size,
prompt_template=[
LLMNodeChatModelMessage(
text=fake_context,
role=PromptMessageRole.SYSTEM,
edition_type="basic",
),
LLMNodeChatModelMessage(
text="{#context#}",
role=PromptMessageRole.USER,
edition_type="basic",
),
LLMNodeChatModelMessage(
text=fake_assistant_prompt,
role=PromptMessageRole.ASSISTANT,
edition_type="basic",
),
],
expected_messages=[
SystemPromptMessage(content=fake_context),
UserPromptMessage(content=fake_context),
AssistantPromptMessage(content=fake_assistant_prompt),
]
+ mock_history[fake_window_size * -2 :]
+ [
UserPromptMessage(content=fake_query),
],
),
LLMNodeTestScenario(
description="User files",
sys_query=fake_query,
sys_files=[
File(
tenant_id="test",
type=FileType.IMAGE,
filename="test1.jpg",
transfer_method=FileTransferMethod.REMOTE_URL,
remote_url=fake_remote_url,
extension=".jpg",
mime_type="image/jpg",
storage_key="",
)
],
vision_enabled=True,
vision_detail=fake_vision_detail,
features=[ModelFeature.VISION],
window_size=fake_window_size,
prompt_template=[
LLMNodeChatModelMessage(
text=fake_context,
role=PromptMessageRole.SYSTEM,
edition_type="basic",
),
LLMNodeChatModelMessage(
text="{#context#}",
role=PromptMessageRole.USER,
edition_type="basic",
),
LLMNodeChatModelMessage(
text=fake_assistant_prompt,
role=PromptMessageRole.ASSISTANT,
edition_type="basic",
),
],
expected_messages=[
SystemPromptMessage(content=fake_context),
UserPromptMessage(content=fake_context),
AssistantPromptMessage(content=fake_assistant_prompt),
]
+ mock_history[fake_window_size * -2 :]
+ [
UserPromptMessage(
content=[
TextPromptMessageContent(data=fake_query),
ImagePromptMessageContent(
url=fake_remote_url, mime_type="image/jpg", format="jpg", detail=fake_vision_detail
),
]
),
],
),
LLMNodeTestScenario(
description="Prompt template with variable selector of File",
sys_query=fake_query,
sys_files=[],
vision_enabled=False,
vision_detail=fake_vision_detail,
features=[ModelFeature.VISION],
window_size=fake_window_size,
prompt_template=[
LLMNodeChatModelMessage(
text="{{#input.image#}}",
role=PromptMessageRole.USER,
edition_type="basic",
),
],
expected_messages=[
UserPromptMessage(
content=[
ImagePromptMessageContent(
url=fake_remote_url, mime_type="image/jpg", format="jpg", detail=fake_vision_detail
),
]
),
]
+ mock_history[fake_window_size * -2 :]
+ [UserPromptMessage(content=fake_query)],
file_variables={
"input.image": File(
tenant_id="test",
type=FileType.IMAGE,
filename="test1.jpg",
transfer_method=FileTransferMethod.REMOTE_URL,
remote_url=fake_remote_url,
extension=".jpg",
mime_type="image/jpg",
storage_key="",
)
},
),
]
# # Test scenarios covering different file input combinations
# test_scenarios = [
# LLMNodeTestScenario(
# description="No files",
# sys_query=fake_query,
# sys_files=[],
# features=[],
# vision_enabled=False,
# vision_detail=None,
# window_size=fake_window_size,
# prompt_template=[
# LLMNodeChatModelMessage(
# text=fake_context,
# role=PromptMessageRole.SYSTEM,
# edition_type="basic",
# ),
# LLMNodeChatModelMessage(
# text="{#context#}",
# role=PromptMessageRole.USER,
# edition_type="basic",
# ),
# LLMNodeChatModelMessage(
# text=fake_assistant_prompt,
# role=PromptMessageRole.ASSISTANT,
# edition_type="basic",
# ),
# ],
# expected_messages=[
# SystemPromptMessage(content=fake_context),
# UserPromptMessage(content=fake_context),
# AssistantPromptMessage(content=fake_assistant_prompt),
# ]
# + mock_history[fake_window_size * -2 :]
# + [
# UserPromptMessage(content=fake_query),
# ],
# ),
# LLMNodeTestScenario(
# description="User files",
# sys_query=fake_query,
# sys_files=[
# File(
# tenant_id="test",
# type=FileType.IMAGE,
# filename="test1.jpg",
# transfer_method=FileTransferMethod.REMOTE_URL,
# remote_url=fake_remote_url,
# extension=".jpg",
# mime_type="image/jpg",
# storage_key="",
# )
# ],
# vision_enabled=True,
# vision_detail=fake_vision_detail,
# features=[ModelFeature.VISION],
# window_size=fake_window_size,
# prompt_template=[
# LLMNodeChatModelMessage(
# text=fake_context,
# role=PromptMessageRole.SYSTEM,
# edition_type="basic",
# ),
# LLMNodeChatModelMessage(
# text="{#context#}",
# role=PromptMessageRole.USER,
# edition_type="basic",
# ),
# LLMNodeChatModelMessage(
# text=fake_assistant_prompt,
# role=PromptMessageRole.ASSISTANT,
# edition_type="basic",
# ),
# ],
# expected_messages=[
# SystemPromptMessage(content=fake_context),
# UserPromptMessage(content=fake_context),
# AssistantPromptMessage(content=fake_assistant_prompt),
# ]
# + mock_history[fake_window_size * -2 :]
# + [
# UserPromptMessage(
# content=[
# TextPromptMessageContent(data=fake_query),
# ImagePromptMessageContent(
# url=fake_remote_url, mime_type="image/jpg", format="jpg", detail=fake_vision_detail
# ),
# ]
# ),
# ],
# ),
# LLMNodeTestScenario(
# description="Prompt template with variable selector of File",
# sys_query=fake_query,
# sys_files=[],
# vision_enabled=False,
# vision_detail=fake_vision_detail,
# features=[ModelFeature.VISION],
# window_size=fake_window_size,
# prompt_template=[
# LLMNodeChatModelMessage(
# text="{{#input.image#}}",
# role=PromptMessageRole.USER,
# edition_type="basic",
# ),
# ],
# expected_messages=[
# UserPromptMessage(
# content=[
# ImagePromptMessageContent(
# url=fake_remote_url, mime_type="image/jpg", format="jpg", detail=fake_vision_detail
# ),
# ]
# ),
# ]
# + mock_history[fake_window_size * -2 :]
# + [UserPromptMessage(content=fake_query)],
# file_variables={
# "input.image": File(
# tenant_id="test",
# type=FileType.IMAGE,
# filename="test1.jpg",
# transfer_method=FileTransferMethod.REMOTE_URL,
# remote_url=fake_remote_url,
# extension=".jpg",
# mime_type="image/jpg",
# storage_key="",
# )
# },
# ),
# ]
for scenario in test_scenarios:
model_config.model_schema.features = scenario.features
# for scenario in test_scenarios:
# model_config.model_schema.features = scenario.features
for k, v in scenario.file_variables.items():
selector = k.split(".")
llm_node.graph_runtime_state.variable_pool.add(selector, v)
# for k, v in scenario.file_variables.items():
# selector = k.split(".")
# llm_node.graph_runtime_state.variable_pool.add(selector, v)
# Call the method under test
prompt_messages, _ = llm_node._fetch_prompt_messages(
sys_query=scenario.sys_query,
sys_files=scenario.sys_files,
context=fake_context,
memory=memory,
model_config=model_config,
prompt_template=scenario.prompt_template,
memory_config=memory_config,
vision_enabled=scenario.vision_enabled,
vision_detail=scenario.vision_detail,
variable_pool=llm_node.graph_runtime_state.variable_pool,
jinja2_variables=[],
)
# # Call the method under test
# prompt_messages, _ = llm_node._fetch_prompt_messages(
# sys_query=scenario.sys_query,
# sys_files=scenario.sys_files,
# context=fake_context,
# memory=memory,
# model_config=model_config,
# prompt_template=scenario.prompt_template,
# memory_config=memory_config,
# vision_enabled=scenario.vision_enabled,
# vision_detail=scenario.vision_detail,
# variable_pool=llm_node.graph_runtime_state.variable_pool,
# jinja2_variables=[],
# )
# Verify the result
assert len(prompt_messages) == len(scenario.expected_messages), f"Scenario failed: {scenario.description}"
assert (
prompt_messages == scenario.expected_messages
), f"Message content mismatch in scenario: {scenario.description}"
# # Verify the result
# assert len(prompt_messages) == len(scenario.expected_messages), f"Scenario failed: {scenario.description}"
# assert (
# prompt_messages == scenario.expected_messages
# ), f"Message content mismatch in scenario: {scenario.description}"
def test_handle_list_messages_basic(llm_node):

88
api/tests/unit_tests/core/workflow/nodes/test_continue_on_error.py
View File

@ -126,7 +126,7 @@ class ContinueOnErrorTestHelper:
},
}
if default_value:
node["data"]["default_value"] = default_value
node.node_data.default_value = default_value
return node
@staticmethod
@ -331,55 +331,55 @@ def test_http_node_fail_branch_continue_on_error():
assert sum(1 for e in events if isinstance(e, NodeRunStreamChunkEvent)) == 1
def test_tool_node_default_value_continue_on_error():
"""Test tool node with default value error strategy"""
graph_config = {
"edges": DEFAULT_VALUE_EDGE,
"nodes": [
{"data": {"title": "start", "type": "start", "variables": []}, "id": "start"},
{"data": {"title": "answer", "type": "answer", "answer": "{{#node.result#}}"}, "id": "answer"},
ContinueOnErrorTestHelper.get_tool_node(
"default-value", [{"key": "result", "type": "string", "value": "default tool result"}]
),
],
}
# def test_tool_node_default_value_continue_on_error():
# """Test tool node with default value error strategy"""
# graph_config = {
# "edges": DEFAULT_VALUE_EDGE,
# "nodes": [
# {"data": {"title": "start", "type": "start", "variables": []}, "id": "start"},
# {"data": {"title": "answer", "type": "answer", "answer": "{{#node.result#}}"}, "id": "answer"},
# ContinueOnErrorTestHelper.get_tool_node(
# "default-value", [{"key": "result", "type": "string", "value": "default tool result"}]
# ),
# ],
# }
graph_engine = ContinueOnErrorTestHelper.create_test_graph_engine(graph_config)
events = list(graph_engine.run())
# graph_engine = ContinueOnErrorTestHelper.create_test_graph_engine(graph_config)
# events = list(graph_engine.run())
assert any(isinstance(e, NodeRunExceptionEvent) for e in events)
assert any(
isinstance(e, GraphRunPartialSucceededEvent) and e.outputs == {"answer": "default tool result"} for e in events
)
assert sum(1 for e in events if isinstance(e, NodeRunStreamChunkEvent)) == 1
# assert any(isinstance(e, NodeRunExceptionEvent) for e in events)
# assert any(
# isinstance(e, GraphRunPartialSucceededEvent) and e.outputs == {"answer": "default tool result"} for e in events # noqa: E501
# )
# assert sum(1 for e in events if isinstance(e, NodeRunStreamChunkEvent)) == 1
def test_tool_node_fail_branch_continue_on_error():
"""Test HTTP node with fail-branch error strategy"""
graph_config = {
"edges": FAIL_BRANCH_EDGES,
"nodes": [
{"data": {"title": "Start", "type": "start", "variables": []}, "id": "start"},
{
"data": {"title": "success", "type": "answer", "answer": "tool execute successful"},
"id": "success",
},
{
"data": {"title": "error", "type": "answer", "answer": "tool execute failed"},
"id": "error",
},
ContinueOnErrorTestHelper.get_tool_node(),
],
}
# def test_tool_node_fail_branch_continue_on_error():
# """Test HTTP node with fail-branch error strategy"""
# graph_config = {
# "edges": FAIL_BRANCH_EDGES,
# "nodes": [
# {"data": {"title": "Start", "type": "start", "variables": []}, "id": "start"},
# {
# "data": {"title": "success", "type": "answer", "answer": "tool execute successful"},
# "id": "success",
# },
# {
# "data": {"title": "error", "type": "answer", "answer": "tool execute failed"},
# "id": "error",
# },
# ContinueOnErrorTestHelper.get_tool_node(),
# ],
# }
graph_engine = ContinueOnErrorTestHelper.create_test_graph_engine(graph_config)
events = list(graph_engine.run())
# graph_engine = ContinueOnErrorTestHelper.create_test_graph_engine(graph_config)
# events = list(graph_engine.run())
assert any(isinstance(e, NodeRunExceptionEvent) for e in events)
assert any(
isinstance(e, GraphRunPartialSucceededEvent) and e.outputs == {"answer": "tool execute failed"} for e in events
)
assert sum(1 for e in events if isinstance(e, NodeRunStreamChunkEvent)) == 1
# assert any(isinstance(e, NodeRunExceptionEvent) for e in events)
# assert any(
# isinstance(e, GraphRunPartialSucceededEvent) and e.outputs == {"answer": "tool execute failed"} for e in events # noqa: E501
# )
# assert sum(1 for e in events if isinstance(e, NodeRunStreamChunkEvent)) == 1
def test_llm_node_default_value_continue_on_error():