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

vellum/types/severity_enum.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.severity_enum import *

vellum/types/type_checker_enum.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.type_checker_enum import *

vellum/types/workflow_sandbox_execute_node_response.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.workflow_sandbox_execute_node_response import *

vellum/utils/files/mixin.py

@@ -1,6 +1,13 @@
+import base64
+import binascii
+import re
 from typing import Any, ContextManager, Iterator, Optional
 
+from pydantic import field_serializer
+
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.utils.files.constants import BASE64_DATA_URL_PATTERN
+from vellum.utils.files.exceptions import InvalidFileSourceError
 from vellum.utils.files.read import read_vellum_file
 from vellum.utils.files.stream import stream_vellum_file
 from vellum.utils.files.upload import upload_vellum_file
@@ -107,3 +114,22 @@ class VellumFileMixin(UniversalBaseModel):
             str: A signed URL for accessing the uploaded file
         """
         return get_signed_url(self, expiry_seconds=expiry_seconds, vellum_client=vellum_client)
+
+    @field_serializer("src", check_fields=False)
+    @classmethod
+    def validate_src_on_serialize(cls, src: str) -> str:
+        """Validate the src field during serialization.
+
+        For PDF base64 data URLs, this validates that the base64 content is properly encoded.
+        """
+        data_url_match = re.match(BASE64_DATA_URL_PATTERN, src)
+        if data_url_match:
+            mime_type = data_url_match.group(1) or ""
+            if mime_type == "application/pdf":
+                base64_content = data_url_match.group(3)
+                try:
+                    base64.b64decode(base64_content, validate=True)
+                except binascii.Error as e:
+                    raise InvalidFileSourceError(f"Invalid base64 encoding in PDF data URL: {e}") from e
+
+        return src

vellum/utils/files/tests/test_mixin.py

@@ -1,6 +1,9 @@
+import pytest
 import base64
 from unittest.mock import Mock, patch
 
+from pydantic_core import PydanticSerializationError
+
 from vellum.utils.files import VellumFileMixin
 
 # Sample content for testing
@@ -203,3 +206,62 @@ def test_vellum_file_mixin_get_signed_url_method_already_uploaded():
         assert result == signed_url
         mock_client.uploaded_files.create.assert_not_called()
         mock_client.uploaded_files.retrieve.assert_called_once_with(file_id, expiry_seconds=604800)
+
+
+def test_vellum_file_mixin_serialization_valid_base64():
+    """Tests that serialization succeeds with valid base64 data URLs."""
+
+    # GIVEN a test document with valid base64 content
+    content = b"Hello, this is a test document!"
+    base64_content = base64.b64encode(content).decode("utf-8")
+    src = f"data:application/pdf;base64,{base64_content}"
+    test_file = TestDocument(src=src)
+
+    # WHEN serializing the document
+    serialized = test_file.model_dump()
+
+    # THEN the serialization should succeed and preserve the src
+    assert serialized["src"] == src
+
+
+def test_vellum_file_mixin_serialization_invalid_base64():
+    """Tests that serialization raises an error for invalid base64 data URLs."""
+
+    # GIVEN a test document with invalid base64 content
+    invalid_base64 = "not-valid-base64!!!"
+    src = f"data:application/pdf;base64,{invalid_base64}"
+    test_file = TestDocument(src=src)
+
+    # WHEN serializing the document
+    # THEN a PydanticSerializationError should be raised with InvalidFileSourceError as the cause
+    with pytest.raises(PydanticSerializationError, match="Invalid base64 encoding in PDF data URL"):
+        test_file.model_dump()
+
+
+def test_vellum_file_mixin_serialization_non_data_url():
+    """Tests that serialization succeeds for non-data URL sources."""
+
+    # GIVEN a test document with a regular URL source
+    src = "https://example.com/document.pdf"
+    test_file = TestDocument(src=src)
+
+    # WHEN serializing the document
+    serialized = test_file.model_dump()
+
+    # THEN the serialization should succeed and preserve the src
+    assert serialized["src"] == src
+
+
+def test_vellum_file_mixin_serialization_vellum_uploaded_file():
+    """Tests that serialization succeeds for vellum:uploaded-file sources."""
+
+    # GIVEN a test document with a vellum:uploaded-file source
+    file_id = "12345678-1234-1234-1234-123456789abc"
+    src = f"vellum:uploaded-file:{file_id}"
+    test_file = TestDocument(src=src)
+
+    # WHEN serializing the document
+    serialized = test_file.model_dump()
+
+    # THEN the serialization should succeed and preserve the src
+    assert serialized["src"] == src

vellum/utils/tests/test_vellum_client.py

@@ -0,0 +1,95 @@
+import pytest
+from unittest.mock import patch
+
+from vellum.utils.vellum_client import create_vellum_client, create_vellum_environment
+
+
+class TestCreateVellumClient:
+    """Tests for create_vellum_client and create_vellum_environment."""
+
+    @pytest.mark.parametrize(
+        ["api_url", "predict_api_url", "expected_predict_url"],
+        [
+            (None, None, "https://predict.vellum.ai"),
+            (None, "https://custom-predict.example.com", "https://custom-predict.example.com"),
+            ("CUSTOM_API_URL", None, "https://predict.vellum.ai"),
+            ("CUSTOM_API_URL", "https://custom-predict.example.com", "https://custom-predict.example.com"),
+        ],
+        ids=[
+            "defaults_to_standard_predict_url",
+            "predict_api_url_overrides_default",
+            "api_url_env_var_does_not_affect_predict_when_not_set",
+            "predict_api_url_takes_precedence_over_api_url",
+        ],
+    )
+    def test_create_vellum_client__predict_url_resolution(self, api_url, predict_api_url, expected_predict_url):
+        """Tests that predict_api_url is resolved correctly based on provided parameters."""
+
+        # GIVEN no environment variables are set
+        with patch.dict("os.environ", {}, clear=True):
+            # WHEN creating a vellum client with the given parameters
+            client = create_vellum_client(
+                api_key="test-api-key",
+                api_url=api_url,
+                predict_api_url=predict_api_url,
+            )
+
+            # THEN the predict URL should be resolved correctly
+            assert client._client_wrapper._environment.predict == expected_predict_url
+
+    @pytest.mark.parametrize(
+        ["env_vars", "predict_api_url", "expected_predict_url"],
+        [
+            ({"VELLUM_PREDICT_API_URL": "https://env-predict.example.com"}, None, "https://env-predict.example.com"),
+            (
+                {"VELLUM_PREDICT_API_URL": "https://env-predict.example.com"},
+                "https://param-predict.example.com",
+                "https://param-predict.example.com",
+            ),
+            ({"VELLUM_API_URL": "https://env-api.example.com"}, None, "https://env-api.example.com"),
+            (
+                {
+                    "VELLUM_API_URL": "https://env-api.example.com",
+                    "VELLUM_PREDICT_API_URL": "https://env-predict.example.com",
+                },
+                None,
+                "https://env-predict.example.com",
+            ),
+        ],
+        ids=[
+            "env_var_predict_api_url_is_used",
+            "param_predict_api_url_overrides_env_var",
+            "env_var_api_url_fallback_for_predict",
+            "env_var_predict_api_url_takes_precedence_over_api_url",
+        ],
+    )
+    def test_create_vellum_client__predict_url_env_var_resolution(
+        self, env_vars, predict_api_url, expected_predict_url
+    ):
+        """Tests that predict_api_url respects environment variables with correct precedence."""
+
+        # GIVEN specific environment variables are set
+        with patch.dict("os.environ", env_vars, clear=True):
+            # WHEN creating a vellum client with the given parameters
+            client = create_vellum_client(
+                api_key="test-api-key",
+                predict_api_url=predict_api_url,
+            )
+
+            # THEN the predict URL should be resolved correctly
+            assert client._client_wrapper._environment.predict == expected_predict_url
+
+    def test_create_vellum_environment__predict_api_url_does_not_affect_other_urls(self):
+        """Tests that predict_api_url only affects the predict endpoint, not default or documents."""
+
+        # GIVEN no environment variables are set
+        with patch.dict("os.environ", {}, clear=True):
+            # WHEN creating a vellum environment with a custom predict_api_url
+            environment = create_vellum_environment(predict_api_url="https://custom-predict.example.com")
+
+            # THEN the predict URL should be customized
+            assert environment.predict == "https://custom-predict.example.com"
+
+            # AND the default and documents URLs should remain unchanged
+            assert environment.default == "https://api.vellum.ai"
+            assert environment.documents == "https://documents.vellum.ai"

vellum/utils/uuid.py

@@ -1,8 +1,25 @@
 import uuid
-from typing import Union
+from typing import Any, Literal, Union, overload
+from typing_extensions import TypeGuard
 
 
-def is_valid_uuid(val: Union[str, uuid.UUID, None]) -> bool:
+@overload
+def is_valid_uuid(val: None) -> Literal[False]: ...
+
+
+@overload
+def is_valid_uuid(val: str) -> TypeGuard[str]: ...
+
+
+@overload
+def is_valid_uuid(val: uuid.UUID) -> TypeGuard[uuid.UUID]: ...
+
+
+@overload
+def is_valid_uuid(val: Any) -> Literal[False]: ...
+
+
+def is_valid_uuid(val: Any) -> TypeGuard[Union[str, uuid.UUID]]:
     try:
         uuid.UUID(str(val))
         return True

vellum/utils/vellum_client.py

@@ -9,22 +9,29 @@ def create_vellum_client(
     api_key: Optional[str] = None,
     api_url: Optional[str] = None,
     api_version: Optional[ApiVersionEnum] = None,
+    predict_api_url: Optional[str] = None,
 ) -> Vellum:
     if api_key is None:
         api_key = os.getenv("VELLUM_API_KEY", default="")
 
     return Vellum(
         api_key=api_key,
-        environment=create_vellum_environment(api_url),
+        environment=create_vellum_environment(api_url, predict_api_url=predict_api_url),
         api_version=api_version,
     )
 
 
-def create_vellum_environment(api_url: Optional[str] = None) -> VellumEnvironment:
+def create_vellum_environment(
+    api_url: Optional[str] = None,
+    predict_api_url: Optional[str] = None,
+) -> VellumEnvironment:
+    predict = predict_api_url or _resolve_env(
+        [api_url, "VELLUM_PREDICT_API_URL", "VELLUM_API_URL"], "https://predict.vellum.ai"
+    )
     return VellumEnvironment(
         default=_resolve_env([api_url, "VELLUM_DEFAULT_API_URL", "VELLUM_API_URL"], "https://api.vellum.ai"),
         documents=_resolve_env([api_url, "VELLUM_DOCUMENTS_API_URL", "VELLUM_API_URL"], "https://documents.vellum.ai"),
-        predict=_resolve_env([api_url, "VELLUM_PREDICT_API_URL", "VELLUM_API_URL"], "https://predict.vellum.ai"),
+        predict=predict,
     )
 
 

vellum/workflows/__init__.py

@@ -26,15 +26,17 @@ from .nodes import (
     WebSearchNode,
 )
 from .nodes.displayable.tool_calling_node import ToolCallingNode
+from .nodes.mocks import MockNodeExecution
 from .ports import Port
 from .references.environment_variable import EnvironmentVariableReference
 from .references.lazy import LazyReference
 from .runner import WorkflowRunner
 from .sandbox import WorkflowSandboxRunner
 from .state.base import BaseState
-from .triggers import IntegrationTrigger, ScheduleTrigger
+from .triggers import ChatMessageTrigger, IntegrationTrigger, ScheduleTrigger
 from .types.core import Json, MergeBehavior
 from .types.definition import DeploymentDefinition, MCPServer, VellumIntegrationToolDefinition
+from .utils.functions import tool, use_tool_inputs
 from .workflows import BaseWorkflow
 
 __all__ = [
@@ -48,6 +50,7 @@ __all__ = [
     "BaseSearchNode",
     "BaseState",
     "BaseWorkflow",
+    "ChatMessageTrigger",
     "CodeExecutionNode",
     "ConditionalNode",
     "DatasetRow",
@@ -64,6 +67,7 @@ __all__ = [
     "MapNode",
     "MCPServer",
     "MergeBehavior",
+    "MockNodeExecution",
     "NodeException",
     "NoteNode",
     "Port",
@@ -73,8 +77,10 @@ __all__ = [
     "SearchNode",
     "SubworkflowDeploymentNode",
     "TemplatingNode",
+    "tool",
     "ToolCallingNode",
     "TryNode",
+    "use_tool_inputs",
     "VellumIntegrationToolDefinition",
     "WebSearchNode",
     "WorkflowErrorCode",

vellum/workflows/descriptors/base.py

@@ -1,5 +1,8 @@
 from typing import TYPE_CHECKING, Any, Generic, Optional, Tuple, Type, TypeVar, Union, cast, overload
 
+from pydantic import GetCoreSchemaHandler, ValidationInfo
+from pydantic_core import core_schema
+
 if TYPE_CHECKING:
     from vellum.workflows.expressions.accessor import AccessorExpression
     from vellum.workflows.expressions.add import AddExpression
@@ -65,6 +68,15 @@ class BaseDescriptor(Generic[_T]):
     def types(self) -> Tuple[Type[_T], ...]:
         return self._types
 
+    @property
+    def normalized_type(self) -> Type:
+        all_types = self._types
+        if len(all_types) == 1:
+            return all_types[0]
+        else:
+            # Union should be considered a type, but mypy doesn't agree
+            return Union[all_types]  # type: ignore[return-value]
+
     @property
     def instance(self) -> Optional[_T]:
         return self._instance
@@ -179,6 +191,17 @@ class BaseDescriptor(Generic[_T]):
 
         return LessThanExpression(lhs=self, rhs=other)
 
+    @overload
+    def __lt__(self, other: "BaseDescriptor[_O]") -> "LessThanExpression[_T, _O]": ...  # type: ignore[misc]
+
+    @overload
+    def __lt__(self, other: _O) -> "LessThanExpression[_T, _O]": ...
+
+    def __lt__(self, other: "Union[BaseDescriptor[_O], _O]") -> "LessThanExpression[_T, _O]":
+        from vellum.workflows.expressions.less_than import LessThanExpression
+
+        return LessThanExpression(lhs=self, rhs=other)
+
     @overload
     def greater_than(self, other: "BaseDescriptor[_O]") -> "GreaterThanExpression[_T, _O]": ...
 
@@ -190,6 +213,17 @@ class BaseDescriptor(Generic[_T]):
 
         return GreaterThanExpression(lhs=self, rhs=other)
 
+    @overload
+    def __gt__(self, other: "BaseDescriptor[_O]") -> "GreaterThanExpression[_T, _O]": ...  # type: ignore[misc]
+
+    @overload
+    def __gt__(self, other: _O) -> "GreaterThanExpression[_T, _O]": ...
+
+    def __gt__(self, other: "Union[BaseDescriptor[_O], _O]") -> "GreaterThanExpression[_T, _O]":
+        from vellum.workflows.expressions.greater_than import GreaterThanExpression
+
+        return GreaterThanExpression(lhs=self, rhs=other)
+
     @overload
     def less_than_or_equal_to(self, other: "BaseDescriptor[_O]") -> "LessThanOrEqualToExpression[_T, _O]": ...
 
@@ -201,6 +235,17 @@ class BaseDescriptor(Generic[_T]):
 
         return LessThanOrEqualToExpression(lhs=self, rhs=other)
 
+    @overload
+    def __le__(self, other: "BaseDescriptor[_O]") -> "LessThanOrEqualToExpression[_T, _O]": ...  # type: ignore[misc]
+
+    @overload
+    def __le__(self, other: _O) -> "LessThanOrEqualToExpression[_T, _O]": ...
+
+    def __le__(self, other: "Union[BaseDescriptor[_O], _O]") -> "LessThanOrEqualToExpression[_T, _O]":
+        from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
+
+        return LessThanOrEqualToExpression(lhs=self, rhs=other)
+
     @overload
     def greater_than_or_equal_to(self, other: "BaseDescriptor[_O]") -> "GreaterThanOrEqualToExpression[_T, _O]": ...
 
@@ -214,6 +259,17 @@ class BaseDescriptor(Generic[_T]):
 
         return GreaterThanOrEqualToExpression(lhs=self, rhs=other)
 
+    @overload
+    def __ge__(self, other: "BaseDescriptor[_O]") -> "GreaterThanOrEqualToExpression[_T, _O]": ...  # type: ignore[misc]
+
+    @overload
+    def __ge__(self, other: _O) -> "GreaterThanOrEqualToExpression[_T, _O]": ...
+
+    def __ge__(self, other: "Union[BaseDescriptor[_O], _O]") -> "GreaterThanOrEqualToExpression[_T, _O]":
+        from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
+
+        return GreaterThanOrEqualToExpression(lhs=self, rhs=other)
+
     @overload
     def contains(self, other: "BaseDescriptor[_O]") -> "ContainsExpression[_T, _O]": ...
 
@@ -439,3 +495,33 @@ class BaseDescriptor(Generic[_T]):
         from vellum.workflows.expressions.minus import MinusExpression
 
         return MinusExpression(lhs=self, rhs=other)
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Type[Any], handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        def validate(value: Any, info: ValidationInfo) -> Any:
+            if isinstance(value, cls):
+                return value
+
+            if not isinstance(value, dict):
+                raise TypeError(f"Value must be an instance of {cls.__name__} or a dict")
+
+            context = info.context
+            if not isinstance(context, dict):
+                raise TypeError(f"Unexpected type for context: {type(context)}")
+
+            validator = context.get("descriptor_validator")
+            if validator and callable(validator):
+                return validator(value)
+
+            from vellum.workflows.references.constant import ConstantValueReference
+
+            return ConstantValueReference(value)
+
+        return core_schema.union_schema(
+            [
+                core_schema.is_instance_schema(cls),
+                core_schema.with_info_after_validator_function(validate, core_schema.dict_schema()),
+            ]
+        )

vellum/workflows/descriptors/tests/test_utils.py

@@ -132,6 +132,11 @@ class DummyNode(BaseNode[FixtureState]):
         (FixtureState.alpha + FixtureState.beta, 3),
         (FixtureState.gamma + FixtureState.delta, "helloel"),
         (FixtureState.theta + FixtureState.theta, ["baz", "baz"]),
+        # Test comparison operators
+        (FixtureState.alpha < FixtureState.beta, True),
+        (FixtureState.alpha > FixtureState.beta, False),
+        (FixtureState.alpha <= FixtureState.beta, True),
+        (FixtureState.alpha >= FixtureState.beta, False),
     ],
     ids=[
         "or",
@@ -194,6 +199,10 @@ class DummyNode(BaseNode[FixtureState]):
         "add_integers",
         "add_strings",
         "add_lists",
+        "lt_operator",
+        "gt_operator",
+        "le_operator",
+        "ge_operator",
     ],
 )
 def test_resolve_value__happy_path(descriptor, expected_value):

vellum/workflows/errors/tests/test_types.py

@@ -0,0 +1,52 @@
+from vellum.workflows.errors.types import WorkflowError, WorkflowErrorCode, workflow_error_to_vellum_error
+
+
+def test_workflow_error_to_vellum_error__preserves_raw_data():
+    """
+    Tests that raw_data is preserved when converting WorkflowError to VellumError.
+    """
+    # GIVEN a WorkflowError with raw_data containing integration details
+    raw_data = {
+        "integration": {
+            "id": "550e8400-e29b-41d4-a716-446655440000",
+            "provider": "COMPOSIO",
+            "name": "GITHUB",
+        }
+    }
+    workflow_error = WorkflowError(
+        message="Integration credentials unavailable",
+        code=WorkflowErrorCode.INTEGRATION_CREDENTIALS_UNAVAILABLE,
+        raw_data=raw_data,
+    )
+
+    # WHEN we convert to VellumError
+    vellum_error = workflow_error_to_vellum_error(workflow_error)
+
+    # THEN the raw_data should be preserved
+    assert vellum_error.raw_data == raw_data
+
+    # AND the integration details should be accessible
+    raw_data_from_error = vellum_error.raw_data
+    assert raw_data_from_error is not None
+    integration = raw_data_from_error.get("integration")
+    assert isinstance(integration, dict)
+    assert integration["name"] == "GITHUB"
+    assert integration["provider"] == "COMPOSIO"
+
+
+def test_workflow_error_to_vellum_error__handles_none_raw_data():
+    """
+    Tests that None raw_data is handled correctly when converting WorkflowError to VellumError.
+    """
+    # GIVEN a WorkflowError with no raw_data
+    workflow_error = WorkflowError(
+        message="Some error",
+        code=WorkflowErrorCode.INTERNAL_ERROR,
+        raw_data=None,
+    )
+
+    # WHEN we convert to VellumError
+    vellum_error = workflow_error_to_vellum_error(workflow_error)
+
+    # THEN the raw_data should be None
+    assert vellum_error.raw_data is None

vellum/workflows/errors/types.py

@@ -115,4 +115,5 @@ def workflow_error_to_vellum_error(error: WorkflowError) -> VellumError:
     return VellumError(
         message=error.message,
         code=_WORKFLOW_ERROR_CODE_TO_VELLUM_ERROR_CODE.get(error.code, "INTERNAL_SERVER_ERROR"),
+        raw_data=error.raw_data,
     )

vellum/workflows/events/node.py

@@ -1,8 +1,11 @@
+import json
 from typing import TYPE_CHECKING, Any, Dict, Generic, List, Literal, Optional, Set, Type, Union
 
 from pydantic import SerializationInfo, field_serializer, model_serializer
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.client.types import SeverityEnum
+from vellum.utils.json_encoder import VellumJsonEncoder
 from vellum.workflows.errors import WorkflowError
 from vellum.workflows.expressions.accessor import AccessorExpression
 from vellum.workflows.outputs.base import BaseOutput
@@ -136,6 +139,15 @@ class NodeExecutionFulfilledEvent(_BaseNodeEvent, Generic[OutputsType]):
             and serialized["body"]["invoked_ports"] is None
         ):
             del serialized["body"]["invoked_ports"]
+
+        if (
+            self._event_max_size is not None
+            and len(json.dumps(serialized, cls=VellumJsonEncoder)) > self._event_max_size
+            and "body" in serialized
+            and isinstance(serialized["body"], dict)
+        ):
+            serialized["body"]["outputs"] = {}
+
         return serialized
 
 
@@ -171,6 +183,17 @@ class NodeExecutionResumedEvent(_BaseNodeEvent):
     body: NodeExecutionResumedBody
 
 
+class NodeExecutionLogBody(_BaseNodeExecutionBody):
+    attributes: Optional[Dict[str, Any]] = None
+    severity: SeverityEnum
+    message: str
+
+
+class NodeExecutionLogEvent(_BaseNodeEvent):
+    name: Literal["node.execution.log"] = "node.execution.log"
+    body: NodeExecutionLogBody
+
+
 NodeEvent = Union[
     NodeExecutionInitiatedEvent,
     NodeExecutionStreamingEvent,
@@ -178,4 +201,5 @@ NodeEvent = Union[
     NodeExecutionRejectedEvent,
     NodeExecutionPausedEvent,
     NodeExecutionResumedEvent,
+    NodeExecutionLogEvent,
 ]