Merge pull request #17803 from BerriAI/litellm_preserve_system_instructions

fix: Preserve systemInstructions for vertex ai generate content request
This commit is contained in:
Sameer Kankute 2025-12-11 10:31:08 +05:30 committed by GitHub
commit 1107febf05
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 97 additions and 4 deletions

View file

@ -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:

View file

@ -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

View file

@ -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:

View file

@ -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,

View file

@ -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
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"