Coverage for open_webui/__init__.py: 30%

1import base64 

2import os 

3import random 

4import sys 

5from pathlib import Path 

6from typing import Annotated 

7 

8import typer 

9import uvicorn 

10 

11app = typer.Typer() 

12 

13KEY_FILE = Path.cwd() / '.webui_secret_key' 

14DEFAULT_SECRET_KEY_LENGTH = 24 

15 

16 

17def version_callback(value: bool) -> None: 

18 if value: 

19 from open_webui.env import VERSION 

20 

21 # LICENSE covers this Open WebUI CLI identifier. 

22 # Do not alter, remove, obscure, or replace it except as LICENSE permits: 

23 # https://docs.openwebui.com/license. 

24 typer.echo(f'Open WebUI version: {VERSION}') 

25 raise typer.Exit() 

26 

27 

28@app.command() 

29def main( 

30 version: Annotated[bool | None, typer.Option('--version', callback=version_callback)] = None, 

31): 

32 pass 

33 

34 

35@app.command() 

36def serve( 

37 host: str = '0.0.0.0', 

38 port: int = 8080, 

39): 

40 os.environ['FROM_INIT_PY'] = 'true' 

41 if os.getenv('WEBUI_SECRET_KEY') is None: 

42 typer.echo('Loading WEBUI_SECRET_KEY from file, not provided as an environment variable.') 

43 if not KEY_FILE.exists(): 

44 key_length = int(os.getenv('WEBUI_SECRET_KEY_LENGTH', DEFAULT_SECRET_KEY_LENGTH)) 

45 if key_length < 1: 

46 raise ValueError('WEBUI_SECRET_KEY_LENGTH must be a positive integer') 

47 typer.echo(f'Generating a new secret key and saving it to {KEY_FILE}') 

48 KEY_FILE.write_bytes(base64.b64encode(random.randbytes(key_length))) 

49 typer.echo(f'Loading WEBUI_SECRET_KEY from {KEY_FILE}') 

50 os.environ['WEBUI_SECRET_KEY'] = KEY_FILE.read_text() 

51 

52 if os.getenv('USE_CUDA_DOCKER', 'false') == 'true': 

53 typer.echo('CUDA is enabled, appending LD_LIBRARY_PATH to include torch/cudnn & cublas libraries.') 

54 LD_LIBRARY_PATH = os.getenv('LD_LIBRARY_PATH', '').split(':') 

55 os.environ['LD_LIBRARY_PATH'] = ':'.join( 

56 LD_LIBRARY_PATH 

57 + [ 

58 '/usr/local/lib/python3.11/site-packages/torch/lib', 

59 '/usr/local/lib/python3.11/site-packages/nvidia/cudnn/lib', 

60 ] 

61 ) 

62 try: 

63 import torch 

64 

65 assert torch.cuda.is_available(), 'CUDA not available' 

66 typer.echo('CUDA seems to be working') 

67 except Exception as e: 

68 typer.echo( 

69 'Error when testing CUDA but USE_CUDA_DOCKER is true. ' 

70 'Resetting USE_CUDA_DOCKER to false and removing ' 

71 f'LD_LIBRARY_PATH modifications: {e}' 

72 ) 

73 os.environ['USE_CUDA_DOCKER'] = 'false' 

74 os.environ['LD_LIBRARY_PATH'] = ':'.join(LD_LIBRARY_PATH) 

75 

76 import open_webui.main # noqa: F401 

77 from open_webui.env import UVICORN_WORKERS, UVICORN_WS_PER_MESSAGE_DEFLATE 

78 

79 # On Windows, uvicorn's default loop factory hardcodes ProactorEventLoop, 

80 # which is incompatible with psycopg v3 async. Setting loop='none' lets 

81 # asyncio.run() respect the WindowsSelectorEventLoopPolicy set in db.py. 

82 loop = 'none' if sys.platform == 'win32' else 'auto' 

83 

84 uvicorn.run( 

85 'open_webui.main:app', 

86 host=host, 

87 port=port, 

88 forwarded_allow_ips='*', 

89 workers=UVICORN_WORKERS, 

90 ws_per_message_deflate=UVICORN_WS_PER_MESSAGE_DEFLATE, 

91 loop=loop, 

92 ) 

93 

94 

95@app.command() 

96def dev( 

97 host: str = '0.0.0.0', 

98 port: int = 8080, 

99 reload: bool = True, 

100): 

101 from open_webui.env import UVICORN_WS_PER_MESSAGE_DEFLATE 

102 

103 uvicorn.run( 

104 'open_webui.main:app', 

105 host=host, 

106 port=port, 

107 reload=reload, 

108 forwarded_allow_ips='*', 

109 ws_per_message_deflate=UVICORN_WS_PER_MESSAGE_DEFLATE, 

110 ) 

111 

112 

113if __name__ == '__main__': 113 ↛ 114line 113 didn't jump to line 114 because the condition on line 113 was never true

114 app()