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

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

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

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

@@ -5,16 +5,16 @@ 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=RlAw1tHeIlnOXGrFQN-w3EOLPZkhp6Dfy6d1r7kU5oc,22573
-workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=Qo8u6mPyRCmE2jamY1yIh8l44hgo4-Nwlq03z61ND5g,27031
-workflow_server/api/workflow_view.py,sha256=_WhjNgimTPoS10C-npRWDfJixzg4eHTJ5xIKACStZf4,21943
-workflow_server/code_exec_runner.py,sha256=lBnMIorPZL8zZBye6TjeCIs06WTJM7P2HR07B1fjJJI,2533
+workflow_server/api/tests/test_workflow_view_stream_workflow_route.py,sha256=FaEIgGsbq8M7ZF2opVjBdGCYeRPF_vqsUhFTruLInxA,31358
+workflow_server/api/workflow_view.py,sha256=UZTxWImM9kmof819SCa3ljJKiYEsCrBFD8vp2_f-zAg,21719
+workflow_server/code_exec_runner.py,sha256=E-HsjAL53L-znSMPg7lDiQNzyCjL6W076ZoWWbrSRrU,2217
 workflow_server/config.py,sha256=qmmTr6ty3ZN5LDOFs3TfUxYshYe6Mmn_LanplHHeE9Q,1796
 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=bNryNvXBbg7IHW3BFzzFaGoCO1rGufeEE4FrY9TDA90,18294
-workflow_server/core/utils.py,sha256=aIpSINstLGslP2PIoDLM82_1GlJ1uC_0AIrP-V7Yobo,3230
-workflow_server/core/workflow_executor_context.py,sha256=w3OhV_AXpgh7AxpjEsc0vo-IJypgJcgr5DXJCqGptOU,1587
+workflow_server/core/executor.py,sha256=5cji5KQSukLrABmihJp9cgKjXS145TocOz2cOcnof04,18962
+workflow_server/core/utils.py,sha256=cmwHbKCfXqtUutBD3akGus0Ga7a1xG3zlOw-jEMx6mI,1795
+workflow_server/core/workflow_executor_context.py,sha256=VafZg74t_GQ3_2DEWVroy38rSy_spcAw4c3NrOOWOKY,3198
 workflow_server/server.py,sha256=QBU12AaAfAgLqfCDBd24qIJl_mbheiq0-hfcWV7rZM4,1234
 workflow_server/start.py,sha256=pkwRcms6I4tkVHP06LdrZY6rG_DFHfBx4ioY5X91W5k,2264
 workflow_server/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -28,7 +28,7 @@ workflow_server/utils/tests/test_sentry_integration.py,sha256=LGmWiaLhFrx-jslrRj
 workflow_server/utils/tests/test_system_utils.py,sha256=_4GwXvVvU5BrATxUEWwQIPg0bzQXMWBtiBmjP8MTxJM,4314
 workflow_server/utils/tests/test_utils.py,sha256=0Nq6du8o-iBtTrip9_wgHES53JSiJbVdSXaBnPobw3s,6930
 workflow_server/utils/utils.py,sha256=ZPoM1Suhid22dpB8oEFLux8wx-9iyzmSfWuYxSCrgWk,4774
-vellum_workflow_server-1.4.1.post1.dist-info/METADATA,sha256=sXR0B1bNZZvRQswy2hK4IqACrHG67fw6XLErZfqRoRE,2273
-vellum_workflow_server-1.4.1.post1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_workflow_server-1.4.1.post1.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
-vellum_workflow_server-1.4.1.post1.dist-info/RECORD,,
+vellum_workflow_server-1.4.1.post3.dist-info/METADATA,sha256=tZiQ-25OwiPIcV3lWGZsF9gRuPuAZaypRCtAVbVUirs,2273
+vellum_workflow_server-1.4.1.post3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_workflow_server-1.4.1.post3.dist-info/entry_points.txt,sha256=uB_0yPkr7YV6RhEXzvFReUM8P4OQBlVXD6TN6eb9-oc,277
+vellum_workflow_server-1.4.1.post3.dist-info/RECORD,,

workflow_server/api/tests/test_workflow_view_stream_workflow_route.py

@@ -10,6 +10,8 @@ from uuid import uuid4
 
 import requests_mock
 
+from vellum.workflows.emitters.base import WorkflowEvent
+from vellum.workflows.emitters.vellum_emitter import VellumEmitter
 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
@@ -371,7 +373,9 @@ from vellum.workflows.inputs import BaseInputs
     # THEN we get a 200 response
     assert status_code == 200, events
 
-    # THEN we get the expected events
+    # THEN we get the expected events: vembda initiated, workflow initiated, workflow rejected, vembda fulfilled
+    assert len(events) == 4
+
     assert events[0] == {
         "id": mock.ANY,
         "trace_id": mock.ANY,
@@ -387,7 +391,15 @@ from vellum.workflows.inputs import BaseInputs
         },
     }
 
-    assert events[1] == {
+    assert events[1]["name"] == "workflow.execution.initiated"
+
+    assert events[2]["name"] == "workflow.execution.rejected"
+    assert events[2]["span_id"] == events[1]["span_id"]
+    assert (
+        "Failed to initialize workflow: unexpected indent (inputs.py, line 3)" in events[2]["body"]["error"]["message"]
+    )
+
+    assert events[3] == {
         "id": mock.ANY,
         "trace_id": events[0]["trace_id"],
         "span_id": str(span_id),
@@ -398,15 +410,75 @@ from vellum.workflows.inputs import BaseInputs
         "name": "vembda.execution.fulfilled",
         "body": mock.ANY,
     }
-    assert events[1]["body"] == {
-        "exit_code": -1,
+    assert events[3]["body"] == {
+        "exit_code": 0,
         "log": "",
-        "stderr": "Failed to initialize workflow: unexpected indent (inputs.py, line 3)",
+        "stderr": "",
         "timed_out": False,
         "container_overhead_latency": mock.ANY,
     }
 
-    assert len(events) == 2
+
+def test_stream_workflow_route__invalid_inputs_initialization_events(both_stream_types):
+    """
+    Tests that invalid inputs initialization gets us back a workflow initiated and workflow rejected event.
+    """
+    # GIVEN a valid request body with valid inputs file but omitting required input to cause
+    # WorkflowInitializationException
+    span_id = uuid4()
+    request_body = {
+        "timeout": 360,
+        "execution_id": str(span_id),
+        "inputs": [
+            # Omit the required input to trigger WorkflowInitializationException
+        ],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.state import BaseState
+from .inputs import Inputs
+
+class Workflow(BaseWorkflow[Inputs, BaseState]):
+    class Outputs(BaseWorkflow.Outputs):
+        foo = "hello"
+""",
+            "inputs.py": """\
+from vellum.workflows.inputs import BaseInputs
+
+class Inputs(BaseInputs):
+    foo: str
+""",
+        },
+    }
+
+    # 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: vembda initiated, workflow initiated, workflow rejected, vembda fulfilled
+    assert len(events) == 4
+
+    # AND the first event should be vembda execution initiated
+    assert events[0]["name"] == "vembda.execution.initiated"
+    assert events[0]["span_id"] == str(span_id)
+
+    # AND the second event should be workflow execution initiated
+    assert events[1]["name"] == "workflow.execution.initiated"
+
+    # AND the third event should be workflow execution rejected
+    assert events[2]["name"] == "workflow.execution.rejected"
+    assert events[1]["span_id"] == events[2]["span_id"]
+    assert "Required input variables foo should have defined value" in events[2]["body"]["error"]["message"]
+
+    # AND the fourth event should be vembda execution fulfilled
+    assert events[3]["name"] == "vembda.execution.fulfilled"
+    assert events[3]["span_id"] == str(span_id)
+    assert events[3]["body"]["exit_code"] == 0
 
 
 @pytest.mark.parametrize(
@@ -912,3 +984,50 @@ class Workflow(BaseWorkflow):
 
     # AND we get the expected timeout error message
     assert response_data == {"detail": "Request timed out trying to initiate the Workflow"}
+
+
+@pytest.mark.parametrize("non_process_stream_types", [code_exec_stream, flask_stream_disable_process_wrapper])
+def test_stream_workflow_route__vembda_emitting_calls_monitoring_api(non_process_stream_types):
+    """
+    Tests that the monitoring API is called when vembda emitting is enabled.
+    """
+
+    # GIVEN a valid request body with vembda emitting enabled
+    span_id = uuid4()
+    request_body = {
+        "execution_id": str(span_id),
+        "inputs": [],
+        "environment_api_key": "test",
+        "module": "workflow",
+        "timeout": 360,
+        "feature_flags": {"vembda-event-emitting-enabled": True},
+        "files": {
+            "__init__.py": "",
+            "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+
+class Workflow(BaseWorkflow):
+    class Outputs(BaseWorkflow.Outputs):
+        foo = "hello"
+""",
+        },
+    }
+    emitted_events = []
+
+    def send_events(self, events: list[WorkflowEvent]) -> None:
+        for event in events:
+            emitted_events.append(event)
+
+    VellumEmitter._send_events = send_events
+
+    # WHEN we call the stream route with mocked monitoring API
+    status_code, events = non_process_stream_types(request_body)
+
+    # THEN we get a 200 response
+    assert status_code == 200, events
+
+    # AND the expected workflow events were emitted
+    event_names = [event.name for event in emitted_events]
+    assert len(event_names) == 2, "Should include 2 events"
+    assert "workflow.execution.initiated" in event_names, "Should include workflow.execution.initiated event"
+    assert "workflow.execution.fulfilled" in event_names, "Should include workflow.execution.fulfilled event"

workflow_server/api/workflow_view.py

@@ -38,7 +38,6 @@ from workflow_server.core.events import (
 )
 from workflow_server.core.executor import stream_node_pebble_timeout, stream_workflow, stream_workflow_process_timeout
 from workflow_server.core.utils import (
-    create_vellum_client,
     create_vembda_rejected_event,
     is_events_emitting_enabled,
     serialize_vembda_rejected_event,
@@ -47,6 +46,7 @@ from workflow_server.core.workflow_executor_context import (
     DEFAULT_TIMEOUT_SECONDS,
     NodeExecutorContext,
     WorkflowExecutorContext,
+    create_vellum_client,
 )
 from workflow_server.utils.oom_killer import get_is_oom_killed
 from workflow_server.utils.system_utils import (
@@ -169,11 +169,6 @@ def stream_workflow_route() -> Response:
                 span_id_emitted = True
                 for event in workflow_iterator:
                     yield event
-            except WorkflowInitializationException as e:
-                if not span_id_emitted:
-                    yield f"{SPAN_ID_EVENT}:{uuid4()}"
-
-                yield serialize_vembda_rejected_event(context, str(e))
             except Exception as e:
                 if not span_id_emitted:
                     yield f"{SPAN_ID_EVENT}:{uuid4()}"

workflow_server/code_exec_runner.py

@@ -6,7 +6,6 @@ from threading import Event as ThreadingEvent
 from uuid import uuid4
 from typing import Optional
 
-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
 from workflow_server.core.executor import stream_workflow
@@ -47,18 +46,14 @@ def run_code_exec_stream() -> None:
 
         print(f"{_EVENT_LINE}{initiated_event}")  # noqa: T201
 
-        try:
-            stream_iterator, span_id = stream_workflow(
-                context,
-                disable_redirect=True,
-                # Timeouts are handled at the code exec level right now so just passing in an unused threading event
-                timeout_signal=ThreadingEvent(),
-            )
-            for line in stream_iterator:
-                print(f"{_EVENT_LINE}{json.dumps(line)}")  # noqa: T201
-        except WorkflowInitializationException as e:
-            fulfilled_event = serialize_vembda_rejected_event(context, str(e))
-            print(f"{_EVENT_LINE}{fulfilled_event}")  # noqa: T201
+        stream_iterator, span_id = stream_workflow(
+            context,
+            disable_redirect=True,
+            # Timeouts are handled at the code exec level right now so just passing in an unused threading event
+            timeout_signal=ThreadingEvent(),
+        )
+        for line in stream_iterator:
+            print(f"{_EVENT_LINE}{json.dumps(line)}")  # noqa: T201
     except Exception as e:
         logger.exception(e)
 

workflow_server/core/executor.py

@@ -20,10 +20,10 @@ 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
 
-from vellum import Vellum
 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.types import BaseEvent
 from vellum.workflows.events.workflow import WorkflowEventDisplayContext
 from vellum.workflows.exceptions import WorkflowInitializationException
@@ -44,7 +44,7 @@ from workflow_server.core.events import (
     VembdaExecutionFulfilledBody,
     VembdaExecutionFulfilledEvent,
 )
-from workflow_server.core.utils import create_vellum_client, is_events_emitting_enabled, serialize_vembda_rejected_event
+from workflow_server.core.utils import is_events_emitting_enabled, serialize_vembda_rejected_event
 from workflow_server.core.workflow_executor_context import (
     DEFAULT_TIMEOUT_SECONDS,
     BaseExecutorContext,
@@ -110,11 +110,6 @@ def _stream_workflow_wrapper(
         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(serialize_vembda_rejected_event(executor_context, str(e)))
     except Exception as e:
         if not span_id_emitted:
             queue.put(f"{SPAN_ID_EVENT}:{uuid4()}")
@@ -122,11 +117,6 @@ def _stream_workflow_wrapper(
         logger.exception(e)
         queue.put(serialize_vembda_rejected_event(executor_context, "Internal Server Error"))
 
-    emitter_thread = next(
-        (t for t in threading.enumerate() if t.name.endswith(".background_thread") and t.is_alive()), None
-    )
-    if emitter_thread:
-        emitter_thread.join()
     queue.put(STREAM_FINISHED_EVENT)
 
     exit(0)
@@ -173,27 +163,28 @@ def stream_workflow(
     disable_redirect: bool = True,
     cancel_signal: Optional[ThreadingEvent] = None,
 ) -> tuple[Iterator[dict], UUID]:
-    workflow, namespace = _gather_workflow(executor_context)
-    workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
-    display_context = _gather_display_context(workflow, namespace)
-    workflow_state = (
-        workflow.deserialize_state(
-            executor_context.state,
-            workflow_inputs=workflow_inputs or BaseInputs(),
+    cancel_watcher_kill_switch = ThreadingEvent()
+    try:
+        workflow, namespace = _gather_workflow(executor_context)
+        workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
+
+        display_context = _gather_display_context(workflow, namespace)
+        workflow_state = (
+            workflow.deserialize_state(
+                executor_context.state,
+                workflow_inputs=workflow_inputs or BaseInputs(),
+            )
+            if executor_context.state
+            else None
+        )
+        run_from_node = _get_run_from_node(executor_context, workflow)
+        node_output_mocks = MockNodeExecution.validate_all(
+            executor_context.node_output_mocks,
+            workflow.__class__,
         )
-        if executor_context.state
-        else None
-    )
-    run_from_node = _get_run_from_node(executor_context, workflow)
-    node_output_mocks = MockNodeExecution.validate_all(
-        executor_context.node_output_mocks,
-        workflow.__class__,
-    )
 
-    cancel_watcher_kill_switch = ThreadingEvent()
-    cancel_signal = cancel_signal or ThreadingEvent()
+        cancel_signal = cancel_signal or ThreadingEvent()
 
-    try:
         stream = workflow.stream(
             inputs=workflow_inputs,
             state=workflow_state,
@@ -203,6 +194,26 @@ def stream_workflow(
             entrypoint_nodes=[run_from_node] if run_from_node else None,
             previous_execution_id=executor_context.previous_execution_id,
         )
+    except WorkflowInitializationException as e:
+        cancel_watcher_kill_switch.set()
+        initialization_exception_stream = stream_initialization_exception(e)
+
+        def _stream_generator() -> Generator[dict[str, Any], Any, None]:
+            for event in initialization_exception_stream:
+                yield _dump_event(
+                    event=event,
+                    executor_context=executor_context,
+                )
+
+        return (
+            _call_stream(
+                executor_context=executor_context,
+                stream_generator=_stream_generator,
+                disable_redirect=disable_redirect,
+                timeout_signal=timeout_signal,
+            ),
+            initialization_exception_stream.span_id,
+        )
     except Exception:
         cancel_watcher_kill_switch.set()
         logger.exception("Failed to generate Workflow Stream")
@@ -235,7 +246,6 @@ def stream_workflow(
                 yield _dump_event(
                     event=event,
                     executor_context=executor_context,
-                    client=workflow.context.vellum_client,
                 )
         except Exception as e:
             logger.exception("Failed to generate event from Workflow Stream")
@@ -243,6 +253,16 @@ def stream_workflow(
         finally:
             cancel_watcher_kill_switch.set()
 
+        emitter_thread = next(
+            (t for t in threading.enumerate() if t.name.endswith(".background_thread") and t.is_alive()), None
+        )
+        if emitter_thread:
+            emitter_thread.join()
+
+        timer_thread = next((t for t in threading.enumerate() if t.name.startswith("Thread-")), None)
+        if timer_thread:
+            timer_thread.join()
+
     return (
         _call_stream(
             executor_context=executor_context,
@@ -372,18 +392,13 @@ def _create_workflow(executor_context: BaseExecutorContext, namespace: str) -> B
 
 
 def _create_workflow_context(executor_context: BaseExecutorContext) -> WorkflowContext:
-    vellum_client = create_vellum_client(
-        api_key=executor_context.environment_api_key,
-        api_version=executor_context.api_version,
-    )
-
     if executor_context.environment_variables:
         os.environ.update(executor_context.environment_variables)
 
     namespace = _get_file_namespace(executor_context)
 
     return WorkflowContext(
-        vellum_client=vellum_client,
+        vellum_client=executor_context.vellum_client,
         execution_context=executor_context.execution_context,
         generated_files=executor_context.files,
         namespace=namespace,
@@ -403,9 +418,11 @@ def _get_file_namespace(executor_context: BaseExecutorContext) -> str:
     )
 
 
-def _dump_event(event: BaseEvent, executor_context: BaseExecutorContext, client: Vellum) -> dict:
+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, client)})
+    dump = event.model_dump(
+        mode="json", context={"event_enricher": lambda event: event_enricher(event, executor_context.vellum_client)}
+    )
     if dump["name"] in {
         "workflow.execution.initiated",
         "workflow.execution.fulfilled",

workflow_server/core/utils.py

@@ -1,10 +1,7 @@
 from datetime import datetime
-import os
 from uuid import uuid4
 from typing import Optional
 
-from vellum import ApiVersionEnum, Vellum, VellumEnvironment
-from workflow_server.config import IS_VPC, VELLUM_API_URL_HOST, VELLUM_API_URL_PORT
 from workflow_server.core.events import VembdaExecutionFulfilledBody, VembdaExecutionFulfilledEvent
 from workflow_server.core.workflow_executor_context import BaseExecutorContext
 
@@ -51,42 +48,3 @@ def is_events_emitting_enabled(executor_context: Optional[BaseExecutorContext])
         return False
 
     return executor_context.feature_flags.get("vembda-event-emitting-enabled") or False
-
-
-def create_vellum_client(
-    api_key: str,
-    api_version: Optional[ApiVersionEnum] = None,
-) -> Vellum:
-    """
-    Create a VellumClient with proper environment configuration.
-
-    Args:
-        api_key: The API key for the Vellum client
-        api_version: Optional API version to use
-
-    Returns:
-        Configured Vellum client instance
-
-    Note: Ideally we replace this with `vellum.workflows.vellum_client.create_vellum_client`
-    """
-    if IS_VPC:
-        environment = VellumEnvironment(
-            default=os.getenv("VELLUM_DEFAULT_API_URL", VellumEnvironment.PRODUCTION.default),
-            documents=os.getenv("VELLUM_DOCUMENTS_API_URL", VellumEnvironment.PRODUCTION.documents),
-            predict=os.getenv("VELLUM_PREDICT_API_URL", VellumEnvironment.PRODUCTION.predict),
-        )
-    elif os.getenv("USE_LOCAL_VELLUM_API") == "true":
-        VELLUM_API_URL = f"http://{VELLUM_API_URL_HOST}:{VELLUM_API_URL_PORT}"
-        environment = VellumEnvironment(
-            default=VELLUM_API_URL,
-            documents=VELLUM_API_URL,
-            predict=VELLUM_API_URL,
-        )
-    else:
-        environment = VellumEnvironment.PRODUCTION
-
-    return Vellum(
-        api_key=api_key,
-        environment=environment,
-        api_version=api_version,
-    )

workflow_server/core/workflow_executor_context.py

@@ -1,14 +1,57 @@
 from dataclasses import field
+import os
 from uuid import UUID
 from typing import Any, Optional
 
+from _pytest.compat import cached_property
+
+from vellum import ApiVersionEnum, Vellum, VellumEnvironment
 from vellum.client.core import UniversalBaseModel
-from vellum.client.types.api_version_enum import ApiVersionEnum
 from vellum.workflows.context import ExecutionContext
+from workflow_server.config import IS_VPC, VELLUM_API_URL_HOST, VELLUM_API_URL_PORT
 
 DEFAULT_TIMEOUT_SECONDS = 60 * 30
 
 
+def create_vellum_client(
+    api_key: str,
+    api_version: Optional[ApiVersionEnum] = None,
+) -> Vellum:
+    """
+    Create a VellumClient with proper environment configuration.
+
+    Args:
+        api_key: The API key for the Vellum client
+        api_version: Optional API version to use
+
+    Returns:
+        Configured Vellum client instance
+
+    Note: Ideally we replace this with `vellum.workflows.vellum_client.create_vellum_client`
+    """
+    if IS_VPC:
+        environment = VellumEnvironment(
+            default=os.getenv("VELLUM_DEFAULT_API_URL", VellumEnvironment.PRODUCTION.default),
+            documents=os.getenv("VELLUM_DOCUMENTS_API_URL", VellumEnvironment.PRODUCTION.documents),
+            predict=os.getenv("VELLUM_PREDICT_API_URL", VellumEnvironment.PRODUCTION.predict),
+        )
+    elif os.getenv("USE_LOCAL_VELLUM_API") == "true":
+        VELLUM_API_URL = f"http://{VELLUM_API_URL_HOST}:{VELLUM_API_URL_PORT}"
+        environment = VellumEnvironment(
+            default=VELLUM_API_URL,
+            documents=VELLUM_API_URL,
+            predict=VELLUM_API_URL,
+        )
+    else:
+        environment = VellumEnvironment.PRODUCTION
+
+    return Vellum(
+        api_key=api_key,
+        environment=environment,
+        api_version=api_version,
+    )
+
+
 class BaseExecutorContext(UniversalBaseModel):
     inputs: dict
     state: Optional[dict] = None
@@ -35,6 +78,13 @@ class BaseExecutorContext(UniversalBaseModel):
     def trace_id(self) -> UUID:
         return self.execution_context.trace_id
 
+    @cached_property
+    def vellum_client(self) -> Vellum:
+        return create_vellum_client(
+            api_key=self.environment_api_key,
+            api_version=self.api_version,
+        )
+
     def __hash__(self) -> int:
         # do we think we need anything else for a unique hash for caching?
         return hash(str(self.execution_id))