vellum-workflow-server 1.4.2.post1__tar.gz → 1.5.0__tar.gz

{vellum_workflow_server-1.4.2.post1 → vellum_workflow_server-1.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-workflow-server
-Version: 1.4.2.post1
+Version: 1.5.0
 Summary: 
 License: AGPL
 Requires-Python: >=3.9.0,<4
@@ -29,7 +29,7 @@ Requires-Dist: pyjwt (==2.10.0)
 Requires-Dist: python-dotenv (==1.0.1)
 Requires-Dist: retrying (==1.3.4)
 Requires-Dist: sentry-sdk[flask] (==2.20.0)
-Requires-Dist: vellum-ai (==1.4.2)
+Requires-Dist: vellum-ai (==1.5.0)
 Description-Content-Type: text/markdown
 
 # Vellum Workflow Runner Server

{vellum_workflow_server-1.4.2.post1 → vellum_workflow_server-1.5.0}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-workflow-server"
 
 [tool.poetry]
 name = "vellum-workflow-server"
-version = "1.4.2.post1"
+version = "1.5.0"
 description = ""
 readme = "README.md"
 authors = []
@@ -45,7 +45,7 @@ flask = "2.3.3"
 orderly-set = "5.2.2"
 pebble = "5.0.7"
 gunicorn = "23.0.0"
-vellum-ai = "1.4.2"
+vellum-ai = "1.5.0"
 python-dotenv = "1.0.1"
 retrying = "1.3.4"
 sentry-sdk = {extras = ["flask"], version = "2.20.0"}

{vellum_workflow_server-1.4.2.post1 → vellum_workflow_server-1.5.0}/src/workflow_server/api/tests/test_workflow_view.py

@@ -1,7 +1,8 @@
+import json
 import logging
 import re
 from unittest.mock import patch
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from deepdiff import DeepDiff
 
@@ -534,3 +535,197 @@ def test_serialize_route__with_invalid_workspace_api_key():
 
     # AND the response should contain the serialization result
     assert "exec_config" in response.json
+
+
+def test_stream_node_route__with_node_id():
+    """
+    Tests that the stream-node endpoint works with node_id.
+    """
+    # GIVEN a valid request body with node_id
+    node_id = uuid4()
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "node_id": str(node_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "from .display import *",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .nodes.test_node import TestNode
+
+class Workflow(BaseWorkflow):
+    graph = TestNode
+
+    class Outputs(BaseWorkflow.Outputs):
+        result = TestNode.Outputs.value
+""",
+            "nodes/__init__.py": "from .test_node import TestNode\n__all__ = ['TestNode']",
+            "nodes/test_node.py": """\
+from vellum.workflows.nodes import BaseNode
+
+class TestNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        value = "test_result"
+""",
+            "display/__init__.py": "from .nodes import *\nfrom .workflow import *",
+            "display/workflow.py": """\
+from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
+""",
+            "display/nodes/__init__.py": "from .test_node import TestNodeDisplay\n__all__ = ['TestNodeDisplay']",
+            "display/nodes/test_node.py": f"""\
+from uuid import UUID
+from vellum_ee.workflows.display.nodes import BaseNodeDisplay
+from ...nodes.test_node import TestNode
+
+class TestNodeDisplay(BaseNodeDisplay[TestNode]):
+    node_id = UUID("{node_id}")
+""",
+        },
+    }
+
+    flask_app = create_app()
+
+    # WHEN we call the stream-node route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream-node", json=request_body)
+
+    # THEN we get a 200 response
+    assert response.status_code == 200
+
+    events = [json.loads(line) for line in response.data.decode().strip().split("\n") if line.strip()]
+    assert len(events) == 4
+    assert events[0]["name"] == "vembda.execution.initiated"
+    assert events[1]["name"] == "node.execution.initiated"
+    assert events[2]["name"] == "node.execution.fulfilled"
+    assert events[3]["name"] == "vembda.execution.fulfilled"
+
+
+def test_stream_node_route__with_node_module_and_name_backward_compatibility():
+    """
+    Tests that the stream-node endpoint still works with node_module and node_name for backward compatibility.
+    """
+    # GIVEN a valid request body with node_module and node_name (old format)
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "node_module": "nodes.test_node",
+        "node_name": "TestNode",
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .nodes.test_node import TestNode
+
+class Workflow(BaseWorkflow):
+    graph = TestNode
+""",
+            "nodes/__init__.py": "from .test_node import TestNode",
+            "nodes/test_node.py": """\
+from vellum.workflows.nodes import BaseNode
+
+class TestNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        value = "test_result"
+""",
+        },
+    }
+
+    flask_app = create_app()
+
+    # WHEN we call the stream-node route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream-node", json=request_body)
+
+    # THEN we get a 200 response
+    assert response.status_code == 200
+
+    events = [json.loads(line) for line in response.data.decode().strip().split("\n") if line.strip()]
+    assert events[0]["name"] == "vembda.execution.initiated"
+    assert events[1]["name"] == "node.execution.initiated", json.dumps(events[1]["body"])
+    assert events[2]["name"] == "node.execution.fulfilled"
+    assert events[3]["name"] == "vembda.execution.fulfilled"
+    assert len(events) == 4
+
+
+def test_stream_node_route__missing_node_info_validation():
+    """
+    Tests that the stream-node endpoint returns validation error when neither
+    node_id nor node_module/node_name are provided.
+    """
+    # GIVEN a request body missing node identification
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {"__init__.py": "", "workflow.py": ""},
+    }
+
+    flask_app = create_app()
+
+    # WHEN we call the stream-node route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream-node", json=request_body)
+
+    # THEN we get a 400 response
+    assert response.status_code == 400
+
+    # AND we get a validation error message
+    assert "Either node_id or both node_module and node_name must be provided" in response.get_json()["detail"]
+
+
+def test_stream_node_route__invalid_node_id():
+    """
+    Tests that the stream-node endpoint returns 404 for invalid node_id.
+    """
+    # GIVEN a request body with invalid node_id
+    invalid_node_id = uuid4()
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "node_id": str(invalid_node_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+
+class Workflow(BaseWorkflow):
+    pass
+""",
+        },
+    }
+
+    flask_app = create_app()
+
+    # WHEN we call the stream-node route
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream-node", json=request_body)
+
+    # THEN we get a 200 response
+    # TODO: In the future, we would want this to return a 4xx response by returning the workflow
+    # instance and the node definition that we want to run as part of request deserialization.
+    assert response.status_code == 200, response.text
+
+    # AND we get an appropriate error message
+    # TODO: In a future where we are returning 4xx responses, we assert the following data:
+    # response_data = response.get_json()
+    # assert "Node with ID" in response_data["detail"]
+    # assert "not found" in response_data["detail"]
+    events = [json.loads(line) for line in response.data.decode().strip().split("\n") if line.strip()]
+    assert events[0]["name"] == "vembda.execution.initiated"
+    assert events[1]["name"] == "vembda.execution.fulfilled"
+    assert len(events) == 2

{vellum_workflow_server-1.4.2.post1 → vellum_workflow_server-1.5.0}/src/workflow_server/api/workflow_view.py

@@ -36,7 +36,7 @@ from workflow_server.core.events import (
     VembdaExecutionInitiatedBody,
     VembdaExecutionInitiatedEvent,
 )
-from workflow_server.core.executor import stream_node_pebble_timeout, stream_workflow, stream_workflow_process_timeout
+from workflow_server.core.executor import stream_node_process_timeout, stream_workflow, stream_workflow_process_timeout
 from workflow_server.core.utils import (
     create_vembda_rejected_event,
     is_events_emitting_enabled,
@@ -374,10 +374,10 @@ def stream_node_route() -> Response:
         parent=None,
     )
 
-    app.logger.debug(f"Node stream received {data.get('execution_id')}")
+    app.logger.debug(f"Node stream started. Trace ID: {context.trace_id}")
 
     pebble_queue: Queue[dict] = Queue()
-    stream_future = stream_node_pebble_timeout(
+    process = stream_node_process_timeout(
         executor_context=context,
         queue=pebble_queue,
     )
@@ -388,14 +388,9 @@ def stream_node_route() -> Response:
                 event = pebble_queue.get(timeout=context.timeout)
 
             except Empty:
-                if stream_future.exception() is not None:
-                    # This happens when theres a problem with the stream function call
-                    # itself not the workflow runner
+                if not process.is_alive():
                     yield create_vembda_rejected_event(context, "Internal Server Error")
-                    app.logger.exception(stream_future.exception())
                     break
-                else:
-                    continue
 
             yield event
             if event.get("name") == VEMBDA_EXECUTION_FULFILLED_EVENT_NAME:

{vellum_workflow_server-1.4.2.post1 → vellum_workflow_server-1.5.0}/src/workflow_server/core/executor.py

@@ -15,7 +15,6 @@ from traceback import format_exc
 from uuid import UUID, uuid4
 from typing import Any, Callable, Generator, Iterator, Optional, Tuple, Type
 
-from pebble import concurrent
 from vellum_ee.workflows.display.utils.events import event_enricher
 from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
 from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
@@ -24,6 +23,14 @@ from vellum.workflows import BaseWorkflow
 from vellum.workflows.emitters.base import BaseWorkflowEmitter
 from vellum.workflows.emitters.vellum_emitter import VellumEmitter
 from vellum.workflows.events.exception_handling import stream_initialization_exception
+from vellum.workflows.events.node import (
+    NodeExecutionFulfilledBody,
+    NodeExecutionFulfilledEvent,
+    NodeExecutionInitiatedBody,
+    NodeExecutionInitiatedEvent,
+    NodeExecutionStreamingBody,
+    NodeExecutionStreamingEvent,
+)
 from vellum.workflows.events.types import BaseEvent
 from vellum.workflows.events.workflow import WorkflowEventDisplayContext
 from vellum.workflows.exceptions import WorkflowInitializationException
@@ -44,9 +51,12 @@ from workflow_server.core.events import (
     VembdaExecutionFulfilledBody,
     VembdaExecutionFulfilledEvent,
 )
-from workflow_server.core.utils import is_events_emitting_enabled, serialize_vembda_rejected_event
+from workflow_server.core.utils import (
+    create_vembda_rejected_event,
+    is_events_emitting_enabled,
+    serialize_vembda_rejected_event,
+)
 from workflow_server.core.workflow_executor_context import (
-    DEFAULT_TIMEOUT_SECONDS,
     BaseExecutorContext,
     NodeExecutorContext,
     WorkflowExecutorContext,
@@ -56,38 +66,31 @@ from workflow_server.utils.log_proxy import redirect_log
 logger = logging.getLogger(__name__)
 
 
-@concurrent.process(timeout=DEFAULT_TIMEOUT_SECONDS)
-# type ignore since pebble annotation changes return type
-def stream_node_pebble_timeout(
+def stream_node_process_timeout(
     executor_context: NodeExecutorContext,
     queue: Queue,
-) -> None:
-    _stream_node_wrapper(
-        executor_context=executor_context,
-        queue=queue,
+) -> Process:
+    node_process = Process(
+        target=_stream_node_wrapper,
+        args=(executor_context, queue),
     )
+    node_process.start()
+
+    if node_process.exitcode is not None:
+        queue.put(create_vembda_rejected_event(executor_context, "Internal Server Error", timed_out=True))
+
+    return node_process
 
 
 def _stream_node_wrapper(executor_context: NodeExecutorContext, queue: Queue) -> None:
     try:
         for event in stream_node(executor_context=executor_context):
             queue.put(event)
+    except WorkflowInitializationException as e:
+        queue.put(create_vembda_rejected_event(executor_context, e.message))
     except Exception as e:
         logger.exception(e)
-        queue.put(
-            VembdaExecutionFulfilledEvent(
-                id=uuid4(),
-                timestamp=datetime.now(),
-                trace_id=executor_context.trace_id,
-                span_id=executor_context.execution_id,
-                body=VembdaExecutionFulfilledBody(
-                    exit_code=-1,
-                    stderr="Internal Server Error",
-                    container_overhead_latency=executor_context.container_overhead_latency,
-                ),
-                parent=None,
-            ).model_dump(mode="json")
-        )
+        queue.put(create_vembda_rejected_event(executor_context, "Internal Server Error"))
 
 
 def _stream_workflow_wrapper(
@@ -125,7 +128,7 @@ def _stream_workflow_wrapper(
 def stream_workflow_process_timeout(
     executor_context: WorkflowExecutorContext,
     queue: Queue,
-    cancel_signal: Optional[ThreadingEvent],
+    cancel_signal: ThreadingEvent,
     timeout_signal: ThreadingEvent,
 ) -> Process:
     workflow_process = Process(
@@ -140,19 +143,7 @@ def stream_workflow_process_timeout(
     workflow_process.start()
 
     if workflow_process.exitcode is not None:
-        vembda_fulfilled_event = VembdaExecutionFulfilledEvent(
-            id=uuid4(),
-            timestamp=datetime.now(),
-            trace_id=executor_context.trace_id,
-            span_id=executor_context.execution_id,
-            body=VembdaExecutionFulfilledBody(
-                exit_code=-1,
-                timed_out=True,
-                container_overhead_latency=executor_context.container_overhead_latency,
-            ),
-            parent=None,
-        )
-        queue.put(vembda_fulfilled_event.model_dump(mode="json"))
+        queue.put(create_vembda_rejected_event(executor_context, "Internal Server Error", timed_out=True))
 
     return workflow_process
 
@@ -276,28 +267,74 @@ def stream_node(
     executor_context: NodeExecutorContext,
     disable_redirect: bool = True,
 ) -> Iterator[dict]:
-    def call_node() -> Generator[dict[str, Any], Any, None]:
-        workflow, namespace = _gather_workflow(executor_context)
-        node_module = importlib.import_module(f"{namespace}.{executor_context.node_module}")
-
-        Node = getattr(node_module, executor_context.node_name)
-
-        workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
-        workflow_state = _get_workflow_state(executor_context, workflow_inputs=workflow_inputs)
-
-        node = Node(
-            state=workflow_state,
-            context=workflow._context,
+    workflow, namespace = _gather_workflow(executor_context)
+    Node: Optional[Type[BaseNode]] = None
+
+    for workflow_node in workflow.get_nodes():
+        if executor_context.node_id and workflow_node.__id__ == executor_context.node_id:
+            Node = workflow_node
+            break
+        elif (
+            executor_context.node_module
+            and executor_context.node_name
+            and workflow_node.__name__ == executor_context.node_name
+            and workflow_node.__module__ == f"{namespace}.{executor_context.node_module}"
+        ):
+            Node = workflow_node
+            break
+
+    if not Node:
+        identifier = executor_context.node_id or f"{executor_context.node_module}.{executor_context.node_name}"
+        raise WorkflowInitializationException(
+            message=f"Node '{identifier}' not found in workflow",
+            workflow_definition=workflow.__class__,
         )
 
+    workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
+    workflow_state = _get_workflow_state(executor_context, workflow_inputs=workflow_inputs)
+    node = Node(
+        state=workflow_state,
+        context=workflow._context,
+    )
+
+    def call_node() -> Generator[dict[str, Any], Any, None]:
         executor_context.stream_start_time = time.time_ns()
+        span_id = uuid4()
+        yield NodeExecutionInitiatedEvent(
+            id=uuid4(),
+            timestamp=datetime.now(),
+            trace_id=executor_context.trace_id,
+            span_id=span_id,
+            body=NodeExecutionInitiatedBody(
+                node_definition=Node,
+                inputs=node._inputs,
+            ),
+        ).model_dump(mode="json")
         node_outputs = node.run()
 
         if isinstance(node_outputs, (Iterator)):
             for node_output in node_outputs:
-                yield json.loads(json.dumps(node_output, default=vars))
+                yield NodeExecutionStreamingEvent(
+                    id=uuid4(),
+                    timestamp=datetime.now(),
+                    trace_id=executor_context.trace_id,
+                    span_id=span_id,
+                    body=NodeExecutionStreamingBody(
+                        node_definition=Node,
+                        output=node_output,
+                    ),
+                ).model_dump(mode="json")
         else:
-            yield json.loads(json.dumps(node_outputs, default=vars))
+            yield NodeExecutionFulfilledEvent(
+                id=uuid4(),
+                timestamp=datetime.now(),
+                trace_id=executor_context.trace_id,
+                span_id=span_id,
+                body=NodeExecutionFulfilledBody(
+                    node_definition=Node,
+                    outputs=node_outputs,
+                ),
+            ).model_dump(mode="json")
 
     return _call_stream(
         executor_context=executor_context,

{vellum_workflow_server-1.4.2.post1 → vellum_workflow_server-1.5.0}/src/workflow_server/core/utils.py

@@ -7,7 +7,7 @@ from workflow_server.core.workflow_executor_context import BaseExecutorContext
 
 
 def _create_vembda_rejected_event_base(
-    executor_context: Optional[BaseExecutorContext], error_message: str
+    executor_context: Optional[BaseExecutorContext], error_message: str, timed_out: bool
 ) -> VembdaExecutionFulfilledEvent:
     if executor_context:
         trace_id = executor_context.trace_id
@@ -25,6 +25,7 @@ def _create_vembda_rejected_event_base(
         span_id=span_id,
         body=VembdaExecutionFulfilledBody(
             exit_code=-1,
+            timed_out=timed_out,
             stderr=error_message,
             container_overhead_latency=container_overhead_latency,
         ),
@@ -32,12 +33,16 @@ def _create_vembda_rejected_event_base(
     )
 
 
-def create_vembda_rejected_event(executor_context: Optional[BaseExecutorContext], error_message: str) -> dict:
-    return _create_vembda_rejected_event_base(executor_context, error_message).model_dump(mode="json")
+def create_vembda_rejected_event(
+    executor_context: Optional[BaseExecutorContext], error_message: str, timed_out: bool = False
+) -> dict:
+    return _create_vembda_rejected_event_base(executor_context, error_message, timed_out).model_dump(mode="json")
 
 
-def serialize_vembda_rejected_event(executor_context: Optional[BaseExecutorContext], error_message: str) -> str:
-    return _create_vembda_rejected_event_base(executor_context, error_message).model_dump_json()
+def serialize_vembda_rejected_event(
+    executor_context: Optional[BaseExecutorContext], error_message: str, timed_out: bool = False
+) -> str:
+    return _create_vembda_rejected_event_base(executor_context, error_message, timed_out).model_dump_json()
 
 
 def is_events_emitting_enabled(executor_context: Optional[BaseExecutorContext]) -> bool:

{vellum_workflow_server-1.4.2.post1 → vellum_workflow_server-1.5.0}/src/workflow_server/core/workflow_executor_context.py

@@ -3,6 +3,9 @@ from functools import cached_property
 import os
 from uuid import UUID
 from typing import Any, Optional
+from typing_extensions import Self
+
+from pydantic import model_validator
 
 from vellum import ApiVersionEnum, Vellum, VellumEnvironment
 from vellum.client.core import UniversalBaseModel
@@ -94,5 +97,12 @@ class WorkflowExecutorContext(BaseExecutorContext):
 
 
 class NodeExecutorContext(BaseExecutorContext):
-    node_module: str
-    node_name: str
+    node_id: Optional[UUID] = None
+    node_module: Optional[str] = None
+    node_name: Optional[str] = None
+
+    @model_validator(mode="after")
+    def validate_node_identification(self) -> Self:
+        if not self.node_id and not (self.node_module and self.node_name):
+            raise ValueError("Either node_id or both node_module and node_name must be provided")
+        return self