From 83ff1d77ea2255981f4fc0bb4ff4f8ad9ad9f2a1 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Sun, 19 Nov 2023 00:41:29 -0800 Subject: [PATCH] feat: set first user to admin by default --- backend/apps/ollama/main.py | 4 ++-- backend/apps/web/main.py | 4 ++-- backend/apps/web/models/users.py | 3 +++ backend/apps/web/routers/auths.py | 3 ++- backend/config.py | 30 +++++++++++++----------------- backend/utils/utils.py | 2 +- compose.yaml | 3 ++- src/lib/constants.ts | 4 +++- 8 files changed, 28 insertions(+), 25 deletions(-) diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py index 023a10861..53daefac0 100644 --- a/backend/apps/ollama/main.py +++ b/backend/apps/ollama/main.py @@ -9,7 +9,7 @@ import json from apps.web.models.users import Users from constants import ERROR_MESSAGES from utils.utils import extract_token_from_auth_header -from config import OLLAMA_API_BASE_URL, OLLAMA_WEBUI_AUTH +from config import OLLAMA_API_BASE_URL, WEBUI_AUTH app = Flask(__name__) CORS( @@ -32,7 +32,7 @@ def proxy(path): headers = dict(request.headers) # Basic RBAC support - if OLLAMA_WEBUI_AUTH: + if WEBUI_AUTH: if "Authorization" in headers: token = extract_token_from_auth_header(headers["Authorization"]) user = Users.get_user_by_token(token) diff --git a/backend/apps/web/main.py b/backend/apps/web/main.py index 30c45d4a0..0380f0c34 100644 --- a/backend/apps/web/main.py +++ b/backend/apps/web/main.py @@ -2,7 +2,7 @@ from fastapi import FastAPI, Request, Depends, HTTPException from fastapi.middleware.cors import CORSMiddleware from apps.web.routers import auths, users -from config import OLLAMA_WEBUI_VERSION, OLLAMA_WEBUI_AUTH +from config import WEBUI_VERSION, WEBUI_AUTH app = FastAPI() @@ -23,4 +23,4 @@ app.include_router(users.router, prefix="/users", tags=["users"]) @app.get("/") async def get_status(): - return {"status": True, "version": OLLAMA_WEBUI_VERSION, "auth": OLLAMA_WEBUI_AUTH} + return {"status": True, "version": WEBUI_VERSION, "auth": WEBUI_AUTH} diff --git a/backend/apps/web/models/users.py b/backend/apps/web/models/users.py index e3c871c75..5f101f655 100644 --- a/backend/apps/web/models/users.py +++ b/backend/apps/web/models/users.py @@ -81,6 +81,9 @@ class UsersTable: ) ] + def get_num_users(self) -> Optional[int]: + return self.table.count_documents({}) + def update_user_by_id(self, id: str, updated: dict) -> Optional[UserModel]: user = self.table.find_one_and_update( {"id": id}, {"$set": updated}, return_document=ReturnDocument.AFTER diff --git a/backend/apps/web/routers/auths.py b/backend/apps/web/routers/auths.py index c306f4cd9..023e1914d 100644 --- a/backend/apps/web/routers/auths.py +++ b/backend/apps/web/routers/auths.py @@ -86,8 +86,9 @@ async def signin(form_data: SigninForm): async def signup(form_data: SignupForm): if not Users.get_user_by_email(form_data.email.lower()): try: + role = "admin" if Users.get_num_users() == 0 else "pending" hashed = get_password_hash(form_data.password) - user = Auths.insert_new_auth(form_data.email, hashed, form_data.name) + user = Auths.insert_new_auth(form_data.email, hashed, form_data.name, role) if user: token = create_token(data={"email": user.email}) diff --git a/backend/config.py b/backend/config.py index 3b67ef7f8..d19ebdbd0 100644 --- a/backend/config.py +++ b/backend/config.py @@ -26,40 +26,36 @@ if ENV == "prod": OLLAMA_API_BASE_URL = "http://host.docker.internal:11434/api" #################################### -# OLLAMA_WEBUI_VERSION +# WEBUI_VERSION #################################### -OLLAMA_WEBUI_VERSION = os.environ.get("OLLAMA_WEBUI_VERSION", "v1.0.0-alpha.9") +WEBUI_VERSION = os.environ.get("WEBUI_VERSION", "v1.0.0-alpha.9") #################################### -# OLLAMA_WEBUI_AUTH +# WEBUI_AUTH #################################### -OLLAMA_WEBUI_AUTH = ( - True if os.environ.get("OLLAMA_WEBUI_AUTH", "TRUE") == "TRUE" else False -) +WEBUI_AUTH = True if os.environ.get("WEBUI_AUTH", "TRUE") == "TRUE" else False -if OLLAMA_WEBUI_AUTH: +if WEBUI_AUTH: #################################### - # OLLAMA_WEBUI_DB + # WEBUI_DB #################################### - OLLAMA_WEBUI_DB_URL = os.environ.get( - "OLLAMA_WEBUI_DB_URL", "mongodb://root:root@localhost:27017/" + WEBUI_DB_URL = os.environ.get( + "WEBUI_DB_URL", "mongodb://root:root@localhost:27017/" ) - DB_CLIENT = MongoClient(f"{OLLAMA_WEBUI_DB_URL}?authSource=admin") + DB_CLIENT = MongoClient(f"{WEBUI_DB_URL}?authSource=admin") DB = DB_CLIENT["ollama-webui"] #################################### - # OLLAMA_WEBUI_JWT_SECRET_KEY + # WEBUI_JWT_SECRET_KEY #################################### - OLLAMA_WEBUI_JWT_SECRET_KEY = os.environ.get( - "OLLAMA_WEBUI_JWT_SECRET_KEY", "t0p-s3cr3t" - ) + WEBUI_JWT_SECRET_KEY = os.environ.get("WEBUI_JWT_SECRET_KEY", "t0p-s3cr3t") if ENV == "prod": - if OLLAMA_WEBUI_JWT_SECRET_KEY == "": - OLLAMA_WEBUI_JWT_SECRET_KEY = str(b64encode(token_bytes(32)).decode()) + if WEBUI_JWT_SECRET_KEY == "": + WEBUI_JWT_SECRET_KEY = str(b64encode(token_bytes(32)).decode()) diff --git a/backend/utils/utils.py b/backend/utils/utils.py index 96f40ddc3..62e6958fb 100644 --- a/backend/utils/utils.py +++ b/backend/utils/utils.py @@ -9,7 +9,7 @@ import jwt import config -JWT_SECRET_KEY = config.OLLAMA_WEBUI_JWT_SECRET_KEY +JWT_SECRET_KEY = config.WEBUI_JWT_SECRET_KEY ALGORITHM = "HS256" ############## diff --git a/compose.yaml b/compose.yaml index 1c1db1ada..43aaf96e5 100644 --- a/compose.yaml +++ b/compose.yaml @@ -46,7 +46,8 @@ services: - 3000:8080 environment: - "OLLAMA_API_BASE_URL=http://ollama:11434/api" - - "OLLAMA_WEBUI_DB_URL=mongodb://root:example@ollama-webui-db:27017/" + - "WEBUI_DB_URL=mongodb://root:example@ollama-webui-db:27017/" + - "WEBUI_AUTH=TRUE" extra_hosts: - host.docker.internal:host-gateway restart: unless-stopped diff --git a/src/lib/constants.ts b/src/lib/constants.ts index 9db787f02..e8cae39d0 100644 --- a/src/lib/constants.ts +++ b/src/lib/constants.ts @@ -3,8 +3,10 @@ import { PUBLIC_API_BASE_URL } from '$env/static/public'; export const OLLAMA_API_BASE_URL = PUBLIC_API_BASE_URL === '' - ? browser + ? dev ? `http://${location.hostname}:8080/ollama/api` + : browser + ? `http://${location.hostname}:11434/api` : `http://localhost:11434/api` : PUBLIC_API_BASE_URL;