Coverage for open_webui/functions.py: 13%

1import asyncio 

2import inspect 

3import logging 

4import sys 

5from typing import AsyncGenerator, Generator, Iterator 

6 

7from fastapi import ( 

8 Depends, 

9 FastAPI, 

10 File, 

11 Form, 

12 HTTPException, 

13 Request, 

14 UploadFile, 

15 status, 

16) 

17from pydantic import BaseModel 

18from starlette.responses import Response, StreamingResponse 

19 

20from open_webui.config import BYPASS_ADMIN_ACCESS_CONTROL 

21from open_webui.constants import ERROR_MESSAGES 

22from open_webui.env import BYPASS_MODEL_ACCESS_CONTROL, ENABLE_PLUGINS, GLOBAL_LOG_LEVEL 

23from open_webui.models.functions import Functions 

24from open_webui.models.models import Models 

25from open_webui.models.users import UserModel 

26from open_webui.socket.main import ( 

27 get_event_call, 

28 get_event_emitter, 

29) 

30from open_webui.utils.access_control import check_model_access 

31from open_webui.utils.json_codec import JSONCodec 

32from open_webui.utils.misc import ( 

33 add_or_update_system_message, 

34 get_last_user_message, 

35 openai_chat_chunk_message_template, 

36 openai_chat_completion_message_template, 

37 prepend_to_first_user_message_content, 

38) 

39from open_webui.utils.payload import ( 

40 apply_model_params_to_body_openai, 

41 apply_system_prompt_to_body, 

42) 

43from open_webui.utils.plugin import ( 

44 get_function_module_from_cache, 

45 load_function_module_by_id, 

46) 

47 

48logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL) 

49log = logging.getLogger(__name__) 

50 

51 

52async def get_function_module_by_id(request: Request, pipe_id: str): 

53 function_module, _, _ = await get_function_module_from_cache(request, pipe_id) 

54 

55 if hasattr(function_module, 'valves') and hasattr(function_module, 'Valves'): 

56 Valves = function_module.Valves 

57 valves = await Functions.get_function_valves_by_id(pipe_id) 

58 

59 if valves: 

60 try: 

61 function_module.valves = Valves(**{k: v for k, v in valves.items() if v is not None}) 

62 except Exception as e: 

63 log.exception(f'Error loading valves for function {pipe_id}: {e}') 

64 raise e 

65 else: 

66 function_module.valves = Valves() 

67 

68 return function_module 

69 

70 

71async def get_function_models(request): 

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

73 return [] 

74 

75 pipes = await Functions.get_functions_by_type('pipe', active_only=True) 

76 pipe_models = [] 

77 

78 for pipe in pipes: 78 ↛ 79line 78 didn't jump to line 79 because the loop on line 78 never started

79 try: 

80 function_module = await get_function_module_by_id(request, pipe.id) 

81 

82 has_user_valves = False 

83 if hasattr(function_module, 'UserValves'): 

84 has_user_valves = True 

85 

86 # Check if function is a manifold 

87 if hasattr(function_module, 'pipes'): 

88 sub_pipes = [] 

89 

90 # Handle pipes being a list, sync function, or async function 

91 try: 

92 if callable(function_module.pipes): 

93 if asyncio.iscoroutinefunction(function_module.pipes): 

94 sub_pipes = await function_module.pipes() 

95 else: 

96 sub_pipes = function_module.pipes() 

97 else: 

98 sub_pipes = function_module.pipes 

99 except Exception as e: 

100 log.exception(e) 

101 sub_pipes = [] 

102 

103 log.debug("get_function_models: function '%s' is a manifold of %s", pipe.id, sub_pipes) 

104 

105 for p in sub_pipes: 

106 sub_pipe_id = f'{pipe.id}.{p["id"]}' 

107 sub_pipe_name = p['name'] 

108 

109 if hasattr(function_module, 'name'): 

110 sub_pipe_name = f'{function_module.name}{sub_pipe_name}' 

111 

112 pipe_flag = {'type': pipe.type} 

113 

114 pipe_models.append( 

115 { 

116 'id': sub_pipe_id, 

117 'name': sub_pipe_name, 

118 'object': 'model', 

119 'created': pipe.created_at, 

120 'owned_by': 'openai', 

121 'pipe': pipe_flag, 

122 'has_user_valves': has_user_valves, 

123 } 

124 ) 

125 else: 

126 pipe_flag = {'type': 'pipe'} 

127 

