mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 18a28f7457 into eb65b8f017
This commit is contained in:
devin-ai-integration[bot]
GitHub
commit
4ec4213b98
23 changed files with 1169 additions and 261 deletions
|
|
@ -3,9 +3,12 @@ import os
|
|||
from collections.abc import Mapping
|
||||
from typing import Final, Literal, cast
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
||||
# Versioned ciphertext marker for AES-256-GCM values.
|
||||
# Format: "v2:gcm:" + base64url(nonce(12) || ciphertext || tag(16)).
|
||||
|
|
@ -296,3 +299,99 @@ def decode_secret_map(value: object, *, key: str) -> Mapping[str, str] | None:
|
|||
return _SECRET_MAP.validate_json(decrypted, strict=True)
|
||||
except ValidationError:
|
||||
raise SecretMapDecodeError(f"Cannot decode encrypted MCP {key}; check LITELLM_SALT_KEY") from None
|
||||
|
||||
|
||||
_JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
def _is_mapping(
|
||||
value: object,
|
||||
) -> TypeIs[Mapping[object, object]]: # guard-ok: isinstance decides it, keys and values stay object
|
||||
return isinstance(value, Mapping)
|
||||
|
||||
|
||||
def json_value(value: object) -> JsonValue:
|
||||
plain: Final[object] = dict(value) if _is_mapping(value) else value
|
||||
try:
|
||||
return _JSON_VALUE.validate_python(plain)
|
||||
except ValidationError:
|
||||
return _JSON_VALUE.validate_json(safe_dumps(plain))
|
||||
|
||||
|
||||
def _plaintext_of(value: str, signing_key: str | None = None) -> str:
|
||||
key: Final = signing_key or get_salt_key()
|
||||
if key is None:
|
||||
return value
|
||||
decrypted: Final = decrypt_if_encrypted_with(value, key)
|
||||
return value if decrypted is None else decrypted
|
||||
|
||||
|
||||
def _encrypted_string(value: str, new_encryption_key: str | None) -> str:
|
||||
if (new_encryption_key or get_salt_key()) is None:
|
||||
return value
|
||||
return _SECRET_STRING.validate_python(
|
||||
encrypt_value_helper(_plaintext_of(value), new_encryption_key=new_encryption_key), strict=True
|
||||
)
|
||||
|
||||
|
||||
def encrypt_json_strings(value: JsonValue, new_encryption_key: str | None = None, depth: int = 0) -> JsonValue:
|
||||
if depth > DEFAULT_MAX_RECURSE_DEPTH:
|
||||
raise ValueError(f"cannot encrypt a value nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels")
|
||||
if isinstance(value, str):
|
||||
return _encrypted_string(value, new_encryption_key)
|
||||
if isinstance(value, list):
|
||||
return [encrypt_json_strings(item, new_encryption_key, depth + 1) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {key: encrypt_json_strings(item, new_encryption_key, depth + 1) for key, item in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def decrypt_json_strings(value: JsonValue, signing_key: str | None = None, depth: int = 0) -> JsonValue:
|
||||
if depth > DEFAULT_MAX_RECURSE_DEPTH:
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
return _plaintext_of(value, signing_key)
|
||||
if isinstance(value, list):
|
||||
return [decrypt_json_strings(item, signing_key, depth + 1) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {key: decrypt_json_strings(item, signing_key, depth + 1) for key, item in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def parse_stored_json_object(value: object) -> Mapping[str, JsonValue] | None:
|
||||
try:
|
||||
return _JSON_OBJECT.validate_json(value) if isinstance(value, str) else _JSON_OBJECT.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _stored_json_object(value: object) -> Mapping[str, JsonValue]:
|
||||
parsed: Final = parse_stored_json_object(value)
|
||||
if parsed is not None:
|
||||
return parsed
|
||||
if value is not None:
|
||||
verbose_proxy_logger.error("Stored value is not a JSON object (%s); reading it as empty", type(value).__name__)
|
||||
return {}
|
||||
|
||||
|
||||
def decrypt_stored_json_object(value: object) -> Mapping[str, JsonValue]:
|
||||
return {key: decrypt_json_strings(item) for key, item in _stored_json_object(value).items()}
|
||||
|
||||
|
||||
ENCRYPTED_CONFIG_SECTIONS: Final = frozenset({"router_settings"})
|
||||
|
||||
|
||||
def encrypt_config_section(
|
||||
section_name: str, section: object, new_encryption_key: str | None = None
|
||||
) -> Mapping[str, JsonValue]:
|
||||
values: Final = _stored_json_object(section)
|
||||
if section_name not in ENCRYPTED_CONFIG_SECTIONS:
|
||||
return values
|
||||
return {key: encrypt_json_strings(item, new_encryption_key=new_encryption_key) for key, item in values.items()}
|
||||
|
||||
|
||||
def decrypt_config_section(section_name: str, value: object) -> Mapping[str, JsonValue]:
|
||||
if section_name in ENCRYPTED_CONFIG_SECTIONS:
|
||||
return decrypt_stored_json_object(value)
|
||||
return _stored_json_object(value)
|
||||
|
|
|
|||
|
|
@ -759,7 +759,6 @@ class InMemoryGuardrailHandler:
|
|||
return self.IN_MEMORY_GUARDRAILS[guardrail_id]
|
||||
|
||||
litellm_params_data: Final = guardrail["litellm_params"]
|
||||
verbose_proxy_logger.debug("litellm_params= %s", litellm_params_data)
|
||||
|
||||
if isinstance(litellm_params_data, dict):
|
||||
if reject_invalid_logging_only_scope:
|
||||
|
|
@ -1098,7 +1097,7 @@ class InMemoryGuardrailHandler:
|
|||
|
||||
# Log differences if any found
|
||||
if changed_fields:
|
||||
verbose_proxy_logger.debug("Guardrail params changed. Differences: %s", changed_fields)
|
||||
verbose_proxy_logger.debug("Guardrail params changed: %s", sorted(changed_fields))
|
||||
|
||||
# Return True if any fields changed
|
||||
return len(changed_fields) > 0
|
||||
|
|
|
|||
|
|
@ -56,6 +56,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( # noqa: F401 #
|
|||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
get_salt_key,
|
||||
parse_stored_json_object,
|
||||
)
|
||||
|
||||
ValueClass = Literal["migrated", "legacy", "plaintext", "undecryptable", "not-a-string"]
|
||||
|
|
@ -470,6 +471,7 @@ def _classify_callback_value(value: object) -> ValueClass:
|
|||
_COVERED_TABLE_SPECS: Final = [
|
||||
("model_table", "litellm_proxymodeltable", ("litellm_params",), ()),
|
||||
("credentials", "litellm_credentialstable", ("credential_values",), ()),
|
||||
("guardrails", "litellm_guardrailstable", ("litellm_params",), ()),
|
||||
("mcp_server", "litellm_mcpservertable", ("credentials", "env_vars", "static_headers", "env"), ()),
|
||||
("mcp_user_credentials", "litellm_mcpusercredentials", (), ("credential_b64",)),
|
||||
("mcp_user_env_vars", "litellm_mcpuserenvvars", (), ("values_b64",)),
|
||||
|
|
@ -502,7 +504,9 @@ def _classify_into_report(report: LocationReport, value: str) -> None:
|
|||
names, base URLs, …) do not decrypt and fall through to ``plaintext``, so
|
||||
over-scanning a column is harmless to the residual count.
|
||||
"""
|
||||
report.count(classify_value(value, key="scan"))
|
||||
from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX
|
||||
|
||||
report.count(classify_value(value.removeprefix(CALLBACK_VAR_ENCRYPTED_PREFIX), key="scan"))
|
||||
|
||||
|
||||
def _classify_secret_map(value: object, key: str) -> ValueClass | None:
|
||||
|
|
@ -554,23 +558,23 @@ async def _scan_one_table(
|
|||
return report
|
||||
|
||||
|
||||
async def _scan_config_env_vars(prisma_client: object) -> LocationReport:
|
||||
"""Scan the ``environment_variables`` config row (``param_value`` dict)."""
|
||||
report: Final = LocationReport(location="config_environment_variables")
|
||||
_CONFIG_SECTION_LOCATIONS: Final = (
|
||||
("environment_variables", "config_environment_variables"),
|
||||
("router_settings", "config_router_settings"),
|
||||
)
|
||||
|
||||
|
||||
async def _scan_config_section(prisma_client: object, param_name: str, location: str) -> LocationReport:
|
||||
"""Scan one encrypted ``LiteLLM_Config`` row (``param_value`` JSON object)."""
|
||||
report: Final = LocationReport(location=location)
|
||||
try:
|
||||
record: Final = await prisma_client.db.litellm_config.find_unique(where={"param_name": "environment_variables"})
|
||||
record: Final = await prisma_client.db.litellm_config.find_unique(where={"param_name": param_name})
|
||||
except Exception as e: # pragma: no cover - defensive
|
||||
verbose_proxy_logger.debug("scan: config env vars unavailable: %s", str(e))
|
||||
verbose_proxy_logger.debug("scan: config %s unavailable: %s", param_name, str(e))
|
||||
return report
|
||||
if record is None or record.param_value is None:
|
||||
return report
|
||||
value = record.param_value
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = _DECODED_JSON.validate_python(json.loads(value))
|
||||
except (ValueError, TypeError):
|
||||
value = {}
|
||||
for s in _iter_encrypted_strings(value):
|
||||
for s in _iter_encrypted_strings(parse_stored_json_object(record.param_value) or {}):
|
||||
_classify_into_report(report, s)
|
||||
return report
|
||||
|
||||
|
|
@ -580,7 +584,8 @@ async def _scan_covered_tables(prisma_client: object) -> list[LocationReport]:
|
|||
reports: Final[list[LocationReport]] = []
|
||||
for location, db_attr, json_cols, scalar_cols in _COVERED_TABLE_SPECS:
|
||||
reports.append(await _scan_one_table(prisma_client, location, db_attr, json_cols, scalar_cols))
|
||||
reports.append(await _scan_config_env_vars(prisma_client))
|
||||
for param_name, location in _CONFIG_SECTION_LOCATIONS:
|
||||
reports.append(await _scan_config_section(prisma_client, param_name, location))
|
||||
return reports
|
||||
|
||||
|
||||
|
|
@ -595,7 +600,8 @@ _CLOUDZERO_SENSITIVE: Final = ["api_key"]
|
|||
|
||||
async def _migrate_covered_tables(prisma_client: object, user_api_key_dict: object) -> list[LocationReport]:
|
||||
"""Re-encrypt the tables already covered by ``_rotate_master_key`` (model
|
||||
table, credentials, MCP credential/env tables, config environment_variables)
|
||||
table, credentials, guardrails, MCP credential/env tables, config
|
||||
environment_variables and router_settings)
|
||||
by running that orchestrator in *same-key* mode. With the AES gate on, the
|
||||
re-encrypt writes land in ``v2:`` format.
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.auth.auth_utils import fallback_target_model_name
|
||||
from litellm.proxy.auth.model_checks import get_all_fallbacks
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_config_section, encrypt_config_section
|
||||
from litellm.proxy.utils import PrismaClient, invalidate_config_param
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -142,7 +143,7 @@ async def _stored_router_settings(prisma_client: PrismaClient) -> Mapping[str, J
|
|||
return {}
|
||||
if row is None or not isinstance(row.param_value, Mapping):
|
||||
return {}
|
||||
return _STORED_ROUTER_SETTINGS.validate_python(row.param_value)
|
||||
return decrypt_config_section(ROUTER_SETTINGS_PARAM, _STORED_ROUTER_SETTINGS.validate_python(row.param_value))
|
||||
|
||||
|
||||
def _editable_rules(
|
||||
|
|
@ -174,7 +175,7 @@ async def _write_rules(
|
|||
|
||||
|
||||
async def _persist_router_settings(prisma_client: PrismaClient, router_settings: Mapping[str, object]) -> None:
|
||||
router_settings_json: Final = json.dumps(router_settings)
|
||||
router_settings_json: Final = json.dumps(encrypt_config_section(ROUTER_SETTINGS_PARAM, router_settings))
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": ROUTER_SETTINGS_PARAM},
|
||||
data={
|
||||
|
|
|
|||
|
|
@ -76,6 +76,11 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
|
|||
coordination_redis_cache,
|
||||
publish_config_change,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
ENCRYPTED_CONFIG_SECTIONS,
|
||||
encrypt_config_section,
|
||||
parse_stored_json_object,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for
|
||||
from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
|
@ -5917,6 +5922,27 @@ async def delete_key_aliases(
|
|||
)
|
||||
|
||||
|
||||
async def _rekey_config_section(
|
||||
prisma_client: PrismaClient, param_name: str, param_value: object, new_master_key: str
|
||||
) -> None:
|
||||
import prisma
|
||||
|
||||
section: Final = parse_stored_json_object(param_value)
|
||||
if section is None:
|
||||
verbose_proxy_logger.error(
|
||||
"Master key rotation skipped config %s: stored value is not a JSON object", param_name
|
||||
)
|
||||
return
|
||||
try:
|
||||
rekeyed: Final = encrypt_config_section(param_name, section, new_encryption_key=new_master_key)
|
||||
except ValueError as too_deep:
|
||||
verbose_proxy_logger.error("Master key rotation skipped config %s: %s", param_name, too_deep)
|
||||
return
|
||||
await _config_table(prisma_client).update(
|
||||
where={"param_name": param_name}, data={"param_value": prisma.Json(rekeyed)}
|
||||
)
|
||||
|
||||
|
||||
async def rotate_master_key(
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
@ -6001,6 +6027,10 @@ async def rotate_master_key(
|
|||
data={"param_value": prisma.Json({**environment_variables_dict, **encrypted_env_vars})},
|
||||
)
|
||||
|
||||
for c in config:
|
||||
if c.param_name in ENCRYPTED_CONFIG_SECTIONS and c.param_value is not None:
|
||||
await _rekey_config_section(prisma_client, c.param_name, c.param_value, new_master_key)
|
||||
|
||||
try:
|
||||
from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
|
||||
|
||||
|
|
|
|||
|
|
@ -74,8 +74,9 @@ from litellm.proxy.common_utils.credential_hydration import (
|
|||
write_touches_federation_surface,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_json_strings,
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
json_value,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import proxy_exception_for
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
|
@ -126,7 +127,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
|||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.chunked_in import find_many_in
|
||||
from litellm.repositories.credentials_repository import CredentialsRepository
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.model_repository import ModelRepository, encrypt_model_litellm_params
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.project_repository import ProjectRepository
|
||||
from litellm.repositories.table_repositories import ModelTableRepository
|
||||
|
|
@ -510,8 +511,9 @@ def _effective_complexity_router_config(
|
|||
incoming: Final = None if incoming_params is None else incoming_params.complexity_router_config
|
||||
existing: Final = None if existing_params is None else existing_params.complexity_router_config
|
||||
config_adapter: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
|
||||
stored: Final = decrypt_json_strings(json_value(config_adapter.validate_python(existing)))
|
||||
return resolve_complexity_router_config_write(
|
||||
config_adapter.validate_python(incoming), config_adapter.validate_python(existing)
|
||||
config_adapter.validate_python(incoming), config_adapter.validate_python(stored)
|
||||
).effective
|
||||
|
||||
|
||||
|
|
@ -1432,19 +1434,16 @@ def update_db_model(
|
|||
k: v for k, v in stored_model_info.items() if k not in echoed_pricing
|
||||
}
|
||||
|
||||
# update litellm params
|
||||
if updated_patch.litellm_params:
|
||||
# Encrypt any sensitive values
|
||||
encrypted_params: Final = {
|
||||
patched_params: Final = {
|
||||
k: (
|
||||
_effective_complexity_router_config(updated_patch.litellm_params, db_model.litellm_params)
|
||||
if k == "complexity_router_config"
|
||||
else encrypt_value_helper(v)
|
||||
else v
|
||||
)
|
||||
for k, v in updated_patch.litellm_params.model_dump(exclude_none=True, exclude_unset=True).items()
|
||||
}
|
||||
|
||||
merged_litellm_params.update(encrypted_params)
|
||||
merged_litellm_params.update(encrypt_model_litellm_params({**merged_litellm_params, **patched_params}))
|
||||
|
||||
# update model info
|
||||
if updated_patch.model_info:
|
||||
|
|
@ -2062,11 +2061,8 @@ async def add_model_to_db(
|
|||
model_params.litellm_params, None
|
||||
)
|
||||
_original_litellm_model_name: Final = model_params.litellm_params.model
|
||||
for k, v in _litellm_params_dict.items():
|
||||
encrypted_value = (
|
||||
encrypt_value_helper(value=v, new_encryption_key=new_encryption_key) if isinstance(v, str) else v
|
||||
)
|
||||
model_params.litellm_params[k] = encrypted_value
|
||||
for k, v in encrypt_model_litellm_params(_litellm_params_dict, new_encryption_key).items():
|
||||
model_params.litellm_params[k] = v
|
||||
_data: Final[dict] = {
|
||||
"model_id": model_params.model_info.id,
|
||||
"model_name": model_params.model_name,
|
||||
|
|
@ -2557,6 +2553,15 @@ def _canonical_session_tags(tags: Sequence[AwsSessionTag]) -> tuple[tuple[str, s
|
|||
return tuple(sorted((tag["Key"], tag["Value"]) for tag in tags))
|
||||
|
||||
|
||||
_SESSION_TAGS: Final = TypeAdapter(list[AwsSessionTag])
|
||||
|
||||
|
||||
def _stored_session_tags(existing_litellm_params: GenericLiteLLMParams | None) -> Sequence[AwsSessionTag] | None:
|
||||
if existing_litellm_params is None or existing_litellm_params.aws_session_tags is None:
|
||||
return None
|
||||
return _SESSION_TAGS.validate_python(decrypt_json_strings(json_value(existing_litellm_params.aws_session_tags)))
|
||||
|
||||
|
||||
class ModelManagementAuthChecks:
|
||||
"""
|
||||
Common auth checks for model management endpoints
|
||||
|
|
@ -2624,7 +2629,7 @@ class ModelManagementAuthChecks:
|
|||
return True
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return True
|
||||
existing_tags: Final = existing_litellm_params.aws_session_tags if existing_litellm_params is not None else None
|
||||
existing_tags: Final = _stored_session_tags(existing_litellm_params)
|
||||
if existing_tags is not None and _canonical_session_tags(existing_tags) == _canonical_session_tags(
|
||||
litellm_params.aws_session_tags or ()
|
||||
):
|
||||
|
|
@ -3369,25 +3374,25 @@ async def update_model(
|
|||
|
||||
_new_litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True)
|
||||
|
||||
### ENCRYPT PARAMS ###
|
||||
encrypted_params: Final = MappingProxyType(
|
||||
requested_params: Final = MappingProxyType(
|
||||
{
|
||||
k: (
|
||||
_effective_complexity_router_config(model_params.litellm_params, deployment.litellm_params)
|
||||
if k == "complexity_router_config"
|
||||
else encrypt_value_helper(value=v)
|
||||
else v
|
||||
)
|
||||
for k, v in _new_litellm_params_dict.items()
|
||||
}
|
||||
)
|
||||
|
||||
### MERGE WITH EXISTING DATA ###
|
||||
_mp: Final[dict[str, object]] = model_params.litellm_params.dict()
|
||||
merged_dictionary: Final = {
|
||||
key: _existing_litellm_params_dict[key] if value is None else encrypted_params[key]
|
||||
for key, value in _mp.items()
|
||||
if value is not None or _existing_litellm_params_dict.get(key) is not None
|
||||
}
|
||||
merged_dictionary: Final = encrypt_model_litellm_params(
|
||||
{
|
||||
key: _existing_litellm_params_dict[key] if value is None else requested_params[key]
|
||||
for key, value in _mp.items()
|
||||
if value is not None or _existing_litellm_params_dict.get(key) is not None
|
||||
}
|
||||
)
|
||||
|
||||
base_update: Final[PrismaCompatibleUpdateDBModel] = {
|
||||
"litellm_params": json.dumps(merged_dictionary),
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from pydantic import TypeAdapter
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth, user_api_key_has_admin_view
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, parse_stored_json_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
tag_cache_key,
|
||||
tag_registry_cache_key,
|
||||
|
|
@ -32,7 +33,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
|
|||
)
|
||||
from litellm.proxy.management_helpers.utils import handle_budget_for_entity
|
||||
from litellm.repositories.daily_activity_buckets import api_key_from_bucket
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.model_repository import ModelRepository, encrypt_model_litellm_params
|
||||
from litellm.repositories.table_repositories import (
|
||||
DailyTagSpendRepository,
|
||||
TagRepository,
|
||||
|
|
@ -365,7 +366,6 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str):
|
|||
raise HTTPException(status_code=500, detail="Database not connected")
|
||||
|
||||
try:
|
||||
# Get current model from database to preserve encrypted fields
|
||||
db_model = await ModelRepository(prisma_client).table.find_unique(where={"model_id": deployment.model_info.id})
|
||||
|
||||
if db_model is None:
|
||||
|
|
@ -374,24 +374,20 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str):
|
|||
detail=f"Model {deployment.model_info.id} not found in database",
|
||||
)
|
||||
|
||||
# Prisma returns litellm_params as dict (already parsed from JSON)
|
||||
existing_params = db_model.litellm_params
|
||||
if isinstance(existing_params, str): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub is str
|
||||
# If it's a string, parse it
|
||||
existing_params = json.loads(existing_params)
|
||||
elif not isinstance(existing_params, dict): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub
|
||||
raise Exception(f"Unexpected litellm_params type: {type(existing_params)}")
|
||||
stored_params: Final = parse_stored_json_object(db_model.litellm_params)
|
||||
if stored_params is None:
|
||||
raise Exception(f"Unexpected litellm_params type: {type(db_model.litellm_params)}")
|
||||
existing_params: Final = {key: decrypt_json_strings(value) for key, value in stored_params.items()}
|
||||
stored_tags: Final = existing_params.get("tags")
|
||||
tags: Final = stored_tags if isinstance(stored_tags, list) else []
|
||||
if tag in tags:
|
||||
return
|
||||
|
||||
# Add tag to tags array (preserve encryption of other fields)
|
||||
if "tags" not in existing_params:
|
||||
existing_params["tags"] = []
|
||||
if tag not in existing_params["tags"]:
|
||||
existing_params["tags"].append(tag)
|
||||
|
||||
# Update database with modified params (keeps encrypted fields encrypted)
|
||||
await ModelRepository(prisma_client).table.update(
|
||||
where={"model_id": deployment.model_info.id},
|
||||
data={"litellm_params": json.dumps(existing_params)},
|
||||
data={
|
||||
"litellm_params": json.dumps(encrypt_model_litellm_params({**existing_params, "tags": [*tags, tag]}))
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error adding tag to deployment: %s", e)
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
|
|||
check_team_member_model_access, # pyright: ignore[reportPrivateUsage] # shared membership authorization owner
|
||||
)
|
||||
from litellm.proxy.auth.team_grants import team_model_aliases
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, decrypt_value_helper, json_value
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.prisma_protocols import DatabaseClient
|
||||
from litellm.repositories.project_repository import ProjectRepository
|
||||
|
|
@ -130,6 +130,14 @@ class _RouterConfigSource(LiteLLMBaseModel):
|
|||
complexity_router_config: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
_ROUTER_CONFIG: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
|
||||
|
||||
|
||||
def _decrypted_router_config(existing: Deployment) -> Mapping[str, object] | None:
|
||||
stored: Final = _RouterConfigSource.model_validate(existing.litellm_params.model_dump()).complexity_router_config
|
||||
return _ROUTER_CONFIG.validate_python(decrypt_json_strings(json_value(stored)))
|
||||
|
||||
|
||||
class _MembershipKey(TypedDict):
|
||||
user_id: ReadOnly[str]
|
||||
team_id: ReadOnly[str]
|
||||
|
|
@ -380,11 +388,7 @@ async def authorize_member_auto_router_write(
|
|||
if existing is not None and incoming.model_name not in (None, public_name, existing.model_name):
|
||||
raise HTTPException(status_code=403, detail="Team members cannot rename an auto router.")
|
||||
supplied_config: Final = _RouterConfigSource.model_validate(params.model_dump()).complexity_router_config
|
||||
stored_config: Final = (
|
||||
_RouterConfigSource.model_validate(existing.litellm_params.model_dump()).complexity_router_config
|
||||
if existing is not None
|
||||
else None
|
||||
)
|
||||
stored_config: Final = _decrypted_router_config(existing) if existing is not None else None
|
||||
resolved_config: Final = resolve_complexity_router_config_write(supplied_config, stored_config)
|
||||
if resolved_config.supplied_connection_fields:
|
||||
raise HTTPException(status_code=403, detail="Team members cannot change classifier connections.")
|
||||
|
|
|
|||
|
|
@ -531,8 +531,13 @@ from litellm.proxy.common_utils.discoverable_model_filter import (
|
|||
undiscoverable_model_names,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
ENCRYPTED_CONFIG_SECTIONS,
|
||||
decrypt_config_section,
|
||||
decrypt_json_strings,
|
||||
decrypt_value_helper,
|
||||
encrypt_config_section,
|
||||
encrypt_value_helper,
|
||||
json_value,
|
||||
)
|
||||
from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id
|
||||
from litellm.proxy.common_utils.fips import enforce_fips_mode
|
||||
|
|
@ -6247,11 +6252,7 @@ class ProxyConfig:
|
|||
existing_row: Final[_ConfigParamRow | None] = await config_table.find_first(where=config_where)
|
||||
existing_value: Final[object] = cast(object, existing_row.param_value) if existing_row is not None else None
|
||||
existing_section: Final[Mapping[str, JsonValue]] = (
|
||||
_CONFIG_SECTION_VALUES.validate_json(existing_value)
|
||||
if isinstance(existing_value, str)
|
||||
else _CONFIG_SECTION_VALUES.validate_python(existing_value)
|
||||
if isinstance(existing_value, Mapping)
|
||||
else MappingProxyType({})
|
||||
MappingProxyType({}) if existing_value is None else decrypt_config_section(section_name, existing_value)
|
||||
)
|
||||
merged_section: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{
|
||||
|
|
@ -6264,7 +6265,7 @@ class ProxyConfig:
|
|||
)
|
||||
if merged_section == existing_section:
|
||||
return None
|
||||
serialized_section: Final = json.dumps(dict(merged_section))
|
||||
serialized_section: Final = json.dumps(encrypt_config_section(section_name, merged_section))
|
||||
config_data: Final[_ConfigParamUpsert] = {
|
||||
"create": {"param_name": section_name, "param_value": serialized_section},
|
||||
"update": {"param_value": serialized_section},
|
||||
|
|
@ -7972,7 +7973,7 @@ class ProxyConfig:
|
|||
|
||||
def _resolve_db_litellm_param(self, key: str, value: object) -> object:
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
return decrypt_json_strings(json_value(value))
|
||||
|
||||
decrypted_value: Final = decrypt_value_helper(value=value, key=key, return_original_value=True)
|
||||
if (
|
||||
|
|
@ -8367,7 +8368,7 @@ class ProxyConfig:
|
|||
where={"param_name": "router_settings"}
|
||||
)
|
||||
db_values: Final = (
|
||||
_as_settings_mapping(db_router_settings.param_value)
|
||||
_as_settings_mapping(decrypt_config_section("router_settings", db_router_settings.param_value))
|
||||
if db_router_settings is not None and db_router_settings.param_value is not None
|
||||
else _EMPTY_SETTINGS_MAPPING
|
||||
)
|
||||
|
|
@ -8835,7 +8836,8 @@ class ProxyConfig:
|
|||
self._warned_unreadable_db_general_settings = True
|
||||
|
||||
def _prepared_db_settings_values(self, section: Section, value: object) -> Mapping[str, SettingsJsonValue]:
|
||||
scrubbed: Final = _scrub_db_overlay_remote_module_loads(section=section, db_value=value)
|
||||
db_value: Final = decrypt_config_section(section, value) if section in ENCRYPTED_CONFIG_SECTIONS else value
|
||||
scrubbed: Final = _scrub_db_overlay_remote_module_loads(section=section, db_value=db_value)
|
||||
return _as_settings_mapping(scrubbed)
|
||||
|
||||
def _apply_db_environment_variables(
|
||||
|
|
@ -20454,10 +20456,10 @@ async def update_config(
|
|||
)
|
||||
if row is None or row.param_value is None:
|
||||
return {}
|
||||
return dict(row.param_value)
|
||||
return dict(decrypt_config_section(param_name, row.param_value))
|
||||
|
||||
async def _upsert_section(param_name: str, value: dict) -> None:
|
||||
serialized: Final = json.dumps(value)
|
||||
serialized: Final = json.dumps(encrypt_config_section(param_name, json.dumps(value)))
|
||||
await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": param_name},
|
||||
data={
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import json
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import ConfigDict, JsonValue, TypeAdapter
|
||||
|
||||
from litellm.models.model import LiteLLM_ProxyModelTable
|
||||
from litellm.repositories.base_repository import BaseRepository, DbRecord, record_to_dict
|
||||
|
|
@ -23,6 +23,55 @@ class _ProxyModelTableRepository(PrismaTableRepository["prisma_models.LiteLLM_Pr
|
|||
table_name = "litellm_proxymodeltable"
|
||||
|
||||
|
||||
CLASSIFIER_CREDENTIAL_FIELDS: Final = frozenset({"api_key", "api_base"})
|
||||
CLASSIFIER_CONFIG_KEYS: Final = frozenset({"opensource_classifier_config", "jev_classifier_config"})
|
||||
|
||||
|
||||
def _classifier_with_encrypted_credentials(classifier: JsonValue, new_encryption_key: str | None) -> JsonValue:
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_json_strings
|
||||
|
||||
if not isinstance(classifier, dict):
|
||||
return classifier
|
||||
return {
|
||||
key: (
|
||||
encrypt_json_strings(value, new_encryption_key=new_encryption_key)
|
||||
if key in CLASSIFIER_CREDENTIAL_FIELDS and isinstance(value, str)
|
||||
else value
|
||||
)
|
||||
for key, value in classifier.items()
|
||||
}
|
||||
|
||||
|
||||
def _complexity_router_config_with_encrypted_classifier_credentials(
|
||||
config: JsonValue, new_encryption_key: str | None
|
||||
) -> JsonValue:
|
||||
if not isinstance(config, dict):
|
||||
return config
|
||||
return {
|
||||
key: (
|
||||
_classifier_with_encrypted_credentials(value, new_encryption_key)
|
||||
if key in CLASSIFIER_CONFIG_KEYS
|
||||
else value
|
||||
)
|
||||
for key, value in config.items()
|
||||
}
|
||||
|
||||
|
||||
def encrypt_model_litellm_params(
|
||||
litellm_params: Mapping[str, object], new_encryption_key: str | None = None
|
||||
) -> Mapping[str, JsonValue]:
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_json_strings, json_value
|
||||
|
||||
return {
|
||||
key: (
|
||||
_complexity_router_config_with_encrypted_classifier_credentials(json_value(value), new_encryption_key)
|
||||
if key == "complexity_router_config"
|
||||
else encrypt_json_strings(json_value(value), new_encryption_key=new_encryption_key)
|
||||
)
|
||||
for key, value in litellm_params.items()
|
||||
}
|
||||
|
||||
|
||||
class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
|
||||
"""Repository for proxy model database operations with encryption support."""
|
||||
|
||||
|
|
@ -39,30 +88,15 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
|
|||
return LiteLLM_ProxyModelTable
|
||||
|
||||
def _encrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Encrypt sensitive values in litellm_params."""
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
|
||||
encrypted: Final = {}
|
||||
for key, value in litellm_params.items():
|
||||
if isinstance(value, str):
|
||||
encrypted[key] = encrypt_value_helper(value, new_encryption_key=self._encryption_key)
|
||||
else:
|
||||
encrypted[key] = value
|
||||
return encrypted
|
||||
return encrypt_model_litellm_params(litellm_params, new_encryption_key=self._encryption_key)
|
||||
|
||||
def _decrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Decrypt sensitive values in litellm_params."""
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, json_value
|
||||
|
||||
decrypted: Final = {}
|
||||
for key, value in litellm_params.items():
|
||||
if isinstance(value, str):
|
||||
decrypted[key] = decrypt_value_helper(
|
||||
value, key=key, exception_type="debug", return_original_value=True
|
||||
)
|
||||
else:
|
||||
decrypted[key] = value
|
||||
return decrypted
|
||||
return {
|
||||
key: decrypt_json_strings(json_value(value), signing_key=self._encryption_key)
|
||||
for key, value in litellm_params.items()
|
||||
}
|
||||
|
||||
def _to_model(self, record: DbRecord | None) -> LiteLLM_ProxyModelTable | None:
|
||||
"""Convert a database record to a Model with decryption."""
|
||||
|
|
|
|||
|
|
@ -85,6 +85,8 @@ IGNORE_FUNCTIONS = [
|
|||
"_redacted_client_field", # bounded by the nesting depth of a JSON-decoded error body (a finite JSON tree, no cycles possible).
|
||||
"_sort_processed_sets", # bounded by the nesting depth of the log-record extra it walks (a finite JSON tree, no cycles possible).
|
||||
"scrub_json_strings", # max depth set (MAX_SCRUB_DEPTH); fails closed by returning "[Filtered]" for anything nested past the cap.
|
||||
"encrypt_json_strings", # max depth set.
|
||||
"decrypt_json_strings", # max depth set.
|
||||
"_without_requester_ip_address", # walks only the fresh tree _sanitize_request_body_for_spend_logs_payload returns (mapping cycles already cut), so it never goes deeper than that sanitizer did.
|
||||
"_merge_session_fields", # max depth set (_SESSION_MERGE_MAX_DEPTH); deeper mappings are replaced wholesale by the later session.update.
|
||||
"_json", # router_rust index views: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising ValueError at the cap.
|
||||
|
|
|
|||
|
|
@ -7,21 +7,36 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130
|
|||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
V2_GCM_PREFIX,
|
||||
decrypt_bearer_token,
|
||||
decrypt_config_section,
|
||||
decrypt_if_encrypted_with,
|
||||
decrypt_json_strings,
|
||||
decrypt_stored_json_object,
|
||||
decrypt_value_helper,
|
||||
encrypt_bearer_token,
|
||||
encrypt_config_section,
|
||||
encrypt_json_strings,
|
||||
encrypt_value,
|
||||
encrypt_value_helper,
|
||||
json_value,
|
||||
parse_stored_json_object,
|
||||
)
|
||||
|
||||
SALT_KEY = "sk-salt-aes-1234"
|
||||
|
||||
|
||||
def _use_aes(monkeypatch):
|
||||
"""Flip the write-time algorithm to AES-256-GCM for the duration of a test."""
|
||||
|
|
@ -31,7 +46,7 @@ def _use_aes(monkeypatch):
|
|||
@pytest.fixture(autouse=True)
|
||||
def _salt_key(monkeypatch):
|
||||
# Dominant convention in the test_litellm/ tree: set the key via env.
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-aes-1234")
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY)
|
||||
# Ensure the legacy default is in force unless a test opts into AES.
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
yield
|
||||
|
|
@ -241,6 +256,183 @@ def test_explicit_key_decrypt_supports_the_empty_master_key():
|
|||
assert decrypt_if_encrypted_with(base64.urlsafe_b64encode(written_with_empty_key).decode(), "") == "stored-secret"
|
||||
|
||||
|
||||
_NESTED_PARAMS = {
|
||||
"model": "openai/gpt-5.4-mini",
|
||||
"extra_headers": {"Authorization": "Bearer gateway-secret", "X-Trace": "trace-1"},
|
||||
"fallbacks": [{"gpt-5.5-mini": ["gpt-5.4-mini"]}],
|
||||
"rpm": 10,
|
||||
"enabled": True,
|
||||
"api_base": None,
|
||||
}
|
||||
|
||||
|
||||
class _StoredNestedParams(TypedDict):
|
||||
model: ReadOnly[str]
|
||||
extra_headers: ReadOnly[dict[str, str]]
|
||||
fallbacks: ReadOnly[list[dict[str, list[str]]]]
|
||||
rpm: ReadOnly[int]
|
||||
enabled: ReadOnly[bool]
|
||||
api_base: ReadOnly[None]
|
||||
|
||||
|
||||
_STORED_NESTED_PARAMS: Final = TypeAdapter(_StoredNestedParams)
|
||||
|
||||
|
||||
def _nested_string_leaves(stored: JsonValue) -> tuple[str, str, str, str]:
|
||||
params: Final = _STORED_NESTED_PARAMS.validate_python(stored)
|
||||
return (
|
||||
params["model"],
|
||||
params["extra_headers"]["Authorization"],
|
||||
params["extra_headers"]["X-Trace"],
|
||||
params["fallbacks"][0]["gpt-5.5-mini"][0],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_aes", [False, True])
|
||||
def test_encrypt_json_strings_encrypts_every_nested_string_leaf(monkeypatch, use_aes: bool):
|
||||
if use_aes:
|
||||
_use_aes(monkeypatch)
|
||||
|
||||
stored = encrypt_json_strings(_NESTED_PARAMS)
|
||||
|
||||
assert (stored["rpm"], stored["enabled"], stored["api_base"]) == (10, True, None)
|
||||
plaintexts = _nested_string_leaves(_NESTED_PARAMS)
|
||||
assert all(leaf != plain for leaf, plain in zip(_nested_string_leaves(stored), plaintexts))
|
||||
assert tuple(decrypt_if_encrypted_with(leaf, SALT_KEY) for leaf in _nested_string_leaves(stored)) == plaintexts
|
||||
assert decrypt_json_strings(stored) == _NESTED_PARAMS
|
||||
|
||||
|
||||
def test_encrypt_json_strings_reencrypts_ciphertext_under_the_new_key_without_double_wrapping():
|
||||
stored = encrypt_json_strings({"api_key": "top-secret", "nested": {"token": encrypt_value_helper("nested-secret")}})
|
||||
|
||||
rotated = encrypt_json_strings(stored, new_encryption_key="sk-rotated")
|
||||
|
||||
assert decrypt_if_encrypted_with(rotated["api_key"], "sk-rotated") == "top-secret"
|
||||
assert decrypt_if_encrypted_with(rotated["nested"]["token"], "sk-rotated") == "nested-secret"
|
||||
assert decrypt_if_encrypted_with(rotated["api_key"], SALT_KEY) is None
|
||||
|
||||
|
||||
def test_encrypt_json_strings_treats_a_legacy_plaintext_row_as_plaintext():
|
||||
legacy_row = {"api_key": "plain-secret", "extra_headers": {"Authorization": "Bearer plain"}}
|
||||
|
||||
assert decrypt_json_strings(legacy_row) == legacy_row
|
||||
assert decrypt_json_strings(encrypt_json_strings(legacy_row)) == legacy_row
|
||||
|
||||
|
||||
def test_encrypt_json_strings_without_any_key_stores_the_value_as_sent(monkeypatch):
|
||||
monkeypatch.delenv("LITELLM_SALT_KEY")
|
||||
monkeypatch.setattr(proxy_server, "master_key", None)
|
||||
|
||||
assert encrypt_json_strings(_NESTED_PARAMS) == _NESTED_PARAMS
|
||||
assert decrypt_json_strings(_NESTED_PARAMS) == _NESTED_PARAMS
|
||||
|
||||
|
||||
def test_encrypt_json_strings_refuses_a_value_nested_past_the_recursion_cap():
|
||||
value = {"secret": "deep-secret"}
|
||||
for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1):
|
||||
value = {"child": value}
|
||||
|
||||
with pytest.raises(ValueError, match=f"nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels"):
|
||||
encrypt_json_strings(value)
|
||||
assert decrypt_json_strings(value) == value
|
||||
|
||||
|
||||
def test_encrypt_json_strings_encrypts_a_value_nested_at_the_recursion_cap():
|
||||
value = "deep-secret"
|
||||
for _ in range(DEFAULT_MAX_RECURSE_DEPTH):
|
||||
value = {"child": value}
|
||||
|
||||
stored = encrypt_json_strings(value)
|
||||
|
||||
deepest = stored
|
||||
for _ in range(DEFAULT_MAX_RECURSE_DEPTH):
|
||||
deepest = deepest["child"]
|
||||
assert deepest != "deep-secret"
|
||||
assert decrypt_json_strings(stored) == value
|
||||
|
||||
|
||||
def test_json_value_coerces_non_json_containers():
|
||||
assert json_value(MappingProxyType({"headers": ("a", "b")})) == {"headers": ["a", "b"]}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stored", ["not json", "[1]", ["a"], None, 3])
|
||||
def test_decrypt_stored_json_object_returns_an_empty_object_for_a_non_object_row(stored: object):
|
||||
assert decrypt_stored_json_object(stored) == {}
|
||||
|
||||
|
||||
def test_decrypt_stored_json_object_reads_json_text_and_parsed_rows_alike():
|
||||
stored = encrypt_json_strings(_NESTED_PARAMS)
|
||||
|
||||
assert decrypt_stored_json_object(json.dumps(stored)) == _NESTED_PARAMS
|
||||
assert decrypt_stored_json_object(stored) == _NESTED_PARAMS
|
||||
|
||||
|
||||
def test_encrypt_config_section_encrypts_router_settings_string_leaves_only():
|
||||
router_settings = {
|
||||
"redis_password": "redis-pw",
|
||||
"num_retries": 2,
|
||||
"fallbacks": [{"gpt-5.5-mini": ["gpt-5.4-mini"]}],
|
||||
}
|
||||
|
||||
stored = encrypt_config_section("router_settings", router_settings)
|
||||
|
||||
assert stored["num_retries"] == 2
|
||||
assert stored["redis_password"] != "redis-pw"
|
||||
assert decrypt_if_encrypted_with(stored["redis_password"], SALT_KEY) == "redis-pw"
|
||||
assert decrypt_if_encrypted_with(stored["fallbacks"][0]["gpt-5.5-mini"][0], SALT_KEY) == "gpt-5.4-mini"
|
||||
assert decrypt_config_section("router_settings", json.dumps(stored)) == router_settings
|
||||
|
||||
|
||||
def test_encrypt_config_section_leaves_other_sections_readable():
|
||||
general_settings = {"alerting": ["slack"], "proxy_batch_write_at": 60}
|
||||
|
||||
assert encrypt_config_section("general_settings", general_settings) == general_settings
|
||||
assert decrypt_config_section("general_settings", json.dumps(general_settings)) == general_settings
|
||||
|
||||
|
||||
def test_encrypt_config_section_rekeys_router_settings_for_master_key_rotation():
|
||||
stored = encrypt_config_section("router_settings", {"redis_password": "redis-pw"})
|
||||
|
||||
rotated = encrypt_config_section("router_settings", stored, new_encryption_key="sk-rotated")
|
||||
|
||||
assert decrypt_if_encrypted_with(rotated["redis_password"], "sk-rotated") == "redis-pw"
|
||||
|
||||
|
||||
def test_decrypt_json_strings_reads_values_encrypted_under_the_given_key():
|
||||
stored = encrypt_json_strings(_NESTED_PARAMS, new_encryption_key="sk-other-key")
|
||||
|
||||
assert decrypt_json_strings(stored, signing_key="sk-other-key") == _NESTED_PARAMS
|
||||
assert decrypt_json_strings(stored)["extra_headers"]["Authorization"] != "Bearer gateway-secret"
|
||||
|
||||
|
||||
def test_encrypt_config_section_rekeys_a_row_stored_as_a_json_string():
|
||||
stored_row = json.dumps(encrypt_json_strings({"redis_password": "redis-secret", "num_retries": 2}))
|
||||
|
||||
rotated = encrypt_config_section("router_settings", stored_row, new_encryption_key="sk-rotated")
|
||||
|
||||
assert decrypt_if_encrypted_with(rotated["redis_password"], "sk-rotated") == "redis-secret"
|
||||
assert rotated["num_retries"] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stored_row", ("not json", json.dumps(["api_key"]), 7, None))
|
||||
def test_parse_stored_json_object_returns_none_for_a_row_that_is_not_a_json_object(stored_row):
|
||||
assert parse_stored_json_object(stored_row) is None
|
||||
|
||||
|
||||
def test_decrypt_stored_json_object_reports_a_malformed_row_without_its_contents(caplog, monkeypatch):
|
||||
import logging
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
monkeypatch.setattr(verbose_proxy_logger, "propagate", True)
|
||||
caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name)
|
||||
|
||||
assert decrypt_stored_json_object("leaked-secret-text") == {}
|
||||
assert decrypt_stored_json_object(None) == {}
|
||||
|
||||
assert [record.levelno for record in caplog.records] == [logging.ERROR]
|
||||
assert "leaked-secret-text" not in caplog.text
|
||||
|
||||
def test_bearer_token_opens_only_under_its_own_prefix():
|
||||
token = encrypt_bearer_token("session", prefix="kind_a_")
|
||||
relabeled = "kind_b_" + token.removeprefix("kind_a_")
|
||||
|
|
|
|||
|
|
@ -99,6 +99,7 @@ def _seeded_tables() -> Tables:
|
|||
"rpm": 10,
|
||||
"use_in_pass_through": False,
|
||||
"api_base": None,
|
||||
"extra_headers": {"Authorization": _encrypted("Bearer nested-secret"), "X-Trace": "plain"},
|
||||
},
|
||||
},
|
||||
{"model_id": "model-2", "litellm_params": {"api_key": _encrypted("other-deployment", UNRELATED_KEY)}},
|
||||
|
|
@ -141,7 +142,7 @@ def _seeded_tables() -> Tables:
|
|||
}
|
||||
|
||||
|
||||
_VALUES_UNDER_THE_PREVIOUS_KEY = 8
|
||||
_VALUES_UNDER_THE_PREVIOUS_KEY = 9
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -160,6 +161,8 @@ async def test_reencryption_moves_every_stored_shape_to_the_new_key_and_nothing_
|
|||
assert decrypt_if_encrypted_with(model_params["model"], NEW_KEY) == "openai/gpt-5.4-mini"
|
||||
assert decrypt_if_encrypted_with(model_params["api_key"], PREVIOUS_KEY) is None
|
||||
assert (model_params["rpm"], model_params["use_in_pass_through"], model_params["api_base"]) == (10, False, None)
|
||||
assert decrypt_if_encrypted_with(model_params["extra_headers"]["Authorization"], NEW_KEY) == "Bearer nested-secret"
|
||||
assert model_params["extra_headers"]["X-Trace"] == "plain"
|
||||
mcp_server = tables["LiteLLM_MCPServerTable"][0]
|
||||
assert decrypt_if_encrypted_with(mcp_server["static_headers"], NEW_KEY) == '{"X-Api-Key": "header-secret"}'
|
||||
assert mcp_server["credentials"]["aws_region_name"] == "us-east-1"
|
||||
|
|
|
|||
|
|
@ -1982,3 +1982,48 @@ async def test_rotate_guardrail_params_gives_up_on_a_row_that_keeps_changing(mon
|
|||
assert rows_updated == 0
|
||||
assert prisma_client.db.litellm_guardrailstable.update_many.await_count == GUARDRAIL_ROTATION_ATTEMPTS
|
||||
assert prisma_client.db.litellm_guardrailstable.find_unique.await_count == GUARDRAIL_ROTATION_ATTEMPTS - 1
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proxy_debug_log(caplog, monkeypatch):
|
||||
import logging
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
monkeypatch.setattr(verbose_proxy_logger, "propagate", True)
|
||||
caplog.set_level(logging.DEBUG, logger=verbose_proxy_logger.name)
|
||||
return caplog
|
||||
|
||||
|
||||
def test_initialize_guardrail_keeps_the_vendor_api_key_out_of_the_logs(proxy_debug_log):
|
||||
registry_module = _register_noop_initializer("quiet_init_test")
|
||||
guardrail = {
|
||||
"guardrail_name": "quiet-init",
|
||||
"litellm_params": {"guardrail": "quiet_init_test", "mode": "pre_call", "api_key": "vendor-secret-init"},
|
||||
}
|
||||
try:
|
||||
assert InMemoryGuardrailHandler().initialize_guardrail(guardrail=guardrail) is not None
|
||||
finally:
|
||||
registry_module.guardrail_initializer_registry.pop("quiet_init_test", None)
|
||||
|
||||
assert "vendor-secret-init" not in proxy_debug_log.text
|
||||
|
||||
|
||||
def test_has_guardrail_params_changed_logs_only_the_changed_field_names(proxy_debug_log):
|
||||
handler = InMemoryGuardrailHandler()
|
||||
handler.IN_MEMORY_GUARDRAILS["g-1"] = {
|
||||
"guardrail_id": "g-1",
|
||||
"guardrail_name": "moderation",
|
||||
"litellm_params": {"guardrail": "openai_moderation", "mode": "pre_call", "api_key": "old-vendor-secret"},
|
||||
}
|
||||
changed = {
|
||||
"guardrail_id": "g-1",
|
||||
"guardrail_name": "moderation",
|
||||
"litellm_params": {"guardrail": "openai_moderation", "mode": "pre_call", "api_key": "new-vendor-secret"},
|
||||
}
|
||||
|
||||
assert handler._has_guardrail_params_changed("g-1", changed) is True
|
||||
|
||||
assert "api_key" in proxy_debug_log.text
|
||||
assert "old-vendor-secret" not in proxy_debug_log.text
|
||||
assert "new-vendor-secret" not in proxy_debug_log.text
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ import pytest
|
|||
|
||||
from litellm._service_logger import ServiceTypes
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
V2_GCM_PREFIX,
|
||||
encrypt_value_helper,
|
||||
|
|
@ -483,6 +484,27 @@ async def test_scan_covered_tables_classifies_search_tool_params(salt_key, monke
|
|||
assert by_loc["search_tools"].plaintext == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scan_covered_tables_classifies_marked_guardrail_params(salt_key, monkeypatch):
|
||||
legacy = _legacy_ct("guardrail-secret", monkeypatch)
|
||||
|
||||
client = MagicMock()
|
||||
_empty_covered_tables(client)
|
||||
client.db.litellm_guardrailstable.find_many = AsyncMock(
|
||||
return_value=[
|
||||
SimpleNamespace(
|
||||
litellm_params={"guardrail": "openai_moderation", "api_key": CALLBACK_VAR_ENCRYPTED_PREFIX + legacy}
|
||||
)
|
||||
]
|
||||
)
|
||||
client.db.litellm_config.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
by_loc = {r.location: r for r in await cm._scan_covered_tables(client)}
|
||||
|
||||
assert by_loc["guardrails"].legacy == 1
|
||||
assert by_loc["guardrails"].plaintext == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("column", ("static_headers", "env"))
|
||||
@pytest.mark.parametrize("algorithm", ("xsalsa20-poly1305", "aes-256-gcm"))
|
||||
|
|
@ -601,6 +623,24 @@ async def test_migrate_covered_tables_reports_real_counts(salt_key, monkeypatch)
|
|||
assert by_loc["model_table"].already_v2 == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scan_covered_tables_reports_router_settings_next_to_environment_variables(salt_key, monkeypatch):
|
||||
legacy = _legacy_ct("sk-openai", monkeypatch)
|
||||
config_rows = {
|
||||
"environment_variables": SimpleNamespace(param_value={"OPENAI_API_KEY": legacy}),
|
||||
"router_settings": SimpleNamespace(param_value={"redis_password": "plain-redis-pw", "num_retries": 2}),
|
||||
}
|
||||
|
||||
client = MagicMock()
|
||||
_empty_covered_tables(client)
|
||||
client.db.litellm_config.find_unique = AsyncMock(side_effect=lambda where: config_rows.get(where["param_name"]))
|
||||
|
||||
by_loc = {r.location: r for r in await cm._scan_covered_tables(client)}
|
||||
|
||||
assert by_loc["config_environment_variables"].legacy == 1
|
||||
assert by_loc["config_router_settings"].plaintext == 1
|
||||
assert by_loc["config_router_settings"].legacy == 0
|
||||
|
||||
def _db_service_hooks() -> tuple[AsyncMock, MagicMock]:
|
||||
success: Final = AsyncMock()
|
||||
service_logging: Final = MagicMock(async_service_success_hook=success, async_service_failure_hook=AsyncMock())
|
||||
|
|
|
|||
|
|
@ -9811,6 +9811,7 @@ async def test_rotate_master_key_reencrypts_model_params_in_place(
|
|||
|
||||
# Mock config table — no env vars
|
||||
mock_prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
# Mock credentials table — no credentials
|
||||
mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock(
|
||||
|
|
@ -19116,6 +19117,7 @@ async def test_rotate_master_key_rotates_sso_identity_assertions(
|
|||
)
|
||||
)
|
||||
mock_prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock(
|
||||
return_value=[]
|
||||
)
|
||||
|
|
@ -23570,6 +23572,139 @@ class TestTeamAdminMemberKeyBudgetUpdate:
|
|||
assert "member_key_budgets" not in str(exc.value.detail)
|
||||
|
||||
|
||||
def _rotation_prisma_client(stored_router_settings: object) -> AsyncMock:
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db = MagicMock()
|
||||
mock_prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.tx = MagicMock(
|
||||
return_value=AsyncMock(__aenter__=AsyncMock(return_value=AsyncMock()), __aexit__=AsyncMock(return_value=False))
|
||||
)
|
||||
mock_prisma_client.db.litellm_config.find_many = AsyncMock(
|
||||
return_value=[SimpleNamespace(param_name="router_settings", param_value=stored_router_settings)]
|
||||
)
|
||||
mock_prisma_client.db.litellm_config.update = AsyncMock()
|
||||
mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.writer_db = mock_prisma_client.db
|
||||
return mock_prisma_client
|
||||
|
||||
|
||||
async def _rotate_to_new_master_key(mock_prisma_client: AsyncMock, monkeypatch) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import rotate_master_key
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key")
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.decrypt_model_list_from_db.return_value = []
|
||||
with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config):
|
||||
await rotate_master_key(
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="test-user"
|
||||
),
|
||||
current_master_key="sk-old-master-key",
|
||||
new_master_key="sk-new-master-key",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key")
|
||||
async def test_rotate_master_key_rekeys_router_settings(mock_rotate_mcp, monkeypatch):
|
||||
import prisma
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_json_strings
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key")
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
stored_router_settings = encrypt_json_strings({"redis_password": "redis-pw", "num_retries": 2})
|
||||
mock_prisma_client = _rotation_prisma_client(stored_router_settings)
|
||||
mock_rotate_mcp.return_value = None
|
||||
|
||||
await _rotate_to_new_master_key(mock_prisma_client, monkeypatch)
|
||||
|
||||
router_update = mock_prisma_client.db.litellm_config.update.call_args
|
||||
assert router_update.kwargs["where"] == {"param_name": "router_settings"}
|
||||
assert isinstance(router_update.kwargs["data"]["param_value"], prisma.Json)
|
||||
rotated_router_settings = router_update.kwargs["data"]["param_value"].data
|
||||
assert decrypt_if_encrypted_with(rotated_router_settings["redis_password"], "sk-new-master-key") == "redis-pw"
|
||||
assert decrypt_if_encrypted_with(rotated_router_settings["redis_password"], "sk-old-master-key") is None
|
||||
assert rotated_router_settings["num_retries"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key")
|
||||
async def test_rotate_master_key_rekeys_router_settings_stored_as_a_json_string(mock_rotate_mcp, monkeypatch):
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_json_strings
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-old-master-key")
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
stored_router_settings = json.dumps(encrypt_json_strings({"redis_password": "redis-pw"}))
|
||||
mock_prisma_client = _rotation_prisma_client(stored_router_settings)
|
||||
mock_rotate_mcp.return_value = None
|
||||
|
||||
await _rotate_to_new_master_key(mock_prisma_client, monkeypatch)
|
||||
|
||||
rotated_router_settings = mock_prisma_client.db.litellm_config.update.call_args.kwargs["data"]["param_value"].data
|
||||
assert decrypt_if_encrypted_with(rotated_router_settings["redis_password"], "sk-new-master-key") == "redis-pw"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key")
|
||||
async def test_rotate_master_key_reports_router_settings_that_are_not_a_json_object_without_its_contents(
|
||||
mock_rotate_mcp, monkeypatch, caplog
|
||||
):
|
||||
import logging
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
monkeypatch.setattr(verbose_proxy_logger, "propagate", True)
|
||||
caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name)
|
||||
mock_prisma_client = _rotation_prisma_client("plaintext-router-secret")
|
||||
mock_rotate_mcp.return_value = None
|
||||
|
||||
await _rotate_to_new_master_key(mock_prisma_client, monkeypatch)
|
||||
|
||||
skipped: Final = [
|
||||
record.getMessage() for record in caplog.records if record.getMessage().startswith("Master key rotation skipped")
|
||||
]
|
||||
assert skipped == ["Master key rotation skipped config router_settings: stored value is not a JSON object"]
|
||||
assert "plaintext-router-secret" not in caplog.text
|
||||
mock_prisma_client.db.litellm_config.update.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("litellm.proxy.management_endpoints.key_management_endpoints.rotate_mcp_server_credentials_master_key")
|
||||
async def test_rotate_master_key_skips_router_settings_nested_past_the_recursion_cap(
|
||||
mock_rotate_mcp, monkeypatch, caplog
|
||||
):
|
||||
import logging
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
|
||||
monkeypatch.setattr(verbose_proxy_logger, "propagate", True)
|
||||
caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name)
|
||||
too_deep: dict[str, object] = {"secret": "vendor-secret"}
|
||||
for _ in range(DEFAULT_MAX_RECURSE_DEPTH + 1):
|
||||
too_deep = {"child": too_deep}
|
||||
mock_prisma_client = _rotation_prisma_client(too_deep)
|
||||
mock_rotate_mcp.return_value = None
|
||||
|
||||
await _rotate_to_new_master_key(mock_prisma_client, monkeypatch)
|
||||
|
||||
skipped: Final = [
|
||||
record.getMessage() for record in caplog.records if record.getMessage().startswith("Master key rotation skipped")
|
||||
]
|
||||
reason: Final = f"cannot encrypt a value nested deeper than {DEFAULT_MAX_RECURSE_DEPTH} levels"
|
||||
assert skipped == [f"Master key rotation skipped config router_settings: {reason}"]
|
||||
assert "vendor-secret" not in caplog.text
|
||||
mock_prisma_client.db.litellm_config.update.assert_not_called()
|
||||
|
||||
|
||||
_MEMBER_KEY_TEAM = LiteLLM_TeamTableCachedObj(
|
||||
team_id="team-1",
|
||||
members_with_roles=[
|
||||
|
|
|
|||
|
|
@ -30,7 +30,11 @@ from litellm.proxy._types import (
|
|||
ReconcileOutcome,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_json_strings,
|
||||
encrypt_value_helper,
|
||||
json_value,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
KubernetesPodsResponse,
|
||||
ModelManagementAuthChecks,
|
||||
|
|
@ -819,6 +823,32 @@ class TestModelManagementAuthChecks:
|
|||
)
|
||||
assert exc_info.value.code == "403"
|
||||
|
||||
def test_can_user_set_aws_session_tags_compares_against_the_decrypted_stored_tags(self, monkeypatch):
|
||||
from litellm.repositories.model_repository import encrypt_model_litellm_params
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-session-tags-salt-8627")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
stored: Final = encrypt_model_litellm_params({"aws_session_tags": [{"Key": "team", "Value": "genai"}]})
|
||||
assert stored["aws_session_tags"] != [{"Key": "team", "Value": "genai"}]
|
||||
existing: Final = LiteLLM_Params(model="bedrock/test_model", aws_session_tags=stored["aws_session_tags"])
|
||||
|
||||
result = ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "genai"}]
|
||||
),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
existing_litellm_params=existing,
|
||||
)
|
||||
assert result is True
|
||||
with pytest.raises(Exception, match="Only a proxy admin can set aws_session_tags"):
|
||||
ModelManagementAuthChecks.can_user_set_aws_session_tags(
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="bedrock/test_model", aws_session_tags=[{"Key": "team", "Value": "platform"}]
|
||||
),
|
||||
user_api_key_dict=self.team_admin_user,
|
||||
existing_litellm_params=existing,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_new_model_rejects_aws_session_tags_for_non_admin(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
|
@ -2045,8 +2075,8 @@ class TestUpdateModel:
|
|||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **kwargs: value,
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -2096,8 +2126,8 @@ class TestUpdateModel:
|
|||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **kwargs: value,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -3140,7 +3170,7 @@ class TestAddAndDeleteModelLifecycle:
|
|||
mock_router.registry.return_value.names.return_value = [model_id]
|
||||
|
||||
_PS = "litellm.proxy.proxy_server"
|
||||
_ENCRYPT = "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper"
|
||||
_ENCRYPT = "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper"
|
||||
with (
|
||||
patch(f"{_PS}.prisma_client", mock_prisma),
|
||||
patch(f"{_PS}.store_model_in_db", True),
|
||||
|
|
@ -4801,7 +4831,7 @@ class TestModelInfoServerDerivedPricingFilter:
|
|||
mock_router.registry.return_value.names.return_value = [model_id]
|
||||
|
||||
_PS = "litellm.proxy.proxy_server"
|
||||
_ENCRYPT = "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper"
|
||||
_ENCRYPT = "litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper"
|
||||
with (
|
||||
patch(f"{_PS}.prisma_client", mock_prisma),
|
||||
patch(f"{_PS}.store_model_in_db", True),
|
||||
|
|
@ -5212,7 +5242,7 @@ class TestUpdateDBModelClearCredentialName:
|
|||
litellm_params=updateLiteLLMParams(litellm_credential_name=None)
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
result: Final = update_db_model(db_model=db_model, updated_patch=update_patch)
|
||||
|
||||
params: Final = json.loads(result["litellm_params"])
|
||||
|
|
@ -5241,7 +5271,7 @@ class TestUpdateDBModelClearCredentialName:
|
|||
)
|
||||
update_patch: Final = updateDeployment(litellm_params=updateLiteLLMParams(tpm=10))
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
result: Final = update_db_model(db_model=db_model, updated_patch=update_patch)
|
||||
|
||||
params: Final = json.loads(result["litellm_params"])
|
||||
|
|
@ -5260,7 +5290,7 @@ class TestUpdateDBModelClearCredentialName:
|
|||
litellm_params=updateLiteLLMParams(litellm_credential_name=None)
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
result: Final = update_db_model(db_model=db_model, updated_patch=update_patch)
|
||||
|
||||
params: Final = json.loads(result["litellm_params"])
|
||||
|
|
@ -5288,7 +5318,7 @@ class TestUpdateDBModelClearCredentialName:
|
|||
)
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
result: Final = update_db_model(db_model=db_model, updated_patch=update_patch)
|
||||
|
||||
params: Final = json.loads(result["litellm_params"])
|
||||
|
|
@ -5315,7 +5345,7 @@ class TestUpdateDBModelClearCredentialName:
|
|||
litellm_params=updateLiteLLMParams(litellm_credential_name="other-credential")
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
with patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value):
|
||||
result: Final = update_db_model(db_model=db_model, updated_patch=update_patch)
|
||||
|
||||
params: Final = json.loads(result["litellm_params"])
|
||||
|
|
@ -5391,7 +5421,7 @@ class TestPatchModelCredentialName:
|
|||
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **kwargs: value,
|
||||
),
|
||||
patch(
|
||||
|
|
@ -7140,7 +7170,7 @@ class TestStrategyRouterWriteValidation:
|
|||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
lambda value, new_encryption_key=None: value,
|
||||
),
|
||||
):
|
||||
|
|
@ -7206,7 +7236,7 @@ class TestStrategyRouterWriteValidation:
|
|||
)
|
||||
with (
|
||||
patch( # test-quality-ok: params are encrypted with the proxy master key, which this test does not configure
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
lambda value, new_encryption_key=None: value,
|
||||
),
|
||||
patch( # test-quality-ok: the team list write is the collaborator whose ordering is asserted
|
||||
|
|
@ -7246,7 +7276,7 @@ class TestStrategyRouterWriteValidation:
|
|||
new=AsyncMock(return_value=None),
|
||||
),
|
||||
patch( # test-quality-ok: params are encrypted before the slot is entered; no master key in this test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
lambda value, new_encryption_key=None: value,
|
||||
),
|
||||
):
|
||||
|
|
@ -8622,7 +8652,7 @@ class TestAccessGroupModelSync:
|
|||
f"{self._MOD}.clear_cache",
|
||||
new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
|
||||
),
|
||||
patch(f"{self._MOD}.encrypt_value_helper", side_effect=lambda value, **kwargs: value),
|
||||
patch("litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper", side_effect=lambda value, **kwargs: value),
|
||||
):
|
||||
stack.enter_context(target)
|
||||
yield stack.enter_context(patch(self._INVALIDATE, new=AsyncMock()))
|
||||
|
|
@ -8990,7 +9020,7 @@ class TestTeamMemberAutoRouterWrites:
|
|||
await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor))
|
||||
written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"]
|
||||
saved: Final = json.loads(written["litellm_params"])["complexity_router_config"]
|
||||
assert saved == (
|
||||
assert decrypt_json_strings(json_value(saved)) == (
|
||||
self._classifier_config({**connection, "timeout_ms": 8100}, False) if change_config else stored_config
|
||||
)
|
||||
|
||||
|
|
@ -9079,7 +9109,8 @@ class TestTeamMemberAutoRouterWrites:
|
|||
await add_new_model(deployment, actor)
|
||||
written: Final = database.db.litellm_proxymodeltable.create.await_args.kwargs["data"]
|
||||
saved: Final = json.loads(written["litellm_params"])["complexity_router_config"]
|
||||
assert saved == {
|
||||
assert "stored-secret" not in written["litellm_params"]
|
||||
assert decrypt_json_strings(json_value(saved)) == {
|
||||
"classifier_type": "oss_classifier",
|
||||
"tiers": {"SIMPLE": "allowed"},
|
||||
"opensource_classifier_config": {**classifier, "provider": "laya" if provider == "laya" else "jev"},
|
||||
|
|
@ -9157,29 +9188,34 @@ class TestTeamMemberAutoRouterWrites:
|
|||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("endpoint", ["patch", "legacy"])
|
||||
@pytest.mark.parametrize("stored_legacy,supplied_legacy", [(True, True), (True, False), (False, True), (False, False)])
|
||||
@pytest.mark.parametrize("change", ["save", "rotate", "move", "move-without-key", "reset", "heuristic"])
|
||||
@pytest.mark.parametrize("stored_as", ["plaintext", "ciphertext"])
|
||||
@pytest.mark.parametrize(
|
||||
"change", ["save", "rotate", "move", "move-without-key", "same-base", "reset", "heuristic"]
|
||||
)
|
||||
async def test_jev_dashboard_save_preserves_server_transport(
|
||||
self, endpoint: str, change: str, stored_legacy: bool, supplied_legacy: bool
|
||||
self, endpoint: str, change: str, stored_legacy: bool, supplied_legacy: bool, stored_as: str
|
||||
) -> None:
|
||||
from litellm.repositories.model_repository import encrypt_model_litellm_params
|
||||
|
||||
original: Final = self._row()
|
||||
transport: Final = {"api_key": "synthetic-original-jev-key", "api_base": "https://jev.example.com"}
|
||||
stored_config: Final = self._classifier_config(
|
||||
{**transport, "instructions": "Old instructions", "timeout_ms": 6100}, stored_legacy
|
||||
)
|
||||
row: Final = original.model_copy(
|
||||
update={
|
||||
"litellm_params": {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": stored_config,
|
||||
},
|
||||
}
|
||||
plaintext_params: Final = {
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": self._classifier_config(
|
||||
{**transport, "instructions": "Old instructions", "timeout_ms": 6100}, stored_legacy
|
||||
),
|
||||
}
|
||||
stored_params: Final = (
|
||||
plaintext_params if stored_as == "plaintext" else encrypt_model_litellm_params(plaintext_params)
|
||||
)
|
||||
row: Final = original.model_copy(update={"litellm_params": stored_params})
|
||||
database: Final = self._database(self._team(), row)
|
||||
overrides: Final = {
|
||||
"save": {},
|
||||
"rotate": {"api_key": "synthetic-replacement-jev-key"},
|
||||
"move": {"api_base": "https://new-jev.example.com", "api_key": "synthetic-replacement-jev-key"},
|
||||
"move-without-key": {"api_base": "https://new-jev.example.com"},
|
||||
"same-base": {"api_base": "https://jev.example.com"},
|
||||
"reset": {"api_key": None, "api_base": None},
|
||||
"heuristic": {},
|
||||
}[change]
|
||||
|
|
@ -9205,17 +9241,23 @@ class TestTeamMemberAutoRouterWrites:
|
|||
await operation
|
||||
written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"]
|
||||
saved: Final = json.loads(written["litellm_params"])["complexity_router_config"]
|
||||
expected_classifier: Final = {**transport, "timeout_ms": 8100, **overrides}
|
||||
expected: Final = (
|
||||
config
|
||||
if change == "heuristic"
|
||||
else {
|
||||
"classifier_type": "oss_classifier",
|
||||
"tiers": {"SIMPLE": "allowed"},
|
||||
"opensource_classifier_config": {**transport, "timeout_ms": 8100, **overrides},
|
||||
"opensource_classifier_config": expected_classifier,
|
||||
}
|
||||
)
|
||||
assert saved == expected
|
||||
assert row.litellm_params["complexity_router_config"] == stored_config
|
||||
stored_classifier: Final = saved.get("opensource_classifier_config")
|
||||
if isinstance(stored_classifier, dict):
|
||||
for field in ("api_key", "api_base"):
|
||||
if isinstance(stored_classifier.get(field), str):
|
||||
assert stored_classifier[field] != expected_classifier[field]
|
||||
assert decrypt_json_strings(json_value(saved)) == expected
|
||||
assert row.litellm_params == stored_params
|
||||
assert request.litellm_params.complexity_router_config == config
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -9282,7 +9324,7 @@ class TestTeamMemberAutoRouterWrites:
|
|||
written: Final = database.db.litellm_proxymodeltable.update.await_args.kwargs["data"]
|
||||
saved: Final = json.loads(written["litellm_params"])["complexity_router_config"]
|
||||
expected_provider: Final = supplied.get("provider", stored_provider)
|
||||
assert saved == {
|
||||
assert decrypt_json_strings(json_value(saved)) == {
|
||||
"classifier_type": "oss_classifier",
|
||||
"tiers": {"SIMPLE": "allowed"},
|
||||
"opensource_classifier_config": {
|
||||
|
|
@ -9361,7 +9403,7 @@ class TestTeamMemberAutoRouterWrites:
|
|||
return
|
||||
response: Final = await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor))
|
||||
written: Final = database.transaction.litellm_proxymodeltable.update.await_args.kwargs["data"]
|
||||
saved_config: Final = json.loads(written["litellm_params"])["complexity_router_config"]
|
||||
saved_config: Final = decrypt_json_strings(json_value(json.loads(written["litellm_params"])["complexity_router_config"]))
|
||||
untouched: Final = config_shape in ("omit-config", "null-config")
|
||||
saved: Final = saved_config["jev_classifier_config" if untouched else "opensource_classifier_config"]
|
||||
assert saved == secret_params["complexity_router_config"]["jev_classifier_config"]
|
||||
|
|
@ -9385,6 +9427,37 @@ class TestTeamMemberAutoRouterWrites:
|
|||
}
|
||||
assert "retained-laya-secret" in row.model_dump_json()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("endpoint", ["patch", "legacy"])
|
||||
async def test_member_save_over_an_encrypted_stored_connection_keeps_it(self, endpoint: str) -> None:
|
||||
from litellm.repositories.model_repository import encrypt_model_litellm_params
|
||||
|
||||
config: Final = {
|
||||
"classifier_type": "jev", "tiers": {"SIMPLE": "allowed"},
|
||||
"jev_classifier_config": {"provider": "laya", "model": "english"},
|
||||
}
|
||||
connection: Final = {"api_key": "retained-laya-secret", "api_base": "https://laya.test"}
|
||||
stored_params: Final = encrypt_model_litellm_params({
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {**config, "jev_classifier_config": {**config["jev_classifier_config"], **connection}},
|
||||
})
|
||||
row: Final = self._row().model_copy(update={"litellm_params": dict(stored_params)})
|
||||
team: Final = self._team().model_copy(update={"models": ["allowed", "laya/english"]})
|
||||
database: Final = self._database(team, row)
|
||||
database.transaction.litellm_proxymodeltable.update.return_value = row
|
||||
request: Final = updateDeployment(
|
||||
litellm_params=updateLiteLLMParams.model_validate({"complexity_router_config": config}),
|
||||
model_info=ModelInfo(id=row.model_id, team_id="member-team"),
|
||||
)
|
||||
actor: Final = UserAPIKeyAuth(
|
||||
user_id="owner", user_role=LitellmUserRoles.INTERNAL_USER, models=["allowed", "laya/english"], config={"timeout": 60},
|
||||
)
|
||||
with self._environment(database, row):
|
||||
await (patch_model(row.model_id, request, actor) if endpoint == "patch" else update_model(request, actor))
|
||||
written: Final = database.transaction.litellm_proxymodeltable.update.await_args.kwargs["data"]
|
||||
saved: Final = decrypt_json_strings(json_value(json.loads(written["litellm_params"])["complexity_router_config"]))
|
||||
assert saved["opensource_classifier_config"] == {**config["jev_classifier_config"], **connection}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("endpoint", ["patch", "legacy"])
|
||||
@pytest.mark.parametrize("access", ["owner", "peer", "limited-key"])
|
||||
|
|
@ -9610,8 +9683,8 @@ class TestModelManagementActorEdges:
|
|||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **kwargs: value,
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -9659,8 +9732,8 @@ class TestModelManagementActorEdges:
|
|||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **kwargs: value,
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -9781,8 +9854,8 @@ class TestModelManagementActorEdges:
|
|||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **kwargs: value,
|
||||
),
|
||||
patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -9840,6 +9913,122 @@ class TestModelManagementActorEdges:
|
|||
prisma.db.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
|
||||
|
||||
class TestNestedLitellmParamsEncryption:
|
||||
SALT_KEY: Final = "sk-nested-salt-1234"
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _salt_key(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", self.SALT_KEY)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_new_model_encrypts_nested_extra_headers(self):
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model
|
||||
|
||||
model_id: Final = "nested-headers-model"
|
||||
db_row: Final = LiteLLM_ProxyModelTable(
|
||||
model_id=model_id,
|
||||
model_name="gateway-model",
|
||||
litellm_params={"model": "openai/gpt-5.4-mini"},
|
||||
model_info={"id": model_id},
|
||||
created_by="test-admin",
|
||||
updated_by="test-admin",
|
||||
)
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.litellm_proxymodeltable = AsyncMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=db_row)
|
||||
mock_proxy_config: Final = MagicMock()
|
||||
mock_proxy_config.add_deployment = AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None))
|
||||
mock_router: Final = MagicMock()
|
||||
mock_router.registry.return_value.names.return_value = [model_id]
|
||||
_PS: Final = "litellm.proxy.proxy_server"
|
||||
with (
|
||||
patch(f"{_PS}.prisma_client", mock_prisma),
|
||||
patch(f"{_PS}.store_model_in_db", True),
|
||||
patch(f"{_PS}.proxy_config", mock_proxy_config),
|
||||
patch(f"{_PS}.proxy_logging_obj", MagicMock()),
|
||||
patch(f"{_PS}.premium_user", True),
|
||||
patch(f"{_PS}.llm_router", mock_router),
|
||||
):
|
||||
await add_new_model(
|
||||
model_params=Deployment(
|
||||
model_name="gateway-model",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="openai/gpt-5.4-mini",
|
||||
api_key="sk-placeholder",
|
||||
extra_headers={"Authorization": "Bearer gateway-secret", "X-Gateway-Token": "gw-token"},
|
||||
rpm=10,
|
||||
),
|
||||
model_info={"id": model_id},
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
|
||||
create_call: Final = mock_prisma.db.litellm_proxymodeltable.create.call_args
|
||||
stored_text: Final = (create_call.kwargs["data"] if "data" in create_call.kwargs else create_call.args[0])[
|
||||
"litellm_params"
|
||||
]
|
||||
assert "gateway-secret" not in stored_text
|
||||
assert "gw-token" not in stored_text
|
||||
stored: Final = json.loads(stored_text)
|
||||
assert decrypt_if_encrypted_with(stored["extra_headers"]["Authorization"], self.SALT_KEY) == "Bearer gateway-secret"
|
||||
assert decrypt_if_encrypted_with(stored["extra_headers"]["X-Gateway-Token"], self.SALT_KEY) == "gw-token"
|
||||
assert decrypt_if_encrypted_with(stored["api_key"], self.SALT_KEY) == "sk-placeholder"
|
||||
assert stored["rpm"] == 10
|
||||
|
||||
def test_update_db_model_encrypts_nested_extra_headers_and_reencrypts_a_legacy_plaintext_row(self):
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
|
||||
|
||||
db_model: Final = Deployment(
|
||||
model_name="gateway-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-5.4-mini", api_key="legacy-plaintext-key"),
|
||||
model_info={"id": "legacy-model"},
|
||||
)
|
||||
update_patch: Final = updateDeployment(
|
||||
litellm_params=updateLiteLLMParams(extra_headers={"Authorization": "Bearer gateway-secret"})
|
||||
)
|
||||
|
||||
result: Final = update_db_model(db_model=db_model, updated_patch=update_patch)
|
||||
|
||||
assert "gateway-secret" not in result["litellm_params"]
|
||||
assert "legacy-plaintext-key" not in result["litellm_params"]
|
||||
stored: Final = json.loads(result["litellm_params"])
|
||||
assert decrypt_if_encrypted_with(stored["extra_headers"]["Authorization"], self.SALT_KEY) == "Bearer gateway-secret"
|
||||
assert decrypt_if_encrypted_with(stored["api_key"], self.SALT_KEY) == "legacy-plaintext-key"
|
||||
assert decrypt_if_encrypted_with(stored["model"], self.SALT_KEY) == "openai/gpt-5.4-mini"
|
||||
|
||||
def test_update_db_model_keeps_an_already_encrypted_value_decryptable_after_a_second_write(self):
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
|
||||
|
||||
first: Final = update_db_model(
|
||||
db_model=Deployment(
|
||||
model_name="gateway-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-5.4-mini", api_key="sk-first"),
|
||||
model_info={"id": "twice-written"},
|
||||
),
|
||||
updated_patch=updateDeployment(
|
||||
litellm_params=updateLiteLLMParams(extra_headers={"Authorization": "Bearer gateway-secret"})
|
||||
),
|
||||
)
|
||||
second: Final = update_db_model(
|
||||
db_model=Deployment(
|
||||
model_name="gateway-model",
|
||||
litellm_params=LiteLLM_Params(**json.loads(first["litellm_params"])),
|
||||
model_info={"id": "twice-written"},
|
||||
),
|
||||
updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(rpm=5)),
|
||||
)
|
||||
|
||||
stored: Final = json.loads(second["litellm_params"])
|
||||
assert decrypt_if_encrypted_with(stored["extra_headers"]["Authorization"], self.SALT_KEY) == "Bearer gateway-secret"
|
||||
assert decrypt_if_encrypted_with(stored["api_key"], self.SALT_KEY) == "sk-first"
|
||||
assert stored["rpm"] == 5
|
||||
|
||||
class TestAddModelToDbBlocked:
|
||||
"""`_add_model_to_db` must thread `blocked` into the initial insert, so the wizard can
|
||||
create a discovered-but-unchecked model already paused instead of active-then-patched."""
|
||||
|
|
@ -10300,8 +10489,8 @@ class TestNonAdminCannotPersistWifFieldsOnModel:
|
|||
True,
|
||||
),
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **_: value,
|
||||
),
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -10441,8 +10630,8 @@ class TestNonAdminCannotPersistWifFieldsOnModel:
|
|||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **_: value,
|
||||
),
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -10476,8 +10665,8 @@ class TestNonAdminCannotPersistWifFieldsOnModel:
|
|||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **_: value,
|
||||
),
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -10649,8 +10838,8 @@ class TestNonAdminCannotPersistWifFieldsOnModel:
|
|||
True,
|
||||
),
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **_: value,
|
||||
),
|
||||
patch( # test-quality-ok: the proxy wiring under test is what this patches
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -10931,8 +11120,8 @@ class TestWifBoundaryReadsTheResultingDeployment:
|
|||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
patch("litellm.proxy.proxy_server.premium_user", True),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **_: value,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
@ -11073,8 +11262,8 @@ class TestFederationGateScopesToWhatTheWriteTouches:
|
|||
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: proxy wiring under test
|
||||
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy wiring under test
|
||||
patch( # test-quality-ok: proxy wiring under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper",
|
||||
side_effect=lambda value: value,
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda value, **_: value,
|
||||
),
|
||||
patch( # test-quality-ok: proxy wiring under test
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
||||
|
|
|
|||
|
|
@ -1498,58 +1498,6 @@ async def test_add_tag_to_deployment_with_string_params():
|
|||
assert updated_params["api_key"] == "encrypted_claude_key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_tag_to_deployment_no_duplicate_tags():
|
||||
"""
|
||||
Test that _add_tag_to_deployment doesn't add duplicate tags
|
||||
"""
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
from litellm.proxy.management_endpoints.tag_management_endpoints import (
|
||||
_add_tag_to_deployment,
|
||||
)
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
# Setup prisma mocks
|
||||
mock_db = Mock()
|
||||
mock_prisma.db = mock_db
|
||||
|
||||
# Mock the database model with existing tags
|
||||
db_model = Mock()
|
||||
db_model.model_id = "model-789"
|
||||
db_model.litellm_params = {
|
||||
"model": "gpt-4",
|
||||
"api_key": "encrypted_key",
|
||||
"tags": ["existing-tag", "another-tag"],
|
||||
}
|
||||
|
||||
# Mock find_unique to return the db model
|
||||
mock_db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_model)
|
||||
|
||||
# Mock update
|
||||
mock_db.litellm_proxymodeltable.update = AsyncMock(return_value=db_model)
|
||||
|
||||
# Create deployment
|
||||
deployment = Deployment(
|
||||
model_name="gpt-4",
|
||||
litellm_params=LiteLLM_Params(model="gpt-4"),
|
||||
model_info=ModelInfo(id="model-789"),
|
||||
)
|
||||
|
||||
# Try to add an existing tag
|
||||
await _add_tag_to_deployment(deployment, "existing-tag")
|
||||
|
||||
# Verify update was called
|
||||
update_call = mock_db.litellm_proxymodeltable.update.call_args
|
||||
updated_params = json.loads(update_call[1]["data"]["litellm_params"])
|
||||
|
||||
# Verify no duplicate tags
|
||||
assert updated_params["tags"].count("existing-tag") == 1
|
||||
assert len(updated_params["tags"]) == 2
|
||||
assert "another-tag" in updated_params["tags"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_tag_to_deployment_model_not_found():
|
||||
"""
|
||||
|
|
@ -1585,6 +1533,40 @@ async def test_add_tag_to_deployment_model_not_found():
|
|||
assert "not found in database" in str(exc_info.value.detail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_tag_to_deployment_merges_into_the_decrypted_tags_and_writes_ciphertext(monkeypatch):
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, encrypt_json_strings
|
||||
from litellm.proxy.management_endpoints.tag_management_endpoints import _add_tag_to_deployment
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-tag-salt-8627")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
plaintext_params = {"model": "openai/gpt-5.4-mini", "api_key": "sk-placeholder", "tags": ["existing-tag"]}
|
||||
stored_params = encrypt_json_strings(plaintext_params)
|
||||
assert stored_params != plaintext_params
|
||||
deployment = Deployment(
|
||||
model_name="gpt-5.4-mini",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-5.4-mini"),
|
||||
model_info=ModelInfo(id="model-8627"),
|
||||
)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
|
||||
mock_db = Mock()
|
||||
mock_prisma.db = mock_db
|
||||
mock_db.litellm_proxymodeltable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(model_id="model-8627", litellm_params=stored_params)
|
||||
)
|
||||
mock_db.litellm_proxymodeltable.update = AsyncMock()
|
||||
|
||||
await _add_tag_to_deployment(deployment, "existing-tag")
|
||||
mock_db.litellm_proxymodeltable.update.assert_not_awaited()
|
||||
|
||||
await _add_tag_to_deployment(deployment, "new-tag")
|
||||
written = json.loads(mock_db.litellm_proxymodeltable.update.call_args.kwargs["data"]["litellm_params"])
|
||||
assert "new-tag" not in written["tags"]
|
||||
assert written["api_key"] != "sk-placeholder"
|
||||
assert decrypt_json_strings(written) == {**plaintext_params, "tags": ["existing-tag", "new-tag"]}
|
||||
|
||||
class _StoredTagTable:
|
||||
def __init__(self, model_info: object) -> None:
|
||||
self.model_info = model_info
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm.proxy._types import (
|
|||
ProxyException,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings
|
||||
from litellm.proxy.management_helpers.auto_router_permissions import (
|
||||
MemberAutoRouterDependencyObjects,
|
||||
authorize_member_auto_router_dependencies,
|
||||
|
|
@ -316,7 +317,7 @@ async def test_member_authorization_and_persistence_resolve_the_same_classifier(
|
|||
grant: Final = await operation
|
||||
persisted: Final = update_db_model(existing, patch)
|
||||
saved: Final = RequestComplexityRouterConfig.model_validate(
|
||||
json.loads(persisted["litellm_params"])["complexity_router_config"]
|
||||
decrypt_json_strings(json.loads(persisted["litellm_params"])["complexity_router_config"])
|
||||
)
|
||||
assert grant.config == saved
|
||||
assert saved.jev_classifier_config is not None
|
||||
|
|
|
|||
|
|
@ -1345,6 +1345,25 @@ async def test_ProxyConfig_save_config_writes_only_changed_router_settings(monke
|
|||
assert table.upserted_param_names == ["router_settings"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_encrypts_router_settings_and_reencrypts_the_legacy_row(monkeypatch):
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-save-config-salt-1234")
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"router_settings": {"redis_host": "redis.internal"}})
|
||||
baseline: Final = {"model_list": [], "router_settings": {"num_retries": 1}}
|
||||
proxy_config.update_config_state(config=baseline)
|
||||
changed: Final = {**baseline, "router_settings": {"num_retries": 1, "redis_password": "redis-canary-password"}}
|
||||
|
||||
await proxy_config.save_config(changed)
|
||||
|
||||
stored: Final = table.rows["router_settings"]
|
||||
assert "redis-canary-password" not in json.dumps(stored)
|
||||
assert "redis.internal" not in json.dumps(stored)
|
||||
assert decrypt_json_strings(stored) == {"redis_host": "redis.internal", "redis_password": "redis-canary-password"}
|
||||
assert table.upserted_param_names == ["router_settings"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_removes_a_key_only_when_the_db_has_it(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from redis.exceptions import NoPermissionError
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_json_strings, encrypt_json_strings
|
||||
from litellm.proxy.config_resolvers.settings_rules import Section
|
||||
from litellm.proxy.management_endpoints.fallback_management_endpoints import (
|
||||
FallbackCreateRequest,
|
||||
|
|
@ -344,9 +345,11 @@ class _RedisThatRefusesDel:
|
|||
class TestFallbackWritesSeeTheLatestStoredRules:
|
||||
"""A write reads the rules the database holds now and leaves no stale copy in the config cache behind"""
|
||||
|
||||
async def _create(self, request: FallbackCreateRequest, stored: _StoredRouterSettings) -> FallbackResponse:
|
||||
async def _create(
|
||||
self, request: FallbackCreateRequest, stored: _StoredRouterSettings, router: Router | None = None
|
||||
) -> FallbackResponse:
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.llm_router", _team_router()),
|
||||
patch("litellm.proxy.proxy_server.llm_router", router or _team_router()),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client),
|
||||
patch("litellm.proxy.proxy_server.proxy_config", stored.proxy_config),
|
||||
patch("litellm.proxy.proxy_server.store_model_in_db", True),
|
||||
|
|
@ -369,6 +372,28 @@ class TestFallbackWritesSeeTheLatestStoredRules:
|
|||
|
||||
assert stored.stored_fallbacks() == [TEAM_RULE, GATEWAY_RULE]
|
||||
|
||||
async def test_a_create_beside_an_encrypted_row_applies_and_stores_its_rules_decrypted(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-fallback-salt-1234")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
stored: Final = _StoredRouterSettings(
|
||||
encrypt_json_strings({"redis_password": "redis-canary-password", "fallbacks": [TEAM_RULE]})
|
||||
)
|
||||
router: Final = _team_router()
|
||||
|
||||
await self._create(
|
||||
FallbackCreateRequest(model="gpt-5.4-mini", fallback_models=["team-fallback"]), stored, router
|
||||
)
|
||||
|
||||
assert stored.row is not None
|
||||
assert "redis-canary-password" not in json.dumps(stored.row.param_value)
|
||||
assert decrypt_json_strings(stored.row.param_value) == {
|
||||
"redis_password": "redis-canary-password",
|
||||
"fallbacks": [TEAM_RULE, GATEWAY_RULE],
|
||||
}
|
||||
assert router.settings().fallbacks == [TEAM_RULE, GATEWAY_RULE]
|
||||
|
||||
async def test_a_second_create_keeps_the_first_rule(self) -> None:
|
||||
stored: Final = _StoredRouterSettings(None)
|
||||
|
||||
|
|
|
|||
|
|
@ -16538,6 +16538,59 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is
|
|||
]
|
||||
|
||||
|
||||
ROUTER_SETTINGS_SALT_KEY = "sk-router-settings-salt-1234"
|
||||
|
||||
|
||||
def test_update_config_router_settings_encrypted_before_write(_update_config_setup, monkeypatch):
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_if_encrypted_with,
|
||||
encrypt_json_strings,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", ROUTER_SETTINGS_SALT_KEY)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
client, prisma, restore = _update_config_setup(
|
||||
initial_rows={"router_settings": encrypt_json_strings({"redis_host": "redis.internal"})}
|
||||
)
|
||||
try:
|
||||
resp = client.post("/config/update", json={"router_settings": {"redis_password": "redis-canary-password"}})
|
||||
assert resp.status_code == 200, resp.text
|
||||
stored = prisma.db.litellm_config.rows["router_settings"]
|
||||
assert "redis-canary-password" not in json.dumps(stored)
|
||||
assert "redis.internal" not in json.dumps(stored)
|
||||
assert decrypt_if_encrypted_with(stored["redis_password"], ROUTER_SETTINGS_SALT_KEY) == "redis-canary-password"
|
||||
assert decrypt_if_encrypted_with(stored["redis_host"], ROUTER_SETTINGS_SALT_KEY) == "redis.internal"
|
||||
finally:
|
||||
restore()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_router_settings_from_db_config_decrypts_the_stored_row(monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_json_strings
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", ROUTER_SETTINGS_SALT_KEY)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
proxy_config = ProxyConfig()
|
||||
mock_router = route_entry_points_to_stubs(MagicMock())
|
||||
mock_router.update_settings = MagicMock()
|
||||
stored_row = MagicMock()
|
||||
stored_row.param_value = encrypt_json_strings({"redis_password": "redis-canary-password", "num_retries": 3})
|
||||
assert stored_row.param_value["redis_password"] != "redis-canary-password"
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=stored_row)
|
||||
proxy_config.router_settings.load_yaml({"routing_strategy": "simple-shuffle"})
|
||||
|
||||
await proxy_config._add_router_settings_from_db_config(llm_router=mock_router, prisma_client=mock_prisma_client)
|
||||
|
||||
combined_settings = mock_router.update_settings.call_args[1]
|
||||
assert combined_settings["redis_password"] == "redis-canary-password"
|
||||
assert combined_settings["num_retries"] == 3
|
||||
assert combined_settings["routing_strategy"] == "simple-shuffle"
|
||||
|
||||
|
||||
def test_feedback_box_is_one_json_record_in_json_mode(
|
||||
monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -431,11 +431,8 @@ class TestModelRepository:
|
|||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda v, **kw: f"encrypted_{v}",
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_create_model_encrypts_params(self, mock_decrypt, mock_encrypt, repo):
|
||||
async def test_create_model_encrypts_params(self, mock_encrypt):
|
||||
repo = ModelRepository(MockPrismaClient(), encryption_key="unit-test-key")
|
||||
model = await repo.create_model(
|
||||
model_name="gpt-4",
|
||||
litellm_params={"api_key": "sk-secret"},
|
||||
|
|
@ -449,11 +446,7 @@ class TestModelRepository:
|
|||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda v, **kw: f"encrypted_{v}",
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_create_model_all_fields(self, mock_decrypt, mock_encrypt, repo):
|
||||
async def test_create_model_all_fields(self, mock_encrypt, repo):
|
||||
model = await repo.create_model(
|
||||
model_name="gpt-4-turbo",
|
||||
litellm_params={
|
||||
|
|
@ -473,11 +466,7 @@ class TestModelRepository:
|
|||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda v, **kw: f"encrypted_{v}",
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_update_model_all_fields(self, mock_decrypt, mock_encrypt, repo):
|
||||
async def test_update_model_all_fields(self, mock_encrypt, repo):
|
||||
repo._prisma_client.db.litellm_proxymodeltable._records["model-full"] = {
|
||||
"model_id": "model-full",
|
||||
"model_name": "old-name",
|
||||
|
|
@ -495,11 +484,7 @@ class TestModelRepository:
|
|||
assert updated.model_name == "new-name"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_find_all(self, mock_decrypt, repo):
|
||||
async def test_find_all(self, repo):
|
||||
repo._prisma_client.db.litellm_proxymodeltable._records = {
|
||||
"m1": {
|
||||
"model_id": "m1",
|
||||
|
|
@ -518,11 +503,7 @@ class TestModelRepository:
|
|||
assert len(models) == 2
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_find_unblocked(self, mock_decrypt, repo):
|
||||
async def test_find_unblocked(self, repo):
|
||||
repo._prisma_client.db.litellm_proxymodeltable._records = {
|
||||
"m1": {
|
||||
"model_id": "m1",
|
||||
|
|
@ -535,11 +516,7 @@ class TestModelRepository:
|
|||
assert len(models) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_find_by_name(self, mock_decrypt, repo):
|
||||
async def test_find_by_name(self, repo):
|
||||
repo._prisma_client.db.litellm_proxymodeltable._records = {
|
||||
"m1": {
|
||||
"model_id": "m1",
|
||||
|
|
@ -555,11 +532,7 @@ class TestModelRepository:
|
|||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_update_model(self, mock_decrypt, mock_encrypt, repo):
|
||||
async def test_update_model(self, mock_encrypt, repo):
|
||||
repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = {
|
||||
"model_id": "m1",
|
||||
"model_name": "gpt-4",
|
||||
|
|
@ -574,11 +547,7 @@ class TestModelRepository:
|
|||
assert updated.blocked is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_delete_model(self, mock_decrypt, repo):
|
||||
async def test_delete_model(self, repo):
|
||||
repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = {
|
||||
"model_id": "m1",
|
||||
"model_name": "gpt-4",
|
||||
|
|
@ -592,11 +561,7 @@ class TestModelRepository:
|
|||
"litellm.proxy.common_utils.encrypt_decrypt_utils.encrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda v, **kw: v,
|
||||
)
|
||||
async def test_block_unblock_model(self, mock_decrypt, mock_encrypt, repo):
|
||||
async def test_block_unblock_model(self, mock_encrypt, repo):
|
||||
repo._prisma_client.db.litellm_proxymodeltable._records["m1"] = {
|
||||
"model_id": "m1",
|
||||
"model_name": "gpt-4",
|
||||
|
|
@ -2056,11 +2021,7 @@ class TestModelRepositoryExtended:
|
|||
return ModelRepository(client)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"litellm.proxy.common_utils.encrypt_decrypt_utils.decrypt_value_helper",
|
||||
side_effect=lambda value, **kw: value,
|
||||
)
|
||||
async def test_find_by_team_id(self, mock_decrypt, repo):
|
||||
async def test_find_by_team_id(self, repo):
|
||||
repo._prisma_client.db.litellm_proxymodeltable._records["model-1"] = {
|
||||
"model_id": "model-1",
|
||||
"model_name": "gpt-4",
|
||||
|
|
@ -2485,6 +2446,91 @@ class TestAutoRouterSessionRepository:
|
|||
_ = AutoRouterSessionRepository(None).table
|
||||
|
||||
|
||||
class TestModelRepositoryNestedEncryption:
|
||||
SALT_KEY: Final = "sk-repository-salt-1234"
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _salt_key(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", self.SALT_KEY)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_model_encrypts_nested_strings_and_find_by_id_decrypts_them(self):
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with
|
||||
|
||||
client: Final = MockPrismaClient()
|
||||
repo: Final = ModelRepository(client, encryption_key=self.SALT_KEY)
|
||||
litellm_params: Final = {
|
||||
"model": "openai/gpt-5.4-mini",
|
||||
"api_key": "sk-placeholder",
|
||||
"extra_headers": {"Authorization": "Bearer gateway-secret", "X-Gateway-Token": "gw-token"},
|
||||
"rpm": 10,
|
||||
}
|
||||
|
||||
created: Final = await repo.create_model(
|
||||
model_name="gateway-model", litellm_params=litellm_params, created_by="admin", model_id="gw-1"
|
||||
)
|
||||
|
||||
stored_text: Final = client.db.litellm_proxymodeltable._records["gw-1"]["litellm_params"]
|
||||
assert "gateway-secret" not in stored_text
|
||||
assert "gw-token" not in stored_text
|
||||
stored: Final = json.loads(stored_text)
|
||||
assert decrypt_if_encrypted_with(stored["extra_headers"]["Authorization"], self.SALT_KEY) == "Bearer gateway-secret"
|
||||
assert decrypt_if_encrypted_with(stored["api_key"], self.SALT_KEY) == "sk-placeholder"
|
||||
assert stored["rpm"] == 10
|
||||
assert created is not None
|
||||
assert created.litellm_params == litellm_params
|
||||
found: Final = await repo.find_by_id("gw-1")
|
||||
assert found is not None
|
||||
assert found.litellm_params == litellm_params
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_model_encrypts_classifier_credentials_and_keeps_classifier_fields_readable_in_sql(self):
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with
|
||||
|
||||
client: Final = MockPrismaClient()
|
||||
repo: Final = ModelRepository(client, encryption_key=self.SALT_KEY)
|
||||
complexity_router_config: Final = {
|
||||
"classifier_type": "jev",
|
||||
"jev_classifier_config": {
|
||||
"api_base": "https://classifier.example",
|
||||
"api_key": "sk-classifier-secret",
|
||||
"instructions": "Rate the question",
|
||||
},
|
||||
}
|
||||
litellm_params: Final = {"model": "auto_router/tuned", "complexity_router_config": complexity_router_config}
|
||||
|
||||
await repo.create_model(model_name="auto-router", litellm_params=litellm_params, created_by="admin", model_id="ar-1")
|
||||
|
||||
stored_text: Final = client.db.litellm_proxymodeltable._records["ar-1"]["litellm_params"]
|
||||
assert "sk-classifier-secret" not in stored_text
|
||||
assert "classifier.example" not in stored_text
|
||||
stored_config: Final = json.loads(stored_text)["complexity_router_config"]
|
||||
assert stored_config["classifier_type"] == "jev"
|
||||
assert stored_config["jev_classifier_config"]["instructions"] == "Rate the question"
|
||||
stored_jev: Final = stored_config["jev_classifier_config"]
|
||||
assert decrypt_if_encrypted_with(stored_jev["api_base"], self.SALT_KEY) == "https://classifier.example"
|
||||
assert decrypt_if_encrypted_with(stored_jev["api_key"], self.SALT_KEY) == "sk-classifier-secret"
|
||||
found: Final = await repo.find_by_id("ar-1")
|
||||
assert found is not None
|
||||
assert found.litellm_params == litellm_params
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reads_use_the_same_injected_key_as_writes(self):
|
||||
client: Final = MockPrismaClient()
|
||||
repo: Final = ModelRepository(client, encryption_key="sk-injected-key-5678")
|
||||
litellm_params: Final = {"model": "openai/gpt-5.4-mini", "extra_headers": {"Authorization": "Bearer gw"}}
|
||||
|
||||
await repo.create_model(model_name="gateway-model", litellm_params=litellm_params, created_by="admin", model_id="gw-2")
|
||||
|
||||
found: Final = await repo.find_by_id("gw-2")
|
||||
assert found is not None
|
||||
assert found.litellm_params == litellm_params
|
||||
env_key_read: Final = await ModelRepository(client).find_by_id("gw-2")
|
||||
assert env_key_read is not None
|
||||
assert env_key_read.litellm_params["extra_headers"]["Authorization"] != "Bearer gw"
|
||||
|
||||
|
||||
class TestSubscriptionAccountRepository:
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_fee_writes_only_the_named_columns(self):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue