Make image writing to and reading from LLMContext more robust; let's allow storing in context image types other than JPEG, meaning not lossily and unnecessarily re-encoding non-JPEG images as JPEG.

This commit is contained in:
Paul Kompfner
2025-12-15 10:34:33 -05:00
parent e604e9b490
commit 54926f390d
7 changed files with 46 additions and 43 deletions

View File

@@ -283,11 +283,14 @@ class AnthropicLLMAdapter(BaseLLMAdapter[AnthropicLLMInvocationParams]):
# handle image_url -> image conversion
if item["type"] == "image_url":
if item["image_url"]["url"].startswith("data:"):
# Extract MIME type from data URL (format: "data:image/jpeg;base64,...")
url = item["image_url"]["url"]
mime_type = url.split(":")[1].split(";")[0]
item["type"] = "image"
item["source"] = {
"type": "base64",
"media_type": "image/jpeg",
"data": item["image_url"]["url"].split(",")[1],
"media_type": mime_type,
"data": url.split(",")[1],
}
del item["image_url"]
elif item["image_url"]["url"].startswith("http"):

View File

@@ -257,14 +257,15 @@ class AWSBedrockLLMAdapter(BaseLLMAdapter[AWSBedrockLLMInvocationParams]):
# handle image_url -> image conversion
if item["type"] == "image_url":
if item["image_url"]["url"].startswith("data:"):
# Extract format from data URL (format: "data:image/jpeg;base64,...")
url = item["image_url"]["url"]
mime_type = url.split(":")[1].split(";")[0]
# Bedrock expects format like "jpeg", "png" etc., not "image/jpeg"
image_format = mime_type.split("/")[1]
new_item = {
"image": {
"format": "jpeg",
"source": {
"bytes": base64.b64decode(
item["image_url"]["url"].split(",")[1]
)
},
"format": image_format,
"source": {"bytes": base64.b64decode(url.split(",")[1])},
}
}
new_content.append(new_item)

View File

@@ -399,11 +399,14 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
if c["type"] == "text":
parts.append(Part(text=c["text"]))
elif c["type"] == "image_url" and c["image_url"]["url"].startswith("data:"):
# Extract MIME type from data URL (format: "data:image/jpeg;base64,...")
url = c["image_url"]["url"]
mime_type = url.split(":")[1].split(";")[0]
parts.append(
Part(
inline_data=Blob(
mime_type="image/jpeg",
data=base64.b64decode(c["image_url"]["url"].split(",")[1]),
mime_type=mime_type,
data=base64.b64decode(url.split(",")[1]),
)
)
)

View File

@@ -227,7 +227,7 @@ class ImageRawFrame:
Parameters:
image: Raw image bytes.
size: Image dimensions as (width, height) tuple.
format: Image format (e.g., 'JPEG', 'PNG').
format: Image format (e.g., 'RGB', 'RGBA').
"""
image: bytes
@@ -1468,16 +1468,19 @@ class UserImageRawFrame(InputImageRawFrame):
@dataclass
class AssistantImageRawFrame(OutputImageRawFrame):
"""Frame containing image generated by the assistant.
"""Frame containing an image generated by the assistant.
An image generated by the assistant. Gets appended to the LLM context.
Contains both the raw frame for display (superclass functionality) as well
as the original image, which can get used directly in LLM contexts.
Parameters:
original_jpeg: The already-JPEG-encoded image bytes, which may be
appended directly to the LLM context without further encoding.
original_data: The original image data, which can get used directly in
an LLM context message without further encoding.
original_mime_type: The MIME type of the original image data.
"""
original_jpeg: Optional[bytes] = None
original_data: Optional[bytes] = None
original_mime_type: Optional[str] = None
@dataclass

View File

