vertex add vertex endpoints

This commit is contained in:
Ishaan Jaff 2024-08-29 15:49:25 -07:00
parent 556f34937d
commit 6491e461d6

View file

@ -84,15 +84,33 @@ async def vertex_proxy_route(
):
encoded_endpoint = httpx.URL(endpoint).path
import re
from litellm.fine_tuning.main import vertex_fine_tuning_apis_instance
verbose_proxy_logger.debug("requested endpoint %s", endpoint)
headers: dict = {}
# Use headers from the incoming request if default_vertex_config is not set
if default_vertex_config is None:
raise ValueError(
"Vertex credentials not added on litellm proxy, please add `default_vertex_config` on your config.yaml"
headers = dict(request.headers) or {}
verbose_proxy_logger.debug(
"default_vertex_config not set, incoming request headers %s", headers
)
vertex_project = default_vertex_config.get("vertex_project", None)
vertex_location = default_vertex_config.get("vertex_location", None)
vertex_credentials = default_vertex_config.get("vertex_credentials", None)
# extract location from endpoint, endpoint
# "v1beta1/projects/adroit-crow-413218/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent"
match = re.search(r"/locations/([^/]+)", endpoint)
vertex_location = match.group(1) if match else None
base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"
headers.pop("content-length", None)
_new_headers = {
"Authorization": headers.get("authorization"),
}
headers = _new_headers
else:
vertex_project = default_vertex_config.get("vertex_project")
vertex_location = default_vertex_config.get("vertex_location")
vertex_credentials = default_vertex_config.get("vertex_credentials")
base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"
auth_header, _ = vertex_fine_tuning_apis_instance._get_token_and_url(