packages/traceloop-sdk/tests/dataset/test_file_operations.py (1)

67-105: Ensure temporary file is always cleaned up with try/finally

If set_file_cell or the assertions fail, os.unlink(test_file.name) won’t run and the temp file may be left behind. Wrapping the test body in a try/finally improves robustness.

-    # Create test file
-    test_file = tempfile.NamedTemporaryFile(suffix=".txt", delete=False)
-    test_file.write(b"test content")
-    test_file.close()
-
-    # Mock S3 upload to succeed
-    with patch.object(row, "_upload_file_to_presigned_url", return_value=True):
-        row.set_file_cell(
-            column_slug="file",
-            file_path=test_file.name,
-            file_type=FileCellType.FILE,
-        )
-
-    # Verify
-    assert row.values["file"]["storage"] == "internal"
-    assert row.values["file"]["status"] == "success"
-    assert "storage_key" in row.values["file"]
-
-    # Clean up
-    os.unlink(test_file.name)
+    # Create test file
+    test_file = tempfile.NamedTemporaryFile(suffix=".txt", delete=False)
+    try:
+        test_file.write(b"test content")
+        test_file.close()
+
+        # Mock S3 upload to succeed
+        with patch.object(row, "_upload_file_to_presigned_url", return_value=True):
+            row.set_file_cell(
+                column_slug="file",
+                file_path=test_file.name,
+                file_type=FileCellType.FILE,
+            )
+
+        # Verify
+        assert row.values["file"]["storage"] == "internal"
+        assert row.values["file"]["status"] == "success"
+        assert "storage_key" in row.values["file"]
+    finally:
+        os.unlink(test_file.name)

packages/traceloop-sdk/traceloop/sdk/dataset/row.py (2)

223-231: Act on _confirm_file_upload result when marking uploads failed/successful

In both the failure path (Line 228–230) and the final confirmation (Line 249–251), _confirm_file_upload is called but its boolean return value is ignored. If the backend returns {"success": false} without raising an HTTP error, the SDK will still proceed as if the operation succeeded.

Consider:

• Checking the return value and raising an exception (or at least logging) when it is False.
• Optionally returning early or not updating self.values[column_slug] if confirmation fails.

For example:

-            if not success:
-                self._confirm_file_upload(column_slug, status="failed", metadata=metadata)
-                raise Exception(f"Failed to upload file to S3: {file_name}")
+            if not success:
+                confirmed = self._confirm_file_upload(
+                    column_slug, status="failed", metadata=metadata
+                )
+                if not confirmed:
+                    raise Exception(
+                        f"Failed to upload file to S3 and backend did not confirm failure for cell {column_slug}"
+                    )
+                raise Exception(f"Failed to upload file to S3: {file_name}")
@@
-            self._confirm_file_upload(
-                column_slug=column_slug, status="success", metadata=final_metadata
-            )
+            if not self._confirm_file_upload(
+                column_slug=column_slug, status="success", metadata=final_metadata
+            ):
+                raise Exception(
+                    f"Upload confirmation failed for cell {column_slug} after successful upload"
+                )

Also applies to: 249-251

---

232-244: Use logging instead of print for thumbnail warnings

The warnings for missing thumbnail file and failed thumbnail upload currently use print, which makes it harder to route or filter in real environments.

Prefer the standard logging module, e.g.:

-import os
+import os
+import logging
@@
-                if not os.path.exists(thumbnail_path):
-                    print(f"Warning: Thumbnail file not found: {thumbnail_path}")
+                if not os.path.exists(thumbnail_path):
+                    logging.warning("Thumbnail file not found: %s", thumbnail_path)
@@
-                    if not thumb_success:
-                        print(f"Warning: Failed to upload thumbnail for {file_name}")
+                    if not thumb_success:
+                        logging.warning(
+                            "Failed to upload thumbnail for %s", file_name
+                        )

packages/traceloop-sdk/traceloop/sdk/dataset/model.py (1)

12-72: File-cell models are consistent; consider tightening status representation

The new FileCellType, FileStorageType, and related models (FileCellMetadata, FileCellValue, UploadURLRequest/Response, UploadStatusRequest, ExternalURLRequest) line up with the row logic and test expectations, so the overall shape looks solid.

If you want stronger validation, you might later replace the free-form status: str fields with a small enum or shared constants (e.g., "in_progress", "success", "failed") to avoid accidental typos, but that’s not required for correctness.

packages/traceloop-sdk/traceloop/sdk/dataset/row.py (2)

122-149: Optional: derive local file-cell state from API response

_set_external_file_url currently ignores the server response (beyond None/error) and reconstructs the local cell dict manually. If the backend ever adds fields (e.g., normalized URL, derived metadata), this method might drift from the canonical representation.

If the API returns the full cell payload, consider updating self.values[column_slug] from that response instead of hand-building the dict, or at least centralizing the shape in a single helper to avoid duplication.

---

245-261: Avoid mutating caller-provided metadata when adding size_bytes

final_metadata = metadata or {} followed by final_metadata["size_bytes"] = ... will mutate the original dict when metadata is provided by the caller, which can be surprising if that dict is reused elsewhere.

Consider copying:

-            final_metadata = metadata or {}
-            final_metadata["size_bytes"] = os.path.getsize(file_path)
+            final_metadata = {**(metadata or {})}
+            final_metadata["size_bytes"] = os.path.getsize(file_path)

This keeps the API safer while preserving behavior.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Disabled knowledge base sources:

• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: nina-kollman
Repo: traceloop/openllmetry PR: 3219
File: packages/traceloop-sdk/traceloop/sdk/datasets/dataset.py:357-360
Timestamp: 2025-08-04T15:35:30.188Z
Learning: In the traceloop SDK Dataset class, the established error handling pattern is that HTTP client methods return None on failure (after catching and logging RequestException), and the Dataset API methods check for None return values and raise Exception with descriptive messages. The update_row_api method is inconsistent with this pattern.

packages/traceloop-sdk/tests/dataset/cassettes/test_dataset_with_attachments/test_create_dataset_with_mixed_attachments.yaml (1)

85-85: Verify AWS credentials in presigned URLs are scrubbed.

Similar to test_create_dataset_with_file_attachments_mocked.yaml, this cassette contains AWS S3 presigned URLs with temporary security tokens. Per coding guidelines, ensure these are properly scrubbed.

packages/traceloop-sdk/tests/dataset/cassettes/test_attachments/test_attachment_upload_with_mock_s3.yaml (1)

85-85: Verify AWS credentials in presigned URLs are scrubbed.

This cassette also contains AWS S3 presigned URLs with temporary security tokens. Ensure compliance with the coding guideline to never commit secrets in VCR cassettes.

Also applies to: 261-261

packages/traceloop-sdk/traceloop/sdk/dataset/__init__.py (1)

8-20: Consider sorting __all__ for consistency.

The static analysis tool suggests applying an isort-style sorting to __all__. While not critical, alphabetical ordering improves maintainability and reduces merge conflicts.

Apply this diff to sort the exports:

 __all__ = [
-    "Dataset",
-    "Column",
-    "Row",
-    "BaseDatasetEntity",
-    "ColumnType",
-    "DatasetMetadata",
-    "FileCellType",
-    "FileStorageType",
     "Attachment",
+    "AttachmentReference",
+    "BaseDatasetEntity",
+    "Column",
+    "ColumnType",
+    "Dataset",
+    "DatasetMetadata",
     "ExternalAttachment",
-    "AttachmentReference",
+    "FileCellType",
+    "FileStorageType",
+    "Row",
 ]

packages/traceloop-sdk/traceloop/sdk/datasets/datasets.py (1)

319-319: Consider using logging instead of print statements.

Library code should use Python's logging module rather than print statements for warnings, allowing consumers to control log output and levels.

Add logging at the module level:

import logging

logger = logging.getLogger(__name__)

Then replace print statements:

-                print(f"Warning: Row index {row_idx} out of range, skipping attachments")
+                logger.warning(f"Row index {row_idx} out of range, skipping attachments")

-                    print(f"Warning: Failed to process attachment for row {row_idx}, column {col_slug}: {e}")
+                    logger.warning(f"Failed to process attachment for row {row_idx}, column {col_slug}: {e}")

Also applies to: 346-346

packages/traceloop-sdk/traceloop/sdk/dataset/attachment.py (1)

465-469: TODO noted: Internal file download URL generation.

The comment acknowledges that getting presigned download URLs for internal files requires backend API support. This is a known limitation.

Do you want me to open a tracking issue for implementing the backend endpoint and completing this functionality?

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Disabled knowledge base sources:

• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/traceloop-sdk/traceloop/sdk/datasets/datasets.py (2)

320-322: Use logging instead of print() for diagnostics.

This warning should use a proper logging mechanism (e.g., logging.warning) for better production diagnostics and log aggregation.

As per past review comments, this has already been flagged.

---

349-351: Use logging instead of print() for diagnostics.

This warning should use a proper logging mechanism (e.g., logging.warning) for better production diagnostics and log aggregation.

As per past review comments, this has already been flagged.

packages/traceloop-sdk/traceloop/sdk/datasets/datasets.py (1)

288-297: Consider using copy.deepcopy for truly nested-safe copying.

The manual iteration creates a shallow copy for non-attachment values. If row values contain nested dicts or lists (e.g., {"col": {"nested": "value"}}), the nested objects would be shared between the original and cleaned request. This could lead to unintended mutations if the original request is modified elsewhere.

If nested values are possible, apply this diff:

+        import copy
+
         # Create a deep copy of rows to avoid modifying the original
-        clean_rows = []
-        for row in request.rows:
-            clean_row = {}
-            for col_slug, value in row.items():
-                if isinstance(value, (Attachment, ExternalAttachment)):
-                    clean_row[col_slug] = None
-                else:
-                    clean_row[col_slug] = value
-            clean_rows.append(clean_row)
+        clean_rows = copy.deepcopy(request.rows)
+        for row in clean_rows:
+            for col_slug, value in list(row.items()):
+                if isinstance(value, (Attachment, ExternalAttachment)):
+                    row[col_slug] = None

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Disabled knowledge base sources:

• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/traceloop-sdk/traceloop/sdk/dataset/attachment.py (2)

109-179: Add timeouts and tighten error handling for upload and thumbnail requests

The requests.put calls for the main file upload (Line 173) and thumbnail upload (Line 150) have no timeout, so a hung S3/presigned endpoint can block indefinitely. Static analysis (S113) and prior review comments already flagged this. Also, _upload_to_s3 swallows all exceptions with except Exception, losing the root cause.

I recommend:

• Add explicit timeouts for both uploads (e.g., 300s for main file, 60s for thumbnail).
• Narrow the exception handler in _upload_to_s3 to requests.RequestException so file/OS errors still surface.

Example diff:

-            if thumb_bytes is not None:
-                requests.put(upload_response.thumbnail_upload_url, data=thumb_bytes)
+            if thumb_bytes is not None:
+                requests.put(
+                    upload_response.thumbnail_upload_url,
+                    data=thumb_bytes,
+                    timeout=60,
+                )

@@
-        try:
-            file_data = self._get_file_data()
-            response = requests.put(
-                upload_url, data=file_data, headers={"Content-Type": self.content_type}
-            )
-            return response.status_code in [200, 201, 204]
-        except Exception:
-            return False
+        try:
+            file_data = self._get_file_data()
+            response = requests.put(
+                upload_url,
+                data=file_data,
+                headers={"Content-Type": self.content_type},
+                timeout=300,
+            )
+            return response.status_code in [200, 201, 204]
+        except requests.RequestException:
+            return False

Optionally, you might also check the result of the upload-status http_client.put and raise if it fails, so callers are aware when confirmation does not persist.

---

255-264: Add timeout and consider a dedicated helper/exception for downloads

The data property issues a raw requests.get(download_url) without a timeout (Line 261), so a stuck endpoint can hang indefinitely. This was previously flagged for _download_data; the same concern applies here, and Ruff reports S113.

Suggested change:

-            response = requests.get(download_url)
+            response = requests.get(download_url, timeout=300)
             response.raise_for_status()
             self._cached_data = response.content

You might also want to:

• Extract this into a small _download_data() helper (as you previously had) to centralize future tweaks (retries, streaming, etc.).
• Replace raise Exception("No download URL available") with a more specific custom exception (e.g., AttachmentError) to make error handling clearer.

packages/traceloop-sdk/traceloop/sdk/dataset/attachment.py (3)

90-107: Consider streaming large files instead of fully loading into memory

_get_file_data and _get_file_size read the entire file into memory (f.read() / len(self.data)), which is fine for small/medium attachments but can become costly for large files. A future improvement would be to support streaming (e.g., passing a file-like object/chunk iterator to the upload path) to avoid holding the whole file in memory at once. This can be deferred if datasets are expected to carry relatively small blobs.

---

186-227: Clarify use of content_type/filename and server response in ExternalAttachment.attach

A couple of minor design points here:

• content_type and filename are stored on the instance but never sent to the API or surfaced in the returned AttachmentReference. If those fields are meant to influence server-side behavior or downstream display, consider wiring them into metadata or the request model; otherwise they are effectively unused.
• The attach method only checks that result is truthy and then ignores any details the server might return (e.g., a storage key or normalized URL). If the API ever exposes richer information, you may want to capture it in the AttachmentReference instead of just echoing self.url and self.metadata.

These are not blockers but worth aligning with how you expect external attachments to be represented long-term.

---

276-281: Internal attachment downloads are currently unsupported; consider making this explicit

get_url returns the external URL for FileStorageType.EXTERNAL, but always None for internal storage, which makes AttachmentReference.data unusable for internal attachments created by Attachment.upload (it will raise “No download URL available”).

If that’s intentional for now, consider:

• Documenting explicitly that AttachmentReference.data/download() only work for external attachments today, or
• Raising a more explicit NotImplementedError/custom exception for INTERNAL storage, or
• Implementing presigned URL generation for internal storage using storage_key and http_client when available.

This will make the API behavior less surprising for callers who obtain an AttachmentReference from Attachment.upload.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Disabled knowledge base sources:

• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/traceloop-sdk/examples/dataset_attachments_example.py (3)

1-10: Align shebang usage with how the script is invoked

Ruff flags EXE001 because the file has a shebang but (in the repo) isn’t executable. If this script is meant to be run as python dataset_attachments_example.py, consider dropping the shebang; if you want ./dataset_attachments_example.py, make sure the file is committed with the executable bit.

Example change if you choose to remove it:

-#!/usr/bin/env python3

---

28-35: Consider centralizing Traceloop.init in main() instead of inside one example

Right now only example_external_attachments() calls Traceloop.init, so calling other example_* functions directly (e.g., from a REPL) will likely fail if Traceloop wasn’t initialized elsewhere. Initializing once in main() and assuming it in the examples would be more robust and clearer.

Suggested change inside example_external_attachments:

-    # Initialize Traceloop
-    Traceloop.init(app_name="attachment-demo")
-    datasets = Datasets()
+    # Datasets client (assumes Traceloop.init(...) was called by the caller)
+    datasets = Datasets()

And then in main() (outside this hunk), initialize once:

def main():
    """Run all examples."""
    print("=" * 60)
    print("Traceloop SDK Attachment Feature Examples")
    print("=" * 60)

    Traceloop.init(app_name="attachment-demo")
    ...

---

334-336: Document API key configuration without implying hardcoding in code

To better align with the “no hardcoded secrets” guideline, it’s clearer to show setting the API key via the shell environment rather than via os.environ in code, even in comments.

-    # Set your API key
-    # os.environ["TRACELOOP_API_KEY"] = "your-api-key-here"
+    # Ensure TRACELOOP_API_KEY is set in your environment, e.g.:
+    # export TRACELOOP_API_KEY="your-api-key-here"

As per coding guidelines, this keeps secrets configuration outside the code.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Disabled knowledge base sources:

• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/traceloop-sdk/examples/dataset_attachments_example.py (1)

339-352: Broad except Exception in main still trips BLE001 (repeat of prior feedback)

The catch‑all around the example runner is reasonable for UX, but it continues to trigger Ruff BLE001 and was already flagged in an earlier review.

To keep the friendly error while satisfying lint, you can explicitly document/suppress the broad catch:

-    except Exception as e:
-        print(f"\nError: {e}")
-        print("Make sure to set TRACELOOP_API_KEY environment variable")
+    except Exception as e:  # noqa: BLE001  # Broad catch is intentional for this example script
+        print(f"\nError: {e}")
+        print("Make sure to set TRACELOOP_API_KEY environment variable")

If the SDK exposes specific exception types for auth/config errors, narrowing the except to those would be even better.

packages/traceloop-sdk/traceloop/sdk/datasets/datasets.py (2)

276-309: Preserve future CreateDatasetRequest fields when cleaning rows

Right now _prepare_request_for_creation reconstructs CreateDatasetRequest by manually listing fields. If new fields are later added to CreateDatasetRequest (e.g., tags, options), they’ll be silently dropped here.

You can make this future‑proof by cloning the model and only updating rows:

-        # Create a new request with cleaned rows
-        return CreateDatasetRequest(
-            slug=request.slug,
-            name=request.name,
-            description=request.description,
-            columns=request.columns,
-            rows=clean_rows,
-        )
+        # Create a new request with cleaned rows, preserving any additional fields
+        return request.model_copy(update={"rows": clean_rows})
+        # If you are still on Pydantic < 2, use instead:
+        # return request.copy(update={"rows": clean_rows})

Please pick model_copy vs copy according to the Pydantic version actually used in this project.

---

311-360: Tighten or explicitly document the broad Exception catch in _process_attachments

Catching bare Exception here triggers Ruff BLE001 and can hide unexpected programmer errors, even though you likely want best‑effort behavior per cell.

A balanced approach is to (a) keep the catch‑all for UX, (b) use structured logging rather than an f-string, and (c) explicitly mark the broad catch as intentional:

-                except Exception as e:
-                    logger.warning(
-                        f"Warning: Failed to process attachment for row {row_idx}, column {col_slug}: {e}"
-                    )
+                except Exception as e:  # noqa: BLE001  # Broad catch is intentional: attachment failures shouldn't break dataset creation
+                    logger.warning(
+                        "Failed to process attachment for row %s, column %s: %s",
+                        row_idx,
+                        col_slug,
+                        e,
+                    )

If the attachment methods expose well-defined exception types, narrowing the except clause to those would be even better; otherwise the above makes the catch‑all explicit and lint‑clean.

packages/traceloop-sdk/examples/dataset_attachments_example.py (1)

1-1: Address shebang vs executability (EXE001)

Ruff flags the shebang because this file isn’t marked executable. For a library example that’s typically run via python dataset_attachments_example.py, it’s simplest to drop the shebang:

-#!/usr/bin/env python3

Alternatively, ensure the file is installed with executable permissions so the shebang is meaningful.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Disabled knowledge base sources:

• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/traceloop-sdk/traceloop/sdk/datasets/datasets.py (2)

282-315: Consider using the request model’s copy functionality instead of rebuilding by hand.

Manually reconstructing CreateDatasetRequest with only slug, name, description, columns, and rows makes _prepare_request_for_creation brittle if the model gains new fields (e.g., tags, settings) in the future. You can future‑proof this by cloning and updating just rows via the Pydantic copy API instead of re‑passing individual fields.

For example (adapt to your Pydantic version):

# Pydantic v1 style:
clean_request = request.copy(update={"rows": clean_rows})

# or Pydantic v2 style:
clean_request = request.model_copy(update={"rows": clean_rows})

This keeps all other fields intact while still stripping attachments from rows.

---

317-366: Align broad except Exception with Ruff BLE001 and improve logging detail.

Catching a bare Exception here is intentional (best‑effort attachment processing), but Ruff flags this as BLE001. Making this explicit and using parameterized logging would both satisfy the linter and improve diagnostics.

One option that preserves behavior:

-                except Exception as e:
-                    logger.warning(
-                        f"Warning: Failed to process attachment for row {row_idx}, column {col_slug}: {e}"
-                    )
+                except Exception as e:  # noqa: BLE001
+                    logger.warning(
+                        "Failed to process attachment for row %s, column %s: %s",
+                        row_idx,
+                        col_slug,
+                        e,
+                    )

This keeps the “catch‑all” semantics while making the linter suppression explicit and using logger‑style formatting. If you know the concrete exception types upload/attach can raise, narrowing the except to those would be even better.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Disabled knowledge base sources:

• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

packages/traceloop-sdk/traceloop/sdk/datasets/datasets.py (2)

262-279: LGTM! Attachment extraction logic is sound.

The method correctly identifies and extracts Attachment and ExternalAttachment objects from row values.

Optional: Consider a more specific return type for improved type safety:

-    ) -> Dict[int, Dict[str, Any]]:
+    ) -> Dict[int, Dict[str, Union[Attachment, ExternalAttachment]]]:

You would need to add Union to the typing imports.

---

316-365: Consider more specific exception handling for better error visibility.

The broad Exception catch on line 356 enables resilient bulk processing (logging failures and continuing), but it may mask unexpected errors.

Consider catching more specific exceptions or adding an exception type check:

-                except Exception as e:
+                except (IOError, OSError, RuntimeError, ValueError) as e:
                     logger.warning(
                         f"Warning: Failed to process attachment for row {row_idx}, column {col_slug}: {e}"
                     )

If you need to catch all exceptions for resilience, document the rationale:

-                except Exception as e:
+                except Exception as e:  # Catch all to continue processing remaining attachments
                     logger.warning(
                         f"Warning: Failed to process attachment for row {row_idx}, column {col_slug}: {e}"
                     )

Alternatively, consider collecting and returning failed attachments so callers can decide how to handle them.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Disabled knowledge base sources:

• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.