vellum-workflow-server 0.14.72.post3__tar.gz → 0.14.72.post4__tar.gz

{vellum_workflow_server-0.14.72.post3 → vellum_workflow_server-0.14.72.post4}/PKG-INFO

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

{vellum_workflow_server-0.14.72.post3 → vellum_workflow_server-0.14.72.post4}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-workflow-server"
 
 [tool.poetry]
 name = "vellum-workflow-server"
-version = "0.14.72.post3"
+version = "0.14.72.post4"
 description = ""
 readme = "README.md"
 authors = []

{vellum_workflow_server-0.14.72.post3 → vellum_workflow_server-0.14.72.post4}/src/workflow_server/api/tests/test_workflow_view_stream_workflow_route.py

@@ -83,6 +83,7 @@ class Workflow(BaseWorkflow):
     # AND the version headers are present
     assert "X-Vellum-SDK-Version" in response.headers
     assert "X-Vellum-Server-Version" in response.headers
+    assert "X-Vellum-Workflow-Span-Id" in response.headers
 
 
 def test_stream_workflow_route__happy_path(both_stream_types):

{vellum_workflow_server-0.14.72.post3 → vellum_workflow_server-0.14.72.post4}/src/workflow_server/api/workflow_view.py

@@ -23,6 +23,7 @@ from vellum.workflows.nodes import BaseNode
 from vellum.workflows.utils.names import pascal_to_title_case
 from workflow_server.config import MEMORY_LIMIT_MB
 from workflow_server.core.events import (
+    SPAN_ID_EVENT,
     STREAM_FINISHED_EVENT,
     VEMBDA_EXECUTION_FULFILLED_EVENT_NAME,
     VembdaExecutionFulfilledBody,
@@ -84,7 +85,12 @@ def stream_workflow_route() -> Response:
         parent=None,
     )
 
-    process_output_queue: Queue[dict] = Queue()
+    process_output_queue: Queue[Union[str, dict]] = Queue()
+
+    headers = {
+        "X-Vellum-SDK-Version": vembda_initiated_event.body.sdk_version,
+        "X-Vellum-Server-Version": vembda_initiated_event.body.server_version,
+    }
 
     # We can exceed the concurrency count currently with long running workflows due to a knative issue. So here
     # if we detect a memory problem just exit us early
@@ -100,10 +106,7 @@ def stream_workflow_route() -> Response:
             ),
             status=200,
             content_type="application/x-ndjson",
-            headers={
-                "X-Vellum-SDK-Version": vembda_initiated_event.body.sdk_version,
-                "X-Vellum-Server-Version": vembda_initiated_event.body.server_version,
-            },
+            headers=headers,
         )
 
     try:
@@ -130,6 +133,19 @@ def stream_workflow_route() -> Response:
         )
         process_output_queue.put(vembda_fulfilled_event.model_dump(mode="json"))
 
+    first_item = process_output_queue.get(timeout=0.1)
+    if isinstance(first_item, str) and first_item.startswith(SPAN_ID_EVENT):
+        span_id = first_item.split(":")[1]
+        headers["X-Vellum-Workflow-Span-Id"] = span_id
+    else:
+        logger.error("Workflow stream did not start with span id event")
+        return Response(
+            json.dumps({"detail": "Internal Server Error"}),
+            status=500,
+            content_type="application/json",
+            headers=headers,
+        )
+
     def process_events(queue: Queue) -> Iterator[Union[str, dict]]:
         event: Union[str, dict]
         loops = 0
@@ -267,10 +283,7 @@ def stream_workflow_route() -> Response:
         stream_with_context(generator()),
         status=200,
         content_type="application/x-ndjson",
-        headers={
-            "X-Vellum-SDK-Version": vembda_initiated_event.body.sdk_version,
-            "X-Vellum-Server-Version": vembda_initiated_event.body.server_version,
-        },
+        headers=headers,
     )
     return resp
 
@@ -347,14 +360,16 @@ def stream_node_route() -> Response:
             yield "\n"
             yield json.dumps(row)
 
+    headers = {
+        "X-Vellum-SDK-Version": vembda_initiated_event.body.sdk_version,
+        "X-Vellum-Server-Version": vembda_initiated_event.body.server_version,
+    }
+
     resp = Response(
         stream_with_context(generator()),
         status=200,
         content_type="application/x-ndjson",
-        headers={
-            "X-Vellum-SDK-Version": vembda_initiated_event.body.sdk_version,
-            "X-Vellum-Server-Version": vembda_initiated_event.body.server_version,
-        },
+        headers=headers,
     )
     return resp
 

{vellum_workflow_server-0.14.72.post3 → vellum_workflow_server-0.14.72.post4}/src/workflow_server/code_exec_runner.py

@@ -53,7 +53,8 @@ def run_code_exec_stream() -> None:
         print(f"{_EVENT_LINE}{initiated_event}")  # noqa: T201
 
         try:
