Coverage for open_webui/utils/plugin.py: 34%

1from __future__ import annotations

3import asyncio

4import logging

5import os

6import re

7import subprocess

8import sys

9import tempfile

10import types

11from importlib import util

12from typing import Any

14from open_webui.env import (

15 ENABLE_PIP_INSTALL_FRONTMATTER_REQUIREMENTS,

16 ENABLE_PLUGINS,

17 OFFLINE_MODE,

18 PIP_OPTIONS,

19 PIP_PACKAGE_INDEX_OPTIONS,

20)

21from open_webui.models.functions import FunctionModel, Functions

22from open_webui.models.tools import Tools

24log = logging.getLogger(__name__)

27def resolve_valves_schema_options(valves_class: type, schema: dict, user: Any = None) -> dict:

28 """

29 Resolve dynamic options in a Valves schema.

31 For properties with `input.options`, this function handles two cases:

32 - List: Used directly as dropdown options

33 - String: Treated as method name, called to get options dynamically

35 Usage in Valves:

36 class UserValves(BaseModel):

37 # Static options

38 priority: str = Field(

39 default="medium",

40 json_schema_extra={

41 "input": {

42 "type": "select",

43 "options": ["low", "medium", "high"]

44 }

45 }

46 )

48 # Dynamic options (method name)

49 model: str = Field(

50 default="",

51 json_schema_extra={

52 "input": {

53 "type": "select",

54 "options": "get_model_options"

55 }

56 }

57 )

59 @classmethod

60 def get_model_options(cls, __user__=None) -> list[dict]:

61 return [{"value": "gpt-4", "label": "GPT-4"}]

63 Args:

64 valves_class: The Valves or UserValves Pydantic model class

65 schema: The JSON schema dict from valves_class.schema()

66 user: Optional user object passed to methods that accept __user__

68 Returns:

69 Modified schema dict with resolved options

70 """

71 if not schema or 'properties' not in schema:

72 return schema

74 # Make a copy to avoid mutating the original

75 schema = dict(schema)

76 schema['properties'] = dict(schema.get('properties', {}))

78 for prop_name, prop_schema in list(schema['properties'].items()):

79 # Get the original field info from the Pydantic model

80 if not hasattr(valves_class, 'model_fields'):

81 continue

83 field_info = valves_class.model_fields.get(prop_name)

84 if not field_info:

85 continue

87 # Check json_schema_extra for options

88 json_schema_extra = field_info.json_schema_extra

89 if not json_schema_extra or not isinstance(json_schema_extra, dict):

90 continue

92 input_config = json_schema_extra.get('input')

93 if not input_config or not isinstance(input_config, dict):

94 continue

96 options = input_config.get('options')

97 if options is None:

98 continue

100 resolved_options = None

101

102 # Case 1: options is already a list - use directly

103 if isinstance(options, list):

104 resolved_options = options

105

106 # Case 2: options is a string - treat as method name

107 elif isinstance(options, str) and options:

108 method = getattr(valves_class, options, None)

109 if method is None or not callable(method):

110 log.warning(f"options '{options}' not found or not callable on {valves_class.__name__}")

111 continue

112

113 try:

114 import inspect

115

116 sig = inspect.signature(method)

117 params = sig.parameters

118

119 # Prepare kwargs based on what the method accepts

120 kwargs = {}

121 if '__user__' in params and user is not None:

122 kwargs['__user__'] = user.model_dump() if hasattr(user, 'model_dump') else user

123 if 'user' in params and user is not None:

124 kwargs['user'] = user.model_dump() if hasattr(user, 'model_dump') else user

125

126 resolved_options = method(**kwargs) if kwargs else method()

127

128 # Validate return type

129 if not isinstance(resolved_options, list):

130 log.warning(f"Method '{options}' did not return a list for {prop_name}")

131 continue

132

133 except Exception as e:

134 log.warning(f'Failed to resolve options for {prop_name}: {e}')

135 continue

136 else:

137 # Invalid options type - skip

138 continue

139

140 # Update the schema with resolved options

141 schema['properties'][prop_name] = dict(prop_schema)

142 if 'input' not in schema['properties'][prop_name]:

143 schema['properties'][prop_name]['input'] = {'type': 'select'}

144 else:

145 schema['properties'][prop_name]['input'] = dict(schema['properties'][prop_name].get('input', {}))

146 schema['properties'][prop_name]['input']['options'] = resolved_options

147

148 return schema

149

150

151def extract_frontmatter(content):

152 """

153 Extract frontmatter as a dictionary from the provided content string.

154 """

155 frontmatter = {}

156 frontmatter_started = False

157 frontmatter_ended = False

158 frontmatter_pattern = re.compile(r'^\s*([a-z_]+):\s*(.*)\s*$', re.IGNORECASE)

159

160 try:

161 lines = content.splitlines()

162 if len(lines) < 1 or lines[0].strip() != '"""': 162 ↛ 166line 162 didn't jump to line 166 because the condition on line 162 was always true

163 # The content doesn't start with triple quotes

164 return {}

165

166 frontmatter_started = True

167

168 for line in lines[1:]:

169 if '"""' in line:

170 if frontmatter_started:

171 frontmatter_ended = True

172 break

173

174 if frontmatter_started and not frontmatter_ended:

175 match = frontmatter_pattern.match(line)

176 if match:

177 key, value = match.groups()

178 frontmatter[key.strip()] = value.strip()

179

180 except Exception as e:

181 log.exception(f'Failed to extract frontmatter: {e}')

182 return {}

183

184 return frontmatter

185

186

187def replace_imports(content):

188 """

189 Replace the import paths in the content.

190 """

191 replacements = {

192 'from utils': 'from open_webui.utils',

193 'from apps': 'from open_webui.apps',

194 'from main': 'from open_webui.main',

195 'from config': 'from open_webui.config',

196 }

197

198 for old, new in replacements.items():

199 content = content.replace(old, new)

200

201 return content

202

203

204# May the intent of the one who wrote it survive every

205# import and transformation, as a deed survives the generations.

206async def load_tool_module_by_id(tool_id, content=None):

207 if not ENABLE_PLUGINS: 207 ↛ 208line 207 didn't jump to line 208 because the condition on line 207 was never true

208 raise RuntimeError('Plugins are disabled by ENABLE_PLUGINS=false')

209

210 frontmatter = None

211 if content is None: 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true

212 tool = await Tools.get_tool_by_id(tool_id)

213 if not tool:

214 raise Exception(f'Toolkit not found: {tool_id}')

215

216 content = tool.content

217

218 content = replace_imports(content)

219 await Tools.update_tool_by_id(tool_id, {'content': content})

220 else:

221 frontmatter = extract_frontmatter(content)

222 # Install required packages found within the frontmatter.

223 # Runs `pip install` via subprocess, which can take a long time;

224 # offload to a thread so it doesn't block the event loop.

225 await asyncio.to_thread(install_frontmatter_requirements, frontmatter.get('requirements', ''))

226

227 module_name = f'tool_{tool_id}'

228 module = types.ModuleType(module_name)

229 sys.modules[module_name] = module

230

231 # Create a temporary file and use it to define `__file__` so

232 # that it works as expected from the module's perspective.

233 temp_file = tempfile.NamedTemporaryFile(delete=False)

234 temp_file.close()

235 try:

236 with open(temp_file.name, 'w', encoding='utf-8') as f:

237 f.write(content)

238 module.__dict__['__file__'] = temp_file.name

239

240 # Executing the modified content in the created module's namespace

241 exec(content, module.__dict__)

242 if frontmatter is None: 242 ↛ 243line 242 didn't jump to line 243 because the condition on line 242 was never true

243 frontmatter = extract_frontmatter(content)

244 log.info('Loaded module: %s', module.__name__)

245

246 # Create and return the object if the class 'Tools' is found in the module

247 if hasattr(module, 'Tools'): 247 ↛ 248line 247 didn't jump to line 248 because the condition on line 247 was never true

248 return module.Tools(), frontmatter

249 else:

250 raise Exception('No Tools class found in the module')

251 except Exception as e:

252 log.error(f'Error loading module: {tool_id}: {e}')

253 del sys.modules[module_name] # Clean up

254 raise e

255 finally:

256 os.unlink(temp_file.name)

257

258

259async def load_function_module_by_id(function_id: str, content: str | None = None):

260 if not ENABLE_PLUGINS: 260 ↛ 261line 260 didn't jump to line 261 because the condition on line 260 was never true

261 raise RuntimeError('Plugins are disabled by ENABLE_PLUGINS=false')

262

263 frontmatter = None

264 if content is None: 264 ↛ 265line 264 didn't jump to line 265 because the condition on line 264 was never true

265 function = await Functions.get_function_by_id(function_id)

266 if not function:

267 raise Exception(f'Function not found: {function_id}')

268 content = function.content

269

270 content = replace_imports(content)

271 await Functions.update_function_by_id(function_id, {'content': content})

272 else:

273 frontmatter = extract_frontmatter(content)

274 # `pip install` via subprocess can block for a long time; offload it.

275 await asyncio.to_thread(install_frontmatter_requirements, frontmatter.get('requirements', ''))

276

277 module_name = f'function_{function_id}'

278 module = types.ModuleType(module_name)

279 sys.modules[module_name] = module

280

281 # Create a temporary file and use it to define `__file__` so

282 # that it works as expected from the module's perspective.

283 temp_file = tempfile.NamedTemporaryFile(delete=False)

284 temp_file.close()

285 try:

286 with open(temp_file.name, 'w', encoding='utf-8') as f:

287 f.write(content)

288 module.__dict__['__file__'] = temp_file.name

289

290 # Execute the modified content in the created module's namespace

291 exec(content, module.__dict__)

292 if frontmatter is None: 292 ↛ 293line 292 didn't jump to line 293 because the condition on line 292 was never true

293 frontmatter = extract_frontmatter(content)

294 log.info('Loaded module: %s', module.__name__)

295

296 # Create appropriate object based on available class type in the module

297 if hasattr(module, 'Pipe'): 297 ↛ 298line 297 didn't jump to line 298 because the condition on line 297 was never true

298 return module.Pipe(), 'pipe', frontmatter

299 elif hasattr(module, 'Filter'): 299 ↛ 300line 299 didn't jump to line 300 because the condition on line 299 was never true

300 return module.Filter(), 'filter', frontmatter

301 elif hasattr(module, 'Action'): 301 ↛ 302line 301 didn't jump to line 302 because the condition on line 301 was never true

302 return module.Action(), 'action', frontmatter

303 elif hasattr(module, 'Event'): 303 ↛ 304line 303 didn't jump to line 304 because the condition on line 303 was never true

304 return module.Event(), 'event', frontmatter

305 else:

306 raise Exception('No Function class found in the module')

307 except Exception as e:

308 log.error(f'Error loading module: {function_id}: {e}')

309 # Cleanup by removing the module in case of error

310 del sys.modules[module_name]

311

312 await Functions.update_function_by_id(function_id, {'is_active': False})

313 raise e

314 finally:

315 os.unlink(temp_file.name)

316

317

318def _state_cache(request, name: str) -> dict:

319 if not hasattr(request.app.state, name):

320 setattr(request.app.state, name, {})

321 return getattr(request.app.state, name)

322

323

324def get_tools_cache(request) -> dict:

325 return _state_cache(request, 'TOOLS')

326

327

328def get_tool_contents_cache(request) -> dict:

329 return _state_cache(request, 'TOOL_CONTENTS')

330

331

332def get_functions_cache(request) -> dict:

333 return _state_cache(request, 'FUNCTIONS')

334

335

336def get_function_contents_cache(request) -> dict:

337 return _state_cache(request, 'FUNCTION_CONTENTS')

338

339

340async def get_tool_module_from_cache(request, tool_id, load_from_db=True):

341 tools_cache = get_tools_cache(request)

342 tool_contents_cache = get_tool_contents_cache(request)

343 content = None

344

345 if load_from_db:

346 # Always load from the database by default

347 tool = await Tools.get_tool_by_id(tool_id)

348 if not tool:

349 raise Exception(f'Tool not found: {tool_id}')

350 content = tool.content

351

352 new_content = replace_imports(content)

353 if new_content != content:

354 content = new_content

355 # Update the tool content in the database

356 await Tools.update_tool_by_id(tool_id, {'content': content})

357

358 if tool_id in tool_contents_cache and tool_id in tools_cache:

359 if tool_contents_cache[tool_id] == content:

360 return tools_cache[tool_id], None

361

362 tool_module, frontmatter = await load_tool_module_by_id(tool_id, content)

363 else:

364 if tool_id in tools_cache:

365 return tools_cache[tool_id], None

366

367 tool_module, frontmatter = await load_tool_module_by_id(tool_id)

368

369 tools_cache[tool_id] = tool_module

370 tool_contents_cache[tool_id] = content

371

372 return tool_module, frontmatter

373

374

375async def get_function_module_from_cache(

376 request, function_id, function: FunctionModel | None = None, load_from_db=True

377):

378 functions_cache = get_functions_cache(request)

379 function_contents_cache = get_function_contents_cache(request)

380 content = None

381

382 if load_from_db:

383 # Always load from the database by default

384 # This is useful for hooks like "inlet" or "outlet" where the content might change

385 # and we want to ensure the latest content is used.

386

387 if function is None:

388 function = await Functions.get_function_by_id(function_id)

389 if not function:

390 raise Exception(f'Function not found: {function_id}')

391 content = function.content

392

393 new_content = replace_imports(content)

394 if new_content != content:

395 content = new_content

396 # Update the function content in the database

397 await Functions.update_function_by_id(function_id, {'content': content})

398

399 if function_id in function_contents_cache and function_id in functions_cache:

400 if function_contents_cache[function_id] == content:

401 return functions_cache[function_id], None, None

402

403 function_module, function_type, frontmatter = await load_function_module_by_id(function_id, content)

404 else:

405 # Load from cache (e.g. "stream" hook)

406 # This is useful for performance reasons

407

408 if function_id in functions_cache:

409 return functions_cache[function_id], None, None

410

411 function_module, function_type, frontmatter = await load_function_module_by_id(function_id)

412

413 functions_cache[function_id] = function_module

414 function_contents_cache[function_id] = content

415

416 return function_module, function_type, frontmatter

417

418

419_installed_requirements = set()

420

421

422def install_frontmatter_requirements(requirements: str):

423 global _installed_requirements

424 if not ENABLE_PIP_INSTALL_FRONTMATTER_REQUIREMENTS: 424 ↛ 425line 424 didn't jump to line 425 because the condition on line 424 was never true

425 log.info('ENABLE_PIP_INSTALL_FRONTMATTER_REQUIREMENTS is disabled, skipping installation of requirements.')

426 return

427

428 if OFFLINE_MODE: 428 ↛ 432line 428 didn't jump to line 432 because the condition on line 428 was always true

429 log.info('Offline mode enabled, skipping installation of requirements.')

430 return

431

432 if requirements:

433 try:

434 req_list = [req.strip() for req in requirements.split(',')]

435 new_reqs = [req for req in req_list if req and req not in _installed_requirements]

436

437 if not new_reqs:

438 return

439

440 log.info('Installing requirements: %s', ' '.join(new_reqs))

441 subprocess.check_call(

442 [sys.executable, '-m', 'pip', 'install'] + PIP_OPTIONS + new_reqs + PIP_PACKAGE_INDEX_OPTIONS

443 )

444 _installed_requirements.update(new_reqs)

445 except Exception as e:

446 log.error(f'Error installing packages: {" ".join(new_reqs)}')

447 raise e

448

449 else:

450 log.info('No requirements found in frontmatter.')

451

452

453async def install_tool_and_function_dependencies():

454 """

455 Install all dependencies for all admin tools and active functions.

456

457 By first collecting all dependencies from the frontmatter of each tool and function,

458 and then installing them using pip. Duplicates or similar version specifications are

459 handled by pip as much as possible.

460 """

461 if not ENABLE_PLUGINS: 461 ↛ 462line 461 didn't jump to line 462 because the condition on line 461 was never true

462 log.info('ENABLE_PLUGINS is disabled, skipping tool and function dependencies.')

463 return

464

465 function_list = await Functions.get_functions(active_only=True)

466 tool_list = await Tools.get_tools()

467

468 all_dependencies = ''

469 try:

470 for function in function_list: 470 ↛ 471line 470 didn't jump to line 471 because the loop on line 470 never started

471 frontmatter = extract_frontmatter(replace_imports(function.content))

472 if dependencies := frontmatter.get('requirements'):

473 all_dependencies += f'{dependencies}, '

474 for tool in tool_list: 474 ↛ 476line 474 didn't jump to line 476 because the loop on line 474 never started

475 # Only install requirements for admin tools

476 if tool.user and tool.user.role == 'admin':

477 frontmatter = extract_frontmatter(replace_imports(tool.content))

478 if dependencies := frontmatter.get('requirements'):

479 all_dependencies += f'{dependencies}, '

480

481 # `pip install` via subprocess can block for a long time; offload it.

482 await asyncio.to_thread(install_frontmatter_requirements, all_dependencies.strip(', '))

483 except Exception as e:

484 log.error(f'Error installing requirements: {e}')