Spaces:
Configuration error
Configuration error
File size: 7,858 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 |
import json
import os
import sys
from datetime import datetime
sys.path.insert(
0, os.path.abspath("../../")
) # Adds the parent directory to the system path
import litellm
import pytest
from datetime import timedelta
from litellm.types.utils import ImageResponse, ImageObject
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
LiteLLMResponseObjectHandler,
)
def test_convert_to_image_response_basic():
# Test basic conversion with minimal input
response_dict = {
"created": 1234567890,
"data": [{"url": "http://example.com/image.jpg"}],
}
result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict)
assert isinstance(result, ImageResponse)
assert result.created == 1234567890
assert result.data[0].url == "http://example.com/image.jpg"
def test_convert_to_image_response_with_hidden_params():
# Test with hidden params
response_dict = {
"created": 1234567890,
"data": [{"url": "http://example.com/image.jpg"}],
}
hidden_params = {"api_key": "test_key"}
result = LiteLLMResponseObjectHandler.convert_to_image_response(
response_dict, hidden_params=hidden_params
)
assert result._hidden_params == {"api_key": "test_key"}
def test_convert_to_image_response_multiple_images():
# Test handling multiple images in response
response_dict = {
"created": 1234567890,
"data": [
{"url": "http://example.com/image1.jpg"},
{"url": "http://example.com/image2.jpg"},
],
}
result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict)
assert len(result.data) == 2
assert result.data[0].url == "http://example.com/image1.jpg"
assert result.data[1].url == "http://example.com/image2.jpg"
def test_convert_to_image_response_with_b64_json():
# Test handling b64_json in response
response_dict = {
"created": 1234567890,
"data": [{"b64_json": "base64encodedstring"}],
}
result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict)
assert result.data[0].b64_json == "base64encodedstring"
def test_convert_to_image_response_with_extra_fields():
response_dict = {
"created": 1234567890,
"data": [
{
"url": "http://example.com/image1.jpg",
"content_filter_results": {"category": "violence", "flagged": True},
},
{
"url": "http://example.com/image2.jpg",
"content_filter_results": {"category": "violence", "flagged": True},
},
],
}
result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict)
assert result.data[0].url == "http://example.com/image1.jpg"
assert result.data[1].url == "http://example.com/image2.jpg"
def test_convert_to_image_response_with_extra_fields_2():
"""
Date from a non-OpenAI API could have some obscure field in addition to the expected ones. This should not break the conversion.
"""
response_dict = {
"created": 1234567890,
"data": [
{
"url": "http://example.com/image1.jpg",
"very_obscure_field": "some_value",
},
{
"url": "http://example.com/image2.jpg",
"very_obscure_field2": "some_other_value",
},
],
}
result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict)
assert result.data[0].url == "http://example.com/image1.jpg"
assert result.data[1].url == "http://example.com/image2.jpg"
def test_convert_to_image_response_with_none_usage_fields():
"""
Test handling of None values in usage fields, specifically for gpt-image-1 responses.
This test verifies the fix for the bug where gpt-image-1 returns None values
for usage statistics fields, which caused Pydantic validation errors.
The fix should clean these None values and let ImageResponse constructor
handle the default values.
"""
response_dict = {
"created": 1234567890,
"data": [{"b64_json": "base64encodedstring"}],
"usage": {
"input_tokens": None, # gpt-image-1 returns None instead of integer
"input_tokens_details": None, # gpt-image-1 returns None instead of object
"output_tokens": None, # gpt-image-1 returns None instead of integer
"total_tokens": None, # gpt-image-1 returns None instead of integer
}
}
# This should not raise a ValidationError
result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict)
assert isinstance(result, ImageResponse)
assert result.created == 1234567890
assert result.data[0].b64_json == "base64encodedstring"
# Usage should be properly initialized with default values
assert result.usage is not None
assert result.usage.input_tokens == 0
assert result.usage.output_tokens == 0
assert result.usage.total_tokens == 0
assert result.usage.input_tokens_details is not None
assert result.usage.input_tokens_details.image_tokens == 0
assert result.usage.input_tokens_details.text_tokens == 0
def test_convert_to_image_response_with_partial_none_usage_fields():
"""
Test handling of mixed None and valid values in usage fields.
"""
response_dict = {
"created": 1234567890,
"data": [{"b64_json": "base64encodedstring"}],
"usage": {
"input_tokens": 10, # Valid value
"input_tokens_details": None, # None value (should be cleaned)
"output_tokens": None, # None value (should be cleaned)
"total_tokens": 10, # Valid value
}
}
# This should not raise a ValidationError
result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict)
assert isinstance(result, ImageResponse)
assert result.created == 1234567890
assert result.data[0].b64_json == "base64encodedstring"
# Usage should be properly initialized with defaults where needed
# Valid values should be preserved, None values should be cleaned and use defaults
assert result.usage is not None
assert result.usage.input_tokens == 10 # Valid value should be preserved
assert result.usage.output_tokens == 0 # None value should become 0
assert result.usage.total_tokens == 10 # Calculated as input_tokens + output_tokens (10 + 0)
assert result.usage.input_tokens_details is not None
assert result.usage.input_tokens_details.image_tokens == 0
assert result.usage.input_tokens_details.text_tokens == 0
def test_convert_to_image_response_with_valid_usage_fields():
"""
Test that valid usage fields are preserved correctly.
"""
response_dict = {
"created": 1234567890,
"data": [{"b64_json": "base64encodedstring"}],
"usage": {
"input_tokens": 50,
"input_tokens_details": {
"image_tokens": 30,
"text_tokens": 20,
},
"output_tokens": 10,
"total_tokens": 60,
}
}
result = LiteLLMResponseObjectHandler.convert_to_image_response(response_dict)
assert isinstance(result, ImageResponse)
assert result.created == 1234567890
assert result.data[0].b64_json == "base64encodedstring"
# Valid usage fields should be preserved
assert result.usage is not None
assert result.usage.input_tokens == 50
assert result.usage.output_tokens == 10
assert result.usage.total_tokens == 60
assert result.usage.input_tokens_details is not None
assert result.usage.input_tokens_details.image_tokens == 30
assert result.usage.input_tokens_details.text_tokens == 20
|