vellum-workflow-server 1.7.15__tar.gz → 1.7.15.post1__tar.gz

{vellum_workflow_server-1.7.15 → vellum_workflow_server-1.7.15.post1}/PKG-INFO

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

{vellum_workflow_server-1.7.15 → vellum_workflow_server-1.7.15.post1}/pyproject.toml

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

{vellum_workflow_server-1.7.15 → vellum_workflow_server-1.7.15.post1}/src/workflow_server/api/tests/test_workflow_view_stream_workflow_route.py

@@ -131,8 +131,9 @@ class Workflow(BaseWorkflow):
         },
     }
 
-    # WHEN we call the stream route
-    status_code, events = both_stream_types(request_body)
+    with mock.patch("builtins.open", mock.mock_open(read_data="104857600")):
+        # 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
@@ -164,9 +165,26 @@ class Workflow(BaseWorkflow):
     assert isinstance(display_context["workflow_inputs"], dict)
     assert isinstance(display_context["workflow_outputs"], dict)
     assert "foo" in display_context["workflow_outputs"]
+
+    # AND the initiated event should have server_metadata with version info and memory usage
+    assert "server_metadata" in events[1]["body"], events[1]["body"]
+    server_metadata = events[1]["body"]["server_metadata"]
+    assert server_metadata is not None, "server_metadata should not be None"
+    assert "server_version" in server_metadata
+    assert "sdk_version" in server_metadata
+    assert "memory_usage_mb" in server_metadata
+    assert isinstance(server_metadata["memory_usage_mb"], (int, float))
+
     assert events[2]["name"] == "workflow.execution.fulfilled", events[2]
     assert events[2]["body"]["workflow_definition"]["module"] == ["test", "workflow"]
 
+    # AND the fulfilled event should have server_metadata with memory usage
+    assert "server_metadata" in events[2]["body"], events[2]["body"]
+    fulfilled_metadata = events[2]["body"]["server_metadata"]
+    assert fulfilled_metadata is not None, "fulfilled server_metadata should not be None"
+    assert "memory_usage_mb" in fulfilled_metadata
+    assert isinstance(fulfilled_metadata["memory_usage_mb"], (int, float))
+
     assert events[3] == {
         "id": mock.ANY,
         "trace_id": events[0]["trace_id"],

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

@@ -22,6 +22,7 @@ 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.types import BaseEvent
+from vellum.workflows.events.workflow import WorkflowEvent
 from vellum.workflows.exceptions import WorkflowInitializationException
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import BaseNode
@@ -51,6 +52,8 @@ from workflow_server.core.workflow_executor_context import (
     WorkflowExecutorContext,
 )
 from workflow_server.utils.log_proxy import redirect_log
+from workflow_server.utils.system_utils import get_memory_in_use_mb
+from workflow_server.utils.utils import get_version
 
 logger = logging.getLogger(__name__)
 
@@ -401,10 +404,43 @@ def get_random_namespace() -> str:
     return "workflow_tmp_" + "".join(random.choice(string.ascii_letters + string.digits) for i in range(14))
 
 
+def _enrich_event(event: WorkflowEvent, vellum_client: Optional[Any]) -> WorkflowEvent:
+    """
+    Enrich an event with metadata based on the event type.
+
+    For initiated events, include server and SDK versions.
+    For fulfilled events with WORKFLOW_DEPLOYMENT parent, include memory usage.
+    """
+    metadata: Optional[dict] = None
+
+    try:
+        is_deployment = event.parent and event.parent.type in ("WORKFLOW_DEPLOYMENT", "EXTERNAL")
+
+        if event.name == "workflow.execution.initiated" and is_deployment:
+            metadata = {
+                **get_version(),
+            }
+
+            memory_mb = get_memory_in_use_mb()
+            if memory_mb is not None:
+                metadata["memory_usage_mb"] = memory_mb
+        elif event.name == "workflow.execution.fulfilled" and is_deployment:
+            metadata = {}
+            memory_mb = get_memory_in_use_mb()
+            if memory_mb is not None:
+                metadata["memory_usage_mb"] = memory_mb
+    except Exception:
+        pass
+
+    return event_enricher(event, vellum_client, metadata=metadata)
+
+
 def _dump_event(event: BaseEvent, executor_context: BaseExecutorContext) -> dict:
     module_base = executor_context.module.split(".")
+
     dump = event.model_dump(
-        mode="json", context={"event_enricher": lambda event: event_enricher(event, executor_context.vellum_client)}
+        mode="json",
+        context={"event_enricher": lambda event: _enrich_event(event, executor_context.vellum_client)},
     )
     if dump["name"] in {
         "workflow.execution.initiated",