mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
add _resolve_vertex_model_from_router
This commit is contained in:
Ishaan Jaffer
parent
df5d5664e2
commit
599e140d78
1 changed files with 89 additions and 25 deletions
|
|
@ -16,6 +16,7 @@ from fastapi.responses import StreamingResponse
|
|||
from starlette.websockets import WebSocketState
|
||||
|
||||
import litellm
|
||||
from litellm import get_llm_provider
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS,
|
||||
|
|
@ -1052,6 +1053,79 @@ async def bedrock_proxy_route(
|
|||
return received_value
|
||||
|
||||
|
||||
def _resolve_vertex_model_from_router(
|
||||
model_id: str,
|
||||
llm_router: Optional[litellm.Router],
|
||||
encoded_endpoint: str,
|
||||
endpoint: str,
|
||||
vertex_project: Optional[str],
|
||||
vertex_location: Optional[str],
|
||||
) -> Tuple[str, str, Optional[str], Optional[str]]:
|
||||
"""
|
||||
Resolve Vertex AI model configuration from router.
|
||||
|
||||
Args:
|
||||
model_id: The model ID extracted from the URL (e.g., "gcp/google/gemini-2.5-flash")
|
||||
llm_router: The LiteLLM router instance
|
||||
encoded_endpoint: The encoded endpoint path
|
||||
endpoint: The original endpoint path
|
||||
vertex_project: Current vertex project (may be from URL)
|
||||
vertex_location: Current vertex location (may be from URL)
|
||||
|
||||
Returns:
|
||||
Tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location)
|
||||
with resolved values from router config
|
||||
"""
|
||||
if not llm_router:
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location
|
||||
|
||||
try:
|
||||
deployment = llm_router.get_available_deployment_for_pass_through(model=model_id)
|
||||
if not deployment:
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location
|
||||
|
||||
litellm_params = deployment.get("litellm_params", {})
|
||||
|
||||
# Always override with router config values (they take precedence over URL values)
|
||||
config_vertex_project = litellm_params.get("vertex_project")
|
||||
config_vertex_location = litellm_params.get("vertex_location")
|
||||
if config_vertex_project:
|
||||
vertex_project = config_vertex_project
|
||||
if config_vertex_location:
|
||||
vertex_location = config_vertex_location
|
||||
|
||||
# Get the actual Vertex AI model name by stripping the provider prefix
|
||||
# e.g., "vertex_ai/gemini-2.0-flash-exp" -> "gemini-2.0-flash-exp"
|
||||
model_from_config = litellm_params.get("model", "")
|
||||
if model_from_config:
|
||||
from litellm.utils import get_llm_provider
|
||||
|
||||
# get_llm_provider returns (model, custom_llm_provider, dynamic_api_key, api_base)
|
||||
# For "vertex_ai/gemini-2.0-flash-exp" it returns:
|
||||
# model="gemini-2.0-flash-exp", custom_llm_provider="vertex_ai"
|
||||
actual_model, custom_llm_provider, _, _ = get_llm_provider(model=model_from_config)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"get_llm_provider returned: actual_model={actual_model}, "
|
||||
f"custom_llm_provider={custom_llm_provider}, model_id={model_id}"
|
||||
)
|
||||
|
||||
if actual_model and model_id != actual_model:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Resolved router model '{model_id}' to '{actual_model}' "
|
||||
f"(provider={custom_llm_provider}) with project={vertex_project}, location={vertex_location}"
|
||||
)
|
||||
encoded_endpoint = encoded_endpoint.replace(model_id, actual_model)
|
||||
endpoint = endpoint.replace(model_id, actual_model)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Error resolving vertex model from router for model {model_id}: {e}"
|
||||
)
|
||||
|
||||
return encoded_endpoint, endpoint, vertex_project, vertex_location
|
||||
|
||||
|
||||
def _is_bedrock_agent_runtime_route(endpoint: str) -> bool:
|
||||
"""
|
||||
Return True, if the endpoint should be routed to the `bedrock-agent-runtime` endpoint.
|
||||
|
|
@ -1512,8 +1586,11 @@ async def _prepare_vertex_auth_headers(
|
|||
if router_credentials is not None:
|
||||
vertex_credentials_str = None
|
||||
elif vertex_credentials is not None:
|
||||
vertex_project = vertex_credentials.vertex_project
|
||||
vertex_location = vertex_credentials.vertex_location
|
||||
# Only override vertex_project and vertex_location if they're not already set from router config
|
||||
if vertex_project is None:
|
||||
vertex_project = vertex_credentials.vertex_project
|
||||
if vertex_location is None:
|
||||
vertex_location = vertex_credentials.vertex_location
|
||||
vertex_credentials_str = vertex_credentials.vertex_credentials
|
||||
else:
|
||||
raise ValueError("No vertex credentials found")
|
||||
|
|
@ -1583,6 +1660,7 @@ async def _base_vertex_proxy_route(
|
|||
get_vertex_model_id_from_url,
|
||||
get_vertex_project_id_from_url,
|
||||
)
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
encoded_endpoint = httpx.URL(endpoint).path
|
||||
verbose_proxy_logger.debug("requested endpoint %s", endpoint)
|
||||
|
|
@ -1613,31 +1691,17 @@ async def _base_vertex_proxy_route(
|
|||
# Check if model is in router config - always do this to resolve custom model names
|
||||
model_id = get_vertex_model_id_from_url(endpoint)
|
||||
if model_id:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
if llm_router:
|
||||
try:
|
||||
# Use the dedicated pass-through deployment selection method to automatically filter use_in_pass_through=True
|
||||
deployment = llm_router.get_available_deployment_for_pass_through(model=model_id)
|
||||
if deployment:
|
||||
litellm_params = deployment.get("litellm_params", {})
|
||||
if vertex_project is None:
|
||||
vertex_project = litellm_params.get("vertex_project")
|
||||
if vertex_location is None:
|
||||
vertex_location = litellm_params.get("vertex_location")
|
||||
|
||||
# Replace custom model name with actual Vertex AI model name in the endpoint
|
||||
# e.g., "gcp/google/gemini-3-pro" -> "gemini-3-pro"
|
||||
actual_model = litellm_params.get("model", "")
|
||||
if "/" in actual_model:
|
||||
actual_model = actual_model.split("/", 1)[1]
|
||||
if actual_model and model_id != actual_model:
|
||||
encoded_endpoint = encoded_endpoint.replace(model_id, actual_model)
|
||||
endpoint = endpoint.replace(model_id, actual_model)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Error getting available deployment for model {model_id}: {e}"
|
||||
)
|
||||
# Resolve model configuration from router
|
||||
encoded_endpoint, endpoint, vertex_project, vertex_location = _resolve_vertex_model_from_router(
|
||||
model_id=model_id,
|
||||
llm_router=llm_router,
|
||||
encoded_endpoint=encoded_endpoint,
|
||||
endpoint=endpoint,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
vertex_credentials = passthrough_endpoint_router.get_vertex_credentials(
|
||||
project_id=vertex_project,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue