diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index a937a35da25..7c97975a54a 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -39,6 +39,15 @@ ################################################# +def _get_tool_config_from_kwargs(kwargs: Dict[str, Any]) -> Optional[Dict[str, Any]]: + """Read toolConfig/tool_config without dropping intentionally empty dicts.""" + if "toolConfig" in kwargs: + return kwargs["toolConfig"] + if "tool_config" in kwargs: + return kwargs["tool_config"] + return None + + class GenerateContentSetupResult(BaseModel): """Internal Type - Result of setting up a generate content call""" @@ -171,12 +180,14 @@ def setup_generate_content_call( system_instruction = kwargs.get("systemInstruction") or kwargs.get( "system_instruction" ) + tool_config = _get_tool_config_from_kwargs(kwargs) request_body = ( generate_content_provider_config.transform_generate_content_request( model=model, contents=contents, tools=tools, generate_content_config_dict=generate_content_config_dict, + tool_config=tool_config, system_instruction=system_instruction, ) ) @@ -323,6 +334,7 @@ def generate_content( system_instruction = kwargs.get("systemInstruction") or kwargs.get( "system_instruction" ) + tool_config = _get_tool_config_from_kwargs(kwargs) # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -354,6 +366,7 @@ def generate_content( _is_async=_is_async, client=kwargs.get("client"), litellm_metadata=kwargs.get("litellm_metadata", {}), + tool_config=tool_config, system_instruction=system_instruction, ) @@ -414,6 +427,7 @@ async def agenerate_content_stream( system_instruction = kwargs.get("systemInstruction") or kwargs.get( "system_instruction" ) + tool_config = _get_tool_config_from_kwargs(kwargs) # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -452,6 +466,7 @@ async def agenerate_content_stream( client=kwargs.get("client"), stream=True, litellm_metadata=kwargs.get("litellm_metadata", {}), + tool_config=tool_config, system_instruction=system_instruction, ) @@ -520,6 +535,10 @@ def generate_content_stream( ) # Call the handler with streaming enabled (sync version) + system_instruction = kwargs.get("systemInstruction") or kwargs.get( + "system_instruction" + ) + tool_config = _get_tool_config_from_kwargs(kwargs) return base_llm_http_handler.generate_content_handler( model=setup_result.model, contents=contents, @@ -536,6 +555,8 @@ def generate_content_stream( client=kwargs.get("client"), stream=True, litellm_metadata=kwargs.get("litellm_metadata", {}), + tool_config=tool_config, + system_instruction=system_instruction, ) except Exception as e: diff --git a/litellm/llms/base_llm/google_genai/transformation.py b/litellm/llms/base_llm/google_genai/transformation.py index e8b3bf1a576..7952e2b0e10 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -152,6 +152,7 @@ def transform_generate_content_request( contents: GenerateContentContentListUnionDict, tools: Optional[ToolConfigDict], generate_content_config_dict: Dict, + tool_config: Optional[Dict[str, Any]] = None, system_instruction: Optional[Any] = None, ) -> dict: """ @@ -161,6 +162,7 @@ def transform_generate_content_request( model: The model name contents: Input contents tools: Tools + tool_config: Tool configuration generate_content_config_dict: Generation config parameters system_instruction: Optional system instruction diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4e6c3cba684..3e7a636640f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -9329,6 +9329,7 @@ def generate_content_handler( client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, + tool_config: Optional[Dict[str, Any]] = None, system_instruction: Optional[Any] = None, ) -> Any: """ @@ -9346,6 +9347,7 @@ def generate_content_handler( generate_content_provider_config=generate_content_provider_config, generate_content_config_dict=generate_content_config_dict, tools=tools, + tool_config=tool_config, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params, logging_obj=logging_obj, @@ -9384,6 +9386,7 @@ def generate_content_handler( model=model, contents=contents, tools=tools, + tool_config=tool_config, generate_content_config_dict=generate_content_config_dict, system_instruction=system_instruction, ) @@ -9456,6 +9459,7 @@ async def async_generate_content_handler( client: Optional[AsyncHTTPHandler] = None, stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, + tool_config: Optional[Dict[str, Any]] = None, system_instruction: Optional[Any] = None, ) -> Any: """ @@ -9493,6 +9497,7 @@ async def async_generate_content_handler( model=model, contents=contents, tools=tools, + tool_config=tool_config, generate_content_config_dict=generate_content_config_dict, system_instruction=system_instruction, ) diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index 7c4c7dba626..24f59b8072d 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -308,6 +308,7 @@ def transform_generate_content_request( contents: GenerateContentContentListUnionDict, tools: Optional[ToolConfigDict], generate_content_config_dict: Dict, + tool_config: Optional[Dict[str, Any]] = None, system_instruction: Optional[Any] = None, ) -> dict: from litellm.types.google_genai.main import ( @@ -326,6 +327,8 @@ def transform_generate_content_request( if system_instruction is not None: request_dict["systemInstruction"] = system_instruction + if tool_config is not None: + request_dict["toolConfig"] = tool_config return request_dict def transform_generate_content_response( diff --git a/litellm/llms/vertex_ai/google_genai/transformation.py b/litellm/llms/vertex_ai/google_genai/transformation.py index d7a4ceeb3e7..18836b164de 100644 --- a/litellm/llms/vertex_ai/google_genai/transformation.py +++ b/litellm/llms/vertex_ai/google_genai/transformation.py @@ -73,6 +73,7 @@ def transform_generate_content_request( contents: Any, tools: Optional[Any], generate_content_config_dict: Dict, + tool_config: Optional[Dict[str, Any]] = None, system_instruction: Optional[Any] = None, ) -> dict: """ @@ -89,8 +90,11 @@ def transform_generate_content_request( if tools: result["tools"] = tools + if tool_config is not None: + result["toolConfig"] = tool_config + # Add systemInstruction if provided - if system_instruction: + if system_instruction is not None: result["systemInstruction"] = system_instruction # Handle generationConfig - Vertex AI expects it in the same format diff --git a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py index 90c2cac18d0..92b4b9af496 100644 --- a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py +++ b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py @@ -174,6 +174,7 @@ async def test_google_gemini_httpx_request_direct(): ], "role": "user" }, + "toolConfig": {"functionCallingConfig": {"mode": "ANY"}}, "config": { # Note: already transformed from generationConfig "temperature": 0, "topP": 1, @@ -240,6 +241,7 @@ async def test_google_gemini_httpx_request_direct(): generate_content_provider_config=provider_config, generate_content_config_dict=sample_payload["config"], tools=None, + tool_config=sample_payload["toolConfig"], custom_llm_provider="gemini", litellm_params=litellm_params, logging_obj=logging_obj, @@ -265,6 +267,7 @@ async def test_google_gemini_httpx_request_direct(): request_data = call_kwargs.get('json') if request_data: assert 'contents' in request_data, "Expected 'contents' in request data" + assert request_data["toolConfig"] == sample_payload["toolConfig"] # The config should be included in the request as generationConfig if 'generationConfig' in request_data: diff --git a/tests/test_litellm/google_genai/test_google_genai_main.py b/tests/test_litellm/google_genai/test_google_genai_main.py index 5854e4b55af..5eb5c6a1177 100644 --- a/tests/test_litellm/google_genai/test_google_genai_main.py +++ b/tests/test_litellm/google_genai/test_google_genai_main.py @@ -1,24 +1,13 @@ #!/usr/bin/env python3 -""" -Test to verify the Google GenAI generate_content adapter functionality -""" -import json -import os -import sys - -import pytest - -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +"""Tests for Google GenAI main entrypoints.""" -import json import os import sys +from unittest.mock import AsyncMock, MagicMock, patch import pytest -import litellm +sys.path.insert(0, os.path.abspath("../../..")) @pytest.mark.asyncio @@ -26,8 +15,6 @@ async def test_agenerate_content_stream(): """ Test that the agenerate_content_stream function works """ - from unittest.mock import AsyncMock, patch - from litellm.google_genai.main import ( agenerate_content_stream, base_llm_http_handler, @@ -36,10 +23,40 @@ async def test_agenerate_content_stream(): with patch.object( base_llm_http_handler, "generate_content_handler", new=AsyncMock() ) as mock_post: - result = await agenerate_content_stream( + await agenerate_content_stream( model="gemini/gemini-2.0-flash-001", contents="Hello, world!", stream=True, ) mock_post.assert_called_once() - mock_post.call_args.kwargs["stream"] == True + assert mock_post.call_args.kwargs["stream"] is True + + +def test_generate_content_stream_forwards_system_instruction(): + """Test that generate_content_stream forwards systemInstruction and toolConfig.""" + from litellm.google_genai.main import ( + base_llm_http_handler, + generate_content_stream, + ) + + mock_response = MagicMock() + tool_config = {"functionCallingConfig": {"mode": "ANY"}} + + with patch.object( + base_llm_http_handler, "generate_content_handler", return_value=mock_response + ) as mock_post: + result = generate_content_stream( + model="gemini/gemini-2.0-flash-001", + contents="Hello, world!", + stream=True, + systemInstruction={"parts": [{"text": "You are helpful"}]}, + toolConfig=tool_config, + ) + + assert result is mock_response + mock_post.assert_called_once() + assert mock_post.call_args.kwargs["stream"] is True + assert mock_post.call_args.kwargs["tool_config"] == tool_config + assert mock_post.call_args.kwargs["system_instruction"] == { + "parts": [{"text": "You are helpful"}] + } diff --git a/tests/test_litellm/google_genai/test_google_genai_transformation.py b/tests/test_litellm/google_genai/test_google_genai_transformation.py index 8943d198dc1..f5f63db819c 100644 --- a/tests/test_litellm/google_genai/test_google_genai_transformation.py +++ b/tests/test_litellm/google_genai/test_google_genai_transformation.py @@ -12,6 +12,9 @@ import pytest from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig +from litellm.llms.vertex_ai.google_genai.transformation import ( + VertexAIGoogleGenAIConfig, +) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -173,6 +176,26 @@ def test_map_generate_content_optional_params_response_mime_type(): assert "responseJsonSchema" in result +@pytest.mark.parametrize( + "config_cls", + [GoogleGenAIConfig, VertexAIGoogleGenAIConfig], +) +def test_transform_generate_content_request_preserves_tool_config(config_cls): + config = config_cls() + tool_config = {"functionCallingConfig": {"mode": "ANY"}} + + result = config.transform_generate_content_request( + model="gemini-3-flash-preview", + contents=[{"role": "user", "parts": [{"text": "hello"}]}], + tools=[{"functionDeclarations": [{"name": "execute_command"}]}], + tool_config=tool_config, + generate_content_config_dict={"temperature": 1}, + system_instruction={"parts": [{"text": "system"}]}, + ) + + assert result["toolConfig"] == tool_config + + def test_responses_api_reasoning_dict_format(): """Test that reasoning parameter with dict format is mapped to reasoning_effort""" from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams @@ -274,6 +297,7 @@ def test_transform_generate_content_request_with_system_instruction(): model="gemini-3-flash-preview", contents=contents, tools=None, + tool_config=None, generate_content_config_dict=generate_content_config_dict, system_instruction=system_instruction, ) @@ -305,6 +329,7 @@ def test_transform_generate_content_request_without_system_instruction(): model="gemini-3-flash-preview", contents=contents, tools=None, + tool_config=None, generate_content_config_dict=generate_content_config_dict, system_instruction=None, ) @@ -356,6 +381,7 @@ def test_transform_generate_content_request_system_instruction_with_tools(): model="gemini-3-flash-preview", contents=contents, tools=tools, + tool_config=None, generate_content_config_dict=generate_content_config_dict, system_instruction=system_instruction, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 1aee1d49658..49e35ae2519 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1385,6 +1385,65 @@ def mock_upsert_side_effect( ) +def test_should_create_budget_with_none_values(): + """Test that should_create_budget returns False when all values are None.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + assert TeamMemberBudgetHandler.should_create_budget() is False + assert ( + TeamMemberBudgetHandler.should_create_budget( + team_member_budget=None, + team_member_rpm_limit=None, + team_member_tpm_limit=None, + team_member_budget_duration=None, + ) + is False + ) + + +def test_should_create_budget_with_zero_budget(): + """Test that should_create_budget returns True for explicit 0 budget. + + 0 is a valid explicit budget (zero dollars allowed). The UI bug was + sending 0 when the user didn't set a value — that's fixed in the UI + by using sanitizeNumeric. The backend correctly treats 0 as intentional. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + assert TeamMemberBudgetHandler.should_create_budget(team_member_budget=0) is True + assert ( + TeamMemberBudgetHandler.should_create_budget(team_member_budget=0.0) is True + ) + + +def test_should_create_budget_with_valid_values(): + """Test that should_create_budget returns True when any value is provided.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + TeamMemberBudgetHandler, + ) + + assert ( + TeamMemberBudgetHandler.should_create_budget(team_member_budget=100.0) is True + ) + assert ( + TeamMemberBudgetHandler.should_create_budget(team_member_rpm_limit=50) is True + ) + assert ( + TeamMemberBudgetHandler.should_create_budget(team_member_tpm_limit=1000) + is True + ) + assert ( + TeamMemberBudgetHandler.should_create_budget( + team_member_budget_duration="30d" + ) + is True + ) + + def test_clean_team_member_fields(): """ Test that _clean_team_member_fields removes all team member fields from a dictionary. diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index d2ce79580da..64bd3fafd69 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -497,7 +497,7 @@ const TeamInfoView: React.FC = ({ updateData.team_member_budget_duration = values.team_member_budget_duration; if (values.team_member_budget !== undefined) { - updateData.team_member_budget = Number(values.team_member_budget); + updateData.team_member_budget = sanitizeNumeric(values.team_member_budget); } if (values.team_member_key_duration !== undefined) {