128 log.debug( 

129 "get_function_models: function '%s' is a single pipe { 'id': %s, 'name': %s }", 

130 pipe.id, 

131 pipe.id, 

132 pipe.name, 

133 ) 

134 

135 pipe_models.append( 

136 { 

137 'id': pipe.id, 

138 'name': pipe.name, 

139 'object': 'model', 

140 'created': pipe.created_at, 

141 'owned_by': 'openai', 

142 'pipe': pipe_flag, 

143 'has_user_valves': has_user_valves, 

144 } 

145 ) 

146 except Exception as e: 

147 log.exception(e) 

148 continue 

149 

150 return pipe_models 

151 

152 

153async def generate_function_chat_completion(request, form_data, user, models: dict | None = None): 

154 if models is None: 

155 models = {} 

156 

157 async def execute_pipe(pipe, params): 

158 if inspect.iscoroutinefunction(pipe): 

159 return await pipe(**params) 

160 else: 

161 return pipe(**params) 

162 

163 async def get_message_content(res: str | Generator | AsyncGenerator) -> str: 

164 if isinstance(res, str): 

165 return res 

166 if isinstance(res, Generator): 

167 return ''.join(map(str, res)) 

168 if isinstance(res, AsyncGenerator): 168 ↛ exitline 168 didn't return from function 'get_message_content' because the condition on line 168 was always true

169 return ''.join([str(stream) async for stream in res]) 

170 

171 def process_line(form_data: dict, line): 

172 if isinstance(line, BaseModel): 

173 line = line.model_dump_json() 

174 line = f'data: {line}' 

175 if isinstance(line, dict): 

176 line = f'data: {JSONCodec.dumps(line)}' 

177 

178 try: 

179 line = line.decode('utf-8') 

180 except Exception: 

181 pass 

182 

183 if line.startswith('data:'): 

184 return f'{line}\n\n' 

185 else: 

186 line = openai_chat_chunk_message_template(form_data['model'], line) 

187 return f'data: {JSONCodec.dumps(line)}\n\n' 

188 

189 def get_pipe_id(form_data: dict) -> str: 

190 pipe_id = form_data['model'] 

191 if '.' in pipe_id: 191 ↛ anywhereline 191 didn't jump anywhere: it always raised an exception.

192 pipe_id, _ = pipe_id.split('.', 1) 

193 return pipe_id 

194 

195 async def get_function_params(function_module, form_data, user, extra_params=None): 

196 if extra_params is None: 

197 extra_params = {} 

198 

199 pipe_id = get_pipe_id(form_data) 

200 

201 # Get the signature of the function 

202 sig = inspect.signature(function_module.pipe) 

203 params = {'body': form_data} | {k: v for k, v in extra_params.items() if k in sig.parameters} 

204 

205 if '__user__' in params and hasattr(function_module, 'UserValves'): 205 ↛ anywhereline 205 didn't jump anywhere: it always raised an exception.

206 user_valves = await Functions.get_user_valves_by_id_and_user_id(pipe_id, user.id) 

207 try: 

208 params['__user__']['valves'] = function_module.UserValves(**user_valves) 

209 except Exception as e: 

210 log.exception(e) 

211 params['__user__']['valves'] = function_module.UserValves() 

212 

213 return params 

214 

215 # Set server-side by utils/chat.py, never by client input. Mirrors the routers. 

216 bypass_system_prompt = getattr(request.state, 'bypass_system_prompt', False) 

217 

218 # Copy so the base-model substitution below doesn't leak into the caller's 

219 # payload, which the tool-call continuation re-submits. Mirrors the routers. 

220 form_data = {**form_data} 

221 

222 model_id = form_data.get('model') 

223 model_info = await Models.get_model_by_id(model_id) 

224 

225 metadata = form_data.pop('metadata', {}) 

226 

227 files = metadata.get('files', []) 

228 tool_ids = metadata.get('tool_ids', []) 

229 # Check if tool_ids is None 

230 if tool_ids is None: 

231 tool_ids = [] 

232 

233 __event_emitter__ = None 

234 __event_call__ = None 

235 __task__ = None 

236 __task_body__ = None 

237 

238 if metadata: 

239 if all(k in metadata for k in ('session_id', 'chat_id', 'message_id')): 

240 __event_emitter__ = await get_event_emitter(metadata) 

241 __event_call__ = await get_event_call(metadata) 

242 __task__ = metadata.get('task', None) 

243 __task_body__ = metadata.get('task_body', None) 

244 

245 oauth_token = None 

246 try: 

