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

vellum/workflows/state/base.py

@@ -5,15 +5,15 @@ from datetime import datetime
 from queue import Queue
 from threading import Lock
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Sequence, Set, Tuple, Type, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Type, cast
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler, field_serializer
 from pydantic_core import core_schema
 
 from vellum.core.pydantic_utilities import UniversalBaseModel
-
 from vellum.workflows.constants import UNDEF
+from vellum.workflows.edges.edge import Edge
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
 from vellum.workflows.types.generics import StateType
@@ -71,58 +71,92 @@ def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> An
 
 
 class NodeExecutionCache:
-    _node_execution_ids: Dict[Type["BaseNode"], Stack[UUID]]
+    _node_executions_fulfilled: Dict[Type["BaseNode"], Stack[UUID]]
     _node_executions_initiated: Dict[Type["BaseNode"], Set[UUID]]
-    _dependencies_invoked: Dict[str, Set[str]]
+    _node_executions_queued: Dict[Type["BaseNode"], List[UUID]]
+    _dependencies_invoked: Dict[UUID, Set[Type["BaseNode"]]]
 
     def __init__(
         self,
         dependencies_invoked: Optional[Dict[str, Sequence[str]]] = None,
-        node_execution_ids: Optional[Dict[str, Sequence[str]]] = None,
+        node_executions_fulfilled: Optional[Dict[str, Sequence[str]]] = None,
         node_executions_initiated: Optional[Dict[str, Sequence[str]]] = None,
+        node_executions_queued: Optional[Dict[str, Sequence[str]]] = None,
     ) -> None:
         self._dependencies_invoked = defaultdict(set)
-        self._node_execution_ids = defaultdict(Stack[UUID])
+        self._node_executions_fulfilled = defaultdict(Stack[UUID])
         self._node_executions_initiated = defaultdict(set)
+        self._node_executions_queued = defaultdict(list)
 
-        for node, dependencies in (dependencies_invoked or {}).items():
-            self._dependencies_invoked[node].update(dependencies)
+        for execution_id, dependencies in (dependencies_invoked or {}).items():
+            self._dependencies_invoked[UUID(execution_id)] = {get_class_by_qualname(dep) for dep in dependencies}
 
-        for node, execution_ids in (node_execution_ids or {}).items():
+        for node, execution_ids in (node_executions_fulfilled or {}).items():
             node_class = get_class_by_qualname(node)
-            self._node_execution_ids[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
+            self._node_executions_fulfilled[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
 
         for node, execution_ids in (node_executions_initiated or {}).items():
             node_class = get_class_by_qualname(node)
             self._node_executions_initiated[node_class].update({UUID(execution_id) for execution_id in execution_ids})
 
-    @property
-    def dependencies_invoked(self) -> Dict[str, Set[str]]:
-        return self._dependencies_invoked
+        for node, execution_ids in (node_executions_queued or {}).items():
+            node_class = get_class_by_qualname(node)
+            self._node_executions_queued[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
 
-    def is_node_initiated(self, node: Type["BaseNode"]) -> bool:
-        return node in self._node_executions_initiated and len(self._node_executions_initiated[node]) > 0
+    def _invoke_dependency(
+        self,
+        execution_id: UUID,
+        node: Type["BaseNode"],
+        dependency: Type["BaseNode"],
+        dependencies: Set["Type[BaseNode]"],
+    ) -> None:
+        self._dependencies_invoked[execution_id].add(dependency)
+        if all(dep in self._dependencies_invoked[execution_id] for dep in dependencies):
+            self._node_executions_queued[node].remove(execution_id)
+
+    def queue_node_execution(
+        self, node: Type["BaseNode"], dependencies: Set["Type[BaseNode]"], invoked_by: Optional[Edge] = None
+    ) -> UUID:
+        execution_id = uuid4()
+        if not invoked_by:
+            return execution_id
+
+        source_node = invoked_by.from_port.node_class
+        for queued_node_execution_id in self._node_executions_queued[node]:
+            if source_node not in self._dependencies_invoked[queued_node_execution_id]:
+                self._invoke_dependency(queued_node_execution_id, node, source_node, dependencies)
+                return queued_node_execution_id
+
+        self._node_executions_queued[node].append(execution_id)
+        self._invoke_dependency(execution_id, node, source_node, dependencies)
+        return execution_id
+
+    def is_node_execution_initiated(self, node: Type["BaseNode"], execution_id: UUID) -> bool:
+        return execution_id in self._node_executions_initiated[node]
 
     def initiate_node_execution(self, node: Type["BaseNode"], execution_id: UUID) -> None:
         self._node_executions_initiated[node].add(execution_id)
 
     def fulfill_node_execution(self, node: Type["BaseNode"], execution_id: UUID) -> None:
-        self._node_executions_initiated[node].remove(execution_id)
-        self._node_execution_ids[node].push(execution_id)
+        self._node_executions_fulfilled[node].push(execution_id)
 
     def get_execution_count(self, node: Type["BaseNode"]) -> int:
-        return self._node_execution_ids[node].size()
+        return self._node_executions_fulfilled[node].size()
 
     def dump(self) -> Dict[str, Any]:
         return {
             "dependencies_invoked": {
-                node: list(dependencies) for node, dependencies in self._dependencies_invoked.items()
+                str(execution_id): [str(dep) for dep in dependencies]
+                for execution_id, dependencies in self._dependencies_invoked.items()
             },
             "node_executions_initiated": {
                 str(node): list(execution_ids) for node, execution_ids in self._node_executions_initiated.items()
             },
-            "node_execution_ids": {
-                str(node): execution_ids.dump() for node, execution_ids in self._node_execution_ids.items()
+            "node_executions_fulfilled": {
+                str(node): execution_ids.dump() for node, execution_ids in self._node_executions_fulfilled.items()
+            },
+            "node_executions_queued": {
+                str(node): execution_ids for node, execution_ids in self._node_executions_queued.items()
             },
         }
 
@@ -158,7 +192,6 @@ class StateMeta(UniversalBaseModel):
     node_outputs: Dict[OutputReference, Any] = field(default_factory=dict)
     node_execution_cache: NodeExecutionCache = field(default_factory=NodeExecutionCache)
     parent: Optional["BaseState"] = None
-    is_terminated: Optional[bool] = None
     __snapshot_callback__: Optional[Callable[[], None]] = field(init=False, default=None)
 
     def model_post_init(self, context: Any) -> None:
@@ -252,7 +285,6 @@ class BaseState(metaclass=_BaseStateMeta):
 {values}
     meta:
         id={self.meta.id}
-        is_terminated={self.meta.is_terminated}
         updated_ts={self.meta.updated_ts}
         node_outputs:{' Empty' if not node_outputs else ''}
 {node_outputs}

vellum/workflows/state/context.py

@@ -1,14 +1,24 @@
 from functools import cached_property
-from typing import Optional
+from queue import Queue
+from typing import TYPE_CHECKING, Optional
 
 from vellum import Vellum
-
+from vellum.workflows.events.types import ParentContext
 from vellum.workflows.vellum_client import create_vellum_client
 
+if TYPE_CHECKING:
+    from vellum.workflows.events.workflow import WorkflowEvent
+
 
 class WorkflowContext:
-    def __init__(self, _vellum_client: Optional[Vellum] = None):
+    def __init__(
+        self,
+        _vellum_client: Optional[Vellum] = None,
+        _parent_context: Optional[ParentContext] = None,
+    ):
         self._vellum_client = _vellum_client
+        self._parent_context = _parent_context
+        self._event_queue: Optional[Queue["WorkflowEvent"]] = None
 
     @cached_property
     def vellum_client(self) -> Vellum:
@@ -16,3 +26,16 @@ class WorkflowContext:
             return self._vellum_client
 
         return create_vellum_client()
+
+    @cached_property
+    def parent_context(self) -> Optional[ParentContext]:
+        if self._parent_context:
+            return self._parent_context
+        return None
+
+    def _emit_subworkflow_event(self, event: "WorkflowEvent") -> None:
+        if self._event_queue:
+            self._event_queue.put(event)
+
+    def _register_event_queue(self, event_queue: Queue["WorkflowEvent"]) -> None:
+        self._event_queue = event_queue

vellum/workflows/types/core.py

@@ -83,6 +83,7 @@ EntityInputsInterface = Dict[
 class MergeBehavior(Enum):
     AWAIT_ALL = "AWAIT_ALL"
     AWAIT_ANY = "AWAIT_ANY"
+    AWAIT_ATTRIBUTES = "AWAIT_ATTRIBUTES"
 
 
 class ConditionType(Enum):

vellum/workflows/types/tests/test_utils.py

@@ -33,6 +33,7 @@ class ExampleGenericClass(Generic[T]):
 class ExampleInheritedClass(ExampleClass):
     theta: int
 
+
 @TryNode.wrap()
 class ExampleNode(BaseNode):
     class Outputs(BaseNode.Outputs):

vellum/workflows/types/utils.py

@@ -1,7 +1,6 @@
 from copy import deepcopy
 from datetime import datetime
 import importlib
-import sys
 from typing import (
     Any,
     ClassVar,

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}")