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

vellum/utils/files/tests/test_urls.py

@@ -0,0 +1,252 @@
+import pytest
+import base64
+from unittest.mock import Mock, patch
+
+from vellum import VellumAudio, VellumDocument, VellumImage, VellumVideo
+from vellum.client.core.api_error import ApiError
+from vellum.utils.files.exceptions import FileRetrievalError, InvalidFileSourceError
+from vellum.utils.files.urls import get_signed_url
+
+# Sample content for testing
+SAMPLE_TEXT_CONTENT = b"This is a sample text document content."
+
+
+@pytest.fixture
+def mock_vellum_client():
+    """Fixture that provides a mock Vellum client."""
+    mock_client = Mock()
+    with patch("vellum.utils.vellum_client.create_vellum_client", return_value=mock_client):
+        yield mock_client
+
+
+@pytest.mark.parametrize("file_type", [VellumDocument, VellumImage, VellumAudio, VellumVideo])
+def test_get_signed_url_from_base64(mock_vellum_client, file_type):
+    """Test getting signed URL for a file with base64 data URL source."""
+
+    # GIVEN a VellumFile with a base64 data URL source
+    base64_content = base64.b64encode(SAMPLE_TEXT_CONTENT).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    vellum_file = file_type(src=src)
+    uploaded_file_id = "12345678-1234-1234-1234-123456789abc"
+    signed_url = "https://example.com/signed-url?token=abc123"
+
+    # Configure mock client
+    mock_vellum_client.uploaded_files.create.return_value = Mock(id=uploaded_file_id)
+    mock_vellum_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+
+    # WHEN getting the signed URL
+    result = get_signed_url(vellum_file)
+
+    # THEN the file should be uploaded and signed URL returned
+    assert result == signed_url
+    mock_vellum_client.uploaded_files.create.assert_called_once()
+    mock_vellum_client.uploaded_files.retrieve.assert_called_once_with(uploaded_file_id, expiry_seconds=604800)
+
+
+@pytest.mark.parametrize("file_type", [VellumDocument, VellumImage, VellumAudio, VellumVideo])
+@patch("vellum.utils.files.upload.requests.get")
+def test_get_signed_url_from_http_url(mock_get, mock_vellum_client, file_type):
+    """Test getting signed URL for a file with HTTP URL source."""
+
+    # GIVEN a VellumFile with an HTTP URL source
+    url = "https://example.com/file.dat"
+    vellum_file = file_type(src=url)
+    uploaded_file_id = "87654321-4321-4321-4321-210987654321"
+    signed_url = "https://storage.vellum.ai/signed-url?token=xyz789"
+
+    # Configure mock HTTP response
+    mock_response = Mock(content=SAMPLE_TEXT_CONTENT, headers={"content-type": "application/octet-stream"})
+    mock_response.raise_for_status = Mock()
+    mock_get.return_value = mock_response
+
+    # Configure mock client
+    mock_vellum_client.uploaded_files.create.return_value = Mock(id=uploaded_file_id)
+    mock_vellum_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+
+    # WHEN getting the signed URL
+    result = get_signed_url(vellum_file)
+
+    # THEN the file should be downloaded, uploaded, and signed URL returned
+    assert result == signed_url
+    mock_get.assert_called_once_with(url, stream=True)
+    mock_vellum_client.uploaded_files.create.assert_called_once()
+    mock_vellum_client.uploaded_files.retrieve.assert_called_once_with(uploaded_file_id, expiry_seconds=604800)
+
+
+@pytest.mark.parametrize("file_type", [VellumDocument, VellumImage, VellumAudio, VellumVideo])
+def test_get_signed_url_already_uploaded(mock_vellum_client, file_type):
+    """Test getting signed URL for a file that's already uploaded."""
+
+    # GIVEN a VellumFile that is already uploaded (has vellum:uploaded-file: src)
+    file_id = "12345678-1234-1234-1234-123456789abc"
+    src = f"vellum:uploaded-file:{file_id}"
+    vellum_file = file_type(src=src)
+    signed_url = "https://storage.vellum.ai/files/signed?id=123"
+
+    # Configure mock client
+    mock_vellum_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+
+    # WHEN getting the signed URL
+    result = get_signed_url(vellum_file)
+
+    # THEN the signed URL should be returned without uploading again
+    assert result == signed_url
+    mock_vellum_client.uploaded_files.create.assert_not_called()
+    mock_vellum_client.uploaded_files.retrieve.assert_called_once_with(file_id, expiry_seconds=604800)
+
+
+@pytest.mark.parametrize(
+    ["src_pattern"],
+    [
+        ("vellum:uploaded-file:12345678-1234-1234-1234-123456789abc",),
+        ("VELLUM:UPLOADED-FILE:12345678-1234-1234-1234-123456789abc",),
+        ("Vellum:Uploaded-File:12345678-1234-1234-1234-123456789abc",),
+    ],
+)
+def test_get_signed_url_case_insensitive_vellum_src(mock_vellum_client, src_pattern):
+    """Test that vellum:uploaded-file: pattern matching is case insensitive."""
+
+    # GIVEN a VellumFile with a vellum:uploaded-file: src in various cases
+    file_id = "12345678-1234-1234-1234-123456789abc"
+    vellum_file = VellumDocument(src=src_pattern)
+    signed_url = "https://storage.vellum.ai/files/signed"
+
+    # Configure mock client
+    mock_vellum_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+
+    # WHEN getting the signed URL
+    result = get_signed_url(vellum_file)
+
+    # THEN the signed URL should be returned regardless of case
+    assert result == signed_url
+    mock_vellum_client.uploaded_files.retrieve.assert_called_once_with(file_id, expiry_seconds=604800)
+
+
+@pytest.mark.parametrize("file_type", [VellumDocument, VellumImage, VellumAudio, VellumVideo])
+@patch("vellum.utils.files.urls.upload_vellum_file")
+def test_get_signed_url_with_custom_client(mock_upload, file_type):
+    """Test getting signed URL with a custom Vellum client."""
+
+    # GIVEN a VellumFile and a custom Vellum client
+    base64_content = base64.b64encode(SAMPLE_TEXT_CONTENT).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    vellum_file = file_type(src=src)
+    uploaded_file_id = "aaaabbbb-cccc-dddd-eeee-111122223333"
+    signed_url = "https://custom.storage.com/signed"
+
+    # Create custom client
+    custom_client = Mock()
+    custom_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+
+    # Mock upload_vellum_file to return an uploaded file
+    mock_upload.return_value = file_type(src=f"vellum:uploaded-file:{uploaded_file_id}")
+
+    # WHEN getting the signed URL with the custom client
+    result = get_signed_url(vellum_file, vellum_client=custom_client)
+
+    # THEN the custom client should be used
+    assert result == signed_url
+    mock_upload.assert_called_once_with(vellum_file, vellum_client=custom_client)
+    custom_client.uploaded_files.retrieve.assert_called_once_with(uploaded_file_id, expiry_seconds=604800)
+
+
+def test_get_signed_url_invalid_uploaded_file_format(mock_vellum_client):
+    """Test that InvalidFileSourceError is raised when uploaded file ID cannot be parsed."""
+
+    # GIVEN a VellumFile with a malformed vellum:uploaded-file: src
+    # This should never happen in practice but we test error handling
+    vellum_file = VellumDocument(src="vellum:uploaded-file:invalid-id-format")
+
+    # Configure mock to return an uploaded file with malformed src
+    mock_vellum_client.uploaded_files.create.return_value = Mock(id="invalid-id")
+
+    # Mock the upload_vellum_file to return a file with invalid src
+    with patch("vellum.utils.files.urls.upload_vellum_file") as mock_upload:
+        mock_upload.return_value = VellumDocument(src="vellum:uploaded-file:not-a-valid-uuid")
+
+        # WHEN getting the signed URL
+        # THEN an InvalidFileSourceError should be raised
+        with pytest.raises(InvalidFileSourceError, match="Failed to determine id of uploaded file"):
+            get_signed_url(vellum_file)
+
+
+def test_get_signed_url_api_error_on_retrieve(mock_vellum_client):
+    """Test handling of API errors when retrieving signed URL."""
+
+    # GIVEN a VellumFile that is already uploaded
+    file_id = "12345678-1234-1234-1234-123456789abc"
+    src = f"vellum:uploaded-file:{file_id}"
+    vellum_file = VellumDocument(src=src)
+
+    # Configure mock client to raise an API error on retrieve
+    mock_vellum_client.uploaded_files.retrieve.side_effect = ApiError(status_code=404, body="File not found")
+
+    # WHEN getting the signed URL
+    # THEN a FileRetrievalError should be raised wrapping the API error
+    with pytest.raises(FileRetrievalError, match="Failed to retrieve file from Vellum"):
+        get_signed_url(vellum_file)
+
+
+def test_get_signed_url_missing_signed_url(mock_vellum_client):
+    """Test that FileRetrievalError is raised when signed URL is None."""
+
+    # GIVEN a VellumFile that is already uploaded
+    file_id = "12345678-1234-1234-1234-123456789abc"
+    src = f"vellum:uploaded-file:{file_id}"
+    vellum_file = VellumDocument(src=src)
+
+    # Configure mock client to return None for file_url
+    mock_vellum_client.uploaded_files.retrieve.return_value = Mock(file_url=None)
+
+    # WHEN getting the signed URL
+    # THEN a FileRetrievalError should be raised
+    with pytest.raises(FileRetrievalError, match="Failed to retrieve signed URL"):
+        get_signed_url(vellum_file)
+
+
+def test_get_signed_url_uses_default_client_when_not_provided(mock_vellum_client):
+    """Test that a default client is created when vellum_client is not provided."""
+
+    # GIVEN a VellumFile and no custom client
+    file_id = "12345678-1234-1234-1234-123456789abc"
+    src = f"vellum:uploaded-file:{file_id}"
+    vellum_file = VellumDocument(src=src)
+    signed_url = "https://storage.vellum.ai/files/signed"
+
+    # Configure mock client
+    mock_vellum_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+
+    # WHEN getting the signed URL without providing a client
+    result = get_signed_url(vellum_file)
+
+    # THEN the default client should be created and used
+    assert result == signed_url
+    mock_vellum_client.uploaded_files.retrieve.assert_called_once_with(file_id, expiry_seconds=604800)
+
+
+@pytest.mark.parametrize("file_type", [VellumDocument, VellumImage, VellumAudio, VellumVideo])
+@patch("vellum.utils.files.urls.upload_vellum_file")
+def test_get_signed_url_preserves_file_type(mock_upload, mock_vellum_client, file_type):
+    """Test that the file type is preserved through upload and retrieval."""
+
+    # GIVEN a VellumFile of a specific type
+    base64_content = base64.b64encode(SAMPLE_TEXT_CONTENT).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    vellum_file = file_type(src=src)
+    uploaded_file_id = "11112222-3333-4444-5555-666677778888"
+    signed_url = "https://storage.vellum.ai/files/typed"
+
+    # Mock upload_vellum_file to return an uploaded file
+    mock_upload.return_value = file_type(src=f"vellum:uploaded-file:{uploaded_file_id}")
+
+    # Configure mock client
+    mock_vellum_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+
+    # WHEN getting the signed URL
+    result = get_signed_url(vellum_file)
+
+    # THEN the signed URL should be returned
+    assert result == signed_url
+    # The file type should be consistent throughout the process
+    mock_upload.assert_called_once_with(vellum_file, vellum_client=None)
+    mock_vellum_client.uploaded_files.retrieve.assert_called_once_with(uploaded_file_id, expiry_seconds=604800)

vellum/utils/files/types.py

@@ -0,0 +1,8 @@
+"""Type definitions for Vellum files."""
+
+from typing import TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    from vellum import VellumAudio, VellumDocument, VellumImage, VellumVideo
+
+VellumFileTypes = Union["VellumDocument", "VellumImage", "VellumVideo", "VellumAudio"]

vellum/utils/files/upload.py

@@ -0,0 +1,151 @@
+"""File upload utilities for uploading files to Vellum."""
+
+import base64
+from io import BytesIO
+import logging
+import re
+from typing import TYPE_CHECKING, Optional
+
+import requests
+
+from vellum.client.core.api_error import ApiError
+from vellum.client.core.file import File
+from vellum.utils.files.constants import BASE64_DATA_URL_PATTERN, URL_PATTERN, VELLUM_FILE_SRC_PATTERN
+from vellum.utils.files.exceptions import FileNotFoundError, FileRetrievalError, InvalidFileSourceError
+from vellum.utils.files.extensions import ensure_filename_with_extension
+from vellum.utils.files.types import VellumFileTypes
+
+if TYPE_CHECKING:
+    from vellum.client import Vellum as VellumClient
+
+logger = logging.getLogger(__name__)
+
+
+def upload_vellum_file(
+    vellum_file: VellumFileTypes,
+    *,
+    filename: Optional[str] = None,
+    vellum_client: Optional["VellumClient"] = None,
+) -> VellumFileTypes:
+    """
+    Upload a file to Vellum and return a new VellumFile with the uploaded source.
+
+    This function takes any VellumFile object (with a src that could be a base64 data URL,
+    HTTP/HTTPS URL, or existing vellum:uploaded-file: identifier), downloads its content,
+    and uploads it to Vellum's storage.
+
+    Args:
+        vellum_file: A VellumDocument, VellumImage, VellumAudio, or VellumVideo instance
+        filename: Optional filename to use when uploading. If not provided, defaults to "file"
+                 with an appropriate extension inferred from the MIME type. If provided without
+                 an extension, the extension will be automatically added based on the MIME type.
+        vellum_client: An optional Vellum client instance. If not provided, a default client will be created.
+
+    Returns:
+        VellumFileTypes: A new VellumFile of the same type with the vellum:uploaded-file:{id} source
+
+    Raises:
+        InvalidFileSourceError: If the file source format is invalid
+        FileRetrievalError: If the file cannot be retrieved from its source
+        FileNotFoundError: If the file is not found at its source
+
+    Example:
+        ```python
+        from vellum import VellumDocument, VellumImage
+        from vellum.utils.files import upload_vellum_file
+
+        # Upload a document from a URL
+        doc = VellumDocument(src="https://example.com/doc.pdf")
+        uploaded_doc = upload_vellum_file(doc, filename="my_document.pdf")
+
+        # Upload a base64-encoded image (filename automatically inferred as "file.png")
+        image = VellumImage(src="data:image/png;base64,iVBORw0KGgo...")
+        uploaded_image = upload_vellum_file(image)
+
+        # Upload with filename without extension (extension automatically added)
+        doc = VellumDocument(src="data:application/pdf;base64,...")
+        uploaded_doc = upload_vellum_file(doc, filename="report")  # becomes "report.pdf"
+        ```
+    """
+    src = vellum_file.src
+
+    # Case 1: Vellum Uploaded File (already uploaded, just return as-is)
+    match = re.match(VELLUM_FILE_SRC_PATTERN, src, re.IGNORECASE)
+    if match:
+        vellum_uploaded_file_id = match.group(1)
+        logger.info(
+            "File already uploaded to Vellum, returning unchanged",
+            extra={"vellum_uploaded_file_id": vellum_uploaded_file_id},
+        )
+        return vellum_file
+
+    if vellum_client is None:
+        from vellum.utils.vellum_client import create_vellum_client
+
+        vellum_client = create_vellum_client()
+
+    # Case 2: Base64 Data URL
+    data_url_match = re.match(BASE64_DATA_URL_PATTERN, src)
+    if data_url_match:
+        mime_type = data_url_match.group(1) or "application/octet-stream"
+        base64_content = data_url_match.group(3)
+        decoded = base64.b64decode(base64_content)
+
+        # Ensure filename has appropriate extension
+        resolved_filename = ensure_filename_with_extension(filename, mime_type)
+        file_content: File = (resolved_filename, BytesIO(decoded), mime_type)
+
+        try:
+            uploaded_file = vellum_client.uploaded_files.create(file=file_content)
+        except ApiError as e:
+            raise FileRetrievalError(f"Failed to upload file to Vellum: {str(e)}") from e
+
+        new_src = f"vellum:uploaded-file:{uploaded_file.id}"
+        return _copy_with_new_src(vellum_file, new_src)
+
+    # Case 3: Direct URL
+    if not re.match(URL_PATTERN, src):
+        raise InvalidFileSourceError(
+            f"Invalid file source: {src}. "
+            "Expected a base64 data URL (data:...;base64,...), "
+            "a Vellum uploaded file ID (vellum:uploaded-file:...), "
+            "or an HTTP/HTTPS URL."
+        )
+
+    file_url = src
+
+    # Download from URL
+    try:
+        response = requests.get(file_url, stream=True)
+        response.raise_for_status()
+    except requests.HTTPError as e:
+        if e.response.status_code == 404:
+            raise FileNotFoundError(f"File not found at URL: {file_url}") from e
+        raise FileRetrievalError(f"Failed to retrieve file from URL: {file_url}") from e
+    except requests.RequestException as e:
+        raise FileRetrievalError(f"Network error while retrieving file: {file_url}") from e
+
+    # Read content and get content type
+    content = response.content
+    content_type = response.headers.get("content-type", "application/octet-stream")
+
+    # Ensure filename has appropriate extension
+    resolved_filename = ensure_filename_with_extension(filename, content_type)
+    file_content = (resolved_filename, BytesIO(content), content_type)
+
+    try:
+        uploaded_file = vellum_client.uploaded_files.create(file=file_content)
+    except ApiError as e:
+        raise FileRetrievalError("Failed to upload file to Vellum") from e
+
+    new_src = f"vellum:uploaded-file:{uploaded_file.id}"
+    return _copy_with_new_src(vellum_file, new_src)
+
+
+def _copy_with_new_src(vellum_file: VellumFileTypes, new_src: str) -> VellumFileTypes:
+    """Helper to create a copy of a VellumFile with an updated src."""
+    # Use model_copy for Pydantic v2, copy for v1
+    if hasattr(vellum_file, "model_copy"):
+        return vellum_file.model_copy(update={"src": new_src})
+    else:
+        return vellum_file.copy(update={"src": new_src})

vellum/utils/files/urls.py

@@ -0,0 +1,71 @@
+"""File upload utilities for uploading files to Vellum."""
+
+import logging
+import re
+from typing import TYPE_CHECKING, Optional
+
+from vellum.client.core.api_error import ApiError
+from vellum.utils.files.constants import VELLUM_FILE_SRC_PATTERN
+from vellum.utils.files.exceptions import FileRetrievalError, InvalidFileSourceError
+from vellum.utils.files.types import VellumFileTypes
+from vellum.utils.files.upload import upload_vellum_file
+
+if TYPE_CHECKING:
+    from vellum.client import Vellum as VellumClient
+
+logger = logging.getLogger(__name__)
+
+
+def get_signed_url(
+    vellum_file: VellumFileTypes,
+    *,
+    expiry_seconds: int = 7 * 24 * 60 * 60,  # 7 days
+    vellum_client: Optional["VellumClient"] = None,
+) -> str:
+    """
+    Retrieved a signed url for a file that's been uploaded to Vellum
+
+    This function takes any VellumFile object (with a src that could be a base64 data URL,
+    HTTP/HTTPS URL, or existing vellum:uploaded-file: identifier), uploads it to Vellum (if not already uploaded),
+    and returns a signed url for accessing the file.
+
+    Args:
+        vellum_file: A VellumDocument, VellumImage, VellumAudio, or VellumVideo instance
+        expiry_seconds: The number of seconds until the signed URL expires. Defaults to 7 days.
+        vellum_client: An optional Vellum client instance. If not provided, a default client will be created.
+
+    Returns:
+        str: A signed URL for accessing the uploaded file
+    """
+
+    # Upload the file to Vellum if it isn't already
+    uploaded_file = upload_vellum_file(vellum_file, vellum_client=vellum_client)
+
+    # Parse out its reference ID
+    vellum_uploaded_file_id: Optional[str] = None
+    src = uploaded_file.src
+    match = re.match(VELLUM_FILE_SRC_PATTERN, src, re.IGNORECASE)
+    if match:
+        vellum_uploaded_file_id = match.group(1)
+
+    if not vellum_uploaded_file_id:
+        raise InvalidFileSourceError("Failed to determine id of uploaded file.")
+
+    if vellum_client is None:
+        from vellum.utils.vellum_client import create_vellum_client
+
+        vellum_client = create_vellum_client()
+
+    # Fetch the signed URL for this file from Vellum
+    try:
+        vellum_uploaded_file = vellum_client.uploaded_files.retrieve(
+            vellum_uploaded_file_id, expiry_seconds=expiry_seconds
+        )
+    except ApiError as e:
+        raise FileRetrievalError("Failed to retrieve file from Vellum") from e
+
+    signed_url = vellum_uploaded_file.file_url
+    if not signed_url:
+        raise FileRetrievalError("Failed to retrieve signed URL for uploaded file.")
+
+    return signed_url

vellum/utils/json_encoder.py

@@ -0,0 +1,95 @@
+from dataclasses import asdict, is_dataclass
+from datetime import datetime
+import enum
+from json import JSONEncoder
+from queue import Queue
+from uuid import UUID
+from typing import Any, Callable, Dict, Protocol, Type
+
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+from pydantic_core import PydanticUndefined
+
+# Handle Pydantic v1 compatibility
+try:
+    from pydantic.fields import Undefined as PydanticV1Undefined  # type: ignore[attr-defined]
+except ImportError:
+    PydanticV1Undefined = None  # type: ignore[assignment,misc]
+
+
+class VellumJsonEncodable(Protocol):
+    """Protocol for objects that can be encoded to JSON by VellumJsonEncoder.
+
+    Classes implementing this protocol must define a __vellum_encode__ method
+    that returns a JSON-serializable representation of the object.
+    """
+
+    def __vellum_encode__(self) -> Any:
+        """Return a JSON-serializable representation of this object."""
+        ...
+
+
+class VellumJsonEncoder(JSONEncoder):
+    """JSON encoder that handles Vellum-specific types.
+
+    This encoder supports:
+    - Objects implementing VellumJsonEncodable protocol
+    - Pydantic BaseModel instances
+    - Pydantic FieldInfo objects (serializes to default value or None)
+    - PydanticUndefined (serializes to None)
+    - UUIDs, sets, datetimes, enums
+    - Dataclasses and Queue objects
+    - Custom encoder registry via encoders dict
+    """
+
+    encoders: Dict[Type, Callable] = {}
+
+    def default(self, obj: Any) -> Any:
+        if hasattr(obj, "__vellum_encode__") and callable(getattr(obj, "__vellum_encode__")):
+            return obj.__vellum_encode__()
+
+        if isinstance(obj, UUID):
+            return str(obj)
+
+        if isinstance(obj, set):
+            return list(obj)
+
+        if isinstance(obj, BaseModel):
+            return obj.model_dump()
+
+        if isinstance(obj, datetime):
+            return obj.isoformat()
+
+        if isinstance(obj, enum.Enum):
+            return obj.value
+
+        if isinstance(obj, Queue):
+            return list(obj.queue)
+
+        if is_dataclass(obj) and not isinstance(obj, type):
+            return asdict(obj)
+
+        if isinstance(obj, type):
+            return str(obj)
+
+        if isinstance(obj, Exception):
+            return str(obj)
+
+        # Handle Pydantic FieldInfo objects
+        if isinstance(obj, FieldInfo):
+            # Serialize the default value if it exists and is not undefined (v1 or v2), otherwise None
+            is_undefined = obj.default is PydanticUndefined or (
+                PydanticV1Undefined is not None and obj.default is PydanticV1Undefined
+            )
+            if not is_undefined and obj.default is not None:
+                return obj.default
+            return None
+
+        # Handle PydanticUndefined value
+        if obj is PydanticUndefined:
+            return None
+
+        if obj.__class__ in self.encoders:
+            return self.encoders[obj.__class__](obj)
+
+        return super().default(obj)

