vellum-ai 0.10.9__py3-none-any.whl → 0.11.0__py3-none-any.whl

vellum/workflows/utils/functions.py

@@ -0,0 +1,74 @@
+import dataclasses
+import inspect
+from typing import Any, Callable, Union, get_args, get_origin
+
+from vellum.client.types.function_definition import FunctionDefinition
+
+type_map = {
+    str: "string",
+    int: "integer",
+    float: "number",
+    bool: "boolean",
+    list: "array",
+    dict: "object",
+    None: "null",
+    type(None): "null",
+}
+
+
+def _compile_annotation(annotation: Any, defs: dict[str, Any]) -> dict:
+    if get_origin(annotation) is Union:
+        return {"anyOf": [_compile_annotation(a, defs) for a in get_args(annotation)]}
+
+    if get_origin(annotation) is dict:
+        _, value_type = get_args(annotation)
+        return {"type": "object", "additionalProperties": _compile_annotation(value_type, defs)}
+
+    if get_origin(annotation) is list:
+        item_type = get_args(annotation)[0]
+        return {"type": "array", "items": _compile_annotation(item_type, defs)}
+
+    if dataclasses.is_dataclass(annotation):
+        if annotation.__name__ not in defs:
+            properties = {}
+            required = []
+            for field in dataclasses.fields(annotation):
+                properties[field.name] = _compile_annotation(field.type, defs)
+                if field.default is dataclasses.MISSING:
+                    required.append(field.name)
+                else:
+                    properties[field.name]["default"] = field.default
+            defs[annotation.__name__] = {"type": "object", "properties": properties, "required": required}
+        return {"$ref": f"#/$defs/{annotation.__name__}"}
+
+    return {"type": type_map[annotation]}
+
+
+def compile_function_definition(function: Callable) -> FunctionDefinition:
+    """
+    Converts a Python function into our Vellum-native FunctionDefinition type.
+    """
+
+    try:
+        signature = inspect.signature(function)
+    except ValueError as e:
+        raise ValueError(f"Failed to get signature for function {function.__name__}: {str(e)}")
+
+    properties = {}
+    required = []
+    defs: dict[str, Any] = {}
+    for param in signature.parameters.values():
+        properties[param.name] = _compile_annotation(param.annotation, defs)
+        if param.default is inspect.Parameter.empty:
+            required.append(param.name)
+        else:
+            properties[param.name]["default"] = param.default
+
+    parameters = {"type": "object", "properties": properties, "required": required}
+    if defs:
+        parameters["$defs"] = defs
+
+    return FunctionDefinition(
+        name=function.__name__,
+        parameters=parameters,
+    )

vellum/workflows/utils/tests/test_functions.py

@@ -0,0 +1,171 @@
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+from vellum.client.types.function_definition import FunctionDefinition
+from vellum.workflows.utils.functions import compile_function_definition
+
+
+def test_compile_function_definition__just_name():
+    # GIVEN a function with just a name
+    def my_function():
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={"type": "object", "properties": {}, "required": []},
+    )
+
+
+def test_compile_function_definition__all_args():
+    # GIVEN a function with args of all base types
+    def my_function(a: str, b: int, c: float, d: bool, e: list, f: dict):
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "string"},
+                "b": {"type": "integer"},
+                "c": {"type": "number"},
+                "d": {"type": "boolean"},
+                "e": {"type": "array"},
+                "f": {"type": "object"},
+            },
+            "required": ["a", "b", "c", "d", "e", "f"],
+        },
+    )
+
+
+def test_compile_function_definition__unions():
+    # GIVEN a function with a union arg
+    def my_function(a: Union[str, int]):
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"anyOf": [{"type": "string"}, {"type": "integer"}]},
+            },
+            "required": ["a"],
+        },
+    )
+
+
+def test_compile_function_definition__optionals():
+    # GIVEN a function with various ways to specify optionals
+    def my_function(
+        a: str,
+        b: Optional[str],
+        c: None,
+        d: str = "hello",
+        e: Optional[str] = None,
+    ):
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "string"},
+                "b": {"anyOf": [{"type": "string"}, {"type": "null"}]},
+                "c": {"type": "null"},
+                "d": {"type": "string", "default": "hello"},
+                "e": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": None},
+            },
+            "required": ["a", "b", "c"],
+        },
+    )
+
+
+def test_compile_function_definition__parameterized_dicts():
+    # GIVEN a function with a parameterized dict
+    def my_function(a: Dict[str, int]):
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "object", "additionalProperties": {"type": "integer"}},
+            },
+            "required": ["a"],
+        },
+    )
+
+
+def test_compile_function_definition__parameterized_lists():
+    # GIVEN a function with a parameterized list
+    def my_function(a: List[int]):
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "array", "items": {"type": "integer"}},
+            },
+            "required": ["a"],
+        },
+    )
+
+
+def test_compile_function_definition__dataclasses():
+    # GIVEN a function with a dataclass
+    @dataclass
+    class MyDataClass:
+        a: int
+        b: str
+
+    def my_function(c: MyDataClass):
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function)
+
+    # THEN it should return the compiled function definition
+    assert compiled_function == FunctionDefinition(
+        name="my_function",
+        parameters={
+            "type": "object",
+            "properties": {"c": {"$ref": "#/$defs/MyDataClass"}},
+            "required": ["c"],
+            "$defs": {
+                "MyDataClass": {
+                    "type": "object",
+                    "properties": {"a": {"type": "integer"}, "b": {"type": "string"}},
+                    "required": ["a", "b"],
+                }
+            },
+        },
+    )

vellum/workflows/utils/tests/test_vellum_variables.py

@@ -2,7 +2,6 @@ import pytest
 from typing import List, Optional
 
 from vellum import ChatMessage, SearchResult
-
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 
 

vellum/workflows/utils/vellum_variables.py

@@ -40,8 +40,8 @@ def primitive_type_to_vellum_variable_type(type_: Union[Type, BaseDescriptor]) -
                 int,
                 float,
                 str,
-                typing.List[typing.ForwardRef('Json')],  # type: ignore [misc]
-                typing.Dict[str, typing.ForwardRef('Json')],  # type: ignore [misc]
+                typing.List[typing.ForwardRef("Json")],  # type: ignore [misc]
+                typing.Dict[str, typing.ForwardRef("Json")],  # type: ignore [misc]
             ]:
                 return "JSON"
             raise ValueError(f"Expected Descriptor to only have one type, got {types}")

vellum/workflows/workflows/base.py

@@ -4,7 +4,6 @@ import importlib
 import inspect
 
 from vellum.plugins.utils import load_runtime_plugins
-from vellum.workflows.events.types import CodeResourceDefinition
 from vellum.workflows.workflows.event_filters import workflow_event_filter
 
 load_runtime_plugins()
@@ -62,6 +61,8 @@ from vellum.workflows.events.workflow import (
     WorkflowExecutionRejectedEvent,
     WorkflowExecutionResumedBody,
     WorkflowExecutionResumedEvent,
+    WorkflowExecutionSnapshottedBody,
+    WorkflowExecutionSnapshottedEvent,
     WorkflowExecutionStreamingBody,
     WorkflowExecutionStreamingEvent,
 )
@@ -82,13 +83,16 @@ from vellum.workflows.types.utils import get_original_base
 class _BaseWorkflowMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         if "graph" not in dct:
-            dct["graph"] = []
+            dct["graph"] = set()
 
         return super().__new__(mcs, name, bases, dct)
 
 
+GraphAttribute = Union[Type[BaseNode], Graph, Set[Type[BaseNode]], Set[Graph]]
+
+
 class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkflowMeta):
-    graph: ClassVar[Union[Type[BaseNode], Graph, Set[Type[BaseNode]], Set[Graph]]]
+    graph: ClassVar[GraphAttribute]
     emitters: List[BaseWorkflowEmitter]
     resolvers: List[BaseWorkflowResolver]
 
@@ -99,6 +103,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
         GenericWorkflowEvent,
         WorkflowExecutionInitiatedEvent[WorkflowInputsType],  # type: ignore[valid-type]
         WorkflowExecutionFulfilledEvent[Outputs],
+        WorkflowExecutionSnapshottedEvent[StateType],  # type: ignore[valid-type]
     ]
 
     TerminalWorkflowEvent = Union[
@@ -118,9 +123,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
     ):
         self._parent_state = parent_state
         self.emitters = emitters or (self.emitters if hasattr(self, "emitters") else [])
-        self.resolvers = resolvers or (
-            self.resolvers if hasattr(self, "resolvers") else []
-        )
+        self.resolvers = resolvers or (self.resolvers if hasattr(self, "resolvers") else [])
         self._context = context or WorkflowContext()
         self._store = Store()
 
@@ -137,8 +140,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             return [original_graph]
         if isinstance(original_graph, set):
             return [
-                subgraph if isinstance(subgraph, Graph) else Graph.from_node(subgraph)
-                for subgraph in original_graph
+                subgraph if isinstance(subgraph, Graph) else Graph.from_node(subgraph) for subgraph in original_graph
             ]
         if issubclass(original_graph, BaseNode):
             return [Graph.from_node(original_graph)]
@@ -186,6 +188,31 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
     ) -> TerminalWorkflowEvent:
+        """
+        Invoke a Workflow, returning the last event emitted, which should be one of:
+        - `WorkflowExecutionFulfilledEvent` if the Workflow Execution was successful
+        - `WorkflowExecutionRejectedEvent` if the Workflow Execution was rejected
+        - `WorkflowExecutionPausedEvent` if the Workflow Execution was paused
+
+        Parameters
+        ----------
+        inputs: Optional[WorkflowInputsType] = None
+            The Inputs instance used to initiate the Workflow Execution.
+
+        state: Optional[StateType] = None
+            The State instance to run the Workflow with. Workflows maintain a global state that can be used to
+            deterministically resume execution from any point.
+
+        entrypoint_nodes: Optional[RunFromNodeArg] = None
+            The entrypoint nodes to run the Workflow with. Useful for resuming execution from a specific node.
+
+        external_inputs: Optional[ExternalInputsArg] = None
+            External inputs to pass to the Workflow. Useful for providing human-in-the-loop behavior to the Workflow.
+
+        cancel_signal: Optional[ThreadingEvent] = None
+            A threading event that can be used to cancel the Workflow Execution.
+        """
+
         events = WorkflowRunner(
             self,
             inputs=inputs,
@@ -195,15 +222,10 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             cancel_signal=cancel_signal,
             parent_context=self._context.parent_context,
         ).stream()
-        first_event: Optional[
-            Union[WorkflowExecutionInitiatedEvent, WorkflowExecutionResumedEvent]
-        ] = None
+        first_event: Optional[Union[WorkflowExecutionInitiatedEvent, WorkflowExecutionResumedEvent]] = None
         last_event = None
         for event in events:
-            if (
-                event.name == "workflow.execution.initiated"
-                or event.name == "workflow.execution.resumed"
-            ):
+            if event.name == "workflow.execution.initiated" or event.name == "workflow.execution.resumed":
                 first_event = event
             last_event = event
 
@@ -254,15 +276,39 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
     def stream(
         self,
-        event_filter: Optional[
-            Callable[[Type["BaseWorkflow"], WorkflowEvent], bool]
-        ] = None,
+        event_filter: Optional[Callable[[Type["BaseWorkflow"], WorkflowEvent], bool]] = None,
         inputs: Optional[WorkflowInputsType] = None,
         state: Optional[StateType] = None,
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
     ) -> WorkflowEventStream:
