소스 검색

feat: model filter frontend integration

Timothy J. Baek 1 년 전
부모
커밋
81daf4ceb9
4개의 변경된 파일83개의 추가작업 그리고 6개의 파일을 삭제
  1. 6 3
      backend/apps/ollama/main.py
  2. 5 3
      backend/apps/openai/main.py
  3. 62 0
      src/lib/apis/index.ts
  4. 10 0
      src/lib/components/admin/Settings/Users.svelte

+ 6 - 3
backend/apps/ollama/main.py

@@ -135,11 +135,14 @@ async def get_ollama_tags(
 ):
 ):
     if url_idx == None:
     if url_idx == None:
         models = await get_all_models()
         models = await get_all_models()
+
         if app.state.MODEL_FILTER_ENABLED:
         if app.state.MODEL_FILTER_ENABLED:
             if user.role == "user":
             if user.role == "user":
-                models["models"] = filter(
-                    lambda model: model["name"] in app.state.MODEL_LIST,
-                    models["models"],
+                models["models"] = list(
+                    filter(
+                        lambda model: model["name"] in app.state.MODEL_LIST,
+                        models["models"],
+                    )
                 )
                 )
                 return models
                 return models
         return models
         return models

+ 5 - 3
backend/apps/openai/main.py

@@ -196,9 +196,11 @@ async def get_models(url_idx: Optional[int] = None, user=Depends(get_current_use
         models = await get_all_models()
         models = await get_all_models()
         if app.state.MODEL_FILTER_ENABLED:
         if app.state.MODEL_FILTER_ENABLED:
             if user.role == "user":
             if user.role == "user":
-                models["data"] = filter(
-                    lambda model: model["id"] in app.state.MODEL_LIST,
-                    models["data"],
+                models["data"] = list(
+                    filter(
+                        lambda model: model["id"] in app.state.MODEL_LIST,
+                        models["data"],
+                    )
                 )
                 )
                 return models
                 return models
         return models
         return models

+ 62 - 0
src/lib/apis/index.ts

@@ -77,3 +77,65 @@ export const getVersionUpdates = async () => {
 
 
 	return res;
 	return res;
 };
 };
+
+export const getModelFilterConfig = async (token: string) => {
+	let error = null;
+
+	const res = await fetch(`${WEBUI_BASE_URL}/api/config/model/filter`, {
+		method: 'GET',
+		headers: {
+			'Content-Type': 'application/json',
+			Authorization: `Bearer ${token}`
+		}
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.log(err);
+			error = err;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
+export const updateModelFilterConfig = async (
+	token: string,
+	enabled: boolean,
+	models: string[]
+) => {
+	let error = null;
+
+	const res = await fetch(`${WEBUI_BASE_URL}/api/config/model/filter`, {
+		method: 'POST',
+		headers: {
+			'Content-Type': 'application/json',
+			Authorization: `Bearer ${token}`
+		},
+		body: JSON.stringify({
+			enabled: enabled,
+			models: models
+		})
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.log(err);
+			error = err;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};

+ 10 - 0
src/lib/components/admin/Settings/Users.svelte

@@ -1,4 +1,5 @@
 <script lang="ts">
 <script lang="ts">
+	import { getModelFilterConfig, updateModelFilterConfig } from '$lib/apis';
 	import { getSignUpEnabledStatus, toggleSignUpEnabledStatus } from '$lib/apis/auths';
 	import { getSignUpEnabledStatus, toggleSignUpEnabledStatus } from '$lib/apis/auths';
 	import { getUserPermissions, updateUserPermissions } from '$lib/apis/users';
 	import { getUserPermissions, updateUserPermissions } from '$lib/apis/users';
 	import { models } from '$lib/stores';
 	import { models } from '$lib/stores';
@@ -16,6 +17,13 @@
 
 
 	onMount(async () => {
 	onMount(async () => {
 		permissions = await getUserPermissions(localStorage.token);
 		permissions = await getUserPermissions(localStorage.token);
+
+		const res = await getModelFilterConfig(localStorage.token);
+		if (res) {
+			whitelistEnabled = res.enabled;
+
+			whitelistModels = res.models.length > 0 ? res.models : [''];
+		}
 	});
 	});
 </script>
 </script>
 
 
@@ -24,6 +32,8 @@
 	on:submit|preventDefault={async () => {
 	on:submit|preventDefault={async () => {
 		// console.log('submit');
 		// console.log('submit');
 		await updateUserPermissions(localStorage.token, permissions);
 		await updateUserPermissions(localStorage.token, permissions);
+
+		await updateModelFilterConfig(localStorage.token, whitelistEnabled, whitelistModels);
 		saveHandler();
 		saveHandler();
 	}}
 	}}
 >
 >