vellum-ai 0.12.3__tar.gz → 0.12.5__tar.gz

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/PKG-INFO

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

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_cli/__init__.py

@@ -125,10 +125,16 @@ Should only be used for debugging purposes.""",
     help="""Exclude the code definition of each Resource from the pull response. \
 Should only be used for debugging purposes.""",
 )
+@click.option(
+    "--strict",
+    is_flag=True,
+    help="""Raises an error immediately if there are any issues with the pulling of the Resource.""",
+)
 def pull(
     ctx: click.Context,
     include_json: Optional[bool],
     exclude_code: Optional[bool],
+    strict: Optional[bool],
 ) -> None:
     """Pull Resources from Vellum"""
 
@@ -136,6 +142,7 @@ def pull(
         pull_command(
             include_json=include_json,
             exclude_code=exclude_code,
+            strict=strict,
         )
 
 
@@ -159,12 +166,18 @@ Should only be used for debugging purposes.""",
     help="""Exclude the code definition of the Workflow from the pull response. \
 Should only be used for debugging purposes.""",
 )
+@click.option(
+    "--strict",
+    is_flag=True,
+    help="""Raises an error immediately if there are any issues with the pulling of the Workflow.""",
+)
 def workflows_pull(
     module: Optional[str],
     include_json: Optional[bool],
     workflow_sandbox_id: Optional[str],
     workflow_deployment: Optional[str],
     exclude_code: Optional[bool],
+    strict: Optional[bool],
 ) -> None:
     """
     Pull Workflows from Vellum. If a module is provided, only the Workflow for that module will be pulled.
@@ -177,6 +190,7 @@ def workflows_pull(
         workflow_sandbox_id=workflow_sandbox_id,
         workflow_deployment=workflow_deployment,
         exclude_code=exclude_code,
+        strict=strict,
     )
 
 

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_cli/pull.py

@@ -13,6 +13,8 @@ from vellum.workflows.vellum_client import create_vellum_client
 from vellum_cli.config import VellumCliConfig, WorkflowConfig, load_vellum_cli_config
 from vellum_cli.logger import load_cli_logger
 
+ERROR_LOG_FILE_NAME = "error.log"
+
 
 def _is_valid_uuid(val: Union[str, UUID, None]) -> bool:
     try:
@@ -81,6 +83,7 @@ def pull_command(
     workflow_deployment: Optional[str] = None,
     include_json: Optional[bool] = None,
     exclude_code: Optional[bool] = None,
+    strict: Optional[bool] = None,
 ) -> None:
     load_dotenv()
     logger = load_cli_logger()
@@ -109,6 +112,8 @@ def pull_command(
         query_parameters["include_json"] = include_json
     if exclude_code:
         query_parameters["exclude_code"] = exclude_code
+    if strict:
+        query_parameters["strict"] = strict
 
     response = client.workflows.pull(
         pk,
@@ -119,6 +124,7 @@ def pull_command(
     zip_buffer = io.BytesIO(zip_bytes)
 
     target_dir = os.path.join(os.getcwd(), *workflow_config.module.split("."))
+    error_content = ""
     with zipfile.ZipFile(zip_buffer) as zip_file:
         # Delete files in target_dir that aren't in the zip file
         if os.path.exists(target_dir):
@@ -146,8 +152,13 @@ def pull_command(
             target_file = os.path.join(target_dir, file_name)
             os.makedirs(os.path.dirname(target_file), exist_ok=True)
             with zip_file.open(file_name) as source, open(target_file, "w") as target:
+                content = source.read().decode("utf-8")
+                if file_name == ERROR_LOG_FILE_NAME:
+                    error_content = content
+                    continue
+
                 logger.info(f"Writing to {target_file}...")
-                target.write(source.read().decode("utf-8"))
+                target.write(content)
 
     if include_json:
         logger.warning(
@@ -158,4 +169,7 @@ Its schema should be considered unstable and subject to change at any time."""
     if save_lock_file:
         config.save()
 
-    logger.info(f"Successfully pulled Workflow into {workflow_config.module}")
+    if error_content:
+        logger.error(error_content)
+    else:
+        logger.info(f"Successfully pulled Workflow into {workflow_config.module}")

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_cli/tests/test_pull.py

@@ -441,3 +441,48 @@ def test_pull__sandbox_id_with_other_workflow_deployment_in_lock(vellum_client,
                 "deployments": [],
             },
         ]
+
+
+def test_pull__handle_error_log(vellum_client, mock_module):
+    # GIVEN a pyproject.toml with a workflow configured
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+
+    # AND the workflow pull API call returns a zip file with an error log
+    vellum_client.workflows.pull.return_value = iter(
+        [_zip_file_map({"workflow.py": "print('hello')", "error.log": "test error"})]
+    )
+
+    # WHEN the user runs the pull command with the new workflow sandbox id
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--workflow-sandbox-id", workflow_sandbox_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the error log file is not written to the module directory
+    assert not os.path.exists(os.path.join(temp_dir, *module.split("."), "error.log"))
+
+    # AND the error log is printed to the console
+    assert result.output.endswith("\x1b[31;20mtest error\x1b[0m\n")
+
+
+def test_pull__strict__with_error(vellum_client, mock_module):
+    # GIVEN a pyproject.toml with a workflow configured
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+
+    # AND the workflow pull API call returns a zip file
+    vellum_client.workflows.pull.return_value = iter([_zip_file_map({"workflow.py": "print('hello')"})])
+
+    # WHEN the user runs the pull command with the new workflow sandbox id
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["workflows", "pull", "--strict", "--workflow-sandbox-id", workflow_sandbox_id])
+
+    # THEN the command returns successfully
+    assert result.exit_code == 0
+
+    # AND the pull api is called with strict=True
+    vellum_client.workflows.pull.assert_called_once()
+    call_args = vellum_client.workflows.pull.call_args.kwargs
+    assert call_args["request_options"]["additional_query_parameters"] == {"strict": True}

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_ee/workflows/display/nodes/base_node_vellum_display.py

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import ClassVar, Dict, Optional
+from typing import ClassVar, Dict, Optional, Union
 
 from vellum.workflows.nodes.utils import get_unadorned_node
 from vellum.workflows.ports import Port
@@ -18,7 +18,7 @@ class BaseNodeVellumDisplay(BaseNodeDisplay[NodeType]):
     target_handle_id: ClassVar[Optional[UUID]] = None
 
     # Used to explicitly set the node input ids by name for a node
-    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
+    node_input_ids_by_name: ClassVar[Dict[str, Union[UUID, str]]] = {}
 
     def _get_node_display_uuid(self, attribute: str) -> UUID:
         explicit_value = self._get_explicit_node_display_attr(attribute, UUID)

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -1,6 +1,6 @@
 import inspect
 from uuid import UUID
-from typing import ClassVar, Dict, Generic, Optional, TypeVar
+from typing import ClassVar, Generic, Optional, TypeVar
 
 from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
 from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path
@@ -21,8 +21,6 @@ class BaseCodeExecutionNodeDisplay(BaseNodeVellumDisplay[_CodeExecutionNodeType]
     output_id: ClassVar[Optional[UUID]] = None
     log_output_id: ClassVar[Optional[UUID]] = None
 
-    node_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
-
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
     ) -> JsonObject:

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_ee/workflows/display/nodes/vellum/conditional_node.py

@@ -40,17 +40,17 @@ _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)
 
 @dataclass
 class RuleIdMap:
-    id: UUID
+    id: str
     lhs: Optional["RuleIdMap"]
     rhs: Optional["RuleIdMap"]
-    field_node_input_id: Optional[UUID]
-    value_node_input_id: Optional[UUID]
+    field_node_input_id: Optional[str]
+    value_node_input_id: Optional[str]
 
 
 @dataclass
 class ConditionId:
-    id: UUID
-    rule_group_id: Optional[UUID]
+    id: str
+    rule_group_id: Optional[str]
 
 
 class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Generic[_ConditionalNodeType]):
