fix(router.py): support reusable credentials via passthrough router (#9758)

* fix(router.py): support reusable credentials via passthrough router

enables reusable vertex credentials to be used in passthrough

* test: fix test

* test(test_router_adding_deployments.py): add unit testing
This commit is contained in:
Krish Dholakia 2025-04-04 18:40:14 -07:00 committed by GitHub
parent 8d76da03fe
commit c555c15ad7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 75 additions and 19 deletions

View file

@ -10,6 +10,7 @@ class CredentialAccessor:
@staticmethod
def get_credential_values(credential_name: str) -> dict:
"""Safe accessor for credentials."""
if not litellm.credential_list:
return {}
for credential in litellm.credential_list:

View file

@ -54,6 +54,7 @@ from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
@ -4506,25 +4507,53 @@ class Router:
passthrough_endpoint_router,
)
if deployment.litellm_params.litellm_credential_name is not None:
credential_values = CredentialAccessor.get_credential_values(
deployment.litellm_params.litellm_credential_name
)
else:
credential_values = {}
if custom_llm_provider == "vertex_ai":
vertex_project = (
credential_values.get("vertex_project")
or deployment.litellm_params.vertex_project
)
vertex_location = (
credential_values.get("vertex_location")
or deployment.litellm_params.vertex_location
)
vertex_credentials = (
credential_values.get("vertex_credentials")
or deployment.litellm_params.vertex_credentials
)
if (
deployment.litellm_params.vertex_project is None
or deployment.litellm_params.vertex_location is None
or deployment.litellm_params.vertex_credentials is None
vertex_project is None
or vertex_location is None
or vertex_credentials is None
):
raise ValueError(
"vertex_project, vertex_location, and vertex_credentials must be set in litellm_params for pass-through endpoints"
)
passthrough_endpoint_router.add_vertex_credentials(
project_id=deployment.litellm_params.vertex_project,
location=deployment.litellm_params.vertex_location,
vertex_credentials=deployment.litellm_params.vertex_credentials,
project_id=vertex_project,
location=vertex_location,
vertex_credentials=vertex_credentials,
)
else:
api_base = (
credential_values.get("api_base")
or deployment.litellm_params.api_base
)
api_key = (
credential_values.get("api_key")
or deployment.litellm_params.api_key
)
passthrough_endpoint_router.set_pass_through_credentials(
custom_llm_provider=custom_llm_provider,
api_base=deployment.litellm_params.api_base,
api_key=deployment.litellm_params.api_key,
api_base=api_base,
api_key=api_key,
)
pass
pass

View file

@ -179,6 +179,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams):
max_retries: Optional[int] = None
organization: Optional[str] = None # for openai orgs
configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None
litellm_credential_name: Optional[str] = None
## LOGGING PARAMS ##
litellm_trace_id: Optional[str] = None

View file

@ -9,23 +9,48 @@ from litellm.router import Deployment, LiteLLM_Params
from unittest.mock import patch
import json
def test_initialize_deployment_for_pass_through_success():
@pytest.mark.parametrize("reusable_credentials", [True, False])
def test_initialize_deployment_for_pass_through_success(reusable_credentials):
"""
Test successful initialization of a Vertex AI pass-through deployment
"""
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.types.utils import CredentialItem
vertex_project="test-project"
vertex_location="us-central1"
vertex_credentials=json.dumps({"type": "service_account", "project_id": "test"})
if not reusable_credentials:
litellm_params = LiteLLM_Params(
model="vertex_ai/test-model",
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_credentials=vertex_credentials,
use_in_pass_through=True,
)
else:
# add credentials to the credential accessor
CredentialAccessor.upsert_credentials([
CredentialItem(
credential_name="vertex_credentials",
credential_values={
"vertex_project": vertex_project,
"vertex_location": vertex_location,
"vertex_credentials": vertex_credentials,
},
credential_info={}
)
])
litellm_params = LiteLLM_Params(
model="vertex_ai/test-model",
litellm_credential_name="vertex_credentials",
use_in_pass_through=True,
)
router = Router(model_list=[])
deployment = Deployment(
model_name="vertex-test",
litellm_params=LiteLLM_Params(
model="vertex_ai/test-model",
vertex_project="test-project",
vertex_location="us-central1",
vertex_credentials=json.dumps(
{"type": "service_account", "project_id": "test"}
),
use_in_pass_through=True,
),
litellm_params=litellm_params,
)
# Test the initialization