vellum-workflow-server 0.14.73.post1__py3-none-any.whl → 0.14.73.post3__py3-none-any.whl

{vellum_workflow_server-0.14.73.post1.dist-info → vellum_workflow_server-0.14.73.post3.dist-info}/METADATA

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

{vellum_workflow_server-0.14.73.post1.dist-info → vellum_workflow_server-0.14.73.post3.dist-info}/RECORD

@@ -5,14 +5,14 @@ workflow_server/api/healthz_view.py,sha256=itiRvBDBXncrw8Kbbc73UZLwqMAhgHOR3uSre
 workflow_server/api/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/api/tests/test_input_display_mapping.py,sha256=drBZqMudFyB5wgiUOcMgRXz7E7ge-Qgxbstw4E4f0zE,2211
 workflow_server/api/tests/test_workflow_view.py,sha256=2nscM_QsYPHkkTG8_JhNbE2LmGL5FQKXEtaLLjXouw0,14591
-workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=wmeH8oGo0lKx3YzEJQ9nGjw7adqQD0EL7-xGLPgYWqc,24872
-workflow_server/api/workflow_view.py,sha256=AOHccJXkhzw-iXBo4Q_EaLRynTuNmeQZaQxB5ttgrY4,15434
-workflow_server/code_exec_runner.py,sha256=XSs4w_b0vDUt0HqSakc26Gxx9aoG2wmwOo-JGVL5QJ4,2388
+workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=cUVF3tqV8Ay91YfFr2lT2oXONWRN823Nk5M3EOC8wE8,24133
+workflow_server/api/workflow_view.py,sha256=9_ZI7TGlTRmeNiO47juTRx5TBpAjVmUEGx5C_VqTYMI,14993
+workflow_server/code_exec_runner.py,sha256=njBK48zVUwhAjvap_KY1so-D0UjKgR4UihDuKow3JnM,2274
 workflow_server/config.py,sha256=K5Tavm7wiqCZt0RWWue7zzb8N6e8aWnFOTNlBqEJPcI,1330
 workflow_server/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 workflow_server/core/cancel_workflow.py,sha256=Ffkc3mzmrdMEUcD-sHfEhX4IwVrka-E--SxKA1dUfIU,2185
 workflow_server/core/events.py,sha256=24MA66DVQuaLJJcZrS8IL1Zq4Ohi9CoouKZ5VgoH3Cs,1402
-workflow_server/core/executor.py,sha256=j78gYOgaeJ5Z25nCefnPg4fUPdTLW0tcwaZeWQf0DP0,16588
+workflow_server/core/executor.py,sha256=AjY4qlTxv8T6gq73KXGlWb8gU97xr-p1RMxgbPSg1O0,16288
 workflow_server/core/utils.py,sha256=lgzxkAEjEXPxGXXQlUYTYuCdHht-eDJJmHj5AhEb3_o,1500
 workflow_server/core/workflow_executor_context.py,sha256=a-v48GJbOWUh4JIf_bNwDX-BvfKkg4xwRSPEyRVQmp4,1373
 workflow_server/server.py,sha256=QBU12AaAfAgLqfCDBd24qIJl_mbheiq0-hfcWV7rZM4,1234
@@ -21,13 +21,14 @@ workflow_server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3h
 workflow_server/utils/exit_handler.py,sha256=_FacDVi4zc3bfTA3D2mJsISePlJ8jpLrnGVo5-xZQFs,743
 workflow_server/utils/log_proxy.py,sha256=nugi6fOgAYKX2X9DIc39TG366rsmmDUPoEtG3gzma_Y,3088
 workflow_server/utils/oom_killer.py,sha256=4Sag_iRQWqbp62iIBn6nKP-pxUHguOF93DdVXZTtJDk,2809