@@ -112,7 +112,7 @@ but the defined conditions have length {len(condition_ids)}"""
                     node_id, f"{current_id}.field", descriptor._expression, display_context, field_node_input_id
                 )
                 node_inputs.append(expression_node_input)
-                field_node_input_id = UUID(expression_node_input.id)
+                field_node_input_id = expression_node_input.id
                 operator = self._convert_descriptor_to_operator(descriptor)
 
             elif isinstance(descriptor, (BetweenExpression, NotBetweenExpression)):
@@ -128,8 +128,8 @@ but the defined conditions have length {len(condition_ids)}"""
                 )
                 node_inputs.extend([field_node_input, value_node_input])
                 operator = self._convert_descriptor_to_operator(descriptor)
-                field_node_input_id = UUID(field_node_input.id)
-                value_node_input_id = UUID(value_node_input.id)
+                field_node_input_id = field_node_input.id
+                value_node_input_id = value_node_input.id
 
             else:
                 lhs = descriptor._lhs  # type: ignore[attr-defined]
@@ -145,19 +145,19 @@ but the defined conditions have length {len(condition_ids)}"""
                         node_id, f"{current_id}.value", rhs, display_context, value_node_input_id
                     )
                     node_inputs.append(rhs_node_input)
-                    value_node_input_id = UUID(rhs_node_input.id)
+                    value_node_input_id = rhs_node_input.id
 
                 operator = self._convert_descriptor_to_operator(descriptor)
-                field_node_input_id = UUID(lhs_node_input.id)
+                field_node_input_id = lhs_node_input.id
 
             return {
-                "id": str(current_id),
+                "id": current_id,
                 "rules": None,
                 "combinator": None,
                 "negated": False,
-                "field_node_input_id": str(field_node_input_id),
+                "field_node_input_id": field_node_input_id,
                 "operator": operator,
-                "value_node_input_id": str(value_node_input_id),
+                "value_node_input_id": value_node_input_id,
             }
 
         conditions = []
@@ -263,7 +263,7 @@ but the defined conditions have length {len(condition_ids)}"""
 
     def get_nested_rule_details_by_path(
         self, rule_ids: List[RuleIdMap], path: List[int]
-    ) -> Union[Tuple[UUID, Optional[UUID], Optional[UUID]], None]:
+    ) -> Union[Tuple[str, Optional[str], Optional[str]], None]:
         current_rule = rule_ids[path[0]]
 
         for step in path[1:]:
@@ -284,11 +284,11 @@ but the defined conditions have length {len(condition_ids)}"""
 
         return None
 
-    def _generate_hash_for_rule_ids(self, node_id, rule_id) -> Tuple[UUID, UUID, UUID]:
+    def _generate_hash_for_rule_ids(self, node_id, rule_id) -> Tuple[str, str, str]:
         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"),
+            str(uuid4_from_hash(f"{node_id}|{rule_id}|current")),
+            str(uuid4_from_hash(f"{node_id}|{rule_id}||field")),
+            str(uuid4_from_hash(f"{node_id}|{rule_id}||value")),
         )
 
     def _get_source_handle_ids(self) -> Dict[int, UUID]:

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_ee/workflows/display/nodes/vellum/search_node.py

@@ -28,7 +28,9 @@ class VariableIdMap:
 
 
 class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_SearchNodeType]):
-    input_variable_ids_by_logical_id: Optional[Dict[str, str]] = None
+    # A mapping between the id of the operand (e.g. "lhs_variable_id" or "rhs_variable_id") and the id of the node input
+    # that the operand is pointing to.
+    metadata_filter_input_id_by_operand_id: Dict[UUID, UUID] = {}
 
     def serialize(
         self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
@@ -149,18 +151,20 @@ class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_Sea
                 variables,
             )
         elif isinstance(logical_expression, VellumValueLogicalConditionRequest):
-            lhs_variable_id = str(logical_expression.lhs_variable.value)
-            rhs_variable_id = str(logical_expression.rhs_variable.value)
-            lhs_query_input_id = (
-                self.input_variable_ids_by_logical_id[lhs_variable_id]
-                if self.input_variable_ids_by_logical_id
-                else str(uuid4_from_hash(f"{self.node_id}|{hash(tuple(path))}"))
-            )
-            rhs_query_input_id = (
-                self.input_variable_ids_by_logical_id[rhs_variable_id]
-                if self.input_variable_ids_by_logical_id
-                else str(uuid4_from_hash(f"{self.node_id}|{hash(tuple(path))}"))
-            )
+            lhs_variable_id = logical_expression.lhs_variable.value
+            if not isinstance(lhs_variable_id, str):
+                raise TypeError(f"Expected lhs_variable_id to be a string, got {type(lhs_variable_id)}")
+
+            rhs_variable_id = logical_expression.rhs_variable.value
+            if not isinstance(rhs_variable_id, str):
+                raise TypeError(f"Expected rhs_variable_id to be a string, got {type(rhs_variable_id)}")
+
+            lhs_query_input_id: UUID = self.metadata_filter_input_id_by_operand_id.get(
+                UUID(lhs_variable_id)
+            ) or uuid4_from_hash(f"{self.node_id}|{hash(tuple(path))}")
+            rhs_query_input_id: UUID = self.metadata_filter_input_id_by_operand_id.get(
+                UUID(rhs_variable_id)
+            ) or uuid4_from_hash(f"{self.node_id}|{hash(tuple(path))}")
 
             return (
                 {
@@ -173,7 +177,7 @@ class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_Sea
                     create_node_input(
                         self.node_id,
                         f"vellum-query-builder-variable-{lhs_variable_id}",
-                        lhs_query_input_id,
+                        str(lhs_query_input_id),
                         display_context,
                         input_id=UUID(lhs_variable_id),
                         pointer_type=InputVariablePointer,
@@ -181,7 +185,7 @@ class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_Sea
                     create_node_input(
                         self.node_id,
                         f"vellum-query-builder-variable-{rhs_variable_id}",
-                        rhs_query_input_id,
+                        str(rhs_query_input_id),
                         display_context,
                         input_id=UUID(rhs_variable_id),
                         pointer_type=InputVariablePointer,

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_ee/workflows/display/nodes/vellum/templating_node.py

@@ -25,7 +25,7 @@ class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Gene
 
         template_input_id = self.template_input_id or next(
             (
-                input_id
+                UUID(input_id) if isinstance(input_id, str) else input_id
                 for input_name, input_id in self.node_input_ids_by_name.items()
                 if input_name == TEMPLATE_INPUT_NAME
             ),

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_ee/workflows/display/nodes/vellum/utils.py

@@ -1,5 +1,5 @@
 from uuid import UUID
-from typing import Any, List, Optional, Type, cast
+from typing import Any, List, Optional, Type, Union, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
@@ -27,10 +27,10 @@ def create_node_input(
     input_name: str,
     value: Any,
     display_context: WorkflowDisplayContext,
-    input_id: Optional[UUID],
+    input_id: Union[Optional[UUID], Optional[str]],
     pointer_type: Optional[Type[NodeInputValuePointerRule]] = ConstantValuePointer,
 ) -> NodeInput:
-    input_id = input_id or uuid4_from_hash(f"{node_id}|{input_name}")
+    input_id = str(input_id) if input_id else str(uuid4_from_hash(f"{node_id}|{input_name}"))
     if (
         isinstance(value, OutputReference)
         and value.outputs_class._node_class
@@ -42,7 +42,7 @@ def create_node_input(
 
     rules = create_node_input_value_pointer_rules(value, display_context, pointer_type=pointer_type)
     return NodeInput(
-        id=str(input_id),
+        id=input_id,
         key=input_name,
         value=NodeInputValuePointer(
             rules=rules,

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/ee/vellum_ee/workflows/display/utils/vellum.py

@@ -9,7 +9,6 @@ from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.node import NodeReference
 from vellum.workflows.references.vellum_secret import VellumSecretReference
-from vellum.workflows.types.core import VellumValuePrimitive
 from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
 from vellum.workflows.vellum_client import create_vellum_client
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -93,7 +92,7 @@ def create_node_input_value_pointer_rule(
     raise ValueError(f"Unsupported descriptor type: {value.__class__.__name__}")
 
 
-def primitive_to_vellum_value(value: VellumValuePrimitive) -> VellumValue:
+def primitive_to_vellum_value(value: Any) -> VellumValue:
     """Converts a python primitive to a VellumVariableValue"""
 
     if isinstance(value, str):

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vellum-ai"
-version = "0.12.3"
+version = "0.12.5"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/src/vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.12.3",
+            "X-Fern-SDK-Version": "0.12.5",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/src/vellum/client/resources/workflows/client.py

@@ -27,7 +27,11 @@ class WorkflowsClient:
         self,
         id: str,
         *,
+        exclude_code: typing.Optional[bool] = None,
         format: typing.Optional[WorkflowsPullRequestFormat] = None,
+        include_json: typing.Optional[bool] = None,
+        include_sandbox: typing.Optional[bool] = None,
+        strict: typing.Optional[bool] = None,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> typing.Iterator[bytes]:
         """
@@ -38,8 +42,16 @@ class WorkflowsClient:
         id : str
             The ID of the Workflow to pull from
 
+        exclude_code : typing.Optional[bool]
+
         format : typing.Optional[WorkflowsPullRequestFormat]
 
+        include_json : typing.Optional[bool]
+
+        include_sandbox : typing.Optional[bool]
+
+        strict : typing.Optional[bool]
+
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
@@ -53,7 +65,11 @@ class WorkflowsClient:
             base_url=self._client_wrapper.get_environment().default,
             method="GET",
             params={
+                "exclude_code": exclude_code,
                 "format": format,
+                "include_json": include_json,
+                "include_sandbox": include_sandbox,
+                "strict": strict,
             },
             request_options=request_options,
         ) as _response:
@@ -164,7 +180,11 @@ class AsyncWorkflowsClient:
         self,
         id: str,
         *,
+        exclude_code: typing.Optional[bool] = None,
         format: typing.Optional[WorkflowsPullRequestFormat] = None,
+        include_json: typing.Optional[bool] = None,
+        include_sandbox: typing.Optional[bool] = None,
+        strict: typing.Optional[bool] = None,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> typing.AsyncIterator[bytes]:
         """
@@ -175,8 +195,16 @@ class AsyncWorkflowsClient:
         id : str
             The ID of the Workflow to pull from
 
+        exclude_code : typing.Optional[bool]
+
         format : typing.Optional[WorkflowsPullRequestFormat]
 
+        include_json : typing.Optional[bool]
+
+        include_sandbox : typing.Optional[bool]
+
+        strict : typing.Optional[bool]
+
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
@@ -190,7 +218,11 @@ class AsyncWorkflowsClient:
             base_url=self._client_wrapper.get_environment().default,
             method="GET",
             params={
+                "exclude_code": exclude_code,
                 "format": format,
+                "include_json": include_json,
+                "include_sandbox": include_sandbox,
+                "strict": strict,
             },
             request_options=request_options,
         ) as _response:

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/src/vellum/client/types/function_definition.py

@@ -4,22 +4,41 @@ from ..core.pydantic_utilities import UniversalBaseModel
 import typing
 from .prompt_block_state import PromptBlockState
 from .ephemeral_prompt_cache_config import EphemeralPromptCacheConfig
-from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
 
 
 class FunctionDefinition(UniversalBaseModel):
     """
