phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(proxy_server.py): enable default new user params

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-23 16:39:50 -08:00
parent 1e769bcb77
commit 0806aa8da1
7 changed files with 462 additions and 389 deletions
Download patch file Download diff file Expand all files Collapse all files

1
litellm/__init__.py
View file

@ -165,6 +165,7 @@ s3_callback_params: Optional[Dict] = None
generic_logger_headers: Optional[Dict] = None
default_key_generate_params: Optional[Dict] = None
upperbound_key_generate_params: Optional[Dict] = None
default_user_params: Optional[Dict] = None
default_team_settings: Optional[List] = None
max_user_budget: Optional[float] = None
#### RELIABILITY ####

60
litellm/proxy/proxy_server.py
View file

@ -730,6 +730,8 @@ async def user_api_key_auth(
"/user",
"/model/info",
"/v2/model/info",
"/models",
"/v1/models",
]
# check if the current route startswith any of the allowed routes
if (
@ -1758,6 +1760,7 @@ async def generate_key_helper_fn(
allowed_cache_controls: Optional[list] = [],
permissions: Optional[dict] = {},
model_max_budget: Optional[dict] = {},
table_name: Optional[Literal["key", "user"]] = None,
):
global prisma_client, custom_db_client, user_api_key_cache
@ -1884,8 +1887,10 @@ async def generate_key_helper_fn(
table_name="user",
update_key_values=update_key_values,
)
if user_id == litellm_proxy_budget_name:
# do not create a key for litellm_proxy_budget_name
if user_id == litellm_proxy_budget_name or (
table_name is not None and table_name == "user"
):
# do not create a key for litellm_proxy_budget_name or if table name is set to just 'user'
# we only need to ensure this exists in the user table
# the LiteLLM_VerificationToken table will increase in size if we don't do this check
return key_data
@ -5461,27 +5466,50 @@ async def auth_callback(request: Request):
user_id_models: List = []
# User might not be already created on first generation of key
# But if it is, we want its models preferences
try:
if prisma_client is not None:
user_info = await prisma_client.get_data(user_id=user_id, table_name="user")
if user_info is not None:
user_id_models = getattr(user_info, "models", [])
except Exception as e:
pass
response = await generate_key_helper_fn(
**{
# But if it is, we want their models preferences
default_ui_key_values = {
"duration": "1hr",
"key_max_budget": 0.01,
"models": user_id_models,
"aliases": {},
"config": {},
"spend": 0,
"user_id": user_id,
"team_id": "litellm-dashboard",
}
user_defined_values = {
"models": user_id_models,
"user_id": user_id,
"user_email": user_email,
} # type: ignore
}
try:
if prisma_client is not None:
user_info = await prisma_client.get_data(user_id=user_id, table_name="user")
verbose_proxy_logger.debug(
f"user_info: {user_info}; litellm.default_user_params: {litellm.default_user_params}"
)
if user_info is not None:
user_defined_values = {
"models": getattr(user_info, "models", []),
"user_id": getattr(user_info, "user_id", user_id),
"user_email": getattr(user_info, "user_id", user_email),
}
elif litellm.default_user_params is not None and isinstance(
litellm.default_user_params, dict
):
user_defined_values = {
"models": litellm.default_user_params.get("models", user_id_models),
"user_id": litellm.default_user_params.get("user_id", user_id),
"user_email": litellm.default_user_params.get(
"user_email", user_email
),
}
except Exception as e:
pass
verbose_proxy_logger.info(
f"user_defined_values for creating ui key: {user_defined_values}"
)
response = await generate_key_helper_fn(
**default_ui_key_values, **user_defined_values # type: ignore
)
key = response["token"] # type: ignore
user_id = response["user_id"] # type: ignore

95
ui/litellm-dashboard/src/components/chat_ui.tsx
View file

@ -1,17 +1,26 @@
import React, { useState, useEffect } from "react";
import ReactMarkdown from "react-markdown";
import { Card, Title, Table, TableHead, TableRow, TableCell, TableBody, Grid, Tab,
import {
Card,
Title,
Table,
TableHead,
TableRow,
TableCell,
TableBody,
Grid,
Tab,
TabGroup,
TabList,
TabPanel,
Metric,
Select,
SelectItem,
TabPanels, } from "@tremor/react";
import { modelInfoCall } from "./networking";
TabPanels,
} from "@tremor/react";
import { modelAvailableCall } from "./networking";
import openai from "openai";
import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
interface ChatUIProps {
accessToken: string | null;
@ -20,11 +29,18 @@ interface ChatUIProps {
userID: string | null;
}
async function generateModelResponse(inputMessage: string, updateUI: (chunk: string) => void, selectedModel: string, accessToken: string) {
async function generateModelResponse(
inputMessage: string,
updateUI: (chunk: string) => void,
selectedModel: string,
accessToken: string
) {
// base url should be the current base_url
const isLocal = process.env.NODE_ENV === "development";
console.log("isLocal:", isLocal);
const proxyBaseUrl = isLocal ? "http://localhost:4000" : window.location.origin;
const proxyBaseUrl = isLocal
? "http://localhost:4000"
: window.location.origin;
const client = new openai.OpenAI({
apiKey: accessToken, // Replace with your OpenAI API key
baseURL: proxyBaseUrl, // Replace with your OpenAI API base URL
@ -36,7 +52,7 @@ async function generateModelResponse(inputMessage: string, updateUI: (chunk: str
stream: true,
messages: [
{
role: 'user',
role: "user",
content: inputMessage,
},
],
@ -50,10 +66,17 @@ async function generateModelResponse(inputMessage: string, updateUI: (chunk: str
}
}
const ChatUI: React.FC<ChatUIProps> = ({ accessToken, token, userRole, userID }) => {
const ChatUI: React.FC<ChatUIProps> = ({
accessToken,
token,
userRole,
userID,
}) => {
const [inputMessage, setInputMessage] = useState("");
const [chatHistory, setChatHistory] = useState<any[]>([]);
const [selectedModel, setSelectedModel] = useState<string | undefined>(undefined);
const [selectedModel, setSelectedModel] = useState<string | undefined>(
undefined
);
const [modelInfo, setModelInfo] = useState<any | null>(null); // Declare modelInfo at the component level
useEffect(() => {
@ -62,12 +85,16 @@ const ChatUI: React.FC<ChatUIProps> = ({ accessToken, token, userRole, userID })
}
// Fetch model info and set the default selected model
const fetchModelInfo = async () => {
const fetchedModelInfo = await modelInfoCall(accessToken, userID, userRole);
console.log("model_info:", fetchedModelInfo);
const fetchedAvailableModels = await modelAvailableCall(
accessToken,
userID,
userRole
);
console.log("model_info:", fetchedAvailableModels);
if (fetchedModelInfo?.data.length > 0) {
setModelInfo(fetchedModelInfo);
setSelectedModel(fetchedModelInfo.data[0].model_name);
if (fetchedAvailableModels?.data.length > 0) {
setModelInfo(fetchedAvailableModels.data);
setSelectedModel(fetchedAvailableModels.data[0].id);
}
};
@ -103,7 +130,12 @@ const ChatUI: React.FC<ChatUIProps> = ({ accessToken, token, userRole, userID })
try {
if (selectedModel) {
await generateModelResponse(inputMessage, (chunk) => updateUI("assistant", chunk), selectedModel, accessToken);
await generateModelResponse(
inputMessage,
(chunk) => updateUI("assistant", chunk),
selectedModel,
accessToken
);
}
} catch (error) {
console.error("Error fetching model response", error);
@ -132,14 +164,21 @@ const ChatUI: React.FC<ChatUIProps> = ({ accessToken, token, userRole, userID })
onChange={(e) => setSelectedModel(e.target.value)}
>
{/* Populate dropdown options from available models */}
{modelInfo?.data.map((element: { model_name: string }) => (
<option key={element.model_name} value={element.model_name}>
{element.model_name}
{modelInfo?.map((element: { id: string }) => (
<option key={element.id} value={element.id}>
{element.id}
</option>
))}
</select>
</div>
<Table className="mt-5" style={{ display: "block", maxHeight: "60vh", overflowY: "auto" }}>
<Table
className="mt-5"
style={{
display: "block",
maxHeight: "60vh",
overflowY: "auto",
}}
>
<TableHead>
<TableRow>
<TableCell>
@ -155,7 +194,10 @@ const ChatUI: React.FC<ChatUIProps> = ({ accessToken, token, userRole, userID })
))}
</TableBody>
</Table>
<div className="mt-3" style={{ position: "absolute", bottom: 5, width: "95%" }}>
<div
className="mt-3"
style={{ position: "absolute", bottom: 5, width: "95%" }}
>
<div className="flex">
<input
type="text"
@ -164,7 +206,10 @@ const ChatUI: React.FC<ChatUIProps> = ({ accessToken, token, userRole, userID })
className="flex-1 p-2 border rounded-md mr-2"
placeholder="Type your message..."
/>
<button onClick={handleSendMessage} className="p-2 bg-blue-500 text-white rounded-md">
<button
onClick={handleSendMessage}
className="p-2 bg-blue-500 text-white rounded-md"
>
Send
</button>
</div>
@ -179,7 +224,6 @@ const ChatUI: React.FC<ChatUIProps> = ({ accessToken, token, userRole, userID })
</TabList>
<TabPanels>
<TabPanel>
<SyntaxHighlighter language="python">
{`
import openai
@ -211,7 +255,6 @@ print(response)
</SyntaxHighlighter>
</TabPanel>
<TabPanel>
<SyntaxHighlighter language="python">
{`
import os, dotenv
@ -248,7 +291,6 @@ print(response)
</SyntaxHighlighter>
</TabPanel>
<TabPanel>
<SyntaxHighlighter language="python">
{`
from langchain.chat_models import ChatOpenAI
@ -290,7 +332,6 @@ print(response)
</TabPanel>
</TabPanels>
</TabGroup>
</TabPanel>
</TabPanels>
</TabGroup>
@ -300,6 +341,4 @@ print(response)
);
};
export default ChatUI;

90
ui/litellm-dashboard/src/components/create_key_button.tsx
View file

@ -1,10 +1,17 @@
"use client";
import React, { useState, useEffect, useRef } from "react";
import { Button, TextInput, Grid, Col } from "@tremor/react";
import { Card, Metric, Text } from "@tremor/react";
import { Button as Button2, Modal, Form, Input, InputNumber, Select, message } from "antd";
import {
Button as Button2,
Modal,
Form,
Input,
InputNumber,
Select,
message,
} from "antd";
import { keyCreateCall } from "./networking";
const { Option } = Select;
@ -50,13 +57,12 @@ const CreateKey: React.FC<CreateKeyProps> = ({
setApiKey(response["key"]);
message.success("API Key Created");
form.resetFields();
localStorage.removeItem("userData" + userID)
localStorage.removeItem("userData" + userID);
} catch (error) {
console.error("Error creating the key:", error);
}
};
return (
<div>
<Button className="mx-auto" onClick={() => setIsModalVisible(true)}>
@ -70,30 +76,26 @@ const CreateKey: React.FC<CreateKeyProps> = ({
onOk={handleOk}
onCancel={handleCancel}
>
<Form form={form} onFinish={handleCreate} labelCol={{ span: 8 }} wrapperCol={{ span: 16 }} labelAlign="left">
{userRole === 'App Owner' || userRole === 'Admin' ? (
<>
<Form.Item
label="Key Name"
name="key_alias"
<Form
form={form}
onFinish={handleCreate}
labelCol={{ span: 8 }}
wrapperCol={{ span: 16 }}
labelAlign="left"
>
{userRole === "App Owner" || userRole === "Admin" ? (
<>
<Form.Item label="Key Name" name="key_alias">
<Input />
</Form.Item>
<Form.Item
label="Team ID"
name="team_id"
>
<Form.Item label="Team ID" name="team_id">
<Input placeholder="ai_team" />
</Form.Item>
<Form.Item
label="Models"
name="models"
>
<Form.Item label="Models" name="models">
<Select
mode="multiple"
placeholder="Select models"
style={{ width: '100%' }}
style={{ width: "100%" }}
>
{userModels.map((model) => (
<Option key={model} value={model}>
@ -102,16 +104,10 @@ const CreateKey: React.FC<CreateKeyProps> = ({
))}
</Select>
</Form.Item>
<Form.Item
label="Max Budget (USD)"
name="max_budget"
>
<Form.Item label="Max Budget (USD)" name="max_budget">
<InputNumber step={0.01} precision={2} width={200} />
</Form.Item>
<Form.Item
label="Tokens per minute Limit (TPM)"
name="tpm_limit"
>
<Form.Item label="Tokens per minute Limit (TPM)" name="tpm_limit">
<InputNumber step={1} width={400} />
</Form.Item>
<Form.Item
@ -120,47 +116,29 @@ const CreateKey: React.FC<CreateKeyProps> = ({
>
<InputNumber step={1} width={400} />
</Form.Item>
<Form.Item
label="Duration (eg: 30s, 30h, 30d)"
name="duration"
>
<Form.Item label="Duration (eg: 30s, 30h, 30d)" name="duration">
<Input />
</Form.Item>
<Form.Item
label="Metadata"
name="metadata"
>
<Form.Item label="Metadata" name="metadata">
<Input.TextArea rows={4} placeholder="Enter metadata as JSON" />
</Form.Item>
</>
) : (
<>
<Form.Item
label="Key Name"
name="key_alias"
>
<Form.Item label="Key Name" name="key_alias">
<Input />
</Form.Item>
<Form.Item
label="Team ID (Contact Group)"
name="team_id"
>
<Form.Item label="Team ID (Contact Group)" name="team_id">
<Input placeholder="ai_team" />
</Form.Item>
<Form.Item
label="Description"
name="description"
>
<Form.Item label="Description" name="description">
<Input.TextArea placeholder="Enter description" rows={4} />
</Form.Item>
</>
)
}
<div style={{ textAlign: 'right', marginTop: '10px' }}>
<Button2 htmlType="submit">
Create Key
</Button2>
)}
<div style={{ textAlign: "right", marginTop: "10px" }}>
<Button2 htmlType="submit">Create Key</Button2>
</div>
</Form>
</Modal>
@ -177,8 +155,8 @@ const CreateKey: React.FC<CreateKeyProps> = ({
<p>
Please save this secret key somewhere safe and accessible. For
security reasons, <b>you will not be able to view it again</b>{" "}
through your LiteLLM account. If you lose this secret key, you will
need to generate a new one.
through your LiteLLM account. If you lose this secret key, you
will need to generate a new one.
</p>
</Col>
<Col numColSpan={1}>

82
ui/litellm-dashboard/src/components/create_user_button.tsx
View file

@ -1,7 +1,7 @@
import React, { useState, useEffect } from "react";
import { Button, Modal, Form, Input, message, Select, InputNumber } from "antd";
import { Button as Button2 } from "@tremor/react";
import { userCreateCall, modelInfoCall } from "./networking";
import { userCreateCall, modelAvailableCall } from "./networking";
const { Option } = Select;
interface CreateuserProps {
@ -20,12 +20,16 @@ const Createuser: React.FC<CreateuserProps> = ({ userID, accessToken }) => {
const fetchData = async () => {
try {
const userRole = "any"; // You may need to get the user role dynamically
const modelDataResponse = await modelInfoCall(accessToken, userID, userRole);
const modelDataResponse = await modelAvailableCall(
accessToken,
userID,
userRole
);
// Assuming modelDataResponse.data contains an array of model objects with a 'model_name' property
const availableModels = [];
for (let i = 0; i < modelDataResponse.data.length; i++) {
const model = modelDataResponse.data[i];
availableModels.push(model.model_name);
availableModels.push(model.id);
}
console.log("Model data response:", modelDataResponse.data);
console.log("Available models:", availableModels);
@ -79,27 +83,24 @@ const Createuser: React.FC<CreateuserProps> = ({ userID, accessToken }) => {
onOk={handleOk}
onCancel={handleCancel}
>
<Form form={form} onFinish={handleCreate} labelCol={{ span: 8 }} wrapperCol={{ span: 16 }} labelAlign="left">
<Form.Item
label="User ID"
name="user_id"
<Form
form={form}
onFinish={handleCreate}
labelCol={{ span: 8 }}
wrapperCol={{ span: 16 }}
labelAlign="left"
>
<Form.Item label="User ID" name="user_id">
<Input placeholder="Enter User ID" />
</Form.Item>
<Form.Item
label="Team ID"
name="team_id"
>
<Form.Item label="Team ID" name="team_id">
<Input placeholder="ai_team" />
</Form.Item>
<Form.Item
label="Models"
name="models"
>
<Form.Item label="Models" name="models">
<Select
mode="multiple"
placeholder="Select models"
style={{ width: '100%' }}
style={{ width: "100%" }}
>
{userModels.map((model) => (
<Option key={model} value={model}>
@ -109,46 +110,23 @@ const Createuser: React.FC<CreateuserProps> = ({ userID, accessToken }) => {
</Select>
</Form.Item>
<Form.Item
label="Max Budget (USD)"
name="max_budget"
>
<Form.Item label="Max Budget (USD)" name="max_budget">
<InputNumber step={0.01} precision={2} width={200} />
</Form.Item>
<Form.Item
label="Tokens per minute Limit (TPM)"
name="tpm_limit"
>
<Form.Item label="Tokens per minute Limit (TPM)" name="tpm_limit">
<InputNumber step={1} width={400} />
</Form.Item>
<Form.Item
label="Requests per minute Limit (RPM)"
name="rpm_limit"
>
<Form.Item label="Requests per minute Limit (RPM)" name="rpm_limit">
<InputNumber step={1} width={400} />
</Form.Item>
<Form.Item
label="Duration (eg: 30s, 30h, 30d)"
name="duration"
>
<Form.Item label="Duration (eg: 30s, 30h, 30d)" name="duration">
<Input />
</Form.Item>
<Form.Item
label="Metadata"
name="metadata"
>
<Form.Item label="Metadata" name="metadata">
<Input.TextArea rows={4} placeholder="Enter metadata as JSON" />
</Form.Item>
<div style={{ textAlign: 'right', marginTop: '10px' }}>
<Button htmlType="submit">
Create User
</Button>
<div style={{ textAlign: "right", marginTop: "10px" }}>
<Button htmlType="submit">Create User</Button>
</div>
</Form>
</Modal>
@ -162,11 +140,15 @@ const Createuser: React.FC<CreateuserProps> = ({ userID, accessToken }) => {
>
<p>
Please save this secret user somewhere safe and accessible. For
security reasons, <b>you will not be able to view it again</b> through
your LiteLLM account. If you lose this secret user, you will need to
generate a new one.
security reasons, <b>you will not be able to view it again</b>{" "}
through your LiteLLM account. If you lose this secret user, you will
need to generate a new one.
</p>
<p>
{apiuser != null
? `API user: ${apiuser}`
: "User being created, this might take 30s"}
</p>
<p>{apiuser != null ? `API user: ${apiuser}` : "User being created, this might take 30s"}</p>
</Modal>
)}
</div>

62
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -69,7 +69,6 @@ export const keyCreateCall = async (
}
};
export const userCreateCall = async (
accessToken: string,
userID: string,
@ -133,7 +132,6 @@ export const userCreateCall = async (
}
};
export const keyDeleteCall = async (accessToken: String, user_key: String) => {
try {
const url = proxyBaseUrl ? `${proxyBaseUrl}/key/delete` : `/key/delete`;
@ -207,13 +205,14 @@ export const userInfoCall = async (
}
};
export const modelInfoCall = async (
accessToken: String,
userID: String,
userRole: String
) => {
/**
* Get all models on proxy
*/
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/v2/model/info` : `/v2/model/info`;
@ -242,6 +241,42 @@ export const modelInfoCall = async (
}
};
export const modelAvailableCall = async (
accessToken: String,
userID: String,
userRole: String
) => {
/**
* Get all the models user has access to
*/
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/models` : `/models`;
message.info("Requesting model data");
const response = await fetch(url, {
method: "GET",
headers: {
Authorization: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
});
if (!response.ok) {
const errorData = await response.text();
message.error(errorData);
throw new Error("Network response was not ok");
}
const data = await response.json();
message.info("Received model 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 create key:", error);
throw error;
}
};
export const keySpendLogsCall = async (accessToken: String, token: String) => {
try {
const url = proxyBaseUrl ? `${proxyBaseUrl}/spend/logs` : `/spend/logs`;
@ -363,12 +398,16 @@ export const spendUsersCall = async (accessToken: String, userID: String) => {
}
};
export const userRequestModelCall = async (accessToken: String, model: String, UserID: String, justification: String) => {
export const userRequestModelCall = async (
accessToken: String,
model: String,
UserID: String,
justification: String
) => {
try {
const url = proxyBaseUrl ? `${proxyBaseUrl}/user/request_model` : `/user/request_model`;
const url = proxyBaseUrl
? `${proxyBaseUrl}/user/request_model`
: `/user/request_model`;
const response = await fetch(url, {
method: "POST",
headers: {
@ -398,10 +437,11 @@ export const userRequestModelCall = async (accessToken: String, model: String, U
}
};
export const userGetRequesedtModelsCall = async (accessToken: String) => {
try {
const url = proxyBaseUrl ? `${proxyBaseUrl}/user/get_requests` : `/user/get_requests`;
const url = proxyBaseUrl
? `${proxyBaseUrl}/user/get_requests`
: `/user/get_requests`;
console.log("in userGetRequesedtModelsCall:", url);
const response = await fetch(url, {
method: "GET",

31
ui/litellm-dashboard/src/components/user_dashboard.tsx
View file

@ -1,6 +1,6 @@
"use client";
import React, { useState, useEffect } from "react";
import { userInfoCall, modelInfoCall } from "./networking";
import { userInfoCall, modelAvailableCall } from "./networking";
import { Grid, Col, Card, Text } from "@tremor/react";
import CreateKey from "./create_key_button";
import ViewKeyTable from "./view_key_table";
@ -48,10 +48,9 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
const token = searchParams.get("token");
const [accessToken, setAccessToken] = useState<string | null>(null);
const [userModels, setUserModels] = useState<string[]>([]);
// check if window is not undefined
if (typeof window !== "undefined") {
window.addEventListener('beforeunload', function() {
window.addEventListener("beforeunload", function () {
// Clear session storage
sessionStorage.clear();
});
@ -78,7 +77,6 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
// Moved useEffect inside the component and used a condition to run fetch only if the params are available
useEffect(() => {
if (token) {
const decoded = jwtDecode(token) as { [key: string]: any };
if (decoded) {
@ -109,32 +107,39 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
const cachedUserModels = sessionStorage.getItem("userModels" + userID);
if (cachedUserModels) {
setUserModels(JSON.parse(cachedUserModels));
} else {
const fetchData = async () => {
try {
const response = await userInfoCall(accessToken, userID, userRole);
setUserSpendData(response["user_info"]);
setData(response["keys"]); // Assuming this is the correct path to your data
sessionStorage.setItem("userData" + userID, JSON.stringify(response["keys"]));
sessionStorage.setItem(
"userData" + userID,
JSON.stringify(response["keys"])
);
sessionStorage.setItem(
"userSpendData" + userID,
JSON.stringify(response["user_info"])
);
const model_info = await modelInfoCall(accessToken, userID, userRole);
console.log("model_info:", model_info);
const model_available = await modelAvailableCall(
accessToken,
userID,
userRole
);
// loop through model_info["data"] and create an array of element.model_name
let available_model_names = model_info["data"].filter((element: { model_name: string; user_access: boolean }) => element.user_access === true).map((element: { model_name: string; }) => element.model_name);
let available_model_names = model_available["data"].map(
(element: { id: string }) => element.id
);
console.log("available_model_names:", available_model_names);
setUserModels(available_model_names);
console.log("userModels:", userModels);
sessionStorage.setItem("userModels" + userID, JSON.stringify(available_model_names));
sessionStorage.setItem(
"userModels" + userID,
JSON.stringify(available_model_names)
);
} catch (error) {
console.error("There was an error fetching the data", error);
// Optionally, update your UI to reflect the error state here as well