From 29a8a613a7a15eff2a0bce9a125a0c07ec2c5681 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 31 Jan 2025 17:54:29 -0800
Subject: [PATCH] (UI) Allow adding custom pricing when adding new model
 (#8165)

* ui custom pricing

* fix validateJSON

* ui allow entering custom pricing

* ui fix add model form

* working edit model modal

* working edit custom pricing

* ui add custom pricing on ui
---
 .../add_model/advanced_settings.tsx           | 230 ++++++++++++------
 .../add_model/handle_add_model_submit.tsx     |  24 ++
 .../edit_model/edit_model_modal.tsx           | 146 +++++++++++
 .../src/components/model_dashboard.tsx        | 142 +----------
 4 files changed, 333 insertions(+), 209 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/edit_model/edit_model_modal.tsx

diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
index 3065125fcb..fc65ce54a0 100644
--- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
@@ -1,6 +1,6 @@
 import React from "react";
-import { Form, Switch } from "antd";
-import { Text, Button, Accordion, AccordionHeader, AccordionBody } from "@tremor/react";
+import { Form, Switch, Select, Input } from "antd";
+import { Text, Button, Accordion, AccordionHeader, AccordionBody, TextInput } from "@tremor/react";
 import { Row, Col, Typography, Card } from "antd";
 import TextArea from "antd/es/input/TextArea";
 const { Link } = Typography;
@@ -15,8 +15,20 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
   setShowAdvancedSettings,
 }) => {
   const [form] = Form.useForm();
+  const [customPricing, setCustomPricing] = React.useState(false);
+  const [pricingModel, setPricingModel] = React.useState<'per_token' | 'per_second'>('per_token');
+
+  // Add validation function for numbers
+  const validateNumber = (_: any, value: string) => {
+    if (!value) {
+      return Promise.resolve();
+    }
+    if (isNaN(Number(value)) || Number(value) < 0) {
+      return Promise.reject('Please enter a valid positive number');
+    }
+    return Promise.resolve();
+  };
 
-  // Add validation function
   const validateJSON = (_: any, value: string) => {
     if (!value) {
       return Promise.resolve();
@@ -29,6 +41,19 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
     }
   };
 
+  // Handle custom pricing changes
+  const handleCustomPricingChange = (checked: boolean) => {
+    setCustomPricing(checked);
+    if (!checked) {
+      // Clear pricing fields when disabled
+      form.setFieldsValue({
+        input_cost_per_token: undefined,
+        output_cost_per_token: undefined,
+        input_cost_per_second: undefined,
+      });
+    }
+  };
+
   const handlePassThroughChange = (checked: boolean) => {
     const currentParams = form.getFieldValue('litellm_extra_params');
     try {
@@ -56,81 +81,136 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
 
   return (
     <>
+      <Accordion className="mt-2 mb-4">
+        <AccordionHeader>
+          <b>Advanced Settings</b>
+        </AccordionHeader>
+        <AccordionBody>
+          <div className="bg-white rounded-lg">
+            <Form.Item
+              label="Custom Pricing"
+              name="custom_pricing"
+              valuePropName="checked"
+              className="mb-4"
+            >
+              <Switch onChange={handleCustomPricingChange} className="bg-gray-600" />
+            </Form.Item>
 
-        <Accordion className="mt-2 mb-4">
-          
-          <AccordionHeader>
-            <b>Advanced Settings</b>
-          </AccordionHeader>
-          <AccordionBody>
-            <div className="bg-white rounded-lg">
-              <Form.Item
-                label="Use in pass through routes"
-                name="use_in_pass_through"
-                valuePropName="checked"
-                className="mb-4 mt-4"
-                tooltip={
-                  <span>
-                    Allow using these credentials in pass through routes.{" "}
-                    <Link href="https://docs.litellm.ai/docs/pass_through/vertex_ai" target="_blank">
-                      Learn more
-                    </Link>
-                  </span>
-                }
-              >
-                <Switch 
-                  onChange={handlePassThroughChange} 
-                  className="bg-gray-600" 
-                />
-              </Form.Item>
-              <Form.Item
-                label="LiteLLM Params"
-                name="litellm_extra_params"
-                tooltip="Optional litellm params used for making a litellm.completion() call."
-                className="mb-4 mt-4"
-                rules={[{ validator: validateJSON }]}
-              >
-                <TextArea
-                  rows={4}
-                  placeholder='{
-                    "rpm": 100,
-                    "timeout": 0,
-                    "stream_timeout": 0
-                  }'
-                />
-              </Form.Item>
-              <Row className="mb-4">
-                <Col span={10}></Col>
-                <Col span={10}>
-                  <Text className="text-gray-600 text-sm">
-                    Pass JSON of litellm supported params{" "}
-                    <Link
-                      href="https://docs.litellm.ai/docs/completion/input"
-                      target="_blank"
+            {customPricing && (
+              <div className="ml-6 pl-4 border-l-2 border-gray-200">
+                <Form.Item
+                  label="Pricing Model"
+                  name="pricing_model"
+                  className="mb-4"
+                >
+                  <Select
+                    defaultValue="per_token"
+                    onChange={(value) => setPricingModel(value)}
+                    options={[
+                      { value: 'per_token', label: 'Per Million Tokens' },
+                      { value: 'per_second', label: 'Per Second' },
+                    ]}
+                  />
+                </Form.Item>
+
+                {pricingModel === 'per_token' ? (
+                  <>
+                    <Form.Item
+                      label="Input Cost (per 1M tokens)"
+                      name="input_cost_per_token"
+                      rules={[{ validator: validateNumber }]}
+                      className="mb-4"
                     >
-                      litellm.completion() call
-                    </Link>
-                  </Text>
-                </Col>
-              </Row>
-              <Form.Item
-                label="Model Info"
-                name="model_info_params"
-                tooltip="Optional model info params. Returned when calling `/model/info` endpoint."
-                className="mb-0"
-                rules={[{ validator: validateJSON }]}
-              >
-                <TextArea
-                  rows={4}
-                  placeholder='{
-                    "mode": "chat"
-                  }'
-                />
-              </Form.Item>
-            </div>
-          </AccordionBody>
-        </Accordion>
+                      <TextInput />
+                    </Form.Item>
+                    <Form.Item
+                      label="Output Cost (per 1M tokens)"
+                      name="output_cost_per_token"
+                      rules={[{ validator: validateNumber }]}
+                      className="mb-4"
+                    >
+                      <TextInput />
+                    </Form.Item>
+                  </>
+                ) : (
+                  <Form.Item
+                    label="Cost Per Second"
+                    name="input_cost_per_second"
+                    rules={[{ validator: validateNumber }]}
+                    className="mb-4"
+                  >
+                    <TextInput />
+                  </Form.Item>
+                )}
+              </div>
+            )}
 
+            <Form.Item
+              label="Use in pass through routes"
+              name="use_in_pass_through"
+              valuePropName="checked"
+              className="mb-4 mt-4"
+              tooltip={
+                <span>
+                  Allow using these credentials in pass through routes.{" "}
+                  <Link href="https://docs.litellm.ai/docs/pass_through/vertex_ai" target="_blank">
+                    Learn more
+                  </Link>
+                </span>
+              }
+            >
+              <Switch 
+                onChange={handlePassThroughChange} 
+                className="bg-gray-600" 
+              />
+            </Form.Item>
+            <Form.Item
+              label="LiteLLM Params"
+              name="litellm_extra_params"
+              tooltip="Optional litellm params used for making a litellm.completion() call."
+              className="mb-4 mt-4"
+                rules={[{ validator: validateJSON }]}
+            >
+              <TextArea
+                rows={4}
+                placeholder='{
+                  "rpm": 100,
+                  "timeout": 0,
+                  "stream_timeout": 0
+                }'
+              />
+            </Form.Item>
+            <Row className="mb-4">
+              <Col span={10}></Col>
+              <Col span={10}>
+                <Text className="text-gray-600 text-sm">
+                  Pass JSON of litellm supported params{" "}
+                  <Link
+                    href="https://docs.litellm.ai/docs/completion/input"
+                    target="_blank"
+                  >
+                    litellm.completion() call
+                  </Link>
+                </Text>
+              </Col>
+            </Row>
+            <Form.Item
+              label="Model Info"
+              name="model_info_params"
+              tooltip="Optional model info params. Returned when calling `/model/info` endpoint."
+              className="mb-0"
+                rules={[{ validator: validateJSON }]}
+            >
+              <TextArea
+                rows={4}
+                placeholder='{
+                  "mode": "chat"
+                }'
+              />
+            </Form.Item>
+          </div>
+        </AccordionBody>
+      </Accordion>
     </>
   );
 };
diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx
index 1aad96a165..95003ed9c1 100644
--- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.tsx
@@ -35,6 +35,16 @@ export const handleAddModelSubmit = async (
         console.log(`litellm_model: ${litellm_model}`);
         const litellmParamsObj: Record<string, any> = {};
         const modelInfoObj: Record<string, any> = {};
+        
+        // Handle pricing conversion before processing other fields
+        if (formValues.input_cost_per_token) {
+          formValues.input_cost_per_token = Number(formValues.input_cost_per_token) / 1000000;
+        }
+        if (formValues.output_cost_per_token) {
+          formValues.output_cost_per_token = Number(formValues.output_cost_per_token) / 1000000;
+        }
+        // Keep input_cost_per_second as is, no conversion needed
+        
         // Iterate through the key-value pairs in formValues
         litellmParamsObj["model"] = litellm_model;
         let modelName: string = "";
@@ -43,6 +53,10 @@ export const handleAddModelSubmit = async (
           if (value === "") {
             continue;
           }
+          // Skip the custom_pricing and pricing_model fields as they're only used for UI control
+          if (key === 'custom_pricing' || key === 'pricing_model') {
+            continue;
+          }
           if (key == "model_name") {
             modelName = modelName + value;
           } else if (key == "custom_llm_provider") {
@@ -97,6 +111,16 @@ export const handleAddModelSubmit = async (
             }
           }
   
+          // Handle the pricing fields
+          else if (key === "input_cost_per_token" || 
+                  key === "output_cost_per_token" || 
+                  key === "input_cost_per_second") {
+            if (value) {
+              litellmParamsObj[key] = Number(value);
+            }
+            continue;
+          }
+  
           // Check if key is any of the specified API related keys
           else {
             // Add key-value pair to litellm_params dictionary
diff --git a/ui/litellm-dashboard/src/components/edit_model/edit_model_modal.tsx b/ui/litellm-dashboard/src/components/edit_model/edit_model_modal.tsx
new file mode 100644
index 0000000000..052dad18f7
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/edit_model/edit_model_modal.tsx
@@ -0,0 +1,146 @@
+import React from "react";
+import { Modal, Form, InputNumber, message } from "antd";
+import { TextInput } from "@tremor/react";
+import { Button as Button2 } from "antd";
+
+interface EditModelModalProps {
+  visible: boolean;
+  onCancel: () => void;
+  model: any;
+  onSubmit: (data: FormData) => void;
+}
+
+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,
+      input_cost_per_token: model.litellm_params?.input_cost_per_token ? (model.litellm_params.input_cost_per_token * 1_000_000) : undefined,
+      output_cost_per_token: model.litellm_params?.output_cost_per_token ? (model.litellm_params.output_cost_per_token * 1_000_000) : undefined,
+    };
+    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) => {
+        const submissionValues = {
+          ...values,
+          input_cost_per_token: values.input_cost_per_token ? Number(values.input_cost_per_token) / 1_000_000 : undefined,
+          output_cost_per_token: values.output_cost_per_token ? Number(values.output_cost_per_token) / 1_000_000 : undefined,
+        };
+        onSubmit(submissionValues);
+        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={onSubmit}
+        initialValues={litellm_params_to_edit}
+        labelCol={{ span: 8 }}
+        wrapperCol={{ span: 16 }}
+        labelAlign="left"
+      >
+        <>
+          <Form.Item
+            label="Input Cost (per 1M tokens)"
+            name="input_cost_per_token"
+            tooltip="float (optional) - Input cost per 1 million tokens"
+          >
+            <TextInput />
+          </Form.Item>
+
+          <Form.Item
+            label="Output Cost (per 1M tokens)"
+            name="output_cost_per_token"
+            tooltip="float (optional) - Output cost per 1 million tokens"
+          >
+            <TextInput />
+          </Form.Item>
+          <Form.Item className="mt-8" label="api_base" name="api_base">
+            <TextInput />
+          </Form.Item>
+          <Form.Item
+            label="organization"
+            name="organization"
+            tooltip="OpenAI Organization ID"
+          >
+            <TextInput />
+          </Form.Item>
+
+          <Form.Item
+            label="tpm"
+            name="tpm"
+            tooltip="int (optional) - Tokens limit for this deployment: in tokens per minute (tpm). Find this information on your model/providers website"
+          >
+            <InputNumber min={0} step={1} />
+          </Form.Item>
+
+          <Form.Item
+            label="rpm"
+            name="rpm"
+            tooltip="int (optional) - Rate limit for this deployment: in requests per minute (rpm). Find this information on your model/providers website"
+          >
+            <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="timeout"
+            name="timeout"
+            tooltip="int (optional) - Timeout in seconds for LLM requests (Defaults to 600 seconds)"
+          >
+            <InputNumber min={0} step={1} />
+          </Form.Item>
+
+          <Form.Item
+            label="stream_timeout"
+            name="stream_timeout"
+            tooltip="int (optional) - Timeout for stream requests (seconds)"
+          >
+            <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>
+  );
+};
+
+export default EditModelModal;
\ No newline at end of file
diff --git a/ui/litellm-dashboard/src/components/model_dashboard.tsx b/ui/litellm-dashboard/src/components/model_dashboard.tsx
index c3d9144782..8e5a13a045 100644
--- a/ui/litellm-dashboard/src/components/model_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/model_dashboard.tsx
@@ -20,6 +20,7 @@ import ConditionalPublicModelName from "./add_model/conditional_public_model_nam
 import LiteLLMModelNameField from "./add_model/litellm_model_name";
 import AdvancedSettings from "./add_model/advanced_settings";
 import { handleAddModelSubmit } from "./add_model/handle_add_model_submit";
+import EditModelModal from "./edit_model/edit_model_modal";
 import {
   TabPanel,
   TabPanels,
@@ -367,134 +368,6 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
     return 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="organization"
-              name="organization"
-              tooltip="OpenAI Organization ID"
-            >
-                <TextInput />
-            </Form.Item>
-
-            <Form.Item
-              label="tpm"
-              name="tpm"
-              tooltip="int (optional) - Tokens limit for this deployment: in tokens per minute (tpm). Find this information on your model/providers website"
-            >
-              <InputNumber min={0} step={1} />
-            </Form.Item>
-
-            <Form.Item
-              label="rpm"
-              name="rpm"
-              tooltip="int (optional) - Rate limit for this deployment: in requests per minute (rpm). Find this information on your model/providers website"
-            >
-              <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="timeout"
-              name="timeout"
-              tooltip="int (optional) - Timeout in seconds for LLM requests (Defaults to 600 seconds)"
-            >
-              <InputNumber min={0} step={1} />
-            </Form.Item>
-
-            <Form.Item
-              label="stream_timeout"
-              name="stream_timeout"
-              tooltip="int (optional) - Timeout for stream requests (seconds)"
-            >
-              <InputNumber min={0} step={1} />
-            </Form.Item>
-
-            <Form.Item
-              label="Input Cost per 1M Tokens"
-              name="input_cost_per_million_tokens"
-              tooltip="float (optional) - Input cost per 1 million tokens"
-            >
-              <InputNumber min={0} step={0.01} />
-            </Form.Item>
-
-            <Form.Item
-              label="Output Cost per 1M Tokens"
-              name="output_cost_per_million_tokens"
-              tooltip="float (optional) - Output cost per 1 million tokens"
-            >
-              <InputNumber min={0} step={0.01} />
-            </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);
@@ -526,14 +399,15 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
     let newLiteLLMParams: Record<string, any> = {};
     let model_info_model_id = null;
 
-    if (formValues.input_cost_per_million_tokens) {
-      formValues.input_cost_per_token = formValues.input_cost_per_million_tokens / 1000000;
-      delete formValues.input_cost_per_million_tokens;
+    if (formValues.input_cost_per_token) {
+      // Convert from per 1M tokens to per token
+      formValues.input_cost_per_token = Number(formValues.input_cost_per_token) / 1_000_000;
     }
-    if (formValues.output_cost_per_million_tokens) {
-      formValues.output_cost_per_token = formValues.output_cost_per_million_tokens / 1000000;
-      delete formValues.output_cost_per_million_tokens;
+    if (formValues.output_cost_per_token) {
+      // Convert from per 1M tokens to per token
+      formValues.output_cost_per_token = Number(formValues.output_cost_per_token) / 1_000_000;
     }
+  
 
     for (const [key, value] of Object.entries(formValues)) {
       if (key !== "model_id") {