vellum-workflow-server 0.14.70__tar.gz → 0.14.70.post2__tar.gz

{vellum_workflow_server-0.14.70 → vellum_workflow_server-0.14.70.post2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-workflow-server
-Version: 0.14.70
+Version: 0.14.70.post2
 Summary: 
 License: AGPL
 Requires-Python: >=3.9.0,<4

{vellum_workflow_server-0.14.70 → vellum_workflow_server-0.14.70.post2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vellum-workflow-server"
-version = "0.14.70"
+version = "0.14.70.post2"
 description = ""
 readme = "README.md"
 authors = []

{vellum_workflow_server-0.14.70 → vellum_workflow_server-0.14.70.post2}/src/workflow_server/api/tests/test_workflow_view_stream_workflow_route.py

@@ -743,3 +743,44 @@ class EndNodeDisplay(BaseNodeDisplay[EndNode]):
     assert len(events) > 3, json.dumps(events[-1])
     assert events[2]["name"] == "node.execution.initiated", json.dumps(events[-1])
     assert events[2]["body"]["inputs"] == {"fruit": "cherry"}
+
+
+def test_stream_workflow_route__with_environment_variables(both_stream_types):
+    # GIVEN a valid request body with environment variables
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "environment_variables": {"TEST_ENV_VAR": "test_value", "ANOTHER_VAR": "another_value"},
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.references import EnvironmentVariableReference
+
+class Workflow(BaseWorkflow):
+    class Outputs(BaseWorkflow.Outputs):
+        env_var_value = EnvironmentVariableReference(name="TEST_ENV_VAR", default="not_found")
+        another_var_value = EnvironmentVariableReference(name="ANOTHER_VAR", default="not_found")
+""",
+        },
+    }
+
+    # WHEN we call the stream route
+    status_code, events = both_stream_types(request_body)
+
+    # THEN we get a 200 response
+    assert status_code == 200, events
+
+    # THEN we get the expected events
+    assert events[0]["name"] == "vembda.execution.initiated"
+    assert events[1]["name"] == "workflow.execution.initiated"
+    assert events[2]["name"] == "workflow.execution.fulfilled"
+
+    # AND the environment variables are accessible in the workflow
+    outputs = events[2]["body"]["outputs"]
+    assert outputs["env_var_value"] == "test_value"
+    assert outputs["another_var_value"] == "another_value"

{vellum_workflow_server-0.14.70 → vellum_workflow_server-0.14.70.post2}/src/workflow_server/core/executor.py

@@ -342,6 +342,9 @@ def _create_workflow_context(executor_context: BaseExecutorContext) -> WorkflowC
     else:
         environment = VellumEnvironment.PRODUCTION
 
+    if executor_context.environment_variables:
+        os.environ.update(executor_context.environment_variables)
+
     return WorkflowContext(
         vellum_client=Vellum(
             api_key=executor_context.environment_api_key,

{vellum_workflow_server-0.14.70 → vellum_workflow_server-0.14.70.post2}/src/workflow_server/core/workflow_executor_context.py

@@ -22,6 +22,7 @@ class BaseExecutorContext(UniversalBaseModel):
     stream_start_time: int = 0
     vembda_public_url: Optional[str] = None
     node_output_mocks: Optional[list[Any]] = None
+    environment_variables: Optional[dict[str, str]] = None
 
     @property
     def container_overhead_latency(self) -> int: