Merge pull request #2902 from BerriAI/litellm_ui_set_get_callbacks

UI view set callbacks

litellm/proxy/proxy_server.py

@@ -1809,17 +1809,15 @@ class ProxyConfig:
             }
 
         ## DB
-        if (
-            prisma_client is not None
-            and litellm.get_secret("SAVE_CONFIG_TO_DB", False) == True
+        if prisma_client is not None and (
+            general_settings.get("store_model_in_db", False) == True
         ):
-            prisma_setup(database_url=None)  # in case it's not been connected yet
             _tasks = []
             keys = [
-                "model_list",
                 "general_settings",
                 "router_settings",
                 "litellm_settings",
+                "environment_variables",
             ]
             for k in keys:
                 response = prisma_client.get_generic_data(
@@ -1828,6 +1826,12 @@ class ProxyConfig:
                 _tasks.append(response)
 
             responses = await asyncio.gather(*_tasks)
+            for response in responses:
+                if response is not None:
+                    param_name = getattr(response, "param_name", None)
+                    param_value = getattr(response, "param_value", None)
+                    if param_name is not None and param_value is not None:
+                        config[param_name] = param_value
 
         return config
 
@@ -1836,11 +1840,6 @@ class ProxyConfig:
         # Load existing config
         backup_config = await self.get_config()
 
-        # Save the updated config
-        ## YAML
-        with open(f"{user_config_file_path}", "w") as config_file:
-            yaml.dump(new_config, config_file, default_flow_style=False)
-
         # update Router - verifies if this is a valid config
         try:
             (
@@ -1862,22 +1861,17 @@ class ProxyConfig:
         - 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 litellm.get_secret("SAVE_CONFIG_TO_DB", default_value=False) == True
+        if prisma_client is not None and (
+            general_settings.get("store_model_in_db", False) == True
         ):
-            ### KEY REMOVAL ###
-            models = new_config.get("model_list", [])
-            for m in models:
-                if m.get("litellm_params", {}).get("api_key", None) is not None:
-                    # pop the key
-                    api_key = m["litellm_params"].pop("api_key")
-                    # store in local env
-                    key_name = f"LITELLM_MODEL_KEY_{uuid.uuid4()}"
-                    os.environ[key_name] = api_key
-                    # save the key name (not the value)
-                    m["litellm_params"]["api_key"] = f"os.environ/{key_name}"
+            # if using - db for config - models are in ModelTable
+            new_config.pop("model_list", None)
             await prisma_client.insert_data(data=new_config, 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)
 
     async def load_team_config(self, team_id: str):
         """
@@ -7922,8 +7916,10 @@ async def update_config(config_info: ConfigYAML):
 
     Currently supports modifying General Settings + LiteLLM settings
     """
-    global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj
+    global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key
     try:
+        import base64
+
         # Load existing config
         config = await proxy_config.get_config()
 
@@ -7943,9 +7939,17 @@ async def update_config(config_info: ConfigYAML):
 
         if config_info.environment_variables is not None:
             config.setdefault("environment_variables", {})
-            updated_environment_variables = config_info.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")
             config["environment_variables"] = {
-                **updated_environment_variables,
+                **_updated_environment_variables,
                 **config["environment_variables"],
             }
 
@@ -7958,6 +7962,25 @@ async def update_config(config_info: ConfigYAML):
                 **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
+
         # Save the updated config
         await proxy_config.save_config(new_config=config)
 
@@ -7987,6 +8010,48 @@ async def update_config(config_info: ConfigYAML):
         )
 
 
+@router.get(
+    "/get/config/callbacks",
+    tags=["config.yaml"],
+    include_in_schema=False,
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def get_config():
+    """
+    For Admin UI - allows admin to view config via UI
+
+    """
+    global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key
+    try:
+
+        config_data = await proxy_config.get_config()
+        _environment_variables = config_data.get("environment_variables", {})
+        config_data = config_data["litellm_settings"]
+
+        # only store the keys and return the values as sk...***
+        for key, value in _environment_variables.items():
+            _environment_variables[key] = value[:5] + "*****"
+        config_data["environment_variables"] = _environment_variables
+        return {"data": config_data, "status": "success"}
+    except Exception as e:
+        traceback.print_exc()
+        if isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "detail", f"Authentication Error({str(e)})"),
+                type="auth_error",
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
+            )
+        elif isinstance(e, ProxyException):
+            raise e
+        raise ProxyException(
+            message="Authentication Error, " + str(e),
+            type="auth_error",
+            param=getattr(e, "param", "None"),
+            code=status.HTTP_400_BAD_REQUEST,
+        )
+
+
 @router.get(
     "/config/yaml",
     tags=["config.yaml"],