tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 09:06:18 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
f5b11b72a6dcc8f4e7a16f00a61bdd124cc171d2
litellm/tests/image_gen_tests/test_image_edits.py
T
Mateo Wang bb448b0031 fix(tests): stabilize image-edit VCR cassettes to stop live gpt-image-1 spend (#28110) 
* fix(tests): stabilize image-edit VCR cassettes to stop live gpt-image-1 spend

The image-edit cassettes for ``gpt-image-1`` were accumulating >50
episodes and being refused by the persister
(``tests/_vcr_redis_persister.py``), so every CI run was hitting the
real OpenAI endpoint. The async parametrize was the clearest tell:
``test_openai_image_edit_litellm_sdk[True]`` cached to 1 entry, but the
``[False]`` (async) sibling grew to 51 entries and never replayed.

Two non-deterministic sources were fueling the growth, both fixed
here. After this patch, the cassettes settle at one episode per
unique call and replay for the 24-hour TTL like every other suite.

1. Pin httpx's multipart boundary at the source. The existing
   ``_normalize_multipart_boundary`` rewrites the boundary in the
   ``Content-Type`` header reliably, but on the async transport path
   the body is not always a contiguous ``bytes`` object when
   ``before_record_request`` runs, so the body-side replacement
   silently no-ops and the recorded cassette retains the random
   ``boundary=<hex>`` string. The next CI run gets a fresh random
   boundary, the ``safe_body`` matcher misses, and
   ``record_mode="new_episodes"`` appends another episode. Wrapping
   ``httpx._multipart.MultipartStream.__init__`` so it always uses
   ``vcr-static-boundary`` when no boundary is supplied eliminates
   the variance for both sync and async paths and leaves the normalizer
   in place as a backstop. Exposed as
   ``pin_httpx_multipart_boundary`` so other multipart-heavy suites
   (audio, ocr, batches) can adopt the same fixture later.

2. Pass raw ``bytes`` (not ``BytesIO`` streams) through the
   image-edit fixtures. A ``BytesIO`` whose file pointer is at EOF
   after the first multipart upload silently encodes an empty image on
   the next SDK / Router retry — yet another divergent body that VCR
   records as a new episode. ``bytes`` are immutable and position-less,
   so retries re-encode an identical payload every time. This is also
   a small production-correctness improvement: a customer passing
   ``BytesIO`` today would hit the same empty-body retry bug. The
   BytesIO-specific smoke test
   (``test_openai_image_edit_with_bytesio``) is preserved by giving
   ``get_test_images_as_bytesio`` its own factory instead of aliasing
   the bytes one.

3. Add ``scripts/flush_image_edit_vcr_cassettes.py`` — a one-shot
   Redis SCAN/DEL helper that clears the bloated pre-fix cassettes
   under ``litellm:vcr:cassette:tests/image_gen_tests/test_image_edits/*``.
   Without this, the next CI run still loads the existing 51-entry
   cassette, the new fixed-boundary body still doesn't match any of
   the stale entries, the persister still refuses to save, and the
   bleed continues. Run once with the production
   ``CASSETTE_REDIS_URL`` after merge (dry-run by default).

* DIAGNOSTIC: log VCR body mismatches + per-episode body hashes

Temporary observability boost so we can root-cause why
``test_image_edits.py`` async parametrizes still record fresh
episodes on every CI run even though the multipart boundary is now
pinned (sync parametrizes cache cleanly as VCR HIT). The matcher
currently raises ``AssertionError("request bodies differ")`` with
zero context, so we cannot tell whether the live body genuinely
varies, the matcher is comparing a bytes object to a stream object,
or the normalizer is silently skipping the body because it is not
bytes/str.

Three logs added; the first two are worth keeping permanently, the
third is intended to be reverted after the diagnosis lands:

1. ``_safe_body_matcher`` now emits a structured stderr block on
   mismatch (type of each side, length, SHA-256, first divergent
   byte offset, ±100-byte window). Always-on -- mismatches are
   signal, not noise, and the existing per-test verdict already
   logs once per test. PERMANENT.

2. ``_normalize_multipart_boundary`` now logs to stderr when the
   body type is not bytes/bytearray/str -- the silent ``else:
   return`` branch was masking exactly the case we suspect is
   firing on async (httpx ``MultipartStream`` handed to vcrpy
   before the body is read). PERMANENT.

3. ``_RedisPersister.save_cassette`` now logs every episode's body
   SHA-256, length, and 120-byte preview at save time. This lets
   two consecutive CI runs be diffed: if the same test records a
   different hash run-to-run, the live body genuinely varies; if
   both runs record the same hash but the matcher still misses, the
   bug is in the matcher itself. TEMPORARY -- revert once the
   async variance is identified and fixed.

Once a single ``image_gen_testing`` CI run produces these logs,
revert this commit (or just the persister hash block) with a force
push so the cassette save path is not noisy in steady-state.

* DIAGNOSTIC: route VCR diagnostics through per-PID files (bypass xdist capture)

Re-push of the diagnostic logging from the previous commit, this
time wired so the output actually survives to the CI log. xdist
captures stdout/stderr from every passing test in the worker
process; the body-matcher and normalizer-skip diagnostics fire from
inside vcrpy machinery during the test, so for any test that
ultimately passes (which is all of them once the cassettes are
recorded), the diagnostic lines are silently swallowed.

Fix: write each diagnostic line to a per-PID file under
``test-results/vcr-diagnostics/<pid>.log`` instead of writing to
stderr. The controller's ``pytest_terminal_summary`` aggregates
those files and writes them through ``terminalreporter.write_line``,
which is not subject to per-test capture. As a bonus,
``test-results/`` is already collected by the ``store_test_results``
step in CircleCI, so the raw per-worker logs survive as build
artifacts even after the test session ends.

Three call sites updated:

1. ``_emit_body_mismatch_diagnostic`` (matcher) -- writes the
   structured type/length/sha/window block via ``vcr_diag_write_line``.
2. ``_normalize_multipart_boundary`` -- logs the silent-skip path
   (body not bytes/bytearray/str) the same way.
3. ``_maybe_log_episode_body_hashes`` (persister) -- replaces the
   ``_log.warning`` calls (which the root-logger config also
   swallows in CI) with ``vcr_diag_write_line``.

Image-gen conftest is the only suite wired to dump the aggregated
log at session end. Other suites can opt in by adding
``emit_vcr_diagnostic_log(terminalreporter)`` to their own
``pytest_terminal_summary``. The diagnostic dir is cleared at the
start of each session (controller-only) so a local rerun does not
mix output from prior runs.

Same revert plan as the previous diagnostic commit: keep the
matcher + normalizer skip diagnostics permanently (they only fire
on signal events), revert the persister body-hash dump once the
async variance is identified.

* fix(tests): coalesce iterable request bodies before matching/recording

Root cause of the residual async image-edit cassette leak. The
diagnostic run for ``ba3915d9`` printed:

  [vcr-safe-body-matcher] request body mismatch
    body[a]: type='list_iterator' length=unknown sha256=N/A
    body[b]: type='list_iterator' length=unknown sha256=N/A

httpx's async transport hands vcrpy a ``request.body`` that is a
``list_iterator`` over multipart chunks rather than a contiguous
``bytes`` blob. Two consequences:

1. ``_safe_body_matcher`` compares the two iterator objects with
   ``==``, which is identity comparison for arbitrary iterators -
   semantically identical multipart bodies never compare equal, and
   ``record_mode="new_episodes"`` appends a new episode on every CI
   run until the cassette crosses ``MAX_EPISODES_PER_CASSETTE`` and
   the persister refuses to save (this is exactly what the OVERFLOW
   warning has been catching).
2. ``_normalize_multipart_boundary`` short-circuits its
   ``else: return`` branch because the body is neither bytes nor
   str, so any residual random boundary characters in the body bytes
   are never rewritten.

Sync requests do not hit this code path: httpx's sync transport
hands vcrpy a single ``bytes`` body, so ``==`` works and the
boundary normalizer runs as intended. That is why
``test_openai_image_edit_litellm_sdk[True]`` records to ``entries=1``
and replays cleanly while ``[False]`` (async) kept growing by one
episode per run.

Fix: add ``_materialize_iterable_body`` which coalesces an iterable
``request.body`` into ``bytes`` in-place. Call it from two places:

* The top of ``_before_record_request``, so the boundary normalizer
  and the cassette serializer both see bytes from then on.
* The top of ``_safe_body_matcher``, as defense in depth in case a
  future vcrpy code path invokes the matcher without first going
  through ``_before_record_request``.

The vcrpy ``Request`` is a wrapper used for matching and recording;
the underlying httpx transport sends its own request body
separately, so replacing the iterator on the vcrpy wrapper does
not starve the live HTTP send.

After this lands the async parametrizes should flip from
``[VCR MISS:RECORDED] entries=N+1`` to ``[VCR HIT] entries=N`` on
the next CI run, matching the sync side and dropping the residual
~$3/day to $0.

* fix(tests): handle bytes_iterator + never leave an exhausted body

Follow-up to 8e08272b. The previous attempt at coalescing iterable
request bodies bailed out (``return`` without writing
``request.body``) whenever it could not classify the chunk type.
That was the wrong failure mode for one critical case: vcrpy
sometimes presents the body as ``iter(some_bytes)``, whose Python
type is ``bytes_iterator`` and which yields ``int`` byte values
(0-255), not byte chunks. The old code saw an ``int`` chunk, hit
the ``else: return`` branch, and left ``request.body`` pointing at
the now-exhausted iterator.

The post-fix diagnostic run made this loud:

  [vcr-safe-body-matcher] request body mismatch
    body[a]: type='bytes_iterator' length=unknown sha256=N/A
    body[b]: type='bytes_iterator' length=unknown sha256=N/A

Every async image-edit test then ballooned from entries=2 to
entries=10 in that single CI run -- the exhausted iterator meant
the live multipart upload went out as an empty body, OpenAI
returned 400, the SDK + flaky retries fired, each retry got a
fresh iterator that my hook exhausted again, and ``new_episodes``
recorded each failed attempt as a new cassette episode.

This patch:

* Recognizes ``bytes_iterator`` (chunks are ``int``) and
  reconstructs the buffer via ``bytes(chunks)``.
* Keeps the existing ``list_iterator``-over-bytes-chunks handling
  via ``b"".join(...)``.
* **Always writes a bytes value back to ``request.body`` after
  consuming the iterator.** If the chunk shape is unrecognized,
  ``request.body`` is set to ``b""`` rather than left as an
  exhausted iterator. That is wrong in the sense of "we lost the
  body" but right in the sense of "the failure mode is now visible
  (live API call sends empty body and fails fast) instead of
  invisible (corrupt cassette grows silently)". Combined with the
  matcher diagnostic, any future regression in this code path will
  surface in the CI log immediately.

Local verification covers ``bytes_iterator``, ``list_iterator``
over bytes chunks, generator over bytes chunks, empty iterator,
already-bytes (idempotent), identical-content iterator equality
in the matcher (now matches), and differing-content iterator
inequality (still raises).

* fix(tests): clear vcrpy's sticky _was_iter flag so materialized bodies stay bytes

Actual root cause of the async image-edit cassette leak. The
previous diagnostic run produced this dead giveaway:

  [vcr-episode-body-hash] ... episode[0]: body type='bytes_iterator'
    is not bytes/bytearray/str -- cannot hash
  [vcr-safe-body-matcher] request body mismatch
    body[a]: type='bytes_iterator' length=unknown sha256=N/A
    body[b]: type='bytes_iterator' length=unknown sha256=N/A

Both sides of the matcher were ``bytes_iterator`` **after** the
materializer had supposedly converted them to bytes. That made no
sense until I read vcrpy's ``Request`` class.

vcrpy's ``Request`` keeps two private flags that are set in
``__init__`` from the original body's type and **never cleared by
the setter**:

  def __init__(self, method, uri, body, headers):
      self._was_file = hasattr(body, "read")
      self._was_iter = _is_nonsequence_iterator(body)
      ...

  @property
  def body(self):
      if self._was_file: return BytesIO(self._body)
      if self._was_iter: return iter(self._body)
      return self._body

  @body.setter
  def body(self, value):
      if isinstance(value, str): value = value.encode("utf-8")
      self._body = value   # <-- does NOT touch _was_iter / _was_file

So when httpx's async transport hands vcrpy an iterator body,
``_was_iter`` becomes ``True`` and stays there forever. Even after
``_materialize_iterable_body`` writes plain bytes via
``request.body = out``, the next read of ``.body`` re-wraps the
stored bytes in ``iter()`` -- producing a fresh ``bytes_iterator``
that compares unequal to any other ``bytes_iterator`` via object
identity. The matcher missed every time, the cassette grew by one
episode per run, and the persister saw the same iterator type when
trying to hash the body for the diagnostic log.

Fix: after writing the materialized bytes, also force
``_was_iter`` and ``_was_file`` to ``False``. vcrpy exposes no
public API for this, so we touch the private flags directly --
acknowledged as a pragmatic test-only hack with a clear unit
boundary (the only call site is ``_materialize_iterable_body``).

Local repro reproduces the exact production setup:
``Request('POST', url, iter(b'multipart-content'), {})`` on two
sides, runs the matcher, asserts HIT. Verified the matcher hits on
identical content and still raises on differing content.

Should be the last fix needed. Existing cassettes that contain
oddly-shaped bodies (lists of int chunks, etc. from the previous
``_was_iter=True`` save path) still match because the materializer
canonicalises both sides to bytes before comparison -- no fourth
re-flush required.

* revert(tests): drop the temp per-episode body-hash diagnostic

Removed now that 1c51ad13 has confirmed the root cause (vcrpy's
sticky ``_was_iter`` flag making the body getter re-wrap stored
bytes in ``iter()`` on every access). The hash dump did its job --
the post-1c51ad13 image_gen_testing run shows all five async
image-edit tests as ``[VCR HIT]`` with stable entry counts and
zero billing errors -- and is too noisy to keep on by default
(over 100 lines per session at steady state).

Kept permanently:

* ``_safe_body_matcher`` mismatch diagnostic in
  ``_vcr_conftest_common.py``. Only fires on a body mismatch,
  which is signal worth surfacing whenever it happens.
* ``_normalize_multipart_boundary`` "skipped" log line. Same
  rationale -- only fires when the body shape is something the
  normalizer cannot rewrite in place.
* The ``test-results/vcr-diagnostics/<pid>.log`` per-PID file
  plumbing (``vcr_diag_write_line`` /
  ``emit_vcr_diagnostic_log``). Useful for any future diagnostic
  that needs to bypass xdist stdout/stderr capture; cheap to keep.

* chore(tests): delete unused flush script + wire VCR diagnostic dump everywhere

* Remove ``scripts/flush_image_edit_vcr_cassettes.py``. It was a
  one-shot helper for the initial cassette flush; the iterator and
  ``_was_iter`` fixes mean no future flush should be required, and
  the script was never run anywhere (the actual flushes happened
  inside the CI conftest via the temp hacks that have since been
  reverted).

* The matcher mismatch + normalizer skip diagnostics already write
  per-PID files for every suite that imports the shared VCR
  plumbing, but ``emit_vcr_diagnostic_log`` -- the controller-side
  dump that surfaces those files into the CI log at session end --
  was only wired into ``image_gen_tests``. Add the one-line call to
  the 12 sibling conftests that already use VCR so the diagnostics
  surface in any suite's terminal output if a body matcher ever
  misses. No new output in steady state -- the dump is a no-op when
  no diagnostics were recorded that session.

* chore(tests): trim non-essential comments per project comment policy

Strips docstrings, inline comments, and block comments that this PR
introduced where the code itself was already self-evident. Keeps the
few lines that document non-obvious behaviour (raw-bytes-not-BytesIO
rationale on the image fixtures, the per-PID-files-bypass-xdist note
on the diagnostic directory). Touches only comments this PR added --
no pre-existing comment is removed.

Net: -161 lines of comment/docstring across 3 files, no code
behaviour change.

* chore(tests): forward **kwargs in pin_httpx_multipart_boundary wrapper

Defensive against future httpx MultipartStream.__init__ adding new
optional kwargs. Without the forward, the wrapper would silently drop
them. No behaviour change today.

* chore(tests): canonicalize VCR matchers and surface shouldn't-happen branches

Bundles the "follow-up cleanup PR" into this one so it does not get
lost. Four small changes:

1. Introduce ``_canonical_body(req) -> (bytes, pre_type)`` and route
   ``_safe_body_matcher`` through it. The matcher now operates on
   bytes by construction; the "compare two iterator objects via
   ``==`` and silently get object-identity semantics" failure mode
   (which cost us this entire PR to diagnose) is structurally
   impossible to reintroduce. ``pre_type`` is the body type *before*
   canonicalization, surfaced by the mismatch diagnostic so a future
   regression involving a new body shape is still visible.

2. Add a structured diagnostic to ``_key_fingerprint_matcher``. It
   was previously raising a bare ``AssertionError("API key
   fingerprints differ")`` with zero context -- exactly the
   anti-pattern the body matcher had before this PR.

3. Surface "shouldn't-happen" branches via ``vcr_diag_write_line``:

   * ``_strip_image_b64_payloads`` -- logs when ``response``,
     ``response['body']``, or ``response['body']['string']`` arrives
     in an unexpected shape (vcrpy contract violation).
   * ``_compute_key_fingerprint`` -- logs the ``"no-key"`` fallback
     with the request method/URL so a stripped-auth-header bug is
     visible instead of masked.
   * ``_canonical_body`` -- logs its own empty-bytes fallback when a
     body has a shape ``_materialize_iterable_body`` did not handle.

4. Re-introduce per-episode body-hash logging in
   ``_RedisPersister.save_cassette`` (was reverted in 927c5548 as
   "noisy"). Quantified cost: ~25 KB of CI log per session at peak,
   ~ms-scale CPU, zero output in steady state (no save = no log).
   Trade-off favours keeping it: lets two consecutive CI runs be
   diffed by body hash, which is how we will spot the next regression
   in the same class.

All call sites still work: local repro confirms iter==iter HIT,
iter!=iter raises, plain-bytes HIT, body-hash log emits via the same
per-PID file plumbing as the matcher diagnostics.

* chore(tests): symmetrize diag-log cleanup across every VCR-using conftest

``image_gen_tests/conftest.py`` was the only suite that cleared
``test-results/vcr-diagnostics/*.log`` at session start. The other 12
VCR-using conftests inherited any stale per-PID logs from a previous
local run and would dump them in the terminal summary -- harmless in
CI (fresh container) but confusing locally when running multiple
suites in sequence.

Extracts the cleanup into a ``reset_vcr_diag_dir`` helper in
``tests/_vcr_conftest_common.py`` and calls it from every VCR-using
conftest's ``pytest_configure``. Same single source of truth, no
inline duplication.

* fix(tests): gate body materialization on __next__ and strip PR comments

aiohttp/vcrpy stores the json kwarg as a dict; _materialize_iterable_body
was iterating it via __iter__ and joining the keys, replacing the request
body with concatenated key names ("textlanguageentities"). Gate on
__next__ so containers (dict/list/tuple) are left alone — only single-use
iterators like httpx's bytes_iterator / list_iterator are materialized.
Log diagnostic line when chunk type is unrecognized.

* fix(tests): JSON-encode dict bodies in canonical_body for stable matching

aiohttp stubs store the json kwarg as a dict; the fallback that compared
all dicts as b"" caused concurrent presidio analyze calls to be served
the wrong cassette episode. JSON-encode with sort_keys for stable bytes.

* fix(tests): guard emit_vcr_diagnostic_log against multi-conftest re-emission

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* fix(tests): globalize multipart-boundary pin + stabilize whisper fixtures

Diagnostic shows audio_testing was silently re-recording 50+ live Whisper
episodes per CI run (over MAX_EPISODES_PER_CASSETTE, so the persister
refused to save). Two changes:

* Move the session-autouse _pin_multipart_boundary fixture into the
  shared _vcr_conftest_common module so every VCR-using suite picks it
  up via a single import. image_gen had it inline; the other 12 suites
  silently lacked it.
* Replace the module-level open("rb") audio file handles in test_whisper
  with cached bytes + a per-call (filename, bytes, mimetype) tuple,
  mirroring the image_edits raw-bytes pattern. Stops the file-pointer-
  at-EOF bug where the second test got an empty multipart body.

* chore(tests): drop per-episode body-hash dump and redundant emit guard

---------

Co-authored-by: shin-berri <shin-laptop@berri.ai>
Co-authored-by: yuneng-jiang <yuneng@berri.ai>
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
2026-05-18 09:15:39 -07:00

776 lines
25 KiB
Python
Raw Blame History

import logging
import os
import sys
import traceback
import asyncio
from typing import Optional
import pytest
import base64
from io import BytesIO
from unittest.mock import patch, AsyncMock
import json
from abc import ABC, abstractmethod
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm.utils import ImageResponse
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import StandardLoggingPayload
# Configure pytest marks to avoid warnings
pytestmark = pytest.mark.asyncio
class TestCustomLogger(CustomLogger):
def __init__(self):
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", None)
pass
class BaseLLMImageEditTest(ABC):
"""
Abstract base test class that enforces a common test across all image edit test classes.
"""
@property
def image_edit_function(self):
return litellm.image_edit
@property
def async_image_edit_function(self):
return litellm.aimage_edit
@abstractmethod
def get_base_image_edit_call_args(self) -> dict:
"""Must return the base image edit call args"""
pass
@pytest.fixture(autouse=True)
def _handle_rate_limits(self):
"""Fixture to handle rate limit errors for all test methods"""
try:
yield
except litellm.RateLimitError:
pytest.skip("Rate limit exceeded")
except litellm.InternalServerError:
pytest.skip("Model is overloaded")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_openai_image_edit_litellm_sdk(self, sync_mode):
"""
Test image edit functionality with both sync and async modes.
"""
litellm._turn_on_debug()
try:
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
call_args = self.get_base_image_edit_call_args()
call_args["prompt"] = prompt
if sync_mode:
result = self.image_edit_function(**call_args)
else:
result = await self.async_image_edit_function(**call_args)
print("result from image edit", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
image_bytes = base64.b64decode(image_base64)
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
except litellm.ContentPolicyViolationError as e:
pass
# Get the current directory of the file being run
pwd = os.path.dirname(os.path.realpath(__file__))
def _read_image_bytes(filename: str) -> bytes:
with open(os.path.join(pwd, filename), "rb") as f:
return f.read()
_ISHAAN_GITHUB_BYTES = _read_image_bytes("ishaan_github.png")
_LITELLM_SITE_BYTES = _read_image_bytes("litellm_site.png")
def _make_test_images() -> list:
return [_ISHAAN_GITHUB_BYTES, _LITELLM_SITE_BYTES]
def _make_single_test_image() -> bytes:
return _ISHAAN_GITHUB_BYTES
def get_test_images_as_bytesio():
return [
BytesIO(_ISHAAN_GITHUB_BYTES),
BytesIO(_LITELLM_SITE_BYTES),
]
class TestOpenAIImageEditGPTImage1(BaseLLMImageEditTest):
"""
Concrete implementation of BaseLLMImageEditTest for OpenAI image edits.
"""
def get_base_image_edit_call_args(self) -> dict:
"""Return base call args for OpenAI image edit"""
return {
"model": "gpt-image-1",
"image": _make_test_images(),
}
class TestAzureAIFlux2ImageEdit(BaseLLMImageEditTest):
"""
Concrete implementation of BaseLLMImageEditTest for Azure AI FLUX 2 image edits.
FLUX 2 uses JSON with base64 image instead of multipart/form-data.
"""
def get_base_image_edit_call_args(self) -> dict:
"""Return base call args for Azure AI FLUX 2 image edit"""
return {
"model": "azure_ai/flux.2-pro",
"image": _make_single_test_image(),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_version": "preview",
}
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_openai_image_edit_litellm_router():
litellm._turn_on_debug()
try:
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
router = litellm.Router(
model_list=[
{
"model_name": "gpt-image-1",
"litellm_params": {
"model": "gpt-image-1",
},
}
]
)
result = await router.aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=_make_test_images(),
)
print("result from image edit", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
image_bytes = base64.b64decode(image_base64)
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
except litellm.ContentPolicyViolationError as e:
pass
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_openai_image_edit_with_bytesio():
"""Test image editing using BytesIO objects instead of file readers"""
from litellm import image_edit, aimage_edit
litellm._turn_on_debug()
try:
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Get images as BytesIO objects
bytesio_images = get_test_images_as_bytesio()
result = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=bytesio_images,
)
print("result from image edit with BytesIO", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
image_bytes = base64.b64decode(image_base64)
# Save the image to a file
with open("test_image_edit_bytesio.png", "wb") as f:
f.write(image_bytes)
except litellm.ContentPolicyViolationError as e:
pass
@pytest.mark.asyncio
async def test_azure_image_edit_litellm_sdk():
"""Test Azure image edit with mocked httpx request to validate request body and URL"""
from litellm import aimage_edit
# Mock response for Azure image edit
mock_response = {
"created": 1589478378,
"data": [
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
],
}
class MockResponse:
def __init__(self, json_data, status_code):
self._json_data = json_data
self.status_code = status_code
self.text = json.dumps(json_data)
def json(self):
return self._json_data
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post:
# Configure the mock to return our response
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
test_api_base = "https://ai-api-gw-uae-north.openai.azure.com"
test_api_key = "test-api-key"
test_api_version = "2025-04-01-preview"
result = await aimage_edit(
prompt=prompt,
model="azure/gpt-image-1",
api_base=test_api_base,
api_key=test_api_key,
api_version=test_api_version,
image=_make_test_images(),
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Check the URL
call_args = mock_post.call_args
expected_url = f"{test_api_base}/openai/deployments/gpt-image-1/images/edits?api-version={test_api_version}"
actual_url = (
call_args.args[0] if call_args.args else call_args.kwargs.get("url")
)
print(f"Expected URL: {expected_url}")
print(f"Actual URL: {actual_url}")
assert (
actual_url == expected_url
), f"URL mismatch. Expected: {expected_url}, Got: {actual_url}"
# Check the request body
if "data" in call_args.kwargs:
# For multipart form data, check the data parameter
form_data = call_args.kwargs["data"]
print(
"Form data keys:",
list(form_data.keys()) if hasattr(form_data, "keys") else "Not a dict",
)
# Deployment is in the URL path; Azure rejects model in multipart for this route.
assert (
"model" not in form_data
), "model must not be in form data for Azure /openai/deployments/.../images/edits"
assert "prompt" in form_data, "prompt should be in the form data"
assert (
prompt.strip() in form_data["prompt"]
), f"Expected prompt to contain '{prompt.strip()}'"
# Check headers
headers = call_args.kwargs.get("headers", {})
print("Request headers:", headers)
assert (
"api-key" in headers
), "Azure image edit must use the api-key header, not Authorization: Bearer"
assert headers["api-key"] == test_api_key
assert (
"Authorization" not in headers
), "Azure image edit must not send an Authorization header when an api_key is provided"
print("result from image edit", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
image_bytes = base64.b64decode(image_base64)
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
@pytest.mark.asyncio
async def test_openai_image_edit_cost_tracking():
"""Test OpenAI image edit cost tracking with custom logger"""
from litellm import image_edit, aimage_edit
test_custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [test_custom_logger]
# Mock response for Azure image edit with usage data for cost tracking
mock_response = {
"created": 1589478378,
"data": [
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
],
"usage": {
"total_tokens": 1100,
"input_tokens": 100,
"input_tokens_details": {"image_tokens": 50, "text_tokens": 50},
"output_tokens": 1000,
},
}
class MockResponse:
def __init__(self, json_data, status_code):
self._json_data = json_data
self.status_code = status_code
self.text = json.dumps(json_data)
def json(self):
return self._json_data
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post:
# Configure the mock to return our response
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
result = await aimage_edit(
prompt=prompt,
model="openai/gpt-image-1",
image=_make_test_images(),
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
image_bytes = base64.b64decode(image_base64)
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
await asyncio.sleep(5)
print(
"standard logging payload",
json.dumps(
test_custom_logger.standard_logging_payload, indent=4, default=str
),
)
# check model
assert test_custom_logger.standard_logging_payload["model"] == "gpt-image-1"
assert (
test_custom_logger.standard_logging_payload["custom_llm_provider"]
== "openai"
)
# check response_cost
assert test_custom_logger.standard_logging_payload["response_cost"] is not None
assert test_custom_logger.standard_logging_payload["response_cost"] > 0
@pytest.mark.asyncio
async def test_azure_image_edit_cost_tracking():
"""Test Azure image edit cost tracking with custom logger"""
from litellm import image_edit, aimage_edit
test_custom_logger = TestCustomLogger()
litellm.logging_callback_manager._reset_all_callbacks()
litellm.callbacks = [test_custom_logger]
# Mock response for Azure image edit with usage data for cost tracking
mock_response = {
"created": 1589478378,
"data": [
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
],
"usage": {
"total_tokens": 1100,
"input_tokens": 100,
"input_tokens_details": {"image_tokens": 50, "text_tokens": 50},
"output_tokens": 1000,
},
}
class MockResponse:
def __init__(self, json_data, status_code):
self._json_data = json_data
self.status_code = status_code
self.text = json.dumps(json_data)
def json(self):
return self._json_data
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post:
# Configure the mock to return our response
mock_post.return_value = MockResponse(mock_response, 200)
litellm._turn_on_debug()
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
# Set up test environment variables
result = await aimage_edit(
prompt=prompt,
model="azure/CUSTOM_AZURE_DEPLOYMENT_NAME",
base_model="azure/gpt-image-1",
image=_make_test_images(),
)
# Verify the request was made correctly
mock_post.assert_called_once()
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_base64 = result.data[0].b64_json
if image_base64:
image_bytes = base64.b64decode(image_base64)
# Save the image to a file
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
await asyncio.sleep(5)
print(
"standard logging payload",
json.dumps(
test_custom_logger.standard_logging_payload, indent=4, default=str
),
)
# check model
assert (
test_custom_logger.standard_logging_payload["model"]
== "CUSTOM_AZURE_DEPLOYMENT_NAME"
)
assert (
test_custom_logger.standard_logging_payload["custom_llm_provider"]
== "azure"
)
# check response_cost
assert test_custom_logger.standard_logging_payload["response_cost"] is not None
assert test_custom_logger.standard_logging_payload["response_cost"] > 0
@pytest.mark.asyncio
@pytest.mark.skip(reason="Recraft image edit API only tested locally")
async def test_recraft_image_edit_api():
from litellm import aimage_edit
import requests
litellm._turn_on_debug()
try:
prompt = """
Create a studio ghibli style image that combines all the reference images. Make sure the person looks like a CTO.
"""
result = await aimage_edit(
prompt=prompt,
model="recraft/recraftv3",
image=_make_test_images(),
)
print("result from image edit", result)
# Validate the response meets expected schema
ImageResponse.model_validate(result)
if isinstance(result, ImageResponse) and result.data:
image_url = result.data[0].url
# download the image
image_bytes = requests.get(image_url).content
with open("test_image_edit.png", "wb") as f:
f.write(image_bytes)
except litellm.ContentPolicyViolationError as e:
pass
def test_recraft_image_edit_config():
"""
Test Recraft image edit configuration parameter mapping and request transformation.
"""
from litellm.llms.recraft.image_edit.transformation import RecraftImageEditConfig
from litellm.types.images.main import ImageEditOptionalRequestParams
from litellm.types.router import GenericLiteLLMParams
config = RecraftImageEditConfig()
# Test supported OpenAI params
supported_params = config.get_supported_openai_params("recraftv3")
expected_params = ["n", "response_format", "style"]
assert supported_params == expected_params
# Test parameter mapping (reuses OpenAI logic with filtering)
image_edit_params = ImageEditOptionalRequestParams(
{
"n": 2,
"response_format": "b64_json",
"style": "realistic_image",
"size": "1024x1024", # Should be dropped
"quality": "high", # Should be dropped
}
)
mapped_params = config.map_openai_params(
image_edit_params, "recraftv3", drop_params=True
)
# Should only contain supported params
assert mapped_params["n"] == 2
assert mapped_params["response_format"] == "b64_json"
assert mapped_params["style"] == "realistic_image"
assert "size" not in mapped_params # Should be dropped
assert "quality" not in mapped_params # Should be dropped
# Test request transformation (reuses OpenAI file handling)
mock_image = b"fake_image_data"
prompt = "winter landscape"
litellm_params = GenericLiteLLMParams(api_key="test_key")
data, files = config.transform_image_edit_request(
model="recraftv3",
prompt=prompt,
image=mock_image,
image_edit_optional_request_params={"strength": 0.7, "n": 1},
litellm_params=litellm_params,
headers={},
)
# Check data structure (like OpenAI but with Recraft additions)
assert data["prompt"] == prompt
assert data["strength"] == 0.7 # Recraft-specific parameter
assert data["model"] == "recraftv3"
# Check file structure (reuses OpenAI logic)
assert len(files) == 1
assert files[0][0] == "image" # Field name (not image[] like OpenAI)
assert files[0][1][1] == mock_image # Image data
assert files[0][1][2] == "image/png" # Content type
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_multiple_vs_single_image_edit(sync_mode):
"""Test that both single and multiple image editing work correctly"""
from litellm import image_edit, aimage_edit
litellm._turn_on_debug()
try:
prompt = "Add a soft blue tint to the image(s)"
# Test single image
if sync_mode:
single_result = image_edit(
prompt=prompt,
model="gpt-image-1",
image=_make_single_test_image(),
)
else:
single_result = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=_make_single_test_image(),
)
print("Single image result:", single_result)
ImageResponse.model_validate(single_result)
# Test multiple images
if sync_mode:
multiple_result = image_edit(
prompt=prompt,
model="gpt-image-1",
image=_make_test_images(),
)
else:
multiple_result = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=_make_test_images(),
)
print("Multiple images result:", multiple_result)
ImageResponse.model_validate(multiple_result)
# Both should return valid responses
assert single_result is not None
assert multiple_result is not None
assert single_result.data is not None
assert multiple_result.data is not None
assert len(single_result.data) > 0
assert len(multiple_result.data) > 0
except litellm.ContentPolicyViolationError as e:
pytest.skip(f"Content policy violation: {e}")
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_multiple_image_edit_with_different_formats():
"""Test multiple images editing with different file formats and types"""
from litellm import aimage_edit
litellm._turn_on_debug()
try:
prompt = "Create a cohesive artistic style across all images"
mixed_images = [
_make_single_test_image(),
get_test_images_as_bytesio()[1],
]
result = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=mixed_images,
)
print("Mixed format images result:", result)
ImageResponse.model_validate(result)
assert result is not None
assert result.data is not None
assert len(result.data) > 0
# Save result if available
if result.data and result.data[0].b64_json:
image_bytes = base64.b64decode(result.data[0].b64_json)
with open("test_multiple_image_edit_mixed.png", "wb") as f:
f.write(image_bytes)
except litellm.ContentPolicyViolationError as e:
pytest.skip(f"Content policy violation: {e}")
@pytest.mark.flaky(retries=3, delay=2)
@pytest.mark.asyncio
async def test_image_edit_array_handling():
"""Test that the image parameter correctly handles both single items and arrays"""
from litellm import aimage_edit
# Mock response
mock_response = {
"created": 1589478378,
"data": [
{
"b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
}
],
}
class MockResponse:
def __init__(self, json_data, status_code):
self._json_data = json_data
self.status_code = status_code
self.text = json.dumps(json_data)
def json(self):
return self._json_data
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post:
mock_post.return_value = MockResponse(mock_response, 200)
prompt = "Test prompt"
# Test 1: Single image (should be converted to list internally)
result1 = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=_make_single_test_image(),
)
# Test 2: Multiple images (already a list)
result2 = await aimage_edit(
prompt=prompt,
model="gpt-image-1",
image=_make_test_images(),
)
# Both valid calls should succeed
ImageResponse.model_validate(result1)
ImageResponse.model_validate(result2)
# Verify that both calls were made to the API
assert mock_post.call_count == 2
Reference in New Issue View Git Blame Copy Permalink