mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
Merge pull request #17803 from BerriAI/litellm_preserve_system_instructions
fix: Preserve systemInstructions for vertex ai generate content request
This commit is contained in:
commit
1107febf05
5 changed files with 97 additions and 4 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
Loading…
Add table
Reference in a new issue