tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 02:23:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Gemini image generation output support (#9646)

Browse Source 
* fix(gemini/transformation.py): make GET request to get uri details, if cannot be inferred

* fix: fix linting errors

* Revert "fix: fix linting errors"

This reverts commit 926a5a527ff27a107b39da8f5a26b0ee8e2d9884.

* fix(gemini/transformation.py): modalities param support

Partially resolves https://github.com/BerriAI/litellm/issues/9237

* feat(google_ai_studio/): add image generation support

Closes https://github.com/BerriAI/litellm/issues/9237

* fix: fix types

* fix: fix ruff check
This commit is contained in:
Krish Dholakia
2025-04-04 20:37:48 -07:00
committed by GitHub
parent 90a4dfab3c
commit af42e5855f
6 changed files with 123 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/gemini/chat/transformation.py
+1
View File
@@ -81,6 +81,7 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
"stop",
"logprobs",
"frequency_penalty",
"modalities",
]
def map_openai_params(
litellm/llms/vertex_ai/gemini/transformation.py
+4 -9
View File
@@ -224,17 +224,12 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
if not file_id:
continue
mime_type = format or _get_image_mime_type_from_url(file_id)
if mime_type is not None:
_part = PartType(
file_data=FileDataType(
file_uri=file_id,
mime_type=mime_type,
)
try:
_part = _process_gemini_image(
image_url=file_id, format=format
)
_parts.append(_part)
else:
except Exception:
raise Exception(
"Unable to determine mime type for file_id: {}, set this explicitly using message[{}].content[{}].file.format".format(
file_id, msg_i, element_idx
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
+86 -49
View File
@@ -208,6 +208,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
"seed",
"logprobs",
"top_logprobs", # Added this to list of supported openAI params
"modalities",
]
def map_tool_choice_values(
@@ -312,6 +313,30 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
old_schema = _build_vertex_schema(parameters=old_schema)
return old_schema
def apply_response_schema_transformation(self, value: dict, optional_params: dict):
# remove 'additionalProperties' from json schema
value = _remove_additional_properties(value)
# remove 'strict' from json schema
value = _remove_strict_from_schema(value)
if value["type"] == "json_object":
optional_params["response_mime_type"] = "application/json"
elif value["type"] == "text":
optional_params["response_mime_type"] = "text/plain"
if "response_schema" in value:
optional_params["response_mime_type"] = "application/json"
optional_params["response_schema"] = value["response_schema"]
elif value["type"] == "json_schema": # type: ignore
if "json_schema" in value and "schema" in value["json_schema"]: # type: ignore
optional_params["response_mime_type"] = "application/json"
optional_params["response_schema"] = value["json_schema"]["schema"] # type: ignore
if "response_schema" in optional_params and isinstance(
optional_params["response_schema"], dict
):
optional_params["response_schema"] = self._map_response_schema(
value=optional_params["response_schema"]
)
def map_openai_params(
self,
non_default_params: Dict,
@@ -322,58 +347,39 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
for param, value in non_default_params.items():
if param == "temperature":
optional_params["temperature"] = value
if param == "top_p":
elif param == "top_p":
optional_params["top_p"] = value
if (
elif (
param == "stream" and value is True
): # sending stream = False, can cause it to get passed unchecked and raise issues
optional_params["stream"] = value
if param == "n":
elif param == "n":
optional_params["candidate_count"] = value
if param == "stop":
elif param == "stop":
if isinstance(value, str):
optional_params["stop_sequences"] = [value]
elif isinstance(value, list):
optional_params["stop_sequences"] = value
if param == "max_tokens" or param == "max_completion_tokens":
elif param == "max_tokens" or param == "max_completion_tokens":
optional_params["max_output_tokens"] = value
if param == "response_format" and isinstance(value, dict): # type: ignore
# remove 'additionalProperties' from json schema
value = _remove_additional_properties(value)
# remove 'strict' from json schema
value = _remove_strict_from_schema(value)
if value["type"] == "json_object":
optional_params["response_mime_type"] = "application/json"
elif value["type"] == "text":
optional_params["response_mime_type"] = "text/plain"
if "response_schema" in value:
optional_params["response_mime_type"] = "application/json"
optional_params["response_schema"] = value["response_schema"]
elif value["type"] == "json_schema": # type: ignore
if "json_schema" in value and "schema" in value["json_schema"]: # type: ignore
optional_params["response_mime_type"] = "application/json"
optional_params["response_schema"] = value["json_schema"]["schema"] # type: ignore
if "response_schema" in optional_params and isinstance(
optional_params["response_schema"], dict
):
optional_params["response_schema"] = self._map_response_schema(
value=optional_params["response_schema"]
)
if param == "frequency_penalty":
elif param == "response_format" and isinstance(value, dict): # type: ignore
self.apply_response_schema_transformation(
value=value, optional_params=optional_params
)
elif param == "frequency_penalty":
optional_params["frequency_penalty"] = value
if param == "presence_penalty":
elif param == "presence_penalty":
optional_params["presence_penalty"] = value
if param == "logprobs":
elif param == "logprobs":
optional_params["responseLogprobs"] = value
if param == "top_logprobs":
elif param == "top_logprobs":
optional_params["logprobs"] = value
if (param == "tools" or param == "functions") and isinstance(value, list):
elif (param == "tools" or param == "functions") and isinstance(value, list):
optional_params["tools"] = self._map_function(value=value)
optional_params["litellm_param_is_function_call"] = (
True if param == "functions" else False
)
if param == "tool_choice" and (
elif param == "tool_choice" and (
isinstance(value, str) or isinstance(value, dict)
):
_tool_choice_value = self.map_tool_choice_values(
@@ -381,8 +387,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
)
if _tool_choice_value is not None:
optional_params["tool_choice"] = _tool_choice_value
if param == "seed":
elif param == "seed":
optional_params["seed"] = value
elif param == "modalities" and isinstance(value, list):
response_modalities = []
for modality in value:
if modality == "text":
response_modalities.append("TEXT")
elif modality == "image":
response_modalities.append("IMAGE")
else:
response_modalities.append("MODALITY_UNSPECIFIED")
optional_params["responseModalities"] = response_modalities
if litellm.vertex_ai_safety_settings is not None:
optional_params["safety_settings"] = litellm.vertex_ai_safety_settings
@@ -493,6 +509,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
for part in parts:
if "text" in part:
_content_str += part["text"]
elif "inlineData" in part: # base64 encoded image
_content_str += "data:{};base64,{}".format(
part["inlineData"]["mimeType"], part["inlineData"]["data"]
)
if _content_str:
return _content_str
return None
@@ -685,7 +706,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
chat_completion_logprobs: Optional[ChoiceLogprobs] = None
tools: Optional[List[ChatCompletionToolCallChunk]] = []
functions: Optional[ChatCompletionToolCallFunctionChunk] = None
for idx, candidate in enumerate(_candidates):
if "content" not in candidate:
continue
@@ -698,16 +719,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if "citationMetadata" in candidate:
citation_metadata.append(candidate["citationMetadata"])
if "parts" in candidate["content"]:
chat_completion_message["content"] = VertexGeminiConfig().get_assistant_content_message(
chat_completion_message[
"content"
] = VertexGeminiConfig().get_assistant_content_message(
parts=candidate["content"]["parts"]
)
functions, tools = self._transform_parts(
parts=candidate["content"]["parts"],
index=candidate.get("index", idx),
is_function_call=litellm_params.get("litellm_param_is_function_call"),
is_function_call=litellm_params.get(
"litellm_param_is_function_call"
),
)
if "logprobsResult" in candidate:
@@ -723,7 +748,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
if functions is not None:
chat_completion_message["function_call"] = functions
choice = litellm.Choices(
finish_reason=candidate.get("finishReason", "stop"),
index=candidate.get("index", idx),
@@ -733,7 +758,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
)
model_response.choices.append(choice)
return grounding_metadata, safety_ratings, citation_metadata
def transform_response(
@@ -785,7 +810,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
_candidates = completion_response.get("candidates")
if _candidates and len(_candidates) > 0:
content_policy_violations = VertexGeminiConfig().get_flagged_finish_reasons()
content_policy_violations = (
VertexGeminiConfig().get_flagged_finish_reasons()
)
if (
"finishReason" in _candidates[0]
and _candidates[0]["finishReason"] in content_policy_violations.keys()
@@ -795,12 +822,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
completion_response=completion_response,
)
model_response.choices = [] # type: ignore
model_response.choices = []
try:
grounding_metadata, safety_ratings, citation_metadata = [], [], []
if _candidates:
grounding_metadata, safety_ratings, citation_metadata = self._process_candidates(
(
grounding_metadata,
safety_ratings,
citation_metadata,
) = self._process_candidates(
_candidates, model_response, litellm_params
)
@@ -809,14 +840,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
## ADD METADATA TO RESPONSE ##
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata)
model_response._hidden_params["vertex_ai_grounding_metadata"] = grounding_metadata
model_response._hidden_params[
"vertex_ai_grounding_metadata"
] = grounding_metadata
setattr(model_response, "vertex_ai_safety_results", safety_ratings)
model_response._hidden_params["vertex_ai_safety_results"] = safety_ratings # older approach - maintaining to prevent regressions
model_response._hidden_params[
"vertex_ai_safety_results"
] = safety_ratings # older approach - maintaining to prevent regressions
## ADD CITATION METADATA ##
setattr(model_response, "vertex_ai_citation_metadata", citation_metadata)
model_response._hidden_params["vertex_ai_citation_metadata"] = citation_metadata # older approach - maintaining to prevent regressions
model_response._hidden_params[
"vertex_ai_citation_metadata"
] = citation_metadata # older approach - maintaining to prevent regressions
except Exception as e:
raise VertexAIError(
litellm/types/llms/vertex_ai.py
+9 -3
View File
@@ -56,12 +56,17 @@ class HttpxCodeExecutionResult(TypedDict):
output: str
class HttpxBlobType(TypedDict):
mimeType: str
data: str
class HttpxPartType(TypedDict, total=False):
text: str
inline_data: BlobType
file_data: FileDataType
inlineData: HttpxBlobType
fileData: FileDataType
functionCall: HttpxFunctionCall
function_response: FunctionResponse
functionResponse: FunctionResponse
executableCode: HttpxExecutableCode
codeExecutionResult: HttpxCodeExecutionResult
@@ -160,6 +165,7 @@ class GenerationConfig(TypedDict, total=False):
seed: int
responseLogprobs: bool
logprobs: int
responseModalities: List[Literal["TEXT", "IMAGE", "AUDIO", "VIDEO"]]
class Tools(TypedDict, total=False):
tests/llm_translation/test_gemini.py
+12 -1
View File
@@ -11,7 +11,8 @@ from base_llm_unit_tests import BaseLLMChatTest
from litellm.llms.vertex_ai.context_caching.transformation import (
separate_cached_messages,
)
import litellm
from litellm import completion
class TestGoogleAIStudioGemini(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
@@ -72,3 +73,13 @@ def test_gemini_context_caching_separate_messages():
print(non_cached_messages)
assert len(cached_messages) > 0, "Cached messages should be present"
assert len(non_cached_messages) > 0, "Non-cached messages should be present"
def test_gemini_image_generation():
# litellm._turn_on_debug()
response = completion(
model="gemini/gemini-2.0-flash-exp-image-generation",
messages=[{"role": "user", "content": "Generate an image of a cat"}],
modalities=["image", "text"],
)
assert response.choices[0].message.content is not None
tests/llm_translation/test_optional_params.py
+11 -1
View File
@@ -1405,6 +1405,17 @@ def test_azure_modalities_param():
assert optional_params["audio"] == {"type": "audio_input", "input": "test.wav"}
def test_gemini_modalities_param():
optional_params = get_optional_params(
model="gemini-1.5-pro",
custom_llm_provider="gemini",
modalities=["text", "image"],
)
assert optional_params["responseModalities"] == ["TEXT", "IMAGE"]
def test_azure_response_format_param():
optional_params = litellm.get_optional_params(
@@ -1430,4 +1441,3 @@ def test_anthropic_unified_reasoning_content(model, provider):
reasoning_effort="high",
)
assert optional_params["thinking"] == {"type": "enabled", "budget_tokens": 4096}