scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(vertex_ai): normalize gemini model resource path

Browse source
This commit is contained in:
basil-ali-khan 2026-07-26 15:53:34 +05:00
parent 24123269cc
commit b3da65ce6f
6 changed files with 83 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/llms/vertex_ai/common_utils.py
View file

@ -463,7 +463,7 @@ def _get_gemini_url(
VertexGeminiConfig,
)
_gemini_model_name = "models/{}".format(model)
_gemini_model_name = normalize_gemini_model_path(model)
api_version = "v1alpha" if VertexGeminiConfig._is_gemini_3_or_newer(model) else "v1beta"
if mode == "chat":
@ -494,6 +494,13 @@ def _get_gemini_url(
return url, endpoint
def normalize_gemini_model_path(model: str) -> str:
"""Return a Gemini resource path with exactly one leading ``models/`` segment."""
if model.startswith("models/"):
return model
return f"models/{model}"
def _check_text_in_content(parts: List[PartType]) -> bool:
"""
check that user_content has 'text' parameter.

3
litellm/llms/vertex_ai/vertex_llm_base.py
View file

@ -24,6 +24,7 @@ from .common_utils import (
all_gemini_url_modes,
get_vertex_base_model_name,
get_vertex_base_url,
normalize_gemini_model_path,
)
GOOGLE_IMPORT_ERROR_MESSAGE = (
@ -637,7 +638,7 @@ class VertexBase:
# For Gemini (Google AI Studio), construct the full path like other providers
if model is None:
raise ValueError("Model parameter is required for Gemini custom API base URLs")
url = "{}/models/{}:{}".format(api_base, model, endpoint)
url = "{}/{}:{}".format(api_base, normalize_gemini_model_path(model), endpoint)
if gemini_api_key is None:
raise ValueError(
"Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."

73
litellm/proxy/management_endpoints/credential_migration.py
View file

@ -130,9 +130,7 @@ def classify_value(value: object, key: str = "scan") -> ValueClass:
return "plaintext"
if value.startswith(_V2_GCM_PREFIX):
return "migrated"
decrypted = decrypt_value_helper(
value=value, key=key, exception_type="debug", return_original_value=False
)
decrypted = decrypt_value_helper(value=value, key=key, exception_type="debug", return_original_value=False)
if decrypted is None:
# Did not decrypt under nacl and has no v2 marker: legacy plaintext.
return "plaintext"
@ -151,9 +149,7 @@ def reencrypt_value(value: object, key: str = "migrate") -> object:
return value
if value.startswith(_V2_GCM_PREFIX):
return value # idempotent: already migrated
decrypted = decrypt_value_helper(
value=value, key=key, exception_type="debug", return_original_value=False
)
decrypted = decrypt_value_helper(value=value, key=key, exception_type="debug", return_original_value=False)
if decrypted is None:
# Either legacy plaintext (no ciphertext to migrate) or corrupt. Either
# way, do not overwrite — preserve the value as stored.
@ -161,9 +157,7 @@ def reencrypt_value(value: object, key: str = "migrate") -> object:
return encrypt_value_helper(decrypted)
def reencrypt_selective_dict(
data: dict[str, object], sensitive_keys: list[str]
) -> dict[str, object]:
def reencrypt_selective_dict(data: dict[str, object], sensitive_keys: list[str]) -> dict[str, object]:
"""Return a copy of ``data`` with only ``sensitive_keys`` re-encrypted.
Non-sensitive fields (e.g. ``base_url``, ``connection_id``) are left as-is.
@ -212,9 +206,7 @@ async def _migrate_config_settings_row(
dict with selected sensitive fields (vantage_settings / cloudzero_settings).
"""
report = LocationReport(location=param_name)
record = await prisma_client.db.litellm_config.find_unique(
where={"param_name": param_name}
)
record = await prisma_client.db.litellm_config.find_unique(where={"param_name": param_name})
if record is None or record.param_value is None:
return report
@ -266,9 +258,7 @@ async def _migrate_sso_config(prisma_client: object, dry_run: bool) -> LocationR
every present string field.
"""
report = LocationReport(location="sso_config")
record = await prisma_client.db.litellm_ssoconfig.find_unique(
where={"id": "sso_config"}
)
record = await prisma_client.db.litellm_ssoconfig.find_unique(where={"id": "sso_config"})
if record is None or record.sso_settings is None:
return report
@ -344,9 +334,7 @@ async def _migrate_callback_vars_table(
rows = await table.find_many()
for row in rows or []:
metadata = getattr(row, "metadata", None)
if not isinstance(metadata, dict) or (
"logging" not in metadata and "callback_settings" not in metadata
):
if not isinstance(metadata, dict) or ("logging" not in metadata and "callback_settings" not in metadata):
continue
# Classify every callback-var value directly (strip the litellm_enc::
@ -534,9 +522,7 @@ async def _scan_config_env_vars(prisma_client: object) -> LocationReport:
"""Scan the ``environment_variables`` config row (``param_value`` dict)."""
report = LocationReport(location="config_environment_variables")
try:
record = await prisma_client.db.litellm_config.find_unique(
where={"param_name": "environment_variables"}
)
record = await prisma_client.db.litellm_config.find_unique(where={"param_name": "environment_variables"})
except Exception as e: # pragma: no cover - defensive
verbose_proxy_logger.debug("scan: config env vars unavailable: %s", str(e))
return report
@ -557,11 +543,7 @@ async def _scan_covered_tables(prisma_client: object) -> list[LocationReport]:
"""Read-only classification of every rotation-covered table. No writes."""
reports: list[LocationReport] = []
for location, db_attr, json_cols, scalar_cols in _COVERED_TABLE_SPECS:
reports.append(
await _scan_one_table(
prisma_client, location, db_attr, json_cols, scalar_cols
)
)
reports.append(await _scan_one_table(prisma_client, location, db_attr, json_cols, scalar_cols))
reports.append(await _scan_config_env_vars(prisma_client))
return reports
@ -575,9 +557,7 @@ _VANTAGE_SENSITIVE = ["api_key", "integration_token"]
_CLOUDZERO_SENSITIVE = ["api_key"]
async def _migrate_covered_tables(
prisma_client: object, user_api_key_dict: object
) -> list[LocationReport]:
async def _migrate_covered_tables(prisma_client: object, user_api_key_dict: object) -> list[LocationReport]:
"""Re-encrypt the tables already covered by ``_rotate_master_key`` (model
table, credentials, MCP credential/env tables, config environment_variables)
by running that orchestrator in *same-key* mode. With the AES gate on, the
@ -597,8 +577,7 @@ async def _migrate_covered_tables(
current_key = _get_salt_key()
if current_key is None:
raise RuntimeError(
"Cannot migrate covered tables: no salt key / master key is set. "
"Set LITELLM_SALT_KEY before migrating."
"Cannot migrate covered tables: no salt key / master key is set. Set LITELLM_SALT_KEY before migrating."
)
await _rotate_master_key(
prisma_client=cast("PrismaClient", prisma_client),
@ -648,19 +627,9 @@ async def migrate_encryption(
# Net-new walkers (items 3, 4, 11, 12, 13).
report.add(await _migrate_callback_vars_table(prisma_client, "team", dry_run))
report.add(
await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run)
)
report.add(
await _migrate_config_settings_row(
prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run
)
)
report.add(
await _migrate_config_settings_row(
prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run
)
)
report.add(await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run))
report.add(await _migrate_config_settings_row(prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run))
report.add(await _migrate_config_settings_row(prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run))
report.add(await _migrate_sso_config(prisma_client, dry_run))
return report
@ -683,20 +652,10 @@ async def check_encryption(prisma_client: object) -> MigrationReport:
# Net-new walker locations, in dry-run (read-only) mode.
report.add(await _migrate_callback_vars_table(prisma_client, "team", dry_run=True))
report.add(await _migrate_callback_vars_table(prisma_client, "verification_token", dry_run=True))
report.add(await _migrate_config_settings_row(prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run=True))
report.add(
await _migrate_callback_vars_table(
prisma_client, "verification_token", dry_run=True
)
)
report.add(
await _migrate_config_settings_row(
prisma_client, "vantage_settings", _VANTAGE_SENSITIVE, dry_run=True
)
)
report.add(
await _migrate_config_settings_row(
prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run=True
)
await _migrate_config_settings_row(prisma_client, "cloudzero_settings", _CLOUDZERO_SENSITIVE, dry_run=True)
)
report.add(await _migrate_sso_config(prisma_client, dry_run=True))
return report

19
tests/test_litellm/google_genai/test_google_genai_adapter.py
View file

@ -1140,13 +1140,18 @@ def test_get_gemini_url_excludes_api_key():
"""
from litellm.llms.vertex_ai.common_utils import _get_gemini_url
for mode in ("chat", "embedding", "batch_embedding", "count_tokens"):
url, _ = _get_gemini_url(
mode=mode,
model="gemini-2.5-flash",
stream=False,
)
assert "key=" not in url, f"API key found in URL for mode={mode}: {url}"
for model in ("gemini-1.5-flash", "models/gemini-1.5-flash"):
for mode in ("chat", "embedding", "batch_embedding", "count_tokens"):
url, _ = _get_gemini_url(
mode=mode,
model=model,
stream=False,
)
assert "key=" not in url, f"API key found in URL for mode={mode}: {url}"
assert "/models/models/" not in url, f"Duplicate models path for mode={mode}: {url}"
url, _ = _get_gemini_url(mode="chat", model=model, stream=False)
assert url == "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent"
# Streaming chat should only have ?alt=sse
url, _ = _get_gemini_url(mode="chat", model="gemini-2.5-flash", stream=True)

40
tests/test_litellm/llms/vertex_ai/test_vertex.py
View file

@ -309,6 +309,46 @@ def test_function_calling_with_gemini():
]
def test_qualified_gemini_model_uses_single_models_path():
from unittest.mock import patch
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
mock_response = MagicMock()
mock_response.json.return_value = {
"candidates": [
{
"content": {
"parts": [{"text": "Hello from Gemini."}],
"role": "model",
},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": 1,
"candidatesTokenCount": 1,
"totalTokenCount": 2,
},
}
with patch.object(client, "post", return_value=mock_response) as mock_post:
response = litellm.completion(
model="gemini/models/gemini-1.5-flash",
messages=[{"role": "user", "content": "Hello"}],
api_key="test-api-key",
client=client,
)
assert response.choices[0].message.content == "Hello from Gemini."
assert mock_post.call_args.kwargs["url"] == (
"https://generativelanguage.googleapis.com/v1beta/"
"models/gemini-1.5-flash:generateContent"
)
assert "/models/models/" not in mock_post.call_args.kwargs["url"]
def test_multiple_function_call():
litellm.set_verbose = True
from litellm.llms.custom_httpx.http_handler import HTTPHandler

5
tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py
View file

@ -875,6 +875,11 @@ class TestVertexBase:
"streamGenerateContent",
"https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent",
),
(
"models/gemini-1.5-flash",
"generateContent",
"https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent",
),
]
for model, endpoint, expected_url in test_cases: