From 9344d29a1523c0fce95dbb9976ff49d662c45685 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 11 Dec 2025 09:52:32 +0530 Subject: [PATCH] fix: Preserve systemInstructions for vertex ai generate content request --- litellm/google_genai/main.py | 11 +++ .../base_llm/google_genai/transformation.py | 6 +- litellm/llms/custom_httpx/llm_http_handler.py | 5 ++ .../gemini/google_genai/transformation.py | 1 + .../test_google_api_endpoints.py | 78 ++++++++++++++++++- 5 files changed, 97 insertions(+), 4 deletions(-) diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 8a9cb809404..b7523ef8c16 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -164,12 +164,15 @@ class GenerateContentHelper: model=model, ) ) + # Extract systemInstruction from kwargs to pass to transform + system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") request_body = ( generate_content_provider_config.transform_generate_content_request( model=model, contents=contents, tools=tools, generate_content_config_dict=generate_content_config_dict, + system_instruction=system_instruction, ) ) @@ -311,6 +314,9 @@ def generate_content( **kwargs, ) + # Extract systemInstruction from kwargs to pass to handler + system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: # Use the adapter to convert to completion format @@ -340,6 +346,7 @@ def generate_content( _is_async=_is_async, client=kwargs.get("client"), litellm_metadata=kwargs.get("litellm_metadata", {}), + system_instruction=system_instruction, ) return response @@ -395,6 +402,9 @@ async def agenerate_content_stream( **kwargs, ) + # Extract systemInstruction from kwargs to pass to handler + system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: # Use the adapter to convert to completion format @@ -428,6 +438,7 @@ async def agenerate_content_stream( client=kwargs.get("client"), stream=True, litellm_metadata=kwargs.get("litellm_metadata", {}), + 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 6dbccaada9a..0a85e127bd7 100644 --- a/litellm/llms/base_llm/google_genai/transformation.py +++ b/litellm/llms/base_llm/google_genai/transformation.py @@ -149,6 +149,7 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): contents: GenerateContentContentListUnionDict, tools: Optional[ToolConfigDict], generate_content_config_dict: Dict, + system_instruction: Optional[Any] = None, ) -> dict: """ Transform the request parameters for the generate content API. @@ -157,9 +158,8 @@ class BaseGoogleGenAIGenerateContentConfig(ABC): model: The model name contents: Input contents tools: Tools - generate_content_request_params: Request parameters - litellm_params: LiteLLM parameters - headers: Request headers + generate_content_config_dict: Generation config parameters + system_instruction: Optional system instruction Returns: Transformed request data diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 72fe0ac7ec1..381d94f0186 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7311,6 +7311,7 @@ class BaseLLMHTTPHandler: client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None, stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, + system_instruction: Optional[Any] = None, ) -> Any: """ Handles Google GenAI generate content requests. @@ -7336,6 +7337,7 @@ class BaseLLMHTTPHandler: client=client if isinstance(client, AsyncHTTPHandler) else None, stream=stream, litellm_metadata=litellm_metadata, + system_instruction=system_instruction, ) if client is None or not isinstance(client, HTTPHandler): @@ -7365,6 +7367,7 @@ class BaseLLMHTTPHandler: contents=contents, tools=tools, generate_content_config_dict=generate_content_config_dict, + system_instruction=system_instruction, ) if extra_body: @@ -7435,6 +7438,7 @@ class BaseLLMHTTPHandler: client: Optional[AsyncHTTPHandler] = None, stream: bool = False, litellm_metadata: Optional[Dict[str, Any]] = None, + system_instruction: Optional[Any] = None, ) -> Any: """ Async version of the generate content handler. @@ -7472,6 +7476,7 @@ class BaseLLMHTTPHandler: contents=contents, tools=tools, generate_content_config_dict=generate_content_config_dict, + system_instruction=system_instruction, ) if extra_body: diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index 2d585769029..bc32aca6554 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -272,6 +272,7 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): contents: GenerateContentContentListUnionDict, tools: Optional[ToolConfigDict], generate_content_config_dict: Dict, + system_instruction: Optional[Any] = None, ) -> dict: from litellm.types.google_genai.main import ( GenerateContentConfigDict, diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index 923cb2c6fb8..11e34cbbea4 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -233,4 +233,80 @@ def test_google_stream_generate_content_with_cost_tracking_metadata(): assert called_data["litellm_metadata"]["user_api_key_user_id"] == "test-user-id" assert called_data["litellm_metadata"]["user_api_key_team_id"] == "test-team-id" # Verify stream is set to True - assert called_data["stream"] is True \ No newline at end of file + assert called_data["stream"] is True + + +def test_google_generate_content_with_system_instruction(): + """ + Test that systemInstruction is correctly passed through from the endpoint to the router. + + This test verifies the fix for systemInstruction being dropped when forwarding + requests to Vertex AI through the Google GenAI endpoint. + """ + try: + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.google_endpoints.endpoints import router as google_router + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + # Create a FastAPI app and include the router + app = FastAPI() + app.include_router(google_router) + + # Create a test client + client = TestClient(app) + + # Mock all required proxy server dependencies + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \ + patch("litellm.proxy.proxy_server.general_settings", {}), \ + patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, \ + patch("litellm.proxy.proxy_server.version", "1.0.0"), \ + patch("litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request") as mock_add_data: + + mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) + + # Mock add_litellm_data_to_request to pass through data unchanged + async def mock_add_litellm_data(data, request, user_api_key_dict, proxy_config, general_settings, version): + return data + + mock_add_data.side_effect = mock_add_litellm_data + + # Define the systemInstruction to test + system_instruction = { + "parts": [{"text": "Your name is Doodle."}] + } + + # Send a request with systemInstruction + response = client.post( + "/v1beta/models/gemini-2.5-pro:generateContent", + json={ + "systemInstruction": system_instruction, + "contents": [ + { + "parts": [{"text": "What is your name?"}], + "role": "user" + } + ] + }, + headers={"Authorization": "Bearer sk-test-key"} + ) + + # Verify the response + assert response.status_code == 200 + + # Verify that agenerate_content was called + mock_router.agenerate_content.assert_called_once() + call_args = mock_router.agenerate_content.call_args + called_data = call_args[1] + + # Verify that systemInstruction is present in the call arguments + assert "systemInstruction" in called_data + assert called_data["systemInstruction"] == system_instruction + assert called_data["systemInstruction"]["parts"][0]["text"] == "Your name is Doodle." + + # Verify contents are also present + assert "contents" in called_data + assert len(called_data["contents"]) == 1 + assert called_data["contents"][0]["role"] == "user" \ No newline at end of file