scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

test: update tests

Browse source
This commit is contained in:
Krrish Dholakia 2025-09-01 18:51:05 -07:00
parent d4cc748b40
commit 6a2da54fa8
2 changed files with 14 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -46,6 +46,7 @@ from litellm.types.llms.openai import (
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolParamFunctionChunk,
ImageURLListItem,
ImageURLObject,
OpenAIChatCompletionFinishReason,
)
@ -795,9 +796,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
def _extract_image_response_from_parts(
self, parts: List[HttpxPartType]
) -> Optional[List[ImageURLObject]]:
) -> Optional[List[ImageURLListItem]]:
"""Extract image response from parts if present"""
images: List[ImageURLObject] = []
images: List[ImageURLListItem] = []
for part in parts:
if "inlineData" in part:
mime_type = part["inlineData"]["mimeType"]
@ -805,7 +806,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if mime_type.startswith("image/"):
# Convert base64 data to data URI format
data_uri = f"data:{mime_type};base64,{data}"
images.append(ImageURLObject(url=data_uri, detail="auto"))
images.append(
ImageURLListItem(
image_url=ImageURLObject(url=data_uri, detail="auto"),
index=0,
type="image_url",
)
)
return images
def _extract_audio_response_from_parts(
@ -1134,7 +1141,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
tools: Optional[List[ChatCompletionToolCallChunk]],
functions: Optional[ChatCompletionToolCallFunctionChunk],
chat_completion_logprobs: Optional[ChoiceLogprobs],
image_response: Optional[List[ImageURLObject]],
image_response: Optional[List[ImageURLListItem]],
) -> StreamingChoices:
"""
Helper method to create a streaming choice object for Vertex AI
@ -1222,7 +1229,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
grounding_metadata: List[dict] = []
url_context_metadata: List[dict] = []
image_response: Optional[List[ImageURLObject]] = None
image_response: Optional[List[ImageURLListItem]] = None
safety_ratings: List = []
citation_metadata: List = []
chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"}

4
tests/llm_translation/test_gemini.py
View file

@ -597,7 +597,7 @@ async def test_gemini_image_generation_async():
CONTENT = response.choices[0].message.content
IMAGE_URL = response.choices[0].message.images[0]
IMAGE_URL = response.choices[0].message.images[0]["image_url"]
print("IMAGE_URL: ", IMAGE_URL)
assert CONTENT is not None, "CONTENT is not None"
@ -629,7 +629,7 @@ async def test_gemini_image_generation_async_stream():
and chunk.choices[0].delta.images is not None
and len(chunk.choices[0].delta.images) > 0
):
model_response_image = chunk.choices[0].delta.images[0]
model_response_image = chunk.choices[0].delta.images[0]["image_url"]
print("MODEL_RESPONSE_IMAGE: ", model_response_image)
assert model_response_image is not None
assert model_response_image["url"].startswith("data:image/png;base64,")