-workflow_server/utils/sentry.py,sha256=Pr3xKvHdk0XFSpXgy-55bWI4J3bbf_36gjDyLOs7oVU,855
+workflow_server/utils/sentry.py,sha256=pqx3X_4W3yOzmz8QMJYUEi39skIKWtrTN5nyFhaPkbk,1597
 workflow_server/utils/system_utils.py,sha256=fTzbdpmZ-0bXiNBLYYQdNJWtFAItZgIH8cLJdoXDuQQ,2114
 workflow_server/utils/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+workflow_server/utils/tests/test_sentry_integration.py,sha256=LGmWiaLhFrx-jslrRjRq9JY6Z5ShLZyx_N_L0-FU6OI,2100
 workflow_server/utils/tests/test_system_utils.py,sha256=MdBxI9gxUOpR_JBAHpEz6dGFY6JjxhMSM2oExpqFvNA,4314
 workflow_server/utils/tests/test_utils.py,sha256=qwK5Rmy3RQyjtlUrYAuGuDlBeRzZKsf1yS-y2IpUizQ,6452
 workflow_server/utils/utils.py,sha256=Wqqn-1l2ugkGgy5paWWdt0AVxAyPMQCYcnRSSOMjXlA,4355
-vellum_workflow_server-0.14.73.post1.dist-info/METADATA,sha256=XP2DjToBtdOhbNvZ29sQqlP9b2Zh1MNTIaqyrlvk5gg,2243
-vellum_workflow_server-0.14.73.post1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_workflow_server-0.14.73.post1.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
-vellum_workflow_server-0.14.73.post1.dist-info/RECORD,,
+vellum_workflow_server-0.14.73.post3.dist-info/METADATA,sha256=WI6P39NDrxKvYn99EqvupmJvsghS7m2bvR-InCMll1U,2243
+vellum_workflow_server-0.14.73.post3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_workflow_server-0.14.73.post3.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
+vellum_workflow_server-0.14.73.post3.dist-info/RECORD,,

workflow_server/api/tests/test_workflow_view_stream_workflow_route.py

@@ -10,6 +10,7 @@ import requests_mock
 
 from workflow_server.code_exec_runner import run_code_exec_stream
 from workflow_server.server import create_app
+from workflow_server.utils.system_utils import get_active_process_count
 
 
 def flask_stream(request_body: dict) -> tuple[int, list]:
@@ -775,44 +776,15 @@ class Workflow(BaseWorkflow):
     # WHEN we call the stream route
     status_code, events = flask_stream(request_body)
 
-    # THEN we get a 200 response
-    assert status_code == 200, events
+    # THEN we get a 429 response
+    assert status_code == 429, events
 
-    # THEN we get the expected events
+    # AND we get a simple JSON error response
+    assert len(events) == 1
     assert events[0] == {
-        "id": mock.ANY,
-        "trace_id": mock.ANY,
-        "span_id": str(span_id),
-        "timestamp": mock.ANY,
-        "api_version": "2024-10-25",
-        "parent": None,
-        "name": "vembda.execution.initiated",
-        "body": {
-            "sdk_version": version("vellum-ai"),
-            "server_version": "local",
-        },
+        "detail": f"Workflow server concurrent request rate exceeded. Process count: {get_active_process_count()}"
     }
 
-    # AND we get a vembda.execution.fulfilled event with error
-    assert events[1] == {
-        "id": mock.ANY,
-        "trace_id": events[0]["trace_id"],
-        "span_id": str(span_id),
-        "timestamp": mock.ANY,
-        "api_version": "2024-10-25",
-        "parent": None,
-        "name": "vembda.execution.fulfilled",
-        "body": {
-            "log": "",
-            "exit_code": -1,
-            "stderr": "Workflow server concurrent request rate exceeded. Process count: 0",
-            "container_overhead_latency": mock.ANY,
-            "timed_out": False,
-        },
-    }
-
-    assert len(events) == 2
-
 
 def test_stream_workflow_route__with_environment_variables(both_stream_types):
     # GIVEN a valid request body with environment variables

workflow_server/api/workflow_view.py

@@ -15,7 +15,6 @@ from typing import Any, Dict, Generator, Iterator, Union, cast
 
 from flask import Blueprint, Response, current_app as app, request, stream_with_context
 from pydantic import ValidationError
