refactor: various small improvements

Author: gustavocidornelas

src/openlayer/lib/data/_upload.py

@@ -97,7 +97,7 @@ def upload_bytes_multipart(
     data: Union[bytes, BinaryIO],
     object_name: str,
     content_type: str,
-    fields: Optional[Dict] = None,
+    fields: Dict[str, str] = {},
 ) -> Response:
     """Upload data using multipart POST (for S3 and local storage).
 
@@ -115,8 +115,7 @@ def upload_bytes_multipart(
     if isinstance(data, bytes):
         data = io.BytesIO(data)
 
-    upload_fields = dict(fields) if fields else {}
-    upload_fields["file"] = (object_name, data, content_type)
+    upload_fields = {"file": (object_name, data, content_type), **fields}
 
     encoder = MultipartEncoder(fields=upload_fields)
     headers = {"Content-Type": encoder.content_type}
@@ -136,7 +135,7 @@ def upload_bytes_put(
     url: str,
     data: Union[bytes, BinaryIO],
     content_type: str,
-    extra_headers: Optional[Dict[str, str]] = None,
+    extra_headers: Dict[str, str] = {},
 ) -> Response:
     """Upload data using PUT request (for GCS and Azure).
 
@@ -149,9 +148,7 @@ def upload_bytes_put(
     Returns:
         The response from the upload request.
     """
-    headers = {"Content-Type": content_type}
-    if extra_headers:
-        headers.update(extra_headers)
+    headers = {"Content-Type": content_type, **extra_headers}
 
     response = requests.put(
         url,

src/openlayer/lib/integrations/openai_tracer.py

@@ -400,10 +400,7 @@ def handle_non_streaming_create(
         processed_messages = extract_chat_completion_messages(kwargs["messages"])
 
         # Check if response contains audio (to sanitize raw_output)
-        has_audio = (
-            hasattr(response.choices[0].message, "audio")
-            and response.choices[0].message.audio is not None
-        )
+        has_audio = bool(getattr(response.choices[0].message, "audio", None))
 
         # Sanitize raw_output to remove heavy base64 data already uploaded as attachments
         raw_output = response.model_dump()
@@ -481,11 +478,9 @@ def extract_chat_completion_messages(
 
         # If content is a list, process each item
         if isinstance(content, list):
-            normalized_content: List[ContentItem] = []
-            for item in content:
-                content_item = _normalize_content_item(item)
-                normalized_content.append(content_item)
-
+            normalized_content: List[ContentItem] = [
+                _normalize_content_item(item) for item in content
+            ]
             processed_msg = {"role": role, "content": normalized_content}
             if name:
                 processed_msg["name"] = name
@@ -594,7 +589,8 @@ def _normalize_content_item(item: Dict[str, Any]) -> ContentItem:
 
     else:
         # Unknown type, return as TextContent with string representation
-        logger.debug("Unknown content item type: %s", item_type)
+        # Log as warning so we can learn about new OpenAI content types
+        logger.warning("Unknown content item type '%s', preserving as text", item_type)
         return TextContent(text=str(item))
 
 
@@ -874,9 +870,7 @@ def handle_responses_non_streaming_create(
         if hasattr(response, "model_dump"):
             raw_output = response.model_dump()
             if has_generated_images:
-                raw_output = _sanitize_raw_output(
-                    raw_output, has_generated_images=True
-                )
+                raw_output = _sanitize_raw_output(raw_output, has_generated_images=True)
         else:
             raw_output = str(response)
 
@@ -1114,6 +1108,10 @@ def parse_responses_output_data(
                     )
                     content_items.append(ImageContent(attachment=attachment))
 
+            else:
+                # Unknown output type - log for future support
+                logger.debug("Unknown Responses API output type: %s", output_type)
+
         # Return appropriate format based on content
         if len(content_items) > 1:
             # Multimodal output - return list of ContentItems
@@ -1127,7 +1125,7 @@ def parse_responses_output_data(
                 return content_items
 
     except Exception as e:
-        logger.debug("Could not parse Responses API output data: %s", e)
+        logger.warning("Could not parse Responses API output data: %s", e)
 
     return None
 
@@ -1230,30 +1228,46 @@ def _sanitize_raw_output(
 
     sanitized = copy.deepcopy(raw_output)
 
-    # Clear audio data from Chat Completions response
     if has_audio:
-        try:
-            for choice in sanitized.get("choices", []):
-                message = choice.get("message", {})
-                if message and "audio" in message and message["audio"]:
-                    if "data" in message["audio"]:
-                        message["audio"]["data"] = "[UPLOADED_TO_STORAGE]"
-        except Exception as e:
-            logger.debug("Could not sanitize audio data from raw_output: %s", e)
+        _sanitize_audio_data(sanitized)
 
-    # Clear image data from Responses API
     if has_generated_images:
-        try:
-            for output_item in sanitized.get("output", []):
-                if output_item.get("type") == "image_generation_call":
-                    if "result" in output_item:
-                        output_item["result"] = "[UPLOADED_TO_STORAGE]"
-        except Exception as e:
-            logger.debug("Could not sanitize image data from raw_output: %s", e)
+        _sanitize_image_data(sanitized)
 
     return sanitized
 
 
+def _sanitize_audio_data(sanitized: Dict[str, Any]) -> None:
+    """Remove audio base64 data from Chat Completions response.
+
+    Args:
+        sanitized: The raw output dict to sanitize (modified in place)
+    """
+    try:
+        for choice in sanitized.get("choices", []):
+            message = choice.get("message", {})
+            if message and "audio" in message and message["audio"]:
+                if "data" in message["audio"]:
+                    message["audio"]["data"] = "[UPLOADED_TO_STORAGE]"
+    except Exception as e:
+        logger.debug("Could not sanitize audio data from raw_output: %s", e)
+
+
+def _sanitize_image_data(sanitized: Dict[str, Any]) -> None:
+    """Remove generated image base64 data from Responses API output.
+
+    Args:
+        sanitized: The raw output dict to sanitize (modified in place)
+    """
+    try:
+        for output_item in sanitized.get("output", []):
+            if output_item.get("type") == "image_generation_call":
+                if "result" in output_item:
+                    output_item["result"] = "[UPLOADED_TO_STORAGE]"
+    except Exception as e:
+        logger.debug("Could not sanitize image data from raw_output: %s", e)
+
+
 def parse_non_streaming_output_data(
     response: "openai.types.chat.chat_completion.ChatCompletion",
 ) -> Union[str, List[ContentItem], Dict[str, Any], None]:

src/openlayer/lib/tracing/attachment_uploader.py

@@ -13,8 +13,8 @@
     from .steps import Step
     from .traces import Trace
 
-from .attachments import Attachment
 from ..data._upload import STORAGE, StorageType, upload_bytes
+from .attachments import Attachment
 
 logger = logging.getLogger(__name__)
 
@@ -59,16 +59,16 @@ class AttachmentUploader:
     S3, GCS, Azure, and local storage backends.
     """
 
-    def __init__(self, client: "Openlayer", storage: Optional[StorageType] = None):
+    def __init__(self, client: "Openlayer", storage: StorageType = STORAGE):
         """Initialize the attachment uploader.
 
         Args:
             client: The Openlayer client instance.
             storage: Storage type override. Defaults to the global STORAGE setting.
         """
         self._client = client
-        self._storage = storage or STORAGE
-        self._upload_cache: Dict[str, str] = {}  # checksum -> storage_uri
+        self._storage = storage
+        self._storage_uri_cache: Dict[str, str] = {}  # checksum -> storage_uri
 
     def upload_attachment(self, attachment: "Attachment") -> "Attachment":
         """Upload a single attachment if needed.
@@ -104,8 +104,11 @@ def upload_attachment(self, attachment: "Attachment") -> "Attachment":
             return attachment
 
         # Check cache by checksum for deduplication
-        if attachment.checksum_md5 and attachment.checksum_md5 in self._upload_cache:
-            attachment.storage_uri = self._upload_cache[attachment.checksum_md5]
+        if (
+            attachment.checksum_md5
+            and attachment.checksum_md5 in self._storage_uri_cache
+        ):
+            attachment.storage_uri = self._storage_uri_cache[attachment.checksum_md5]
             logger.debug(
                 "Using cached storage_uri for attachment %s (checksum: %s)",
                 attachment.name,
@@ -146,7 +149,9 @@ def upload_attachment(self, attachment: "Attachment") -> "Attachment":
 
             # Cache for deduplication
             if attachment.checksum_md5:
-                self._upload_cache[attachment.checksum_md5] = attachment.storage_uri
+                self._storage_uri_cache[attachment.checksum_md5] = (
+                    attachment.storage_uri
+                )
 
             # Clear data after upload (no longer needed, avoid duplicating in JSON)
             attachment._pending_bytes = None
@@ -213,11 +218,11 @@ def upload_trace_attachments(self, trace: "Trace") -> int:
         Returns:
             The number of attachments uploaded.
         """
-        upload_count = 0
         seen_ids: set = set()
 
-        def process_step(step: "Step") -> None:
-            nonlocal upload_count
+        def process_step(step: "Step") -> int:
+            """Process a step and return the number of attachments uploaded."""
+            step_upload_count = 0
 
             # Collect attachments from all sources
             all_attachments: List[Attachment] = list(step.attachments)
@@ -233,14 +238,15 @@ def process_step(step: "Step") -> None:
                 if not attachment.is_uploaded() and attachment.has_data():
                     self.upload_attachment(attachment)
                     if attachment.is_uploaded():
-                        upload_count += 1
+                        step_upload_count += 1
 
             # Process nested steps recursively
             for nested_step in step.steps:
-                process_step(nested_step)
+                step_upload_count += process_step(nested_step)
+
+            return step_upload_count
 
-        for step in trace.steps:
-            process_step(step)
+        upload_count = sum(process_step(step) for step in trace.steps)
 
         if upload_count > 0:
             logger.info("Uploaded %d attachment(s) for trace", upload_count)

src/openlayer/lib/tracing/steps.py

@@ -93,15 +93,20 @@ def attach(
                 name=name,
                 media_type=media_type,
             )
-        else:
-            # File-like object
+        elif hasattr(data, "read"):
+            # File-like object (BinaryIO, BufferedReader, etc.)
             file_bytes = data.read()
             inferred_name = name or getattr(data, "name", None) or "attachment"
             attachment = Attachment.from_bytes(
                 data=file_bytes,
                 name=inferred_name,
                 media_type=media_type or "application/octet-stream",
             )
+        else:
+            raise TypeError(
+                f"Unsupported data type for attach(): {type(data).__name__}. "
+                "Expected bytes, str, Path, file-like object, or Attachment."
+            )
 
         if metadata:
             attachment.metadata.update(metadata)