vellum-workflow-server 1.4.1.post2__tar.gz → 1.4.1.post3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of vellum-workflow-server might be problematic. Click here for more details.

Files changed (34) hide show
  1. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/PKG-INFO +1 -1
  2. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/pyproject.toml +1 -1
  3. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/api/tests/test_workflow_view_stream_workflow_route.py +76 -6
  4. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/api/workflow_view.py +1 -6
  5. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/code_exec_runner.py +8 -13
  6. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/core/executor.py +46 -34
  7. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/core/utils.py +0 -42
  8. vellum_workflow_server-1.4.1.post3/src/workflow_server/core/workflow_executor_context.py +99 -0
  9. vellum_workflow_server-1.4.1.post2/src/workflow_server/core/workflow_executor_context.py +0 -49
  10. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/README.md +0 -0
  11. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/__init__.py +0 -0
  12. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/api/__init__.py +0 -0
  13. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/api/auth_middleware.py +0 -0
  14. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/api/healthz_view.py +0 -0
  15. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/api/tests/__init__.py +0 -0
  16. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/api/tests/test_input_display_mapping.py +0 -0
  17. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/api/tests/test_workflow_view.py +0 -0
  18. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/config.py +0 -0
  19. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/core/__init__.py +0 -0
  20. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/core/cancel_workflow.py +0 -0
  21. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/core/events.py +0 -0
  22. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/server.py +0 -0
  23. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/start.py +0 -0
  24. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/__init__.py +0 -0
  25. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/exit_handler.py +0 -0
  26. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/log_proxy.py +0 -0
  27. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/oom_killer.py +0 -0
  28. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/sentry.py +0 -0
  29. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/system_utils.py +0 -0
  30. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/tests/__init__.py +0 -0
  31. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/tests/test_sentry_integration.py +0 -0
  32. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/tests/test_system_utils.py +0 -0
  33. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/tests/test_utils.py +0 -0
  34. {vellum_workflow_server-1.4.1.post2 → vellum_workflow_server-1.4.1.post3}/src/workflow_server/utils/utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: vellum-workflow-server
3
- Version: 1.4.1.post2
3
+ Version: 1.4.1.post3
4
4
  Summary:
5
5
  License: AGPL
6
6
  Requires-Python: >=3.9.0,<4
@@ -3,7 +3,7 @@ name = "vellum-workflow-server"
3
3
 
4
4
  [tool.poetry]
5
5
  name = "vellum-workflow-server"
6
- version = "1.4.1.post2"
6
+ version = "1.4.1.post3"
7
7
  description = ""
8
8
  readme = "README.md"
9
9
  authors = []
@@ -373,7 +373,9 @@ from vellum.workflows.inputs import BaseInputs
373
373
  # THEN we get a 200 response
374
374
  assert status_code == 200, events
375
375
 
376
- # THEN we get the expected events
376
+ # THEN we get the expected events: vembda initiated, workflow initiated, workflow rejected, vembda fulfilled
377
+ assert len(events) == 4
378
+
377
379
  assert events[0] == {
378
380
  "id": mock.ANY,
379
381
  "trace_id": mock.ANY,
@@ -389,7 +391,15 @@ from vellum.workflows.inputs import BaseInputs
389
391
  },
390
392
  }
391
393
 
