From 150d0adea26e29bfc34af814902f4e63309cf1e5 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 14 Nov 2024 18:57:25 -0800 Subject: [PATCH] refac: access_control field --- backend/open_webui/apps/webui/models/knowledge.py | 8 ++++++++ backend/open_webui/apps/webui/models/models.py | 10 +++++++++- backend/open_webui/apps/webui/models/prompts.py | 9 ++++++++- backend/open_webui/apps/webui/models/tools.py | 11 ++++++++++- 4 files changed, 35 insertions(+), 3 deletions(-) diff --git a/backend/open_webui/apps/webui/models/knowledge.py b/backend/open_webui/apps/webui/models/knowledge.py index 269ad8cc3..d5d270069 100644 --- a/backend/open_webui/apps/webui/models/knowledge.py +++ b/backend/open_webui/apps/webui/models/knowledge.py @@ -34,6 +34,12 @@ class Knowledge(Base): data = Column(JSON, nullable=True) meta = Column(JSON, nullable=True) + access_control = Column(JSON, nullable=True) # Controls data access levels. + # NULL for public access (open to all users with "user" role). + # {} for individual access (private to the owner). + # {"group_ids": ["group_id1", "group_id2"]} for access restricted to specific groups. + # {"user_ids": ["user_id1", "user_id2"]} for access restricted to specific users. + created_at = Column(BigInteger) updated_at = Column(BigInteger) @@ -50,6 +56,8 @@ class KnowledgeModel(BaseModel): data: Optional[dict] = None meta: Optional[dict] = None + access_control = Optional[dict] = None + created_at: int # timestamp in epoch updated_at: int # timestamp in epoch diff --git a/backend/open_webui/apps/webui/models/models.py b/backend/open_webui/apps/webui/models/models.py index 9bdffb9bc..8dfb7b9ee 100644 --- a/backend/open_webui/apps/webui/models/models.py +++ b/backend/open_webui/apps/webui/models/models.py @@ -5,7 +5,7 @@ from typing import Optional from open_webui.apps.webui.internal.db import Base, JSONField, get_db from open_webui.env import SRC_LOG_LEVELS from pydantic import BaseModel, ConfigDict -from sqlalchemy import BigInteger, Column, Text +from sqlalchemy import BigInteger, Column, Text, JSON log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["MODELS"]) @@ -67,6 +67,12 @@ class Model(Base): Holds a JSON encoded blob of metadata, see `ModelMeta`. """ + access_control = Column(JSON, nullable=True) # Controls data access levels. + # NULL for public access (open to all users with "user" role). + # {} for individual access (private to the owner). + # {"group_ids": ["group_id1", "group_id2"]} for access restricted to specific groups. + # {"user_ids": ["user_id1", "user_id2"]} for access restricted to specific users. + updated_at = Column(BigInteger) created_at = Column(BigInteger) @@ -80,6 +86,8 @@ class ModelModel(BaseModel): params: ModelParams meta: ModelMeta + access_control = Optional[dict] = None + updated_at: int # timestamp in epoch created_at: int # timestamp in epoch diff --git a/backend/open_webui/apps/webui/models/prompts.py b/backend/open_webui/apps/webui/models/prompts.py index 6b98e5c53..11f0f7353 100644 --- a/backend/open_webui/apps/webui/models/prompts.py +++ b/backend/open_webui/apps/webui/models/prompts.py @@ -3,7 +3,7 @@ from typing import Optional from open_webui.apps.webui.internal.db import Base, get_db from pydantic import BaseModel, ConfigDict -from sqlalchemy import BigInteger, Column, String, Text +from sqlalchemy import BigInteger, Column, String, Text, JSON #################### # Prompts DB Schema @@ -19,6 +19,12 @@ class Prompt(Base): content = Column(Text) timestamp = Column(BigInteger) + access_control = Column(JSON, nullable=True) # Controls data access levels. + # NULL for public access (open to all users with "user" role). + # {} for individual access (private to the owner). + # {"group_ids": ["group_id1", "group_id2"]} for access restricted to specific groups. + # {"user_ids": ["user_id1", "user_id2"]} for access restricted to specific users. + class PromptModel(BaseModel): command: str @@ -27,6 +33,7 @@ class PromptModel(BaseModel): content: str timestamp: int # timestamp in epoch + access_control = Optional[dict] = None model_config = ConfigDict(from_attributes=True) diff --git a/backend/open_webui/apps/webui/models/tools.py b/backend/open_webui/apps/webui/models/tools.py index e06f83452..77f1858f6 100644 --- a/backend/open_webui/apps/webui/models/tools.py +++ b/backend/open_webui/apps/webui/models/tools.py @@ -6,7 +6,7 @@ from open_webui.apps.webui.internal.db import Base, JSONField, get_db from open_webui.apps.webui.models.users import Users from open_webui.env import SRC_LOG_LEVELS from pydantic import BaseModel, ConfigDict -from sqlalchemy import BigInteger, Column, String, Text +from sqlalchemy import BigInteger, Column, String, Text, JSON log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["MODELS"]) @@ -26,6 +26,13 @@ class Tool(Base): specs = Column(JSONField) meta = Column(JSONField) valves = Column(JSONField) + + access_control = Column(JSON, nullable=True) # Controls data access levels. + # NULL for public access (open to all users with "user" role). + # {} for individual access (private to the owner). + # {"group_ids": ["group_id1", "group_id2"]} for access restricted to specific groups. + # {"user_ids": ["user_id1", "user_id2"]} for access restricted to specific users. + updated_at = Column(BigInteger) created_at = Column(BigInteger) @@ -42,6 +49,8 @@ class ToolModel(BaseModel): content: str specs: list[dict] meta: ToolMeta + access_control = Optional[dict] = None + updated_at: int # timestamp in epoch created_at: int # timestamp in epoch