Merge pull request #16 from workflowai/guillaume/fix-file-type

Add url to file type

Author: guillaq

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "workflowai"
-version = "0.3.1"
+version = "0.3.2"
 description = ""
 authors = ["Guillaume Aquilina <guillaume@workflowai.com>"]
 readme = "README.md"

workflowai/core/fields/file.py

@@ -1,10 +1,33 @@
-from pydantic import BaseModel, Field
+from typing import Optional
+
+from pydantic import BaseModel, Field, model_validator
 
 
 class File(BaseModel):
-    name: str = Field(..., description="An optional name for the image")
-    content_type: str = Field(..., description="The content type of the image", examples=["image/png", "image/jpeg"])
-    data: str = Field(..., description="The base64 encoded data of the image")
+    name: Optional[str] = Field(
+        default=None,
+        description="An optional name for the file [no longer used]",
+        deprecated=True,
+    )
+    content_type: Optional[str] = Field(
+        default=None,
+        description="The content type of the file. Not needed if content type can be inferred from the URL.",
+        examples=["image/png", "image/jpeg"],
+    )
+    data: Optional[str] = Field(
+        default=None,
+        description="The base64 encoded data of the file. Required if no URL is provided.",
+    )
+    url: Optional[str] = Field(
+        default=None,
+        description="The URL of the file. Required if no data is provided.",
+    )
 
     def to_url(self) -> str:
         return f"data:{self.content_type};base64,{self.data}"
+
+    @model_validator(mode="after")
+    def validate_data_or_url(self):
+        if self.url is None and (self.data is None or self.content_type is None):
+            raise ValueError("Either data or url must be provided")
+        return self

workflowai/core/fields/file_test.py

@@ -0,0 +1,21 @@
+import pytest
+from pydantic import ValidationError
+
+from workflowai.core.fields.file import File
+
+
+class TestFile:
+    def test_validate_data_or_url(self):
+        with pytest.raises(ValidationError):
+            File(data=None, url=None)
+
+    def test_with_valid_data(self):
+        file = File(data="data", content_type="image/png")
+        assert file.data == "data"
+        assert file.content_type == "image/png"
+        assert file.url is None
+
+    def test_with_valid_url(self):
+        file = File(url="https://example.com/image.png")
+        assert file.data is None
+        assert file.content_type is None