Truncate image data logs

src/exo/shared/types/api.py

@@ -1,4 +1,5 @@
 import time
+from collections.abc import Generator
 from typing import Any, Literal
 
 from fastapi import UploadFile
@@ -223,12 +224,26 @@ class ImageEditsInternalParams(BaseModel):
     size: str | None = "1024x1024"
     image_strength: float = 0.7
 
+    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+        for name, value in super().__repr_args__():
+            if name == "image_data":
+                yield name, f"<{len(self.image_data)} chars>"
+            elif name is not None:
+                yield name, value
+
 
 class ImageData(BaseModel):
     b64_json: str | None = None
     url: str | None = None
     revised_prompt: str | None = None
 
+    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+        for name, value in super().__repr_args__():
+            if name == "b64_json" and value is not None:
+                yield name, f"<{len(value)} chars>"
+            elif name is not None:
+                yield name, value
+
 
 class ImageGenerationResponse(BaseModel):
     created: int = Field(default_factory=lambda: int(time.time()))

src/exo/shared/types/chunks.py

@@ -1,4 +1,6 @@
+from collections.abc import Generator
 from enum import Enum
+from typing import Any
 
 from exo.utils.pydantic_ext import TaggedModel
 
@@ -29,6 +31,13 @@ class ImageChunk(BaseChunk):
     total_chunks: int
     image_index: int
 
+    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+        for name, value in super().__repr_args__():
+            if name == "data":
+                yield name, f"<{len(self.data)} chars>"
+            elif name is not None:
+                yield name, value
+
 
 class InputImageChunk(BaseChunk):
     command_id: CommandId
@@ -36,5 +45,12 @@ class InputImageChunk(BaseChunk):
     chunk_index: int
     total_chunks: int
 
+    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+        for name, value in super().__repr_args__():
+            if name == "data":
+                yield name, f"<{len(self.data)} chars>"
+            elif name is not None:
+                yield name, value
+
 
 GenerationChunk = TokenChunk | ImageChunk

src/exo/shared/types/worker/runner_response.py

@@ -1,4 +1,6 @@
-from typing import Literal
+from collections.abc import Generator
+from typing import Any, Literal
+
 from exo.shared.types.api import FinishReason
 from exo.utils.pydantic_ext import TaggedModel
 
@@ -22,6 +24,13 @@ class ImageGenerationResponse(BaseRunnerResponse):
     image_data: bytes
     format: Literal["png", "jpeg", "webp"] = "png"
 
+    def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
+        for name, value in super().__repr_args__():
+            if name == "image_data":
+                yield name, f"<{len(self.image_data)} bytes>"
+            elif name is not None:
+                yield name, value
+
 
 class FinishedResponse(BaseRunnerResponse):
     pass