mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(google_genai): propagate user from kwargs to logging obj in agenerate_content
The Google-native /v1beta/models/{model}:generateContent route silently
drops the `user` field from the spend log. The handler in
google_endpoints/endpoints.py reads `user` from the request body via
add_litellm_data_to_request, but setup_generate_content_call passes
kwargs to litellm_logging_obj.update_from_kwargs without including
`user`. As a result, `model_call_details["user"]` is None and the
spend log row shows an empty `user` field, blocking per-end-user
attribution on this route.
Pass `user=kwargs.get("user")` so the logging object reflects what the
client sent. OpenAI-compat /v1/chat/completions has always done this
via litellm.completion calling logging.update_environment_variables
with the resolved user value.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
850fe595ac
commit
b221afb702
2 changed files with 66 additions and 0 deletions
|
|
@ -188,6 +188,7 @@ class GenerateContentHelper:
|
|||
litellm_logging_obj.update_from_kwargs(
|
||||
kwargs=kwargs,
|
||||
model=model,
|
||||
user=kwargs.get("user"),
|
||||
optional_params=dict(generate_content_config_dict),
|
||||
litellm_params={
|
||||
"litellm_call_id": litellm_call_id,
|
||||
|
|
|
|||
|
|
@ -43,3 +43,68 @@ async def test_agenerate_content_stream():
|
|||
)
|
||||
mock_post.assert_called_once()
|
||||
mock_post.call_args.kwargs["stream"] == True
|
||||
|
||||
|
||||
def test_setup_generate_content_call_propagates_user_to_logging_obj():
|
||||
"""
|
||||
Regression test for Google-native /v1beta/.../:generateContent route
|
||||
dropping the `user` field from the spend log.
|
||||
|
||||
`setup_generate_content_call` must pass `user` from kwargs to
|
||||
`litellm_logging_obj.update_from_kwargs` so that `logging_obj.user` and
|
||||
`model_call_details["user"]` reflect what the client sent (via header or
|
||||
body) instead of defaulting to None / "".
|
||||
"""
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from litellm.google_genai.main import GenerateContentHelper
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
# Use a real Logging instance (required by Pydantic model validation),
|
||||
# but stub out update_from_kwargs so we can assert what it was called with.
|
||||
real_logging_obj = Logging(
|
||||
model="gemini-2.5-pro",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="agenerate_content",
|
||||
start_time=__import__("datetime").datetime.now(),
|
||||
litellm_call_id="call-id-123",
|
||||
function_id="",
|
||||
)
|
||||
mock_update = MagicMock()
|
||||
real_logging_obj.update_from_kwargs = mock_update # type: ignore[method-assign]
|
||||
|
||||
# Mock the provider config so we hit the update_from_kwargs path (not the
|
||||
# adapter early-return path). Use spec= on a real subclass so Pydantic's
|
||||
# is_instance_of check passes.
|
||||
from litellm.llms.gemini.google_genai.transformation import (
|
||||
GoogleGenAIConfig,
|
||||
)
|
||||
|
||||
mock_provider_config = MagicMock(spec=GoogleGenAIConfig())
|
||||
mock_provider_config.map_generate_content_optional_params.return_value = {}
|
||||
mock_provider_config.transform_generate_content_request.return_value = {}
|
||||
|
||||
with patch(
|
||||
"litellm.get_llm_provider",
|
||||
return_value=("gemini-2.5-pro", "gemini", None, None),
|
||||
), patch(
|
||||
"litellm.utils.ProviderConfigManager.get_provider_google_genai_generate_content_config",
|
||||
return_value=mock_provider_config,
|
||||
):
|
||||
GenerateContentHelper.setup_generate_content_call(
|
||||
model="gemini-2.5-pro",
|
||||
contents=[{"parts": [{"text": "hi"}], "role": "user"}],
|
||||
config={},
|
||||
custom_llm_provider="gemini",
|
||||
litellm_logging_obj=real_logging_obj,
|
||||
litellm_call_id="call-id-123",
|
||||
user="my-end-user-uuid-456",
|
||||
metadata={"tags": ["scan_id=abc"]},
|
||||
)
|
||||
|
||||
mock_update.assert_called_once()
|
||||
call_kwargs = mock_update.call_args.kwargs
|
||||
assert (
|
||||
call_kwargs.get("user") == "my-end-user-uuid-456"
|
||||
), f"Expected user to be propagated, got: {call_kwargs.get('user')!r}"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue