vellum-ai 0.10.4__tar.gz → 0.10.7__tar.gz

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.10.4
+Version: 0.10.7
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -63,7 +63,6 @@ Description-Content-Type: text/markdown
 
 # Introduction
 
-
 [Vellum](https://www.vellum.ai/) is the end-to-end development platform for building production-grade AI applications
 
 ### Core Features
@@ -85,22 +84,21 @@ Description-Content-Type: text/markdown
 - [Contributing](#contributing)
 - [Open-source vs paid](#open-source-vs-paid)
 
-
 ## Get Started
+
 Most functionality within the SDKs here requires a Vellum account and API key. To sign up, [talk to us](https://www.vellum.ai/landing-pages/request-demo)
 or visit our [pricing page](https://www.vellum.ai/pricing).
 
 Even without a Vellum account, you can use the Workflows SDK to define the control flow of your AI systems. [Learn
 more below](#workflows-sdk).
 
-
-
 ## Client SDK
+
 The Vellum Client SDK, found within `src/client` is a low-level client used to interact directly with the Vellum API.
 Learn more and get started by visiting the [Vellum Client SDK README](/src/vellum/client/README.md).
 
-
 ## Workflows SDK
+
 The Vellum Workflows SDK is a high-level framework for defining and debugging the control flow of AI systems. At
 it's core, it's a powerful workflow engine with syntactic sugar for intuitively defining graphs, the nodes within,
 and the relationships between them.
@@ -111,6 +109,9 @@ and debugging via a UI.
 
 To learn more and get started, visit the [Vellum Workflows SDK README](/src/vellum/workflows/README.md).
 
+## Contributing
+
+See the [CONTRIBUTING.md](/CONTRIBUTING.md) for information on how to contribute to the Vellum SDKs.
 
 ## Open-Source vs. Paid
 
@@ -118,3 +119,4 @@ This repo is available under the [MIT expat license](https://github.com/vellum-a
 for the `ee` directory (which has its [license here](https://github.com/vellum-ai/vellum-python-sdks/blob/main/ee/LICENSE)) if applicable.
 
 To learn more, [book a demo](https://www.vellum.ai/landing-pages/request-demo) or see our [pricing page](https://www.vellum.ai/pricing).
+

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/README.md

@@ -23,7 +23,6 @@
 
 # Introduction
 
-
 [Vellum](https://www.vellum.ai/) is the end-to-end development platform for building production-grade AI applications
 
 ### Core Features
@@ -45,22 +44,21 @@
 - [Contributing](#contributing)
 - [Open-source vs paid](#open-source-vs-paid)
 
-
 ## Get Started
+
 Most functionality within the SDKs here requires a Vellum account and API key. To sign up, [talk to us](https://www.vellum.ai/landing-pages/request-demo)
 or visit our [pricing page](https://www.vellum.ai/pricing).
 
 Even without a Vellum account, you can use the Workflows SDK to define the control flow of your AI systems. [Learn
 more below](#workflows-sdk).
 
-
-
 ## Client SDK
+
 The Vellum Client SDK, found within `src/client` is a low-level client used to interact directly with the Vellum API.
 Learn more and get started by visiting the [Vellum Client SDK README](/src/vellum/client/README.md).
 
-
 ## Workflows SDK
+
 The Vellum Workflows SDK is a high-level framework for defining and debugging the control flow of AI systems. At
 it's core, it's a powerful workflow engine with syntactic sugar for intuitively defining graphs, the nodes within,
 and the relationships between them.
@@ -71,10 +69,13 @@ and debugging via a UI.
 
 To learn more and get started, visit the [Vellum Workflows SDK README](/src/vellum/workflows/README.md).
 
+## Contributing
+
+See the [CONTRIBUTING.md](/CONTRIBUTING.md) for information on how to contribute to the Vellum SDKs.
 
 ## Open-Source vs. Paid
 
 This repo is available under the [MIT expat license](https://github.com/vellum-ai/vellum-python-sdks/blob/main/LICENSE), except
 for the `ee` directory (which has its [license here](https://github.com/vellum-ai/vellum-python-sdks/blob/main/ee/LICENSE)) if applicable.
 
-To learn more, [book a demo](https://www.vellum.ai/landing-pages/request-demo) or see our [pricing page](https://www.vellum.ai/pricing).
+To learn more, [book a demo](https://www.vellum.ai/landing-pages/request-demo) or see our [pricing page](https://www.vellum.ai/pricing).

vellum_ai-0.10.7/ee/vellum_cli/CONTRIBUTING.md

@@ -0,0 +1,66 @@
+<!-- START doctoc generated TOC please keep comment here to allow auto update -->
+<!-- DON'T EDIT THIS SECTION, INSTEAD RE-RUN doctoc TO UPDATE -->
+
+- [Vellum CLI](#vellum-cli)
+- [Updating this README](#updating-this-readme)
+- [Development](#development)
+  - [vellum pull](#vellum-pull)
+    - [Step 1. Run the Codegen Service Locally](#step-1-run-the-codegen-service-locally)
+    - [Step 2. Run the Vellum Webserver Locally](#step-2-run-the-vellum-webserver-locally)
+    - [Step 3. Run the Vellum CLI Locally](#step-3-run-the-vellum-cli-locally)
+
+<!-- END doctoc generated TOC please keep comment here to allow auto update -->
+
+# Vellum CLI
+
+This contribution guide is for the Vellum CLI, which is the command-line interface connecting Vellum SDKs to the Vellum application.
+
+**Note:** This directory is _not_ accepting contributions from anyone besides Vellum employees. Please open an issue for feedback or feature requests.
+
+# Updating this README
+
+This README is generated using [doctoc](https://github.com/thlorenz/doctoc)
+
+To install doctoc, run `npm install -g doctoc`
+
+To update the table of contents, run `doctoc README.md`
+
+# Development
+
+## vellum pull
+
+If you want to run `vellum pull` locally end to end, here are the steps you'll need to follow.
+
+### Step 1. Run the Codegen Service Locally
+
+- Pull down the private `vellum-ai/codegen-service` repo [here](https://github.com/vellum-ai/codegen-service)
+- Follow its README to run the service locally. This will spin up a local webserver running on `localhost:5111`.
+- By default, when you run that service, it'll use the latest published `@vellum-ai/vellum-codegen` npm package.
+  However, you'll often want to instead point it to the local version defined here in this repo under the `codegen/`
+  directory (which is the source of where @vellum-ai/vellum-codegen is published from). The README in the
+  codegen-service repo should show you how to do that.
+
+### Step 2. Run the Vellum Webserver Locally
+
+- Pull down the private `vellum-ai/vellum` repo [here](https://github.com/vellum-ai/vellum)
+- Follow its README to set up the Vellum backend
+- Run `make run-backend-dev`. This will start up the Vellum backend on `localhost:8000` and have it point to the
+  local codegen-service running on `localhost:5111`.
+
+### Step 3. Run the Vellum CLI Locally
+
+- From within the root of this repo, run `pip install -e .` to install the Vellum CLI locally.
+- Ensure you have a valid workflow config in your `pyproject.toml` file. For example:
+  ```toml
+  [[tool.vellum.workflows]]
+  module = "examples.basic_rag_workflow"
+  workflow_sandbox_id = "<your-workflow-sandbox-id>"
+  ```
+- Ensure you have an API key stored in your `.env` file:
+  ```
+  VELLUM_API_KEY="<your-api-key>"
+  ```
+- Run the cli command and point it to the local Vellum backend running on `localhost:8000`. For example:
+  ```
+  VELLUM_API_URL="http://localhost:8000" vellum pull examples.basic_rag_workflow
+  ```

vellum_ai-0.10.7/ee/vellum_cli/README.md

@@ -0,0 +1,3 @@
+# Vellum CLI
+
+<!-- TODO: To be generated from the generator repo and shared with docs -->

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/ee/vellum_ee/workflows/display/base.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from uuid import UUID
-from typing import TypeVar
+from typing import Optional, TypeVar
 
 
 @dataclass
@@ -20,6 +20,7 @@ WorkflowMetaDisplayOverridesType = TypeVar("WorkflowMetaDisplayOverridesType", b
 @dataclass
 class WorkflowInputsDisplayOverrides:
     id: UUID
+    color: Optional[str] = None
 
 
 @dataclass

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -1,11 +1,8 @@
 from functools import cached_property
 import inspect
 from uuid import UUID
-from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar, get_args
+from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar, get_args, get_origin, cast
 
-from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
-from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
-from vellum_ee.workflows.display.vellum import CodeResourceDefinition, NodeDefinition
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
 from vellum.workflows.ports import Port
@@ -14,6 +11,9 @@ from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.generics import NodeType
 from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.names import pascal_to_title_case
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
+from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
+from vellum_ee.workflows.display.vellum import CodeResourceDefinition, NodeDefinition
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -92,6 +92,29 @@ class BaseNodeDisplay(Generic[NodeType]):
         if node_display_attribute is None:
             return None
 
+        origin = get_origin(attribute_type)
+        args = get_args(attribute_type)
+
+        if origin is not None:
+            # Handle Dict
+            if origin is dict and isinstance(node_display_attribute, dict):
+                if len(args) == 2:
+                    key_type, value_type = args
+                    if all(isinstance(k, key_type) and isinstance(v, value_type) for k, v in
+                           node_display_attribute.items()):
+                        return cast(_NodeDisplayAttrType, node_display_attribute)
+                raise ValueError(f"Node {cls.__name__} must define an explicit {attribute} of type {attribute_type}.")
+
+            # Handle List
+            elif origin is list and isinstance(node_display_attribute, list):
+                if len(args) == 1:
+                    item_type = args[0]
+                    if all(isinstance(item, item_type) for item in node_display_attribute):
+                        return cast(_NodeDisplayAttrType, node_display_attribute)
+                raise ValueError(f"Node {cls.__name__} must define an explicit {attribute} of type {attribute_type}.")
+
+            raise ValueError(f"Node {cls.__name__} must define an explicit {attribute} of type {attribute_type}.")
+
         if isinstance(node_display_attribute, attribute_type):
             return node_display_attribute
 

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/ee/vellum_ee/workflows/display/nodes/vellum/__init__.py

@@ -7,6 +7,7 @@ from .inline_prompt_node import BaseInlinePromptNodeDisplay
 from .inline_subworkflow_node import BaseInlineSubworkflowNodeDisplay
 from .map_node import BaseMapNodeDisplay
 from .merge_node import BaseMergeNodeDisplay
+from .note_node import BaseNoteNodeDisplay
 from .prompt_deployment_node import BasePromptDeploymentNodeDisplay
 from .search_node import BaseSearchNodeDisplay
 from .subworkflow_deployment_node import BaseSubworkflowDeploymentNodeDisplay
@@ -23,6 +24,7 @@ __all__ = [
     "BaseAPINodeDisplay",
     "BaseMapNodeDisplay",
     "BaseMergeNodeDisplay",
+    "BaseNoteNodeDisplay",
     "BaseSearchNodeDisplay",
     "BaseSubworkflowDeploymentNodeDisplay",
     "BaseTemplatingNodeDisplay",

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/ee/vellum_ee/workflows/display/nodes/vellum/conditional_node.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, List, Optional, TypeVar
+from typing import Any, ClassVar, Dict, Generic, List, Optional, TypeVar, Union, Tuple
 
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -29,7 +29,7 @@ from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.expressions.not_in import NotInExpression
 from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.nodes.displayable import ConditionalNode
-from vellum.workflows.types.core import JsonObject
+from vellum.workflows.types.core import JsonObject, ConditionType
 
 _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)
 
@@ -39,11 +39,13 @@ class RuleIdMap:
     id: UUID
     lhs: Optional["RuleIdMap"]
     rhs: Optional["RuleIdMap"]
+    field_node_input_id: Optional[UUID]
+    value_node_input_id: Optional[UUID]
 
 @dataclass
 class ConditionId:
     id: UUID
-    rule_group_id: UUID
+    rule_group_id: Optional[UUID]
 
 
 class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Generic[_ConditionalNodeType]):
@@ -73,8 +75,11 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
                 if rule_id_map is not None
                 else str(uuid4_from_hash(f"{node_id}|rule|{','.join(str(p) for p in path)}"))
             )
-            field_node_input_id = None
-            value_node_input_id = None
+
+            result = self.get_nested_rule_details_by_path(rule_ids, path) if rule_ids else None
+            if result is None:
+                result = self._generate_hash_for_rule_ids(node_id, rule_id)
+            current_id, field_node_input_id, value_node_input_id = result
 
             # Recursive step: Keep recursing until we hit the other descriptors
             if isinstance(descriptor, (AndExpression, OrExpression)):
@@ -97,79 +102,89 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
             # Base cases for other descriptors
             elif isinstance(descriptor, (IsNullExpression, IsNotNullExpression)):
                 expression_node_input = create_node_input(
-                    node_id, f"{rule_id}.field", descriptor._expression, display_context, None
+                    node_id, f"{current_id}.field", descriptor._expression, display_context, field_node_input_id
                 )
                 node_inputs.append(expression_node_input)
-                field_node_input_id = expression_node_input.id
+                field_node_input_id = UUID(expression_node_input.id)
                 operator = self._convert_descriptor_to_operator(descriptor)
 
             elif isinstance(descriptor, (BetweenExpression, NotBetweenExpression)):
                 field_node_input = create_node_input(
-                    node_id, f"{rule_id}.field", descriptor._value, display_context, None
+                    node_id, f"{current_id}.field", descriptor._value, display_context, field_node_input_id
                 )
                 value_node_input = create_node_input(
-                    node_id, f"{rule_id}.value", f"{descriptor._start},{descriptor._end}", display_context, None
+                    node_id, f"{current_id}.value", f"{descriptor._start},{descriptor._end}", display_context, value_node_input_id
                 )
                 node_inputs.extend([field_node_input, value_node_input])
                 operator = self._convert_descriptor_to_operator(descriptor)
-                field_node_input_id = field_node_input.id
-                value_node_input_id = value_node_input.id
+                field_node_input_id = UUID(field_node_input.id)
+                value_node_input_id = UUID(value_node_input.id)
 
             else:
                 lhs = descriptor._lhs  # type: ignore[attr-defined]
                 rhs = descriptor._rhs  # type: ignore[attr-defined]
 
-                lhs_node_input = create_node_input(node_id, f"{rule_id}.field", lhs, display_context, None)
+                lhs_node_input = create_node_input(node_id, f"{current_id}.field", lhs, display_context, field_node_input_id)
                 node_inputs.append(lhs_node_input)
 
                 if descriptor._rhs is not None:  # type: ignore[attr-defined]
-                    rhs_node_input = create_node_input(node_id, f"{rule_id}.value", rhs, display_context, None)
+                    rhs_node_input = create_node_input(node_id, f"{current_id}.value", rhs, display_context, value_node_input_id)
                     node_inputs.append(rhs_node_input)
-                    value_node_input_id = rhs_node_input.id
+                    value_node_input_id = UUID(rhs_node_input.id)
 
                 operator = self._convert_descriptor_to_operator(descriptor)
-                field_node_input_id = lhs_node_input.id
+                field_node_input_id = UUID(lhs_node_input.id)
 
             return {
-                "id": rule_id,
+                "id": str(current_id),
                 "rules": None,
                 "combinator": None,
                 "negated": False,
-                "field_node_input_id": field_node_input_id,
+                "field_node_input_id": str(field_node_input_id),
                 "operator": operator,
-                "value_node_input_id": value_node_input_id,
+                "value_node_input_id": str(value_node_input_id),
             }
 
         conditions = []
         for idx, port in enumerate(node.Ports):
-            if port._condition is None or port._condition_type is None:
-                continue
-
 
             condition_id = str(condition_ids[idx].id if condition_ids else uuid4_from_hash(f"{node_id}|conditions|{idx}"))
             rule_group_id = str(condition_ids[idx].rule_group_id if condition_ids else uuid4_from_hash(f"{condition_id}|rule_group"))
             source_handle_id = str(source_handle_ids.get(idx) or uuid4_from_hash(f"{node_id}|handles|{idx}"))
 
-            rule_ids = self._get_rule_ids()
-            condition_rule = serialize_rule(port._condition, [idx], rule_ids[idx] if len(rule_ids) > idx else None)
-            rules = condition_rule["rules"] if condition_rule["rules"] else [condition_rule]
-
-            conditions.append(
-                {
-                    "id": condition_id,
-                    "type": port._condition_type.value,
-                    "source_handle_id": source_handle_id,
-                    "data": {
-                        "id": rule_group_id,
-                        "rules": rules,
-                        "combinator": "AND",
-                        "negated": False,
-                        "field_node_input_id": None,
-                        "operator": None,
-                        "value_node_input_id": None,
-                    },
-                }
-            )
+            if port._condition is None:
+                if port._condition_type == ConditionType.ELSE:
+                    conditions.append({
+                        "id": condition_id,
+                        "type": ConditionType.ELSE.value,
+                        "source_handle_id": source_handle_id,
+                        "data": None
+                    })
+                else:
+                    continue
+
+            else:
+                rule_ids = self._get_rule_ids()
+                condition_rule = serialize_rule(port._condition, [idx], rule_ids[idx] if len(rule_ids) > idx else None)
+                rules = condition_rule["rules"] if condition_rule["rules"] else [condition_rule]
+                if port._condition_type is None:
+                    raise ValueError("Condition type is None, but a valid ConditionType is expected.")
+                conditions.append(
+                    {
+                        "id": condition_id,
+                        "type": port._condition_type.value,
+                        "source_handle_id": source_handle_id,
+                        "data": {
+                            "id": rule_group_id,
+                            "rules": rules,
+                            "combinator": "AND",
+                            "negated": False,
+                            "field_node_input_id": None,
+                            "operator": None,
+                            "value_node_input_id": None,
+                        },
+                    }
+                )
 
         return {
             "id": str(node_id),
@@ -225,6 +240,36 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
         else:
             raise ValueError(f"Unsupported descriptor type: {descriptor}")
 
+    def get_nested_rule_details_by_path(
+            self, rule_ids: List[RuleIdMap], path: List[int]
+    ) -> Union[Tuple[UUID, Optional[UUID], Optional[UUID]], None]:
+        current_rule = rule_ids[path[0]]
+
+        for step in path[1:]:
+            if step == 0 and current_rule.lhs:
+                current_rule = current_rule.lhs
+            elif step == 1 and current_rule.rhs:
+                current_rule = current_rule.rhs
+            else:
+                return None
+
+        # This is essentially a leaf
+        if current_rule.lhs and current_rule.lhs.lhs is None and current_rule.lhs.rhs is None:
+            return (
+                current_rule.lhs.id,
+                current_rule.lhs.field_node_input_id,
+                current_rule.lhs.value_node_input_id,
+            )
+
+        return None
+
+    def _generate_hash_for_rule_ids(self, node_id, rule_id) -> Tuple[UUID, UUID, UUID]:
+        return (
+            uuid4_from_hash(f"{node_id}|{rule_id}|current"),
+            uuid4_from_hash(f"{node_id}|{rule_id}||field"),
+            uuid4_from_hash(f"{node_id}|{rule_id}||value")
+        )
+
     def _get_source_handle_ids(self) -> Dict[int, UUID]:
         return self._get_explicit_node_display_attr("source_handle_ids", Dict[int, UUID]) or {}
 

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/ee/vellum_ee/workflows/display/nodes/vellum/final_output_node.py

@@ -1,14 +1,16 @@
 from uuid import UUID
 from typing import Any, ClassVar, Generic, Optional, TypeVar
 
+from vellum.workflows.nodes.core.map_node.node import MapNode
+from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
+from vellum.workflows.references.output import OutputReference
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import to_kebab_case
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
-from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
-from vellum.workflows.types.core import JsonObject
 
 _FinalOutputNodeType = TypeVar("_FinalOutputNodeType", bound=FinalOutputNode)
 

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/ee/vellum_ee/workflows/display/nodes/vellum/guardrail_node.py

@@ -1,12 +1,12 @@
 from uuid import UUID
 from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
 
+from vellum.workflows.nodes import GuardrailNode
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum.workflows.nodes import GuardrailNode
-from vellum.workflows.types.core import JsonObject
 
 _GuardrailNodeType = TypeVar("_GuardrailNodeType", bound=GuardrailNode)
 
@@ -15,7 +15,7 @@ class BaseGuardrailNodeDisplay(BaseNodeVellumDisplay[_GuardrailNodeType], Generi
     metric_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/ee/vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py

@@ -2,7 +2,9 @@ from uuid import UUID
 from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
 
 from vellum import PromptBlock, RichTextChildBlock, VellumVariable
-
+from vellum.workflows.nodes import InlinePromptNode
+from vellum.workflows.references import OutputReference
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -10,9 +12,6 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import NodeInput
-from vellum.workflows.nodes import InlinePromptNode
-from vellum.workflows.references import OutputReference
-from vellum.workflows.types.core import JsonObject
 
 _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
 
@@ -23,7 +22,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
     prompt_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id

{vellum_ai-0.10.4 → vellum_ai-0.10.7}/ee/vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py

@@ -1,17 +1,16 @@
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, cast
 
 from vellum import VellumVariable
-
+from vellum.workflows.nodes import InlineSubworkflowNode
+from vellum.workflows.types.core import JsonObject
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
-from vellum_ee.workflows.display.vellum import NodeInput
+from vellum_ee.workflows.display.vellum import NodeInput, WorkflowOutputVellumDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
-from vellum.workflows.nodes import InlineSubworkflowNode
-from vellum.workflows.types.core import JsonObject
 
 _InlineSubworkflowNodeType = TypeVar("_InlineSubworkflowNodeType", bound=InlineSubworkflowNode)
 
@@ -22,19 +21,19 @@ class BaseInlineSubworkflowNodeDisplay(
     workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
 
     def serialize(
-        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
+        self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:
         node = self._node
         node_id = self.node_id
 
         node_inputs, workflow_inputs = self._generate_node_and_workflow_inputs(node_id, node, display_context)
-        workflow_outputs = self._generate_workflow_outputs(node)
 
         subworkflow_display = get_workflow_display(
             base_display_class=display_context.workflow_display_class,
             workflow_class=raise_if_descriptor(node.subworkflow),
             parent_display_context=display_context,
         )
+        workflow_outputs = self._generate_workflow_outputs(node, subworkflow_display.display_context)
         serialized_subworkflow = subworkflow_display.serialize()
 
         return {
@@ -87,13 +86,14 @@ class BaseInlineSubworkflowNodeDisplay(
     def _generate_workflow_outputs(
         self,
         node: Type[InlineSubworkflowNode],
+        display_context: WorkflowDisplayContext,
     ) -> List[VellumVariable]:
         workflow_outputs: List[VellumVariable] = []
         for output_descriptor in raise_if_descriptor(node.subworkflow).Outputs:  # type: ignore[union-attr]
-            node_output_display = self.get_node_output_display(output_descriptor)
+            workflow_output_display = cast(WorkflowOutputVellumDisplay, display_context.workflow_output_displays[output_descriptor])
             output_type = infer_vellum_variable_type(output_descriptor)
             workflow_outputs.append(
-                VellumVariable(id=str(node_output_display.id), key=node_output_display.name, type=output_type)
+                VellumVariable(id=str(workflow_output_display.id), key=workflow_output_display.name, type=output_type)
             )
 
         return workflow_outputs