vellum-ai 0.12.17__py3-none-any.whl → 0.13.1__py3-none-any.whl

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

@@ -74,6 +74,9 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
 
         compiled_inputs: List[PromptDeploymentInputRequest] = []
 
+        if not self.prompt_inputs:
+            return compiled_inputs
+
         for input_name, input_value in self.prompt_inputs.items():
             if isinstance(input_value, str):
                 compiled_inputs.append(

vellum/workflows/nodes/displayable/bases/search_node.py

@@ -15,6 +15,7 @@ from vellum.core import ApiError, RequestOptions
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.displayable.bases.types import SearchFilters
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.types.generics import StateType
 
@@ -33,7 +34,11 @@ class BaseSearchNode(BaseNode[StateType], Generic[StateType]):
     document_index: Union[UUID, str] - Either the UUID or name of the Vellum Document Index that you'd like to search
         against
     query: str - The query to search for
-    options: Optional[SearchRequestOptionsRequest] = None - Runtime configuration for the search
+    limit: Optional[int] = None - The maximum number of results to return.
+    weights: Optional[SearchWeightsRequest] = None - The weights to use for the search. Must add up to 1.0.
+    result_merging: Optional[SearchResultMergingRequest] = None - The configuration for merging results.
+    filters: Optional[SearchFiltersRequest] = None - The filters to apply to the search.
+    options: Optional[SearchRequestOptionsRequest] = None - [DEPRECATED] Runtime configuration for the search
     request_options: Optional[RequestOptions] = None - The request options to use for the search
     """
 
@@ -43,11 +48,24 @@ class BaseSearchNode(BaseNode[StateType], Generic[StateType]):
     # The Document Index to Search against. Identified by either its UUID or its name.
     document_index: ClassVar[Union[UUID, str]]
 
+    # The maximum number of results to return.
+    limit: ClassVar[Optional[int]] = None
+
+    # The weights to use for the search. Must add up to 1.0.
+    weights: ClassVar[Optional[SearchWeightsRequest]] = None
+
+    # The configuration for merging results.
+    result_merging: ClassVar[Optional[SearchResultMergingRequest]] = None
+
+    # The filters to apply to the search.
+    filters: ClassVar[Optional[SearchFilters]] = None
+
     # Ideally we could reuse node descriptors to derive other node descriptor values. Two action items are
     # blocking us from doing so in this use case:
     # 1. Node Descriptor resolution during runtime - https://app.shortcut.com/vellum/story/4781
     # 2. Math operations between descriptors - https://app.shortcut.com/vellum/story/4782
     # search_weights = DEFAULT_SEARCH_WEIGHTS
+    # Deprecated: Use the top level `limit`, `weights`, `result_merging`, and `filters` attributes instead
     options = SearchRequestOptionsRequest(
         limit=DEFAULT_SEARCH_LIMIT,
         weights=SearchWeightsRequest(
@@ -77,7 +95,7 @@ class BaseSearchNode(BaseNode[StateType], Generic[StateType]):
             return self._context.vellum_client.search(
                 query=self.query,
                 document_index=str(self.document_index),
-                options=self.options,
+                options=self._get_options_request(),
             )
         except NotFoundError:
             raise NodeException(
@@ -90,6 +108,23 @@ class BaseSearchNode(BaseNode[StateType], Generic[StateType]):
                 code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
+    def _get_options_request(self) -> SearchRequestOptionsRequest:
+        return SearchRequestOptionsRequest(
+            limit=self.limit if self.limit is not None else self.options.limit,
+            weights=self.weights if self.weights is not None else self.options.weights,
+            result_merging=self.result_merging if self.result_merging is not None else self.options.result_merging,
+            filters=self._get_filters_request(),
+        )
+
+    def _get_filters_request(self) -> Optional[SearchFiltersRequest]:
+        if self.filters is None:
+            return self.options.filters
+
+        return SearchFiltersRequest(
+            external_ids=self.filters.external_ids,
+            metadata=self.filters.metadata.to_request() if self.filters.metadata is not None else None,
+        )
+
     def run(self) -> Outputs:
         response = self._perform_search()
         return self.Outputs(results=response.results)

vellum/workflows/nodes/displayable/bases/tests/test_utils.py

@@ -0,0 +1,61 @@
+import pytest
+import enum
+
+from vellum.client.types.chat_history_vellum_value import ChatHistoryVellumValue
+from vellum.client.types.chat_message import ChatMessage
+from vellum.client.types.json_vellum_value import JsonVellumValue
+from vellum.client.types.number_vellum_value import NumberVellumValue
+from vellum.client.types.search_result import SearchResult
+from vellum.client.types.search_result_document import SearchResultDocument
+from vellum.client.types.search_results_vellum_value import SearchResultsVellumValue
+from vellum.client.types.string_vellum_value import StringVellumValue
+from vellum.client.types.string_vellum_value_request import StringVellumValueRequest
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value, primitive_to_vellum_value_request
+
+
+class MockEnum(enum.Enum):
+    FOO = "foo"
+
+
+@pytest.mark.parametrize(
+    ["value", "expected_output"],
+    [
+        ("hello", StringVellumValue(value="hello")),
+        (MockEnum.FOO, StringVellumValue(value="foo")),
+        (1, NumberVellumValue(value=1)),
+        (1.0, NumberVellumValue(value=1.0)),
+        (
+            [ChatMessage(role="USER", text="hello")],
+            ChatHistoryVellumValue(value=[ChatMessage(role="USER", text="hello")]),
+        ),
+        (
+            [
+                SearchResult(
+                    text="Search query",
+                    score="0.0",
+                    keywords=["keywords"],
+                    document=SearchResultDocument(label="label"),
+                )
+            ],
+            SearchResultsVellumValue(
+                value=[
+                    SearchResult(
+                        text="Search query",
+                        score="0.0",
+                        keywords=["keywords"],
+                        document=SearchResultDocument(label="label"),
+                    )
+                ]
+            ),
+        ),
+        (StringVellumValue(value="hello"), StringVellumValue(value="hello")),
+        (StringVellumValueRequest(value="hello"), StringVellumValueRequest(value="hello")),
+        ({"foo": "bar"}, JsonVellumValue(value={"foo": "bar"})),
+    ],
+)
+def test_primitive_to_vellum_value(value, expected_output):
+    assert primitive_to_vellum_value(value) == expected_output
+
+
+def test_primitive_to_vellum_value_request():
+    assert primitive_to_vellum_value_request("hello") == StringVellumValueRequest(value="hello")

vellum/workflows/nodes/displayable/bases/types.py

@@ -0,0 +1,42 @@
+from typing import Any, List, Optional, Union
+
+from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.client.types.condition_combinator import ConditionCombinator
+from vellum.client.types.logical_operator import LogicalOperator
+from vellum.client.types.vellum_value_logical_condition_group_request import VellumValueLogicalConditionGroupRequest
+from vellum.client.types.vellum_value_logical_condition_request import VellumValueLogicalConditionRequest
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value_request
+
+
+class MetadataLogicalConditionGroup(UniversalBaseModel):
+    combinator: ConditionCombinator
+    negated: bool
+    conditions: List["MetadataLogicalExpression"]
+
+    def to_request(self) -> VellumValueLogicalConditionGroupRequest:
+        return VellumValueLogicalConditionGroupRequest(
+            combinator=self.combinator,
+            negated=self.negated,
+            conditions=[c.to_request() for c in self.conditions],
+        )
+
+
+class MetadataLogicalCondition(UniversalBaseModel):
+    lhs_variable: Any
+    operator: LogicalOperator
+    rhs_variable: Any
+
+    def to_request(self) -> VellumValueLogicalConditionRequest:
+        return VellumValueLogicalConditionRequest(
+            lhs_variable=primitive_to_vellum_value_request(self.lhs_variable),
+            operator=self.operator,
+            rhs_variable=primitive_to_vellum_value_request(self.rhs_variable),
+        )
+
+
+MetadataLogicalExpression = Union[MetadataLogicalConditionGroup, MetadataLogicalCondition]
+
+
+class SearchFilters(UniversalBaseModel):
+    external_ids: Optional[List[str]] = None
+    metadata: Optional[MetadataLogicalConditionGroup] = None

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

@@ -0,0 +1,112 @@
+import enum
+import json
+from typing import Any, List, Union, cast
+
+from vellum.client.types.array_vellum_value import ArrayVellumValue
+from vellum.client.types.array_vellum_value_request import ArrayVellumValueRequest
+from vellum.client.types.audio_vellum_value import AudioVellumValue
+from vellum.client.types.audio_vellum_value_request import AudioVellumValueRequest
+from vellum.client.types.chat_history_vellum_value import ChatHistoryVellumValue
+from vellum.client.types.chat_history_vellum_value_request import ChatHistoryVellumValueRequest
+from vellum.client.types.chat_message import ChatMessage
+from vellum.client.types.error_vellum_value import ErrorVellumValue
+from vellum.client.types.error_vellum_value_request import ErrorVellumValueRequest
+from vellum.client.types.function_call_vellum_value import FunctionCallVellumValue
+from vellum.client.types.function_call_vellum_value_request import FunctionCallVellumValueRequest
+from vellum.client.types.image_vellum_value import ImageVellumValue
+from vellum.client.types.image_vellum_value_request import ImageVellumValueRequest
+from vellum.client.types.json_vellum_value import JsonVellumValue
+from vellum.client.types.json_vellum_value_request import JsonVellumValueRequest
+from vellum.client.types.number_vellum_value import NumberVellumValue
+from vellum.client.types.number_vellum_value_request import NumberVellumValueRequest
+from vellum.client.types.search_result import SearchResult
+from vellum.client.types.search_result_request import SearchResultRequest
+from vellum.client.types.search_results_vellum_value import SearchResultsVellumValue
+from vellum.client.types.search_results_vellum_value_request import SearchResultsVellumValueRequest
+from vellum.client.types.string_vellum_value import StringVellumValue
+from vellum.client.types.string_vellum_value_request import StringVellumValueRequest
+from vellum.client.types.vellum_value import VellumValue
+from vellum.client.types.vellum_value_request import VellumValueRequest
+
+VELLUM_VALUE_REQUEST_TUPLE = (
+    StringVellumValueRequest,
+    NumberVellumValueRequest,
+    JsonVellumValueRequest,
+    ImageVellumValueRequest,
+    AudioVellumValueRequest,
+    FunctionCallVellumValueRequest,
+    ErrorVellumValueRequest,
+    ArrayVellumValueRequest,
+    ChatHistoryVellumValueRequest,
+    SearchResultsVellumValueRequest,
+)
+
+
+def primitive_to_vellum_value(value: Any) -> VellumValue:
+    """Converts a python primitive to a VellumValue"""
+
+    if isinstance(value, str):
+        return StringVellumValue(value=value)
+    elif isinstance(value, enum.Enum):
+        return StringVellumValue(value=value.value)
+    elif isinstance(value, (int, float)):
+        return NumberVellumValue(value=value)
+    elif isinstance(value, list) and (
+        all(isinstance(message, ChatMessage) for message in value)
+        or all(isinstance(message, ChatMessage) for message in value)
+    ):
+        chat_messages = cast(Union[List[ChatMessage], List[ChatMessage]], value)
+        return ChatHistoryVellumValue(value=chat_messages)
+    elif isinstance(value, list) and (
+        all(isinstance(search_result, SearchResultRequest) for search_result in value)
+        or all(isinstance(search_result, SearchResult) for search_result in value)
+    ):
+        search_results = cast(Union[List[SearchResultRequest], List[SearchResult]], value)
+        return SearchResultsVellumValue(value=search_results)
+    elif isinstance(
+        value,
+        (
+            StringVellumValue,
+            NumberVellumValue,
+            JsonVellumValue,
+            ImageVellumValue,
+            AudioVellumValue,
+            FunctionCallVellumValue,
+            ErrorVellumValue,
+            ArrayVellumValue,
+            ChatHistoryVellumValue,
+            SearchResultsVellumValue,
+        ),
+    ):
+        return value
+    elif isinstance(
+        value,
+        VELLUM_VALUE_REQUEST_TUPLE,
+    ):
+        # This type ignore is safe because consumers of this function won't care the difference between
+        # XVellumValue and XVellumValueRequest. Hopefully in the near future, neither will we
+        return value  # type: ignore
+
+    try:
+        json_value = json.dumps(value)
+    except json.JSONDecodeError:
+        raise ValueError(f"Unsupported variable type: {value.__class__.__name__}")
+
+    return JsonVellumValue(value=json.loads(json_value))
+
+
+def primitive_to_vellum_value_request(value: Any) -> VellumValueRequest:
+    vellum_value = primitive_to_vellum_value(value)
+    vellum_value_request_class = next(
+        (
+            vellum_value_request_class
+            for vellum_value_request_class in VELLUM_VALUE_REQUEST_TUPLE
+            if vellum_value_request_class.__name__.startswith(vellum_value.__class__.__name__)
+        ),
+        None,
+    )
+
+    if vellum_value_request_class is None:
+        raise ValueError(f"Unsupported variable type: {vellum_value.__class__.__name__}")
+
+    return vellum_value_request_class.model_validate(vellum_value.model_dump())

vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py

@@ -76,7 +76,6 @@ def test_inline_prompt_node__function_definitions(vellum_adhoc_prompt_client):
     class MyNode(InlinePromptNode):
         ml_model = "gpt-4o"
         functions = [my_function]
-        prompt_inputs = {}
         blocks = []
 
     # AND a known response from invoking an inline prompt

vellum/workflows/nodes/displayable/search_node/tests/test_node.py

@@ -0,0 +1,164 @@
+from vellum import SearchResponse, SearchResult, SearchResultDocument
+from vellum.client.types.search_filters_request import SearchFiltersRequest
+from vellum.client.types.search_request_options_request import SearchRequestOptionsRequest
+from vellum.client.types.search_result_merging_request import SearchResultMergingRequest
+from vellum.client.types.search_weights_request import SearchWeightsRequest
+from vellum.client.types.string_vellum_value_request import StringVellumValueRequest
+from vellum.client.types.vellum_value_logical_condition_group_request import VellumValueLogicalConditionGroupRequest
+from vellum.client.types.vellum_value_logical_condition_request import VellumValueLogicalConditionRequest
+from vellum.workflows.nodes.displayable.bases.types import (
+    MetadataLogicalCondition,
+    MetadataLogicalConditionGroup,
+    SearchFilters,
+)
+from vellum.workflows.nodes.displayable.search_node.node import SearchNode
+
+
+def test_run_workflow__happy_path(vellum_client):
+    """Confirm that we can successfully invoke a Workflow with the new option attributes"""
+
+    # GIVEN a workflow that's set up run a Search Node
+    class MySearchNode(SearchNode):
+        query = "Search query"
+        document_index = "document_index"
+        limit = 1
+        weights = SearchWeightsRequest(
+            semantic_similarity=0.8,
+            keywords=0.2,
+        )
+        result_merging = SearchResultMergingRequest(
+            enabled=True,
+        )
+        filters = SearchFilters(
+            external_ids=["external_id"],
+            metadata=MetadataLogicalConditionGroup(
+                combinator="AND",
+                negated=False,
+                conditions=[
+                    MetadataLogicalCondition(
+                        lhs_variable="TYPE",
+                        operator="=",
+                        rhs_variable="COMPANY",
+                    )
+                ],
+            ),
+        )
+
+    # AND a Search request that will return a 200 ok resposne
+    search_response = SearchResponse(
+        results=[
+            SearchResult(
+                text="Search query", score="0.0", keywords=["keywords"], document=SearchResultDocument(label="label")
+            )
+        ]
+    )
+
+    vellum_client.search.return_value = search_response
+
+    # WHEN we run the workflow
+    outputs = MySearchNode().run()
+
+    # THEN the workflow should have completed successfully
+    assert outputs.text == "Search query"
+
+    # AND the options should be as expected
+    assert vellum_client.search.call_args.kwargs["options"] == SearchRequestOptionsRequest(
+        limit=1,
+        weights=SearchWeightsRequest(
+            semantic_similarity=0.8,
+            keywords=0.2,
+        ),
+        result_merging=SearchResultMergingRequest(
+            enabled=True,
+        ),
+        filters=SearchFiltersRequest(
+            external_ids=["external_id"],
+            metadata=VellumValueLogicalConditionGroupRequest(
+                combinator="AND",
+                negated=False,
+                conditions=[
+                    VellumValueLogicalConditionRequest(
+                        lhs_variable=StringVellumValueRequest(value="TYPE"),
+                        operator="=",
+                        rhs_variable=StringVellumValueRequest(value="COMPANY"),
+                    )
+                ],
+            ),
+        ),
+    )
+
+
+def test_run_workflow__happy_path__options_attribute(vellum_client):
+    """Confirm that we can successfully invoke a single Search node with the legacy options attribute"""
+
+    # GIVEN a workflow that's set up run a Search Node
+    class MySearchNode(SearchNode):
+        query = "Search query"
+        document_index = "document_index"
+        options = SearchRequestOptionsRequest(
+            limit=1,
+            weights=SearchWeightsRequest(
+                semantic_similarity=0.8,
+                keywords=0.2,
+            ),
+            result_merging=SearchResultMergingRequest(
+                enabled=True,
+            ),
+            filters=SearchFiltersRequest(
+                external_ids=["external_id"],
+                metadata=VellumValueLogicalConditionGroupRequest(
+                    combinator="AND",
+                    negated=False,
+                    conditions=[
+                        VellumValueLogicalConditionRequest(
+                            lhs_variable=StringVellumValueRequest(value="TYPE"),
+                            operator="=",
+                            rhs_variable=StringVellumValueRequest(value="COMPANY"),
+                        )
+                    ],
+                ),
+            ),
+        )
+
+    # AND a Search request that will return a 200 ok resposne
+    search_response = SearchResponse(
+        results=[
+            SearchResult(
+                text="Search query", score="0.0", keywords=["keywords"], document=SearchResultDocument(label="label")
+            )
+        ]
+    )
+
+    vellum_client.search.return_value = search_response
+
+    # WHEN we run the workflow
+    outputs = MySearchNode().run()
+
+    # THEN the workflow should have completed successfully
+    assert outputs.text == "Search query"
+
+    # AND the options should be as expected
+    assert vellum_client.search.call_args.kwargs["options"] == SearchRequestOptionsRequest(
+        limit=1,
+        weights=SearchWeightsRequest(
+            semantic_similarity=0.8,
+            keywords=0.2,
+        ),
+        result_merging=SearchResultMergingRequest(
+            enabled=True,
+        ),
+        filters=SearchFiltersRequest(
+            external_ids=["external_id"],
+            metadata=VellumValueLogicalConditionGroupRequest(
+                combinator="AND",
+                negated=False,
+                conditions=[
+                    VellumValueLogicalConditionRequest(
+                        lhs_variable=StringVellumValueRequest(value="TYPE"),
+                        operator="=",
+                        rhs_variable=StringVellumValueRequest(value="COMPANY"),
+                    )
+                ],
+            ),
+        ),
+    )

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

@@ -1,3 +1,4 @@
+from unittest import mock
 from uuid import uuid4
 from typing import Any, Iterator, List
 
@@ -33,7 +34,6 @@ def test_inline_text_prompt_node__basic(vellum_adhoc_prompt_client):
 
     class MyInlinePromptNode(InlinePromptNode):
         ml_model = "gpt-4o"
-        prompt_inputs = {}
         blocks = []
 
     # AND a known response from invoking an inline prompt
@@ -90,7 +90,7 @@ def test_inline_text_prompt_node__basic(vellum_adhoc_prompt_client):
             logit_bias=None,
             custom_parameters=None,
         ),
-        request_options=None,
+        request_options=mock.ANY,
     )
 
 
@@ -107,7 +107,6 @@ def test_inline_text_prompt_node__catch_provider_error(vellum_adhoc_prompt_clien
     @TryNode.wrap(on_error_code=WorkflowErrorCode.PROVIDER_ERROR)
     class MyInlinePromptNode(InlinePromptNode):
         ml_model = "gpt-4o"
-        prompt_inputs = {}
         blocks = []
 
     # AND a known response from invoking an inline prompt that fails

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

@@ -27,7 +27,6 @@ def test_text_prompt_deployment_node__basic(vellum_client):
 
     class MyPromptDeploymentNode(PromptDeploymentNode):
         deployment = "my-deployment"
-        prompt_inputs = {}
 
     # AND a known response from invoking a deployed prompt
     expected_outputs: List[PromptOutput] = [

vellum/workflows/runner/runner.py

@@ -4,7 +4,21 @@ 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, List, Optional, Sequence, Set, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
 
 from vellum.workflows.constants import UNDEF
 from vellum.workflows.context import execution_context, get_parent_context
@@ -76,6 +90,7 @@ class WorkflowRunner(Generic[StateType]):
         cancel_signal: Optional[ThreadingEvent] = None,
         node_output_mocks: Optional[List[BaseOutputs]] = None,
         parent_context: Optional[ParentContext] = None,
+        max_concurrency: Optional[int] = None,
     ):
         if state and external_inputs:
             raise ValueError("Can only run a Workflow providing one of state or external inputs, not both")
@@ -120,6 +135,9 @@ class WorkflowRunner(Generic[StateType]):
         # This queue is responsible for sending events from the inner worker threads to WorkflowRunner
         self._workflow_event_inner_queue: Queue[WorkflowEvent] = Queue()
 
+        self._max_concurrency = max_concurrency
+        self._concurrency_queue: Queue[Tuple[StateType, Type[BaseNode], Optional[Edge]]] = Queue()
+
         # This queue is responsible for sending events from WorkflowRunner to the background thread
         # for user defined emitters
         self._background_thread_queue: Queue[BackgroundThreadItem] = Queue()
@@ -350,7 +368,19 @@ class WorkflowRunner(Generic[StateType]):
                 else:
                     next_state = state
 
-                self._run_node_if_ready(next_state, edge.to_node, edge)
+                if self._max_concurrency:
+                    self._concurrency_queue.put((next_state, edge.to_node, edge))
+                else:
+                    self._run_node_if_ready(next_state, edge.to_node, edge)
+
+        if self._max_concurrency:
+            num_nodes_to_run = self._max_concurrency - len(self._active_nodes_by_execution_id)
+            for _ in range(num_nodes_to_run):
+                if self._concurrency_queue.empty():
+                    break
+
+                next_state, node_class, invoked_edge = self._concurrency_queue.get()
+                self._run_node_if_ready(next_state, node_class, invoked_edge)
 
     def _run_node_if_ready(
         self,
@@ -513,8 +543,11 @@ class WorkflowRunner(Generic[StateType]):
         )
         for node_cls in self._entrypoints:
             try:
-                with execution_context(parent_context=current_parent):
-                    self._run_node_if_ready(self._initial_state, node_cls)
+                if not self._max_concurrency or len(self._active_nodes_by_execution_id) < self._max_concurrency:
+                    with execution_context(parent_context=current_parent):
+                        self._run_node_if_ready(self._initial_state, node_cls)
+                else:
+                    self._concurrency_queue.put((self._initial_state, node_cls, None))
             except NodeException as e:
                 self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
                 return

vellum/workflows/types/tests/test_utils.py

@@ -21,6 +21,7 @@ class ExampleClass:
     zeta: ClassVar[str]
     eta: List[str]
     kappa: Any
+    mu: list[str]
 
 
 T = TypeVar("T")
@@ -56,6 +57,7 @@ class ExampleNode(BaseNode):
         (ExampleInheritedClass, "beta", (int,)),
         (ExampleNode.Outputs, "iota", (str,)),
         (ExampleClass, "kappa", (Any,)),
+        (ExampleClass, "mu", (list[str],)),
     ],
     ids=[
         "str",
@@ -71,6 +73,7 @@ class ExampleNode(BaseNode):
         "inherited_parent_class_var",
         "try_node_output",
         "any",
+        "list_str_generic",
     ],
 )
 def test_infer_types(cls, attr_name, expected_type):
@@ -80,9 +83,9 @@ def test_infer_types(cls, attr_name, expected_type):
 @pytest.mark.parametrize(
     "cls, expected_attr_names",
     [
-        (ExampleClass, {"alpha", "beta", "gamma", "epsilon", "zeta", "eta", "kappa"}),
+        (ExampleClass, {"alpha", "beta", "gamma", "epsilon", "zeta", "eta", "kappa", "mu"}),
         (ExampleGenericClass, {"delta"}),
-        (ExampleInheritedClass, {"alpha", "beta", "gamma", "epsilon", "zeta", "eta", "theta", "kappa"}),
+        (ExampleInheritedClass, {"alpha", "beta", "gamma", "epsilon", "zeta", "eta", "theta", "kappa", "mu"}),
     ],
 )
 def test_class_attr_names(cls, expected_attr_names):

vellum/workflows/types/utils.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 from datetime import datetime
 import importlib
+from types import GenericAlias
 from typing import (
     Any,
     ClassVar,
@@ -77,6 +78,9 @@ def infer_types(object_: Type, attr_name: str, localns: Optional[Dict[str, Any]]
                 return (type_hint,)
             if isinstance(type_hint, SpecialGenericAlias):
                 return (type_hint,)
+            if isinstance(type_hint, GenericAlias):
+                # In future versions of python, list[str] will be a `GenericAlias`
+                return (cast(Type, type_hint),)
             if isinstance(type_hint, TypeVar):
                 if type_hint in type_var_mapping:
                     return (type_var_mapping[type_hint],)