Feat/environment variables in workflow (#6515)

Co-authored-by: JzoNg <jzongcode@gmail.com>

api/tests/unit_tests/app/test_segment.py

@@ -0,0 +1,53 @@
+from core.app.segments import SecretVariable, parser
+from core.helper import encrypter
+from core.workflow.entities.node_entities import SystemVariable
+from core.workflow.entities.variable_pool import VariablePool
+
+
+def test_segment_group_to_text():
+    variable_pool = VariablePool(
+        system_variables={
+            SystemVariable('user_id'): 'fake-user-id',
+        },
+        user_inputs={},
+        environment_variables=[
+            SecretVariable(name='secret_key', value='fake-secret-key'),
+        ],
+    )
+    variable_pool.add(('node_id', 'custom_query'), 'fake-user-query')
+    template = (
+        'Hello, {{#sys.user_id#}}! Your query is {{#node_id.custom_query#}}. And your key is {{#env.secret_key#}}.'
+    )
+    segments_group = parser.convert_template(template=template, variable_pool=variable_pool)
+
+    assert segments_group.text == 'Hello, fake-user-id! Your query is fake-user-query. And your key is fake-secret-key.'
+    assert (
+        segments_group.log
+        == f"Hello, fake-user-id! Your query is fake-user-query. And your key is {encrypter.obfuscated_token('fake-secret-key')}."
+    )
+
+
+def test_convert_constant_to_segment_group():
+    variable_pool = VariablePool(
+        system_variables={},
+        user_inputs={},
+        environment_variables=[],
+    )
+    template = 'Hello, world!'
+    segments_group = parser.convert_template(template=template, variable_pool=variable_pool)
+    assert segments_group.text == 'Hello, world!'
+    assert segments_group.log == 'Hello, world!'
+
+
+def test_convert_variable_to_segment_group():
+    variable_pool = VariablePool(
+        system_variables={
+            SystemVariable('user_id'): 'fake-user-id',
+        },
+        user_inputs={},
+        environment_variables=[],
+    )
+    template = '{{#sys.user_id#}}'
+    segments_group = parser.convert_template(template=template, variable_pool=variable_pool)
+    assert segments_group.text == 'fake-user-id'
+    assert segments_group.log == 'fake-user-id'

api/tests/unit_tests/app/test_variables.py

@@ -0,0 +1,91 @@
+import pytest
+from pydantic import ValidationError
+
+from core.app.segments import (
+    FloatVariable,
+    IntegerVariable,
+    SecretVariable,
+    SegmentType,
+    StringVariable,
+    factory,
+)
+
+
+def test_string_variable():
+    test_data = {'value_type': 'string', 'name': 'test_text', 'value': 'Hello, World!'}
+    result = factory.build_variable_from_mapping(test_data)
+    assert isinstance(result, StringVariable)
+
+
+def test_integer_variable():
+    test_data = {'value_type': 'number', 'name': 'test_int', 'value': 42}
+    result = factory.build_variable_from_mapping(test_data)
+    assert isinstance(result, IntegerVariable)
+
+
+def test_float_variable():
+    test_data = {'value_type': 'number', 'name': 'test_float', 'value': 3.14}
+    result = factory.build_variable_from_mapping(test_data)
+    assert isinstance(result, FloatVariable)
+
+
+def test_secret_variable():
+    test_data = {'value_type': 'secret', 'name': 'test_secret', 'value': 'secret_value'}
+    result = factory.build_variable_from_mapping(test_data)
+    assert isinstance(result, SecretVariable)
+
+
+def test_invalid_value_type():
+    test_data = {'value_type': 'unknown', 'name': 'test_invalid', 'value': 'value'}
+    with pytest.raises(ValueError):
+        factory.build_variable_from_mapping(test_data)
+
+
+def test_frozen_variables():
+    var = StringVariable(name='text', value='text')
+    with pytest.raises(ValidationError):
+        var.value = 'new value'
+
+    int_var = IntegerVariable(name='integer', value=42)
+    with pytest.raises(ValidationError):
+        int_var.value = 100
+
+    float_var = FloatVariable(name='float', value=3.14)
+    with pytest.raises(ValidationError):
+        float_var.value = 2.718
+
+    secret_var = SecretVariable(name='secret', value='secret_value')
+    with pytest.raises(ValidationError):
+        secret_var.value = 'new_secret_value'
+
+
+def test_variable_value_type_immutable():
+    with pytest.raises(ValidationError):
+        StringVariable(value_type=SegmentType.ARRAY, name='text', value='text')
+
+    with pytest.raises(ValidationError):
+        StringVariable.model_validate({'value_type': 'not text', 'name': 'text', 'value': 'text'})
+
+    var = IntegerVariable(name='integer', value=42)
+    with pytest.raises(ValidationError):
+        IntegerVariable(value_type=SegmentType.ARRAY, name=var.name, value=var.value)
+
+    var = FloatVariable(name='float', value=3.14)
+    with pytest.raises(ValidationError):
+        FloatVariable(value_type=SegmentType.ARRAY, name=var.name, value=var.value)
+
+    var = SecretVariable(name='secret', value='secret_value')
+    with pytest.raises(ValidationError):
+        SecretVariable(value_type=SegmentType.ARRAY, name=var.name, value=var.value)
+
+
+def test_build_a_blank_string():
+    result = factory.build_variable_from_mapping(
+        {
+            'value_type': 'string',
+            'name': 'blank',
+            'value': '',
+        }
+    )
+    assert isinstance(result, StringVariable)
+    assert result.value == ''

api/tests/unit_tests/configs/test_dify_config.py

@@ -1,3 +1,4 @@
+import os
 from textwrap import dedent
 
 import pytest
@@ -48,7 +49,9 @@ def test_dify_config(example_env_file):
 # This is due to `pymilvus` loading all the variables from the `.env` file into `os.environ`.
 def test_flask_configs(example_env_file):
     flask_app = Flask('app')
-    flask_app.config.from_mapping(DifyConfig(_env_file=example_env_file).model_dump())
+    # clear system environment variables
+    os.environ.clear()
+    flask_app.config.from_mapping(DifyConfig(_env_file=example_env_file).model_dump())  # pyright: ignore
     config = flask_app.config
 
     # configs read from pydantic-settings

api/tests/unit_tests/core/workflow/nodes/test_answer.py

@@ -31,9 +31,9 @@ def test_execute_answer():
     pool = VariablePool(system_variables={
         SystemVariable.FILES: [],
         SystemVariable.USER_ID: 'aaa'
-    }, user_inputs={})
-    pool.append_variable(node_id='start', variable_key_list=['weather'], value='sunny')
-    pool.append_variable(node_id='llm', variable_key_list=['text'], value='You are a helpful AI.')
+    }, user_inputs={}, environment_variables=[])
+    pool.add(['start', 'weather'], 'sunny')
+    pool.add(['llm', 'text'], 'You are a helpful AI.')
 
     # Mock db.session.close()
     db.session.close = MagicMock()

