mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
refactor(proxy): extract _add_config_models helper and sanitize error logging
This commit is contained in:
Ayush Arora
parent
28969ffe07
commit
94fa04c1cd
2 changed files with 73 additions and 65 deletions
|
|
@ -4818,7 +4818,7 @@ class ProxyConfig:
|
|||
# _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"]:
|
||||
if config_to_save.get("environment_variables"):
|
||||
config_to_save["environment_variables"] = self._encrypt_env_variables_for_db(
|
||||
environment_variables=config_to_save["environment_variables"]
|
||||
)
|
||||
|
|
@ -6423,6 +6423,63 @@ class ProxyConfig:
|
|||
return get_secret(decrypted_value)
|
||||
return decrypted_value
|
||||
|
||||
def _add_config_models(self, config_models: list | None = None) -> int:
|
||||
if config_models is None:
|
||||
config_state = self.get_config_state()
|
||||
if isinstance(config_state, dict):
|
||||
config_models = config_state.get("model_list", None)
|
||||
if not config_models and user_config_file_path and os.path.exists(user_config_file_path):
|
||||
try:
|
||||
with open(user_config_file_path, "r") as f:
|
||||
yaml_cfg = yaml.safe_load(f)
|
||||
if isinstance(yaml_cfg, dict):
|
||||
config_models = yaml_cfg.get("model_list", None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if not config_models:
|
||||
return 0
|
||||
|
||||
added_count = 0
|
||||
for model in config_models:
|
||||
try:
|
||||
raw_litellm_params = copy.deepcopy(model.get("litellm_params", {}))
|
||||
for k, v in raw_litellm_params.items():
|
||||
if isinstance(v, str) and v.startswith("os.environ/"):
|
||||
raw_litellm_params[k] = get_secret(v)
|
||||
|
||||
model_info_dict = copy.deepcopy(model.get("model_info", {}))
|
||||
model_id = model_info_dict.get("id", None)
|
||||
if model_id is None:
|
||||
model_id = llm_router.generate_model_id(
|
||||
model_group=model["model_name"],
|
||||
litellm_params=raw_litellm_params,
|
||||
)
|
||||
else:
|
||||
model_id = str(model_id)
|
||||
model_info_dict["id"] = model_id
|
||||
model_info_dict["db_model"] = False
|
||||
|
||||
_model_info = RouterModelInfo(**model_info_dict)
|
||||
_litellm_params = LiteLLM_Params.model_validate(raw_litellm_params)
|
||||
|
||||
added = llm_router.upsert_deployment(
|
||||
deployment=Deployment(
|
||||
model_name=model["model_name"],
|
||||
litellm_params=_litellm_params,
|
||||
model_info=_model_info,
|
||||
)
|
||||
)
|
||||
if added is not None:
|
||||
added_count += 1
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"Error adding config model to llm_router: %s. model_name=%s",
|
||||
e,
|
||||
model.get("model_name"),
|
||||
)
|
||||
return added_count
|
||||
|
||||
def _add_deployment(self, db_models: list, config_models: list | None = None) -> int:
|
||||
"""
|
||||
Iterate through db models and config models
|
||||
|
|
@ -6465,55 +6522,7 @@ class ProxyConfig:
|
|||
added_models += 1
|
||||
|
||||
## ADD CONFIG MODEL LOGIC
|
||||
if config_models is None:
|
||||
config_state = self.get_config_state()
|
||||
if isinstance(config_state, dict):
|
||||
config_models = config_state.get("model_list", None)
|
||||
if not config_models and user_config_file_path and os.path.exists(user_config_file_path):
|
||||
try:
|
||||
with open(user_config_file_path, "r") as f:
|
||||
yaml_cfg = yaml.safe_load(f)
|
||||
if isinstance(yaml_cfg, dict):
|
||||
config_models = yaml_cfg.get("model_list", None)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if config_models:
|
||||
for model in config_models:
|
||||
try:
|
||||
raw_litellm_params = copy.deepcopy(model.get("litellm_params", {}))
|
||||
for k, v in raw_litellm_params.items():
|
||||
if isinstance(v, str) and v.startswith("os.environ/"):
|
||||
raw_litellm_params[k] = get_secret(v)
|
||||
|
||||
model_info_dict = copy.deepcopy(model.get("model_info", {}))
|
||||
model_id = model_info_dict.get("id", None)
|
||||
if model_id is None:
|
||||
model_id = llm_router.generate_model_id(
|
||||
model_group=model["model_name"],
|
||||
litellm_params=raw_litellm_params,
|
||||
)
|
||||
else:
|
||||
model_id = str(model_id)
|
||||
model_info_dict["id"] = model_id
|
||||
model_info_dict["db_model"] = False
|
||||
|
||||
_model_info = RouterModelInfo(**model_info_dict)
|
||||
_litellm_params = LiteLLM_Params.model_validate(raw_litellm_params)
|
||||
|
||||
added = llm_router.upsert_deployment(
|
||||
deployment=Deployment(
|
||||
model_name=model["model_name"],
|
||||
litellm_params=_litellm_params,
|
||||
model_info=_model_info,
|
||||
)
|
||||
)
|
||||
if added is not None:
|
||||
added_models += 1
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"Error adding config model to llm_router: %s. model=%s", e, model
|
||||
)
|
||||
added_models += self._add_config_models(config_models=config_models)
|
||||
|
||||
return added_models
|
||||
|
||||
|
|
@ -11163,11 +11172,11 @@ async def completion(
|
|||
if _data.get("stream", None) is not None and _data["stream"] is True:
|
||||
_text_response: Final = litellm.ModelResponse()
|
||||
# Set text attribute dynamically for text completion format
|
||||
setattr(_text_response.choices[0], "text", e.message)
|
||||
_text_response.choices[0].text = e.message
|
||||
_text_response.model = e.model
|
||||
_usage = _blocked_response_usage(e.original_response)
|
||||
# Set usage attribute dynamically (ModelResponse accepts usage in __init__ but it's not in type definition)
|
||||
setattr(_text_response, "usage", _usage)
|
||||
_text_response.usage = _usage
|
||||
_iterator = litellm.utils.ModelResponseIterator(model_response=_text_response, convert_to_delta=True)
|
||||
_streaming_response = litellm.TextCompletionStreamWrapper(
|
||||
completion_stream=_iterator,
|
||||
|
|
@ -16229,17 +16238,15 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str:
|
|||
|
||||
response: Final = await generate_key_helper_fn(
|
||||
request_type="key",
|
||||
**{
|
||||
"user_role": user_obj.user_role,
|
||||
"duration": LITELLM_UI_SESSION_DURATION,
|
||||
"key_max_budget": litellm.max_ui_session_budget,
|
||||
"models": [],
|
||||
"aliases": {},
|
||||
"config": {},
|
||||
"spend": 0,
|
||||
"user_id": user_obj.user_id,
|
||||
"team_id": UI_TEAM_ID,
|
||||
},
|
||||
user_role=user_obj.user_role,
|
||||
duration=LITELLM_UI_SESSION_DURATION,
|
||||
key_max_budget=litellm.max_ui_session_budget,
|
||||
models=[],
|
||||
aliases={},
|
||||
config={},
|
||||
spend=0,
|
||||
user_id=user_obj.user_id,
|
||||
team_id=UI_TEAM_ID,
|
||||
)
|
||||
key: Final = response["token"]
|
||||
|
||||
|
|
|
|||
|
|
@ -8,9 +8,10 @@ router, because the exception propagated up and was caught by the
|
|||
catch-all handler in _update_llm_router.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
||||
|
|
@ -294,7 +295,7 @@ class TestDeleteDeploymentKeepsPluginConfigModels:
|
|||
@pytest.mark.asyncio
|
||||
async def test_config_model_updated_params_reconciles_successfully(self, tmp_path):
|
||||
import yaml
|
||||
import litellm
|
||||
|
||||
from litellm.router import Router
|
||||
|
||||
initial_config = {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue