diff --git a/src/xai_sdk/aio/files.py b/src/xai_sdk/aio/files.py index 655c5fc..1260853 100644 --- a/src/xai_sdk/aio/files.py +++ b/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,17 +94,15 @@ 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): @@ -106,9 +110,21 @@ async def upload( 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. diff --git a/src/xai_sdk/sync/files.py b/src/xai_sdk/sync/files.py index d0fad71..87e5b89 100644 --- a/src/xai_sdk/sync/files.py +++ b/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,17 +93,15 @@ 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): @@ -105,9 +109,21 @@ def upload( 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. diff --git a/tests/aio/files_test.py b/tests/aio/files_test.py index e6a08db..e0999bd 100644 --- a/tests/aio/files_test.py +++ b/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 diff --git a/tests/sync/files_test.py b/tests/sync/files_test.py index c3a89b1..48527e2 100644 --- a/tests/sync/files_test.py +++ b/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