Files
litellm/tests/image_gen_tests/base_image_generation_test.py
T
Krish Dholakia beae5cfea9 Litellm staging 05 10 2025 - openai pdf url support + sagemaker chat content length error fix (#10724)
* Support pdf url's to openai (#10640)

* fix(gpt_transformation.py): support pdf url input to openai

pass as base64 as openai doesn't support image url's

* fix(openai.py): support async message transformation

allows async get request to convert url to base64

* fix(gpt_transformation.py): fix linting errrors and use common components across sync + async flows

* fix: fix linting errors

* fix(openai.py): pop correct var

* Fix sagemaker chat calls - content length error  (#10607)

* fix(sagemaker_chat/): support passing dynamic aws params

previously being ignored

* refactor(sagemaker/chat): more refactoring

* fix(sagemaker_chat/): make sure streaming is correctly handled post-refactor

* refactor: more refactoring to support using signed json str

* fix(sagemaker/chat): working sync streaming post refactor

* fix(sagemaker/chat): support async streaming post refactor

* fix(llm_http_handler.py): await async function

* fix: remove print statements

* test: update test

* test: update test

* fix(llm_http_handler.py): retain passing in data as json str

* test: update test

* fix(base_model_iterator.py): fix linting error

* test: test auth

* fix: fix linting error

* test: update test

* test: update translation test

* fix(gpt_transformation.py): handle awaitable/non-awaitable object

* fix: handle async flow for message transformation on openai compatible api's

* test: cleanup testing

* test: update test

* test(test_router.py): use model with higher quota

* test: simplify test

* test: update test
2025-05-10 17:41:57 -07:00

90 lines
3.3 KiB
Python

import asyncio
import httpx
import json
import pytest
import sys
from typing import Any, Dict, List, Optional
from unittest.mock import MagicMock, Mock, patch
import os
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm.exceptions import BadRequestError
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.utils import CustomStreamWrapper
from openai.types.image import Image
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import StandardLoggingPayload
class TestCustomLogger(CustomLogger):
def __init__(self):
super().__init__()
self.standard_logging_payload: Optional[StandardLoggingPayload] = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.standard_logging_payload = kwargs.get("standard_logging_object")
pass
# test_example.py
from abc import ABC, abstractmethod
class BaseImageGenTest(ABC):
"""
Abstract base test class that enforces a common test across all test classes.
"""
@abstractmethod
def get_base_image_generation_call_args(self) -> dict:
"""Must return the base image generation call args"""
pass
@pytest.mark.asyncio(scope="module")
async def test_basic_image_generation(self):
"""Test basic image generation"""
try:
custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [custom_logger]
base_image_generation_call_args = self.get_base_image_generation_call_args()
litellm.set_verbose = True
response = await litellm.aimage_generation(
**base_image_generation_call_args, prompt="A image of a otter"
)
print(response)
await asyncio.sleep(1)
# assert response._hidden_params["response_cost"] is not None
# assert response._hidden_params["response_cost"] > 0
# print("response_cost", response._hidden_params["response_cost"])
logged_standard_logging_payload = custom_logger.standard_logging_payload
print("logged_standard_logging_payload", logged_standard_logging_payload)
assert logged_standard_logging_payload is not None
assert logged_standard_logging_payload["response_cost"] is not None
assert logged_standard_logging_payload["response_cost"] > 0
from openai.types.images_response import ImagesResponse
ImagesResponse.model_validate(response.model_dump())
for d in response.data:
assert isinstance(d, Image)
print("data in response.data", d)
assert d.b64_json is not None or d.url is not None
except litellm.RateLimitError as e:
pass
except litellm.ContentPolicyViolationError:
pass # Azure randomly raises these errors - skip when they occur
except litellm.InternalServerError:
pass
except Exception as e:
if "Your task failed as a result of our safety system." in str(e):
pass
else:
pytest.fail(f"An exception occurred - {str(e)}")