mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-24 18:24:20 +00:00
Merge 0a1bebd5c8 into f5996b2f6b
This commit is contained in:
PradyMagal
GitHub
commit
a4c3f7cca5
4 changed files with 302 additions and 2 deletions
|
|
@ -844,6 +844,21 @@ def exception_type( # type: ignore # noqa: PLR0915
|
|||
llm_provider="bedrock",
|
||||
response=getattr(original_exception, "response", None),
|
||||
)
|
||||
elif (
|
||||
"Client error '400 Bad Request'" in error_str
|
||||
and (
|
||||
"files are not supported" in error_str
|
||||
or "Unsupported image format" in error_str
|
||||
or "Supported formats are:" in error_str
|
||||
)
|
||||
):
|
||||
exception_mapping_worked = True
|
||||
raise BadRequestError(
|
||||
message=f"BedrockException - {error_str}",
|
||||
model=model,
|
||||
llm_provider="bedrock",
|
||||
response=getattr(original_exception, "response", None),
|
||||
)
|
||||
elif (
|
||||
"Unable to locate credentials" in error_str
|
||||
or "The security token included in the request is invalid"
|
||||
|
|
|
|||
|
|
@ -2352,8 +2352,9 @@ class BedrockImageProcessor:
|
|||
]
|
||||
|
||||
if not valid_extensions:
|
||||
document_type = mime_type.split('/')[-1].upper() # e.g., "PDF" from "application/pdf"
|
||||
raise ValueError(
|
||||
f"No supported extensions for MIME type: {mime_type}. Supported formats: {supported_doc_formats}"
|
||||
f"Client error '400 Bad Request': {document_type} files are not supported. Supported formats are: {supported_doc_formats}"
|
||||
)
|
||||
|
||||
# Use first valid extension instead of provided image_format
|
||||
|
|
@ -2361,7 +2362,7 @@ class BedrockImageProcessor:
|
|||
else:
|
||||
if image_format not in supported_image_formats:
|
||||
raise ValueError(
|
||||
f"Unsupported image format: {image_format}. Supported formats: {supported_image_formats}"
|
||||
f"Client error '400 Bad Request': Unsupported image format: {image_format}. Supported formats are: {supported_image_formats}"
|
||||
)
|
||||
return image_format
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,194 @@
|
|||
import pytest
|
||||
import requests
|
||||
import base64
|
||||
import logging
|
||||
|
||||
# Set up logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PROXY_URL = "http://0.0.0.0:4000/chat/completions"
|
||||
|
||||
def create_base64_pdf():
|
||||
"""Create a minimal valid PDF in base64 format"""
|
||||
minimal_pdf = b"%PDF-1.0\n1 0 obj<</Type/Catalog/Pages 2 0 R>>endobj\n2 0 obj<</Type/Pages/Kids[3 0 R]/Count 1>>endobj\n3 0 obj<</Type/Page/MediaBox[0 0 3 3]>>endobj\nxref\n0 4\n0000000000 65535 f\n0000000010 00000 n\n0000000053 00000 n\n0000000102 00000 n\ntrailer<</Size 4/Root 1 0 R>>\nstartxref\n149\n%EOF"
|
||||
return base64.b64encode(minimal_pdf).decode()
|
||||
|
||||
def create_base64_png():
|
||||
"""Create a simple PNG in base64 format (red dot)"""
|
||||
return "iVBORw0KGgoAAAANSUhEUgAAAAUAAAAFCAYAAACNbyblAAAAHElEQVQI12P4//8/w38GIAXDIBKE0DHxgljNBAAO9TXL0Y4OHwAAAABJRU5ErkJggg=="
|
||||
|
||||
def create_base64_jpeg():
|
||||
"""Create a simple JPEG in base64 format (white pixel)"""
|
||||
return "/9j/4AAQSkZJRgABAQEAYABgAAD/2wBDAP//////////////////////////////////////////////////////////////////////////////////////2wBDAf//////////////////////////////////////////////////////////////////////////////////////wAARCAABAAEDASIAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwD3+iiigD//2Q=="
|
||||
|
||||
def test_pdf_content_raises_400():
|
||||
"""Test that sending PDF content raises a 400 error with correct message"""
|
||||
pdf_base64 = create_base64_pdf()
|
||||
|
||||
payload = {
|
||||
"model": "claude-3",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:application/pdf;base64,{pdf_base64}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in this document?"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
PROXY_URL,
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
response_data = response.json()
|
||||
|
||||
logger.info(f"Error response for PDF test: {response_data}")
|
||||
|
||||
error_message = response_data["error"]["message"]
|
||||
assert "litellm.BadRequestError: AnthropicException" in error_message
|
||||
assert "claude-3-haiku-20240307" in error_message
|
||||
assert "does not support PDF input" in error_message
|
||||
|
||||
def test_png_content_succeeds():
|
||||
"""Test that sending PNG content works successfully"""
|
||||
png_base64 = create_base64_png()
|
||||
|
||||
payload = {
|
||||
"model": "claude-3",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{png_base64}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in this image?"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
PROXY_URL,
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
response_data = response.json()
|
||||
logger.info(f"Successful PNG response status: {response.status_code}")
|
||||
|
||||
# Check response has expected structure
|
||||
assert "choices" in response_data
|
||||
assert len(response_data["choices"]) > 0
|
||||
assert "message" in response_data["choices"][0]
|
||||
assert "content" in response_data["choices"][0]["message"]
|
||||
|
||||
def test_jpeg_content_succeeds():
|
||||
"""Test that sending JPEG content works successfully"""
|
||||
jpeg_base64 = create_base64_jpeg()
|
||||
|
||||
payload = {
|
||||
"model": "claude-3",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{jpeg_base64}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in this image?"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
PROXY_URL,
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
response_data = response.json()
|
||||
logger.info(f"Successful JPEG response status: {response.status_code}")
|
||||
|
||||
# Check response has expected structure
|
||||
assert "choices" in response_data
|
||||
assert len(response_data["choices"]) > 0
|
||||
assert "message" in response_data["choices"][0]
|
||||
assert "content" in response_data["choices"][0]["message"]
|
||||
|
||||
def test_multiple_images_succeeds():
|
||||
"""Test that sending multiple supported images in one request works"""
|
||||
png_base64 = create_base64_png()
|
||||
jpeg_base64 = create_base64_jpeg()
|
||||
|
||||
payload = {
|
||||
"model": "claude-3",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{png_base64}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/jpeg;base64,{jpeg_base64}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What's in these images?"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
PROXY_URL,
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
response_data = response.json()
|
||||
logger.info(f"Successful multiple images response status: {response.status_code}")
|
||||
|
||||
# Check response has expected structure
|
||||
assert "choices" in response_data
|
||||
assert len(response_data["choices"]) > 0
|
||||
assert "message" in response_data["choices"][0]
|
||||
assert "content" in response_data["choices"][0]["message"]
|
||||
|
|
@ -0,0 +1,90 @@
|
|||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
import mimetypes
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import BedrockImageProcessor
|
||||
|
||||
# Mock the configurations and classes we need
|
||||
class MockAmazonConfig:
|
||||
def get_supported_image_types(self):
|
||||
return ['png', 'jpeg', 'gif', 'webp']
|
||||
|
||||
def get_supported_document_types(self):
|
||||
return ['pdf', 'docx']
|
||||
|
||||
|
||||
def test_error_mapping_ValueError_to_BadRequest():
|
||||
"""Test that ValueError gets mapped to BadRequestError with correct message"""
|
||||
print("\n--- Running test_error_mapping_ValueError_to_BadRequest ---")
|
||||
try:
|
||||
# Simulate the ValueError being raised
|
||||
error_message = "Client error '400 Bad Request': PDF files are not supported. Supported formats are: ['docx']"
|
||||
print(f"Raising ValueError with message: {error_message}")
|
||||
|
||||
raise ValueError(error_message)
|
||||
except ValueError as e:
|
||||
print(f"Caught ValueError: {e}")
|
||||
|
||||
with pytest.raises(BadRequestError) as exc_info:
|
||||
print("Calling exception_type...")
|
||||
exception_type(
|
||||
model="model-name",
|
||||
original_exception=e,
|
||||
custom_llm_provider="bedrock"
|
||||
)
|
||||
|
||||
error = exc_info.value
|
||||
print(f"Caught BadRequestError: {error}")
|
||||
print(f"Error status code: {error.status_code}")
|
||||
print(f"Error message: {str(error)}")
|
||||
|
||||
assert error.status_code == 400
|
||||
assert "PDF files are not supported" in str(error)
|
||||
print("Assertions passed successfully!")
|
||||
|
||||
def test_error_mapping_413_error():
|
||||
"""Test that 413 status code gets mapped correctly"""
|
||||
print("\n--- Running test_error_mapping_413_error ---")
|
||||
mock_exception = Mock()
|
||||
mock_exception.status_code = 413
|
||||
mock_exception.message = "File too large"
|
||||
|
||||
print(f"Mock exception status code: {mock_exception.status_code}")
|
||||
print(f"Mock exception message: {mock_exception.message}")
|
||||
|
||||
with pytest.raises(BadRequestError) as exc_info:
|
||||
exception_type(
|
||||
model="model-name",
|
||||
original_exception=mock_exception,
|
||||
custom_llm_provider="replicate"
|
||||
)
|
||||
|
||||
error = exc_info.value
|
||||
print(f"Caught BadRequestError: {error}")
|
||||
print(f"Error status code: {error.status_code}")
|
||||
print(f"Error message: {str(error)}")
|
||||
|
||||
assert error.status_code == 400
|
||||
assert "ReplicateException" in str(error)
|
||||
assert "File too large" in str(error)
|
||||
print("Assertions passed successfully!")
|
||||
|
||||
def test_validation_edge_cases():
|
||||
"""Test edge cases in format validation"""
|
||||
print("\n--- Running test_validation_edge_cases ---")
|
||||
with patch('litellm.AmazonConverseConfig', return_value=MockAmazonConfig()):
|
||||
# Test empty mime type
|
||||
print("Testing empty mime type")
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
BedrockImageProcessor._validate_format("", "")
|
||||
print(f"Caught ValueError: {exc_info.value}")
|
||||
assert "Client error '400 Bad Request'" in str(exc_info.value)
|
||||
|
||||
# Test invalid mime type
|
||||
print("Testing invalid mime type")
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
BedrockImageProcessor._validate_format("invalid/type", "format")
|
||||
print(f"Caught ValueError: {exc_info.value}")
|
||||
assert "Client error '400 Bad Request'" in str(exc_info.value)
|
||||
print("Assertions passed successfully!")
|
||||
Loading…
Add table
Add a link
Reference in a new issue