-            for line in stream_workflow(context, disable_redirect=True):
+            stream_iterator, span_id = stream_workflow(context, disable_redirect=True)
+            for line in stream_iterator:
                 print(f"{_EVENT_LINE}{json.dumps(line)}")  # noqa: T201
         except WorkflowInitializationException as e:
             fulfilled_event = VembdaExecutionFulfilledEvent(

{vellum_workflow_server-0.14.72.post3 → vellum_workflow_server-0.14.72.post4}/src/workflow_server/core/events.py

@@ -8,6 +8,7 @@ from vellum.workflows.nodes import BaseNode
 VEMBDA_EXECUTION_INITIATED_EVENT_NAME = "vembda.execution.initiated"
 VEMBDA_EXECUTION_FULFILLED_EVENT_NAME = "vembda.execution.fulfilled"
 STREAM_FINISHED_EVENT = "STREAM_FINISHED"
+SPAN_ID_EVENT = "SPAN_ID_EVENT"
 
 
 class VembdaExecutionFulfilledBody(UniversalBaseModel):

{vellum_workflow_server-0.14.72.post3 → vellum_workflow_server-0.14.72.post4}/src/workflow_server/core/executor.py

@@ -11,7 +11,7 @@ import sys
 from threading import Event as ThreadingEvent
 import time
 from traceback import format_exc
-from uuid import uuid4
+from uuid import UUID, uuid4
 from typing import Any, Callable, Generator, Iterator, Optional, Tuple, Type
 
 from pebble import concurrent
@@ -35,6 +35,7 @@ from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 from workflow_server.config import IS_VPC, VELLUM_API_URL_HOST, VELLUM_API_URL_PORT
 from workflow_server.core.cancel_workflow import CancelWorkflowWatcherThread
 from workflow_server.core.events import (
+    SPAN_ID_EVENT,
     STREAM_FINISHED_EVENT,
     VembdaExecutionFulfilledBody,
     VembdaExecutionFulfilledEvent,
@@ -86,11 +87,20 @@ def _stream_node_wrapper(executor_context: NodeExecutorContext, queue: Queue) ->
 
 
 def _stream_workflow_wrapper(executor_context: WorkflowExecutorContext, queue: Queue) -> None:
+    span_id_emitted = False
     try:
-        for event in stream_workflow(executor_context=executor_context):
+        stream_iterator, span_id = stream_workflow(executor_context=executor_context)
+
+        queue.put(f"{SPAN_ID_EVENT}:{span_id}")
+        span_id_emitted = True
+
+        for event in stream_iterator:
             queue.put(json.dumps(event))
 
     except WorkflowInitializationException as e:
+        if not span_id_emitted:
+            queue.put(f"{SPAN_ID_EVENT}:{uuid4()}")
+
         queue.put(
             VembdaExecutionFulfilledEvent(
                 id=uuid4(),
@@ -106,6 +116,9 @@ def _stream_workflow_wrapper(executor_context: WorkflowExecutorContext, queue: Q
             ).model_dump(mode="json")
         )
     except Exception as e:
+        if not span_id_emitted:
+            queue.put(f"{SPAN_ID_EVENT}:{uuid4()}")
+
         sentry_sdk.set_tag("vellum_trace_id", str(executor_context.trace_id))
         logger.exception(e)
         queue.put(
@@ -161,7 +174,7 @@ def stream_workflow_process_timeout(
 def stream_workflow(
     executor_context: WorkflowExecutorContext,
     disable_redirect: bool = True,
-) -> Iterator[dict]:
+) -> tuple[Iterator[dict], UUID]:
     workflow, namespace = _gather_workflow(executor_context)
     workflow_inputs = _get_workflow_inputs(executor_context)
     display_context = _gather_display_context(workflow, namespace)
@@ -179,30 +192,38 @@ def stream_workflow(
         workflow.__class__,
     )
 
-    def call_workflow() -> Generator[dict[str, Any], Any, None]:
-        cancel_watcher_kill_switch = ThreadingEvent()
-        cancel_signal = ThreadingEvent()
-        cancel_watcher = CancelWorkflowWatcherThread(
-            kill_switch=cancel_watcher_kill_switch,
-            execution_id=executor_context.execution_id,
-            timeout_seconds=executor_context.timeout,
-            vembda_public_url=executor_context.vembda_public_url,
+    cancel_watcher_kill_switch = ThreadingEvent()
+    cancel_signal = ThreadingEvent()
+    cancel_watcher = CancelWorkflowWatcherThread(
+        kill_switch=cancel_watcher_kill_switch,
+        execution_id=executor_context.execution_id,
+        timeout_seconds=executor_context.timeout,
+        vembda_public_url=executor_context.vembda_public_url,
+        cancel_signal=cancel_signal,
+    )
+
+    try:
+        if executor_context.vembda_public_url:
+            cancel_watcher.start()
+    except Exception:
+        logger.exception("Failed to start cancel watcher")
+
+    try:
+        stream = workflow.stream(
+            inputs=workflow_inputs,
+            state=workflow_state,
+            node_output_mocks=node_output_mocks,
+            event_filter=all_workflow_event_filter,
             cancel_signal=cancel_signal,
+            entrypoint_nodes=[run_from_node] if run_from_node else None,
         )
+    except Exception:
+        cancel_watcher_kill_switch.set()
+        logger.exception("Failed to generate Workflow Stream")
+        raise
 
+    def call_workflow() -> Generator[dict[str, Any], Any, None]:
         try:
-            if executor_context.vembda_public_url:
-                cancel_watcher.start()
-
-            stream = workflow.stream(
-                inputs=workflow_inputs,
-                state=workflow_state,
-                node_output_mocks=node_output_mocks,
-                event_filter=all_workflow_event_filter,
-                cancel_signal=cancel_signal,
-                entrypoint_nodes=[run_from_node] if run_from_node else None,
-            )
-
             first = True
             for event in stream:
                 if first:
@@ -222,13 +243,19 @@ def stream_workflow(
                     event=event,
                     executor_context=executor_context,
                 )
+        except Exception as e:
+            logger.exception("Failed to generate event from Workflow Stream")
+            raise e
         finally:
             cancel_watcher_kill_switch.set()
 
-    return _call_stream(
-        executor_context=executor_context,
-        stream_generator=call_workflow,
-        disable_redirect=disable_redirect,
+    return (
+        _call_stream(
+            executor_context=executor_context,
+            stream_generator=call_workflow,
+            disable_redirect=disable_redirect,
+        ),
+        stream.span_id,
     )