247 oauth_session_id = request.cookies.get('oauth_session_id', None) 

248 if oauth_session_id: 

249 oauth_token = await request.app.state.oauth_manager.get_oauth_token( 

250 user.id, 

251 oauth_session_id, 

252 ) 

253 

254 # Fallback: no cookie (automation, API key, etc.) — use most recent session 

255 if oauth_token is None: 

256 from open_webui.models.oauth_sessions import OAuthSessions 

257 

258 sessions = await OAuthSessions.get_sessions_by_user_id(user.id) 

259 if sessions: 

260 best = max(sessions, key=lambda s: s.updated_at) 

261 oauth_token = await request.app.state.oauth_manager.get_oauth_token( 

262 user.id, 

263 best.id, 

264 ) 

265 except Exception as e: 

266 log.error(f'Error getting OAuth token: {e}') 

267 

268 extra_params = { 

269 '__event_emitter__': __event_emitter__, 

270 '__event_call__': __event_call__, 

271 '__chat_id__': metadata.get('chat_id', None), 

272 '__session_id__': metadata.get('session_id', None), 

273 '__message_id__': metadata.get('message_id', None), 

274 '__task__': __task__, 

275 '__task_body__': __task_body__, 

276 '__files__': files, 

277 '__user__': user.model_dump() if isinstance(user, UserModel) else {}, 

278 '__metadata__': metadata, 

279 '__oauth_token__': oauth_token, 

280 '__request__': request, 

281 } 

282 extra_params['__tools__'] = metadata.get('tools', {}) 

283 

284 if model_info: 

285 if model_info.base_model_id: 

286 form_data['model'] = model_info.base_model_id 

287 

288 if not BYPASS_MODEL_ACCESS_CONTROL: 

289 bypass = isinstance(user, UserModel) and user.role == 'admin' and BYPASS_ADMIN_ACCESS_CONTROL 

290 await check_model_access(user if isinstance(user, UserModel) else UserModel(**user), model_info, bypass) 

291 

292 params = model_info.params.model_dump() 

293 

294 if params: 

295 system = params.pop('system', None) 

296 form_data = apply_model_params_to_body_openai(params, form_data) 

297 if not bypass_system_prompt: 

298 form_data = await apply_system_prompt_to_body(system, form_data, metadata, user) 

299 

300 pipe_id = get_pipe_id(form_data) 

301 function_module = await get_function_module_by_id(request, pipe_id) 

302 

303 pipe = function_module.pipe 

304 params = await get_function_params(function_module, form_data, user, extra_params) 

305 

306 if form_data.get('stream', False): 

307 

308 async def stream_content(): 

309 try: 

310 res = await execute_pipe(pipe, params) 

311 

312 # Directly return if the response is a StreamingResponse 

313 if isinstance(res, StreamingResponse): 

314 async for data in res.body_iterator: 

315 yield data 

316 return 

317 if isinstance(res, dict): 

318 yield f'data: {JSONCodec.dumps(res)}\n\n' 

319 return 

320 

321 except Exception as e: 

322 log.error(f'Error: {e}') 

323 yield f'data: {JSONCodec.dumps({"error": {"detail": str(e)}})}\n\n' 

324 return 

325 

326 if isinstance(res, str): 

327 message = openai_chat_chunk_message_template(form_data['model'], res) 

328 yield f'data: {JSONCodec.dumps(message)}\n\n' 

329 

330 if isinstance(res, Iterator): 

331 for line in res: 

332 yield process_line(form_data, line) 

333 

334 if isinstance(res, AsyncGenerator): 

335 async for line in res: 

336 yield process_line(form_data, line) 

337 

338 finish_message = openai_chat_chunk_message_template(form_data['model'], '') 

339 finish_message['choices'][0]['finish_reason'] = 'stop' 

340 yield f'data: {JSONCodec.dumps(finish_message)}\n\n' 

341 yield 'data: [DONE]' 

342 

343 return StreamingResponse(stream_content(), media_type='text/event-stream') 

344 else: 

345 try: 

346 res = await execute_pipe(pipe, params) 

347 

348 except Exception as e: 

349 log.error(f'Error: {e}') 

350 return {'error': {'detail': str(e)}} 

351 

352 if isinstance(res, StreamingResponse) or isinstance(res, dict): 

353 return res 

354 if isinstance(res, BaseModel): 

355 return res.model_dump() 

356 

357 message = await get_message_content(res) 

358 return openai_chat_completion_message_template(form_data['model'], message)