vellum-ai 0.14.3__py3-none-any.whl → 0.14.5__py3-none-any.whl

vellum/workflows/nodes/displayable/code_execution_node/utils.py

@@ -67,6 +67,49 @@ def _clean_for_dict_wrapper(obj):
     return obj
 
 
+def _get_type_name(obj: Any) -> str:
+    if isinstance(obj, type):
+        return obj.__name__
+
+    if get_origin(obj) is Union:
+        children = [_get_type_name(child) for child in get_args(obj)]
+        return " | ".join(children)
+
+    return str(obj)
+
+
+def _cast_to_output_type(result: Any, output_type: Any) -> Any:
+    is_valid_output_type = isinstance(output_type, type)
+    if get_origin(output_type) is Union:
+        allowed_types = get_args(output_type)
+        for allowed_type in allowed_types:
+            try:
+                return _cast_to_output_type(result, allowed_type)
+            except NodeException:
+                continue
+    elif get_origin(output_type) is list:
+        allowed_item_type = get_args(output_type)[0]
+        if isinstance(result, list):
+            return [_cast_to_output_type(item, allowed_item_type) for item in result]
+    elif is_valid_output_type and issubclass(output_type, BaseModel) and not isinstance(result, output_type):
+        try:
+            return output_type.model_validate(result)
+        except ValidationError as e:
+            raise NodeException(
+                code=WorkflowErrorCode.INVALID_OUTPUTS,
+                message=re.sub(r"\s+For further information visit [^\s]+", "", str(e)),
+            ) from e
+    elif is_valid_output_type and isinstance(result, output_type):
+        return result
+
+    output_type_name = _get_type_name(output_type)
+    result_type_name = _get_type_name(type(result))
+    raise NodeException(
+        code=WorkflowErrorCode.INVALID_OUTPUTS,
+        message=f"Expected an output of type '{output_type_name}', but received '{result_type_name}'",
+    )
+
+
 def run_code_inline(
     code: str,
     inputs: EntityInputsInterface,
@@ -112,25 +155,6 @@ __arg__out = main({", ".join(run_args)})
     result = exec_globals["__arg__out"]
 
     if output_type != Any:
-        if get_origin(output_type) is Union:
-            allowed_types = get_args(output_type)
-            if not isinstance(result, allowed_types):
-                raise NodeException(
-                    code=WorkflowErrorCode.INVALID_OUTPUTS,
-                    message=f"Expected output to be in types {allowed_types}, but received '{type(result).__name__}'",
-                )
-        elif issubclass(output_type, BaseModel) and not isinstance(result, output_type):
-            try:
-                result = output_type.model_validate(result)
-            except ValidationError as e:
-                raise NodeException(
-                    code=WorkflowErrorCode.INVALID_OUTPUTS,
-                    message=re.sub(r"\s+For further information visit [^\s]+", "", str(e)),
-                ) from e
-        elif not isinstance(result, output_type):
-            raise NodeException(
-                code=WorkflowErrorCode.INVALID_OUTPUTS,
-                message=f"Expected an output of type '{output_type.__name__}', but received '{type(result).__name__}'",
-            )
+        result = _cast_to_output_type(result, output_type)
 
     return logs, result

vellum/workflows/nodes/displayable/prompt_deployment_node/node.py

@@ -1,3 +1,4 @@
+import json
 from typing import Iterator
 
 from vellum.workflows.errors import WorkflowErrorCode
@@ -44,13 +45,19 @@ class PromptDeploymentNode(BasePromptDeploymentNode[StateType]):
                 code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
-        string_output = next((output for output in outputs if output.type == "STRING"), None)
-        if not string_output or string_output.value is None:
-            output_types = {output.type for output in outputs}
-            is_plural = len(output_types) > 1
-            raise NodeException(
-                message=f"Expected to receive a non-null string output from Prompt. Only found outputs of type{'s' if is_plural else ''}: {', '.join(output_types)}",  # noqa: E501
-                code=WorkflowErrorCode.INTERNAL_ERROR,
-            )
-
-        yield BaseOutput(name="text", value=string_output.value)
+        string_outputs = []
+        for output in outputs:
+            if output.value is None:
+                continue
+
+            if output.type == "STRING":
+                string_outputs.append(output.value)
+            elif output.type == "JSON":
+                string_outputs.append(json.dumps(output.value, indent=4))
+            elif output.type == "FUNCTION_CALL":
+                string_outputs.append(output.value.model_dump_json(indent=4))
+            else:
+                string_outputs.append(output.value.message)
+
+        value = "\n".join(string_outputs)
+        yield BaseOutput(name="text", value=value)

vellum/workflows/nodes/displayable/tests/test_inline_text_prompt_node.py

@@ -91,6 +91,7 @@ def test_inline_text_prompt_node__basic(vellum_adhoc_prompt_client):
             custom_parameters=None,
         ),
         request_options=mock.ANY,
+        settings=None,
     )
 
 

vellum/workflows/tests/test_undefined.py

@@ -0,0 +1,12 @@
+import pytest
+
+from vellum.workflows.constants import undefined
+
+
+def test_undefined__ensure_sensible_error_messages():
+    # WHEN we invoke an invalid operation on `undefined`
+    with pytest.raises(Exception) as e:
+        len(undefined)
+
+    # THEN we get a sensible error message
+    assert str(e.value) == "object of type 'undefined' has no len()"

vellum/workflows/workflows/base.py

@@ -100,6 +100,30 @@ class _BaseWorkflowMeta(type):
                     new_dct,
                 )
 
+        def collect_nodes(graph_item: Union[GraphAttribute, Set[GraphAttribute]]) -> Set[Type[BaseNode]]:
+            nodes: Set[Type[BaseNode]] = set()
+            if isinstance(graph_item, Graph):
+                nodes.update(node for node in graph_item.nodes)
+            elif isinstance(graph_item, set):
+                for item in graph_item:
+                    if isinstance(item, Graph):
+                        nodes.update(node for node in item.nodes)
+                    elif inspect.isclass(item) and issubclass(item, BaseNode):
+                        nodes.add(item)
+            elif issubclass(graph_item, BaseNode):
+                nodes.add(graph_item)
+            else:
+                raise ValueError(f"Unexpected graph type: {graph_item.__class__}")
+            return nodes
+
+        graph_nodes = collect_nodes(dct.get("graph", set()))
+        unused_nodes = collect_nodes(dct.get("unused_graphs", set()))
+
+        overlap = graph_nodes & unused_nodes
+        if overlap:
+            node_names = [node.__name__ for node in overlap]
+            raise ValueError(f"Node(s) {', '.join(node_names)} cannot appear in both graph and unused_graphs")
+
         cls = super().__new__(mcs, name, bases, dct)
         workflow_class = cast(Type["BaseWorkflow"], cls)
         workflow_class.__id__ = uuid4_from_hash(workflow_class.__qualname__)
@@ -112,6 +136,7 @@ GraphAttribute = Union[Type[BaseNode], Graph, Set[Type[BaseNode]], Set[Graph]]
 class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
     __id__: UUID = uuid4_from_hash(__qualname__)
     graph: ClassVar[GraphAttribute]
+    unused_graphs: ClassVar[Set[GraphAttribute]]  # nodes or graphs that are defined but not used in the graph
     emitters: List[BaseWorkflowEmitter]
     resolvers: List[BaseWorkflowResolver]
 
@@ -196,6 +221,57 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
                     nodes.add(node)
                     yield node
 
+    @classmethod
+    def get_unused_subgraphs(cls) -> List[Graph]:
+        """
+        Returns a list of subgraphs that are defined but not used in the graph
+        """
+        if not hasattr(cls, "unused_graphs"):
+            return []
+
+        graphs = []
+        for item in cls.unused_graphs:
+            if isinstance(item, Graph):
+                graphs.append(item)
+            elif isinstance(item, set):
+                for subitem in item:
+                    if isinstance(subitem, Graph):
+                        graphs.append(subitem)
+                    elif issubclass(subitem, BaseNode):
+                        graphs.append(Graph.from_node(subitem))
+            elif issubclass(item, BaseNode):
+                graphs.append(Graph.from_node(item))
+        return graphs
+
+    @classmethod
+    def get_unused_nodes(cls) -> Iterator[Type[BaseNode]]:
+        """
+        Returns an iterator over the nodes that are defined but not used in the graph.
+        """
+        if not hasattr(cls, "unused_graphs"):
+            yield from ()
+        else:
+            nodes = set()
+            subgraphs = cls.get_unused_subgraphs()
+            for subgraph in subgraphs:
+                for node in subgraph.nodes:
+                    if node not in nodes:
+                        nodes.add(node)
+                        yield node
+
+    @classmethod
+    def get_unused_edges(cls) -> Iterator[Edge]:
+        """
+        Returns an iterator over edges that are defined but not used in the graph.
+        """
+        edges = set()
+        subgraphs = cls.get_unused_subgraphs()
+        for subgraph in subgraphs:
+            for edge in subgraph.edges:
+                if edge not in edges:
+                    edges.add(edge)
+                    yield edge
+
     @classmethod
     def get_entrypoints(cls) -> Iterable[Type[BaseNode]]:
         return iter({e for g in cls.get_subgraphs() for e in g.entrypoints})

vellum/workflows/workflows/tests/test_base_workflow.py

@@ -1,3 +1,6 @@
+import pytest
+
+from vellum.workflows.edges.edge import Edge
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.inline_subworkflow_node.node import InlineSubworkflowNode
@@ -78,3 +81,135 @@ def test_subworkflow__inherit_base_outputs():
     # TEST that the outputs are correct
     assert terminal_event.name == "workflow.execution.fulfilled", terminal_event
     assert terminal_event.outputs == {"output": "bar"}
+
+
+def test_workflow__nodes_not_in_graph():
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    class NodeC(BaseNode):
+        pass
+
+    # WHEN we create a workflow with multiple unused nodes
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = NodeA
+        unused_graphs = {NodeB, NodeC}
+
+    # TEST that all nodes from unused_graphs are collected
+    unused_graphs = set(TestWorkflow.get_unused_nodes())
+    assert unused_graphs == {NodeB, NodeC}
+
+
+def test_workflow__unused_graphs():
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    class NodeC(BaseNode):
+        pass
+
+    class NodeD(BaseNode):
+        pass
+
+    class NodeE(BaseNode):
+        pass
+
+    class NodeF(BaseNode):
+        pass
+
+    # WHEN we create a workflow with unused nodes in a graph
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = NodeA
+        unused_graphs = {NodeB >> {NodeC >> NodeD}, NodeE, NodeF}
+
+    # TEST that all nodes from unused_graphs are collected
+    unused_graphs = set(TestWorkflow.get_unused_nodes())
+    assert unused_graphs == {NodeB, NodeC, NodeD, NodeE, NodeF}
+
+
+def test_workflow__no_unused_nodes():
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    # WHEN we create a workflow with no unused nodes
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = NodeA >> NodeB
+
+    # TEST that nodes not in the graph are empty
+    nodes = set(TestWorkflow.get_unused_nodes())
+    assert nodes == set()
+
+
+def test_workflow__node_in_both_graph_and_unused():
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    class NodeC(BaseNode):
+        pass
+
+    # WHEN we try to create a workflow where NodeA appears in both graph and unused
+    with pytest.raises(ValueError) as exc_info:
+
+        class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+            graph = NodeA >> NodeB
+            unused_graphs = {NodeA >> NodeC}
+
+    # THEN it should raise an error
+    assert "Node(s) NodeA cannot appear in both graph and unused_graphs" in str(exc_info.value)
+
+
+def test_workflow__get_unused_edges():
+    """
+    Test that get_unused_edges correctly identifies edges that are defined but not used in the workflow graph.
+    """
+
+    class NodeA(BaseNode):
+        pass
+
+    class NodeB(BaseNode):
+        pass
+
+    class NodeC(BaseNode):
+        pass
+
+    class NodeD(BaseNode):
+        pass
+
+    class NodeE(BaseNode):
+        pass
+
+    class NodeF(BaseNode):
+        pass
+
+    class NodeG(BaseNode):
+        pass
+
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = NodeA >> NodeB
+        unused_graphs = {NodeC >> {NodeD >> NodeE, NodeF} >> NodeG}
+
+    edge_c_to_d = Edge(from_port=NodeC.Ports.default, to_node=NodeD)
+    edge_c_to_f = Edge(from_port=NodeC.Ports.default, to_node=NodeF)
+    edge_d_to_e = Edge(from_port=NodeD.Ports.default, to_node=NodeE)
+    edge_e_to_g = Edge(from_port=NodeE.Ports.default, to_node=NodeG)
+    edge_f_to_g = Edge(from_port=NodeF.Ports.default, to_node=NodeG)
+
+    # Collect unused edges
+    unused_edges = set(TestWorkflow.get_unused_edges())
+
+    # Expected unused edges
+    expected_unused_edges = {edge_c_to_d, edge_c_to_f, edge_d_to_e, edge_e_to_g, edge_f_to_g}
+
+    # TEST that unused edges are correctly identified
+    assert unused_edges == expected_unused_edges, f"Expected {expected_unused_edges}, but got {unused_edges}"

vellum/workflows/workflows/tests/test_context.py

@@ -0,0 +1,60 @@
+from uuid import UUID, uuid4
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.context import execution_context, get_execution_context
+from vellum.workflows.events.types import NodeParentContext, WorkflowParentContext
+from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.nodes import BaseNode
+from vellum.workflows.references import VellumSecretReference
+from vellum.workflows.state import BaseState
+
+
+class MockInputs(BaseInputs):
+    foo: str
+
+
+class MockNode(BaseNode):
+    node_foo = MockInputs.foo
+    node_secret = VellumSecretReference("secret")
+
+    class Outputs(BaseNode.Outputs):
+        example: str
+
+
+class MockWorkflow(BaseWorkflow[MockInputs, BaseState]):
+    graph = MockNode
+
+
+def test_context_trace_and_parent():
+    trace_id = uuid4()
+    parent_context = NodeParentContext(
+        node_definition=MockNode,
+        span_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+        parent=WorkflowParentContext(
+            workflow_definition=MockWorkflow,
+            span_id=UUID("123e4567-e89b-12d3-a456-426614174000"),
+        ),
+    )
+    second_parent_context = WorkflowParentContext(
+        workflow_definition=MockWorkflow, span_id=uuid4(), parent=parent_context
+    )
+    # When using execution context , if we set trace id within
+    with execution_context(parent_context=parent_context, trace_id=trace_id):
+        test = get_execution_context()
+        assert test.trace_id == trace_id
+        assert test.parent_context == parent_context
+        with execution_context(parent_context=second_parent_context):
+            test1 = get_execution_context()
+            assert test1.trace_id == trace_id
+            assert test1.parent_context == second_parent_context
+            # then we can assume trace id will not change
+            with execution_context(trace_id=uuid4()):
+                test3 = get_execution_context()
+                assert test3.trace_id == trace_id
+            with execution_context(parent_context=parent_context, trace_id=uuid4()):
+                test3 = get_execution_context()
+                assert test3.trace_id == trace_id
+    # and if we have a new context, the trace will differ
+    with execution_context(parent_context=parent_context, trace_id=uuid4()):
+        test = get_execution_context()
+        assert test.trace_id != trace_id

{vellum_ai-0.14.3.dist-info → vellum_ai-0.14.5.dist-info}/METADATA

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