fix(vertex_ai/common_utils.py): fix handling constructed url with default vertex config

This commit is contained in:
Krrish Dholakia 2025-03-22 11:32:01 -07:00
parent b44b3bd36b
commit 06e69a414e
4 changed files with 98 additions and 17 deletions

View file

@ -303,11 +303,26 @@ def get_vertex_location_from_url(url: str) -> Optional[str]:
return match.group(1) if match else None
def replace_project_and_location_in_route(
requested_route: str, vertex_project: str, vertex_location: str
) -> str:
"""
Replace project and location values in the route with the provided values
"""
# Replace project and location values while keeping route structure
modified_route = re.sub(
r"/projects/[^/]+/locations/[^/]+/",
f"/projects/{vertex_project}/locations/{vertex_location}/",
requested_route,
)
return modified_route
def construct_target_url(
base_url: str,
requested_route: str,
default_vertex_location: Optional[str],
default_vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_project: Optional[str],
) -> httpx.URL:
"""
Allow user to specify their own project id / location.
@ -321,8 +336,12 @@ def construct_target_url(
"""
new_base_url = httpx.URL(base_url)
if "locations" in requested_route: # contains the target project id + location
updated_url = new_base_url.copy_with(path=requested_route)
return updated_url
if vertex_project and vertex_location:
requested_route = replace_project_and_location_in_route(
requested_route, vertex_project, vertex_location
)
return new_base_url.copy_with(path=requested_route)
"""
- Add endpoint version (e.g. v1beta for cachedContent, v1 for rest)
- Add default project id
@ -333,7 +352,7 @@ def construct_target_url(
vertex_version = "v1beta1"
base_requested_route = "{}/projects/{}/locations/{}".format(
vertex_version, default_vertex_project, default_vertex_location
vertex_version, vertex_project, vertex_location
)
updated_requested_route = "/" + base_requested_route + requested_route