-import sentry_sdk
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
@@ -97,16 +96,14 @@ def stream_workflow_route() -> Response:
     # if we detect a memory problem just exit us early
     if not wait_for_available_process():
         return Response(
-            stream_with_context(
-                startup_error_generator(
-                    context=context,
-                    message=f"Workflow server concurrent request rate exceeded. "
-                    f"Process count: {get_active_process_count()}",
-                    vembda_initiated_event=vembda_initiated_event,
-                )
+            json.dumps(
+                {
+                    "detail": f"Workflow server concurrent request rate exceeded. "
+                    f"Process count: {get_active_process_count()}"
+                }
             ),
-            status=200,
-            content_type="application/x-ndjson",
+            status=429,
+            content_type="application/json",
             headers=headers,
         )
 
@@ -117,7 +114,6 @@ def stream_workflow_route() -> Response:
         )
         increment_process_count(1)
     except Exception as e:
-        sentry_sdk.set_tag("vellum_trace_id", str(context.trace_id))
         logger.exception(e)
 
         process_output_queue.put(create_vembda_rejected_event(context, traceback.format_exc()))
@@ -204,7 +200,6 @@ def stream_workflow_route() -> Response:
 
                 continue
             except Exception as e:
-                sentry_sdk.set_tag("vellum_trace_id", str(context.trace_id))
                 logger.exception(e)
                 break
 
@@ -302,7 +297,6 @@ def stream_node_route() -> Response:
                     # This happens when theres a problem with the stream function call
                     # itself not the workflow runner
                     yield create_vembda_rejected_event(context, "Internal Server Error")
-                    sentry_sdk.set_tag("vellum_trace_id", str(context.trace_id))
                     app.logger.exception(stream_future.exception())
                     break
                 else:
@@ -379,7 +373,6 @@ def get_version_route() -> tuple[dict, int]:
 
         resp["nodes"] = nodes
     except Exception as e:
-        sentry_sdk.set_tag("vellum_trace_id", "unknown")
         logger.exception(f"Failed to discover nodes: {str(e)}")
         resp["nodes"] = []
 

workflow_server/code_exec_runner.py

@@ -5,8 +5,6 @@ import os
 from uuid import uuid4
 from typing import Optional
 
-import sentry_sdk
-
 from vellum.workflows.exceptions import WorkflowInitializationException
 from workflow_server.api.workflow_view import get_workflow_request_context
 from workflow_server.core.events import VembdaExecutionInitiatedBody, VembdaExecutionInitiatedEvent
@@ -56,7 +54,6 @@ def run_code_exec_stream() -> None:
             fulfilled_event = serialize_vembda_rejected_event(context, str(e))
             print(f"{_EVENT_LINE}{fulfilled_event}")  # noqa: T201
     except Exception as e:
-        sentry_sdk.set_tag("vellum_trace_id", str(context.trace_id) if context else "unknown")
         logger.exception(e)
 
         event = serialize_vembda_rejected_event(context, "Internal Server Error")

workflow_server/core/executor.py

@@ -15,7 +15,6 @@ from uuid import UUID, uuid4
 from typing import Any, Callable, Generator, Iterator, Optional, Tuple, Type
 
 from pebble import concurrent
-import sentry_sdk
 from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
 from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
 
@@ -68,7 +67,6 @@ def _stream_node_wrapper(executor_context: NodeExecutorContext, queue: Queue) ->
         for event in stream_node(executor_context=executor_context):
             queue.put(event)
     except Exception as e:
