litellm-mirror/litellm/proxy/vertex_ai_endpoints/vertex_passthrough_router.py

import json
import re
from typing import Dict, Optional

from litellm._logging import verbose_proxy_logger
from litellm.proxy.vertex_ai_endpoints.vertex_endpoints import (
    VertexPassThroughCredentials,
)


class VertexPassThroughRouter:
    """
    Vertex Pass Through Router for Vertex AI pass-through endpoints


    - if request specifies a project-id, location -> use credentials corresponding to the project-id, location
    - if request does not specify a project-id, location -> use credentials corresponding to the DEFAULT_VERTEXAI_PROJECT, DEFAULT_VERTEXAI_LOCATION
    """

    def __init__(self):
        """
        Initialize the VertexPassThroughRouter
        Stores the vertex credentials for each deployment key
        ```
        {
            "project_id-location": VertexPassThroughCredentials,
            "adroit-crow-us-central1": VertexPassThroughCredentials,
        }
        ```
        """
        self.deployment_key_to_vertex_credentials: Dict[
            str, VertexPassThroughCredentials
        ] = {}
        pass

    def get_vertex_credentials(
        self, project_id: Optional[str], location: Optional[str]
    ) -> VertexPassThroughCredentials:
        """
        Get the vertex credentials for the given project-id, location
        """
        from litellm.proxy.vertex_ai_endpoints.vertex_endpoints import (
            default_vertex_config,
        )

        deployment_key = self._get_deployment_key(
            project_id=project_id,
            location=location,
        )
        if deployment_key is None:
            return default_vertex_config
        if deployment_key in self.deployment_key_to_vertex_credentials:
            return self.deployment_key_to_vertex_credentials[deployment_key]
        else:
            return default_vertex_config

    def add_vertex_credentials(
        self,
        project_id: str,
        location: str,
        vertex_credentials: str,
    ):
        """
        Add the vertex credentials for the given project-id, location
        """
        from litellm.proxy.vertex_ai_endpoints.vertex_endpoints import (
            _set_default_vertex_config,
        )

        deployment_key = self._get_deployment_key(
            project_id=project_id,
            location=location,
        )
        if deployment_key is None:
            verbose_proxy_logger.debug(
                "No deployment key found for project-id, location"
            )
            return
        vertex_pass_through_credentials = VertexPassThroughCredentials(
            vertex_project=project_id,
            vertex_location=location,
            vertex_credentials=vertex_credentials,
        )
        self.deployment_key_to_vertex_credentials[deployment_key] = (
            vertex_pass_through_credentials
        )
        verbose_proxy_logger.debug(
            f"self.deployment_key_to_vertex_credentials: {json.dumps(self.deployment_key_to_vertex_credentials, indent=4, default=str)}"
        )
        _set_default_vertex_config(vertex_pass_through_credentials)

    def _get_deployment_key(
        self, project_id: Optional[str], location: Optional[str]
    ) -> Optional[str]:
        """
        Get the deployment key for the given project-id, location
        """
        if project_id is None or location is None:
            return None
        return f"{project_id}-{location}"

    @staticmethod
    def _get_vertex_project_id_from_url(url: str) -> Optional[str]:
        """
        Get the vertex project id from the url

        `https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/publishers/google/models/${MODEL_ID}:streamGenerateContent`
        """
        match = re.search(r"/projects/([^/]+)", url)
        return match.group(1) if match else None

    @staticmethod
    def _get_vertex_location_from_url(url: str) -> Optional[str]:
        """
        Get the vertex location from the url

        `https://${LOCATION}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${LOCATION}/publishers/google/models/${MODEL_ID}:streamGenerateContent`
        """
        match = re.search(r"/locations/([^/]+)", url)
        return match.group(1) if match else None