add comments

This commit is contained in:
Xi Yan 2025-01-30 15:01:18 -08:00
parent ff73f0a17b
commit 5288b602ec
4 changed files with 61 additions and 25 deletions

View file

@ -2446,10 +2446,12 @@
"type": "object", "type": "object",
"properties": { "properties": {
"url": { "url": {
"$ref": "#/components/schemas/URL" "$ref": "#/components/schemas/URL",
"description": "A URL of the image or data URL in the format of data:image/{type};base64,{data}. Note that URL could have length limits."
}, },
"data": { "data": {
"type": "string" "type": "string",
"description": "base64 encoded image data as string"
} }
}, },
"additionalProperties": false, "additionalProperties": false,

View file

@ -1473,8 +1473,12 @@ components:
properties: properties:
url: url:
$ref: '#/components/schemas/URL' $ref: '#/components/schemas/URL'
description: >-
A URL of the image or data URL in the format of data:image/{type};base64,{data}.
Note that URL could have length limits.
data: data:
type: string type: string
description: base64 encoded image data as string
additionalProperties: false additionalProperties: false
description: >- description: >-
Image as a base64 encoded string or an URL Image as a base64 encoded string or an URL

View file

@ -19,8 +19,16 @@ class URL(BaseModel):
class _URLOrData(BaseModel): class _URLOrData(BaseModel):
"""
A URL or a base64 encoded string
:param url: A URL of the image or data URL in the format of data:image/{type};base64,{data}. Note that URL could have length limits.
:param data: base64 encoded image data as string
"""
url: Optional[URL] = None url: Optional[URL] = None
# data is a base64 encoded string # data is a base64 encoded string
# TODO: annotate with contentEncoding="base64" in OpenAPI schema
data: Optional[str] = None data: Optional[str] = None
@model_validator(mode="before") @model_validator(mode="before")

View file

@ -5,7 +5,7 @@
# the root directory of this source tree. # the root directory of this source tree.
import base64 import base64
import os import pathlib
import pytest import pytest
from pydantic import BaseModel from pydantic import BaseModel
@ -48,14 +48,31 @@ def get_weather_tool_definition():
} }
# @pytest.fixture
# def base64_image_url():
# image_path = os.path.join(os.path.dirname(__file__), "dog.png")
# with open(image_path, "rb") as image_file:
# # Convert the image to base64
# base64_string = base64.b64encode(image_file.read()).decode("utf-8")
# base64_url = f"data:image/png;base64,{base64_string}"
# return base64_url
@pytest.fixture @pytest.fixture
def base64_image_url(): def image_path():
image_path = os.path.join(os.path.dirname(__file__), "dog.png") return pathlib.Path(__file__).parent / "dog.png"
with open(image_path, "rb") as image_file:
# Convert the image to base64
base64_string = base64.b64encode(image_file.read()).decode("utf-8") @pytest.fixture
base64_url = f"data:image/png;base64,{base64_string}" def base64_image_data(image_path):
return base64_url # Convert the image to base64
return base64.b64encode(image_path.read_bytes()).decode("utf-8")
@pytest.fixture
def base64_image_url(base64_image_data, image_path):
# suffix includes the ., so we remove it
return f"data:image/{image_path.suffix[1:]};base64,{base64_image_data}"
def test_text_completion_non_streaming(llama_stack_client, text_model_id): def test_text_completion_non_streaming(llama_stack_client, text_model_id):
@ -353,25 +370,30 @@ def test_image_chat_completion_streaming(llama_stack_client, vision_model_id):
assert any(expected in streamed_content for expected in {"dog", "puppy", "pup"}) assert any(expected in streamed_content for expected in {"dog", "puppy", "pup"})
def test_image_chat_completion_base64_url( @pytest.mark.parametrize("type_", ["url", "data"])
llama_stack_client, vision_model_id, base64_image_url def test_image_chat_completion_base64(
llama_stack_client, vision_model_id, base64_image_data, base64_image_url, type_
): ):
message = { image_spec = {
"role": "user", "url": {
"content": [ "type": "image",
{ "image": {
"type": "image", "url": {
"image": { "uri": base64_image_url,
"url": {
"uri": base64_image_url,
},
}, },
}, },
{ },
"type": "text", "data": {
"text": "Describe what is in this image.", "type": "image",
"image": {
"data": base64_image_data,
}, },
], },
}[type_]
message = {
"role": "user",
"content": [image_spec],
} }
response = llama_stack_client.inference.chat_completion( response = llama_stack_client.inference.chat_completion(
model_id=vision_model_id, model_id=vision_model_id,