-    A block that represents a function definition in a prompt template.
+    The definition of a Function (aka "Tool Call") that a Prompt/Model has access to.
     """
 
     state: typing.Optional[PromptBlockState] = None
     cache_config: typing.Optional[EphemeralPromptCacheConfig] = None
-    name: typing.Optional[str] = None
-    description: typing.Optional[str] = None
-    parameters: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
-    forced: typing.Optional[bool] = None
-    strict: typing.Optional[bool] = None
+    name: typing.Optional[str] = pydantic.Field(default=None)
+    """
+    The name identifying the function.
+    """
+
+    description: typing.Optional[str] = pydantic.Field(default=None)
+    """
+    A description to help guide the model when to invoke this function.
+    """
+
+    parameters: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = pydantic.Field(default=None)
+    """
+    An OpenAPI specification of parameters that are supported by this function.
+    """
+
+    forced: typing.Optional[bool] = pydantic.Field(default=None)
+    """
+    Set this option to true to force the model to return a function call of this function.
+    """
+
+    strict: typing.Optional[bool] = pydantic.Field(default=None)
+    """
+    Set this option to use strict schema decoding when available.
+    """
 
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2

{vellum_ai-0.12.3 → vellum_ai-0.12.5}/src/vellum/plugins/vellum_mypy.py

@@ -172,41 +172,39 @@ class VellumMypyPlugin(Plugin):
 
     def _dynamic_output_node_class_hook(self, ctx: ClassDefContext, attribute_name: str) -> None:
         """
-        We use this hook to properly annotate the Outputs class for Templating Node using the resolved type
-        of the TemplatingNode's class _OutputType generic.
+        We use this hook to properly annotate the Outputs class for Templating Node and
+        Code Execution Node using the resolved type
+        of the TemplatingNode's and CodeExecutionNode's class _OutputType generic.
         """
-
-        templating_node_info = ctx.cls.info
-        templating_node_bases = ctx.cls.info.bases
-        if not templating_node_bases:
+        node_info = ctx.cls.info
+        node_bases = ctx.cls.info.bases
+        if not node_bases:
             return
-        if not isinstance(templating_node_bases[0], Instance):
+        if not isinstance(node_bases[0], Instance):
             return
 
-        base_templating_args = templating_node_bases[0].args
-        base_templating_node = templating_node_bases[0].type
-        if not _is_subclass(base_templating_node, "vellum.workflows.nodes.core.templating_node.node.TemplatingNode"):
-            return
+        base_args = node_bases[0].args
+        base_node = node_bases[0].type
 
-        if len(base_templating_args) != 2:
+        if len(base_args) != 2:
             return
 
-        base_templating_node_resolved_type = base_templating_args[1]
-        if isinstance(base_templating_node_resolved_type, AnyType):
-            base_templating_node_resolved_type = ctx.api.named_type("builtins.str")
+        base_node_resolved_type = base_args[1]
+        if isinstance(base_node_resolved_type, AnyType):
+            base_node_resolved_type = ctx.api.named_type("builtins.str")
 
-        base_templating_node_outputs = base_templating_node.names.get("Outputs")
-        if not base_templating_node_outputs:
+        base_node_outputs = base_node.names.get("Outputs")
+        if not base_node_outputs:
             return
 
-        current_templating_node_outputs = templating_node_info.names.get("Outputs")
-        if not current_templating_node_outputs:
-            templating_node_info.names["Outputs"] = base_templating_node_outputs.copy()
-            new_outputs_sym = templating_node_info.names["Outputs"].node
+        current_node_outputs = node_info.names.get("Outputs")
+        if not current_node_outputs:
+            node_info.names["Outputs"] = base_node_outputs.copy()
+            new_outputs_sym = node_info.names["Outputs"].node
             if isinstance(new_outputs_sym, TypeInfo):
                 result_sym = new_outputs_sym.names[attribute_name].node
                 if isinstance(result_sym, Var):
-                    result_sym.type = base_templating_node_resolved_type
+                    result_sym.type = base_node_resolved_type
 
     def _base_node_class_hook(self, ctx: ClassDefContext) -> None:
         """
@@ -233,7 +231,6 @@ class VellumMypyPlugin(Plugin):
             type_ = sym.node.type
             if not type_:
                 continue
-
             sym.node.type = self._get_resolvable_type(
                 lambda fullname, types: ctx.api.named_type(fullname, types), type_
             )