mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
Merge pull request #41748 from BerriAI/litellm_wt_0917_7e12
fix(proxy): persist only the keys a caller changed in save_config
This commit is contained in:
yuneng-jiang
GitHub
commit
fc49ad7042
6 changed files with 612 additions and 89 deletions
17
litellm/proxy/config_resolvers/changed_section_keys.py
Normal file
17
litellm/proxy/config_resolvers/changed_section_keys.py
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
||||
def changed_section_keys(
|
||||
baseline: Mapping[str, JsonValue], new: Mapping[str, JsonValue]
|
||||
) -> tuple[Mapping[str, JsonValue], frozenset[str]]:
|
||||
changed: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{key: value for key, value in new.items() if key not in baseline or baseline[key] != value}
|
||||
)
|
||||
removed: Final = frozenset(baseline).difference(new)
|
||||
return changed, removed
|
||||
|
|
@ -27,6 +27,7 @@ from collections.abc import (
|
|||
Sequence,
|
||||
)
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from itertools import chain
|
||||
from types import MappingProxyType, UnionType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -436,6 +437,7 @@ from litellm.proxy.config_resolvers.alerting import (
|
|||
MS_TEAMS_DESCRIPTORS,
|
||||
SLACK_DESCRIPTORS,
|
||||
)
|
||||
from litellm.proxy.config_resolvers.changed_section_keys import changed_section_keys
|
||||
from litellm.proxy.container_endpoints.endpoints import router as container_router
|
||||
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
|
|
@ -4757,13 +4759,56 @@ def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:
|
|||
return any(str(obj) == object_type_str for obj in supported_db_objects)
|
||||
|
||||
|
||||
_CONFIG_PERSISTED_SECTIONS: Final = ("general_settings", "router_settings", "litellm_settings")
|
||||
_CONFIG_UNMANAGED_EXCLUSIONS: Final = frozenset(("environment_variables", "model_list"))
|
||||
_CONFIG_SECTION_VALUES: Final = TypeAdapter(Mapping[str, JsonValue])
|
||||
_CONFIG_SECTION_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtext($1))"
|
||||
|
||||
|
||||
class _ConfigParamWhere(TypedDict):
|
||||
param_name: ReadOnly[str]
|
||||
|
||||
|
||||
class _ConfigParamCreate(TypedDict):
|
||||
param_name: ReadOnly[str]
|
||||
param_value: ReadOnly[str]
|
||||
|
||||
|
||||
class _ConfigParamUpdate(TypedDict):
|
||||
param_value: ReadOnly[str]
|
||||
|
||||
|
||||
class _ConfigParamUpsert(TypedDict):
|
||||
create: ReadOnly[_ConfigParamCreate]
|
||||
update: ReadOnly[_ConfigParamUpdate]
|
||||
|
||||
|
||||
class _EnvironmentVariablesConfigData(TypedDict):
|
||||
environment_variables: ReadOnly[object]
|
||||
|
||||
|
||||
class _ConfigWithBaseline(dict[str, object]):
|
||||
def __init__(self, config: Mapping[str, object]) -> None:
|
||||
super().__init__(config)
|
||||
self._baseline: Mapping[str, object] = MappingProxyType(
|
||||
{key: copy.deepcopy(value) for key, value in config.items()}
|
||||
)
|
||||
|
||||
@property
|
||||
def baseline(self) -> Mapping[str, object]:
|
||||
return self._baseline
|
||||
|
||||
def update_baseline(self, config: Mapping[str, object]) -> None:
|
||||
self._baseline = MappingProxyType({key: copy.deepcopy(value) for key, value in config.items()})
|
||||
|
||||
|
||||
class ProxyConfig:
|
||||
"""
|
||||
Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.config: dict[str, Any] = {}
|
||||
self.config: Mapping[str, object] = MappingProxyType({})
|
||||
self._last_semantic_filter_config: dict[str, object] | None = None
|
||||
self._last_hashicorp_vault_config: dict[str, object] | None = None
|
||||
self._last_cyberark_config: dict[str, object] | None = None # mutable-ok: change-detection cache
|
||||
|
|
@ -4870,50 +4915,130 @@ class ProxyConfig:
|
|||
|
||||
return await resolve_includes(config=config, location=config_file_path, resolve=resolve, read=read_included)
|
||||
|
||||
async def save_config(self, new_config: dict, include_env_vars: bool = False):
|
||||
async def save_config(self, new_config: Mapping[str, object], include_env_vars: bool = False) -> None:
|
||||
global prisma_client, general_settings, user_config_file_path, store_model_in_db
|
||||
# Load existing config
|
||||
## DB - writes valid config to db
|
||||
"""
|
||||
- Do not write restricted params like 'api_key' to the database
|
||||
- if api_key is passed, save that to the local environment or connected secret manage (maybe expose `litellm.save_secret()`)
|
||||
"""
|
||||
|
||||
if prisma_client is not None and (
|
||||
general_settings.get("store_model_in_db", False) is True or store_model_in_db
|
||||
):
|
||||
# if using - db for config - models are in ModelTable
|
||||
|
||||
# Make a copy to avoid mutating the original config
|
||||
config_to_save: Final = new_config.copy()
|
||||
|
||||
# environment_variables are persisted to the DB only when a caller
|
||||
# explicitly opts in. Most callers reach save_config after
|
||||
# get_config() merged YAML + OS env into new_config (with
|
||||
# os.environ/ placeholders already resolved to plaintext), so
|
||||
# persisting them here would snapshot file/container env vars into
|
||||
# a config row that then shadows those sources on every restart.
|
||||
# The dedicated /config/update path writes env vars directly, so
|
||||
# no current caller needs include_env_vars=True.
|
||||
if not include_env_vars:
|
||||
config_to_save.pop("environment_variables", None)
|
||||
|
||||
# SECURITY: Always encrypt environment_variables before DB write.
|
||||
# _encrypt_env_variables_for_db is idempotent — a caller that
|
||||
# already encrypted the values (or re-submitted ciphertext read
|
||||
# back from the DB) will not get a stacked second layer.
|
||||
if "environment_variables" in config_to_save and config_to_save["environment_variables"]:
|
||||
config_to_save["environment_variables"] = self._encrypt_env_variables_for_db(
|
||||
environment_variables=config_to_save["environment_variables"]
|
||||
baseline: Final[Mapping[str, object]] = (
|
||||
new_config.baseline if isinstance(new_config, _ConfigWithBaseline) else self.get_config_state()
|
||||
)
|
||||
for section_name in _CONFIG_PERSISTED_SECTIONS:
|
||||
await self._save_changed_config_section(
|
||||
section_name=section_name,
|
||||
baseline=baseline,
|
||||
new_config=new_config,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
config_to_save.pop("model_list", None)
|
||||
await prisma_client.insert_data(data=config_to_save, table_name="config")
|
||||
else:
|
||||
# Save the updated config - if user is not using a dB
|
||||
## YAML
|
||||
with open(f"{user_config_file_path}", "w") as config_file:
|
||||
yaml.dump(new_config, config_file, default_flow_style=False)
|
||||
unmanaged_config: Final[Mapping[str, object]] = MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in new_config.items()
|
||||
if key not in _CONFIG_PERSISTED_SECTIONS
|
||||
and key not in _CONFIG_UNMANAGED_EXCLUSIONS
|
||||
and (key not in baseline or baseline[key] != value)
|
||||
}
|
||||
)
|
||||
if unmanaged_config:
|
||||
await prisma_client.insert_data(data=unmanaged_config, table_name="config")
|
||||
|
||||
environment_variables: Final = new_config.get("environment_variables")
|
||||
if include_env_vars and environment_variables is not None:
|
||||
encrypted_environment_variables: Final = (
|
||||
self._encrypt_env_variables_for_db(environment_variables=environment_variables)
|
||||
if isinstance(environment_variables, dict) and environment_variables
|
||||
else environment_variables
|
||||
)
|
||||
environment_variables_data: Final[_EnvironmentVariablesConfigData] = {
|
||||
"environment_variables": encrypted_environment_variables
|
||||
}
|
||||
await prisma_client.insert_data(data=environment_variables_data, table_name="config")
|
||||
next_config: Final[Mapping[str, object]] = MappingProxyType({**baseline, **new_config})
|
||||
self.update_config_state(config=next_config)
|
||||
if isinstance(new_config, _ConfigWithBaseline):
|
||||
new_config.update_baseline(config=next_config)
|
||||
return
|
||||
|
||||
with open(f"{user_config_file_path}", "w") as config_file:
|
||||
yaml.dump(
|
||||
dict(new_config), config_file, default_flow_style=False
|
||||
) # mutable-ok: YAML must serialize a plain dict
|
||||
|
||||
async def _save_changed_config_section(
|
||||
self,
|
||||
*,
|
||||
section_name: str,
|
||||
baseline: Mapping[str, object],
|
||||
new_config: Mapping[str, object],
|
||||
prisma_client: PrismaClient,
|
||||
) -> None:
|
||||
if section_name not in new_config:
|
||||
return
|
||||
baseline_value: Final = baseline.get(section_name)
|
||||
new_value: Final = new_config[section_name]
|
||||
baseline_section: Final[Mapping[str, JsonValue]] = (
|
||||
_CONFIG_SECTION_VALUES.validate_python(baseline_value)
|
||||
if isinstance(baseline_value, Mapping)
|
||||
else MappingProxyType({})
|
||||
)
|
||||
new_section: Final[Mapping[str, JsonValue]] = (
|
||||
_CONFIG_SECTION_VALUES.validate_python(new_value)
|
||||
if isinstance(new_value, Mapping)
|
||||
else MappingProxyType({})
|
||||
)
|
||||
changed_keys, removed_keys = changed_section_keys(baseline_section, new_section)
|
||||
if not changed_keys and not removed_keys:
|
||||
return
|
||||
wrote_section: Final = await self._upsert_changed_config_section(
|
||||
section_name=section_name,
|
||||
changed_keys=changed_keys,
|
||||
removed_keys=removed_keys,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
if not wrote_section:
|
||||
return
|
||||
await invalidate_config_param(section_name)
|
||||
|
||||
async def _upsert_changed_config_section(
|
||||
self,
|
||||
*,
|
||||
section_name: str,
|
||||
changed_keys: Mapping[str, JsonValue],
|
||||
removed_keys: frozenset[str],
|
||||
prisma_client: PrismaClient,
|
||||
) -> bool:
|
||||
async with prisma_client.tx() as tx:
|
||||
await tx.query_raw(_CONFIG_SECTION_LOCK_SQL, section_name)
|
||||
config_table: Final = cast("TableActions[_ConfigParamRow]", tx.litellm_config)
|
||||
config_where: Final[_ConfigParamWhere] = {"param_name": section_name}
|
||||
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({})
|
||||
)
|
||||
merged_section: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in chain(
|
||||
((key, value) for key, value in existing_section.items() if key not in removed_keys),
|
||||
changed_keys.items(),
|
||||
)
|
||||
}
|
||||
)
|
||||
if merged_section == existing_section:
|
||||
return False
|
||||
serialized_section: Final = json.dumps(dict(merged_section)) # mutable-ok: JSON encoder requires a dict
|
||||
config_data: Final[_ConfigParamUpsert] = {
|
||||
"create": {"param_name": section_name, "param_value": serialized_section},
|
||||
"update": {"param_value": serialized_section},
|
||||
}
|
||||
await config_table.upsert(where=config_where, data=config_data)
|
||||
return True
|
||||
|
||||
async def save_environment_variables(self, updates: dict[str, str | None]) -> None:
|
||||
"""Persist specific environment variables to the DB config row.
|
||||
|
|
@ -5265,26 +5390,26 @@ class ProxyConfig:
|
|||
|
||||
self.update_config_state(config=config)
|
||||
|
||||
return config
|
||||
return _ConfigWithBaseline(config)
|
||||
|
||||
def update_config_state(self, config: dict):
|
||||
self.config = config
|
||||
def update_config_state(self, config: Mapping[str, object]) -> None:
|
||||
self.config = MappingProxyType({key: copy.deepcopy(value) for key, value in config.items()})
|
||||
|
||||
def get_config_state(self):
|
||||
def get_config_state(self) -> Mapping[str, object]:
|
||||
"""
|
||||
Returns a deep copy of the config,
|
||||
|
||||
Do this, to avoid mutating the config state outside of allowed methods
|
||||
"""
|
||||
try:
|
||||
return copy.deepcopy(self.config)
|
||||
return MappingProxyType({key: copy.deepcopy(value) for key, value in self.config.items()})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
"ProxyConfig:get_config_state(): Error returning copy of config state. self.config=%s\nError: %s",
|
||||
self.config,
|
||||
e,
|
||||
)
|
||||
return {}
|
||||
return MappingProxyType({})
|
||||
|
||||
def load_credential_list(self, config: dict) -> list[CredentialItem]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@
|
|||
- {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"}
|
||||
- {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"}
|
||||
- {id: mgmt.router_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "router_settings_endpoints.py", rationale: "Router config (smoke)"}
|
||||
- {id: mgmt.config.allowed_ip.changed_key_only, module: mgmt, tier: P2, surface: api, assertions: [persists], source: "proxy_setting_endpoints.py:496", rationale: "An allowed-IP change leaves unrelated file settings out of the DB row"}
|
||||
- {id: mgmt.jwt_key_mapping.new.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "jwt_key_mapping_endpoints.py", rationale: "JWT->key mapping (smoke)"}
|
||||
- {id: mgmt.compliance.gdpr.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "compliance_endpoints.py", rationale: "GDPR ops (smoke)"}
|
||||
- {id: mgmt.tool_management.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "tool_management_endpoints.py", rationale: "Tool inventory (smoke)"}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
general_settings:
|
||||
max_parallel_requests: 100
|
||||
proxy_batch_write_at: 5
|
||||
enable_jwt_auth: true
|
||||
litellm_jwtauth:
|
||||
|
|
|
|||
|
|
@ -21,12 +21,13 @@ from __future__ import annotations
|
|||
import math
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import NoBody, Success, unwrap, unwrap_status
|
||||
from e2e_http import NoBody, Success, UnknownApiError, unwrap, unwrap_status
|
||||
from lifecycle import ResourceManager
|
||||
from management_client import ManagementClient
|
||||
from models import KeyGenerateBody, LiteLLMParamsBody, TeamNewBody
|
||||
|
|
@ -198,6 +199,19 @@ class ConfigUpdateResponse(BaseModel):
|
|||
message: str
|
||||
|
||||
|
||||
class AllowedIpBody(BaseModel):
|
||||
ip: str
|
||||
|
||||
|
||||
class ConfigFieldInfoParams(BaseModel):
|
||||
field_name: str
|
||||
|
||||
|
||||
class ConfigFieldInfoResponse(BaseModel):
|
||||
field_name: str
|
||||
field_value: JsonValue
|
||||
|
||||
|
||||
class RouterCurrentValues(BaseModel):
|
||||
num_retries: int | None = None
|
||||
|
||||
|
|
@ -516,6 +530,45 @@ class TestRouterSettings:
|
|||
)
|
||||
|
||||
|
||||
class TestConfigPersistence:
|
||||
@pytest.mark.covers("mgmt.config.allowed_ip.changed_key_only")
|
||||
def test_add_allowed_ip_does_not_store_unrelated_config_value(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
allowed_ip: Final = "127.0.0.1"
|
||||
added: Final = unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/add/allowed_ip",
|
||||
headers=client.proxy.transport.master,
|
||||
json=AllowedIpBody(ip=allowed_ip),
|
||||
response_type=ConfigUpdateResponse,
|
||||
)
|
||||
)
|
||||
resources.defer(
|
||||
lambda: unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/delete/allowed_ip",
|
||||
headers=client.proxy.transport.master,
|
||||
json=AllowedIpBody(ip=allowed_ip),
|
||||
response_type=ConfigUpdateResponse,
|
||||
)
|
||||
)
|
||||
)
|
||||
assert added.message == f"IP {allowed_ip} address added successfully"
|
||||
|
||||
field_info: Final = client.proxy.transport.get(
|
||||
"/config/field/info",
|
||||
headers=client.proxy.transport.master,
|
||||
params=ConfigFieldInfoParams(field_name="max_parallel_requests"),
|
||||
response_type=ConfigFieldInfoResponse,
|
||||
)
|
||||
match field_info:
|
||||
case UnknownApiError(status_code=400, body=body):
|
||||
assert "not in DB" in body
|
||||
case _:
|
||||
pytest.fail(f"expected max_parallel_requests to remain absent from the DB row, got {field_info}")
|
||||
|
||||
|
||||
class TestMcpServerSubmission:
|
||||
@pytest.mark.covers("mgmt.mcp_server.register.happy_path")
|
||||
def test_register_submits_pending_server(self, client: ManagementClient, resources: ResourceManager) -> None:
|
||||
|
|
|
|||
|
|
@ -13,8 +13,11 @@ import json
|
|||
import logging
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Dict
|
||||
from typing import Any, Dict, Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -35,7 +38,7 @@ from litellm.proxy.proxy_server import (
|
|||
)
|
||||
|
||||
from .conftest import normalize
|
||||
from pydantic import ValidationError
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _is_remote_module_url
|
||||
|
|
@ -853,6 +856,314 @@ async def test_ProxyConfig__process_includes_terminates_on_a_cycle(tmp_path):
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_CONFIG_VALUE: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ConfigRow:
|
||||
param_value: dict[str, JsonValue] | str
|
||||
|
||||
|
||||
class _ConfigTable:
|
||||
def __init__(self, rows: Mapping[str, Mapping[str, JsonValue] | str]) -> None:
|
||||
self.rows = {
|
||||
param_name: value if isinstance(value, str) else _CONFIG_VALUE.validate_python(value)
|
||||
for param_name, value in rows.items()
|
||||
}
|
||||
self.upserted_param_names: list[str] = []
|
||||
self._section_lock = asyncio.Lock()
|
||||
|
||||
async def find_first(self, *, where: Mapping[str, str]) -> _ConfigRow | None:
|
||||
value: Final = self.rows.get(where["param_name"])
|
||||
await asyncio.sleep(0)
|
||||
return _ConfigRow(param_value=value) if value is not None else None
|
||||
|
||||
async def upsert(
|
||||
self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]
|
||||
) -> _ConfigRow:
|
||||
param_name: Final = where["param_name"]
|
||||
value: Final = _CONFIG_VALUE.validate_json(data["update"]["param_value"])
|
||||
self.rows[param_name] = value
|
||||
self.upserted_param_names.append(param_name)
|
||||
return _ConfigRow(param_value=value)
|
||||
|
||||
|
||||
class _ConfigTransaction:
|
||||
def __init__(self, table: _ConfigTable) -> None:
|
||||
self.litellm_config: Final = table
|
||||
self._section_lock: Final = table._section_lock
|
||||
self._locked = False
|
||||
|
||||
async def __aenter__(self) -> _ConfigTransaction:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_: object) -> None:
|
||||
if self._locked:
|
||||
self._section_lock.release()
|
||||
|
||||
async def query_raw(self, _: str, __: str) -> None:
|
||||
await self._section_lock.acquire()
|
||||
self._locked = True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ConfigDb:
|
||||
litellm_config: _ConfigTable
|
||||
|
||||
def tx(self) -> _ConfigTransaction:
|
||||
return _ConfigTransaction(self.litellm_config)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ConfigPrisma:
|
||||
db: _ConfigDb
|
||||
|
||||
def tx(self) -> _ConfigTransaction:
|
||||
return self.db.tx()
|
||||
|
||||
async def insert_data(self, *, data: Mapping[str, object], table_name: str) -> None:
|
||||
if table_name != "config":
|
||||
raise AssertionError(f"Expected config write, got {table_name}")
|
||||
for param_name, value in data.items():
|
||||
self.db.litellm_config.rows[param_name] = _CONFIG_VALUE.validate_python(value)
|
||||
self.db.litellm_config.upserted_param_names.append(param_name)
|
||||
|
||||
|
||||
def _db_backed_proxy_config(monkeypatch, rows: Mapping[str, Mapping[str, JsonValue]]) -> tuple[ProxyConfig, _ConfigTable]:
|
||||
table: Final = _ConfigTable(rows)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _ConfigPrisma(db=_ConfigDb(litellm_config=table)))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"store_model_in_db": True})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.invalidate_config_param", AsyncMock())
|
||||
return ProxyConfig(), table
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_merges_changed_keys_without_copying_file_settings(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"db_only": "stored"}})
|
||||
baseline: Final = {
|
||||
"model_list": [],
|
||||
"general_settings": {"max_parallel_requests": 5, "file_only": "yaml", "allowed_ips": []},
|
||||
"router_settings": {"num_retries": 1},
|
||||
"litellm_settings": {"drop_params": True},
|
||||
}
|
||||
proxy_config.update_config_state(config=baseline)
|
||||
changed: Final = {
|
||||
**baseline,
|
||||
"general_settings": {**baseline["general_settings"], "allowed_ips": ["127.0.0.1"]},
|
||||
}
|
||||
|
||||
await proxy_config.save_config(changed)
|
||||
|
||||
assert table.rows == {"general_settings": {"db_only": "stored", "allowed_ips": ["127.0.0.1"]}}
|
||||
assert table.upserted_param_names == ["general_settings"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_skips_unchanged_config(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"db_only": "stored"}})
|
||||
baseline: Final = {
|
||||
"model_list": [],
|
||||
"general_settings": {"max_parallel_requests": 5},
|
||||
"router_settings": {"num_retries": 1},
|
||||
"litellm_settings": {"drop_params": True},
|
||||
}
|
||||
proxy_config.update_config_state(config=baseline)
|
||||
|
||||
await proxy_config.save_config(baseline)
|
||||
|
||||
assert table.rows == {"general_settings": {"db_only": "stored"}}
|
||||
assert table.upserted_param_names == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_skips_unchanged_unmanaged_values(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
|
||||
baseline: Final = {"general_settings": {}, "guardrails": {"enabled": True}}
|
||||
proxy_config.update_config_state(config=baseline)
|
||||
|
||||
await proxy_config.save_config(baseline)
|
||||
|
||||
assert table.rows == {}
|
||||
assert table.upserted_param_names == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_leaves_omitted_sections_unchanged(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(
|
||||
monkeypatch,
|
||||
{"general_settings": {"allowed_ips": ["10.0.0.1"], "db_only": "stored"}},
|
||||
)
|
||||
proxy_config.update_config_state(
|
||||
config={"general_settings": {"allowed_ips": ["10.0.0.1"]}, "router_settings": {"num_retries": 1}}
|
||||
)
|
||||
|
||||
await proxy_config.save_config({"router_settings": {"num_retries": 2}})
|
||||
|
||||
assert table.rows == {
|
||||
"general_settings": {"allowed_ips": ["10.0.0.1"], "db_only": "stored"},
|
||||
"router_settings": {"num_retries": 2},
|
||||
}
|
||||
assert table.upserted_param_names == ["router_settings"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_decodes_a_serialized_config_row(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": '{"db_only":"stored"}'})
|
||||
proxy_config.update_config_state(config={"general_settings": {"allowed_ips": []}})
|
||||
|
||||
await proxy_config.save_config({"general_settings": {"allowed_ips": ["127.0.0.1"]}})
|
||||
|
||||
assert table.rows == {"general_settings": {"db_only": "stored", "allowed_ips": ["127.0.0.1"]}}
|
||||
assert table.upserted_param_names == ["general_settings"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_serializes_concurrent_changes_to_one_section(monkeypatch):
|
||||
first, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"a": 0, "b": 0}})
|
||||
second: Final = ProxyConfig()
|
||||
baseline: Final = {"general_settings": {"a": 0, "b": 0}}
|
||||
first.update_config_state(config=baseline)
|
||||
second.update_config_state(config=baseline)
|
||||
|
||||
await asyncio.gather(
|
||||
first.save_config({"general_settings": {"a": 1, "b": 0}}),
|
||||
second.save_config({"general_settings": {"a": 0, "b": 1}}),
|
||||
)
|
||||
|
||||
assert table.rows == {"general_settings": {"a": 1, "b": 1}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_updates_the_baseline_after_a_save(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
|
||||
proxy_config.update_config_state(config={"general_settings": {}})
|
||||
|
||||
await proxy_config.save_config({"general_settings": {"removed_key": True}})
|
||||
await proxy_config.save_config({"general_settings": {}})
|
||||
|
||||
assert table.rows == {"general_settings": {}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_keeps_omitted_sections_in_its_next_baseline(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"allowed_ips": ["10.0.0.1"]}})
|
||||
proxy_config.update_config_state(
|
||||
config={"general_settings": {"allowed_ips": ["10.0.0.1"]}, "router_settings": {"num_retries": 1}}
|
||||
)
|
||||
|
||||
await proxy_config.save_config({"router_settings": {"num_retries": 2}})
|
||||
await proxy_config.save_config({"general_settings": {}})
|
||||
|
||||
assert table.rows == {"general_settings": {}, "router_settings": {"num_retries": 2}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_uses_the_baseline_from_the_loaded_config(tmp_path, monkeypatch):
|
||||
config_file: Final = tmp_path / "config.yaml"
|
||||
config_file.write_text("general_settings:\n yaml_only: true\n")
|
||||
proxy_config: Final = ProxyConfig()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
first: Final = await proxy_config.get_config(config_file_path=str(config_file))
|
||||
second: Final = await proxy_config.get_config(config_file_path=str(config_file))
|
||||
table: Final = _ConfigTable({})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _ConfigPrisma(db=_ConfigDb(litellm_config=table)))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"store_model_in_db": True})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.invalidate_config_param", AsyncMock())
|
||||
first["general_settings"]["first"] = True
|
||||
second["general_settings"]["second"] = True
|
||||
|
||||
await proxy_config.save_config(second)
|
||||
await proxy_config.save_config(first)
|
||||
|
||||
assert table.rows == {"general_settings": {"second": True, "first": True}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_accepts_non_json_model_metadata(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
|
||||
proxy_config.update_config_state(config={"general_settings": {"allowed_ips": []}})
|
||||
config: Final = {
|
||||
"model_list": [{"model_name": "date-model", "model_info": {"created_at": datetime(2026, 1, 1)}}],
|
||||
"general_settings": {"allowed_ips": ["127.0.0.1"]},
|
||||
}
|
||||
|
||||
await proxy_config.save_config(config)
|
||||
|
||||
assert table.rows == {"general_settings": {"allowed_ips": ["127.0.0.1"]}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_writes_only_changed_router_settings(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"router_settings": {"db_only": "stored"}})
|
||||
baseline: Final = {
|
||||
"model_list": [],
|
||||
"general_settings": {"max_parallel_requests": 5},
|
||||
"router_settings": {"num_retries": 1},
|
||||
"litellm_settings": {"drop_params": True},
|
||||
}
|
||||
proxy_config.update_config_state(config=baseline)
|
||||
changed: Final = {**baseline, "router_settings": {"num_retries": 2}}
|
||||
|
||||
await proxy_config.save_config(changed)
|
||||
|
||||
assert table.rows == {"router_settings": {"db_only": "stored", "num_retries": 2}}
|
||||
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(
|
||||
monkeypatch, {"general_settings": {"removed_key": "db", "db_only": "stored"}}
|
||||
)
|
||||
baseline: Final = {"general_settings": {"removed_key": "yaml", "file_only": "yaml"}}
|
||||
proxy_config.update_config_state(config=baseline)
|
||||
changed: Final = {"general_settings": {"file_only": "yaml"}}
|
||||
|
||||
await proxy_config.save_config(changed)
|
||||
|
||||
assert table.rows == {"general_settings": {"db_only": "stored"}}
|
||||
assert table.upserted_param_names == ["general_settings"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_keeps_an_unstored_removed_key_as_a_noop(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"db_only": "stored"}})
|
||||
baseline: Final = {"general_settings": {"file_only": "yaml"}}
|
||||
proxy_config.update_config_state(config=baseline)
|
||||
|
||||
await proxy_config.save_config({"general_settings": {}})
|
||||
|
||||
assert table.rows == {"general_settings": {"db_only": "stored"}}
|
||||
assert table.upserted_param_names == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_get_config_keeps_state_separate_from_returned_config(tmp_path, monkeypatch):
|
||||
config_file: Final = tmp_path / "config.yaml"
|
||||
config_file.write_text("general_settings:\n max_parallel_requests: 5\n")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file))
|
||||
|
||||
proxy_config: Final = ProxyConfig()
|
||||
loaded: Final = await proxy_config.get_config(config_file_path=str(config_file))
|
||||
loaded["general_settings"]["max_parallel_requests"] = 6
|
||||
|
||||
assert proxy_config.get_config_state()["general_settings"]["max_parallel_requests"] == 5
|
||||
|
||||
|
||||
def test_ProxyConfig_update_config_state_keeps_a_copy_of_its_input():
|
||||
source: Final = {"general_settings": {"max_parallel_requests": 5}}
|
||||
proxy_config: Final = ProxyConfig()
|
||||
proxy_config.update_config_state(config=source)
|
||||
source["general_settings"]["max_parallel_requests"] = 6
|
||||
|
||||
assert proxy_config.get_config_state()["general_settings"]["max_parallel_requests"] == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_writes_yaml_when_no_db(tmp_path, monkeypatch):
|
||||
target = tmp_path / "out.yaml"
|
||||
|
|
@ -869,6 +1180,25 @@ async def test_ProxyConfig_save_config_writes_yaml_when_no_db(tmp_path, monkeypa
|
|||
assert loaded == cfg
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_writes_a_loadable_yaml_for_a_loaded_config(tmp_path, monkeypatch):
|
||||
config_file: Final = tmp_path / "config.yaml"
|
||||
config_file.write_text("general_settings:\n max_parallel_requests: 5\n")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
proxy_config: Final = ProxyConfig()
|
||||
loaded_config: Final = await proxy_config.get_config(config_file_path=str(config_file))
|
||||
loaded_config["general_settings"]["max_parallel_requests"] = 6
|
||||
|
||||
await proxy_config.save_config(loaded_config)
|
||||
|
||||
import yaml as _yaml
|
||||
|
||||
assert _yaml.safe_load(config_file.read_text()) == {"general_settings": {"max_parallel_requests": 6}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_invalid_path_raises(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
|
|
@ -885,58 +1215,54 @@ async def test_ProxyConfig_save_config_invalid_path_raises(monkeypatch):
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_db_omits_environment_variables_by_default(monkeypatch):
|
||||
"""A save_config after get_config() (which resolves os.environ/ placeholders
|
||||
to plaintext and merges the environment_variables section) must not snapshot
|
||||
those env vars into the DB config row. Persisting them would make a stale DB
|
||||
row shadow YAML/container env on every subsequent restart."""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.insert_data = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
# a valid salt so the env-var encryption path (reached only if the pop
|
||||
# regresses) runs cleanly, making this fail on the assertion below rather
|
||||
# than on an incidental encryption crash
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-key")
|
||||
|
||||
pc = ProxyConfig()
|
||||
cfg = {
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
|
||||
baseline: Final = {"model_list": [], "litellm_settings": {}}
|
||||
proxy_config.update_config_state(config=baseline)
|
||||
config: Final = {
|
||||
"model_list": [{"model_name": "gpt-4o"}],
|
||||
"litellm_settings": {"success_callback": ["langfuse"]},
|
||||
"environment_variables": {"OPENAI_API_KEY": "sk-from-yaml"},
|
||||
}
|
||||
await pc.save_config(cfg)
|
||||
|
||||
mock_prisma.insert_data.assert_awaited_once()
|
||||
written = mock_prisma.insert_data.await_args.kwargs["data"]
|
||||
assert "environment_variables" not in written
|
||||
# unrelated sections are still persisted; model_list is stripped as before
|
||||
assert written["litellm_settings"] == {"success_callback": ["langfuse"]}
|
||||
assert "model_list" not in written
|
||||
# the caller's dict is not mutated (save_config works on a copy)
|
||||
assert cfg["environment_variables"] == {"OPENAI_API_KEY": "sk-from-yaml"}
|
||||
await proxy_config.save_config(config)
|
||||
|
||||
assert table.rows == {"litellm_settings": {"success_callback": ["langfuse"]}}
|
||||
assert table.upserted_param_names == ["litellm_settings"]
|
||||
assert config["environment_variables"] == {"OPENAI_API_KEY": "sk-from-yaml"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_db_persists_environment_variables_when_opted_in(monkeypatch):
|
||||
"""The explicit opt-in path (include_env_vars=True) still persists env vars,
|
||||
encrypted, so the dedicated config-update flow can write them."""
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.insert_data = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
|
||||
proxy_config.update_config_state(config={"litellm_settings": {}})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-key")
|
||||
config: Final = {"litellm_settings": {}, "environment_variables": {"OPENAI_API_KEY": "sk-explicit"}}
|
||||
|
||||
await proxy_config.save_config(config, include_env_vars=True)
|
||||
|
||||
assert set(table.rows["environment_variables"]) == {"OPENAI_API_KEY"}
|
||||
assert table.rows["environment_variables"]["OPENAI_API_KEY"] != "sk-explicit"
|
||||
assert table.upserted_param_names == ["environment_variables"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_save_config_persists_unchanged_environment_variables_when_opted_in(monkeypatch):
|
||||
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
|
||||
config: Final = {
|
||||
"litellm_settings": {},
|
||||
"environment_variables": {"OPENAI_API_KEY": "sk-explicit"},
|
||||
}
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-key")
|
||||
|
||||
pc = ProxyConfig()
|
||||
cfg = {"litellm_settings": {}, "environment_variables": {"OPENAI_API_KEY": "sk-explicit"}}
|
||||
await pc.save_config(cfg, include_env_vars=True)
|
||||
await proxy_config.save_config(config)
|
||||
|
||||
mock_prisma.insert_data.assert_awaited_once()
|
||||
written = mock_prisma.insert_data.await_args.kwargs["data"]
|
||||
assert set(written["environment_variables"].keys()) == {"OPENAI_API_KEY"}
|
||||
# value is encrypted at rest, not the plaintext it came in as
|
||||
assert written["environment_variables"]["OPENAI_API_KEY"] != "sk-explicit"
|
||||
assert table.rows == {}
|
||||
assert table.upserted_param_names == []
|
||||
|
||||
await proxy_config.save_config(config, include_env_vars=True)
|
||||
|
||||
assert set(table.rows["environment_variables"]) == {"OPENAI_API_KEY"}
|
||||
assert table.upserted_param_names == ["environment_variables"]
|
||||
|
||||
|
||||
def _install_fake_config_repo(monkeypatch, existing_row):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue