mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
Merge 8a48c1d665 into 2dccc0dc79
This commit is contained in:
Kolade Fajimi
GitHub
commit
192396948b
4 changed files with 337 additions and 30 deletions
146
litellm/llms/vertex_ai/credentials_source.py
Normal file
146
litellm/llms/vertex_ai/credentials_source.py
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
"""
|
||||
Resolve the `vertex_credentials` config value into the JSON object google-auth needs.
|
||||
|
||||
The value is either the credentials JSON itself or a path to a file holding it. Which one
|
||||
it is decides what a failure means, so the two are told apart by the shape of the value and
|
||||
each failure is returned as its own case instead of collapsing into one parse error.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, NoReturn, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VertexCredentialsJson:
|
||||
value: Mapping[str, object]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VertexCredentialsFileUnreadable:
|
||||
path: str
|
||||
reason: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VertexCredentialsFileNotJson:
|
||||
path: str
|
||||
detail: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VertexCredentialsInlineNotJson:
|
||||
detail: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _NotAJsonObject:
|
||||
detail: str
|
||||
|
||||
|
||||
VertexCredentialsFailure: TypeAlias = (
|
||||
VertexCredentialsFileUnreadable | VertexCredentialsFileNotJson | VertexCredentialsInlineNotJson
|
||||
)
|
||||
VertexCredentialsSource: TypeAlias = VertexCredentialsJson | VertexCredentialsFailure
|
||||
_VertexCredentialsFile: TypeAlias = (
|
||||
VertexCredentialsJson | VertexCredentialsFileUnreadable | VertexCredentialsFileNotJson
|
||||
)
|
||||
|
||||
|
||||
_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _parse_json_object(raw: str) -> Mapping[str, object] | _NotAJsonObject:
|
||||
"""Parse *raw*, describing any failure without echoing it: the input can be key material."""
|
||||
try:
|
||||
return _JSON_OBJECT_ADAPTER.validate_json(raw)
|
||||
except ValidationError as e:
|
||||
# Only "msg" is reported. Pydantic keeps the offending value under "input", and that
|
||||
# value is the credential.
|
||||
return _NotAJsonObject("; ".join(detail["msg"] for detail in e.errors()))
|
||||
|
||||
|
||||
def _read_json_file(path: str) -> _VertexCredentialsFile:
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
contents: Final = f.read()
|
||||
except OSError as e:
|
||||
return VertexCredentialsFileUnreadable(path, f"{e.strerror or e} ({type(e).__name__})")
|
||||
except UnicodeDecodeError:
|
||||
return VertexCredentialsFileNotJson(path, "file is not UTF-8 text")
|
||||
except ValueError as e:
|
||||
# open() rejects a few paths before touching the filesystem, e.g. "embedded null byte".
|
||||
return VertexCredentialsFileUnreadable(path, f"{e} ({type(e).__name__})")
|
||||
|
||||
parsed: Final = _parse_json_object(contents)
|
||||
if isinstance(parsed, _NotAJsonObject):
|
||||
return VertexCredentialsFileNotJson(path, parsed.detail)
|
||||
return VertexCredentialsJson(parsed)
|
||||
|
||||
|
||||
def _is_inline_credentials_json(credentials: str) -> bool:
|
||||
"""Whether *credentials* carries the JSON itself rather than a path to a file holding it."""
|
||||
return credentials.lstrip().startswith("{")
|
||||
|
||||
|
||||
def load_vertex_credentials_source(credentials: str) -> VertexCredentialsSource:
|
||||
"""Read *credentials* as the credentials JSON when it is shaped like one, else as a path to it.
|
||||
|
||||
Telling the two apart by shape rather than by `os.path.exists()` is what keeps a read
|
||||
failure recognisable: `os.path.exists()` answers False for an unreadable path as well as
|
||||
an absent one, so both used to reach the inline branch and be reported as malformed JSON.
|
||||
"""
|
||||
inline_first: Final = _is_inline_credentials_json(credentials)
|
||||
verbose_logger.debug(
|
||||
"Vertex: Loading vertex credentials, is_file_path=%s, current dir %s", not inline_first, os.getcwd()
|
||||
)
|
||||
if not inline_first:
|
||||
return _read_json_file(credentials)
|
||||
|
||||
inline: Final = _parse_json_object(credentials)
|
||||
if not isinstance(inline, _NotAJsonObject):
|
||||
return VertexCredentialsJson(inline)
|
||||
|
||||
# A file can legitimately be named "{vertex}.json", so a brace-prefixed value that does
|
||||
# not parse is still given the file read it used to get before dispatch moved to shape.
|
||||
from_file: Final = _read_json_file(credentials)
|
||||
if isinstance(from_file, VertexCredentialsJson):
|
||||
return from_file
|
||||
return VertexCredentialsInlineNotJson(inline.detail)
|
||||
|
||||
|
||||
def raise_vertex_credentials_failure(failure: VertexCredentialsFailure) -> NoReturn:
|
||||
"""Map a load failure onto the ValueError the auth flow already surfaces to the caller.
|
||||
|
||||
The caller is told which of the three faults it was; the path goes to the proxy log
|
||||
instead, because the message reaches whoever sent the request and the operator who can
|
||||
act on the path is reading the log anyway. The path is redacted on the way there too, so
|
||||
a credential misconfigured into this field does not become a log entry.
|
||||
"""
|
||||
match failure:
|
||||
case VertexCredentialsFileUnreadable(path=path, reason=reason):
|
||||
verbose_logger.error("Vertex: cannot read the credentials file at %s: %s", redact_string(path), reason)
|
||||
raise ValueError(
|
||||
f"Unable to read the vertex credentials file: {reason}. The proxy log names the path. "
|
||||
"Set `vertex_credentials` to a readable file path, or to the credentials JSON itself."
|
||||
)
|
||||
case VertexCredentialsFileNotJson(path=path, detail=detail):
|
||||
verbose_logger.error("Vertex: credentials file at %s is not valid JSON: %s", redact_string(path), detail)
|
||||
raise ValueError(
|
||||
f"The vertex credentials file is not valid JSON: {detail}. The proxy log names the path. "
|
||||
"Check for unescaped newlines in private_key."
|
||||
)
|
||||
case VertexCredentialsInlineNotJson(detail=detail):
|
||||
raise ValueError(
|
||||
f"The inline `vertex_credentials` value is not valid JSON: {detail}. "
|
||||
"Check for unescaped newlines in private_key."
|
||||
)
|
||||
case _: # pragma: no cover - exhaustiveness guard, unreachable while the union holds
|
||||
assert_never(failure) # pragma: no cover
|
||||
|
|
@ -6,7 +6,6 @@ Handles Authentication and generating request urls for Vertex AI and Google AI S
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
|
||||
|
|
@ -26,6 +25,11 @@ from .common_utils import (
|
|||
get_vertex_base_model_name,
|
||||
get_vertex_base_url,
|
||||
)
|
||||
from .credentials_source import (
|
||||
VertexCredentialsJson,
|
||||
load_vertex_credentials_source,
|
||||
raise_vertex_credentials_failure,
|
||||
)
|
||||
|
||||
|
||||
def _graft_default_vertex_path(api_base: str, default_url: str) -> str:
|
||||
|
|
@ -127,27 +131,10 @@ class VertexBase:
|
|||
) -> tuple[_VertexCredentialsObject | None, str]:
|
||||
if credentials is not None:
|
||||
if isinstance(credentials, str):
|
||||
_is_path: Final = os.path.exists(
|
||||
credentials
|
||||
) # credentials is from server config (litellm_params), not user input
|
||||
verbose_logger.debug(
|
||||
"Vertex: Loading vertex credentials, is_file_path=%s, current dir %s",
|
||||
_is_path,
|
||||
os.getcwd(),
|
||||
)
|
||||
|
||||
try:
|
||||
if _is_path:
|
||||
with open(credentials) as f:
|
||||
json_obj = json.load(f)
|
||||
else:
|
||||
json_obj = json.loads(credentials)
|
||||
except Exception as e:
|
||||
raise Exception(
|
||||
"Unable to load vertex credentials from environment. "
|
||||
"Ensure the JSON is valid (check for unescaped newlines in private_key). "
|
||||
f"Parse error: {type(e).__name__}"
|
||||
)
|
||||
source: Final = load_vertex_credentials_source(credentials)
|
||||
if not isinstance(source, VertexCredentialsJson):
|
||||
raise_vertex_credentials_failure(source)
|
||||
json_obj: Mapping[str, object] = source.value
|
||||
elif isinstance(credentials, dict):
|
||||
json_obj = credentials
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -2193,3 +2194,162 @@ class TestVertexBase:
|
|||
|
||||
assert token == "cached-token"
|
||||
assert not mock_get_lock.called, "Fast path should not acquire lock"
|
||||
|
||||
|
||||
class TestVertexCredentialsSource:
|
||||
"""A `vertex_credentials` path that cannot be read must not be reported as malformed JSON.
|
||||
|
||||
Regression for https://github.com/BerriAI/litellm/issues/40166: dispatching on
|
||||
os.path.exists() sent a missing or unreadable path down the inline-JSON branch,
|
||||
where the path string itself was parsed and every failure came back as a
|
||||
JSONDecodeError blaming the key material.
|
||||
"""
|
||||
|
||||
def test_missing_credentials_file_names_the_read_failure(self, tmp_path, caplog):
|
||||
missing = tmp_path / "vertexai.json"
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
with pytest.raises(ValueError, match="No such file or directory") as exc_info:
|
||||
VertexBase().load_auth(credentials=str(missing), project_id="p")
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "not valid JSON" not in message
|
||||
assert str(missing) not in message
|
||||
assert str(missing) in caplog.text
|
||||
|
||||
def test_unreadable_credentials_file_names_the_read_failure(self, tmp_path, caplog):
|
||||
# Present but not openable as a file, the same shape as a path on a mount
|
||||
# that has stopped serving reads.
|
||||
unreadable = tmp_path / "vertexai.json"
|
||||
unreadable.mkdir()
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
with pytest.raises(ValueError, match="Unable to read the vertex credentials file") as exc_info:
|
||||
VertexBase().load_auth(credentials=str(unreadable), project_id="p")
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "not valid JSON" not in message
|
||||
assert str(unreadable) not in message
|
||||
assert str(unreadable) in caplog.text
|
||||
|
||||
def test_malformed_credentials_file_keeps_the_private_key_hint(self, tmp_path, caplog):
|
||||
malformed = tmp_path / "vertexai.json"
|
||||
malformed.write_text('{"type": "service_account", "private_key": "-----BEGIN\nPRIVATE KEY-----"}')
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
with pytest.raises(ValueError, match="is not valid JSON") as exc_info:
|
||||
VertexBase().load_auth(credentials=str(malformed), project_id="p")
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "private_key" in message
|
||||
assert str(malformed) not in message
|
||||
assert str(malformed) in caplog.text
|
||||
|
||||
def test_malformed_inline_credentials_do_not_echo_the_credential(self):
|
||||
inline = (
|
||||
'{"type": "service_account", "private_key": '
|
||||
'"-----BEGIN PRIVATE KEY-----\nMIIEvQIBADA\n-----END PRIVATE KEY-----"}'
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="The inline `vertex_credentials` value is not valid JSON"
|
||||
) as exc_info:
|
||||
VertexBase().load_auth(credentials=inline, project_id="p")
|
||||
|
||||
assert "MIIEvQIBADA" not in str(exc_info.value)
|
||||
|
||||
def test_a_credential_misconfigured_as_a_path_is_not_echoed(self):
|
||||
"""A value that is neither a path nor JSON must not come back in the message."""
|
||||
pem = "-----BEGIN PRIVATE KEY-----\nMIIEvQIBADA\n-----END PRIVATE KEY-----"
|
||||
|
||||
with pytest.raises(ValueError, match="Unable to read the vertex credentials file") as exc_info:
|
||||
VertexBase().load_auth(credentials=pem, project_id="p")
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "MIIEvQIBADA" not in message
|
||||
assert "BEGIN PRIVATE KEY" not in message
|
||||
|
||||
def test_readable_credentials_file_is_still_loaded(self, tmp_path):
|
||||
creds_file = tmp_path / "vertexai.json"
|
||||
creds_file.write_text(json.dumps({"type": "service_account", "project_id": "from-file"}))
|
||||
vertex_base = VertexBase()
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.project_id = "from-file"
|
||||
|
||||
with (
|
||||
patch.object(vertex_base, "_credentials_from_service_account", return_value=mock_creds) as from_sa,
|
||||
patch.object(vertex_base, "refresh_auth"),
|
||||
):
|
||||
creds, project_id = vertex_base.load_auth(credentials=str(creds_file), project_id=None)
|
||||
|
||||
assert creds is mock_creds
|
||||
assert project_id == "from-file"
|
||||
assert from_sa.call_args.args[0] == {"type": "service_account", "project_id": "from-file"}
|
||||
|
||||
def test_inline_credentials_json_is_still_parsed(self):
|
||||
vertex_base = VertexBase()
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.project_id = "from-inline"
|
||||
|
||||
with (
|
||||
patch.object(vertex_base, "_credentials_from_service_account", return_value=mock_creds) as from_sa,
|
||||
patch.object(vertex_base, "refresh_auth"),
|
||||
):
|
||||
creds, project_id = vertex_base.load_auth(
|
||||
credentials=json.dumps({"type": "service_account", "project_id": "from-inline"}),
|
||||
project_id=None,
|
||||
)
|
||||
|
||||
assert creds is mock_creds
|
||||
assert project_id == "from-inline"
|
||||
assert from_sa.call_args.args[0] == {"type": "service_account", "project_id": "from-inline"}
|
||||
|
||||
def test_a_credential_misconfigured_as_a_path_is_not_logged_either(self, caplog):
|
||||
"""The path reaches the log, so a credential put in that field must be scrubbed first."""
|
||||
pem = "-----BEGIN PRIVATE KEY-----\nMIIEvQIBADA\n-----END PRIVATE KEY-----"
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM"):
|
||||
with pytest.raises(ValueError, match="Unable to read the vertex credentials file"):
|
||||
VertexBase().load_auth(credentials=pem, project_id="p")
|
||||
|
||||
assert "MIIEvQIBADA" not in caplog.text
|
||||
assert "REDACTED" in caplog.text
|
||||
|
||||
def test_a_path_open_rejects_outright_is_reported_not_raised_raw(self):
|
||||
"""open() rejects some paths with a bare ValueError before any filesystem call."""
|
||||
with pytest.raises(ValueError, match="Unable to read the vertex credentials file") as exc_info:
|
||||
VertexBase().load_auth(credentials="creds" + chr(0) + ".json", project_id="p")
|
||||
|
||||
assert "embedded null" in str(exc_info.value)
|
||||
|
||||
def test_a_file_whose_name_starts_with_a_brace_is_still_read(self, tmp_path, monkeypatch):
|
||||
"""Shape dispatch must not strand a real file that happens to be named like JSON.
|
||||
|
||||
The value has to be relative for this to bite: an absolute path never starts with
|
||||
a brace, so only "{vertex}.json" reaches the inline branch and needs the fallback.
|
||||
"""
|
||||
braced = tmp_path / "{vertex}.json"
|
||||
braced.write_text(json.dumps({"type": "service_account", "project_id": "from-braced-file"}))
|
||||
monkeypatch.chdir(tmp_path)
|
||||
vertex_base = VertexBase()
|
||||
mock_creds = MagicMock()
|
||||
mock_creds.project_id = "from-braced-file"
|
||||
|
||||
with (
|
||||
patch.object(vertex_base, "_credentials_from_service_account", return_value=mock_creds) as from_sa,
|
||||
patch.object(vertex_base, "refresh_auth"),
|
||||
):
|
||||
creds, project_id = vertex_base.load_auth(credentials="{vertex}.json", project_id=None)
|
||||
|
||||
assert creds is mock_creds
|
||||
assert project_id == "from-braced-file"
|
||||
assert from_sa.call_args.args[0] == {"type": "service_account", "project_id": "from-braced-file"}
|
||||
|
||||
def test_a_credentials_file_that_is_not_text_is_reported_as_such(self, tmp_path):
|
||||
not_text = tmp_path / "vertexai.json"
|
||||
not_text.write_bytes(bytes([0xFF, 0xFE, 0x00, 0x01]))
|
||||
|
||||
with pytest.raises(ValueError, match="is not valid JSON") as exc_info:
|
||||
VertexBase().load_auth(credentials=str(not_text), project_id="p")
|
||||
|
||||
assert "UTF-8" in str(exc_info.value)
|
||||
|
|
|
|||
|
|
@ -467,15 +467,29 @@ def test_service_account_blob_fully_redacted():
|
|||
|
||||
|
||||
def test_vertex_error_message_no_credential_leak():
|
||||
"""The old Vertex error format leaked the full credential JSON.
|
||||
The new format must not contain any credential material."""
|
||||
new_msg = (
|
||||
"Unable to load vertex credentials from environment. "
|
||||
"Ensure the JSON is valid (check for unescaped newlines in private_key). "
|
||||
"Parse error: JSONDecodeError"
|
||||
"""The old Vertex error format leaked the full credential JSON. Every message the
|
||||
credential loader raises today must survive redaction unchanged, which it only can
|
||||
if it never carried credential material in the first place."""
|
||||
from litellm.llms.vertex_ai.credentials_source import (
|
||||
VertexCredentialsFileNotJson,
|
||||
VertexCredentialsFileUnreadable,
|
||||
VertexCredentialsInlineNotJson,
|
||||
raise_vertex_credentials_failure,
|
||||
)
|
||||
result = _redact_string(new_msg)
|
||||
assert result == new_msg # nothing to redact
|
||||
|
||||
path = "/etc/litellm/vertexai.json"
|
||||
failures = (
|
||||
VertexCredentialsFileUnreadable(path, "No such file or directory (FileNotFoundError)"),
|
||||
VertexCredentialsFileNotJson(path, "Invalid control character at: line 1 column 55"),
|
||||
VertexCredentialsInlineNotJson("Expecting value: line 1 column 1 (char 0)"),
|
||||
)
|
||||
|
||||
for failure in failures:
|
||||
with pytest.raises(ValueError, match="vertex") as exc_info:
|
||||
raise_vertex_credentials_failure(failure)
|
||||
message = str(exc_info.value)
|
||||
assert path not in message # the path goes to the log, not to the API caller
|
||||
assert _redact_string(message) == message # nothing to redact
|
||||
|
||||
|
||||
def test_vertex_traceback_redacts_pem():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue