fix: correct image size limit enforcement and vertex_location None passthrough

token_counter.py: the previous size-limit raises were inside except Exception: pass,
so they were silently swallowed. The post-read raise was worse — img_data was already
assigned the full body before the raise, so the oversized value was used downstream.
Restructured to only assign img_data when the body is within bounds.

vertex_ai/common_utils.py and llm_passthrough_endpoints.py: the is-not-None guard
skipped validation for None, falling through to produce "https://None-aiplatform..."
Added explicit None check that raises before the regex guard.
This commit is contained in:
Yuneng Jiang 2026-04-23 11:20:20 -07:00
parent 7ebe7cc99b
commit 994e35135d
No known key found for this signature in database
3 changed files with 11 additions and 10 deletions

View file

@ -219,10 +219,11 @@ def get_image_dimensions(
max_bytes = int(MAX_IMAGE_URL_DOWNLOAD_SIZE_MB * 1024 * 1024)
content_length = response.headers.get("Content-Length")
if content_length is not None and int(content_length) > max_bytes:
raise ValueError("Image response exceeds size limit")
img_data = response.read()
if len(img_data) > max_bytes:
raise ValueError("Image response exceeds size limit")
pass # skip download; img_data stays None
else:
body = response.read()
if len(body) <= max_bytes:
img_data = body
except Exception:
pass
if img_data is None:

View file

@ -232,9 +232,9 @@ def get_vertex_base_url(
"""
if vertex_location == "global":
return "https://aiplatform.googleapis.com"
if vertex_location is not None and not re.match(
r"^[a-z][a-z0-9-]*$", vertex_location
):
if vertex_location is None:
raise ValueError("vertex_location is required")
if not re.match(r"^[a-z][a-z0-9-]*$", vertex_location):
raise ValueError("Invalid vertex_location format")
return f"https://{vertex_location}-aiplatform.googleapis.com"

View file

@ -1501,9 +1501,9 @@ def get_vertex_base_url(vertex_location: Optional[str]) -> str:
"""
if vertex_location == "global":
return "https://aiplatform.googleapis.com/"
if vertex_location is not None and not re.match(
r"^[a-z][a-z0-9-]*$", vertex_location
):
if vertex_location is None:
raise ValueError("vertex_location is required")
if not re.match(r"^[a-z][a-z0-9-]*$", vertex_location):
raise ValueError("Invalid vertex_location format")
return f"https://{vertex_location}-aiplatform.googleapis.com/"