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

Handle more gemini tool calling edge cases + support bedrock 'stable-image-core' (#10351)

Browse Source 
* test(test_amazing_vertex_completion.py): try to repro https://github.com/BerriAI/litellm/issues/10319

* fix(common_utils.py): handle edge case on tools

Fixes https://github.com/BerriAI/litellm/issues/10319

* test: add unit testing for infinite loops

* fix(amazon_stability3_transformation.py): support 'stable-image-core' transformation

Fixes https://github.com/BerriAI/litellm/issues/8488

* test: add unit testing for stable image core model

* test: update test
This commit is contained in:
Krish Dholakia
2025-04-28 14:22:29 -07:00
committed by GitHub
parent a5ef8a9556
commit bf9382a182
6 changed files with 409 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/bedrock/image/amazon_stability3_transformation.py
+1 -1
View File
@@ -60,7 +60,7 @@ class AmazonStability3Config:
if model:
if "sd3" in model or "sd3.5" in model:
return True
if "stable-image-ultra-v1" in model:
if "stable-image" in model:
return True
return False
litellm/llms/vertex_ai/common_utils.py
+31
View File
@@ -192,6 +192,9 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
# * https://stackoverflow.com/a/58841311
# * https://github.com/pydantic/pydantic/discussions/4872
convert_anyof_null_to_nullable(parameters)
# Handle empty items objects
process_items(parameters)
add_object_type(parameters)
# Postprocessing
# Filter out fields that don't exist in Schema
@@ -199,9 +202,27 @@ def _build_vertex_schema(parameters: dict, add_property_ordering: bool = False):
if add_property_ordering:
set_schema_property_ordering(parameters)
return parameters
def process_items(schema, depth=0):
if depth > DEFAULT_MAX_RECURSE_DEPTH:
raise ValueError(
f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting."
)
if isinstance(schema, dict):
if "items" in schema and schema["items"] == {}:
schema["items"] = {"type": "object"}
for key, value in schema.items():
if isinstance(value, dict):
process_items(value, depth + 1)
elif isinstance(value, list):
for item in value:
if isinstance(item, dict):
process_items(item, depth + 1)
def set_schema_property_ordering(
schema: Dict[str, Any], depth: int = 0
) -> Dict[str, Any]:
@@ -285,6 +306,9 @@ def convert_anyof_null_to_nullable(schema, depth=0):
# remove null type
anyof.remove(atype)
contains_null = True
elif "type" not in atype and len(atype) == 0:
# Handle empty object case
atype["type"] = "object"
if len(anyof) == 0:
# Edge case: response schema with only null type present is invalid in Vertex AI
@@ -296,6 +320,13 @@ def convert_anyof_null_to_nullable(schema, depth=0):
if contains_null:
# set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python
for atype in anyof:
# Remove items field if type is array and items is empty
if (
atype.get("type") == "array"
and "items" in atype
and not atype["items"]
):
atype.pop("items")
atype["nullable"] = True
properties = schema.get("properties", None)
tests/code_coverage_tests/recursive_detector.py
+1
View File
@@ -19,6 +19,7 @@ IGNORE_FUNCTIONS = [
"_sanitize_request_body_for_spend_logs_payload", # testing added for circular reference
"_sanitize_value", # testing added for circular reference
"set_schema_property_ordering", # testing added for infinite recursion
"process_items", # testing added for infinite recursion + max depth set.
]
tests/litellm/llms/bedrock/image/test_amazon_stability3_transformation.py
+20
View File
@@ -0,0 +1,20 @@
import json
import os
import sys
import pytest
from fastapi.testclient import TestClient
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from unittest.mock import MagicMock, patch
from litellm.llms.bedrock.image.amazon_stability3_transformation import (
AmazonStability3Config,
)
def test_stability_image_core_is_v3_model():
model = "stability.stable-image-core-v1:1"
assert AmazonStability3Config._is_stability_3_model(model)
tests/litellm/llms/vertex_ai/test_vertex_ai_common_utils.py
+128
View File
@@ -164,3 +164,131 @@ def test_set_schema_property_ordering_with_excessive_nesting():
match=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting.",
):
set_schema_property_ordering(schema)
def test_build_vertex_schema():
"""Test build_vertex_schema with a sample schema"""
from litellm.llms.vertex_ai.common_utils import _build_vertex_schema
parameters = {
"properties": {
"state": {
"properties": {
"messages": {"items": {}, "type": "array"},
"conversation_id": {"type": "string"},
},
"required": ["messages", "conversation_id"],
"type": "object",
},
"config": {
"description": "Configuration for a Runnable.",
"properties": {
"tags": {"items": {"type": "string"}, "type": "array"},
"metadata": {"type": "object"},
"callbacks": {
"anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]
},
"run_name": {"type": "string"},
"max_concurrency": {
"anyOf": [{"type": "integer"}, {"type": "null"}]
},
"recursion_limit": {"type": "integer"},
"configurable": {"type": "object"},
"run_id": {
"anyOf": [
{"format": "uuid", "type": "string"},
{"type": "null"},
]
},
},
"type": "object",
},
"kwargs": {"default": None, "type": "object"},
},
"required": ["state", "config"],
"type": "object",
}
expected_output = {
"properties": {
"state": {
"properties": {
"messages": {"items": {"type": "object"}, "type": "array"},
"conversation_id": {"type": "string"},
},
"required": ["messages", "conversation_id"],
"type": "object",
},
"config": {
"description": "Configuration for a Runnable.",
"properties": {
"tags": {"items": {"type": "string"}, "type": "array"},
"metadata": {"type": "object"},
"callbacks": {
"anyOf": [
{"type": "array", "nullable": True},
{"type": "object", "nullable": True},
]
},
"run_name": {"type": "string"},
"max_concurrency": {
"anyOf": [{"type": "integer", "nullable": True}]
},
"recursion_limit": {"type": "integer"},
"configurable": {"type": "object"},
"run_id": {
"anyOf": [
{"format": "uuid", "type": "string", "nullable": True}
]
},
},
"type": "object",
},
"kwargs": {"default": None, "type": "object"},
},
"required": ["state", "config"],
"type": "object",
}
assert _build_vertex_schema(parameters) == expected_output
def test_process_items_with_excessive_nesting():
"""Test process_items with excessive nesting > max levels +1 deep."""
# generate a schema with excessive nesting
from litellm.llms.vertex_ai.common_utils import process_items
schema = {"type": "object", "properties": {}}
current = schema
for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1):
current["properties"] = {"nested": {"type": "object", "properties": {}}}
current = current["properties"]["nested"]
with pytest.raises(
ValueError,
match=f"Max depth of {DEFAULT_MAX_RECURSE_DEPTH} exceeded while processing schema. Please check the schema for excessive nesting.",
):
process_items(schema)
def test_process_items_basic():
"""Test basic functionality of process_items."""
from litellm.llms.vertex_ai.common_utils import process_items
# Test empty items
schema = {"type": "array", "items": {}}
process_items(schema)
assert schema["items"] == {"type": "object"}
# Test nested items
schema = {"type": "array", "items": {"type": "array", "items": {}}}
process_items(schema)
assert schema["items"]["items"] == {"type": "object"}
# Test items in properties
schema = {
"type": "object",
"properties": {"nested": {"type": "array", "items": {}}},
}
process_items(schema)
assert schema["properties"]["nested"]["items"] == {"type": "object"}
tests/local_testing/test_amazing_vertex_completion.py
+228
View File
@@ -3478,3 +3478,231 @@ def test_litellm_api_base(monkeypatch, provider, route):
mock_client.assert_called()
assert mock_client.call_args.kwargs["url"].startswith("https://litellm.com")
def test_gemini_tool_calling_working_demo():
load_vertex_ai_credentials()
litellm._turn_on_debug()
args = {
"messages": [
{
"content": "\n You are a helpful assistant who can help with questions on customers business or personal finances.\n Use the results from the available tools to answer the question.\n ",
"role": "system"
},
{
"content": "Hello",
"role": "user"
}
],
"max_completion_tokens": 1000,
"temperature": 0.0,
"tools": [
{
"type": "function",
"function": {
"name": "test_agent",
"description": "This tool helps find relevant help content",
"parameters": {
"properties": {
"state": {
"properties": {
"messages": {
"items": {
"type": "object"
},
"type": "array"
},
"conversation_id": {
"type": "string"
}
},
"required": [
"messages",
"conversation_id"
],
"type": "object"
},
"config": {
"description": "Configuration for a Runnable.",
"properties": {
"tags": {
"items": {
"type": "string"
},
"type": "array"
},
"metadata": {
"type": "object"
},
"callbacks": {
"anyOf": [
{"type": "array"},
{"type": "object"},
{"type": "null"}
],
},
"run_name": {
"type": "string"
},
"max_concurrency": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
]
},
"recursion_limit": {
"type": "integer"
},
"configurable": {
"type": "object"
},
"run_id": {
"anyOf": [
{
"format": "uuid",
"type": "string"
},
{
"type": "null"
}
]
}
},
"type": "object"
},
"kwargs": {
"default": None,
"type": "object"
}
},
"required": [
"state",
"config"
],
"type": "object"
}
}
}
]
}
response = completion(model="vertex_ai/gemini-2.0-flash", **args)
print(response)
def test_gemini_tool_calling_not_working():
load_vertex_ai_credentials()
litellm._turn_on_debug()
args = {
"messages": [
{
"content": "\n You are a helpful assistant who can help with questions on customers business or personal finances.\n Use the results from the available tools to answer the question.\n ",
"role": "system"
},
{
"content": "Hello",
"role": "user"
}
],
"max_completion_tokens": 1000,
"temperature": 0.0,
"tools": [
{
"type": "function",
"function": {
"name": "test_agent",
"description": "This tool helps find relevant help content",
"parameters": {
"properties": {
"state": {
"properties": {
"messages": {
"items": {},
"type": "array"
},
"conversation_id": {
"type": "string"
}
},
"required": [
"messages",
"conversation_id"
],
"type": "object"
},
"config": {
"description": "Configuration for a Runnable.",
"properties": {
"tags": {
"items": {
"type": "string"
},
"type": "array"
},
"metadata": {
"type": "object"
},
"callbacks": {
"anyOf": [
{
"items": {},
"type": "array"
},
{},
{
"type": "null"
}
]
},
"run_name": {
"type": "string"
},
"max_concurrency": {
"anyOf": [
{
"type": "integer"
},
{
"type": "null"
}
]
},
"recursion_limit": {
"type": "integer"
},
"configurable": {
"type": "object"
},
"run_id": {
"anyOf": [
{
"format": "uuid",
"type": "string"
},
{
"type": "null"
}
]
}
},
"type": "object"
},
"kwargs": {
"default": None,
"type": "object"
}
},
"required": [
"state",
"config"
],
"type": "object"
}
}
}
]
}
response = completion(model="vertex_ai/gemini-2.0-flash", **args)
print(response)