vellum-ai 0.14.6__py3-none-any.whl → 0.14.8__py3-none-any.whl

vellum/__init__.py

@@ -74,6 +74,8 @@ from .types import (
     DeploymentReleaseTagDeploymentHistoryItem,
     DeploymentReleaseTagRead,
     DockerServiceToken,
+    DocumentChatMessageContent,
+    DocumentChatMessageContentRequest,
     DocumentDocumentToDocumentIndex,
     DocumentIndexChunking,
     DocumentIndexChunkingRequest,
@@ -81,8 +83,11 @@ from .types import (
     DocumentIndexIndexingConfigRequest,
     DocumentIndexRead,
     DocumentProcessingState,
+    DocumentPromptBlock,
     DocumentRead,
     DocumentStatus,
+    DocumentVellumValue,
+    DocumentVellumValueRequest,
     EnrichedNormalizedCompletion,
     EntityStatus,
     EntityVisibility,
@@ -438,6 +443,8 @@ from .types import (
     VariablePromptBlock,
     VellumAudio,
     VellumAudioRequest,
+    VellumDocument,
+    VellumDocumentRequest,
     VellumError,
     VellumErrorCodeEnum,
     VellumErrorRequest,
@@ -607,6 +614,8 @@ __all__ = [
     "DeploymentReleaseTagRead",
     "DeploymentsListRequestStatus",
     "DockerServiceToken",
+    "DocumentChatMessageContent",
+    "DocumentChatMessageContentRequest",
     "DocumentDocumentToDocumentIndex",
     "DocumentIndexChunking",
     "DocumentIndexChunkingRequest",
@@ -615,8 +624,11 @@ __all__ = [
     "DocumentIndexRead",
     "DocumentIndexesListRequestStatus",
     "DocumentProcessingState",
+    "DocumentPromptBlock",
     "DocumentRead",
     "DocumentStatus",
+    "DocumentVellumValue",
+    "DocumentVellumValueRequest",
     "EnrichedNormalizedCompletion",
     "EntityStatus",
     "EntityVisibility",
@@ -980,6 +992,8 @@ __all__ = [
     "Vellum",
     "VellumAudio",
     "VellumAudioRequest",
+    "VellumDocument",
+    "VellumDocumentRequest",
     "VellumEnvironment",
     "VellumError",
     "VellumErrorCodeEnum",

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.14.6",
+            "X-Fern-SDK-Version": "0.14.8",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/client/types/__init__.py

@@ -78,6 +78,8 @@ from .deployment_read import DeploymentRead
 from .deployment_release_tag_deployment_history_item import DeploymentReleaseTagDeploymentHistoryItem
 from .deployment_release_tag_read import DeploymentReleaseTagRead
 from .docker_service_token import DockerServiceToken
+from .document_chat_message_content import DocumentChatMessageContent
+from .document_chat_message_content_request import DocumentChatMessageContentRequest
 from .document_document_to_document_index import DocumentDocumentToDocumentIndex
 from .document_index_chunking import DocumentIndexChunking
 from .document_index_chunking_request import DocumentIndexChunkingRequest
@@ -85,8 +87,11 @@ from .document_index_indexing_config import DocumentIndexIndexingConfig
 from .document_index_indexing_config_request import DocumentIndexIndexingConfigRequest
 from .document_index_read import DocumentIndexRead
 from .document_processing_state import DocumentProcessingState
+from .document_prompt_block import DocumentPromptBlock
 from .document_read import DocumentRead
 from .document_status import DocumentStatus
+from .document_vellum_value import DocumentVellumValue
+from .document_vellum_value_request import DocumentVellumValueRequest
 from .enriched_normalized_completion import EnrichedNormalizedCompletion
 from .entity_status import EntityStatus
 from .entity_visibility import EntityVisibility
@@ -462,6 +467,8 @@ from .upsert_test_suite_test_case_request import UpsertTestSuiteTestCaseRequest
 from .variable_prompt_block import VariablePromptBlock
 from .vellum_audio import VellumAudio
 from .vellum_audio_request import VellumAudioRequest
+from .vellum_document import VellumDocument
+from .vellum_document_request import VellumDocumentRequest
 from .vellum_error import VellumError
 from .vellum_error_code_enum import VellumErrorCodeEnum
 from .vellum_error_request import VellumErrorRequest
@@ -596,6 +603,8 @@ __all__ = [
     "DeploymentReleaseTagDeploymentHistoryItem",
     "DeploymentReleaseTagRead",
     "DockerServiceToken",
+    "DocumentChatMessageContent",
+    "DocumentChatMessageContentRequest",
     "DocumentDocumentToDocumentIndex",
     "DocumentIndexChunking",
     "DocumentIndexChunkingRequest",
@@ -603,8 +612,11 @@ __all__ = [
     "DocumentIndexIndexingConfigRequest",
     "DocumentIndexRead",
     "DocumentProcessingState",
+    "DocumentPromptBlock",
     "DocumentRead",
     "DocumentStatus",
+    "DocumentVellumValue",
+    "DocumentVellumValueRequest",
     "EnrichedNormalizedCompletion",
     "EntityStatus",
     "EntityVisibility",
@@ -960,6 +972,8 @@ __all__ = [
     "VariablePromptBlock",
     "VellumAudio",
     "VellumAudioRequest",
+    "VellumDocument",
+    "VellumDocumentRequest",
     "VellumError",
     "VellumErrorCodeEnum",
     "VellumErrorRequest",

vellum/client/types/array_chat_message_content_item.py

@@ -5,7 +5,12 @@ from .string_chat_message_content import StringChatMessageContent
 from .function_call_chat_message_content import FunctionCallChatMessageContent
 from .image_chat_message_content import ImageChatMessageContent
 from .audio_chat_message_content import AudioChatMessageContent
+from .document_chat_message_content import DocumentChatMessageContent
 
 ArrayChatMessageContentItem = typing.Union[
-    StringChatMessageContent, FunctionCallChatMessageContent, ImageChatMessageContent, AudioChatMessageContent
+    StringChatMessageContent,
+    FunctionCallChatMessageContent,
+    ImageChatMessageContent,
+    AudioChatMessageContent,
+    DocumentChatMessageContent,
 ]

vellum/client/types/array_chat_message_content_item_request.py

@@ -5,10 +5,12 @@ from .string_chat_message_content_request import StringChatMessageContentRequest
 from .function_call_chat_message_content_request import FunctionCallChatMessageContentRequest
 from .image_chat_message_content_request import ImageChatMessageContentRequest
 from .audio_chat_message_content_request import AudioChatMessageContentRequest
+from .document_chat_message_content_request import DocumentChatMessageContentRequest
 
 ArrayChatMessageContentItemRequest = typing.Union[
     StringChatMessageContentRequest,
     FunctionCallChatMessageContentRequest,
     ImageChatMessageContentRequest,
     AudioChatMessageContentRequest,
+    DocumentChatMessageContentRequest,
 ]

vellum/client/types/chat_message_content.py

@@ -6,6 +6,7 @@ from .function_call_chat_message_content import FunctionCallChatMessageContent
 from .array_chat_message_content import ArrayChatMessageContent
 from .image_chat_message_content import ImageChatMessageContent
 from .audio_chat_message_content import AudioChatMessageContent
+from .document_chat_message_content import DocumentChatMessageContent
 
 ChatMessageContent = typing.Union[
     StringChatMessageContent,
@@ -13,4 +14,5 @@ ChatMessageContent = typing.Union[
     ArrayChatMessageContent,
     ImageChatMessageContent,
     AudioChatMessageContent,
+    DocumentChatMessageContent,
 ]

vellum/client/types/chat_message_content_request.py

@@ -6,6 +6,7 @@ from .function_call_chat_message_content_request import FunctionCallChatMessageC
 from .array_chat_message_content_request import ArrayChatMessageContentRequest
 from .image_chat_message_content_request import ImageChatMessageContentRequest
 from .audio_chat_message_content_request import AudioChatMessageContentRequest
+from .document_chat_message_content_request import DocumentChatMessageContentRequest
 
 ChatMessageContentRequest = typing.Union[
     StringChatMessageContentRequest,
@@ -13,4 +14,5 @@ ChatMessageContentRequest = typing.Union[
     ArrayChatMessageContentRequest,
     ImageChatMessageContentRequest,
     AudioChatMessageContentRequest,
+    DocumentChatMessageContentRequest,
 ]

vellum/client/types/document_chat_message_content.py

@@ -0,0 +1,25 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from .vellum_document import VellumDocument
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+
+
+class DocumentChatMessageContent(UniversalBaseModel):
+    """
+    A document value that is used in a chat message.
+    """
+
+    type: typing.Literal["DOCUMENT"] = "DOCUMENT"
+    value: VellumDocument
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/document_chat_message_content_request.py

@@ -0,0 +1,25 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from .vellum_document_request import VellumDocumentRequest
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+
+
+class DocumentChatMessageContentRequest(UniversalBaseModel):
+    """
+    A document value that is used in a chat message.
+    """
+
+    type: typing.Literal["DOCUMENT"] = "DOCUMENT"
+    value: VellumDocumentRequest
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/document_prompt_block.py

@@ -0,0 +1,29 @@
+# This file was auto-generated by Fern from our API Definition.
+
+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
+
+
+class DocumentPromptBlock(UniversalBaseModel):
+    """
+    A block that represents a document in a prompt template.
+    """
+
+    block_type: typing.Literal["DOCUMENT"] = "DOCUMENT"
+    state: typing.Optional[PromptBlockState] = None
+    cache_config: typing.Optional[EphemeralPromptCacheConfig] = None
+    src: str
+    metadata: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/document_vellum_value.py

@@ -0,0 +1,25 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from .vellum_document import VellumDocument
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+
+
+class DocumentVellumValue(UniversalBaseModel):
+    """
+    A base Vellum primitive value representing a document.
+    """
+
+    type: typing.Literal["DOCUMENT"] = "DOCUMENT"
+    value: typing.Optional[VellumDocument] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/document_vellum_value_request.py

@@ -0,0 +1,25 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from .vellum_document_request import VellumDocumentRequest
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+
+
+class DocumentVellumValueRequest(UniversalBaseModel):
+    """
+    A base Vellum primitive value representing a document.
+    """
+
+    type: typing.Literal["DOCUMENT"] = "DOCUMENT"
+    value: typing.Optional[VellumDocumentRequest] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/prompt_block.py

@@ -8,6 +8,7 @@ from .rich_text_prompt_block import RichTextPromptBlock
 from .audio_prompt_block import AudioPromptBlock
 from .function_call_prompt_block import FunctionCallPromptBlock
 from .image_prompt_block import ImagePromptBlock
+from .document_prompt_block import DocumentPromptBlock
 import typing
 
 if typing.TYPE_CHECKING:
@@ -20,4 +21,5 @@ PromptBlock = typing.Union[
     AudioPromptBlock,
     FunctionCallPromptBlock,
     ImagePromptBlock,
+    DocumentPromptBlock,
 ]

vellum/client/types/vellum_document.py

@@ -0,0 +1,20 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+
+
+class VellumDocument(UniversalBaseModel):
+    src: str
+    metadata: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/vellum_document_request.py

@@ -0,0 +1,20 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from ..core.pydantic_utilities import UniversalBaseModel
+import typing
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+
+
+class VellumDocumentRequest(UniversalBaseModel):
+    src: str
+    metadata: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
+
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
+
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

vellum/client/types/vellum_value.py

@@ -7,6 +7,7 @@ from .number_vellum_value import NumberVellumValue
 from .json_vellum_value import JsonVellumValue
 from .image_vellum_value import ImageVellumValue
 from .audio_vellum_value import AudioVellumValue
+from .document_vellum_value import DocumentVellumValue
 from .function_call_vellum_value import FunctionCallVellumValue
 from .error_vellum_value import ErrorVellumValue
 from .chat_history_vellum_value import ChatHistoryVellumValue
@@ -21,6 +22,7 @@ VellumValue = typing.Union[
     JsonVellumValue,
     ImageVellumValue,
     AudioVellumValue,
+    DocumentVellumValue,
     FunctionCallVellumValue,
     ErrorVellumValue,
     "ArrayVellumValue",

vellum/client/types/vellum_value_request.py

@@ -7,6 +7,7 @@ from .number_vellum_value_request import NumberVellumValueRequest
 from .json_vellum_value_request import JsonVellumValueRequest
 from .image_vellum_value_request import ImageVellumValueRequest
 from .audio_vellum_value_request import AudioVellumValueRequest
+from .document_vellum_value_request import DocumentVellumValueRequest
 from .function_call_vellum_value_request import FunctionCallVellumValueRequest
 from .error_vellum_value_request import ErrorVellumValueRequest
 from .chat_history_vellum_value_request import ChatHistoryVellumValueRequest
@@ -21,6 +22,7 @@ VellumValueRequest = typing.Union[
     JsonVellumValueRequest,
     ImageVellumValueRequest,
     AudioVellumValueRequest,
+    DocumentVellumValueRequest,
     FunctionCallVellumValueRequest,
     ErrorVellumValueRequest,
     "ArrayVellumValueRequest",

vellum/client/types/vellum_variable_type.py

@@ -14,6 +14,7 @@ VellumVariableType = typing.Union[
         "FUNCTION_CALL",
         "IMAGE",
         "AUDIO",
+        "DOCUMENT",
         "NULL",
     ],
     typing.Any,

vellum/types/document_chat_message_content.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.document_chat_message_content import *

vellum/types/document_chat_message_content_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.document_chat_message_content_request import *

vellum/types/document_prompt_block.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.document_prompt_block import *

vellum/types/document_vellum_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.document_vellum_value import *

vellum/types/document_vellum_value_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.document_vellum_value_request import *

vellum/types/vellum_document.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.vellum_document import *

vellum/types/vellum_document_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.vellum_document_request import *

vellum/workflows/exceptions.py

@@ -17,3 +17,21 @@ class NodeException(Exception):
     @staticmethod
     def of(workflow_error: WorkflowError) -> "NodeException":
         return NodeException(message=workflow_error.message, code=workflow_error.code)
+
+
+class WorkflowInitializationException(Exception):
+    def __init__(self, message: str, code: WorkflowErrorCode = WorkflowErrorCode.INVALID_INPUTS):
+        self.message = message
+        self.code = code
+        super().__init__(message)
+
+    @property
+    def error(self) -> WorkflowError:
+        return WorkflowError(
+            message=self.message,
+            code=self.code,
+        )
+
+    @staticmethod
+    def of(workflow_error: WorkflowError) -> "WorkflowInitializationException":
+        return WorkflowInitializationException(message=workflow_error.message, code=workflow_error.code)

vellum/workflows/inputs/base.py

@@ -1,9 +1,11 @@
-from typing import Any, Iterator, Tuple, Type
+from typing import Any, Iterator, Tuple, Type, Union, get_args, get_origin
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler
 from pydantic_core import core_schema
 
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import WorkflowInitializationException
 from vellum.workflows.references import ExternalInputReference, WorkflowInputReference
 from vellum.workflows.references.input import InputReference
 from vellum.workflows.types.utils import get_class_attr_names, infer_types
@@ -40,7 +42,20 @@ class BaseInputs(metaclass=_BaseInputsMeta):
     __parent_class__: Type = type(None)
 
     def __init__(self, **kwargs: Any) -> None:
+        for name, field_type in self.__class__.__annotations__.items():
+            if name not in kwargs and name not in vars(self.__class__):
+                origin = get_origin(field_type)
+                args = get_args(field_type)
+                if not (origin is Union and type(None) in args):
+                    raise WorkflowInitializationException(
+                        message="Required input variables should have defined value",
+                        code=WorkflowErrorCode.INVALID_INPUTS,
+                    )
+
         for name, value in kwargs.items():
+            field_type = self.__class__.__annotations__.get(name)
+            if field_type:
+                self._validate_input(value, field_type)
             setattr(self, name, value)
 
     def __iter__(self) -> Iterator[Tuple[InputReference, Any]]:
@@ -48,6 +63,17 @@ class BaseInputs(metaclass=_BaseInputsMeta):
             if hasattr(self, input_descriptor.name):
                 yield (input_descriptor, getattr(self, input_descriptor.name))
 
+    def _validate_input(self, value: Any, field_type: Any) -> None:
+        if value is None:
+            # Check if field_type is Optional
+            origin = get_origin(field_type)
+            args = get_args(field_type)
+            if not (origin is Union and type(None) in args):
+                raise WorkflowInitializationException(
+                    message="Required input variables should have defined value",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                )
+
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Type[Any], handler: GetCoreSchemaHandler

vellum/workflows/inputs/tests/test_inputs.py

@@ -0,0 +1,49 @@
+import pytest
+from typing import Optional
+
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.exceptions import WorkflowInitializationException
+from vellum.workflows.inputs import BaseInputs
+
+
+def test_base_inputs_happy_path():
+    # GIVEN some input class with required and optional fields
+    class TestInputs(BaseInputs):
+        required_string: str
+        required_int: int
+        optional_string: Optional[str]
+
+    # WHEN we assign the inputs some valid values
+    inputs = TestInputs(required_string="hello", required_int=42, optional_string=None)
+
+    # THEN the inputs should have the correct values
+    assert inputs.required_string == "hello"
+    assert inputs.required_int == 42
+    assert inputs.optional_string is None
+
+
+def test_base_inputs_empty_value():
+    # GIVEN some input class with required and optional string fields
+    class TestInputs(BaseInputs):
+        required_string: str
+        optional_string: Optional[str]
+
+    # WHEN we try to omit a required field
+    with pytest.raises(WorkflowInitializationException) as exc_info:
+        TestInputs(optional_string="ok")  # type: ignore
+
+    # THEN it should raise a NodeException with the correct error message and code
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "Required input variables should have defined value" in str(exc_info.value)
+
+
+def test_base_inputs_with_default():
+    # GIVEN some input class with a field that has a default value
+    class TestInputs(BaseInputs):
+        string_with_default: str = "default_value"
+
+    # WHEN we create an instance without providing the field
+    inputs = TestInputs()
+
+    # THEN it should use the default value
+    assert inputs.string_with_default == "default_value"

vellum/workflows/nodes/core/inline_subworkflow_node/node.py

@@ -70,7 +70,7 @@ class InlineSubworkflowNode(
     subworkflow_inputs: ClassVar[Union[EntityInputsInterface, BaseInputs, Type[undefined]]] = undefined
 
     def run(self) -> Iterator[BaseOutput]:
-        with execution_context(parent_context=get_parent_context() or self._context.parent_context):
+        with execution_context(parent_context=get_parent_context()):
             subworkflow = self.subworkflow(
                 parent_state=self.state,
                 context=WorkflowContext(vellum_client=self._context.vellum_client),

vellum/workflows/nodes/core/map_node/node.py

@@ -16,10 +16,9 @@ from typing import (
     overload,
 )
 
-from vellum.workflows.context import execution_context, get_parent_context
+from vellum.workflows.context import ExecutionContext, execution_context, get_execution_context
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode
-from vellum.workflows.events.types import ParentContext
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
@@ -76,13 +75,13 @@ class MapNode(BaseAdornmentNode[StateType], Generic[StateType, MapNodeItemType])
         fulfilled_iterations: List[bool] = []
         for index, item in enumerate(self.items):
             fulfilled_iterations.append(False)
-            parent_context = get_parent_context() or self._context.parent_context
+            current_execution_context = get_execution_context()
             thread = Thread(
                 target=self._context_run_subworkflow,
                 kwargs={
                     "item": item,
                     "index": index,
-                    "parent_context": parent_context,
+                    "current_execution_context": current_execution_context,
                 },
             )
             if self.max_concurrency is None:
@@ -143,10 +142,11 @@ class MapNode(BaseAdornmentNode[StateType], Generic[StateType, MapNodeItemType])
             yield BaseOutput(name=output_name, value=output_list)
 
     def _context_run_subworkflow(
-        self, *, item: MapNodeItemType, index: int, parent_context: Optional[ParentContext] = None
+        self, *, item: MapNodeItemType, index: int, current_execution_context: ExecutionContext
     ) -> None:
-        parent_context = parent_context or self._context.parent_context
-        with execution_context(parent_context=parent_context):
+        parent_context = current_execution_context.parent_context
+        trace_id = current_execution_context.trace_id
+        with execution_context(parent_context=parent_context, trace_id=trace_id):
             self._run_subworkflow(item=item, index=index)
 
     def _run_subworkflow(self, *, item: MapNodeItemType, index: int) -> None:

vellum/workflows/nodes/core/try_node/node.py

@@ -27,7 +27,7 @@ class TryNode(BaseAdornmentNode[StateType], Generic[StateType]):
         error: Optional[WorkflowError] = None
 
     def run(self) -> Iterator[BaseOutput]:
-        parent_context = get_parent_context() or self._context.parent_context
+        parent_context = get_parent_context()
         with execution_context(parent_context=parent_context):
             subworkflow = self.subworkflow(
                 parent_state=self.state,