allow using known path for setting up pass throughs

This commit is contained in:
Ishaan Jaff 2024-11-25 23:39:37 -08:00
parent c61c429c44
commit 7dfc25f894
2 changed files with 57 additions and 7 deletions

View file

@ -28,14 +28,46 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
create_pass_through_route, create_pass_through_route,
) )
from litellm.secret_managers.main import get_secret_str
from litellm.types.passthrough_endpoints.vertex_ai import *
router = APIRouter() router = APIRouter()
default_vertex_config = None
default_vertex_config: VertexPassThroughCredentials = VertexPassThroughCredentials()
def set_default_vertex_config(config): def _get_vertex_env_vars() -> VertexPassThroughCredentials:
"""
Helper to get vertex pass through config from environment variables
The following environment variables are used:
- DEFAULT_VERTEXAI_PROJECT (project id)
- DEFAULT_VERTEXAI_LOCATION (location)
- DEFAULT_GOOGLE_APPLICATION_CREDENTIALS (path to credentials file)
"""
return VertexPassThroughCredentials(
vertex_project=get_secret_str("DEFAULT_VERTEXAI_PROJECT"),
vertex_location=get_secret_str("DEFAULT_VERTEXAI_LOCATION"),
vertex_credentials=get_secret_str("DEFAULT_GOOGLE_APPLICATION_CREDENTIALS"),
)
def set_default_vertex_config(config: Optional[dict]):
"""Sets vertex configuration from provided config and/or environment variables
Args:
config (Optional[dict]): Configuration dictionary
Example: {
"vertex_project": "my-project-123",
"vertex_location": "us-central1",
"vertex_credentials": "os.environ/GOOGLE_CREDS"
}
"""
global default_vertex_config global default_vertex_config
# Initialize config dictionary if None
if config is None: if config is None:
default_vertex_config = _get_vertex_env_vars()
return return
if not isinstance(config, dict): if not isinstance(config, dict):
@ -46,7 +78,7 @@ def set_default_vertex_config(config):
if isinstance(value, str) and value.startswith("os.environ/"): if isinstance(value, str) and value.startswith("os.environ/"):
config[key] = litellm.get_secret(value) config[key] = litellm.get_secret(value)
default_vertex_config = config default_vertex_config = VertexPassThroughCredentials(**config)
def exception_handler(e: Exception): def exception_handler(e: Exception):
@ -140,7 +172,7 @@ async def vertex_proxy_route(
vertex_project = None vertex_project = None
vertex_location = None vertex_location = None
# Use headers from the incoming request if default_vertex_config is not set # Use headers from the incoming request if default_vertex_config is not set
if default_vertex_config is None: if default_vertex_config.vertex_project is None:
headers = dict(request.headers) or {} headers = dict(request.headers) or {}
verbose_proxy_logger.debug( verbose_proxy_logger.debug(
"default_vertex_config not set, incoming request headers %s", headers "default_vertex_config not set, incoming request headers %s", headers
@ -153,9 +185,9 @@ async def vertex_proxy_route(
headers.pop("content-length", None) headers.pop("content-length", None)
headers.pop("host", None) headers.pop("host", None)
else: else:
vertex_project = default_vertex_config.get("vertex_project") vertex_project = default_vertex_config.vertex_project
vertex_location = default_vertex_config.get("vertex_location") vertex_location = default_vertex_config.vertex_location
vertex_credentials = default_vertex_config.get("vertex_credentials") vertex_credentials = default_vertex_config.vertex_credentials
base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/" base_target_url = f"https://{vertex_location}-aiplatform.googleapis.com/"

View file

@ -0,0 +1,18 @@
"""
Used for /vertex_ai/ pass through endpoints
"""
from typing import Optional
from pydantic import BaseModel
class VertexPassThroughCredentials(BaseModel):
# Example: vertex_project = "my-project-123"
vertex_project: Optional[str] = None
# Example: vertex_location = "us-central1"
vertex_location: Optional[str] = None
# Example: vertex_credentials = "/path/to/credentials.json" or "os.environ/GOOGLE_CREDS"
vertex_credentials: Optional[str] = None