@@ -150,15 +150,17 @@ class LLMContext:
Args:
role: The role of this message (defaults to "user").
format: Image format (e.g., 'RGB', 'RGBA').
format: Image format (e.g., 'RGB', 'RGBA', or, if already encoded,
the MIME type like 'image/jpeg').
size: Image dimensions as (width, height) tuple.
image: Raw image bytes.
text: Optional text to include with the image.
"""
# Format is a mime type: image is already encoded
image_already_encoded = format.startswith("image/")
def encode_image():
if format == "JPEG":
# Already JPEG-encoded
if image_already_encoded:
bytes = image
else:
# Encode to JPEG
@@ -170,7 +172,7 @@ class LLMContext:
encoded_image = await asyncio.to_thread(encode_image)
url = f"data:image/jpeg;base64,{encoded_image}"
url = f"data:{format if image_already_encoded else 'image/jpeg'};base64,{encoded_image}"
return LLMContext.create_image_url_message(role=role, url=url, text=text)
@@ -351,8 +353,8 @@ class LLMContext:
"""Add a message containing an image frame.
Args:
format: Image format (e.g., 'RGB', 'RGBA', or, if already
JPEG-encoded, "JPEG").
format: Image format (e.g., 'RGB', 'RGBA', or, if already encoded,
the MIME type like 'image/jpeg').
size: Image dimensions as (width, height) tuple.
image: Raw image bytes.
text: Optional text to include with the image.

View File

@@ -833,11 +833,11 @@ class LLMAssistantAggregator(LLMContextAggregator):
async def _handle_assistant_image_frame(self, frame: AssistantImageRawFrame):
logger.debug(f"{self} Appending AssistantImageRawFrame to LLM context (size: {frame.size})")
if frame.original_jpeg:
if frame.original_data and frame.original_mime_type:
await self._context.add_image_frame_message(
format="JPEG",
format=frame.original_mime_type,
size=frame.size, # Technically doesn't matter, since already encoded
image=frame.original_jpeg,
image=frame.original_data,
role="assistant",
)
else:

View File

@@ -479,11 +479,16 @@ class GoogleLLMContext(OpenAILLMContext):
if c["type"] == "text":
parts.append(Part(text=c["text"]))
elif c["type"] == "image_url":
# Extract MIME type from data URL (format: "data:image/jpeg;base64,...")
url = c["image_url"]["url"]
mime_type = (
url.split(":")[1].split(";")[0] if url.startswith("data:") else "image/jpeg"
)
parts.append(
Part(
inline_data=Blob(
mime_type="image/jpeg",
data=base64.b64decode(c["image_url"]["url"].split(",")[1]),
mime_type=mime_type,
data=base64.b64decode(url.split(",")[1]),
)
)
)
@@ -995,21 +1000,13 @@ class GoogleLLMService(LLMService):
elif part.inline_data and part.inline_data.data:
# Here we assume that inline_data is an image.
image = Image.open(io.BytesIO(part.inline_data.data))
# NOTE: Gemini 3 Pro Image seems to always give
# JPEGs. It expects us to send back the
# original JPEG data in the context, along with
# the corresponding thought signature. JPEG
# happens to be the format our universal
# context uses for images, so we can just pass
# it through as-is.
await self.push_frame(
AssistantImageRawFrame(
image=image.tobytes(),
size=image.size,
format="RGB",
original_jpeg=part.inline_data.data
if part.inline_data.mime_type == "image/jpeg"
else None,
original_data=part.inline_data.data,
original_mime_type=part.inline_data.mime_type,
)
)
@@ -1037,12 +1034,6 @@ class GoogleLLMService(LLMService):
if part.function_call:
bookmark["function_call"] = function_call_id
elif part.inline_data and part.inline_data.data:
# With Gemini 3 Pro (where sending the
# thought signature is required for images)
# this is the JPEG-encoded image data that
# we sent to be written to the context
# as-is, so it is usable as a bookmark (it
# will match the context data).
bookmark["inline_data"] = part.inline_data
elif part.text is not None:
# Account for Gemini 3 Pro trailing