+        """
+        Invoke a Workflow, yielding events as they are emitted.
+
+        Parameters
+        ----------
+        event_filter: Optional[Callable[[Type["BaseWorkflow"], WorkflowEvent], bool]] = None
+            A filter that can be used to filter events based on the Workflow Class and the event itself. If the method
+            returns `False`, the event will not be yielded.
+
+        inputs: Optional[WorkflowInputsType] = None
+            The Inputs instance used to initiate the Workflow Execution.
+
+        state: Optional[StateType] = None
+            The State instance to run the Workflow with. Workflows maintain a global state that can be used to
+            deterministically resume execution from any point.
+
+        entrypoint_nodes: Optional[RunFromNodeArg] = None
+            The entrypoint nodes to run the Workflow with. Useful for resuming execution from a specific node.
+
+        external_inputs: Optional[ExternalInputsArg] = None
+            External inputs to pass to the Workflow. Useful for providing human-in-the-loop behavior to the Workflow.
+
+        cancel_signal: Optional[ThreadingEvent] = None
+            A threading event that can be used to cancel the Workflow Execution.
+        """
+
         should_yield = event_filter or workflow_event_filter
         for event in WorkflowRunner(
             self,
@@ -299,14 +345,10 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             state_type = BaseState
 
         if not issubclass(inputs_type, BaseInputs):
-            raise ValueError(
-                f"Expected first type to be a subclass of BaseInputs, was: {inputs_type}"
-            )
+            raise ValueError(f"Expected first type to be a subclass of BaseInputs, was: {inputs_type}")
 
         if not issubclass(state_type, BaseState):
-            raise ValueError(
-                f"Expected second type to be a subclass of BaseState, was: {state_type}"
-            )
+            raise ValueError(f"Expected second type to be a subclass of BaseState, was: {state_type}")
 
         return (inputs_type, state_type)
 
@@ -321,9 +363,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
     def get_default_inputs(self) -> WorkflowInputsType:
         return self.get_inputs_class()()
 
-    def get_default_state(
-        self, workflow_inputs: Optional[WorkflowInputsType] = None
-    ) -> StateType:
+    def get_default_state(self, workflow_inputs: Optional[WorkflowInputsType] = None) -> StateType:
         return self.get_state_class()(
             meta=StateMeta(
                 parent=self._parent_state,
@@ -334,10 +374,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
     def get_state_at_node(self, node: Type[BaseNode]) -> StateType:
         event_ts = datetime.min
         for event in self._store.events:
-            if (
-                event.name == "node.execution.initiated"
-                and event.node_definition == node
-            ):
+            if event.name == "node.execution.initiated" and event.node_definition == node:
                 event_ts = event.timestamp
 
         most_recent_state_snapshot: Optional[StateType] = None
@@ -359,9 +396,7 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
             next_state = cast(StateType, snapshot)
             if not most_recent_state_snapshot:
                 most_recent_state_snapshot = next_state
-            elif (
-                next_state.meta.updated_ts >= most_recent_state_snapshot.meta.updated_ts
-            ):
+            elif next_state.meta.updated_ts >= most_recent_state_snapshot.meta.updated_ts:
                 most_recent_state_snapshot = next_state
 
         if not most_recent_state_snapshot:
@@ -376,6 +411,9 @@ class BaseWorkflow(Generic[WorkflowInputsType, StateType], metaclass=_BaseWorkfl
 
         workflows: List[Type[BaseWorkflow]] = []
         for name in dir(module):
+            if name.startswith("__"):
+                continue
+
             attr = getattr(module, name)
             if (
                 inspect.isclass(attr)
@@ -399,6 +437,7 @@ WorkflowExecutionRejectedBody.model_rebuild()
 WorkflowExecutionPausedBody.model_rebuild()
 WorkflowExecutionResumedBody.model_rebuild()
 WorkflowExecutionStreamingBody.model_rebuild()
+WorkflowExecutionSnapshottedBody.model_rebuild()
 
 NodeExecutionInitiatedBody.model_rebuild()
 NodeExecutionFulfilledBody.model_rebuild()
@@ -413,6 +452,7 @@ WorkflowExecutionRejectedEvent.model_rebuild()
 WorkflowExecutionPausedEvent.model_rebuild()
 WorkflowExecutionResumedEvent.model_rebuild()
 WorkflowExecutionStreamingEvent.model_rebuild()
+WorkflowExecutionSnapshottedEvent.model_rebuild()
 
 NodeExecutionInitiatedEvent.model_rebuild()
 NodeExecutionFulfilledEvent.model_rebuild()

vellum/workflows/workflows/event_filters.py

@@ -7,9 +7,7 @@ if TYPE_CHECKING:
     from vellum.workflows.workflows.base import BaseWorkflow
 
 
-def workflow_event_filter(
-    workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent"
-) -> bool:
+def workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent") -> bool:
     """
     Filters for only Workflow events that were emitted by the `workflow_definition` parameter.
     """
@@ -27,9 +25,7 @@ def workflow_event_filter(
     return False
 
 
-def root_workflow_event_filter(
-    workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent"
-) -> bool:
+def root_workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent") -> bool:
     """
     Filters for Workflow and Node events that were emitted by the `workflow_definition` parameter.
     """
@@ -50,12 +46,8 @@ def root_workflow_event_filter(
     if event.parent.type != "WORKFLOW":
         return False
 
-    return event.parent.workflow_definition == CodeResourceDefinition.encode(
-        workflow_definition
-    )
+    return event.parent.workflow_definition == CodeResourceDefinition.encode(workflow_definition)
 
 
-def all_workflow_event_filter(
-    workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent"
-) -> bool:
+def all_workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent") -> bool:
     return True

{vellum_ai-0.10.9.dist-info → vellum_ai-0.11.0.dist-info}/METADATA

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