(Admin UI) - Test Key Tab - Allow using UI Session instead of manua…

…lly creating a virtual key (BerriAI#7348) * ui fix - allow searching model list + fix bug on filtering * qa fix - use correct provider name for azure_text * ui wrap content onto next line * ui fix - allow selecting current UI session when logging in * ui session budgets
kp-forks · Dec 21, 2024 · ce41cd9 · ce41cd9
1 parent b527834
commit ce41cd9
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 8 deletions.
diff --git a/litellm/__init__.py b/litellm/__init__.py
@@ -296,6 +296,7 @@
 max_user_budget: Optional[float] = None
 default_max_internal_user_budget: Optional[float] = None
 max_internal_user_budget: Optional[float] = None
+max_ui_session_budget: Optional[float] = 10  # $10 USD budgets for UI Chat sessions
 internal_user_budget_duration: Optional[str] = None
 tag_budget_config: Optional[Dict[str, BudgetConfig]] = None
 max_end_user_budget: Optional[float] = None

diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py
@@ -449,7 +449,7 @@ def response_convertor(response, client):
     # But if it is, we want their models preferences
     default_ui_key_values = {
         "duration": "24hr",
-        "key_max_budget": 0.01,
+        "key_max_budget": litellm.max_ui_session_budget,
         "aliases": {},
         "config": {},
         "spend": 0,

diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
@@ -3,6 +3,13 @@ model_list:
     litellm_params:
       model: openai/*
       api_key: os.environ/OPENAI_API_KEY
+  - model_name: anthropic/*
+    litellm_params:
+      model: anthropic/*
+      api_key: os.environ/ANTHROPIC_API_KEY
+  - model_name: bedrock/*
+    litellm_params:
+      model: bedrock/*
 
 
 # for /files endpoints

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
@@ -7774,7 +7774,7 @@ async def login(request: Request):  # noqa: PLR0915
                 **{
                     "user_role": LitellmUserRoles.PROXY_ADMIN,
                     "duration": "24hr",
-                    "key_max_budget": 5,
+                    "key_max_budget": litellm.max_ui_session_budget,
                     "models": [],
                     "aliases": {},
                     "config": {},
@@ -7841,7 +7841,7 @@ async def login(request: Request):  # noqa: PLR0915
                     **{  # type: ignore
                         "user_role": user_role,
                         "duration": "24hr",
-                        "key_max_budget": 5,
+                        "key_max_budget": litellm.max_ui_session_budget,
                         "models": [],
                         "aliases": {},
                         "config": {},
@@ -7972,7 +7972,7 @@ async def onboarding(invite_link: str):
         **{
             "user_role": user_obj.user_role,
             "duration": "24hr",
-            "key_max_budget": 5,
+            "key_max_budget": litellm.max_ui_session_budget,
             "models": [],
             "aliases": {},
             "config": {},

diff --git a/ui/litellm-dashboard/src/components/chat_ui.tsx b/ui/litellm-dashboard/src/components/chat_ui.tsx
@@ -88,6 +88,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
   userRole,
   userID,
 }) => {
+  const [apiKeySource, setApiKeySource] = useState<'session' | 'custom'>('session');
   const [apiKey, setApiKey] = useState("");
   const [inputMessage, setInputMessage] = useState("");
   const [chatHistory, setChatHistory] = useState<any[]>([]);
@@ -167,7 +168,14 @@ const ChatUI: React.FC<ChatUIProps> = ({
   const handleSendMessage = async () => {
     if (inputMessage.trim() === "") return;
 
-    if (!apiKey || !token || !userRole || !userID) {
+    if (!token || !userRole || !userID) {
+      return;
+    }
+
+    const effectiveApiKey = apiKeySource === 'session' ? accessToken : apiKey;
+
+    if (!effectiveApiKey) {
+      message.error("Please provide an API key or select Current UI Session");
       return;
     }
 
@@ -182,7 +190,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
           inputMessage,
           (chunk) => updateUI("assistant", chunk),
           selectedModel,
-          apiKey
+          effectiveApiKey
         );
       }
     } catch (error) {
@@ -223,8 +231,25 @@ const ChatUI: React.FC<ChatUIProps> = ({
               <div className="sm:max-w-2xl">
           <Grid numItems={2}>
             <Col>
-            <Text>API Key</Text>
-              <TextInput placeholder="Type API Key here" type="password" onValueChange={setApiKey} value={apiKey}/>
+              <Text>API Key Source</Text>
+              <Select
+                defaultValue="session"
+                style={{ width: "100%" }}
+                onChange={(value) => setApiKeySource(value as "session" | "custom")}
+                options={[
+                  { value: 'session', label: 'Current UI Session' },
+                  { value: 'custom', label: 'Virtual Key' },
+                ]}
+              />
+              {apiKeySource === 'custom' && (
+                <TextInput
+                  className="mt-2"
+                  placeholder="Enter custom API key"
+                  type="password"
+                  onValueChange={setApiKey}
+                  value={apiKey}
+                />
+              )}
             </Col>
             <Col className="mx-2">
             <Text>Select Model:</Text>