|
19 | 19 |
|
20 | 20 |
|
21 | 21 | from fastapi import Request, HTTPException |
22 | | -from starlette.responses import Response, StreamingResponse |
| 22 | +from starlette.responses import Response, StreamingResponse, JSONResponse |
23 | 23 |
|
24 | 24 |
|
25 | 25 | from open_webui.models.chats import Chats |
@@ -1254,91 +1254,111 @@ async def background_tasks_handler(): |
1254 | 1254 | # Non-streaming response |
1255 | 1255 | if not isinstance(response, StreamingResponse): |
1256 | 1256 | if event_emitter: |
1257 | | - if "error" in response: |
1258 | | - error = response["error"].get("detail", response["error"]) |
1259 | | - Chats.upsert_message_to_chat_by_id_and_message_id( |
1260 | | - metadata["chat_id"], |
1261 | | - metadata["message_id"], |
1262 | | - { |
1263 | | - "error": {"content": error}, |
1264 | | - }, |
1265 | | - ) |
1266 | | - |
1267 | | - if "selected_model_id" in response: |
1268 | | - Chats.upsert_message_to_chat_by_id_and_message_id( |
1269 | | - metadata["chat_id"], |
1270 | | - metadata["message_id"], |
1271 | | - { |
1272 | | - "selectedModelId": response["selected_model_id"], |
1273 | | - }, |
1274 | | - ) |
1275 | | - |
1276 | | - choices = response.get("choices", []) |
1277 | | - if choices and choices[0].get("message", {}).get("content"): |
1278 | | - content = response["choices"][0]["message"]["content"] |
1279 | | - |
1280 | | - if content: |
| 1257 | + if isinstance(response, dict) or isinstance(response, JSONResponse): |
1281 | 1258 |
|
1282 | | - await event_emitter( |
1283 | | - { |
1284 | | - "type": "chat:completion", |
1285 | | - "data": response, |
1286 | | - } |
1287 | | - ) |
1288 | | - |
1289 | | - title = Chats.get_chat_title_by_id(metadata["chat_id"]) |
1290 | | - |
1291 | | - await event_emitter( |
| 1259 | + if isinstance(response, JSONResponse) and isinstance( |
| 1260 | + response.body, bytes |
| 1261 | + ): |
| 1262 | + try: |
| 1263 | + response_data = json.loads(response.body.decode("utf-8")) |
| 1264 | + except json.JSONDecodeError: |
| 1265 | + response_data = {"error": {"detail": "Invalid JSON response"}} |
| 1266 | + else: |
| 1267 | + response_data = response |
| 1268 | + |
| 1269 | + if "error" in response_data: |
| 1270 | + error = response_data["error"].get("detail", response_data["error"]) |
| 1271 | + Chats.upsert_message_to_chat_by_id_and_message_id( |
| 1272 | + metadata["chat_id"], |
| 1273 | + metadata["message_id"], |
1292 | 1274 | { |
1293 | | - "type": "chat:completion", |
1294 | | - "data": { |
1295 | | - "done": True, |
1296 | | - "content": content, |
1297 | | - "title": title, |
1298 | | - }, |
1299 | | - } |
| 1275 | + "error": {"content": error}, |
| 1276 | + }, |
1300 | 1277 | ) |
1301 | 1278 |
|
1302 | | - # Save message in the database |
| 1279 | + if "selected_model_id" in response_data: |
1303 | 1280 | Chats.upsert_message_to_chat_by_id_and_message_id( |
1304 | 1281 | metadata["chat_id"], |
1305 | 1282 | metadata["message_id"], |
1306 | 1283 | { |
1307 | | - "role": "assistant", |
1308 | | - "content": content, |
| 1284 | + "selectedModelId": response_data["selected_model_id"], |
1309 | 1285 | }, |
1310 | 1286 | ) |
1311 | 1287 |
|
1312 | | - # Send a webhook notification if the user is not active |
1313 | | - if not get_active_status_by_user_id(user.id): |
1314 | | - webhook_url = Users.get_user_webhook_url_by_id(user.id) |
1315 | | - if webhook_url: |
1316 | | - post_webhook( |
1317 | | - request.app.state.WEBUI_NAME, |
1318 | | - webhook_url, |
1319 | | - f"{title} - {request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}\n\n{content}", |
1320 | | - { |
1321 | | - "action": "chat", |
1322 | | - "message": content, |
| 1288 | + choices = response_data.get("choices", []) |
| 1289 | + if choices and choices[0].get("message", {}).get("content"): |
| 1290 | + content = response_data["choices"][0]["message"]["content"] |
| 1291 | + |
| 1292 | + if content: |
| 1293 | + await event_emitter( |
| 1294 | + { |
| 1295 | + "type": "chat:completion", |
| 1296 | + "data": response_data, |
| 1297 | + } |
| 1298 | + ) |
| 1299 | + |
| 1300 | + title = Chats.get_chat_title_by_id(metadata["chat_id"]) |
| 1301 | + |
| 1302 | + await event_emitter( |
| 1303 | + { |
| 1304 | + "type": "chat:completion", |
| 1305 | + "data": { |
| 1306 | + "done": True, |
| 1307 | + "content": content, |
1323 | 1308 | "title": title, |
1324 | | - "url": f"{request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}", |
1325 | 1309 | }, |
1326 | | - ) |
| 1310 | + } |
| 1311 | + ) |
1327 | 1312 |
|
1328 | | - await background_tasks_handler() |
| 1313 | + # Save message in the database |
| 1314 | + Chats.upsert_message_to_chat_by_id_and_message_id( |
| 1315 | + metadata["chat_id"], |
| 1316 | + metadata["message_id"], |
| 1317 | + { |
| 1318 | + "role": "assistant", |
| 1319 | + "content": content, |
| 1320 | + }, |
| 1321 | + ) |
1329 | 1322 |
|
1330 | | - if events and isinstance(events, list) and isinstance(response, dict): |
1331 | | - extra_response = {} |
1332 | | - for event in events: |
1333 | | - if isinstance(event, dict): |
1334 | | - extra_response.update(event) |
1335 | | - else: |
1336 | | - extra_response[event] = True |
| 1323 | + # Send a webhook notification if the user is not active |
| 1324 | + if not get_active_status_by_user_id(user.id): |
| 1325 | + webhook_url = Users.get_user_webhook_url_by_id(user.id) |
| 1326 | + if webhook_url: |
| 1327 | + post_webhook( |
| 1328 | + request.app.state.WEBUI_NAME, |
| 1329 | + webhook_url, |
| 1330 | + f"{title} - {request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}\n\n{content}", |
| 1331 | + { |
| 1332 | + "action": "chat", |
| 1333 | + "message": content, |
| 1334 | + "title": title, |
| 1335 | + "url": f"{request.app.state.config.WEBUI_URL}/c/{metadata['chat_id']}", |
| 1336 | + }, |
| 1337 | + ) |
1337 | 1338 |
|
1338 | | - response = { |
1339 | | - **extra_response, |
1340 | | - **response, |
1341 | | - } |
| 1339 | + await background_tasks_handler() |
| 1340 | + |
| 1341 | + if events and isinstance(events, list): |
| 1342 | + extra_response = {} |
| 1343 | + for event in events: |
| 1344 | + if isinstance(event, dict): |
| 1345 | + extra_response.update(event) |
| 1346 | + else: |
| 1347 | + extra_response[event] = True |
| 1348 | + |
| 1349 | + response_data = { |
| 1350 | + **extra_response, |
| 1351 | + **response_data, |
| 1352 | + } |
| 1353 | + |
| 1354 | + if isinstance(response, dict): |
| 1355 | + response = response_data |
| 1356 | + if isinstance(response, JSONResponse): |
| 1357 | + response = JSONResponse( |
| 1358 | + content=response_data, |
| 1359 | + headers=response.headers, |
| 1360 | + status_code=response.status_code, |
| 1361 | + ) |
1342 | 1362 |
|
1343 | 1363 | return response |
1344 | 1364 | else: |
|
0 commit comments