refac

backend/open_webui/__init__.py

@@ -11,6 +11,7 @@ import uvicorn
 app = typer.Typer()
 
 KEY_FILE = Path.cwd() / '.webui_secret_key'
+DEFAULT_SECRET_KEY_LENGTH = 24
 
 
 def version_callback(value: bool) -> None:
@@ -37,8 +38,11 @@ def serve(
     if os.getenv('WEBUI_SECRET_KEY') is None:
         typer.echo('Loading WEBUI_SECRET_KEY from file, not provided as an environment variable.')
         if not KEY_FILE.exists():
+            key_length = int(os.getenv('WEBUI_SECRET_KEY_LENGTH', DEFAULT_SECRET_KEY_LENGTH))
+            if key_length < 1:
+                raise ValueError('WEBUI_SECRET_KEY_LENGTH must be a positive integer')
             typer.echo(f'Generating a new secret key and saving it to {KEY_FILE}')
-            KEY_FILE.write_bytes(base64.b64encode(random.randbytes(12)))
+            KEY_FILE.write_bytes(base64.b64encode(random.randbytes(key_length)))
         typer.echo(f'Loading WEBUI_SECRET_KEY from {KEY_FILE}')
         os.environ['WEBUI_SECRET_KEY'] = KEY_FILE.read_text()