From 9dae7d0572831c564c9afdfb36bacb1ac9e74bcf Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Fri, 15 Nov 2024 18:21:41 -0800 Subject: [PATCH] refac: models --- .../open_webui/apps/webui/models/models.py | 26 +++- .../open_webui/apps/webui/routers/models.py | 35 ++++++ .../versions/922e7a387820_add_group_table.py | 14 +++ src/lib/apis/models/index.ts | 37 ++++++ .../workspace/Knowledge/Collection.svelte | 19 ++- .../Knowledge/CreateCollection.svelte | 5 + src/lib/components/workspace/Models.svelte | 117 ++++-------------- .../workspace/Models/ModelEditor.svelte | 8 +- .../workspace/Prompts/PromptEditor.svelte | 5 + .../workspace/Tools/ToolkitEditor.svelte | 16 +-- ...sionsModal.svelte => AccessControl.svelte} | 23 ++-- 11 files changed, 186 insertions(+), 119 deletions(-) rename src/lib/components/workspace/common/{AccessPermissionsModal.svelte => AccessControl.svelte} (89%) diff --git a/backend/open_webui/apps/webui/models/models.py b/backend/open_webui/apps/webui/models/models.py index 77b7c5f67..6434cfb16 100644 --- a/backend/open_webui/apps/webui/models/models.py +++ b/backend/open_webui/apps/webui/models/models.py @@ -12,7 +12,7 @@ from pydantic import BaseModel, ConfigDict from sqlalchemy import or_, and_, func from sqlalchemy.dialects import postgresql, sqlite -from sqlalchemy import BigInteger, Column, Text, JSON +from sqlalchemy import BigInteger, Column, Text, JSON, Boolean from open_webui.utils.utils import has_access @@ -95,6 +95,8 @@ class Model(Base): # } # } + is_active = Column(Boolean, default=True) + updated_at = Column(BigInteger) created_at = Column(BigInteger) @@ -110,6 +112,7 @@ class ModelModel(BaseModel): access_control: Optional[dict] = None + is_active: bool updated_at: int # timestamp in epoch created_at: int # timestamp in epoch @@ -131,6 +134,8 @@ class ModelResponse(BaseModel): meta: ModelMeta access_control: Optional[dict] = None + + is_active: bool updated_at: int # timestamp in epoch created_at: int # timestamp in epoch @@ -141,6 +146,8 @@ class ModelForm(BaseModel): name: str meta: ModelMeta params: ModelParams + access_control: Optional[dict] = None + is_active: bool = True class ModelsTable: @@ -200,6 +207,23 @@ class ModelsTable: except Exception: return None + def toggle_model_by_id(self, id: str) -> Optional[ModelModel]: + with get_db() as db: + try: + is_active = db.query(Model).filter_by(id=id).first().is_active + + db.query(Model).filter_by(id=id).update( + { + "is_active": not is_active, + "updated_at": int(time.time()), + } + ) + db.commit() + + return self.get_model_by_id(id) + except Exception: + return None + def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]: try: with get_db() as db: diff --git a/backend/open_webui/apps/webui/routers/models.py b/backend/open_webui/apps/webui/routers/models.py index 86b8515fd..7ba8d8190 100644 --- a/backend/open_webui/apps/webui/routers/models.py +++ b/backend/open_webui/apps/webui/routers/models.py @@ -79,6 +79,41 @@ async def get_model_by_id(id: str, user=Depends(get_verified_user)): ) +############################ +# ToggelModelById +############################ + + +@router.post("/id/{id}/toggle", response_model=Optional[ModelResponse]) +async def toggle_model_by_id(id: str, user=Depends(get_verified_user)): + model = Models.get_model_by_id(id) + if model: + if ( + user.role == "admin" + or model.user_id == user.id + or has_access(user.id, "write", model.access_control) + ): + model = Models.toggle_model_by_id(id) + + if model: + return model + else: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=ERROR_MESSAGES.DEFAULT("Error updating function"), + ) + else: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=ERROR_MESSAGES.UNAUTHORIZED, + ) + else: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail=ERROR_MESSAGES.NOT_FOUND, + ) + + ############################ # UpdateModelById ############################ diff --git a/backend/open_webui/migrations/versions/922e7a387820_add_group_table.py b/backend/open_webui/migrations/versions/922e7a387820_add_group_table.py index f349e3593..bdab303a7 100644 --- a/backend/open_webui/migrations/versions/922e7a387820_add_group_table.py +++ b/backend/open_webui/migrations/versions/922e7a387820_add_group_table.py @@ -35,6 +35,17 @@ def upgrade(): sa.Column("access_control", sa.JSON(), nullable=True), ) + # Add 'is_active' column to 'model' table + op.add_column( + "model", + sa.Column( + "is_active", + sa.Boolean(), + nullable=False, + server_default=sa.sql.expression.true(), + ), + ) + # Add 'access_control' column to 'knowledge' table op.add_column( "knowledge", @@ -60,6 +71,9 @@ def downgrade(): # Drop 'access_control' column from 'model' table op.drop_column("model", "access_control") + # Drop 'is_active' column from 'model' table + op.drop_column("model", "is_active") + # Drop 'access_control' column from 'knowledge' table op.drop_column("knowledge", "access_control") diff --git a/src/lib/apis/models/index.ts b/src/lib/apis/models/index.ts index 14014ce99..86aeb2d89 100644 --- a/src/lib/apis/models/index.ts +++ b/src/lib/apis/models/index.ts @@ -100,6 +100,43 @@ export const getModelById = async (token: string, id: string) => { return res; }; + +export const toggleModelById = async (token: string, id: string) => { + let error = null; + + const searchParams = new URLSearchParams(); + searchParams.append('id', id); + + const res = await fetch(`${WEBUI_API_BASE_URL}/models/id/${id}/toggle`, { + method: 'POST', + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + authorization: `Bearer ${token}` + } + }) + .then(async (res) => { + if (!res.ok) throw await res.json(); + return res.json(); + }) + .then((json) => { + return json; + }) + .catch((err) => { + error = err; + + console.log(err); + return null; + }); + + if (error) { + throw error; + } + + return res; +}; + + export const updateModelById = async (token: string, id: string, model: object) => { let error = null; diff --git a/src/lib/components/workspace/Knowledge/Collection.svelte b/src/lib/components/workspace/Knowledge/Collection.svelte index cf69f7f37..d2f0318bf 100644 --- a/src/lib/components/workspace/Knowledge/Collection.svelte +++ b/src/lib/components/workspace/Knowledge/Collection.svelte @@ -39,6 +39,7 @@ import Drawer from '$lib/components/common/Drawer.svelte'; import ChevronLeft from '$lib/components/icons/ChevronLeft.svelte'; import MenuLines from '$lib/components/icons/MenuLines.svelte'; + import AccessControl from '../common/AccessControl.svelte'; let largeScreen = true; @@ -687,7 +688,17 @@ /> {:else} -
{$i18n.t('No content found')}
+
+
+ {$i18n.t('No content found')} +
+ +
+ {$i18n.t('Drag and drop a file to upload or select a file to view')} +
+
{/if} @@ -753,7 +764,7 @@ {:else} -
+
@@ -784,8 +795,8 @@
-
- {$i18n.t('Select a file to view or drag and drop a file to upload')} +
+
{/if} diff --git a/src/lib/components/workspace/Knowledge/CreateCollection.svelte b/src/lib/components/workspace/Knowledge/CreateCollection.svelte index 89aaca30c..8700f59cf 100644 --- a/src/lib/components/workspace/Knowledge/CreateCollection.svelte +++ b/src/lib/components/workspace/Knowledge/CreateCollection.svelte @@ -6,6 +6,7 @@ import { createNewKnowledge, getKnowledgeItems } from '$lib/apis/knowledge'; import { toast } from 'svelte-sonner'; import { knowledge } from '$lib/stores'; + import AccessControl from '../common/AccessControl.svelte'; let loading = false; @@ -103,6 +104,10 @@
+
+ +
+
- - +
+ {#if shiftKey} + +
+ + { + toggleModelById(localStorage.token, model.id); + _models.set(await getModels(localStorage.token)); + }} + /> + +
{/if}
@@ -492,43 +456,6 @@
- - {#if localModelfiles.length > 0} -
-
- {localModelfiles.length} Local Modelfiles Detected -
- -
- -
-
- {/if} {/if} diff --git a/src/lib/components/workspace/Models/ModelEditor.svelte b/src/lib/components/workspace/Models/ModelEditor.svelte index d5f8b2671..baff78e40 100644 --- a/src/lib/components/workspace/Models/ModelEditor.svelte +++ b/src/lib/components/workspace/Models/ModelEditor.svelte @@ -17,7 +17,7 @@ import { getTools } from '$lib/apis/tools'; import { getFunctions } from '$lib/apis/functions'; import { getKnowledgeItems } from '$lib/apis/knowledge'; - import AccessPermissions from '../common/AccessPermissionsModal.svelte'; + import AccessControl from '../common/AccessControl.svelte'; const i18n = getContext('i18n'); @@ -79,6 +79,8 @@ let filterIds = []; let actionIds = []; + let accessControl = null; + const addUsage = (base_model_id) => { const baseModel = $models.find((m) => m.id === base_model_id); @@ -208,6 +210,8 @@ capabilities.usage = false; } + accessControl = model?.access_control ?? null; + info = { ...info, ...JSON.parse( @@ -641,7 +645,7 @@
- +
diff --git a/src/lib/components/workspace/Prompts/PromptEditor.svelte b/src/lib/components/workspace/Prompts/PromptEditor.svelte index 51a06dd5f..fa215d236 100644 --- a/src/lib/components/workspace/Prompts/PromptEditor.svelte +++ b/src/lib/components/workspace/Prompts/PromptEditor.svelte @@ -4,6 +4,7 @@ import Textarea from '$lib/components/common/Textarea.svelte'; import { toast } from 'svelte-sonner'; import Tooltip from '$lib/components/common/Tooltip.svelte'; + import AccessControl from '../common/AccessControl.svelte'; export let onSubmit: Function; export let edit = false; @@ -135,6 +136,10 @@
+
+ +
+
{:else} - + diff --git a/src/lib/components/workspace/common/AccessPermissionsModal.svelte b/src/lib/components/workspace/common/AccessControl.svelte similarity index 89% rename from src/lib/components/workspace/common/AccessPermissionsModal.svelte rename to src/lib/components/workspace/common/AccessControl.svelte index 65b45578a..4f323dbe2 100644 --- a/src/lib/components/workspace/common/AccessPermissionsModal.svelte +++ b/src/lib/components/workspace/common/AccessControl.svelte @@ -1,10 +1,19 @@
@@ -14,7 +23,7 @@
- {#if type === 'private'} + {#if access === 'private'}
- {#if type === 'private'} + {#if access === 'private'} {$i18n.t('Only select users and groups with permission can access')} - {:else if type === 'public'} + {:else if access === 'public'} {$i18n.t('Accessible to all users')} {/if}
@@ -69,7 +78,7 @@
- {#if type === 'private'} + {#if access === 'private'}
{$i18n.t('People with access')}