fix(create_key_button.tsx): retrieve all available team models on team select

This commit is contained in:
Krrish Dholakia 2025-03-14 19:50:49 -07:00
parent fe148e3dec
commit db4be0c24a
2 changed files with 42 additions and 25 deletions

View file

@ -91,28 +91,31 @@ const getPredefinedTags = (data: any[] | null) => {
return uniqueTags;
}
export const getTeamModels = (team: Team | null, allAvailableModels: string[]): string[] => {
let tempModelsToPick = [];
if (team) {
if (team.models.length > 0) {
if (team.models.includes("all-proxy-models")) {
// if the team has all-proxy-models show all available models
tempModelsToPick = allAvailableModels;
} else {
// show team models
tempModelsToPick = team.models;
}
} else {
// show all available models if the team has no models set
tempModelsToPick = allAvailableModels;
export const fetchTeamModels = async (userID: string, userRole: string, accessToken: string, teamID: string): Promise<string[]> => {
try {
if (userID === null || userRole === null) {
return [];
}
} else {
// no team set, show all available models
tempModelsToPick = allAvailableModels;
}
return unfurlWildcardModelsInList(tempModelsToPick, allAvailableModels);
if (accessToken !== null) {
const model_available = await modelAvailableCall(
accessToken,
userID,
userRole,
true,
teamID
);
let available_model_names = model_available["data"].map(
(element: { id: string }) => element.id
);
console.log("available_model_names:", available_model_names);
return available_model_names;
}
return [];
} catch (error) {
console.error("Error fetching user models:", error);
return [];
}
};
export const fetchUserModels = async (userID: string, userRole: string, accessToken: string, setUserModels: (models: string[]) => void) => {
@ -182,6 +185,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({
}
}, [accessToken, userID, userRole]);
useEffect(() => {
const fetchGuardrails = async () => {
try {
@ -277,10 +281,14 @@ const CreateKey: React.FC<CreateKeyProps> = ({
};
useEffect(() => {
const models = getTeamModels(selectedCreateKeyTeam, userModels);
setModelsToPick(models);
if (userID && userRole && accessToken && selectedCreateKeyTeam) {
fetchTeamModels(userID, userRole, accessToken, selectedCreateKeyTeam.team_id).then((models) => {
let allModels = Array.from(new Set([...selectedCreateKeyTeam.models, ...models]));
setModelsToPick(allModels);
});
}
form.setFieldValue('models', []);
}, [selectedCreateKeyTeam, userModels]);
}, [selectedCreateKeyTeam]);
// Add a callback function to handle user creation
const handleUserCreated = (userId: string) => {

View file

@ -1284,6 +1284,7 @@ export const modelInfoV1Call = async (accessToken: String, modelId: String) => {
}
};
export const modelHubCall = async (accessToken: String) => {
/**
* Get all models on proxy
@ -1581,7 +1582,8 @@ export const modelAvailableCall = async (
accessToken: String,
userID: String,
userRole: String,
return_wildcard_routes: boolean = false
return_wildcard_routes: boolean = false,
teamID: String | null = null
) => {
/**
* Get all the models user has access to
@ -1589,8 +1591,15 @@ export const modelAvailableCall = async (
console.log("in /models calls, globalLitellmHeaderName", globalLitellmHeaderName)
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/models` : `/models`;
const params = new URLSearchParams();
if (return_wildcard_routes === true) {
url += `?return_wildcard_routes=True`;
params.append('return_wildcard_routes', 'True');
}
if (teamID) {
params.append('team_id', teamID.toString());
}
if (params.toString()) {
url += `?${params.toString()}`;
}
//message.info("Requesting model data");