vellum-ai 0.12.11__py3-none-any.whl → 0.12.13__py3-none-any.whl

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.
@@ -18,7 +18,7 @@ class BaseClientWrapper:
18
18
  headers: typing.Dict[str, str] = {
19
19
  "X-Fern-Language": "Python",
20
20
  "X-Fern-SDK-Name": "vellum-ai",
21
- "X-Fern-SDK-Version": "0.12.11",
21
+ "X-Fern-SDK-Version": "0.12.13",
22
22
  }
23
23
  headers["X_API_KEY"] = self.api_key
24
24
  return headers
@@ -2,7 +2,7 @@ from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
2
2
 
3
3
 
4
4
  class NodeException(Exception):
5
- def __init__(self, message: str, code: WorkflowErrorCode):
5
+ def __init__(self, message: str, code: WorkflowErrorCode = WorkflowErrorCode.INTERNAL_ERROR):
6
6
  self.message = message
7
7
  self.code = code
8
8
  super().__init__(message)
@@ -4,7 +4,7 @@ import logging
4
4
  from queue import Empty, Queue
5
5
  from threading import Event as ThreadingEvent, Thread
6
6
  from uuid import UUID
7
- from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Type, Union
7
+ from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, List, Optional, Sequence, Set, Type, Union
8
8
 
9
9
  from vellum.workflows.constants import UNDEF
10
10
  from vellum.workflows.context import execution_context, get_parent_context
@@ -72,6 +72,7 @@ class WorkflowRunner(Generic[StateType]):
72
72
  entrypoint_nodes: Optional[RunFromNodeArg] = None,
73
73
  external_inputs: Optional[ExternalInputsArg] = None,
74
74
  cancel_signal: Optional[ThreadingEvent] = None,
75
+ node_output_mocks: Optional[List[BaseOutputs]] = None,
75
76
  parent_context: Optional[ParentContext] = None,
76
77
  ):
77
78
  if state and external_inputs:
@@ -123,6 +124,9 @@ class WorkflowRunner(Generic[StateType]):
123
124
 
124
125
  self._dependencies: Dict[Type[BaseNode], Set[Type[BaseNode]]] = defaultdict(set)
125
126
  self._state_forks: Set[StateType] = {self._initial_state}
127
+ self._mocks_by_node_outputs_class = (
128
+ {mock.__class__: mock for mock in node_output_mocks} if node_output_mocks else {}
129
+ )
126
130
 
127
131
  self._active_nodes_by_execution_id: Dict[UUID, BaseNode[StateType]] = {}
128
132
  self._cancel_signal = cancel_signal
@@ -178,8 +182,11 @@ class WorkflowRunner(Generic[StateType]):
178
182
  node_definition=node.__class__,
179
183
  parent=parent_context,
180
184
  )
181
- with execution_context(parent_context=updated_parent_context):
182
- node_run_response = node.run()
185
+ if node.Outputs not in self._mocks_by_node_outputs_class:
186
+ with execution_context(parent_context=updated_parent_context):
187
+ node_run_response = node.run()
188
+ else:
189
+ node_run_response = self._mocks_by_node_outputs_class[node.Outputs]
183
190
  ports = node.Ports()
184
191
  if not isinstance(node_run_response, (BaseOutputs, Iterator)):
185
192
  raise NodeException(
@@ -194,6 +194,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
194
194
  entrypoint_nodes: Optional[RunFromNodeArg] = None,
195
195
  external_inputs: Optional[ExternalInputsArg] = None,
196
196
  cancel_signal: Optional[ThreadingEvent] = None,
197
+ node_output_mocks: Optional[List[BaseOutputs]] = None,
197
198
  ) -> TerminalWorkflowEvent:
198
199
  """
199
200
  Invoke a Workflow, returning the last event emitted, which should be one of:
@@ -218,6 +219,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
218
219
 
219
220
  cancel_signal: Optional[ThreadingEvent] = None
220
221
  A threading event that can be used to cancel the Workflow Execution.
222
+
223
+ node_output_mocks: Optional[List[Outputs]] = None
224
+ A list of Outputs to mock for Nodes during Workflow Execution.
221
225
  """
222
226
 
223
227
  events = WorkflowRunner(
@@ -227,6 +231,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
227
231
  entrypoint_nodes=entrypoint_nodes,
228
232
  external_inputs=external_inputs,
229
233
  cancel_signal=cancel_signal,
234
+ node_output_mocks=node_output_mocks,
230
235
  parent_context=self._context.parent_context,
231
236
  ).stream()
232
237
  first_event: Optional[Union[WorkflowExecutionInitiatedEvent, WorkflowExecutionResumedEvent]] = None
@@ -290,6 +295,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
290
295
  entrypoint_nodes: Optional[RunFromNodeArg] = None,
291
296
  external_inputs: Optional[ExternalInputsArg] = None,
292
297
  cancel_signal: Optional[ThreadingEvent] = None,
298
+ node_output_mocks: Optional[List[BaseOutputs]] = None,
293
299
  ) -> WorkflowEventStream:
294
300
  """
295
301
  Invoke a Workflow, yielding events as they are emitted.
@@ -315,6 +321,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
315
321
 
316
322
  cancel_signal: Optional[ThreadingEvent] = None
317
323
  A threading event that can be used to cancel the Workflow Execution.
324
+
325
+ node_output_mocks: Optional[List[Outputs]] = None
326
+ A list of Outputs to mock for Nodes during Workflow Execution.
318
327
  """
319
328
 
320
329
  should_yield = event_filter or workflow_event_filter
@@ -325,6 +334,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
325
334
  entrypoint_nodes=entrypoint_nodes,
326
335
  external_inputs=external_inputs,
327
336
  cancel_signal=cancel_signal,
337
+ node_output_mocks=node_output_mocks,
328
338
  parent_context=self.context.parent_context,
329
339
  ).stream():
330
340
  if should_yield(self.__class__, event):
@@ -416,7 +426,6 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
416
426
  def load_from_module(module_path: str) -> Type["BaseWorkflow"]:
417
427
  workflow_path = f"{module_path}.workflow"
418
428
  module = importlib.import_module(workflow_path)
419
-
420
429
  workflows: List[Type[BaseWorkflow]] = []
421
430
  for name in dir(module):
422
431
  if name.startswith("__"):
@@ -435,9 +444,40 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
435
444
  raise ValueError(f"No workflows found in {module_path}")
436
445
  elif len(workflows) > 1:
437
446
  raise ValueError(f"Multiple workflows found in {module_path}")
447
+ try:
448
+ BaseWorkflow.import_node_display(module_path)
449
+ except ModuleNotFoundError:
450
+ pass
438
451
 
439
452
  return workflows[0]
440
453
 
454
+ @staticmethod
455
+ def import_node_display(module_path):
456
+ # Import the nodes package
457
+ nodes_package = importlib.import_module(f"{module_path}.display.nodes")
458
+ # Use the loader to get the code
459
+ if hasattr(nodes_package, "__spec__") and nodes_package.__spec__ and nodes_package.__spec__.loader:
460
+ loader = nodes_package.__spec__.loader
461
+
462
+ # Check if the loader has a code attribute
463
+ if hasattr(loader, "code"):
464
+ code = loader.code
465
+
466
+ # Parse the code to find import statements
467
+ import_lines = [line.strip() for line in code.splitlines() if line.startswith("from ")]
468
+
469
+ # Import each module specified in the code
470
+ for line in import_lines:
471
+ try:
472
+ # Extract module name from the import line
473
+ module_name = line.split(" ")[1]
474
+ full_module_path = f"{module_path}.display.nodes{module_name}"
475
+ importlib.import_module(full_module_path)
476
+ except Exception:
477
+ continue
478
+ # Also import from workflow.py
479
+ importlib.import_module(f"{module_path}.display.workflow")
480
+
441
481
 
442
482
  WorkflowExecutionInitiatedBody.model_rebuild()
443
483
  WorkflowExecutionFulfilledBody.model_rebuild()
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: vellum-ai
3
- Version: 0.12.11
3
+ Version: 0.12.13
4
4
  Summary:
5
5
  License: MIT
6
6
  Requires-Python: >=3.9,<4.0
@@ -77,7 +77,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
77
77
  vellum/client/__init__.py,sha256=z59nOGe27vMDqsU-ljfULBwC5J4nyrqFunhmo8xnxbU,111521
78
78
  vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
79
79
  vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
80
- vellum/client/core/client_wrapper.py,sha256=XMPyn7D-uucGK4Ut387xUgRrFjGQOVJfXarC4py5MbY,1869
80
+ vellum/client/core/client_wrapper.py,sha256=vjueLRJ8AxJ_TUVQkuXwAW-kVfBfyveLY3BAgl-dZqE,1869
81
81
  vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
82
82
  vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
83
83
  vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -1234,7 +1234,7 @@ vellum/workflows/events/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5N
1234
1234
  vellum/workflows/events/tests/test_event.py,sha256=izB6Y9U5ROgmHBBpLNUY2navK4-qFp6hdJqJNz6Beek,13350
1235
1235
  vellum/workflows/events/types.py,sha256=cjRE8WL8tYCFradd9NOGl_H0mN3LiWWnA1uHmyT2Q0Q,3412
1236
1236
  vellum/workflows/events/workflow.py,sha256=l5tXes0sg7iWaA1ZUE5dtAqNnGQ8iy6trVbOU9meu7U,5240
1237
- vellum/workflows/exceptions.py,sha256=gXQvDL919cK3pwXc8BhLGKmR-YOskYDLvL1zPSSmLS4,579
1237
+ vellum/workflows/exceptions.py,sha256=l-FLGvXywxg6ivolCts71b8pcsYAWoB1cmUR4Jx7N8g,614
1238
1238
  vellum/workflows/expressions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
1239
1239
  vellum/workflows/expressions/accessor.py,sha256=OFvAHAVABr-k7GceIhtzIurV4OuV_yHft7JPRsq87Es,1472
1240
1240
  vellum/workflows/expressions/and_.py,sha256=I7lNqrUM3-m_5hmjjiMhaHhJtKcLj39kEFVWPDOqwfo,916
@@ -1367,7 +1367,7 @@ vellum/workflows/references/workflow_input.py,sha256=86IuhlBz-9cGxeUzizyjdp482aj
1367
1367
  vellum/workflows/resolvers/__init__.py,sha256=eH6hTvZO4IciDaf_cf7aM2vs-DkBDyJPycOQevJxQnI,82
1368
1368
  vellum/workflows/resolvers/base.py,sha256=WHra9LRtlTuB1jmuNqkfVE2JUgB61Cyntn8f0b0WZg4,411
1369
1369
  vellum/workflows/runner/__init__.py,sha256=i1iG5sAhtpdsrlvwgH6B-m49JsINkiWyPWs8vyT-bqM,72
1370
- vellum/workflows/runner/runner.py,sha256=RXnLEmSJFbp0u4vKF7rvD2fscuYfcRYkspIJINnvFAI,27607
1370
+ vellum/workflows/runner/runner.py,sha256=DuFvMyejwdIdapQDHrfVtouaifLEZ02WYPARmK3UqdE,28006
1371
1371
  vellum/workflows/sandbox.py,sha256=GVJzVjMuYzOBnSrboB0_6MMRZWBluAyQ2o7syeaeBd0,2235
1372
1372
  vellum/workflows/state/__init__.py,sha256=yUUdR-_Vl7UiixNDYQZ-GEM_kJI9dnOia75TtuNEsnE,60
1373
1373
  vellum/workflows/state/base.py,sha256=jpSzF1OQd3-fqi6dMGlNsQl-7JnJxCdzWIigmX8Wz-I,14425
@@ -1397,10 +1397,10 @@ vellum/workflows/utils/uuids.py,sha256=DFzPv9RCvsKhvdTEIQyfSek2A31D6S_QcmeLPbgrg
1397
1397
  vellum/workflows/utils/vellum_variables.py,sha256=g5xHYB8etfHE32ek19nP6Anf8NyjhmUtOwO2KmQ5xZU,3111
1398
1398
  vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528t75s,683
1399
1399
  vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
1400
- vellum/workflows/workflows/base.py,sha256=FxK6ZNrQtuqcgZxWXy4_ole1LeeTiZI8VlxavRd3VbA,17119
1400
+ vellum/workflows/workflows/base.py,sha256=TLk5NiDB8N_Ytkzj_T8o-E4WpTUH8qJYZcoOk2xR3mo,18986
1401
1401
  vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
1402
- vellum_ai-0.12.11.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
1403
- vellum_ai-0.12.11.dist-info/METADATA,sha256=JJLnOJAxDsQ4U2evGTWRI8vJiJBvs2lR_rQ-Eno1GSg,5161
1404
- vellum_ai-0.12.11.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1405
- vellum_ai-0.12.11.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
1406
- vellum_ai-0.12.11.dist-info/RECORD,,
1402
+ vellum_ai-0.12.13.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
1403
+ vellum_ai-0.12.13.dist-info/METADATA,sha256=uXAzy9xHbkmecOGR1GasDk6auS-Hbl12VwdAgqAuo5c,5161
1404
+ vellum_ai-0.12.13.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1405
+ vellum_ai-0.12.13.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
1406
+ vellum_ai-0.12.13.dist-info/RECORD,,