diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 898ac1b59..4931a2456 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -1729,6 +1729,16 @@ CODE_INTERPRETER_JUPYTER_TIMEOUT = PersistentConfig( ), ) +CODE_INTERPRETER_BLACKLISTED_MODULES = PersistentConfig( + "CODE_INTERPRETER_BLACKLISTED_MODULES", + "code_interpreter.blacklisted_modules", + [ + library.strip() + for library in os.environ.get("CODE_INTERPRETER_BLACKLISTED_MODULES", "").split(",") + if library.strip() + ], +) + DEFAULT_CODE_INTERPRETER_PROMPT = """ #### Tools Available diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index b1e69db26..70eb76c61 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -3,6 +3,7 @@ import logging import sys import os import base64 +import textwrap import asyncio from aiocache import cached @@ -89,6 +90,7 @@ from open_webui.config import ( CACHE_DIR, DEFAULT_TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE, DEFAULT_CODE_INTERPRETER_PROMPT, + CODE_INTERPRETER_BLACKLISTED_MODULES, ) from open_webui.env import ( SRC_LOG_LEVELS, @@ -2227,6 +2229,25 @@ async def process_chat_response( try: if content_blocks[-1]["attributes"].get("type") == "code": code = content_blocks[-1]["content"] + if CODE_INTERPRETER_BLACKLISTED_MODULES: + blocking_code = textwrap.dedent(f""" + import builtins + + BLACKLISTED_MODULES = {CODE_INTERPRETER_BLACKLISTED_MODULES} + + _real_import = builtins.__import__ + def restricted_import(name, globals=None, locals=None, fromlist=(), level=0): + if name.split('.')[0] in BLACKLISTED_MODULES: + importer_name = globals.get('__name__') if globals else None + if importer_name == '__main__': + raise ImportError( + f"Direct import of module {{name}} is restricted." + ) + return _real_import(name, globals, locals, fromlist, level) + + builtins.__import__ = restricted_import + """) + code = blocking_code + "\n" + code if ( request.app.state.config.CODE_INTERPRETER_ENGINE