diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py new file mode 100644 index 0000000000..4757122017 --- /dev/null +++ b/tests/test_litellm/llms/anthropic/test_anthropic_files_and_batches.py @@ -0,0 +1,639 @@ +""" +Test Anthropic Files Handler and Batch Retrieval + +Tests for: +1. AnthropicFilesHandler.afile_content() - retrieving batch results +2. AnthropicBatchesConfig.transform_retrieve_batch_response() - transforming batch responses +3. Transformation of Anthropic batch results to OpenAI format +""" + +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +sys.path.insert(0, os.path.abspath("../../../../")) + +import httpx +import pytest + +from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig +from litellm.llms.anthropic.files.handler import AnthropicFilesHandler +from litellm.types.llms.openai import FileContentRequest, HttpxBinaryResponseContent + + +class TestAnthropicFilesHandler: + """Test Anthropic Files Handler for batch results retrieval""" + + @pytest.fixture + def handler(self): + """Create AnthropicFilesHandler instance""" + return AnthropicFilesHandler() + + @pytest.fixture + def mock_anthropic_batch_results_succeeded(self): + """Mock Anthropic batch results with succeeded status""" + return json.dumps({ + "custom_id": "test-request-1", + "result": { + "type": "succeeded", + "message": { + "id": "msg_123", + "model": "claude-3-5-sonnet-20241022", + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Hello, world!" + } + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": { + "input_tokens": 10, + "output_tokens": 5 + } + } + } + }).encode("utf-8") + + @pytest.fixture + def mock_anthropic_batch_results_errored(self): + """Mock Anthropic batch results with errored status""" + return json.dumps({ + "custom_id": "test-request-2", + "result": { + "type": "errored", + "error": { + "error": { + "type": "invalid_request_error", + "message": "Invalid request" + }, + "request_id": "req_456" + } + } + }).encode("utf-8") + + @pytest.fixture + def mock_anthropic_batch_results_canceled(self): + """Mock Anthropic batch results with canceled status""" + return json.dumps({ + "custom_id": "test-request-3", + "result": { + "type": "canceled" + } + }).encode("utf-8") + + @pytest.fixture + def mock_anthropic_batch_results_mixed(self): + """Mock Anthropic batch results with multiple result types""" + lines = [ + json.dumps({ + "custom_id": "test-request-1", + "result": { + "type": "succeeded", + "message": { + "id": "msg_123", + "model": "claude-3-5-sonnet-20241022", + "role": "assistant", + "content": [{"type": "text", "text": "Success"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 5} + } + } + }), + json.dumps({ + "custom_id": "test-request-2", + "result": { + "type": "errored", + "error": { + "error": { + "type": "rate_limit_error", + "message": "Rate limit exceeded" + }, + "request_id": "req_456" + } + } + }), + json.dumps({ + "custom_id": "test-request-3", + "result": { + "type": "expired" + } + }) + ] + return "\n".join(lines).encode("utf-8") + + @pytest.mark.asyncio + async def test_afile_content_success(self, handler, mock_anthropic_batch_results_succeeded): + """Test successful file content retrieval and transformation""" + file_content_request: FileContentRequest = { + "file_id": "batch_123", + "extra_headers": None, + "extra_body": None + } + + # Mock the httpx client + mock_response = httpx.Response( + status_code=200, + content=mock_anthropic_batch_results_succeeded, + headers={"content-type": "application/json"}, + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123/results") + ) + + with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): + with patch.object(handler.anthropic_model_info, "get_api_base", return_value="https://api.anthropic.com"): + result = await handler.afile_content( + file_content_request=file_content_request, + api_key="test-api-key" + ) + + # Verify result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.status_code == 200 + + # Verify transformation to OpenAI format + content = result.response.content.decode("utf-8") + lines = [line for line in content.strip().split("\n") if line.strip()] + assert len(lines) == 1 + + transformed_result = json.loads(lines[0]) + assert transformed_result["custom_id"] == "test-request-1" + assert transformed_result["response"]["status_code"] == 200 + assert "body" in transformed_result["response"] + # Verify body has required OpenAI format fields + assert "id" in transformed_result["response"]["body"] + assert transformed_result["response"]["body"]["object"] == "chat.completion" + assert "choices" in transformed_result["response"]["body"] + # Verify request_id matches the original message id + assert transformed_result["response"]["request_id"] == "msg_123" + + @pytest.mark.asyncio + async def test_afile_content_with_prefix(self, handler, mock_anthropic_batch_results_succeeded): + """Test file content retrieval with anthropic_batch_results: prefix""" + file_content_request: FileContentRequest = { + "file_id": "anthropic_batch_results:batch_123", + "extra_headers": None, + "extra_body": None + } + + mock_response = httpx.Response( + status_code=200, + content=mock_anthropic_batch_results_succeeded, + headers={"content-type": "application/json"}, + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123/results") + ) + + with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): + with patch.object(handler.anthropic_model_info, "get_api_base", return_value="https://api.anthropic.com"): + result = await handler.afile_content( + file_content_request=file_content_request, + api_key="test-api-key" + ) + + assert isinstance(result, HttpxBinaryResponseContent) + # Verify the URL was constructed correctly (batch_id extracted from prefix) + mock_client.get.assert_called_once() + call_url = mock_client.get.call_args[1]["url"] + assert "batch_123" in call_url + + @pytest.mark.asyncio + async def test_afile_content_errored_result(self, handler, mock_anthropic_batch_results_errored): + """Test transformation of errored batch results""" + file_content_request: FileContentRequest = { + "file_id": "batch_123", + "extra_headers": None, + "extra_body": None + } + + mock_response = httpx.Response( + status_code=200, + content=mock_anthropic_batch_results_errored, + headers={"content-type": "application/json"}, + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123/results") + ) + + with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): + with patch.object(handler.anthropic_model_info, "get_api_base", return_value="https://api.anthropic.com"): + result = await handler.afile_content( + file_content_request=file_content_request, + api_key="test-api-key" + ) + + content = result.response.content.decode("utf-8") + lines = [line for line in content.strip().split("\n") if line.strip()] + assert len(lines) == 1 + + transformed_result = json.loads(lines[0]) + assert transformed_result["custom_id"] == "test-request-2" + assert transformed_result["response"]["status_code"] == 400 # invalid_request_error maps to 400 + assert transformed_result["response"]["body"]["error"]["type"] == "invalid_request_error" + assert transformed_result["response"]["body"]["error"]["message"] == "Invalid request" + + @pytest.mark.asyncio + async def test_afile_content_canceled_result(self, handler, mock_anthropic_batch_results_canceled): + """Test transformation of canceled batch results""" + file_content_request: FileContentRequest = { + "file_id": "batch_123", + "extra_headers": None, + "extra_body": None + } + + mock_response = httpx.Response( + status_code=200, + content=mock_anthropic_batch_results_canceled, + headers={"content-type": "application/json"}, + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123/results") + ) + + with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): + with patch.object(handler.anthropic_model_info, "get_api_base", return_value="https://api.anthropic.com"): + result = await handler.afile_content( + file_content_request=file_content_request, + api_key="test-api-key" + ) + + content = result.response.content.decode("utf-8") + lines = [line for line in content.strip().split("\n") if line.strip()] + assert len(lines) == 1 + + transformed_result = json.loads(lines[0]) + assert transformed_result["custom_id"] == "test-request-3" + assert transformed_result["response"]["status_code"] == 400 + assert "Batch request was canceled" in transformed_result["response"]["body"]["error"]["message"] + + @pytest.mark.asyncio + async def test_afile_content_mixed_results(self, handler, mock_anthropic_batch_results_mixed): + """Test transformation of mixed batch results (succeeded, errored, expired)""" + file_content_request: FileContentRequest = { + "file_id": "batch_123", + "extra_headers": None, + "extra_body": None + } + + mock_response = httpx.Response( + status_code=200, + content=mock_anthropic_batch_results_mixed, + headers={"content-type": "application/json"}, + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123/results") + ) + + with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): + with patch.object(handler.anthropic_model_info, "get_api_base", return_value="https://api.anthropic.com"): + result = await handler.afile_content( + file_content_request=file_content_request, + api_key="test-api-key" + ) + + content = result.response.content.decode("utf-8") + lines = [line for line in content.strip().split("\n") if line.strip()] + assert len(lines) == 3 + + # Check first result (succeeded) + result1 = json.loads(lines[0]) + assert result1["response"]["status_code"] == 200 + + # Check second result (errored) + result2 = json.loads(lines[1]) + assert result2["response"]["status_code"] == 429 # rate_limit_error maps to 429 + + # Check third result (expired) + result3 = json.loads(lines[2]) + assert result3["response"]["status_code"] == 400 + assert "expired" in result3["response"]["body"]["error"]["message"] + + @pytest.mark.asyncio + async def test_afile_content_missing_api_key(self, handler): + """Test file content retrieval with missing API key""" + file_content_request: FileContentRequest = { + "file_id": "batch_123", + "extra_headers": None, + "extra_body": None + } + + with patch.object(handler.anthropic_model_info, "get_api_key", return_value=None): + with pytest.raises(ValueError, match="Missing Anthropic API Key"): + await handler.afile_content( + file_content_request=file_content_request, + api_key=None + ) + + @pytest.mark.asyncio + async def test_afile_content_missing_file_id(self, handler): + """Test file content retrieval with missing file_id""" + file_content_request: FileContentRequest = { + "file_id": None, + "extra_headers": None, + "extra_body": None + } + + with pytest.raises(ValueError, match="file_id is required"): + await handler.afile_content( + file_content_request=file_content_request, + api_key="test-api-key" + ) + + @pytest.mark.asyncio + async def test_afile_content_http_error(self, handler): + """Test file content retrieval with HTTP error""" + file_content_request: FileContentRequest = { + "file_id": "batch_123", + "extra_headers": None, + "extra_body": None + } + + mock_response = httpx.Response( + status_code=404, + content=b"Not Found", + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123/results") + ) + mock_response.raise_for_status = MagicMock(side_effect=httpx.HTTPStatusError("Not Found", request=mock_response.request, response=mock_response)) + + with patch("litellm.llms.anthropic.files.handler.get_async_httpx_client") as mock_get_client: + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + with patch.object(handler.anthropic_model_info, "get_api_key", return_value="test-api-key"): + with patch.object(handler.anthropic_model_info, "get_api_base", return_value="https://api.anthropic.com"): + with pytest.raises(httpx.HTTPStatusError): + await handler.afile_content( + file_content_request=file_content_request, + api_key="test-api-key" + ) + + +class TestAnthropicBatchesConfig: + """Test Anthropic Batches Config for batch retrieval transformation""" + + @pytest.fixture + def config(self): + """Create AnthropicBatchesConfig instance""" + return AnthropicBatchesConfig() + + @pytest.fixture + def mock_anthropic_batch_response_in_progress(self): + """Mock Anthropic batch response with in_progress status""" + return { + "id": "batch_123", + "processing_status": "in_progress", + "created_at": "2024-01-01T00:00:00Z", + "expires_at": "2024-01-02T00:00:00Z", + "request_counts": { + "processing": 5, + "succeeded": 3, + "errored": 1, + "canceled": 0, + "expired": 0 + } + } + + @pytest.fixture + def mock_anthropic_batch_response_completed(self): + """Mock Anthropic batch response with completed status""" + return { + "id": "batch_456", + "processing_status": "ended", + "created_at": "2024-01-01T00:00:00Z", + "ended_at": "2024-01-01T12:00:00Z", + "expires_at": "2024-01-02T00:00:00Z", + "request_counts": { + "processing": 0, + "succeeded": 10, + "errored": 0, + "canceled": 0, + "expired": 0 + } + } + + @pytest.fixture + def mock_anthropic_batch_response_canceling(self): + """Mock Anthropic batch response with canceling status""" + return { + "id": "batch_789", + "processing_status": "canceling", + "created_at": "2024-01-01T00:00:00Z", + "cancel_initiated_at": "2024-01-01T06:00:00Z", + "ended_at": "2024-01-01T07:00:00Z", + "expires_at": "2024-01-02T00:00:00Z", + "request_counts": { + "processing": 0, + "succeeded": 5, + "errored": 0, + "canceled": 3, + "expired": 0 + } + } + + def test_get_retrieve_batch_url(self, config): + """Test URL construction for batch retrieval""" + url = config.get_retrieve_batch_url( + api_base="https://api.anthropic.com", + batch_id="batch_123", + optional_params={}, + litellm_params={} + ) + assert url == "https://api.anthropic.com/v1/messages/batches/batch_123" + + # Test with trailing slash + url = config.get_retrieve_batch_url( + api_base="https://api.anthropic.com/", + batch_id="batch_123", + optional_params={}, + litellm_params={} + ) + assert url == "https://api.anthropic.com/v1/messages/batches/batch_123" + + def test_transform_retrieve_batch_response_in_progress(self, config, mock_anthropic_batch_response_in_progress): + """Test transformation of in_progress batch response""" + mock_response = httpx.Response( + status_code=200, + content=json.dumps(mock_anthropic_batch_response_in_progress).encode("utf-8"), + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123") + ) + + logging_obj = MagicMock() + batch = config.transform_retrieve_batch_response( + model="claude-3-5-sonnet-20241022", + raw_response=mock_response, + logging_obj=logging_obj, + litellm_params={} + ) + + assert batch.id == "batch_123" + assert batch.object == "batch" + assert batch.status == "in_progress" + assert batch.endpoint == "/v1/messages" + assert batch.output_file_id == "batch_123" + assert batch.request_counts.total == 9 # 5 + 3 + 1 + assert batch.request_counts.completed == 3 + assert batch.request_counts.failed == 1 + assert batch.in_progress_at is not None + assert batch.completed_at is None + + def test_transform_retrieve_batch_response_completed(self, config, mock_anthropic_batch_response_completed): + """Test transformation of completed batch response""" + mock_response = httpx.Response( + status_code=200, + content=json.dumps(mock_anthropic_batch_response_completed).encode("utf-8"), + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_456") + ) + + logging_obj = MagicMock() + batch = config.transform_retrieve_batch_response( + model="claude-3-5-sonnet-20241022", + raw_response=mock_response, + logging_obj=logging_obj, + litellm_params={} + ) + + assert batch.id == "batch_456" + assert batch.status == "completed" + assert batch.completed_at is not None + assert batch.request_counts.total == 10 + assert batch.request_counts.completed == 10 + assert batch.request_counts.failed == 0 + + def test_transform_retrieve_batch_response_canceling(self, config, mock_anthropic_batch_response_canceling): + """Test transformation of canceling batch response""" + mock_response = httpx.Response( + status_code=200, + content=json.dumps(mock_anthropic_batch_response_canceling).encode("utf-8"), + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_789") + ) + + logging_obj = MagicMock() + batch = config.transform_retrieve_batch_response( + model="claude-3-5-sonnet-20241022", + raw_response=mock_response, + logging_obj=logging_obj, + litellm_params={} + ) + + assert batch.id == "batch_789" + assert batch.status == "cancelling" + assert batch.cancelling_at is not None + assert batch.cancelled_at is not None + assert batch.request_counts.total == 8 # 5 + 3 + + def test_transform_retrieve_batch_response_invalid_json(self, config): + """Test transformation with invalid JSON response""" + mock_response = httpx.Response( + status_code=200, + content=b"invalid json", + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123") + ) + + logging_obj = MagicMock() + with pytest.raises(ValueError, match="Failed to parse Anthropic batch response"): + config.transform_retrieve_batch_response( + model="claude-3-5-sonnet-20241022", + raw_response=mock_response, + logging_obj=logging_obj, + litellm_params={} + ) + + def test_transform_retrieve_batch_response_timestamp_parsing(self, config): + """Test timestamp parsing in batch response""" + batch_data = { + "id": "batch_123", + "processing_status": "ended", + "created_at": "2024-01-01T12:00:00Z", + "ended_at": "2024-01-01T13:30:45Z", + "expires_at": "2024-01-02T12:00:00Z", + "archived_at": "2024-01-03T00:00:00Z", + "request_counts": { + "processing": 0, + "succeeded": 1, + "errored": 0, + "canceled": 0, + "expired": 0 + } + } + + mock_response = httpx.Response( + status_code=200, + content=json.dumps(batch_data).encode("utf-8"), + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123") + ) + + logging_obj = MagicMock() + batch = config.transform_retrieve_batch_response( + model="claude-3-5-sonnet-20241022", + raw_response=mock_response, + logging_obj=logging_obj, + litellm_params={} + ) + + # Verify timestamps are parsed correctly + assert batch.created_at is not None + assert batch.completed_at is not None + assert batch.expires_at is not None + assert batch.expired_at is not None + + # Verify timestamps are integers (Unix timestamps) + assert isinstance(batch.created_at, int) + assert isinstance(batch.completed_at, int) + assert isinstance(batch.expires_at, int) + assert isinstance(batch.expired_at, int) + + def test_transform_retrieve_batch_response_missing_fields(self, config): + """Test transformation with missing optional fields""" + batch_data = { + "id": "batch_123", + "processing_status": "in_progress", + "request_counts": { + "processing": 1, + "succeeded": 0, + "errored": 0, + "canceled": 0, + "expired": 0 + } + } + + mock_response = httpx.Response( + status_code=200, + content=json.dumps(batch_data).encode("utf-8"), + request=httpx.Request(method="GET", url="https://api.anthropic.com/v1/messages/batches/batch_123") + ) + + logging_obj = MagicMock() + batch = config.transform_retrieve_batch_response( + model="claude-3-5-sonnet-20241022", + raw_response=mock_response, + logging_obj=logging_obj, + litellm_params={} + ) + + # Should still work with missing optional fields + assert batch.id == "batch_123" + assert batch.status == "in_progress" + assert batch.created_at is not None # Should default to current time if missing + assert batch.expires_at is None + assert batch.completed_at is None +