vellum/utils/templating/custom_filters.py

@@ -4,7 +4,7 @@ from typing import Any, Union
 from jinja2 import Undefined
 from jinja2.utils import htmlsafe_json_dumps
 
-from vellum.workflows.state.encoder import DefaultStateEncoder
+from vellum.utils.json_encoder import VellumJsonEncoder
 
 
 def is_valid_json_string(value: Union[str, bytes]) -> bool:
@@ -19,11 +19,11 @@ def is_valid_json_string(value: Union[str, bytes]) -> bool:
 
 def replace(s: Any, old: Any, new: Any) -> str:
     def encode_to_str(obj: Any) -> str:
-        """Encode an object for template rendering using DefaultStateEncoder."""
+        """Encode an object for template rendering using VellumJsonEncoder."""
         try:
             if isinstance(obj, str):
                 return obj
-            return json.dumps(obj, cls=DefaultStateEncoder)
+            return json.dumps(obj, cls=VellumJsonEncoder)
         except TypeError:
             return str(obj)
 
@@ -39,4 +39,4 @@ def replace(s: Any, old: Any, new: Any) -> str:
 def safe_tojson(value: Any) -> str:
     if isinstance(value, Undefined):
         return ""
-    return htmlsafe_json_dumps(value, cls=DefaultStateEncoder)
+    return htmlsafe_json_dumps(value, cls=VellumJsonEncoder)

vellum/utils/templating/render.py

@@ -4,17 +4,17 @@ from typing import Any, Dict, Optional
 from jinja2.sandbox import SandboxedEnvironment
 from pydantic import BaseModel
 
+from vellum.utils.json_encoder import VellumJsonEncoder
 from vellum.utils.templating.constants import FilterFunc
 from vellum.utils.templating.exceptions import JinjaTemplateError
-from vellum.workflows.state.encoder import DefaultStateEncoder
 
 
 def finalize(obj: Any) -> str:
     if isinstance(obj, (dict, list)):
-        return json.dumps(obj, cls=DefaultStateEncoder)
+        return json.dumps(obj, cls=VellumJsonEncoder)
 
     if isinstance(obj, BaseModel):
-        return json.dumps(obj.model_dump(), cls=DefaultStateEncoder)
+        return json.dumps(obj.model_dump(), cls=VellumJsonEncoder)
 
     return str(obj)
 
@@ -33,7 +33,7 @@ def render_sandboxed_jinja_template(
             finalize=finalize,
         )
         environment.policies["json.dumps_kwargs"] = {
-            "cls": DefaultStateEncoder,
+            "cls": VellumJsonEncoder,
         }
 
         if jinja_custom_filters: