vellum-ai 0.14.41__py3-none-any.whl → 1.11.2__py3-none-any.whl

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

@@ -3,19 +3,31 @@ from uuid import UUID
 from typing import Any, ClassVar, Dict, Generator, Generic, Iterator, List, Optional, Sequence, Set, Union
 
 from vellum import (
+    AudioInputRequest,
     ChatHistoryInputRequest,
     ChatMessage,
+    DocumentInputRequest,
     ExecutePromptEvent,
+    ImageInputRequest,
     JsonInputRequest,
     PromptDeploymentExpandMetaRequest,
     PromptDeploymentInputRequest,
     PromptOutput,
     RawPromptExecutionOverridesRequest,
     StringInputRequest,
+    VellumAudio,
+    VellumAudioRequest,
+    VellumDocument,
+    VellumDocumentRequest,
+    VellumImage,
+    VellumImageRequest,
+    VellumVideo,
+    VellumVideoRequest,
+    VideoInputRequest,
 )
 from vellum.client import ApiError, RequestOptions
 from vellum.client.types.chat_message_request import ChatMessageRequest
-from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
+from vellum.workflows.constants import LATEST_RELEASE_TAG
 from vellum.workflows.context import get_execution_context
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.errors.types import vellum_error_to_workflow_error
@@ -48,13 +60,13 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
     deployment: ClassVar[Union[UUID, str]]
 
     release_tag: str = LATEST_RELEASE_TAG
-    external_id: Optional[str] = OMIT
+    external_id: Optional[str] = None
 
-    expand_meta: Optional[PromptDeploymentExpandMetaRequest] = OMIT
-    raw_overrides: Optional[RawPromptExecutionOverridesRequest] = OMIT
-    expand_raw: Optional[Sequence[str]] = OMIT
-    metadata: Optional[Dict[str, Optional[Any]]] = OMIT
-    ml_model_fallbacks: Optional[Sequence[str]] = OMIT
+    expand_meta: Optional[PromptDeploymentExpandMetaRequest] = PromptDeploymentExpandMetaRequest(finish_reason=True)
+    raw_overrides: Optional[RawPromptExecutionOverridesRequest] = None
+    expand_raw: Optional[Sequence[str]] = None
+    metadata: Optional[Dict[str, Optional[Any]]] = None
+    ml_model_fallbacks: Optional[Sequence[str]] = None
 
     class Trigger(BasePromptNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
@@ -103,12 +115,7 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
                 prompt_event_stream = self._get_prompt_event_stream()
                 next(prompt_event_stream)
             except ApiError as e:
-                if (
-                    e.status_code
-                    and e.status_code < 500
-                    and self.ml_model_fallbacks is not OMIT
-                    and self.ml_model_fallbacks is not None
-                ):
+                if e.status_code and e.status_code < 500 and self.ml_model_fallbacks is not None:
                     prompt_event_stream = self._retry_prompt_stream_with_fallbacks(tried_fallbacks)
                 else:
                     self._handle_api_error(e)
@@ -127,7 +134,6 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
                     if (
                         event.error
                         and event.error.code == WorkflowErrorCode.PROVIDER_ERROR.value
-                        and self.ml_model_fallbacks is not OMIT
                         and self.ml_model_fallbacks is not None
                     ):
                         try:
@@ -177,6 +183,10 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
             return compiled_inputs
 
         for input_name, input_value in self.prompt_inputs.items():
+            # Exclude inputs that resolved to be null. This ensure that we don't pass input values
+            # to optional prompt inputs whose values were unresolved.
+            if input_value is None:
+                continue
             if isinstance(input_value, str):
                 compiled_inputs.append(
                     StringInputRequest(
@@ -204,6 +214,66 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
                         value=chat_history,
                     )
                 )
+            elif isinstance(input_value, (VellumAudio, VellumAudioRequest)):
+                audio_value = (
+                    input_value
+                    if isinstance(input_value, VellumAudioRequest)
+                    else VellumAudioRequest.model_validate(input_value.model_dump())
+                )
+                compiled_inputs.append(
+                    AudioInputRequest(
+                        name=input_name,
+                        value=audio_value,
+                    )
+                )
+            elif isinstance(input_value, (VellumImage, VellumImageRequest)):
+                image_value = (
+                    input_value
+                    if isinstance(input_value, VellumImageRequest)
+                    else VellumImageRequest.model_validate(input_value.model_dump())
+                )
+                compiled_inputs.append(
+                    ImageInputRequest(
+                        name=input_name,
+                        value=image_value,
+                    )
+                )
+            elif isinstance(input_value, (VellumDocument, VellumDocumentRequest)):
+                document_value = (
+                    input_value
+                    if isinstance(input_value, VellumDocumentRequest)
+                    else VellumDocumentRequest.model_validate(input_value.model_dump())
+                )
+                compiled_inputs.append(
+                    DocumentInputRequest(
+                        name=input_name,
+                        value=document_value,
+                    )
+                )
+            elif (
+                isinstance(input_value, dict)
+                and "src" in input_value
+                and isinstance(input_value.get("src"), str)
+                and input_value["src"].endswith(".pdf")
+            ):
+                compiled_inputs.append(
+                    DocumentInputRequest(
+                        name=input_name,
+                        value=VellumDocumentRequest.model_validate(input_value),
+                    )
+                )
+            elif isinstance(input_value, (VellumVideo, VellumVideoRequest)):
+                video_value = (
+                    input_value
+                    if isinstance(input_value, VellumVideoRequest)
+                    else VellumVideoRequest.model_validate(input_value.model_dump())
+                )
+                compiled_inputs.append(
+                    VideoInputRequest(
+                        name=input_name,
+                        value=video_value,
+                    )
+                )
             else:
                 try:
                     input_value = default_serializer(input_value)

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

@@ -2,6 +2,8 @@ from decimal import Decimal
 from uuid import UUID
 from typing import ClassVar, Generic, List, Optional, Union
 
+import httpx
+
 from vellum import (
     NotFoundError,
     SearchFiltersRequest,
@@ -11,7 +13,7 @@ from vellum import (
     SearchResultMergingRequest,
     SearchWeightsRequest,
 )
-from vellum.core import ApiError, RequestOptions
+from vellum.client.core import ApiError, RequestOptions
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
@@ -106,11 +108,32 @@ class BaseSearchNode(BaseNode[StateType], Generic[StateType]):
                 message=f"Document Index '{self.document_index}' not found",
                 code=WorkflowErrorCode.INVALID_INPUTS,
             )
-        except ApiError:
+        except httpx.TransportError:
             raise NodeException(
-                message=f"An error occurred while searching against Document Index '{self.document_index}'",  # noqa: E501
+                message="Failed to connect to Vellum server",
                 code=WorkflowErrorCode.INTERNAL_ERROR,
             )
+        except ApiError as e:
+            raw_data = e.body if isinstance(e.body, dict) else None
+            if e.status_code and e.status_code == 403 and isinstance(e.body, dict):
+                raise NodeException(
+                    message=e.body.get("detail", "Provider credentials is missing or unavailable"),
+                    code=WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
+                    raw_data=raw_data,
+                )
+            elif e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
+                raise NodeException(
+                    message=e.body.get(
+                        "detail", f"An error occurred while searching against Document Index '{self.document_index}'"
+                    ),
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                    raw_data=raw_data,
+                ) from e
+            raise NodeException(
+                message=f"An error occurred while searching against Document Index '{self.document_index}'",
+                code=WorkflowErrorCode.INTERNAL_ERROR,
+                raw_data=raw_data,
+            ) from e
 
     def _get_options_request(self) -> SearchRequestOptionsRequest:
         return SearchRequestOptionsRequest(

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

@@ -0,0 +1,90 @@
+import pytest
+
+from vellum import (
+    AudioInputRequest,
+    DocumentInputRequest,
+    ImageInputRequest,
+    VellumAudio,
+    VellumAudioRequest,
+    VellumDocument,
+    VellumDocumentRequest,
+    VellumImage,
+    VellumImageRequest,
+    VellumVideo,
+    VellumVideoRequest,
+    VideoInputRequest,
+)
+from vellum.workflows.nodes import PromptDeploymentNode
+
+
+@pytest.mark.parametrize(
+    [
+        "raw_input",
+        "expected_compiled_inputs",
+    ],
+    [
+        (
+            VellumAudio(src="data:audio/wav;base64,mockaudio"),
+            [AudioInputRequest(name="file_input", value=VellumAudioRequest(src="data:audio/wav;base64,mockaudio"))],
+        ),
+        (
+            VellumImage(src="data:image/png;base64,mockimage"),
+            [ImageInputRequest(name="file_input", value=VellumImageRequest(src="data:image/png;base64,mockimage"))],
+        ),
+        (
+            VellumVideo(src="data:video/mp4;base64,mockvideo"),
+            [VideoInputRequest(name="file_input", value=VellumVideoRequest(src="data:video/mp4;base64,mockvideo"))],
+        ),
+        (
+            VellumDocument(src="mockdocument"),
+            [DocumentInputRequest(name="file_input", value=VellumDocumentRequest(src="mockdocument"))],
+        ),
+        (
+            VellumAudioRequest(src="data:audio/wav;base64,mockaudio"),
+            [AudioInputRequest(name="file_input", value=VellumAudioRequest(src="data:audio/wav;base64,mockaudio"))],
+        ),
+        (
+            VellumImageRequest(src="data:image/png;base64,mockimage"),
+            [ImageInputRequest(name="file_input", value=VellumImageRequest(src="data:image/png;base64,mockimage"))],
+        ),
+        (
+            VellumVideoRequest(src="data:video/mp4;base64,mockvideo"),
+            [VideoInputRequest(name="file_input", value=VellumVideoRequest(src="data:video/mp4;base64,mockvideo"))],
+        ),
+        (
+            VellumDocumentRequest(src="mockdocument"),
+            [DocumentInputRequest(name="file_input", value=VellumDocumentRequest(src="mockdocument"))],
+        ),
+        (
+            {"src": "https://example.com/document.pdf"},
+            [
+                DocumentInputRequest(
+                    name="file_input", value=VellumDocumentRequest(src="https://example.com/document.pdf")
+                )
+            ],
+        ),
+        (
+            {"src": "https://example.com/document.pdf", "metadata": {"author": "test"}},
+            [
+                DocumentInputRequest(
+                    name="file_input",
+                    value=VellumDocumentRequest(src="https://example.com/document.pdf", metadata={"author": "test"}),
+                )
+            ],
+        ),
+    ],
+)
+def test_file_input_compilation(raw_input, expected_compiled_inputs):
+    """
+    Tests that file inputs are correctly compiled to the appropriate input type.
+    """
+
+    class MyPromptDeploymentNode(PromptDeploymentNode):
+        deployment = "test-deployment"
+        prompt_inputs = {"file_input": raw_input}
+
+    # WHEN we compile the inputs
+    compiled_inputs = MyPromptDeploymentNode()._compile_prompt_inputs()
+
+    # THEN we should get the correct input type
+    assert compiled_inputs == expected_compiled_inputs

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

@@ -1,7 +1,9 @@
 import enum
 import json
-from typing import Any, List, Union, cast
+import re
+from typing import Any, List, Optional, Tuple, Union, cast
 
+from vellum import PromptOutput
 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
@@ -28,15 +30,18 @@ from vellum.client.types.string_vellum_value_request import StringVellumValueReq
 from vellum.client.types.vellum_error import VellumError
 from vellum.client.types.vellum_value import VellumValue
 from vellum.client.types.vellum_value_request import VellumValueRequest
+from vellum.client.types.video_vellum_value import VideoVellumValue
+from vellum.client.types.video_vellum_value_request import VideoVellumValueRequest
+from vellum.utils.json_encoder import VellumJsonEncoder
 from vellum.workflows.errors.types import WorkflowError, workflow_error_to_vellum_error
-from vellum.workflows.state.encoder import DefaultStateEncoder
 
 VELLUM_VALUE_REQUEST_TUPLE = (
     StringVellumValueRequest,
     NumberVellumValueRequest,
     JsonVellumValueRequest,
-    ImageVellumValueRequest,
     AudioVellumValueRequest,
+    VideoVellumValueRequest,
+    ImageVellumValueRequest,
     FunctionCallVellumValueRequest,
     ErrorVellumValueRequest,
     ArrayVellumValueRequest,
@@ -52,6 +57,8 @@ def primitive_to_vellum_value(value: Any) -> VellumValue:
         return StringVellumValue(value=value)
     elif isinstance(value, enum.Enum):
         return StringVellumValue(value=value.value)
+    elif isinstance(value, bool):
+        return JsonVellumValue(value=value)
     elif isinstance(value, (int, float)):
         return NumberVellumValue(value=value)
     elif isinstance(value, list) and (
@@ -76,8 +83,9 @@ def primitive_to_vellum_value(value: Any) -> VellumValue:
             StringVellumValue,
             NumberVellumValue,
             JsonVellumValue,
-            ImageVellumValue,
             AudioVellumValue,
+            VideoVellumValue,
+            ImageVellumValue,
             FunctionCallVellumValue,
             ErrorVellumValue,
             ArrayVellumValue,
@@ -95,7 +103,7 @@ def primitive_to_vellum_value(value: Any) -> VellumValue:
         return value  # type: ignore
 
     try:
-        json_value = json.dumps(value, cls=DefaultStateEncoder)
+        json_value = json.dumps(value, cls=VellumJsonEncoder)
     except json.JSONDecodeError:
         raise ValueError(f"Unsupported variable type: {value.__class__.__name__}")
 
@@ -117,3 +125,72 @@ def primitive_to_vellum_value_request(value: Any) -> VellumValueRequest:
         raise ValueError(f"Unsupported variable type: {vellum_value.__class__.__name__}")
 
     return vellum_value_request_class.model_validate(vellum_value.model_dump())
+
+
+def _strip_markdown_json(text: str) -> str:
+    """
+    Strip markdown code block formatting from JSON text.
+
+    Handles cases like:
+    ```json
+    {"key": "value"}
+    ```
+
+    Or with surrounding text:
+    Here is the json you asked for!
+
+    ```json
+    {"key": "value"}
+    ```
+
+    Args:
+        text: The text potentially containing markdown-formatted JSON
+
+    Returns:
+        The stripped JSON text, or the original text if no markdown formatting found
+    """
+    pattern = r"```(?:json)?\s*\n(.*?)\n```"
+
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+
+    return text
+
+
+def process_additional_prompt_outputs(outputs: List[PromptOutput]) -> Tuple[str, Optional[Any]]:
+    """
+    Process prompt outputs using the same logic as prompt nodes to determine text output.
+
+    Args:
+        outputs: List of prompt outputs to process
+
+    Returns:
+        The text representation of the outputs joined with newlines
+        The JSON representation of the outputs
+    """
+    string_outputs = []
+    json_output = None
+    for output in outputs:
+        if output.value is None:
+            continue
+
+        if output.type == "STRING":
+            string_outputs.append(output.value)
+            for value_to_parse in [output.value, _strip_markdown_json(output.value)]:
+                try:
+                    json_output = json.loads(value_to_parse)
+                    break
+                except (json.JSONDecodeError, TypeError):
+                    continue
+        elif output.type == "JSON":
+            string_outputs.append(json.dumps(output.value, indent=4))
+            json_output = output.value
+        elif output.type == "FUNCTION_CALL":
+            string_outputs.append(output.value.model_dump_json(indent=4))
+        elif output.type == "THINKING":
+            continue
+        else:
+            string_outputs.append(output.value.message)
+
+    return "\n".join(string_outputs), json_output

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

@@ -1,4 +1,5 @@
 import inspect
+import os
 from typing import Any, ClassVar, Dict, Generic, List, Optional, Sequence, Tuple, Type, TypeVar, cast, get_args
 
 from vellum import (
@@ -19,13 +20,15 @@ from vellum import (
     VellumError,
     VellumValue,
 )
+from vellum.client.core import RequestOptions
 from vellum.client.core.api_error import ApiError
 from vellum.client.types.code_executor_secret_input import CodeExecutorSecretInput
-from vellum.core import RequestOptions
+from vellum.workflows.constants import undefined
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path, run_code_inline
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior, VellumSecret
@@ -75,6 +78,10 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
     request_options: Optional[RequestOptions] = None - The request options to use for the custom script.
     """
 
+    class Display(BaseNode.Display):
+        icon = "vellum:icon:rectangle-code"
+        color = "lime"
+
     filepath: ClassVar[Optional[str]] = None
     code: ClassVar[Optional[str]] = None
 
@@ -95,9 +102,9 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
 
     def run(self) -> Outputs:
         output_type = self.__class__.get_output_type()
-        code = self._resolve_code()
-        if not self.packages and self.runtime == "PYTHON_3_11_6":
-            logs, result = run_code_inline(code, self.code_inputs, output_type)
+        code, filepath = self._resolve_code()
+        if not self.packages and self.runtime == "PYTHON_3_11_6" and not self._has_secrets_in_code_inputs():
+            logs, result = run_code_inline(code, self.code_inputs, output_type, filepath)
             return self.Outputs(result=result, log=logs)
 
         else:
@@ -114,13 +121,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
                     request_options=self.request_options,
                 )
             except ApiError as e:
-                if e.status_code == 400 and isinstance(e.body, dict) and "message" in e.body:
-                    raise NodeException(
-                        message=e.body["message"],
-                        code=WorkflowErrorCode.INVALID_INPUTS,
-                    )
-
-                raise
+                self._handle_api_error(e)
 
             if code_execution_result.output.type != expected_output_type:
                 actual_type = code_execution_result.output.type
@@ -131,6 +132,31 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
 
             return self.Outputs(result=code_execution_result.output.value, log=code_execution_result.log)
 
+    def _handle_api_error(self, e: ApiError) -> None:
+        if e.status_code and e.status_code == 403 and isinstance(e.body, dict):
+            raise NodeException(
+                message=e.body.get("detail", "Provider credentials is missing or unavailable"),
+                code=WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
+            ) from e
+
+        if e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
+            raise NodeException(
+                message=e.body.get("message", e.body.get("detail", "Failed to execute code")),
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            ) from e
+
+        raise NodeException(
+            message="Failed to execute code",
+            code=WorkflowErrorCode.INTERNAL_ERROR,
+        ) from e
+
+    def _has_secrets_in_code_inputs(self) -> bool:
+        """Check if any code_inputs contain VellumSecret instances that require API execution."""
+        for input_value in self.code_inputs.values():
+            if isinstance(input_value, VellumSecret):
+                return True
+        return False
+
     def _compile_code_inputs(self) -> List[CodeExecutorInput]:
         # TODO: We may want to consolidate with prompt deployment input compilation
         # https://app.shortcut.com/vellum/story/4117
@@ -138,6 +164,8 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
         compiled_inputs: List[CodeExecutorInput] = []
 
         for input_name, input_value in self.code_inputs.items():
+            if input_value is undefined:
+                continue
             if isinstance(input_value, str):
                 compiled_inputs.append(
                     StringInput(
@@ -168,10 +196,13 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
                         )
                     )
                 else:
+                    # Convert primitive values to VellumValue objects
+                    vellum_values: List[VellumValue] = [primitive_to_vellum_value(item) for item in input_value]
+
                     compiled_inputs.append(
                         ArrayInput(
                             name=input_name,
-                            value=cast(List[VellumValue], input_value),
+                            value=vellum_values,
                         )
                     )
             elif isinstance(input_value, dict):
@@ -210,7 +241,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
 
         return compiled_inputs
 
-    def _resolve_code(self) -> str:
+    def _resolve_code(self) -> Tuple[str, str]:
         if self.code and self.filepath:
             raise NodeException(
                 message="Cannot specify both `code` and `filepath` for a CodeExecutionNode",
@@ -218,7 +249,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
             )
 
         if self.code:
-            return self.code
+            return self.code, f"{self.__class__.__name__}.code.py"
 
         if not self.filepath:
             raise NodeException(
@@ -235,4 +266,4 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
                 code=WorkflowErrorCode.INVALID_INPUTS,
             )
 
-        return code
+        return code, os.path.join(os.path.dirname(root), self.filepath)