Merge pull request #3277 from BerriAI/litellm_update_deployments

[UI] V0 - Edit Model tpm, rpm, api_base

litellm/proxy/proxy_server.py

@@ -106,7 +106,7 @@ import pydantic
 from litellm.proxy._types import *
 from litellm.caching import DualCache, RedisCache
 from litellm.proxy.health_check import perform_health_check
-from litellm.router import LiteLLM_Params, Deployment
+from litellm.router import LiteLLM_Params, Deployment, updateDeployment
 from litellm.router import ModelInfo as RouterModelInfo
 from litellm._logging import verbose_router_logger, verbose_proxy_logger
 from litellm.proxy.auth.handle_jwt import JWTHandler
@@ -7243,6 +7243,89 @@ async def add_new_model(
         )
 
 
+#### MODEL MANAGEMENT ####
+@router.post(
+    "/model/update",
+    description="Edit existing model params",
+    tags=["model management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def update_model(
+    model_params: updateDeployment,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config, prisma_client, master_key, store_model_in_db, proxy_logging_obj
+    try:
+        import base64
+
+        global prisma_client
+
+        if prisma_client is None:
+            raise HTTPException(
+                status_code=500,
+                detail={
+                    "error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys"
+                },
+            )
+        # update DB
+        if store_model_in_db == True:
+            _model_id = None
+            _model_info = getattr(model_params, "model_info", None)
+            if _model_info is None:
+                raise Exception("model_info not provided")
+
+            _model_id = _model_info.id
+            if _model_id is None:
+                raise Exception("model_info.id not provided")
+            _existing_litellm_params = (
+                await prisma_client.db.litellm_proxymodeltable.find_unique(
+                    where={"model_id": _model_id}
+                )
+            )
+            if _existing_litellm_params is None:
+                raise Exception("model not found")
+            _existing_litellm_params_dict = dict(
+                _existing_litellm_params.litellm_params
+            )
+
+            if model_params.litellm_params is None:
+                raise Exception("litellm_params not provided")
+
+            _new_litellm_params_dict = model_params.litellm_params.dict(
+                exclude_none=True
+            )
+
+            for key, value in _existing_litellm_params_dict.items():
+                if key in _new_litellm_params_dict:
+                    _existing_litellm_params_dict[key] = _new_litellm_params_dict[key]
+
+            _data: dict = {
+                "litellm_params": json.dumps(_existing_litellm_params_dict),  # type: ignore
+                "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+            }
+            model_response = await prisma_client.db.litellm_proxymodeltable.update(
+                where={"model_id": _model_id},
+                data=_data,  # type: ignore
+            )
+    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(
     "/v2/model/info",
     description="v2 - returns all the models set on the config.yaml, shows 'user_access' = True if the user has access to the model. Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)",

litellm/router.py

@@ -35,7 +35,14 @@ from litellm.utils import (
 import copy
 from litellm._logging import verbose_router_logger
 import logging
-from litellm.types.router import Deployment, ModelInfo, LiteLLM_Params, RouterErrors
+from litellm.types.router import (
+    Deployment,
+    ModelInfo,
+    LiteLLM_Params,
+    RouterErrors,
+    updateDeployment,
+    updateLiteLLMParams,
+)
 from litellm.integrations.custom_logger import CustomLogger
 
 

litellm/tests/test_add_update_models.py

@@ -0,0 +1,191 @@
+import sys, os
+import traceback
+from dotenv import load_dotenv
+from fastapi import Request
+from datetime import datetime
+
+load_dotenv()
+import os, io, time
+
+# this file is to test litellm/proxy
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import pytest, logging, asyncio
+import litellm, asyncio
+from litellm.proxy.proxy_server import add_new_model, update_model
+from litellm._logging import verbose_proxy_logger
+from litellm.proxy.utils import PrismaClient, ProxyLogging
+
+verbose_proxy_logger.setLevel(level=logging.DEBUG)
+from litellm.proxy.utils import DBClient
+from litellm.caching import DualCache
+from litellm.router import (
+    Deployment,
+    updateDeployment,
+    LiteLLM_Params,
+    ModelInfo,
+    updateLiteLLMParams,
+)
+
+from litellm.proxy._types import (
+    UserAPIKeyAuth,
+)
+
+proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
+
+
+@pytest.fixture
+def prisma_client():
+    from litellm.proxy.proxy_cli import append_query_params
+
+    ### add connection pool + pool timeout args
+    params = {"connection_limit": 100, "pool_timeout": 60}
+    database_url = os.getenv("DATABASE_URL")
+    modified_url = append_query_params(database_url, params)
+    os.environ["DATABASE_URL"] = modified_url
+    os.environ["STORE_MODEL_IN_DB"] = "true"
+
+    # Assuming DBClient is a class that needs to be instantiated
+    prisma_client = PrismaClient(
+        database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
+    )
+
+    # Reset litellm.proxy.proxy_server.prisma_client to None
+    litellm.proxy.proxy_server.custom_db_client = None
+    litellm.proxy.proxy_server.litellm_proxy_budget_name = (
+        f"litellm-proxy-budget-{time.time()}"
+    )
+    litellm.proxy.proxy_server.user_custom_key_generate = None
+
+    return prisma_client
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip(reason="new feature, tests passing locally")
+async def test_add_new_model(prisma_client):
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    setattr(litellm.proxy.proxy_server, "store_model_in_db", True)
+
+    await litellm.proxy.proxy_server.prisma_client.connect()
+    from litellm.proxy.proxy_server import user_api_key_cache
+    import uuid
+
+    _new_model_id = f"local-test-{uuid.uuid4().hex}"
+
+    await add_new_model(
+        model_params=Deployment(
+            model_name="test_model",
+            litellm_params=LiteLLM_Params(
+                model="azure/gpt-3.5-turbo",
+                api_key="test_api_key",
+                api_base="test_api_base",
+                rpm=1000,
+                tpm=1000,
+            ),
+            model_info=ModelInfo(
+                id=_new_model_id,
+            ),
+        ),
+        user_api_key_dict=UserAPIKeyAuth(
+            user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+        ),
+    )
+
+    _new_models = await prisma_client.db.litellm_proxymodeltable.find_many()
+    print("_new_models: ", _new_models)
+
+    _new_model_in_db = None
+    for model in _new_models:
+        print("current model: ", model)
+        if model.model_info["id"] == _new_model_id:
+            print("FOUND MODEL: ", model)
+            _new_model_in_db = model
+
+    assert _new_model_in_db is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip(reason="new feature, tests passing locally")
+async def test_add_update_model(prisma_client):
+    # test that existing litellm_params are not updated
+    # only new / updated params get updated
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    setattr(litellm.proxy.proxy_server, "store_model_in_db", True)
+
+    await litellm.proxy.proxy_server.prisma_client.connect()
+    from litellm.proxy.proxy_server import user_api_key_cache
+    import uuid
+
+    _new_model_id = f"local-test-{uuid.uuid4().hex}"
+
+    await add_new_model(
+        model_params=Deployment(
+            model_name="test_model",
+            litellm_params=LiteLLM_Params(
+                model="azure/gpt-3.5-turbo",
+                api_key="test_api_key",
+                api_base="test_api_base",
+                rpm=1000,
+                tpm=1000,
+            ),
+            model_info=ModelInfo(
+                id=_new_model_id,
+            ),
+        ),
+        user_api_key_dict=UserAPIKeyAuth(
+            user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+        ),
+    )
+
+    _new_models = await prisma_client.db.litellm_proxymodeltable.find_many()
+    print("_new_models: ", _new_models)
+
+    _new_model_in_db = None
+    for model in _new_models:
+        print("current model: ", model)
+        if model.model_info["id"] == _new_model_id:
+            print("FOUND MODEL: ", model)
+            _new_model_in_db = model
+
+    assert _new_model_in_db is not None
+
+    _original_model = _new_model_in_db
+    _original_litellm_params = _new_model_in_db.litellm_params
+    print("_original_litellm_params: ", _original_litellm_params)
+    print("now updating the tpm for model")
+    # run update to update "tpm"
+    await update_model(
+        model_params=updateDeployment(
+            litellm_params=updateLiteLLMParams(tpm=123456),
+            model_info=ModelInfo(
+                id=_new_model_id,
+            ),
+        ),
+        user_api_key_dict=UserAPIKeyAuth(
+            user_role="proxy_admin", api_key="sk-1234", user_id="1234"
+        ),
+    )
+
+    _new_models = await prisma_client.db.litellm_proxymodeltable.find_many()
+
+    _new_model_in_db = None
+    for model in _new_models:
+        if model.model_info["id"] == _new_model_id:
+            print("\nFOUND MODEL: ", model)
+            _new_model_in_db = model
+
+    # assert all other litellm params are identical to _original_litellm_params
+    for key, value in _original_litellm_params.items():
+        if key == "tpm":
+            # assert that tpm actually got updated
+            assert _new_model_in_db.litellm_params[key] == 123456
+        else:
+            assert _new_model_in_db.litellm_params[key] == value
+
+    assert _original_model.model_id == _new_model_in_db.model_id
+    assert _original_model.model_name == _new_model_in_db.model_name
+    assert _original_model.model_info == _new_model_in_db.model_info

litellm/types/router.py

@@ -155,6 +155,36 @@ class LiteLLM_Params(BaseModel):
         setattr(self, key, value)
 
 
+class updateLiteLLMParams(BaseModel):
+    # This class is used to update the LiteLLM_Params
+    # only differece is model is optional
+    model: Optional[str] = None
+    tpm: Optional[int] = None
+    rpm: Optional[int] = None
+    api_key: Optional[str] = None
+    api_base: Optional[str] = None
+    api_version: Optional[str] = None
+    timeout: Optional[Union[float, str]] = None  # if str, pass in as os.environ/
+    stream_timeout: Optional[Union[float, str]] = (
+        None  # timeout when making stream=True calls, if str, pass in as os.environ/
+    )
+    max_retries: int = 2  # follows openai default of 2
+    organization: Optional[str] = None  # for openai orgs
+    ## VERTEX AI ##
+    vertex_project: Optional[str] = None
+    vertex_location: Optional[str] = None
+    ## AWS BEDROCK / SAGEMAKER ##
+    aws_access_key_id: Optional[str] = None
+    aws_secret_access_key: Optional[str] = None
+    aws_region_name: Optional[str] = None
+
+
+class updateDeployment(BaseModel):
+    model_name: Optional[str] = None
+    litellm_params: Optional[updateLiteLLMParams] = None
+    model_info: Optional[ModelInfo] = None
+
+
 class Deployment(BaseModel):
     model_name: str
     litellm_params: LiteLLM_Params

ui/litellm-dashboard/src/components/model_dashboard.tsx

@@ -18,7 +18,7 @@ import {
 } from "@tremor/react";
 import { TabPanel, TabPanels, TabGroup, TabList, Tab, TextInput, Icon } from "@tremor/react";
 import { Select, SelectItem, MultiSelect, MultiSelectItem } from "@tremor/react";
-import { modelInfoCall, userGetRequesedtModelsCall, modelCreateCall, Model, modelCostMap, modelDeleteCall, healthCheckCall } from "./networking";
+import { modelInfoCall, userGetRequesedtModelsCall, modelCreateCall, Model, modelCostMap, modelDeleteCall, healthCheckCall, modelUpdateCall } from "./networking";
 import { BarChart } from "@tremor/react";
 import {
   Button as Button2,
@@ -51,6 +51,13 @@ interface ModelDashboardProps {
   userID: string | null;
 }
 
+interface EditModelModalProps {
+  visible: boolean;
+  onCancel: () => void;
+  model: any; // Assuming TeamType is a type representing your team object
+  onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted
+}
+
 //["OpenAI", "Azure OpenAI", "Anthropic", "Gemini (Google AI Studio)", "Amazon Bedrock", "OpenAI-Compatible Endpoints (Groq, Together AI, Mistral AI, etc.)"]
 
 enum Providers {
@@ -183,6 +190,169 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
   
   const [selectedProvider, setSelectedProvider] = useState<String>("OpenAI");
   const [healthCheckResponse, setHealthCheckResponse] = useState<string>('');
+  const [editModalVisible, setEditModalVisible] = useState<boolean>(false);
+  const [selectedModel, setSelectedModel] = useState<any>(null);
+
+  const EditModelModal: React.FC<EditModelModalProps> = ({ visible, onCancel, model, onSubmit }) => {
+    const [form] = Form.useForm();
+    let litellm_params_to_edit: Record<string, any> = {}
+    let model_name = "";
+    let model_id = "";
+    if (model) {
+      litellm_params_to_edit = model.litellm_params
+      model_name = model.model_name;
+      let model_info = model.model_info;
+      if (model_info ) {
+        model_id = model_info.id;
+        console.log(`model_id: ${model_id}`)
+        litellm_params_to_edit.model_id = model_id;
+      }
+    }
+   
+  
+    const handleOk = () => {
+      form
+        .validateFields()
+        .then((values) => {
+          onSubmit(values);
+          form.resetFields();
+        })
+        .catch((error) => {
+          console.error("Validation failed:", error);
+        });
+  };
+  
+    return (
+        <Modal
+              title={"Edit Model " + model_name}
+              visible={visible}
+              width={800}
+              footer={null}
+              onOk={handleOk}
+              onCancel={onCancel}
+            >
+        <Form
+          form={form}
+          onFinish={handleEditSubmit}
+          initialValues={litellm_params_to_edit} // Pass initial values here
+          labelCol={{ span: 8 }}
+          wrapperCol={{ span: 16 }}
+          labelAlign="left"
+        >
+                <>
+
+                <Form.Item className="mt-8"
+                    label="api_base" 
+                    name="api_base"
+                    
+                  >
+                  <TextInput/>
+
+                  </Form.Item>
+                  
+                  <Form.Item 
+                    label="tpm" 
+                    name="tpm"
+                  >
+                  <InputNumber min={0} step={1} />
+
+                  </Form.Item>
+
+                  <Form.Item 
+                    label="rpm" 
+                    name="rpm"
+                  >
+                  <InputNumber min={0} step={1} />
+                  </Form.Item>
+
+                  <Form.Item 
+                    label="max_retries" 
+                    name="max_retries"
+                  >
+
+                    
+                  <InputNumber min={0} step={1} />
+
+                  </Form.Item>
+
+                  
+
+                  <Form.Item 
+                    label="model_id" 
+                    name="model_id"
+                    hidden={true}
+                  >
+                  </Form.Item>
+
+                  
+
+                  
+                </>
+                <div style={{ textAlign: "right", marginTop: "10px" }}>
+                  <Button2 htmlType="submit">Save</Button2>
+                </div>
+              </Form>
+      </Modal>
+    );
+  };
+  
+
+  const handleEditClick = (model: any) => {
+    setSelectedModel(model);
+    setEditModalVisible(true);
+  };
+
+  const handleEditCancel = () => {
+    setEditModalVisible(false);
+    setSelectedModel(null);
+  };
+
+
+const handleEditSubmit = async (formValues: Record<string, any>) => {
+  // Call API to update team with teamId and values
+  
+  console.log("handleEditSubmit:", formValues);
+  if (accessToken == null) {
+    return;
+  }
+
+  let newLiteLLMParams: Record<string, any> = {}
+  let model_info_model_id = null;
+
+  for (const [key, value] of Object.entries(formValues)) {
+    if (key !== "model_id") {
+      newLiteLLMParams[key] = value;
+    } else {
+      model_info_model_id = value;
+    }
+  }
+
+  let payload = {
+    litellm_params: newLiteLLMParams, 
+    model_info: {
+      "id": model_info_model_id
+    }
+  }
+
+  console.log("handleEditSubmit payload:", payload);
+
+  let newModelValue = await modelUpdateCall(accessToken, payload);
+
+  // Update the teams state with the updated team data
+  // if (teams) {
+  //   const updatedTeams = teams.map((team) =>
+  //     team.team_id === teamId ? newTeamValues.data : team
+  //   );
+  //   setTeams(updatedTeams);
+  // }
+  message.success("Model updated successfully, restart server to see updates");
+
+  setEditModalVisible(false);
+  setSelectedModel(null);
+};
+
+
+  
 
 
   const props: UploadProps = {
@@ -510,15 +680,27 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
                   <TableCell>{model.output_cost}</TableCell>
                   <TableCell>{model.max_tokens}</TableCell>
                   <TableCell>
+                        <Icon
+                            icon={PencilAltIcon}
+                            size="sm"
+                            onClick={() => handleEditClick(model)}
+                          />
                           <DeleteModelButton modelID={model.model_info.id} accessToken={accessToken} />
                         </TableCell>
                 </TableRow>
               ))}
             </TableBody>
           </Table>
+
         </Card>
 
       </Grid>
+      <EditModelModal
+          visible={editModalVisible}
+          onCancel={handleEditCancel}
+          model={selectedModel}
+          onSubmit={handleEditSubmit}
+        />
       </TabPanel>
       <TabPanel className="h-full">
       <Title2 level={2}>Add new model</Title2>

ui/litellm-dashboard/src/components/networking.tsx

@@ -1015,6 +1015,41 @@ export const teamUpdateCall = async (
   }
 };
 
+export const modelUpdateCall = async (
+  accessToken: string,
+  formValues: Record<string, any> // Assuming formValues is an object
+) => {
+  try {
+    console.log("Form Values in modelUpateCall:", formValues); // Log the form values before making the API call
+
+    const url = proxyBaseUrl ? `${proxyBaseUrl}/model/update` : `/model/update`;
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        ...formValues, // Include formValues in the request body
+      }),
+    });
+
+    if (!response.ok) {
+      const errorData = await response.text();
+      message.error("Failed to update model: " + errorData, 20);
+      console.error("Error update from the server:", errorData);
+      throw new Error("Network response was not ok");
+    }
+    const data = await response.json();
+    console.log("Update model Response:", data);
+    return data;
+    // Handle success - you might want to update some state or UI based on the created key
+  } catch (error) {
+    console.error("Failed to update model:", error);
+    throw error;
+  }
+};
+
 export interface Member {
   role: string;
   user_id: string | null;