diff --git a/api/core/workflow/workflow_entry.py b/api/core/workflow/workflow_entry.py index b69a9971b5..7aab4037c6 100644 --- a/api/core/workflow/workflow_entry.py +++ b/api/core/workflow/workflow_entry.py @@ -367,7 +367,7 @@ class WorkflowEntry: raise ValueError(f"Variable key {node_variable} not found in user inputs.") # environment variable already exist in variable pool, not from user inputs - if variable_pool.get(variable_selector): + if variable_pool.get(variable_selector) and variable_selector[0] == ENVIRONMENT_VARIABLE_NODE_ID: continue # fetch variable node id from variable selector diff --git a/api/tests/unit_tests/core/workflow/test_workflow_entry.py b/api/tests/unit_tests/core/workflow/test_workflow_entry.py new file mode 100644 index 0000000000..324f58abf6 --- /dev/null +++ b/api/tests/unit_tests/core/workflow/test_workflow_entry.py @@ -0,0 +1,456 @@ +import pytest + +from core.file.enums import FileType +from core.file.models import File, FileTransferMethod +from core.variables.variables import StringVariable +from core.workflow.constants import ( + CONVERSATION_VARIABLE_NODE_ID, + ENVIRONMENT_VARIABLE_NODE_ID, +) +from core.workflow.entities.variable_pool import VariablePool +from core.workflow.system_variable import SystemVariable +from core.workflow.workflow_entry import WorkflowEntry + + +class TestWorkflowEntry: + """Test WorkflowEntry class methods.""" + + def test_mapping_user_inputs_to_variable_pool_with_system_variables(self): + """Test mapping system variables from user inputs to variable pool.""" + # Initialize variable pool with system variables + variable_pool = VariablePool( + system_variables=SystemVariable( + user_id="test_user_id", + app_id="test_app_id", + workflow_id="test_workflow_id", + ), + user_inputs={}, + ) + + # Define variable mapping - sys variables mapped to other nodes + variable_mapping = { + "node1.input1": ["node1", "input1"], # Regular mapping + "node2.query": ["node2", "query"], # Regular mapping + "sys.user_id": ["output_node", "user"], # System variable mapping + } + + # User inputs including sys variables + user_inputs = { + "node1.input1": "new_user_id", + "node2.query": "test query", + "sys.user_id": "system_user", + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify variables were added to pool + # Note: variable_pool.get returns Variable objects, not raw values + node1_var = variable_pool.get(["node1", "input1"]) + assert node1_var is not None + assert node1_var.value == "new_user_id" + + node2_var = variable_pool.get(["node2", "query"]) + assert node2_var is not None + assert node2_var.value == "test query" + + # System variable gets mapped to output node + output_var = variable_pool.get(["output_node", "user"]) + assert output_var is not None + assert output_var.value == "system_user" + + def test_mapping_user_inputs_to_variable_pool_with_env_variables(self): + """Test mapping environment variables from user inputs to variable pool.""" + # Initialize variable pool with environment variables + env_var = StringVariable(name="API_KEY", value="existing_key") + variable_pool = VariablePool( + system_variables=SystemVariable.empty(), + environment_variables=[env_var], + user_inputs={}, + ) + + # Add env variable to pool (simulating initialization) + variable_pool.add([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], env_var) + + # Define variable mapping - env variables should not be overridden + variable_mapping = { + "node1.api_key": [ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], + "node2.new_env": [ENVIRONMENT_VARIABLE_NODE_ID, "NEW_ENV"], + } + + # User inputs + user_inputs = { + "node1.api_key": "user_provided_key", # This should not override existing env var + "node2.new_env": "new_env_value", + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify env variable was not overridden + env_value = variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"]) + assert env_value is not None + assert env_value.value == "existing_key" # Should remain unchanged + + # New env variables from user input should not be added + assert variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "NEW_ENV"]) is None + + def test_mapping_user_inputs_to_variable_pool_with_conversation_variables(self): + """Test mapping conversation variables from user inputs to variable pool.""" + # Initialize variable pool with conversation variables + conv_var = StringVariable(name="last_message", value="Hello") + variable_pool = VariablePool( + system_variables=SystemVariable.empty(), + conversation_variables=[conv_var], + user_inputs={}, + ) + + # Add conversation variable to pool + variable_pool.add([CONVERSATION_VARIABLE_NODE_ID, "last_message"], conv_var) + + # Define variable mapping + variable_mapping = { + "node1.message": ["node1", "message"], # Map to regular node + "conversation.context": ["chat_node", "context"], # Conversation var to regular node + } + + # User inputs + user_inputs = { + "node1.message": "Updated message", + "conversation.context": "New context", + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify variables were added to their target nodes + node1_var = variable_pool.get(["node1", "message"]) + assert node1_var is not None + assert node1_var.value == "Updated message" + + chat_var = variable_pool.get(["chat_node", "context"]) + assert chat_var is not None + assert chat_var.value == "New context" + + def test_mapping_user_inputs_to_variable_pool_with_regular_variables(self): + """Test mapping regular node variables from user inputs to variable pool.""" + # Initialize empty variable pool + variable_pool = VariablePool( + system_variables=SystemVariable.empty(), + user_inputs={}, + ) + + # Define variable mapping for regular nodes + variable_mapping = { + "input_node.text": ["input_node", "text"], + "llm_node.prompt": ["llm_node", "prompt"], + "code_node.input": ["code_node", "input"], + } + + # User inputs + user_inputs = { + "input_node.text": "User input text", + "llm_node.prompt": "Generate a summary", + "code_node.input": {"key": "value"}, + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify regular variables were added + text_var = variable_pool.get(["input_node", "text"]) + assert text_var is not None + assert text_var.value == "User input text" + + prompt_var = variable_pool.get(["llm_node", "prompt"]) + assert prompt_var is not None + assert prompt_var.value == "Generate a summary" + + input_var = variable_pool.get(["code_node", "input"]) + assert input_var is not None + assert input_var.value == {"key": "value"} + + def test_mapping_user_inputs_with_file_handling(self): + """Test mapping file inputs from user inputs to variable pool.""" + variable_pool = VariablePool( + system_variables=SystemVariable.empty(), + user_inputs={}, + ) + + # Define variable mapping + variable_mapping = { + "file_node.file": ["file_node", "file"], + "file_node.files": ["file_node", "files"], + } + + # User inputs with file data - using remote_url which doesn't require upload_file_id + user_inputs = { + "file_node.file": { + "type": "document", + "transfer_method": "remote_url", + "url": "http://example.com/test.pdf", + }, + "file_node.files": [ + { + "type": "image", + "transfer_method": "remote_url", + "url": "http://example.com/image1.jpg", + }, + { + "type": "image", + "transfer_method": "remote_url", + "url": "http://example.com/image2.jpg", + }, + ], + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify file was converted and added + file_var = variable_pool.get(["file_node", "file"]) + assert file_var is not None + assert file_var.value.type == FileType.DOCUMENT + assert file_var.value.transfer_method == FileTransferMethod.REMOTE_URL + + # Verify file list was converted and added + files_var = variable_pool.get(["file_node", "files"]) + assert files_var is not None + assert isinstance(files_var.value, list) + assert len(files_var.value) == 2 + assert all(isinstance(f, File) for f in files_var.value) + assert files_var.value[0].type == FileType.IMAGE + assert files_var.value[1].type == FileType.IMAGE + assert files_var.value[0].type == FileType.IMAGE + assert files_var.value[1].type == FileType.IMAGE + + def test_mapping_user_inputs_missing_variable_error(self): + """Test that mapping raises error when required variable is missing.""" + variable_pool = VariablePool( + system_variables=SystemVariable.empty(), + user_inputs={}, + ) + + # Define variable mapping + variable_mapping = { + "node1.required_input": ["node1", "required_input"], + } + + # User inputs without required variable + user_inputs = { + "node1.other_input": "some value", + } + + # Should raise ValueError for missing variable + with pytest.raises(ValueError, match="Variable key node1.required_input not found in user inputs"): + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + def test_mapping_user_inputs_with_alternative_key_format(self): + """Test mapping with alternative key format (without node prefix).""" + variable_pool = VariablePool( + system_variables=SystemVariable.empty(), + user_inputs={}, + ) + + # Define variable mapping + variable_mapping = { + "node1.input": ["node1", "input"], + } + + # User inputs with alternative key format + user_inputs = { + "input": "value without node prefix", # Alternative format without node prefix + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify variable was added using alternative key + input_var = variable_pool.get(["node1", "input"]) + assert input_var is not None + assert input_var.value == "value without node prefix" + + def test_mapping_user_inputs_with_complex_selectors(self): + """Test mapping with complex node variable keys.""" + variable_pool = VariablePool( + system_variables=SystemVariable.empty(), + user_inputs={}, + ) + + # Define variable mapping - selectors can only have 2 elements + variable_mapping = { + "node1.data.field1": ["node1", "data_field1"], # Complex key mapped to simple selector + "node2.config.settings.timeout": ["node2", "timeout"], # Complex key mapped to simple selector + } + + # User inputs + user_inputs = { + "node1.data.field1": "nested value", + "node2.config.settings.timeout": 30, + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify variables were added with simple selectors + data_var = variable_pool.get(["node1", "data_field1"]) + assert data_var is not None + assert data_var.value == "nested value" + + timeout_var = variable_pool.get(["node2", "timeout"]) + assert timeout_var is not None + assert timeout_var.value == 30 + + def test_mapping_user_inputs_invalid_node_variable(self): + """Test that mapping handles invalid node variable format.""" + variable_pool = VariablePool( + system_variables=SystemVariable.empty(), + user_inputs={}, + ) + + # Define variable mapping with single element node variable (at least one dot is required) + variable_mapping = { + "singleelement": ["node1", "input"], # No dot separator + } + + user_inputs = {"singleelement": "some value"} # Must use exact key + + # Should NOT raise error - function accepts it and uses direct key + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify it was added + var = variable_pool.get(["node1", "input"]) + assert var is not None + assert var.value == "some value" + + def test_mapping_all_variable_types_together(self): + """Test mapping all four types of variables in one operation.""" + # Initialize variable pool with some existing variables + env_var = StringVariable(name="API_KEY", value="existing_key") + conv_var = StringVariable(name="session_id", value="session123") + + variable_pool = VariablePool( + system_variables=SystemVariable( + user_id="test_user", + app_id="test_app", + query="initial query", + ), + environment_variables=[env_var], + conversation_variables=[conv_var], + user_inputs={}, + ) + + # Add existing variables to pool + variable_pool.add([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"], env_var) + variable_pool.add([CONVERSATION_VARIABLE_NODE_ID, "session_id"], conv_var) + + # Define comprehensive variable mapping + variable_mapping = { + # System variables mapped to regular nodes + "sys.user_id": ["start", "user"], + "sys.app_id": ["start", "app"], + # Environment variables (won't be overridden) + "env.API_KEY": ["config", "api_key"], + # Conversation variables mapped to regular nodes + "conversation.session_id": ["chat", "session"], + # Regular variables + "input.text": ["input", "text"], + "process.data": ["process", "data"], + } + + # User inputs + user_inputs = { + "sys.user_id": "new_user", + "sys.app_id": "new_app", + "env.API_KEY": "attempted_override", # Should not override env var + "conversation.session_id": "new_session", + "input.text": "user input text", + "process.data": {"value": 123, "status": "active"}, + } + + # Execute mapping + WorkflowEntry.mapping_user_inputs_to_variable_pool( + variable_mapping=variable_mapping, + user_inputs=user_inputs, + variable_pool=variable_pool, + tenant_id="test_tenant", + ) + + # Verify system variables were added to their target nodes + start_user = variable_pool.get(["start", "user"]) + assert start_user is not None + assert start_user.value == "new_user" + + start_app = variable_pool.get(["start", "app"]) + assert start_app is not None + assert start_app.value == "new_app" + + # Verify env variable was not overridden (still has original value) + env_value = variable_pool.get([ENVIRONMENT_VARIABLE_NODE_ID, "API_KEY"]) + assert env_value is not None + assert env_value.value == "existing_key" + + # Environment variables get mapped to other nodes even when they exist in env pool + # But the original env value remains unchanged + config_api_key = variable_pool.get(["config", "api_key"]) + assert config_api_key is not None + assert config_api_key.value == "attempted_override" + + # Verify conversation variable was mapped to target node + chat_session = variable_pool.get(["chat", "session"]) + assert chat_session is not None + assert chat_session.value == "new_session" + + # Verify regular variables were added + input_text = variable_pool.get(["input", "text"]) + assert input_text is not None + assert input_text.value == "user input text" + + process_data = variable_pool.get(["process", "data"]) + assert process_data is not None + assert process_data.value == {"value": 123, "status": "active"}