Feat: conversation variable & variable assigner node (#7222)

Signed-off-by: -LAN- <laipz8200@outlook.com>
Co-authored-by: Joel <iamjoel007@gmail.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>

api/tests/unit_tests/core/app/segments/test_factory.py

@@ -7,15 +7,16 @@ from core.app.segments import (
     ArrayNumberVariable,
     ArrayObjectVariable,
     ArrayStringVariable,
+    FileSegment,
     FileVariable,
     FloatVariable,
     IntegerVariable,
-    NoneSegment,
     ObjectSegment,
     SecretVariable,
     StringVariable,
     factory,
 )
+from core.app.segments.exc import VariableError
 
 
 def test_string_variable():
@@ -44,7 +45,7 @@ def test_secret_variable():
 
 def test_invalid_value_type():
     test_data = {'value_type': 'unknown', 'name': 'test_invalid', 'value': 'value'}
-    with pytest.raises(ValueError):
+    with pytest.raises(VariableError):
         factory.build_variable_from_mapping(test_data)
 
 
@@ -77,26 +78,14 @@ def test_object_variable():
         'name': 'test_object',
         'description': 'Description of the variable.',
         'value': {
-            'key1': {
-                'id': str(uuid4()),
-                'value_type': 'string',
-                'name': 'text',
-                'value': 'text',
-                'description': 'Description of the variable.',
-            },
-            'key2': {
-                'id': str(uuid4()),
-                'value_type': 'number',
-                'name': 'number',
-                'value': 1,
-                'description': 'Description of the variable.',
-            },
+            'key1': 'text',
+            'key2': 2,
         },
     }
     variable = factory.build_variable_from_mapping(mapping)
     assert isinstance(variable, ObjectSegment)
-    assert isinstance(variable.value['key1'], StringVariable)
-    assert isinstance(variable.value['key2'], IntegerVariable)
+    assert isinstance(variable.value['key1'], str)
+    assert isinstance(variable.value['key2'], int)
 
 
 def test_array_string_variable():
@@ -106,26 +95,14 @@ def test_array_string_variable():
         'name': 'test_array',
         'description': 'Description of the variable.',
         'value': [
-            {
-                'id': str(uuid4()),
-                'value_type': 'string',
-                'name': 'text',
-                'value': 'text',
-                'description': 'Description of the variable.',
-            },
-            {
-                'id': str(uuid4()),
-                'value_type': 'string',
-                'name': 'text',
-                'value': 'text',
-                'description': 'Description of the variable.',
-            },
+            'text',
+            'text',
         ],
     }
     variable = factory.build_variable_from_mapping(mapping)
     assert isinstance(variable, ArrayStringVariable)
-    assert isinstance(variable.value[0], StringVariable)
-    assert isinstance(variable.value[1], StringVariable)
+    assert isinstance(variable.value[0], str)
+    assert isinstance(variable.value[1], str)
 
 
 def test_array_number_variable():
@@ -135,26 +112,14 @@ def test_array_number_variable():
         'name': 'test_array',
         'description': 'Description of the variable.',
         'value': [
-            {
-                'id': str(uuid4()),
-                'value_type': 'number',
-                'name': 'number',
-                'value': 1,
-                'description': 'Description of the variable.',
-            },
-            {
-                'id': str(uuid4()),
-                'value_type': 'number',
-                'name': 'number',
-                'value': 2.0,
-                'description': 'Description of the variable.',
-            },
+            1,
+            2.0,
         ],
     }
     variable = factory.build_variable_from_mapping(mapping)
     assert isinstance(variable, ArrayNumberVariable)
-    assert isinstance(variable.value[0], IntegerVariable)
-    assert isinstance(variable.value[1], FloatVariable)
+    assert isinstance(variable.value[0], int)
+    assert isinstance(variable.value[1], float)
 
 
 def test_array_object_variable():
@@ -165,59 +130,23 @@ def test_array_object_variable():
         'description': 'Description of the variable.',
         'value': [
             {
-                'id': str(uuid4()),
-                'value_type': 'object',
-                'name': 'object',
-                'description': 'Description of the variable.',
-                'value': {
-                    'key1': {
-                        'id': str(uuid4()),
-                        'value_type': 'string',
-                        'name': 'text',
-                        'value': 'text',
-                        'description': 'Description of the variable.',
-                    },
-                    'key2': {
-                        'id': str(uuid4()),
-                        'value_type': 'number',
-                        'name': 'number',
-                        'value': 1,
-                        'description': 'Description of the variable.',
-                    },
-                },
+                'key1': 'text',
+                'key2': 1,
             },
             {
-                'id': str(uuid4()),
-                'value_type': 'object',
-                'name': 'object',
-                'description': 'Description of the variable.',
-                'value': {
-                    'key1': {
-                        'id': str(uuid4()),
-                        'value_type': 'string',
-                        'name': 'text',
-                        'value': 'text',
-                        'description': 'Description of the variable.',
-                    },
-                    'key2': {
-                        'id': str(uuid4()),
-                        'value_type': 'number',
-                        'name': 'number',
-                        'value': 1,
-                        'description': 'Description of the variable.',
-                    },
-                },
+                'key1': 'text',
+                'key2': 1,
             },
         ],
     }
     variable = factory.build_variable_from_mapping(mapping)
     assert isinstance(variable, ArrayObjectVariable)
