fix(proxy_server.py): fix /config/update/

allows updating router config via UI and having the change be propogated across all proxy instances by persisting config changes to the db
This commit is contained in:
Krrish Dholakia 2024-04-24 16:42:42 -07:00
parent bae6f41017
commit f54510b6ee
6 changed files with 50 additions and 90 deletions

View file

@ -4,6 +4,7 @@ import enum
from typing import Optional, List, Union, Dict, Literal, Any from typing import Optional, List, Union, Dict, Literal, Any
from datetime import datetime from datetime import datetime
import uuid, json, sys, os import uuid, json, sys, os
from litellm.types.router import UpdateRouterConfig
def hash_token(token: str): def hash_token(token: str):
@ -750,7 +751,7 @@ class ConfigYAML(LiteLLMBase):
description="litellm Module settings. See __init__.py for all, example litellm.drop_params=True, litellm.set_verbose=True, litellm.api_base, litellm.cache", description="litellm Module settings. See __init__.py for all, example litellm.drop_params=True, litellm.set_verbose=True, litellm.api_base, litellm.cache",
) )
general_settings: Optional[ConfigGeneralSettings] = None general_settings: Optional[ConfigGeneralSettings] = None
router_settings: Optional[dict] = Field( router_settings: Optional[UpdateRouterConfig] = Field(
None, None,
description="litellm router object settings. See router.py __init__ for all, example router.num_retries=5, router.timeout=5, router.max_retries=5, router.retry_after=5", description="litellm router object settings. See router.py __init__ for all, example router.num_retries=5, router.timeout=5, router.max_retries=5, router.retry_after=5",
) )

View file

@ -2636,9 +2636,16 @@ class ProxyConfig:
] ]
# router settings # router settings
if llm_router is not None: if llm_router is not None and prisma_client is not None:
_router_settings = config_data.get("router_settings", {}) db_router_settings = await prisma_client.db.litellm_config.find_first(
llm_router.update_settings(**_router_settings) where={"param_name": "router_settings"}
)
if (
db_router_settings is not None
and db_router_settings.param_value is not None
):
_router_settings = db_router_settings.param_value
llm_router.update_settings(**_router_settings)
async def add_deployment( async def add_deployment(
self, self,
@ -8406,96 +8413,29 @@ async def update_config(config_info: ConfigYAML):
""" """
global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client
try: try:
import base64
# Load existing config """
config = await proxy_config.get_config() - Update the ConfigTable DB
verbose_proxy_logger.debug("Loaded config: %s", config) - Run 'add_deployment'
"""
if prisma_client is None:
raise Exception("No DB Connected")
# update the general settings updated_settings = config_info.json(exclude_none=True)
if config_info.general_settings is not None: updated_settings = prisma_client.jsonify_object(updated_settings)
config.setdefault("general_settings", {}) for k, v in updated_settings.items():
updated_general_settings = config_info.general_settings.dict( await prisma_client.db.litellm_config.upsert(
exclude_none=True where={"param_name": k},
data={
"create": {"param_name": k, "param_value": v},
"update": {"param_value": v},
},
) )
_existing_settings = config["general_settings"] await proxy_config.add_deployment(
for k, v in updated_general_settings.items(): prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
# overwrite existing settings with updated values )
_existing_settings[k] = v
config["general_settings"] = _existing_settings
if config_info.environment_variables is not None:
config.setdefault("environment_variables", {})
_updated_environment_variables = config_info.environment_variables
# encrypt updated_environment_variables #
for k, v in _updated_environment_variables.items():
if isinstance(v, str):
encrypted_value = encrypt_value(value=v, master_key=master_key) # type: ignore
_updated_environment_variables[k] = base64.b64encode(
encrypted_value
).decode("utf-8")
_existing_env_variables = config["environment_variables"]
for k, v in _updated_environment_variables.items():
# overwrite existing env variables with updated values
_existing_env_variables[k] = _updated_environment_variables[k]
# update the litellm settings
if config_info.litellm_settings is not None:
config.setdefault("litellm_settings", {})
updated_litellm_settings = config_info.litellm_settings
config["litellm_settings"] = {
**updated_litellm_settings,
**config["litellm_settings"],
}
# if litellm.success_callback in updated_litellm_settings and config["litellm_settings"]
if (
"success_callback" in updated_litellm_settings
and "success_callback" in config["litellm_settings"]
):
# check both success callback are lists
if isinstance(
config["litellm_settings"]["success_callback"], list
) and isinstance(updated_litellm_settings["success_callback"], list):
combined_success_callback = (
config["litellm_settings"]["success_callback"]
+ updated_litellm_settings["success_callback"]
)
combined_success_callback = list(set(combined_success_callback))
config["litellm_settings"][
"success_callback"
] = combined_success_callback
# router settings
if config_info.router_settings is not None:
config.setdefault("router_settings", {})
_updated_router_settings = config_info.router_settings
config["router_settings"] = {
**config["router_settings"],
**_updated_router_settings,
}
# Save the updated config
await proxy_config.save_config(new_config=config)
# make sure the change is instantly rolled out for langfuse
if prisma_client is not None:
await proxy_config.add_deployment(
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
# Test new connections
## Slack
if "slack" in config.get("general_settings", {}).get("alerting", []):
await proxy_logging_obj.alerting_handler(
message="This is a test", level="Low"
)
return {"message": "Config updated successfully"} return {"message": "Config updated successfully"}
except Exception as e: except Exception as e:
traceback.print_exc() traceback.print_exc()

View file

@ -2553,6 +2553,8 @@ class Router:
"timeout", "timeout",
"max_retries", "max_retries",
"retry_after", "retry_after",
"fallbacks",
"context_window_fallbacks",
] ]
_int_settings = [ _int_settings = [

View file

@ -48,6 +48,23 @@ class RouterConfig(BaseModel):
protected_namespaces = () protected_namespaces = ()
class UpdateRouterConfig(BaseModel):
"""
Set of params that you can modify via `router.update_settings()`.
"""
routing_strategy_args: Optional[dict] = None
routing_strategy: Optional[str] = None
allowed_fails: Optional[int] = None
cooldown_time: Optional[float] = None
num_retries: Optional[int] = None
timeout: Optional[float] = None
max_retries: Optional[int] = None
retry_after: Optional[float] = None
fallbacks: Optional[List[dict]] = None
context_window_fallbacks: Optional[List[dict]] = None
class ModelInfo(BaseModel): class ModelInfo(BaseModel):
id: Optional[ id: Optional[
str str

View file

@ -125,7 +125,7 @@ const GeneralSettings: React.FC<GeneralSettingsPageProps> = ({
</TableRow> </TableRow>
</TableHead> </TableHead>
<TableBody> <TableBody>
{Object.entries(routerSettings).map(([param, value]) => ( {Object.entries(routerSettings).filter(([param, value]) => param != "fallbacks" && param != "context_window_fallbacks").map(([param, value]) => (
<TableRow key={param}> <TableRow key={param}>
<TableCell> <TableCell>
<Text>{param}</Text> <Text>{param}</Text>