Instrument File Methods

src/xai_sdk/aio/files.py

@@ -3,6 +3,8 @@
 from asyncio import Semaphore
 from typing import BinaryIO, Optional, Sequence, Union
 
+from opentelemetry.trace import SpanKind
+
 from ..files import (
     BaseClient,
     BatchUploadCallback,
@@ -16,6 +18,10 @@
     _sort_by_to_pb,
 )
 from ..proto import files_pb2
+from ..telemetry import get_tracer
+from ..telemetry.config import should_disable_sensitive_attributes
+
+tracer = get_tracer(__name__)
 
 
 class Client(BaseClient):
@@ -88,27 +94,37 @@ async def upload(
             if not os.path.exists(file):
                 raise FileNotFoundError(f"File not found: {file}")
             chunks = _async_chunk_file_from_path(file_path=file, progress=on_progress)
-            return await self._stub.UploadFile(chunks)
 
         # Handle bytes
-        if isinstance(file, bytes | bytearray):
+        elif isinstance(file, bytes | bytearray):
             if not filename:
                 raise ValueError("filename is required when uploading bytes")
             chunks = _async_chunk_file_data(filename=filename, data=bytes(file), progress=on_progress)
-            return await self._stub.UploadFile(chunks)
 
         # Handle file-like object (BinaryIO)
-        if hasattr(file, "read"):
+        elif hasattr(file, "read"):
             # Try to get filename from the file object if not provided
             if not filename:
                 if hasattr(file, "name") and isinstance(file.name, str):
                     filename = os.path.basename(file.name)
                 else:
                     raise ValueError("filename is required when uploading a file-like object without a .name attribute")
             chunks = _async_chunk_file_from_fileobj(file_obj=file, filename=filename, progress=on_progress)
-            return await self._stub.UploadFile(chunks)
-
-        raise ValueError(f"Unsupported file type: {type(file)}")
+        else:
+            raise ValueError(f"Unsupported file type: {type(file)}")
+        with tracer.start_as_current_span(
+            name="file.upload",
+            kind=SpanKind.CLIENT,
+            attributes={
+                "operation.name": "upload_file",
+                "provider.name": "xai",
+            },
+        ) as span:
+            res = await self._stub.UploadFile(chunks)
+            if not should_disable_sensitive_attributes():
+                span.set_attribute("file.id", res.id)
+                span.set_attribute("file.filename", res.filename)
+            return res
 
     async def batch_upload(
         self,
@@ -245,7 +261,18 @@ async def delete(self, file_id: str) -> files_pb2.DeleteFileResponse:
             A DeleteFileResponse indicating whether the deletion was successful.
         """
         request = files_pb2.DeleteFileRequest(file_id=file_id)
-        return await self._stub.DeleteFile(request)
+        with tracer.start_as_current_span(
+            name="file.delete",
+            kind=SpanKind.CLIENT,
+            attributes={
+                "operation.name": "delete",
+                "provider.name": "xai",
+            },
+        ) as span:
+            res = await self._stub.DeleteFile(request)
+            if not should_disable_sensitive_attributes():
+                span.set_attribute("file.id", file_id)
+            return res
 
     async def content(self, file_id: str) -> bytes:
         """Get the complete content of a file asynchronously.

src/xai_sdk/sync/files.py

@@ -2,6 +2,8 @@
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import BinaryIO, Optional, Sequence, Union
 
+from opentelemetry.trace import SpanKind
+
 from ..files import (
     BaseClient,
     BatchUploadCallback,
@@ -15,6 +17,10 @@
     _sort_by_to_pb,
 )
 from ..proto import files_pb2
+from ..telemetry import get_tracer
+from ..telemetry.config import should_disable_sensitive_attributes
+
+tracer = get_tracer(__name__)
 
 
 class Client(BaseClient):
@@ -87,27 +93,37 @@ def upload(
             if not os.path.exists(file):
                 raise FileNotFoundError(f"File not found: {file}")
             chunks = _chunk_file_from_path(file_path=file, progress=on_progress)
-            return self._stub.UploadFile(chunks)
 
         # Handle bytes
-        if isinstance(file, bytes | bytearray):
+        elif isinstance(file, bytes | bytearray):
             if not filename:
                 raise ValueError("filename is required when uploading bytes")
             chunks = _chunk_file_data(filename=filename, data=bytes(file), progress=on_progress)
-            return self._stub.UploadFile(chunks)
 
         # Handle file-like object (BinaryIO)
-        if hasattr(file, "read"):
+        elif hasattr(file, "read"):
             # Try to get filename from the file object if not provided
             if not filename:
                 if hasattr(file, "name") and isinstance(file.name, str):
                     filename = os.path.basename(file.name)
                 else:
                     raise ValueError("filename is required when uploading a file-like object without a .name attribute")
             chunks = _chunk_file_from_fileobj(file_obj=file, filename=filename, progress=on_progress)
-            return self._stub.UploadFile(chunks)
-
-        raise ValueError(f"Unsupported file type: {type(file)}")
+        else:
+            raise ValueError(f"Unsupported file type: {type(file)}")
+        with tracer.start_as_current_span(
+            name="file.upload",
+            kind=SpanKind.CLIENT,
+            attributes={
+                "operation.name": "upload_file",
+                "provider.name": "xai",
+            },
+        ) as span:
+            res = self._stub.UploadFile(chunks)
+            if not should_disable_sensitive_attributes():
+                span.set_attribute("file.id", res.id)
+                span.set_attribute("file.name", res.filename)
+            return res
 
     def batch_upload(
         self,
@@ -246,7 +262,18 @@ def delete(self, file_id: str) -> files_pb2.DeleteFileResponse:
             A DeleteFileResponse indicating whether the deletion was successful.
         """
         request = files_pb2.DeleteFileRequest(file_id=file_id)
-        return self._stub.DeleteFile(request)
+        with tracer.start_as_current_span(
+            name="file.delete",
+            kind=SpanKind.CLIENT,
+            attributes={
+                "operation.name": "delete",
+                "provider.name": "xai",
+            },
+        ) as span:
+            res = self._stub.DeleteFile(request)
+            if not should_disable_sensitive_attributes():
+                span.set_attribute("file.id", res.id)
+            return res
 
     def content(self, file_id: str) -> bytes:
         """Get the complete content of a file.

tests/aio/files_test.py

@@ -8,6 +8,7 @@
 import pytest
 import pytest_asyncio
 from google.protobuf import timestamp_pb2
+from opentelemetry.trace import SpanKind
 
 from xai_sdk import AsyncClient
 from xai_sdk.files import _chunk_file_from_path
@@ -729,3 +730,63 @@ async def test_batch_upload_empty_list(client_with_mock_stub: AsyncClient):
     """Test batch upload with empty file list."""
     with pytest.raises(ValueError):
         await client_with_mock_stub.files.batch_upload([])
+
+
+@mock.patch("xai_sdk.aio.files.tracer")
+@pytest.mark.asyncio
+async def test_upload_creates_span_with_correct_attributes(
+    mock_tracer: mock.MagicMock, client_with_mock_stub: AsyncClient, mock_stub
+):
+    mock_span = mock.MagicMock()
+    mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
+
+    mock_response = files_pb2.File(id="file-123", filename="test.txt", size=12, team_id="team-456")
+
+    async def async_return():
+        return mock_response
+
+    mock_stub.UploadFile.return_value = async_return()
+
+    result = await client_with_mock_stub.files.upload(b"data", filename="test.txt")
+
+    mock_tracer.start_as_current_span.assert_called_once_with(
+        name="file.upload",
+        kind=SpanKind.CLIENT,
+        attributes={
+            "operation.name": "upload_file",
+            "provider.name": "xai",
+        },
+    )
+
+    mock_span.set_attribute.assert_any_call("file.id", result.id)
+    mock_span.set_attribute.assert_any_call("file.filename", result.filename)
+
+
+@mock.patch("xai_sdk.aio.files.tracer")
+@pytest.mark.asyncio
+async def test_delete_creates_span_with_correct_attributes(
+    mock_tracer: mock.MagicMock, client_with_mock_stub: AsyncClient, mock_stub
+):
+    mock_span = mock.MagicMock()
+    mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
+
+    mock_response = files_pb2.DeleteFileResponse(id="file-456", deleted=True)
+
+    async def async_return():
+        return mock_response
+
+    mock_stub.DeleteFile.return_value = async_return()
+
+    result = await client_with_mock_stub.files.delete("file-456")
+
+    mock_tracer.start_as_current_span.assert_called_once_with(
+        name="file.delete",
+        kind=SpanKind.CLIENT,
+        attributes={
+            "operation.name": "delete",
+            "provider.name": "xai",
+        },
+    )
+
+    mock_span.set_attribute.assert_called_once_with("file.id", result.id)
+    assert result.deleted is True

tests/sync/files_test.py

@@ -8,6 +8,7 @@
 
 import pytest
 from google.protobuf import timestamp_pb2
+from opentelemetry.trace import SpanKind
 
 from xai_sdk import Client
 from xai_sdk.files import _chunk_file_data, _chunk_file_from_path, _order_to_pb, _sort_by_to_pb
@@ -668,3 +669,53 @@ def test_batch_upload_empty_list(client_with_mock_stub: Client):
     """Test batch upload with empty file list."""
     with pytest.raises(ValueError):
         client_with_mock_stub.files.batch_upload([])
+
+
+@mock.patch("xai_sdk.sync.files.tracer")
+def test_upload_creates_span_with_correct_attributes(
+    mock_tracer: mock.MagicMock, client_with_mock_stub: Client, mock_stub
+):
+    mock_span = mock.MagicMock()
+    mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
+
+    mock_response = files_pb2.File(id="file-123", filename="test.txt", size=12, team_id="team-456")
+    mock_stub.UploadFile.return_value = mock_response
+
+    result = client_with_mock_stub.files.upload(b"data", filename="test.txt")
+
+    mock_tracer.start_as_current_span.assert_called_once_with(
+        name="file.upload",
+        kind=SpanKind.CLIENT,
+        attributes={
+            "operation.name": "upload_file",
+            "provider.name": "xai",
+        },
+    )
+
+    mock_span.set_attribute.assert_any_call("file.id", result.id)
+    mock_span.set_attribute.assert_any_call("file.name", result.filename)
+
+
+@mock.patch("xai_sdk.sync.files.tracer")
+def test_delete_creates_span_with_correct_attributes(
+    mock_tracer: mock.MagicMock, client_with_mock_stub: Client, mock_stub
+):
+    mock_span = mock.MagicMock()
+    mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
+
+    mock_response = files_pb2.DeleteFileResponse(id="file-456", deleted=True)
+    mock_stub.DeleteFile.return_value = mock_response
+
+    result = client_with_mock_stub.files.delete("file-456")
+
+    mock_tracer.start_as_current_span.assert_called_once_with(
+        name="file.delete",
+        kind=SpanKind.CLIENT,
+        attributes={
+            "operation.name": "delete",
+            "provider.name": "xai",
+        },
+    )
+
+    mock_span.set_attribute.assert_called_once_with("file.id", result.id)
+    assert result.deleted is True