From fe5c4b95d5c72f84725694791918aeba7f438436 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Wed, 19 Feb 2025 17:05:37 -0800 Subject: [PATCH] enh: configurable jupyter execution timeout --- backend/open_webui/main.py | 4 + backend/open_webui/routers/configs.py | 12 ++ backend/open_webui/routers/utils.py | 1 + backend/open_webui/utils/middleware.py | 1 + .../admin/Settings/CodeExecution.svelte | 184 +++++++++++------- .../components/chat/Messages/CodeBlock.svelte | 14 +- 6 files changed, 138 insertions(+), 78 deletions(-) diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index ced294469..b4ea2e1ec 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -106,6 +106,7 @@ from open_webui.config import ( CODE_EXECUTION_JUPYTER_AUTH, CODE_EXECUTION_JUPYTER_AUTH_TOKEN, CODE_EXECUTION_JUPYTER_AUTH_PASSWORD, + CODE_EXECUTION_JUPYTER_TIMEOUT, ENABLE_CODE_INTERPRETER, CODE_INTERPRETER_ENGINE, CODE_INTERPRETER_PROMPT_TEMPLATE, @@ -113,6 +114,7 @@ from open_webui.config import ( CODE_INTERPRETER_JUPYTER_AUTH, CODE_INTERPRETER_JUPYTER_AUTH_TOKEN, CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD, + CODE_INTERPRETER_JUPYTER_TIMEOUT, # Image AUTOMATIC1111_API_AUTH, AUTOMATIC1111_BASE_URL, @@ -644,6 +646,7 @@ app.state.config.CODE_EXECUTION_JUPYTER_AUTH_TOKEN = CODE_EXECUTION_JUPYTER_AUTH app.state.config.CODE_EXECUTION_JUPYTER_AUTH_PASSWORD = ( CODE_EXECUTION_JUPYTER_AUTH_PASSWORD ) +app.state.config.CODE_EXECUTION_JUPYTER_TIMEOUT = CODE_EXECUTION_JUPYTER_TIMEOUT app.state.config.ENABLE_CODE_INTERPRETER = ENABLE_CODE_INTERPRETER app.state.config.CODE_INTERPRETER_ENGINE = CODE_INTERPRETER_ENGINE @@ -657,6 +660,7 @@ app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_TOKEN = ( app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD = ( CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD ) +app.state.config.CODE_INTERPRETER_JUPYTER_TIMEOUT = CODE_INTERPRETER_JUPYTER_TIMEOUT ######################################## # diff --git a/backend/open_webui/routers/configs.py b/backend/open_webui/routers/configs.py index d460ae670..388c44f9c 100644 --- a/backend/open_webui/routers/configs.py +++ b/backend/open_webui/routers/configs.py @@ -75,6 +75,7 @@ class CodeInterpreterConfigForm(BaseModel): CODE_EXECUTION_JUPYTER_AUTH: Optional[str] CODE_EXECUTION_JUPYTER_AUTH_TOKEN: Optional[str] CODE_EXECUTION_JUPYTER_AUTH_PASSWORD: Optional[str] + CODE_EXECUTION_JUPYTER_TIMEOUT: Optional[int] ENABLE_CODE_INTERPRETER: bool CODE_INTERPRETER_ENGINE: str CODE_INTERPRETER_PROMPT_TEMPLATE: Optional[str] @@ -82,6 +83,7 @@ class CodeInterpreterConfigForm(BaseModel): CODE_INTERPRETER_JUPYTER_AUTH: Optional[str] CODE_INTERPRETER_JUPYTER_AUTH_TOKEN: Optional[str] CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD: Optional[str] + CODE_INTERPRETER_JUPYTER_TIMEOUT: Optional[int] @router.get("/code_execution", response_model=CodeInterpreterConfigForm) @@ -92,6 +94,7 @@ async def get_code_execution_config(request: Request, user=Depends(get_admin_use "CODE_EXECUTION_JUPYTER_AUTH": request.app.state.config.CODE_EXECUTION_JUPYTER_AUTH, "CODE_EXECUTION_JUPYTER_AUTH_TOKEN": request.app.state.config.CODE_EXECUTION_JUPYTER_AUTH_TOKEN, "CODE_EXECUTION_JUPYTER_AUTH_PASSWORD": request.app.state.config.CODE_EXECUTION_JUPYTER_AUTH_PASSWORD, + "CODE_EXECUTION_JUPYTER_TIMEOUT": request.app.state.config.CODE_EXECUTION_JUPYTER_TIMEOUT, "ENABLE_CODE_INTERPRETER": request.app.state.config.ENABLE_CODE_INTERPRETER, "CODE_INTERPRETER_ENGINE": request.app.state.config.CODE_INTERPRETER_ENGINE, "CODE_INTERPRETER_PROMPT_TEMPLATE": request.app.state.config.CODE_INTERPRETER_PROMPT_TEMPLATE, @@ -99,6 +102,7 @@ async def get_code_execution_config(request: Request, user=Depends(get_admin_use "CODE_INTERPRETER_JUPYTER_AUTH": request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH, "CODE_INTERPRETER_JUPYTER_AUTH_TOKEN": request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_TOKEN, "CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD": request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD, + "CODE_INTERPRETER_JUPYTER_TIMEOUT": request.app.state.config.CODE_INTERPRETER_JUPYTER_TIMEOUT, } @@ -120,6 +124,9 @@ async def set_code_execution_config( request.app.state.config.CODE_EXECUTION_JUPYTER_AUTH_PASSWORD = ( form_data.CODE_EXECUTION_JUPYTER_AUTH_PASSWORD ) + request.app.state.config.CODE_EXECUTION_JUPYTER_TIMEOUT = ( + form_data.CODE_EXECUTION_JUPYTER_TIMEOUT + ) request.app.state.config.ENABLE_CODE_INTERPRETER = form_data.ENABLE_CODE_INTERPRETER request.app.state.config.CODE_INTERPRETER_ENGINE = form_data.CODE_INTERPRETER_ENGINE @@ -141,6 +148,9 @@ async def set_code_execution_config( request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD = ( form_data.CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD ) + request.app.state.config.CODE_INTERPRETER_JUPYTER_TIMEOUT = ( + form_data.CODE_INTERPRETER_JUPYTER_TIMEOUT + ) return { "CODE_EXECUTION_ENGINE": request.app.state.config.CODE_EXECUTION_ENGINE, @@ -148,6 +158,7 @@ async def set_code_execution_config( "CODE_EXECUTION_JUPYTER_AUTH": request.app.state.config.CODE_EXECUTION_JUPYTER_AUTH, "CODE_EXECUTION_JUPYTER_AUTH_TOKEN": request.app.state.config.CODE_EXECUTION_JUPYTER_AUTH_TOKEN, "CODE_EXECUTION_JUPYTER_AUTH_PASSWORD": request.app.state.config.CODE_EXECUTION_JUPYTER_AUTH_PASSWORD, + "CODE_EXECUTION_JUPYTER_TIMEOUT": request.app.state.config.CODE_EXECUTION_JUPYTER_TIMEOUT, "ENABLE_CODE_INTERPRETER": request.app.state.config.ENABLE_CODE_INTERPRETER, "CODE_INTERPRETER_ENGINE": request.app.state.config.CODE_INTERPRETER_ENGINE, "CODE_INTERPRETER_PROMPT_TEMPLATE": request.app.state.config.CODE_INTERPRETER_PROMPT_TEMPLATE, @@ -155,6 +166,7 @@ async def set_code_execution_config( "CODE_INTERPRETER_JUPYTER_AUTH": request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH, "CODE_INTERPRETER_JUPYTER_AUTH_TOKEN": request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_TOKEN, "CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD": request.app.state.config.CODE_INTERPRETER_JUPYTER_AUTH_PASSWORD, + "CODE_INTERPRETER_JUPYTER_TIMEOUT": request.app.state.config.CODE_INTERPRETER_JUPYTER_TIMEOUT, } diff --git a/backend/open_webui/routers/utils.py b/backend/open_webui/routers/utils.py index 61863bda5..fb1dc8272 100644 --- a/backend/open_webui/routers/utils.py +++ b/backend/open_webui/routers/utils.py @@ -56,6 +56,7 @@ async def execute_code( if request.app.state.config.CODE_EXECUTION_JUPYTER_AUTH == "password" else None ), + request.app.state.config.CODE_EXECUTION_JUPYTER_TIMEOUT, ) return output diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index baa4d49a4..7ec764fc0 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -1761,6 +1761,7 @@ async def process_chat_response( == "password" else None ), + request.app.state.config.CODE_INTERPRETER_JUPYTER_TIMEOUT, ) else: output = { diff --git a/src/lib/components/admin/Settings/CodeExecution.svelte b/src/lib/components/admin/Settings/CodeExecution.svelte index cbce8e0e9..c83537455 100644 --- a/src/lib/components/admin/Settings/CodeExecution.svelte +++ b/src/lib/components/admin/Settings/CodeExecution.svelte @@ -91,45 +91,65 @@ -
-
- {$i18n.t('Jupyter Auth')} -
+
+
+
+ {$i18n.t('Jupyter Auth')} +
-
- -
-
- - {#if config.CODE_EXECUTION_JUPYTER_AUTH} -
-
- {#if config.CODE_EXECUTION_JUPYTER_AUTH === 'password'} - - {:else} - - {/if} +
+
- {/if} + + {#if config.CODE_EXECUTION_JUPYTER_AUTH} +
+
+ {#if config.CODE_EXECUTION_JUPYTER_AUTH === 'password'} + + {:else} + + {/if} +
+
+ {/if} +
+ +
+
+ {$i18n.t('Code Execution Timeout')} +
+ +
+ + + +
+
{/if}
@@ -197,45 +217,65 @@
-
-
- {$i18n.t('Jupyter Auth')} -
+
+
+
+ {$i18n.t('Jupyter Auth')} +
-
- -
-
- - {#if config.CODE_INTERPRETER_JUPYTER_AUTH} -
-
- {#if config.CODE_INTERPRETER_JUPYTER_AUTH === 'password'} - - {:else} - - {/if} +
+
- {/if} + + {#if config.CODE_INTERPRETER_JUPYTER_AUTH} +
+
+ {#if config.CODE_INTERPRETER_JUPYTER_AUTH === 'password'} + + {:else} + + {/if} +
+
+ {/if} +
+ +
+
+ {$i18n.t('Code Execution Timeout')} +
+ +
+ + + +
+
{/if}
diff --git a/src/lib/components/chat/Messages/CodeBlock.svelte b/src/lib/components/chat/Messages/CodeBlock.svelte index 7d318dc79..a5d08356f 100644 --- a/src/lib/components/chat/Messages/CodeBlock.svelte +++ b/src/lib/components/chat/Messages/CodeBlock.svelte @@ -123,6 +123,12 @@ }; const executePython = async (code) => { + result = null; + stdout = null; + stderr = null; + + executing = true; + if ($config?.code?.engine === 'jupyter') { const output = await executeCode(localStorage.token, code).catch((error) => { toast.error(`${error}`); @@ -190,18 +196,14 @@ output['stderr'] && (stderr = output['stderr']); } + + executing = false; } else { executePythonAsWorker(code); } }; const executePythonAsWorker = async (code) => { - result = null; - stdout = null; - stderr = null; - - executing = true; - let packages = [ code.includes('requests') ? 'requests' : null, code.includes('bs4') ? 'beautifulsoup4' : null,