api/tests/unit_tests/core/workflow/nodes/test_if_else.py

@@ -121,24 +121,24 @@ def test_execute_if_else_result_true():
     pool = VariablePool(system_variables={
         SystemVariable.FILES: [],
         SystemVariable.USER_ID: 'aaa'
-    }, user_inputs={})
-    pool.append_variable(node_id='start', variable_key_list=['array_contains'], value=['ab', 'def'])
-    pool.append_variable(node_id='start', variable_key_list=['array_not_contains'], value=['ac', 'def'])
-    pool.append_variable(node_id='start', variable_key_list=['contains'], value='cabcde')
-    pool.append_variable(node_id='start', variable_key_list=['not_contains'], value='zacde')
-    pool.append_variable(node_id='start', variable_key_list=['start_with'], value='abc')
-    pool.append_variable(node_id='start', variable_key_list=['end_with'], value='zzab')
-    pool.append_variable(node_id='start', variable_key_list=['is'], value='ab')
-    pool.append_variable(node_id='start', variable_key_list=['is_not'], value='aab')
-    pool.append_variable(node_id='start', variable_key_list=['empty'], value='')
-    pool.append_variable(node_id='start', variable_key_list=['not_empty'], value='aaa')
-    pool.append_variable(node_id='start', variable_key_list=['equals'], value=22)
-    pool.append_variable(node_id='start', variable_key_list=['not_equals'], value=23)
-    pool.append_variable(node_id='start', variable_key_list=['greater_than'], value=23)
-    pool.append_variable(node_id='start', variable_key_list=['less_than'], value=21)
-    pool.append_variable(node_id='start', variable_key_list=['greater_than_or_equal'], value=22)
-    pool.append_variable(node_id='start', variable_key_list=['less_than_or_equal'], value=21)
-    pool.append_variable(node_id='start', variable_key_list=['not_null'], value='1212')
+    }, user_inputs={}, environment_variables=[])
+    pool.add(['start', 'array_contains'], ['ab', 'def'])
+    pool.add(['start', 'array_not_contains'], ['ac', 'def'])
+    pool.add(['start', 'contains'], 'cabcde')
+    pool.add(['start', 'not_contains'], 'zacde')
+    pool.add(['start', 'start_with'], 'abc')
+    pool.add(['start', 'end_with'], 'zzab')
+    pool.add(['start', 'is'], 'ab')
+    pool.add(['start', 'is_not'], 'aab')
+    pool.add(['start', 'empty'], '')
+    pool.add(['start', 'not_empty'], 'aaa')
+    pool.add(['start', 'equals'], 22)
+    pool.add(['start', 'not_equals'], 23)
+    pool.add(['start', 'greater_than'], 23)
+    pool.add(['start', 'less_than'], 21)
+    pool.add(['start', 'greater_than_or_equal'], 22)
+    pool.add(['start', 'less_than_or_equal'], 21)
+    pool.add(['start', 'not_null'], '1212')
 
     # Mock db.session.close()
     db.session.close = MagicMock()
@@ -184,9 +184,9 @@ def test_execute_if_else_result_false():
     pool = VariablePool(system_variables={
         SystemVariable.FILES: [],
         SystemVariable.USER_ID: 'aaa'
-    }, user_inputs={})
-    pool.append_variable(node_id='start', variable_key_list=['array_contains'], value=['1ab', 'def'])
-    pool.append_variable(node_id='start', variable_key_list=['array_not_contains'], value=['ab', 'def'])
+    }, user_inputs={}, environment_variables=[])
+    pool.add(['start', 'array_contains'], ['1ab', 'def'])
+    pool.add(['start', 'array_not_contains'], ['ab', 'def'])
 
     # Mock db.session.close()
     db.session.close = MagicMock()

api/tests/unit_tests/models/test_workflow.py

@@ -0,0 +1,95 @@
+from unittest import mock
+from uuid import uuid4
+
+import contexts
+from constants import HIDDEN_VALUE
+from core.app.segments import FloatVariable, IntegerVariable, SecretVariable, StringVariable
+from models.workflow import Workflow
+
+
+def test_environment_variables():
+    contexts.tenant_id.set('tenant_id')
+
+    # Create a Workflow instance
+    workflow = Workflow()
+
+    # Create some EnvironmentVariable instances
+    variable1 = StringVariable.model_validate({'name': 'var1', 'value': 'value1', 'id': str(uuid4())})
+    variable2 = IntegerVariable.model_validate({'name': 'var2', 'value': 123, 'id': str(uuid4())})
+    variable3 = SecretVariable.model_validate({'name': 'var3', 'value': 'secret', 'id': str(uuid4())})
+    variable4 = FloatVariable.model_validate({'name': 'var4', 'value': 3.14, 'id': str(uuid4())})
+
+    with (
+        mock.patch('core.helper.encrypter.encrypt_token', return_value='encrypted_token'),
+        mock.patch('core.helper.encrypter.decrypt_token', return_value='secret'),
+    ):
+        # Set the environment_variables property of the Workflow instance
+        variables = [variable1, variable2, variable3, variable4]
+        workflow.environment_variables = variables
+
+        # Get the environment_variables property and assert its value
+        assert workflow.environment_variables == variables
+
+
+def test_update_environment_variables():
+    contexts.tenant_id.set('tenant_id')
+
+    # Create a Workflow instance
+    workflow = Workflow()
+
+    # Create some EnvironmentVariable instances
+    variable1 = StringVariable.model_validate({'name': 'var1', 'value': 'value1', 'id': str(uuid4())})
+    variable2 = IntegerVariable.model_validate({'name': 'var2', 'value': 123, 'id': str(uuid4())})
+    variable3 = SecretVariable.model_validate({'name': 'var3', 'value': 'secret', 'id': str(uuid4())})
+    variable4 = FloatVariable.model_validate({'name': 'var4', 'value': 3.14, 'id': str(uuid4())})
+
+    with (
+        mock.patch('core.helper.encrypter.encrypt_token', return_value='encrypted_token'),
+        mock.patch('core.helper.encrypter.decrypt_token', return_value='secret'),
+    ):
+        variables = [variable1, variable2, variable3, variable4]
+
+        # Set the environment_variables property of the Workflow instance
+        workflow.environment_variables = variables
+        assert workflow.environment_variables == [variable1, variable2, variable3, variable4]
+
+        # Update the name of variable3 and keep the value as it is
+        variables[2] = variable3.model_copy(
+            update={
+                'name': 'new name',
+                'value': HIDDEN_VALUE,
+            }
+        )
+
+        workflow.environment_variables = variables
+        assert workflow.environment_variables[2].name == 'new name'
+        assert workflow.environment_variables[2].value == variable3.value
+
+
+def test_to_dict():
+    contexts.tenant_id.set('tenant_id')
+
+    # Create a Workflow instance
+    workflow = Workflow()
+    workflow.graph = '{}'
+    workflow.features = '{}'
+
+    # Create some EnvironmentVariable instances
+
+    with (
+        mock.patch('core.helper.encrypter.encrypt_token', return_value='encrypted_token'),
+        mock.patch('core.helper.encrypter.decrypt_token', return_value='secret'),
+    ):
+        # Set the environment_variables property of the Workflow instance
+        workflow.environment_variables = [
+            SecretVariable.model_validate({'name': 'secret', 'value': 'secret', 'id': str(uuid4())}),
+            StringVariable.model_validate({'name': 'text', 'value': 'text', 'id': str(uuid4())}),
+        ]
+
+        workflow_dict = workflow.to_dict()
+        assert workflow_dict['environment_variables'][0]['value'] == ''
+        assert workflow_dict['environment_variables'][1]['value'] == 'text'
+
+        workflow_dict = workflow.to_dict(include_secret=True)
+        assert workflow_dict['environment_variables'][0]['value'] == 'secret'
+        assert workflow_dict['environment_variables'][1]['value'] == 'text'