-        sentry_sdk.set_tag("vellum_trace_id", str(executor_context.trace_id))
         logger.exception(e)
         queue.put(
             VembdaExecutionFulfilledEvent(
@@ -106,7 +104,6 @@ def _stream_workflow_wrapper(executor_context: WorkflowExecutorContext, queue: Q
         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(serialize_vembda_rejected_event(executor_context, "Internal Server Error"))
     queue.put(STREAM_FINISHED_EVENT)
@@ -456,7 +453,6 @@ def _gather_workflow(context: WorkflowExecutorContext) -> Tuple[BaseWorkflow, st
         )
         return workflow, namespace
     except Exception as e:
-        sentry_sdk.set_tag("vellum_trace_id", str(context.trace_id))
         logger.exception("Failed to initialize Workflow")
         raise WorkflowInitializationException(f"Failed to initialize workflow: {e}") from e
 
@@ -465,6 +461,5 @@ def _gather_display_context(workflow: BaseWorkflow, namespace: str) -> Optional[
     try:
         return BaseWorkflowDisplay.gather_event_display_context(namespace, workflow.__class__)
     except Exception:
-        sentry_sdk.set_tag("vellum_trace_id", "unknown")
         logger.exception("Unable to Parse Workflow Display Context")
         return None

workflow_server/utils/sentry.py

@@ -5,10 +5,44 @@ import sentry_sdk
 from sentry_sdk.integrations.logging import LoggingIntegration
 
 
+def _tag_trace_id(event: dict) -> None:
+    if "request" not in event:
+        return
+
+    if not isinstance(event["request"], dict):
+        return
+
+    url = event["request"].get("url")
+    if not isinstance(url, str):
+        return
+
+    if not url.endswith("/workflow/stream"):
+        return
+
+    body = event["request"].get("data")
+    if not isinstance(body, dict):
+        return
+
+    execution_context = body.get("execution_context")
+    if not isinstance(execution_context, dict):
+        return
+
+    trace_id = execution_context.get("trace_id")
+    if not isinstance(trace_id, str):
+        return
+
+    if "tags" not in event:
+        event["tags"] = {}
+
+    event["tags"]["vellum_trace_id"] = trace_id
+
+
 def before_send(event: dict, hint: dict) -> Optional[dict]:
     if "exc_info" in hint:
         _, _, _ = hint["exc_info"]
 
+    _tag_trace_id(event)
+
     return event
 
 

workflow_server/utils/tests/test_sentry_integration.py

@@ -0,0 +1,69 @@
+import pytest
+from uuid import uuid4
+
+from workflow_server.server import create_app
+
+
+@pytest.fixture
+def mock_sentry_capture_envelope(mocker):
+    mock_transport = mocker.patch("sentry_sdk.client.make_transport")
+    return mock_transport.return_value.capture_envelope
+
+
+def test_sentry_integration_with_workflow_endpoints(monkeypatch, mock_sentry_capture_envelope):
+    # GIVEN sentry is configured
+    monkeypatch.setenv("SENTRY_DSN", "https://test-dsn@sentry.io/1234567890")
+
+    # AND our /workflow/stream endpoint raises an exception
+    def mock_get_version():
+        raise Exception("Test exception")
+
+    monkeypatch.setattr("workflow_server.api.workflow_view.get_version", mock_get_version)
+
+    # AND we have a mock trace_id
+    trace_id = str(uuid4())
+
+    # AND we have a mock request body
+    body = {
+        "execution_id": uuid4(),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+
+class Workflow(BaseWorkflow):
+    pass
+""",
+        },
+        "execution_context": {
+            "trace_id": trace_id,
+            "parent_context": {
+                "type": "API_REQUEST",
+                "span_id": str(uuid4()),
+                "parent": None,
+            },
+        },
+    }
+
+    # WHEN we call the /workflow/version endpoint
+    flask_app = create_app()
+
+    with flask_app.test_client() as test_client:
+        response = test_client.post("/workflow/stream", json=body)
+
+        # THEN we get a 500 error
+        assert response.status_code == 500
+
+        # AND sentry captures the error with the correct data
+        assert mock_sentry_capture_envelope.call_count == 1
+        envelope = mock_sentry_capture_envelope.call_args[0][0]
+        event = envelope.get_event()
+        assert event["level"] == "error"
+        assert "Test exception" in event["exception"]["values"][0]["value"]
+
+        # AND the trace_id is tagged
+        assert event["tags"]["vellum_trace_id"] == trace_id