392
- assert events[1] == {
394
+ assert events[1]["name"] == "workflow.execution.initiated"
395
+
396
+ assert events[2]["name"] == "workflow.execution.rejected"
397
+ assert events[2]["span_id"] == events[1]["span_id"]
398
+ assert (
399
+ "Failed to initialize workflow: unexpected indent (inputs.py, line 3)" in events[2]["body"]["error"]["message"]
400
+ )
401
+
402
+ assert events[3] == {
393
403
  "id": mock.ANY,
394
404
  "trace_id": events[0]["trace_id"],
395
405
  "span_id": str(span_id),
@@ -400,15 +410,75 @@ from vellum.workflows.inputs import BaseInputs
400
410
  "name": "vembda.execution.fulfilled",
401
411
  "body": mock.ANY,
402
412
  }
403
- assert events[1]["body"] == {
404
- "exit_code": -1,
413
+ assert events[3]["body"] == {
414
+ "exit_code": 0,
405
415
  "log": "",
406
- "stderr": "Failed to initialize workflow: unexpected indent (inputs.py, line 3)",
416
+ "stderr": "",
407
417
  "timed_out": False,
408
418
  "container_overhead_latency": mock.ANY,
409
419
  }
410
420
 
411
- assert len(events) == 2
421
+
422
+ def test_stream_workflow_route__invalid_inputs_initialization_events(both_stream_types):
423
+ """
424
+ Tests that invalid inputs initialization gets us back a workflow initiated and workflow rejected event.
425
+ """
426
+ # GIVEN a valid request body with valid inputs file but omitting required input to cause
427
+ # WorkflowInitializationException
428
+ span_id = uuid4()
429
+ request_body = {
430
+ "timeout": 360,
431
+ "execution_id": str(span_id),
432
+ "inputs": [
433
+ # Omit the required input to trigger WorkflowInitializationException
434
+ ],
435
+ "environment_api_key": "test",
436
+ "module": "workflow",
437
+ "files": {
438
+ "__init__.py": "",
439
+ "workflow.py": """\
440
+ from vellum.workflows import BaseWorkflow
441
+ from vellum.workflows.state import BaseState
442
+ from .inputs import Inputs
443
+
444
+ class Workflow(BaseWorkflow[Inputs, BaseState]):
445
+ class Outputs(BaseWorkflow.Outputs):
446
+ foo = "hello"
447
+ """,
448
+ "inputs.py": """\
449
+ from vellum.workflows.inputs import BaseInputs
450
+
451
+ class Inputs(BaseInputs):
452
+ foo: str
453
+ """,
454
+ },
455
+ }
456
+
457
+ # WHEN we call the stream route
458
+ status_code, events = both_stream_types(request_body)
459
+
460
+ # THEN we get a 200 response
461
+ assert status_code == 200, events
462
+
463
+ # THEN we get the expected events: vembda initiated, workflow initiated, workflow rejected, vembda fulfilled
464
+ assert len(events) == 4
465
+
466
+ # AND the first event should be vembda execution initiated
467
+ assert events[0]["name"] == "vembda.execution.initiated"
468
+ assert events[0]["span_id"] == str(span_id)
469
+
470
+ # AND the second event should be workflow execution initiated
471
+ assert events[1]["name"] == "workflow.execution.initiated"
472
+
473
+ # AND the third event should be workflow execution rejected
474
+ assert events[2]["name"] == "workflow.execution.rejected"
475
+ assert events[1]["span_id"] == events[2]["span_id"]
476
+ assert "Required input variables foo should have defined value" in events[2]["body"]["error"]["message"]
477
+
478
+ # AND the fourth event should be vembda execution fulfilled
479
+ assert events[3]["name"] == "vembda.execution.fulfilled"
480
+ assert events[3]["span_id"] == str(span_id)
481
+ assert events[3]["body"]["exit_code"] == 0
412
482
 
413
483
 
414
484
  @pytest.mark.parametrize(
@@ -38,7 +38,6 @@ from workflow_server.core.events import (
38
38
  )
39
39
  from workflow_server.core.executor import stream_node_pebble_timeout, stream_workflow, stream_workflow_process_timeout
40
40
  from workflow_server.core.utils import (
41
- create_vellum_client,
42
41
  create_vembda_rejected_event,
43
42
  is_events_emitting_enabled,
44
43
  serialize_vembda_rejected_event,
@@ -47,6 +46,7 @@ from workflow_server.core.workflow_executor_context import (
47
46
  DEFAULT_TIMEOUT_SECONDS,
48
47
  NodeExecutorContext,
49
48
  WorkflowExecutorContext,
49
+ create_vellum_client,
50
50
  )
51
51
  from workflow_server.utils.oom_killer import get_is_oom_killed
52
52
  from workflow_server.utils.system_utils import (
@@ -169,11 +169,6 @@ def stream_workflow_route() -> Response:
169
169
  span_id_emitted = True
170
170
  for event in workflow_iterator:
171
171
  yield event
172
- except WorkflowInitializationException as e:
173
- if not span_id_emitted:
174
- yield f"{SPAN_ID_EVENT}:{uuid4()}"
175
-
176
- yield serialize_vembda_rejected_event(context, str(e))
177
172
  except Exception as e:
178
173
  if not span_id_emitted:
179
174
  yield f"{SPAN_ID_EVENT}:{uuid4()}"
@@ -6,7 +6,6 @@ from threading import Event as ThreadingEvent
6
6
  from uuid import uuid4
7
7
  from typing import Optional
8
8
 
9
- from vellum.workflows.exceptions import WorkflowInitializationException
10
9
  from workflow_server.api.workflow_view import get_workflow_request_context
11
10
  from workflow_server.core.events import VembdaExecutionInitiatedBody, VembdaExecutionInitiatedEvent
12
11
  from workflow_server.core.executor import stream_workflow
@@ -47,18 +46,14 @@ def run_code_exec_stream() -> None:
47
46
 
48
47
  print(f"{_EVENT_LINE}{initiated_event}") # noqa: T201
49
48
 
50
- try:
51
- stream_iterator, span_id = stream_workflow(
52
- context,
53
- disable_redirect=True,
54
- # Timeouts are handled at the code exec level right now so just passing in an unused threading event
55
- timeout_signal=ThreadingEvent(),
56
- )
57
- for line in stream_iterator:
58
- print(f"{_EVENT_LINE}{json.dumps(line)}") # noqa: T201
59
- except WorkflowInitializationException as e:
60
- fulfilled_event = serialize_vembda_rejected_event(context, str(e))
61
- print(f"{_EVENT_LINE}{fulfilled_event}") # noqa: T201
49
+ stream_iterator, span_id = stream_workflow(
50
+ context,
51
+ disable_redirect=True,
52
+ # Timeouts are handled at the code exec level right now so just passing in an unused threading event
53
+ timeout_signal=ThreadingEvent(),
54
+ )
55
+ for line in stream_iterator:
56
+ print(f"{_EVENT_LINE}{json.dumps(line)}") # noqa: T201
62
57
  except Exception as e:
63
58
  logger.exception(e)
64
59
 
@@ -20,10 +20,10 @@ from vellum_ee.workflows.display.utils.events import event_enricher
20
20
  from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
21
21
  from vellum_ee.workflows.server.virtual_file_loader import VirtualFileFinder
22
22
 
23
- from vellum import Vellum
24
23
  from vellum.workflows import BaseWorkflow
25
24
  from vellum.workflows.emitters.base import BaseWorkflowEmitter
26
25
  from vellum.workflows.emitters.vellum_emitter import VellumEmitter
26
+ from vellum.workflows.events.exception_handling import stream_initialization_exception
27
27
  from vellum.workflows.events.types import BaseEvent
28
28
  from vellum.workflows.events.workflow import WorkflowEventDisplayContext
29
29
  from vellum.workflows.exceptions import WorkflowInitializationException
@@ -44,7 +44,7 @@ from workflow_server.core.events import (
44
44
  VembdaExecutionFulfilledBody,
45
45
  VembdaExecutionFulfilledEvent,
46
46
  )
47
- from workflow_server.core.utils import create_vellum_client, is_events_emitting_enabled, serialize_vembda_rejected_event
47
+ from workflow_server.core.utils import is_events_emitting_enabled, serialize_vembda_rejected_event
48
48
  from workflow_server.core.workflow_executor_context import (
49
49
  DEFAULT_TIMEOUT_SECONDS,
50
50
  BaseExecutorContext,
@@ -110,11 +110,6 @@ def _stream_workflow_wrapper(
110
110
  for event in stream_iterator:
111
111
  queue.put(json.dumps(event))
112
112
 
113
- except WorkflowInitializationException as e:
114
- if not span_id_emitted:
115
- queue.put(f"{SPAN_ID_EVENT}:{uuid4()}")
116
-
117
- queue.put(serialize_vembda_rejected_event(executor_context, str(e)))
118
113
  except Exception as e:
119
114
  if not span_id_emitted:
120
115
  queue.put(f"{SPAN_ID_EVENT}:{uuid4()}")
@@ -168,27 +163,28 @@ def stream_workflow(
168
163
  disable_redirect: bool = True,
169
164
  cancel_signal: Optional[ThreadingEvent] = None,
170
165
  ) -> tuple[Iterator[dict], UUID]:
171
- workflow, namespace = _gather_workflow(executor_context)
172
- workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
173
- display_context = _gather_display_context(workflow, namespace)
174
- workflow_state = (
175
- workflow.deserialize_state(
176
- executor_context.state,
177
- workflow_inputs=workflow_inputs or BaseInputs(),
166
+ cancel_watcher_kill_switch = ThreadingEvent()
167
+ try:
168
+ workflow, namespace = _gather_workflow(executor_context)
169
+ workflow_inputs = _get_workflow_inputs(executor_context, workflow.__class__)
170
+
171
+ display_context = _gather_display_context(workflow, namespace)
172
+ workflow_state = (
173
+ workflow.deserialize_state(
174
+ executor_context.state,
175
+ workflow_inputs=workflow_inputs or BaseInputs(),
176
+ )
177
+ if executor_context.state
178
+ else None
179
+ )
180
+ run_from_node = _get_run_from_node(executor_context, workflow)
181
+ node_output_mocks = MockNodeExecution.validate_all(
182
+ executor_context.node_output_mocks,
183
+ workflow.__class__,
178
184
  )
179
- if executor_context.state
180
- else None
181
- )
182
- run_from_node = _get_run_from_node(executor_context, workflow)
183
- node_output_mocks = MockNodeExecution.validate_all(
184
- executor_context.node_output_mocks,
185
- workflow.__class__,
186
- )
187
185
 
188
- cancel_watcher_kill_switch = ThreadingEvent()
189
- cancel_signal = cancel_signal or ThreadingEvent()
186
+ cancel_signal = cancel_signal or ThreadingEvent()
190
187
 
191
- try:
192
188
  stream = workflow.stream(
193
189
  inputs=workflow_inputs,
194
190
  state=workflow_state,
@@ -198,6 +194,26 @@ def stream_workflow(
198
194
  entrypoint_nodes=[run_from_node] if run_from_node else None,
199
195
  previous_execution_id=executor_context.previous_execution_id,
200
196
  )
197
+ except WorkflowInitializationException as e:
198
+ cancel_watcher_kill_switch.set()
199
+ initialization_exception_stream = stream_initialization_exception(e)
200
+
201
+ def _stream_generator() -> Generator[dict[str, Any], Any, None]:
202
+ for event in initialization_exception_stream:
203
+ yield _dump_event(
204
+ event=event,
205
+ executor_context=executor_context,
206
+ )
207
+
208
+ return (
209
+ _call_stream(
210
+ executor_context=executor_context,
211
+ stream_generator=_stream_generator,
212
+ disable_redirect=disable_redirect,
213
+ timeout_signal=timeout_signal,
214
+ ),
215
+ initialization_exception_stream.span_id,
216
+ )
201
217
  except Exception:
202
218
  cancel_watcher_kill_switch.set()
203
219
  logger.exception("Failed to generate Workflow Stream")
@@ -230,7 +246,6 @@ def stream_workflow(
230
246
  yield _dump_event(
231
247
  event=event,
232
248
  executor_context=executor_context,
233
- client=workflow.context.vellum_client,
234
249
  )
235
250
  except Exception as e:
236
251
  logger.exception("Failed to generate event from Workflow Stream")
@@ -377,18 +392,13 @@ def _create_workflow(executor_context: BaseExecutorContext, namespace: str) -> B
377
392
 
378
393
 
379
394
  def _create_workflow_context(executor_context: BaseExecutorContext) -> WorkflowContext:
380
- vellum_client = create_vellum_client(
381
- api_key=executor_context.environment_api_key,
382
- api_version=executor_context.api_version,
383
- )
384
-
385
395
  if executor_context.environment_variables:
386
396
  os.environ.update(executor_context.environment_variables)
387
397
 
388
398
  namespace = _get_file_namespace(executor_context)
389
399
 
390
400
  return WorkflowContext(
391
- vellum_client=vellum_client,
401
+ vellum_client=executor_context.vellum_client,
392
402
  execution_context=executor_context.execution_context,
393
403
  generated_files=executor_context.files,
394
404
  namespace=namespace,
@@ -408,9 +418,11 @@ def _get_file_namespace(executor_context: BaseExecutorContext) -> str:
408
418
  )
409
419
 
410
420
 
411
- def _dump_event(event: BaseEvent, executor_context: BaseExecutorContext, client: Vellum) -> dict:
421
+ def _dump_event(event: BaseEvent, executor_context: BaseExecutorContext) -> dict:
412
422
  module_base = executor_context.module.split(".")
413
- dump = event.model_dump(mode="json", context={"event_enricher": lambda event: event_enricher(event, client)})
423
+ dump = event.model_dump(
424
+ mode="json", context={"event_enricher": lambda event: event_enricher(event, executor_context.vellum_client)}
425
+ )
414
426
  if dump["name"] in {
415
427
  "workflow.execution.initiated",
416
428
  "workflow.execution.fulfilled",
@@ -1,10 +1,7 @@
1
1
  from datetime import datetime
2
- import os
3
2
  from uuid import uuid4
4
3
  from typing import Optional
5
4
 
6
- from vellum import ApiVersionEnum, Vellum, VellumEnvironment
7
- from workflow_server.config import IS_VPC, VELLUM_API_URL_HOST, VELLUM_API_URL_PORT
8
5
  from workflow_server.core.events import VembdaExecutionFulfilledBody, VembdaExecutionFulfilledEvent
9
6
  from workflow_server.core.workflow_executor_context import BaseExecutorContext
10
7
 
@@ -51,42 +48,3 @@ def is_events_emitting_enabled(executor_context: Optional[BaseExecutorContext])
51
48
  return False
52
49
 
53
50
  return executor_context.feature_flags.get("vembda-event-emitting-enabled") or False
54
-
55
-
56
- def create_vellum_client(
57
- api_key: str,
58
- api_version: Optional[ApiVersionEnum] = None,
59
- ) -> Vellum:
60
- """
61
- Create a VellumClient with proper environment configuration.
62
-
63
- Args:
64
- api_key: The API key for the Vellum client
65
- api_version: Optional API version to use
66
-
67
- Returns:
68
- Configured Vellum client instance
69
-
70
- Note: Ideally we replace this with `vellum.workflows.vellum_client.create_vellum_client`
71
- """
72
- if IS_VPC:
73
- environment = VellumEnvironment(
74
- default=os.getenv("VELLUM_DEFAULT_API_URL", VellumEnvironment.PRODUCTION.default),
75
- documents=os.getenv("VELLUM_DOCUMENTS_API_URL", VellumEnvironment.PRODUCTION.documents),
76
- predict=os.getenv("VELLUM_PREDICT_API_URL", VellumEnvironment.PRODUCTION.predict),
77
- )
78
- elif os.getenv("USE_LOCAL_VELLUM_API") == "true":
79
- VELLUM_API_URL = f"http://{VELLUM_API_URL_HOST}:{VELLUM_API_URL_PORT}"
80
- environment = VellumEnvironment(
81
- default=VELLUM_API_URL,
82
- documents=VELLUM_API_URL,
83
- predict=VELLUM_API_URL,
84
- )
85
- else:
86
- environment = VellumEnvironment.PRODUCTION
87
-
88
- return Vellum(
89
- api_key=api_key,
90
- environment=environment,
91
- api_version=api_version,
92
- )
@@ -0,0 +1,99 @@
1
+ from dataclasses import field
2
+ import os
3
+ from uuid import UUID
4
+ from typing import Any, Optional
5
+
6
+ from _pytest.compat import cached_property
7
+
8
+ from vellum import ApiVersionEnum, Vellum, VellumEnvironment
9
+ from vellum.client.core import UniversalBaseModel
10
+ from vellum.workflows.context import ExecutionContext
11
+ from workflow_server.config import IS_VPC, VELLUM_API_URL_HOST, VELLUM_API_URL_PORT
12
+
13
+ DEFAULT_TIMEOUT_SECONDS = 60 * 30
14
+
15
+
16
+ def create_vellum_client(
17
+ api_key: str,
18
+ api_version: Optional[ApiVersionEnum] = None,
19
+ ) -> Vellum:
20
+ """
21
+ Create a VellumClient with proper environment configuration.
22
+
23
+ Args:
24
+ api_key: The API key for the Vellum client
25
+ api_version: Optional API version to use
26
+
27
+ Returns:
28
+ Configured Vellum client instance
29
+
30
+ Note: Ideally we replace this with `vellum.workflows.vellum_client.create_vellum_client`
31
+ """
32
+ if IS_VPC:
33
+ environment = VellumEnvironment(
34
+ default=os.getenv("VELLUM_DEFAULT_API_URL", VellumEnvironment.PRODUCTION.default),
35
+ documents=os.getenv("VELLUM_DOCUMENTS_API_URL", VellumEnvironment.PRODUCTION.documents),
36
+ predict=os.getenv("VELLUM_PREDICT_API_URL", VellumEnvironment.PRODUCTION.predict),
37
+ )
38
+ elif os.getenv("USE_LOCAL_VELLUM_API") == "true":
39
+ VELLUM_API_URL = f"http://{VELLUM_API_URL_HOST}:{VELLUM_API_URL_PORT}"
40
+ environment = VellumEnvironment(
41
+ default=VELLUM_API_URL,
42
+ documents=VELLUM_API_URL,
43
+ predict=VELLUM_API_URL,
44
+ )
45
+ else:
46
+ environment = VellumEnvironment.PRODUCTION
47
+
48
+ return Vellum(
49
+ api_key=api_key,
50
+ environment=environment,
51
+ api_version=api_version,
52
+ )
53
+
54
+
55
+ class BaseExecutorContext(UniversalBaseModel):
56
+ inputs: dict
57
+ state: Optional[dict] = None
58
+ timeout: int = DEFAULT_TIMEOUT_SECONDS
59
+ files: dict[str, str]
60
+ environment_api_key: str
61
+ api_version: Optional[ApiVersionEnum] = None
62
+ execution_id: UUID
63
+ module: str
64
+ execution_context: ExecutionContext = field(default_factory=ExecutionContext)
65
+ request_start_time: int
66
+ stream_start_time: int = 0
67
+ vembda_public_url: Optional[str] = None
68
+ node_output_mocks: Optional[list[Any]] = None
69
+ environment_variables: Optional[dict[str, str]] = None
70
+ previous_execution_id: Optional[UUID] = None
71
+ feature_flags: Optional[dict[str, bool]] = None
72
+
73
+ @property
74
+ def container_overhead_latency(self) -> int:
75
+ return self.stream_start_time - self.request_start_time if self.stream_start_time else -1
76
+
77
+ @property
78
+ def trace_id(self) -> UUID:
79
+ return self.execution_context.trace_id
80
+
81
+ @cached_property
82
+ def vellum_client(self) -> Vellum:
83
+ return create_vellum_client(
84
+ api_key=self.environment_api_key,
85
+ api_version=self.api_version,
86
+ )
87
+
88
+ def __hash__(self) -> int:
89
+ # do we think we need anything else for a unique hash for caching?
90
+ return hash(str(self.execution_id))
91
+
92
+
93
+ class WorkflowExecutorContext(BaseExecutorContext):
94
+ node_id: Optional[UUID] = None # Sent during run from node UX
95
+
96
+
97
+ class NodeExecutorContext(BaseExecutorContext):
98
+ node_module: str
99
+ node_name: str
@@ -1,49 +0,0 @@
1
- from dataclasses import field
2
- from uuid import UUID
3
- from typing import Any, Optional
4
-
5
- from vellum.client.core import UniversalBaseModel
6
- from vellum.client.types.api_version_enum import ApiVersionEnum
7
- from vellum.workflows.context import ExecutionContext
8
-
9
- DEFAULT_TIMEOUT_SECONDS = 60 * 30
10
-
11
-
12
- class BaseExecutorContext(UniversalBaseModel):
13
- inputs: dict
14
- state: Optional[dict] = None
15
- timeout: int = DEFAULT_TIMEOUT_SECONDS
16
- files: dict[str, str]
17
- environment_api_key: str
18
- api_version: Optional[ApiVersionEnum] = None
19
- execution_id: UUID
20
- module: str
21
- execution_context: ExecutionContext = field(default_factory=ExecutionContext)
22
- request_start_time: int
23
- stream_start_time: int = 0
24
- vembda_public_url: Optional[str] = None
25
- node_output_mocks: Optional[list[Any]] = None
26
- environment_variables: Optional[dict[str, str]] = None
27
- previous_execution_id: Optional[UUID] = None
28
- feature_flags: Optional[dict[str, bool]] = None
29
-
30
- @property
31
- def container_overhead_latency(self) -> int:
32
- return self.stream_start_time - self.request_start_time if self.stream_start_time else -1
33
-
34
- @property
35
- def trace_id(self) -> UUID:
36
- return self.execution_context.trace_id
37
-
38
- def __hash__(self) -> int:
39
- # do we think we need anything else for a unique hash for caching?
40
- return hash(str(self.execution_id))
41
-
42
-
43
- class WorkflowExecutorContext(BaseExecutorContext):
44
- node_id: Optional[UUID] = None # Sent during run from node UX
45
-
46
-
47
- class NodeExecutorContext(BaseExecutorContext):
48
- node_module: str
49
- node_name: str