mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-24 18:24:20 +00:00
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:
parent
8d76da03fe
commit
c555c15ad7
4 changed files with 75 additions and 19 deletions
|
@ -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:
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue