Coverage for open_webui/routers/utils.py: 85%

1from __future__ import annotations 

2 

3import logging 

4 

5import black 

6from fastapi import APIRouter, Depends, HTTPException, Request, status 

7from open_webui.config import DATA_DIR, ENABLE_ADMIN_EXPORT 

8from open_webui.constants import ERROR_MESSAGES 

9from open_webui.models.config import Config 

10from open_webui.utils.auth import get_admin_user, get_verified_user 

11from open_webui.utils.code_interpreter import execute_code_jupyter 

12from open_webui.utils.misc import get_gravatar_url 

13from pydantic import BaseModel 

14from starlette.responses import FileResponse 

15 

16log = logging.getLogger(__name__) 

17 

18router = APIRouter() 

19 

20 

21@router.get('/gravatar') 

22async def get_gravatar(email: str, user=Depends(get_verified_user)): 

23 return get_gravatar_url(email) 

24 

25 

26class CodeForm(BaseModel): 

27 code: str 

28 

29 

30@router.post('/code/format') 

31async def format_code(form_data: CodeForm, user=Depends(get_admin_user)): 

32 try: 

33 formatted_code = black.format_str(form_data.code, mode=black.Mode()) 

34 return {'code': formatted_code} 

35 except black.NothingChanged: 

36 return {'code': form_data.code} 

37 except Exception as e: 

38 raise HTTPException(status_code=400, detail=str(e)) 

39 

40 

41@router.post('/code/execute') 

42async def execute_code(request: Request, form_data: CodeForm, user=Depends(get_verified_user)): 

43 if not await Config.get('code_execution.enable'): 

44 raise HTTPException( 

45 status_code=403, 

46 detail=ERROR_MESSAGES.FEATURE_DISABLED('Code execution'), 

47 ) 

48 

49 if await Config.get('code_execution.engine') == 'jupyter': 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true

50 output = await execute_code_jupyter( 

51 await Config.get('code_execution.jupyter.url'), 

52 form_data.code, 

53 ( 

54 await Config.get('code_execution.jupyter.auth_token') 

55 if await Config.get('code_execution.jupyter.auth') == 'token' 

56 else None 

57 ), 

58 ( 

59 await Config.get('code_execution.jupyter.auth_password') 

60 if await Config.get('code_execution.jupyter.auth') == 'password' 

61 else None 

62 ), 

63 await Config.get('code_execution.jupyter.timeout'), 

64 ) 

65 

66 return output 

67 else: 

68 raise HTTPException( 

69 status_code=400, 

70 detail=ERROR_MESSAGES.DEFAULT('Code execution engine not supported'), 

71 ) 

72 

73 

74@router.get('/db/download') 

75async def download_db(user=Depends(get_admin_user)): 

76 """Download the raw SQLite database file (admin-only, SQLite deployments only).""" 

77 if not ENABLE_ADMIN_EXPORT: 77 ↛ 78line 77 didn't jump to line 78 because the condition on line 77 was never true

78 raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.ACCESS_PROHIBITED) 

79 

80 # Lazy import avoids circular dependency at module load time 

81 from open_webui.internal.db import engine 

82 

83 if engine.name != 'sqlite': 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true

84 raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DB_NOT_SQLITE) 

85 

86 return FileResponse( 

87 str(engine.url.database), 

88 media_type='application/octet-stream', 

89 filename='webui.db', 

90 )