Add new /vertex_ai/discovery route - enables calling AgentBuilder API routes (#10084)

* feat(llm_passthrough_endpoints.py): expose new `/vertex_ai/discovery/` endpoint

Allows calling vertex ai discovery endpoints via passthrough

 For agentbuilder api calls

* refactor(llm_passthrough_endpoints.py): use common _base_vertex_proxy_route

Prevents duplicate code

* feat(llm_passthrough_endpoints.py): add vertex endpoint specific passthrough handlers

litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py

@@ -578,26 +578,70 @@ async def azure_proxy_route(
     )
 
 
-@router.api_route(
-    "/vertex-ai/{endpoint:path}",
-    methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
-    tags=["Vertex AI Pass-through", "pass-through"],
-    include_in_schema=False,
-)
-@router.api_route(
-    "/vertex_ai/{endpoint:path}",
-    methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
-    tags=["Vertex AI Pass-through", "pass-through"],
-)
-async def vertex_proxy_route(
+from abc import ABC, abstractmethod
+
+
+class BaseVertexAIPassThroughHandler(ABC):
+    @staticmethod
+    @abstractmethod
+    def get_default_base_target_url(vertex_location: Optional[str]) -> str:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def update_base_target_url_with_credential_location(
+        base_target_url: str, vertex_location: Optional[str]
+    ) -> str:
+        pass
+
+
+class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler):
+    @staticmethod
+    def get_default_base_target_url(vertex_location: Optional[str]) -> str:
+        return "https://discoveryengine.googleapis.com/"
+
+    @staticmethod
+    def update_base_target_url_with_credential_location(
+        base_target_url: str, vertex_location: Optional[str]
+    ) -> str:
+        return base_target_url
+
+
+class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
+    @staticmethod
+    def get_default_base_target_url(vertex_location: Optional[str]) -> str:
+        return f"https://{vertex_location}-aiplatform.googleapis.com/"
+
+    @staticmethod
+    def update_base_target_url_with_credential_location(
+        base_target_url: str, vertex_location: Optional[str]
+    ) -> str:
+        return f"https://{vertex_location}-aiplatform.googleapis.com/"
+
+
+def get_vertex_pass_through_handler(
+    call_type: Literal["discovery", "aiplatform"]
+) -> BaseVertexAIPassThroughHandler:
+    if call_type == "discovery":
+        return VertexAIDiscoveryPassThroughHandler()
+    elif call_type == "aiplatform":
+        return VertexAIPassThroughHandler()
+    else:
+        raise ValueError(f"Invalid call type: {call_type}")
+
+
+async def _base_vertex_proxy_route(
     endpoint: str,
     request: Request,
     fastapi_response: Response,
+    get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler,
+    user_api_key_dict: Optional[UserAPIKeyAuth] = None,
 ):
     """
-    Call LiteLLM proxy via Vertex AI SDK.
+    Base function for Vertex AI passthrough routes.
+    Handles common logic for all Vertex AI services.
 
-    [Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai)
+    Default base_target_url is `https://{vertex_location}-aiplatform.googleapis.com/`
     """
     from litellm.llms.vertex_ai.common_utils import (
         construct_target_url,
@@ -613,6 +657,14 @@ async def vertex_proxy_route(
         request=request,
         api_key=api_key_to_use,
     )
+
+    if user_api_key_dict is None:
+        api_key_to_use = get_litellm_virtual_key(request=request)
+        user_api_key_dict = await user_api_key_auth(
+            request=request,
+            api_key=api_key_to_use,
+        )
+
     vertex_project: Optional[str] = get_vertex_project_id_from_url(endpoint)
     vertex_location: Optional[str] = get_vertex_location_from_url(endpoint)
     vertex_credentials = passthrough_endpoint_router.get_vertex_credentials(
@@ -620,6 +672,10 @@ async def vertex_proxy_route(
         location=vertex_location,
     )
 
+    base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(
+        vertex_location
+    )
+
     headers_passed_through = False
     # Use headers from the incoming request if no vertex credentials are found
     if vertex_credentials is None or vertex_credentials.vertex_project is None:
@@ -628,7 +684,6 @@ async def vertex_proxy_route(
         verbose_proxy_logger.debug(
             "default_vertex_config  not set, incoming request headers %s", headers
         )
-        base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"
         headers.pop("content-length", None)
         headers.pop("host", None)
     else:
@@ -636,9 +691,6 @@ async def vertex_proxy_route(
         vertex_location = vertex_credentials.vertex_location
         vertex_credentials_str = vertex_credentials.vertex_credentials
 
-        # Construct base URL for the target endpoint
-        base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"
-
         _auth_header, vertex_project = await vertex_llm_base._ensure_access_token_async(
             credentials=vertex_credentials_str,
             project_id=vertex_project,
@@ -661,6 +713,13 @@ async def vertex_proxy_route(
             "Authorization": f"Bearer {auth_header}",
         }
 
+        base_target_url = get_vertex_pass_through_handler.update_base_target_url_with_credential_location(
+            base_target_url, vertex_location
+        )
+
+    if base_target_url is None:
+        base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"
+
     request_route = encoded_endpoint
     verbose_proxy_logger.debug("request_route %s", request_route)
 
@@ -707,6 +766,66 @@ async def vertex_proxy_route(
     return received_value
 
 
+@router.api_route(
+    "/vertex_ai/discovery/{endpoint:path}",
+    methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
+    tags=["Vertex AI Pass-through", "pass-through"],
+)
+async def vertex_discovery_proxy_route(
+    endpoint: str,
+    request: Request,
+    fastapi_response: Response,
+):
+    """
+    Call any vertex discovery endpoint using the proxy.
+
+    Just use `{PROXY_BASE_URL}/vertex_ai/discovery/{endpoint:path}`
+
+    Target url: `https://discoveryengine.googleapis.com`
+    """
+
+    discovery_handler = get_vertex_pass_through_handler(call_type="discovery")
+    return await _base_vertex_proxy_route(
+        endpoint=endpoint,
+        request=request,
+        fastapi_response=fastapi_response,
+        get_vertex_pass_through_handler=discovery_handler,
+    )
+
+
+@router.api_route(
+    "/vertex-ai/{endpoint:path}",
+    methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
+    tags=["Vertex AI Pass-through", "pass-through"],
+    include_in_schema=False,
+)
+@router.api_route(
+    "/vertex_ai/{endpoint:path}",
+    methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
+    tags=["Vertex AI Pass-through", "pass-through"],
+)
+async def vertex_proxy_route(
+    endpoint: str,
+    request: Request,
+    fastapi_response: Response,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Call LiteLLM proxy via Vertex AI SDK.
+
+    [Docs](https://docs.litellm.ai/docs/pass_through/vertex_ai)
+    """
+    ai_platform_handler = get_vertex_pass_through_handler(call_type="aiplatform")
+
+    return await _base_vertex_proxy_route(
+        endpoint=endpoint,
+        request=request,
+        fastapi_response=fastapi_response,
+        get_vertex_pass_through_handler=ai_platform_handler,
+        user_api_key_dict=user_api_key_dict,
+    )
+
+
 @router.api_route(
     "/openai/{endpoint:path}",
     methods=["GET", "POST", "PUT", "DELETE", "PATCH"],