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

vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.12.11",
+            "X-Fern-SDK-Version": "0.12.13",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/workflows/exceptions.py

@@ -2,7 +2,7 @@ from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
 
 
 class NodeException(Exception):
-    def __init__(self, message: str, code: WorkflowErrorCode):
+    def __init__(self, message: str, code: WorkflowErrorCode = WorkflowErrorCode.INTERNAL_ERROR):
         self.message = message
         self.code = code
         super().__init__(message)

vellum/workflows/runner/runner.py

@@ -4,7 +4,7 @@ import logging
 from queue import Empty, Queue
 from threading import Event as ThreadingEvent, Thread
 from uuid import UUID
-from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, List, Optional, Sequence, Set, Type, Union
 
 from vellum.workflows.constants import UNDEF
 from vellum.workflows.context import execution_context, get_parent_context
@@ -72,6 +72,7 @@ class WorkflowRunner(Generic[StateType]):
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
+        node_output_mocks: Optional[List[BaseOutputs]] = None,
         parent_context: Optional[ParentContext] = None,
     ):
         if state and external_inputs:
@@ -123,6 +124,9 @@ class WorkflowRunner(Generic[StateType]):
 
         self._dependencies: Dict[Type[BaseNode], Set[Type[BaseNode]]] = defaultdict(set)
         self._state_forks: Set[StateType] = {self._initial_state}
+        self._mocks_by_node_outputs_class = (
+            {mock.__class__: mock for mock in node_output_mocks} if node_output_mocks else {}
+        )
 
         self._active_nodes_by_execution_id: Dict[UUID, BaseNode[StateType]] = {}
         self._cancel_signal = cancel_signal
@@ -178,8 +182,11 @@ class WorkflowRunner(Generic[StateType]):
                 node_definition=node.__class__,
                 parent=parent_context,
             )
-            with execution_context(parent_context=updated_parent_context):
-                node_run_response = node.run()
+            if node.Outputs not in self._mocks_by_node_outputs_class:
+                with execution_context(parent_context=updated_parent_context):
+                    node_run_response = node.run()
+            else:
+                node_run_response = self._mocks_by_node_outputs_class[node.Outputs]
             ports = node.Ports()
             if not isinstance(node_run_response, (BaseOutputs, Iterator)):
                 raise NodeException(

vellum/workflows/workflows/base.py

@@ -194,6 +194,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
+        node_output_mocks: Optional[List[BaseOutputs]] = None,
     ) -> TerminalWorkflowEvent:
         """
         Invoke a Workflow, returning the last event emitted, which should be one of:
@@ -218,6 +219,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
         cancel_signal: Optional[ThreadingEvent] = None
             A threading event that can be used to cancel the Workflow Execution.
+
+        node_output_mocks: Optional[List[Outputs]] = None
+            A list of Outputs to mock for Nodes during Workflow Execution.
         """
 
         events = WorkflowRunner(
@@ -227,6 +231,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             entrypoint_nodes=entrypoint_nodes,
             external_inputs=external_inputs,
             cancel_signal=cancel_signal,
+            node_output_mocks=node_output_mocks,
             parent_context=self._context.parent_context,
         ).stream()
         first_event: Optional[Union[WorkflowExecutionInitiatedEvent, WorkflowExecutionResumedEvent]] = None
@@ -290,6 +295,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
+        node_output_mocks: Optional[List[BaseOutputs]] = None,
     ) -> WorkflowEventStream:
         """
         Invoke a Workflow, yielding events as they are emitted.
@@ -315,6 +321,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
         cancel_signal: Optional[ThreadingEvent] = None
             A threading event that can be used to cancel the Workflow Execution.
+
+        node_output_mocks: Optional[List[Outputs]] = None
+            A list of Outputs to mock for Nodes during Workflow Execution.
         """
 
         should_yield = event_filter or workflow_event_filter
@@ -325,6 +334,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             entrypoint_nodes=entrypoint_nodes,
             external_inputs=external_inputs,
             cancel_signal=cancel_signal,
+            node_output_mocks=node_output_mocks,
             parent_context=self.context.parent_context,
         ).stream():
             if should_yield(self.__class__, event):
@@ -416,7 +426,6 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
     def load_from_module(module_path: str) -> Type["BaseWorkflow"]:
         workflow_path = f"{module_path}.workflow"
         module = importlib.import_module(workflow_path)
-
         workflows: List[Type[BaseWorkflow]] = []
         for name in dir(module):
             if name.startswith("__"):
@@ -435,9 +444,40 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             raise ValueError(f"No workflows found in {module_path}")
         elif len(workflows) > 1:
             raise ValueError(f"Multiple workflows found in {module_path}")
+        try:
+            BaseWorkflow.import_node_display(module_path)
+        except ModuleNotFoundError:
+            pass
 
         return workflows[0]
 
+    @staticmethod
+    def import_node_display(module_path):
+        # Import the nodes package
+        nodes_package = importlib.import_module(f"{module_path}.display.nodes")
+        # Use the loader to get the code
+        if hasattr(nodes_package, "__spec__") and nodes_package.__spec__ and nodes_package.__spec__.loader:
+            loader = nodes_package.__spec__.loader
+
+            # Check if the loader has a code attribute
+            if hasattr(loader, "code"):
+                code = loader.code
+
+                # Parse the code to find import statements
+                import_lines = [line.strip() for line in code.splitlines() if line.startswith("from ")]
+
+                # Import each module specified in the code
+                for line in import_lines:
+                    try:
+                        # Extract module name from the import line
+                        module_name = line.split(" ")[1]
+                        full_module_path = f"{module_path}.display.nodes{module_name}"
+                        importlib.import_module(full_module_path)
+                    except Exception:
+                        continue
+        # Also import from workflow.py
+        importlib.import_module(f"{module_path}.display.workflow")
+
 
 WorkflowExecutionInitiatedBody.model_rebuild()
 WorkflowExecutionFulfilledBody.model_rebuild()

{vellum_ai-0.12.11.dist-info → vellum_ai-0.12.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.12.11
+Version: 0.12.13
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-0.12.11.dist-info → vellum_ai-0.12.13.dist-info}/RECORD

@@ -77,7 +77,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=z59nOGe27vMDqsU-ljfULBwC5J4nyrqFunhmo8xnxbU,111521
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=XMPyn7D-uucGK4Ut387xUgRrFjGQOVJfXarC4py5MbY,1869
+vellum/client/core/client_wrapper.py,sha256=vjueLRJ8AxJ_TUVQkuXwAW-kVfBfyveLY3BAgl-dZqE,1869
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 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
 vellum/workflows/events/tests/test_event.py,sha256=izB6Y9U5ROgmHBBpLNUY2navK4-qFp6hdJqJNz6Beek,13350
 vellum/workflows/events/types.py,sha256=cjRE8WL8tYCFradd9NOGl_H0mN3LiWWnA1uHmyT2Q0Q,3412
 vellum/workflows/events/workflow.py,sha256=l5tXes0sg7iWaA1ZUE5dtAqNnGQ8iy6trVbOU9meu7U,5240
-vellum/workflows/exceptions.py,sha256=gXQvDL919cK3pwXc8BhLGKmR-YOskYDLvL1zPSSmLS4,579
+vellum/workflows/exceptions.py,sha256=l-FLGvXywxg6ivolCts71b8pcsYAWoB1cmUR4Jx7N8g,614
 vellum/workflows/expressions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/expressions/accessor.py,sha256=OFvAHAVABr-k7GceIhtzIurV4OuV_yHft7JPRsq87Es,1472
 vellum/workflows/expressions/and_.py,sha256=I7lNqrUM3-m_5hmjjiMhaHhJtKcLj39kEFVWPDOqwfo,916
@@ -1367,7 +1367,7 @@ vellum/workflows/references/workflow_input.py,sha256=86IuhlBz-9cGxeUzizyjdp482aj
 vellum/workflows/resolvers/__init__.py,sha256=eH6hTvZO4IciDaf_cf7aM2vs-DkBDyJPycOQevJxQnI,82
 vellum/workflows/resolvers/base.py,sha256=WHra9LRtlTuB1jmuNqkfVE2JUgB61Cyntn8f0b0WZg4,411
 vellum/workflows/runner/__init__.py,sha256=i1iG5sAhtpdsrlvwgH6B-m49JsINkiWyPWs8vyT-bqM,72
-vellum/workflows/runner/runner.py,sha256=RXnLEmSJFbp0u4vKF7rvD2fscuYfcRYkspIJINnvFAI,27607
+vellum/workflows/runner/runner.py,sha256=DuFvMyejwdIdapQDHrfVtouaifLEZ02WYPARmK3UqdE,28006
 vellum/workflows/sandbox.py,sha256=GVJzVjMuYzOBnSrboB0_6MMRZWBluAyQ2o7syeaeBd0,2235
 vellum/workflows/state/__init__.py,sha256=yUUdR-_Vl7UiixNDYQZ-GEM_kJI9dnOia75TtuNEsnE,60
 vellum/workflows/state/base.py,sha256=jpSzF1OQd3-fqi6dMGlNsQl-7JnJxCdzWIigmX8Wz-I,14425
@@ -1397,10 +1397,10 @@ vellum/workflows/utils/uuids.py,sha256=DFzPv9RCvsKhvdTEIQyfSek2A31D6S_QcmeLPbgrg
 vellum/workflows/utils/vellum_variables.py,sha256=g5xHYB8etfHE32ek19nP6Anf8NyjhmUtOwO2KmQ5xZU,3111
 vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528t75s,683
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
-vellum/workflows/workflows/base.py,sha256=FxK6ZNrQtuqcgZxWXy4_ole1LeeTiZI8VlxavRd3VbA,17119
+vellum/workflows/workflows/base.py,sha256=TLk5NiDB8N_Ytkzj_T8o-E4WpTUH8qJYZcoOk2xR3mo,18986
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
-vellum_ai-0.12.11.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.12.11.dist-info/METADATA,sha256=JJLnOJAxDsQ4U2evGTWRI8vJiJBvs2lR_rQ-Eno1GSg,5161
-vellum_ai-0.12.11.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.12.11.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.12.11.dist-info/RECORD,,
+vellum_ai-0.12.13.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.12.13.dist-info/METADATA,sha256=uXAzy9xHbkmecOGR1GasDk6auS-Hbl12VwdAgqAuo5c,5161
+vellum_ai-0.12.13.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.12.13.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.12.13.dist-info/RECORD,,