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

vellum/utils/files/mixin.py

@@ -0,0 +1,135 @@
+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
+from vellum.utils.files.urls import get_signed_url
+
+
+class VellumFileMixin(UniversalBaseModel):
+    """Mixin class that provides stream, read, and upload methods for Vellum file types.
+
+    This class extends UniversalBaseModel, so file types can inherit from this
+    instead of both UniversalBaseModel and VellumFileMixin.
+    """
+
+    def stream(
+        self: Any,
+        *,
+        chunk_size: int = 8192,
+        vellum_client: Optional[Any] = None,
+    ) -> ContextManager[Iterator[bytes]]:
+        """
+        Stream the file content in chunks using a context manager.
+
+        This is the recommended way to read large files, as it avoids loading
+        the entire file into memory at once.
+
+        Args:
+            chunk_size: Size of chunks to yield in bytes (default 8KB)
+            vellum_client: An optional Vellum client instance. If not provided, a default client will be created.
+
+        Returns:
+            A context manager that yields an iterator of file content chunks.
+
+        Example:
+            ```python
+            with file.stream() as chunks:
+                for chunk in chunks:
+                    process(chunk)
+            ```
+        """
+        return stream_vellum_file(self, chunk_size=chunk_size, vellum_client=vellum_client)
+
+    def read(
+        self: Any,
+        *,
+        vellum_client: Optional[Any] = None,
+    ) -> bytes:
+        """
+        Read the entire file content into memory.
+
+        For large files (e.g., videos, large datasets), consider using
+        stream() directly to avoid memory issues.
+
+        Args:
+            vellum_client: An optional Vellum client instance. If not provided, a default client will be created.
+
+        Returns:
+            bytes: The complete file content
+        """
+        return read_vellum_file(self, vellum_client=vellum_client)
+
+    def upload(
+        self: Any,
+        *,
+        filename: Optional[str] = None,
+        vellum_client: Optional[Any] = None,
+    ) -> Any:
+        """
+        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:
+            A new instance with the vellum:uploaded-file:{id} source
+        """
+        return upload_vellum_file(self, filename=filename, vellum_client=vellum_client)
+
+    def get_signed_url(
+        self: Any,
+        *,
+        expiry_seconds: int = 7 * 24 * 60 * 60,  # 7 days
+        vellum_client: Optional[Any] = 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:
+            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
+        """
+        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/read.py

@@ -0,0 +1,41 @@
+"""Convenience utilities for reading Vellum files."""
+
+from typing import TYPE_CHECKING, Optional
+
+from vellum.utils.files.stream import stream_vellum_file
+from vellum.utils.files.types import VellumFileTypes
+
+if TYPE_CHECKING:
+    from vellum.client import Vellum as VellumClient
+
+
+def read_vellum_file(vellum_file: VellumFileTypes, *, vellum_client: Optional["VellumClient"] = None) -> bytes:
+    """
+    Convenience function that reads the entire file content into memory.
+
+    This function is implemented using stream_vellum_file() under the hood.
+    For large files (e.g., videos, large datasets), consider using
+    stream_vellum_file() directly to avoid memory issues.
+
+    Args:
+        vellum_file: A VellumDocument, VellumImage, VellumAudio, or VellumVideo instance
+        vellum_client: An optional Vellum client instance. If not provided, a default client will be created.
+
+    Returns:
+        bytes: The complete file content
+
+    Example:
+        ```python
+        from vellum import VellumDocument
+        from vellum.utils.files import read_vellum_file
+
+        document = VellumDocument(src="data:text/plain;base64,SGVsbG8gV29ybGQ=")
+        content = read_vellum_file(document)
+        print(content.decode('utf-8'))  # "Hello World"
+        ```
+    """
+    chunks = []
+    with stream_vellum_file(vellum_file, vellum_client=vellum_client) as chunk_iter:
+        for chunk in chunk_iter:
+            chunks.append(chunk)
+    return b"".join(chunks)

vellum/utils/files/stream.py

@@ -0,0 +1,135 @@
+"""Streaming utilities for reading Vellum files."""
+
+import base64
+from contextlib import contextmanager
+from io import BytesIO
+import re
+from typing import TYPE_CHECKING, Generator, Iterator, Optional
+
+import requests
+
+from vellum.client.core.api_error import ApiError
+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.types import VellumFileTypes
+
+if TYPE_CHECKING:
+    from vellum.client import Vellum as VellumClient
+
+
+@contextmanager
+def stream_vellum_file(
+    vellum_file: VellumFileTypes,
+    *,
+    chunk_size: int = 8192,
+    vellum_client: Optional["VellumClient"] = None,
+) -> Generator[Iterator[bytes], None, None]:
+    """
+    Stream the file content in chunks using a context manager.
+
+    This is the recommended way to read large files, as it avoids loading
+    the entire file into memory at once.
+
+    Args:
+        vellum_file: A VellumDocument, VellumImage, VellumAudio, or VellumVideo instance
+        chunk_size: Size of chunks to yield in bytes (default 8KB)
+        vellum_client: An optional Vellum client instance. If not provided, a default client will be created.
+
+    Yields:
+        Iterator[bytes]: Chunks of file content
+
+    Example:
+        ```python
+        from vellum import VellumVideo
+        from vellum.utils.files import stream_vellum_file
+
+        # Stream a large video file
+        video = VellumVideo(src="https://example.com/video.mp4")
+        with stream_vellum_file(video, chunk_size=10*1024*1024) as chunks:  # 10MB chunks
+            with open('output.mp4', 'wb') as f:
+                for chunk in chunks:
+                    f.write(chunk)
+
+        # Calculate hash without loading entire file
+        import hashlib
+        hasher = hashlib.sha256()
+        with stream_vellum_file(vellum_file) as chunks:
+            for chunk in chunks:
+                hasher.update(chunk)
+        hash_value = hasher.hexdigest()
+        ```
+    """
+    src = vellum_file.src
+
+    # Case 1: Base64 Data URL
+    # Note: Base64 content is already in memory, but we still provide
+    # a streaming interface for consistency
+    data_url_match = re.match(BASE64_DATA_URL_PATTERN, src)
+    if data_url_match:
+        base64_content = data_url_match.group(3)
+        decoded = base64.b64decode(base64_content)
+        stream = BytesIO(decoded)
+        try:
+            yield _chunk_iterator(stream, chunk_size)
+        finally:
+            stream.close()
+        return
+
+    file_url: str
+
+    # Case 2: Vellum Uploaded File
+    match = re.match(VELLUM_FILE_SRC_PATTERN, src, re.IGNORECASE)
+    if match:
+        vellum_uploaded_file_id = match.group(1)
+        if vellum_client is None:
+            from vellum.utils.vellum_client import create_vellum_client
+
+            vellum_client = create_vellum_client()
+
+        try:
+            uploaded_file = vellum_client.uploaded_files.retrieve(vellum_uploaded_file_id)
+        except ApiError as e:
+            if e.status_code == 404:
+                raise FileNotFoundError(f"Uploaded file not found: {vellum_uploaded_file_id}") from e
+            raise FileRetrievalError(f"Failed to retrieve uploaded file: {vellum_uploaded_file_id}") from e
+
+        if not uploaded_file.file_url:
+            raise FileRetrievalError(f"Uploaded file has no accessible URL: {vellum_uploaded_file_id}")
+
+        file_url = uploaded_file.file_url
+
+    # Case 3: Direct URL
+    elif re.match(URL_PATTERN, src):
+        file_url = src
+    else:
+        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."
+        )
+
+    # Stream 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
+
+    try:
+        yield response.iter_content(chunk_size=chunk_size)
+    finally:
+        response.close()
+
+
+def _chunk_iterator(stream: BytesIO, chunk_size: int) -> Iterator[bytes]:
+    """Helper to read a BytesIO stream in chunks."""
+    while True:
+        chunk = stream.read(chunk_size)
+        if not chunk:
+            break
+        yield chunk

vellum/utils/files/tests/test_extensions.py

@@ -0,0 +1,54 @@
+"""Tests for file extension inference utilities."""
+
+import pytest
+
+from vellum.utils.files.extensions import ensure_filename_with_extension
+
+
+@pytest.mark.parametrize(
+    ["filename", "mime_type", "expected"],
+    [
+        # None filename cases - should default to 'file' with inferred extension
+        pytest.param(None, "application/pdf", "file.pdf", id="none_filename_pdf"),
+        pytest.param(None, "image/png", "file.png", id="none_filename_png"),
+        pytest.param(None, "image/jpeg", "file.jpg", id="none_filename_jpeg_normalized"),
+        pytest.param(None, "text/plain", "file.txt", id="none_filename_txt"),
+        pytest.param(None, "video/mp4", "file.mp4", id="none_filename_mp4"),
+        pytest.param(None, "audio/mpeg", "file.mp3", id="none_filename_mp3"),
+        pytest.param(None, "text/markdown", "file.md", id="none_filename_markdown"),
+        pytest.param(None, "application/x-yaml", "file.yaml", id="none_filename_yaml_x"),
+        pytest.param(None, "text/yaml", "file.yaml", id="none_filename_yaml_text"),
+        pytest.param(None, "application/gzip", "file.gz", id="none_filename_gzip"),
+        pytest.param(None, "application/xml", "file.xml", id="none_filename_xml_normalized"),
+        pytest.param(None, "application/octet-stream", "file.bin", id="none_filename_octet_stream"),
+        pytest.param(None, "application/x-unknown-type", "file.bin", id="none_filename_unknown_mime"),
+        # Filename without extension - should add appropriate extension
+        pytest.param("document", "application/pdf", "document.pdf", id="document_without_ext_pdf"),
+        pytest.param("report", "application/pdf", "report.pdf", id="report_without_ext_pdf"),
+        pytest.param(
+            "report",
+            "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+            "report.docx",
+            id="report_without_ext_docx",
+        ),
+        pytest.param("config", "application/x-yaml", "config.yaml", id="config_without_ext_yaml_x"),
+        pytest.param("config", "text/yaml", "config.yaml", id="config_without_ext_yaml_text"),
+        pytest.param("README", "text/markdown", "README.md", id="readme_without_ext_markdown"),
+        pytest.param("data", "text/csv", "data.csv", id="data_without_ext_csv"),
+        pytest.param("data", "application/octet-stream", "data.bin", id="data_without_ext_bin"),
+        pytest.param("archive", "application/gzip", "archive.gz", id="archive_without_ext_gzip"),
+        pytest.param("file", "application/pdf", "file.pdf", id="file_without_ext_pdf"),
+        # Filename with extension - should keep existing extension
+        pytest.param("document.pdf", "application/pdf", "document.pdf", id="document_pdf_with_ext"),
+        pytest.param("image.png", "image/jpeg", "image.png", id="image_png_mismatched_mime"),
+        pytest.param("data.csv", "text/csv", "data.csv", id="data_csv_with_ext"),
+        pytest.param("config.yaml", "application/x-yaml", "config.yaml", id="config_yaml_with_ext"),
+        # MIME type with parameters - should strip parameters
+        pytest.param(None, "text/html; charset=utf-8", "file.html", id="none_filename_html_with_charset"),
+        pytest.param("page", "text/html; charset=utf-8", "page.html", id="page_without_ext_html_with_charset"),
+    ],
+)
+def test_ensure_filename_with_extension(filename, mime_type, expected):
+    """Test filename extension inference for various filename and MIME type combinations."""
+    result = ensure_filename_with_extension(filename, mime_type)
+    assert result == expected

vellum/utils/files/tests/test_mixin.py

@@ -0,0 +1,267 @@
+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
+SAMPLE_TEXT_CONTENT = b"Hello, this is a test document!"
+
+
+class TestDocument(VellumFileMixin):
+    """Test class that demonstrates how VellumFileMixin is used in postprocess.
+
+    This mirrors the pattern used in postprocess.ts where file types inherit from
+    VellumFileMixin (which extends UniversalBaseModel), replacing UniversalBaseModel.
+    """
+
+    src: str
+
+
+def test_vellum_file_mixin_stream_method():
+    """Test that VellumFileMixin.stream() method works."""
+
+    # GIVEN a test document with base64 content using the mixin (like postprocess does)
+    content = b"Hello, this is a test document!"
+    base64_content = base64.b64encode(content).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    test_file = TestDocument(src=src)
+
+    # WHEN streaming using the mixin's .stream() method
+    chunks = []
+    with test_file.stream() as chunk_iter:
+        for chunk in chunk_iter:
+            chunks.append(chunk)
+
+    # THEN the content should match
+    result = b"".join(chunks)
+    assert result == content
+
+
+def test_vellum_file_mixin_stream_method_with_custom_chunk_size():
+    """Test that VellumFileMixin.stream() respects custom chunk_size parameter."""
+
+    # GIVEN a test document with content using the mixin (like postprocess does)
+    content = b"This is a longer document that will be split into multiple chunks."
+    base64_content = base64.b64encode(content).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    test_file = TestDocument(src=src)
+
+    # WHEN streaming with a small chunk size using the mixin
+    chunk_size = 10
+    chunks = []
+    with test_file.stream(chunk_size=chunk_size) as chunk_iter:
+        for chunk in chunk_iter:
+            chunks.append(chunk)
+            # Each chunk should be <= chunk_size (last chunk may be smaller)
+            assert len(chunk) <= chunk_size
+
+    # THEN the content should match when reassembled
+    result = b"".join(chunks)
+    assert result == content
+    # AND we should have multiple chunks for content larger than chunk_size
+    assert len(chunks) > 1
+
+
+def test_vellum_file_mixin_stream_method_from_url():
+    """Test that VellumFileMixin.stream() method works for files from URLs."""
+
+    # GIVEN a test document with a URL source using the mixin (like postprocess does)
+    content = b"URL fetched content that will be streamed"
+    url = "https://example.com/file.dat"
+    test_file = TestDocument(src=url)
+
+    # WHEN streaming using the mixin's .stream() method
+    with patch("vellum.utils.files.stream.requests.get") as mock_get:
+        mock_response = Mock()
+        # Simulate streaming response with iter_content
+        mock_response.iter_content = Mock(return_value=iter([content[:20], content[20:]]))
+        mock_response.raise_for_status = Mock()
+        mock_response.close = Mock()
+        mock_get.return_value = mock_response
+
+        chunks = []
+        with test_file.stream() as chunk_iter:
+            for chunk in chunk_iter:
+                chunks.append(chunk)
+
+        result = b"".join(chunks)
+        assert result == content
+        mock_get.assert_called_once_with(url, stream=True)
+        mock_response.raise_for_status.assert_called_once()
+        mock_response.close.assert_called_once()
+
+
+def test_vellum_file_mixin_read_method():
+    """Test that VellumFileMixin.read() method works."""
+
+    # GIVEN a test document with base64 content using the mixin (like postprocess does)
+    content = b"Hello, this is a test document!"
+    base64_content = base64.b64encode(content).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    test_file = TestDocument(src=src)
+
+    # WHEN reading using the mixin's .read() method
+    result = test_file.read()
+
+    # THEN the content should match
+    assert result == content
+
+
+def test_vellum_file_mixin_upload_method():
+    """Test that VellumFileMixin.upload() method works."""
+
+    # GIVEN a test document with base64 content using the mixin (like postprocess does)
+    content = b"Hello, this is a test document!"
+    base64_content = base64.b64encode(content).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    test_file = TestDocument(src=src)
+    uploaded_file_id = "12345678-1234-1234-1234-123456789abc"
+
+    # WHEN uploading using the mixin's .upload() method
+    with patch("vellum.utils.vellum_client.create_vellum_client") as mock_create_client:
+        mock_client = Mock()
+        mock_client.uploaded_files.create.return_value = Mock(id=uploaded_file_id)
+        mock_create_client.return_value = mock_client
+
+        result = test_file.upload()
+
+        # THEN the file should be uploaded and return a new instance with uploaded source
+        assert result.src == f"vellum:uploaded-file:{uploaded_file_id}"
+        assert isinstance(result, TestDocument)
+        mock_client.uploaded_files.create.assert_called_once()
+
+
+def test_vellum_file_mixin_get_signed_url_method():
+    """Test that VellumFileMixin.get_signed_url() method works."""
+
+    # GIVEN a test document with base64 content using the mixin (like postprocess does)
+    content = b"Hello, this is a test document!"
+    base64_content = base64.b64encode(content).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    test_file = TestDocument(src=src)
+    uploaded_file_id = "12345678-1234-1234-1234-123456789abc"
+    signed_url = "https://storage.vellum.ai/files/signed?token=abc123"
+
+    # WHEN getting the signed URL using the mixin's .get_signed_url() method
+    with patch("vellum.utils.vellum_client.create_vellum_client") as mock_create_client:
+        mock_client = Mock()
+        mock_client.uploaded_files.create.return_value = Mock(id=uploaded_file_id)
+        mock_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+        mock_create_client.return_value = mock_client
+
+        result = test_file.get_signed_url()
+
+        # THEN the file should be uploaded and signed URL returned
+        assert result == signed_url
+        mock_client.uploaded_files.create.assert_called_once()
+        mock_client.uploaded_files.retrieve.assert_called_once_with(uploaded_file_id, expiry_seconds=604800)
+
+
+def test_vellum_file_mixin_get_signed_url_method_with_custom_client():
+    """Test that VellumFileMixin.get_signed_url() respects custom vellum_client parameter."""
+
+    # GIVEN a test document with base64 content using the mixin (like postprocess does)
+    content = b"Hello, this is a test document!"
+    base64_content = base64.b64encode(content).decode("utf-8")
+    src = f"data:text/plain;base64,{base64_content}"
+    test_file = TestDocument(src=src)
+    uploaded_file_id = "aaaa1111-bbbb-2222-cccc-3333dddd4444"
+    signed_url = "https://custom.storage.com/signed"
+
+    # Create custom client
+    custom_client = Mock()
+    custom_client.uploaded_files.create.return_value = Mock(id=uploaded_file_id)
+    custom_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+
+    # WHEN getting the signed URL using the mixin with a custom client
+    result = test_file.get_signed_url(vellum_client=custom_client)
+
+    # THEN the custom client should be used
+    assert result == signed_url
+    custom_client.uploaded_files.create.assert_called_once()
+    custom_client.uploaded_files.retrieve.assert_called_once_with(uploaded_file_id, expiry_seconds=604800)
+
+
+def test_vellum_file_mixin_get_signed_url_method_already_uploaded():
+    """Test that VellumFileMixin.get_signed_url() works for already uploaded files."""
+
+    # GIVEN a test document that's already uploaded using the mixin
+    file_id = "12345678-1234-1234-1234-123456789abc"
+    src = f"vellum:uploaded-file:{file_id}"
+    test_file = TestDocument(src=src)
+    signed_url = "https://storage.vellum.ai/files/signed"
+
+    # WHEN getting the signed URL using the mixin's .get_signed_url() method
+    with patch("vellum.utils.vellum_client.create_vellum_client") as mock_create_client:
+        mock_client = Mock()
+        mock_client.uploaded_files.retrieve.return_value = Mock(file_url=signed_url)
+        mock_create_client.return_value = mock_client
+
+        result = test_file.get_signed_url()
+
+        # THEN the signed URL should be returned without uploading again
+        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