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:
@@ -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"):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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]),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user