-    assert isinstance(variable.value[0], ObjectSegment)
-    assert isinstance(variable.value[1], ObjectSegment)
-    assert isinstance(variable.value[0].value['key1'], StringVariable)
-    assert isinstance(variable.value[0].value['key2'], IntegerVariable)
-    assert isinstance(variable.value[1].value['key1'], StringVariable)
-    assert isinstance(variable.value[1].value['key2'], IntegerVariable)
+    assert isinstance(variable.value[0], dict)
+    assert isinstance(variable.value[1], dict)
+    assert isinstance(variable.value[0]['key1'], str)
+    assert isinstance(variable.value[0]['key2'], int)
+    assert isinstance(variable.value[1]['key1'], str)
+    assert isinstance(variable.value[1]['key2'], int)
 
 
 def test_file_variable():
@@ -257,51 +186,53 @@ def test_array_file_variable():
         'value': [
             {
                 'id': str(uuid4()),
-                'name': 'file',
-                'value_type': 'file',
-                'value': {
-                    'id': str(uuid4()),
-                    'tenant_id': 'tenant_id',
-                    'type': 'image',
-                    'transfer_method': 'local_file',
-                    'url': 'url',
-                    'related_id': 'related_id',
-                    'extra_config': {
-                        'image_config': {
-                            'width': 100,
-                            'height': 100,
-                        },
+                'tenant_id': 'tenant_id',
+                'type': 'image',
+                'transfer_method': 'local_file',
+                'url': 'url',
+                'related_id': 'related_id',
+                'extra_config': {
+                    'image_config': {
+                        'width': 100,
+                        'height': 100,
                     },
-                    'filename': 'filename',
-                    'extension': 'extension',
-                    'mime_type': 'mime_type',
                 },
+                'filename': 'filename',
+                'extension': 'extension',
+                'mime_type': 'mime_type',
             },
             {
                 'id': str(uuid4()),
-                'name': 'file',
-                'value_type': 'file',
-                'value': {
-                    'id': str(uuid4()),
-                    'tenant_id': 'tenant_id',
-                    'type': 'image',
-                    'transfer_method': 'local_file',
-                    'url': 'url',
-                    'related_id': 'related_id',
-                    'extra_config': {
-                        'image_config': {
-                            'width': 100,
-                            'height': 100,
-                        },
+                'tenant_id': 'tenant_id',
+                'type': 'image',
+                'transfer_method': 'local_file',
+                'url': 'url',
+                'related_id': 'related_id',
+                'extra_config': {
+                    'image_config': {
+                        'width': 100,
+                        'height': 100,
                     },
-                    'filename': 'filename',
-                    'extension': 'extension',
-                    'mime_type': 'mime_type',
                 },
+                'filename': 'filename',
+                'extension': 'extension',
+                'mime_type': 'mime_type',
             },
         ],
     }
     variable = factory.build_variable_from_mapping(mapping)
     assert isinstance(variable, ArrayFileVariable)
-    assert isinstance(variable.value[0], FileVariable)
-    assert isinstance(variable.value[1], FileVariable)
+    assert isinstance(variable.value[0], FileSegment)
+    assert isinstance(variable.value[1], FileSegment)
+
+
+def test_variable_cannot_large_than_5_kb():
+    with pytest.raises(VariableError):
+        factory.build_variable_from_mapping(
+            {
+                'id': str(uuid4()),
+                'value_type': 'string',
+                'name': 'test_text',
+                'value': 'a' * 1024 * 6,
+            }
+        )

api/tests/unit_tests/core/prompt/test_advanced_prompt_transform.py

@@ -2,8 +2,8 @@ from unittest.mock import MagicMock
 
 import pytest
 
-from core.app.app_config.entities import FileExtraConfig, ModelConfigEntity
-from core.file.file_obj import FileTransferMethod, FileType, FileVar
+from core.app.app_config.entities import ModelConfigEntity
+from core.file.file_obj import FileExtraConfig, FileTransferMethod, FileType, FileVar
 from core.memory.token_buffer_memory import TokenBufferMemory
 from core.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessageRole, UserPromptMessage
 from core.prompt.advanced_prompt_transform import AdvancedPromptTransform

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

@@ -0,0 +1,150 @@
+from unittest import mock
+from uuid import uuid4
+
+from core.app.entities.app_invoke_entities import InvokeFrom
+from core.app.segments import ArrayStringVariable, StringVariable
+from core.workflow.entities.node_entities import SystemVariable
+from core.workflow.entities.variable_pool import VariablePool
+from core.workflow.nodes.base_node import UserFrom
+from core.workflow.nodes.variable_assigner import VariableAssignerNode, WriteMode
+
+DEFAULT_NODE_ID = 'node_id'
+
+
+def test_overwrite_string_variable():
+    conversation_variable = StringVariable(
+        id=str(uuid4()),
+        name='test_conversation_variable',
+        value='the first value',
+    )
+
+    input_variable = StringVariable(
+        id=str(uuid4()),
+        name='test_string_variable',
+        value='the second value',
+    )
+
+    node = VariableAssignerNode(
+        tenant_id='tenant_id',
+        app_id='app_id',
+        workflow_id='workflow_id',
+        user_id='user_id',
+        user_from=UserFrom.ACCOUNT,
+        invoke_from=InvokeFrom.DEBUGGER,
+        config={
+            'id': 'node_id',
+            'data': {
+                'assigned_variable_selector': ['conversation', conversation_variable.name],
+                'write_mode': WriteMode.OVER_WRITE.value,
+                'input_variable_selector': [DEFAULT_NODE_ID, input_variable.name],
+            },
+        },
+    )
+
+    variable_pool = VariablePool(
+        system_variables={SystemVariable.CONVERSATION_ID: 'conversation_id'},
+        user_inputs={},
+        environment_variables=[],
+        conversation_variables=[conversation_variable],
+    )
+    variable_pool.add(
+        [DEFAULT_NODE_ID, input_variable.name],
+        input_variable,
+    )
+
+    with mock.patch('core.workflow.nodes.variable_assigner.update_conversation_variable') as mock_run:
+        node.run(variable_pool)
+        mock_run.assert_called_once()
+
+    got = variable_pool.get(['conversation', conversation_variable.name])
+    assert got is not None
+    assert got.value == 'the second value'
+    assert got.to_object() == 'the second value'
+
+
+def test_append_variable_to_array():
+    conversation_variable = ArrayStringVariable(
+        id=str(uuid4()),
+        name='test_conversation_variable',
+        value=['the first value'],
+    )
+
+    input_variable = StringVariable(
+        id=str(uuid4()),
+        name='test_string_variable',
+        value='the second value',
+    )
+
+    node = VariableAssignerNode(
+        tenant_id='tenant_id',
+        app_id='app_id',
+        workflow_id='workflow_id',
+        user_id='user_id',
+        user_from=UserFrom.ACCOUNT,
+        invoke_from=InvokeFrom.DEBUGGER,
+        config={
+            'id': 'node_id',
+            'data': {
+                'assigned_variable_selector': ['conversation', conversation_variable.name],
+                'write_mode': WriteMode.APPEND.value,
+                'input_variable_selector': [DEFAULT_NODE_ID, input_variable.name],
+            },
+        },
+    )
+
+    variable_pool = VariablePool(
+        system_variables={SystemVariable.CONVERSATION_ID: 'conversation_id'},
+        user_inputs={},
+        environment_variables=[],
+        conversation_variables=[conversation_variable],
+    )
+    variable_pool.add(
+        [DEFAULT_NODE_ID, input_variable.name],
+        input_variable,
+    )
+
+    with mock.patch('core.workflow.nodes.variable_assigner.update_conversation_variable') as mock_run:
+        node.run(variable_pool)
+        mock_run.assert_called_once()
+
+    got = variable_pool.get(['conversation', conversation_variable.name])
+    assert got is not None
+    assert got.to_object() == ['the first value', 'the second value']
+
+
+def test_clear_array():
+    conversation_variable = ArrayStringVariable(
+        id=str(uuid4()),
+        name='test_conversation_variable',
+        value=['the first value'],
+    )
+
+    node = VariableAssignerNode(
+        tenant_id='tenant_id',
+        app_id='app_id',
+        workflow_id='workflow_id',
+        user_id='user_id',
+        user_from=UserFrom.ACCOUNT,
+        invoke_from=InvokeFrom.DEBUGGER,
+        config={
+            'id': 'node_id',
+            'data': {
+                'assigned_variable_selector': ['conversation', conversation_variable.name],
+                'write_mode': WriteMode.CLEAR.value,
+                'input_variable_selector': [],
+            },
+        },
+    )
+
+    variable_pool = VariablePool(
+        system_variables={SystemVariable.CONVERSATION_ID: 'conversation_id'},
+        user_inputs={},
+        environment_variables=[],
+        conversation_variables=[conversation_variable],
+    )
+
+    node.run(variable_pool)
+
+    got = variable_pool.get(['conversation', conversation_variable.name])
+    assert got is not None
+    assert got.to_object() == []