From ca4b839e6d9e62866d23e294cfd7c7bdacce0f4f Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Mon, 3 Feb 2025 00:24:09 -0800 Subject: [PATCH] enh: code interpreter --- backend/open_webui/utils/middleware.py | 52 +++++++++++++++----------- 1 file changed, 31 insertions(+), 21 deletions(-) diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index c93c840bf..49318b790 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -1124,6 +1124,8 @@ async def process_chat_response( return content def tag_content_handler(content_type, tags, content, content_blocks): + end_flag = False + def extract_attributes(tag_content): """Extract attributes from a tag if they exist.""" attributes = {} @@ -1173,6 +1175,7 @@ async def process_chat_response( end_tag_pattern, "", block_content ).strip() if block_content: + end_flag = True content_blocks[-1]["content"] = block_content content_blocks[-1]["ended_at"] = time.time() content_blocks[-1]["duration"] = int( @@ -1196,7 +1199,7 @@ async def process_chat_response( else: # Remove the block if content is empty content_blocks.pop() - return content, content_blocks + return content, content_blocks, end_flag message = Chats.get_message_by_id_and_message_id( metadata["chat_id"], metadata["message_id"] @@ -1280,21 +1283,28 @@ async def process_chat_response( ) if DETECT_REASONING: - content, content_blocks = tag_content_handler( - "reasoning", - reasoning_tags, - content, - content_blocks, + content, content_blocks, _ = ( + tag_content_handler( + "reasoning", + reasoning_tags, + content, + content_blocks, + ) ) if DETECT_CODE_INTERPRETER: - content, content_blocks = tag_content_handler( - "code_interpreter", - code_interpreter_tags, - content, - content_blocks, + content, content_blocks, end = ( + tag_content_handler( + "code_interpreter", + code_interpreter_tags, + content, + content_blocks, + ) ) + if end: + break + if ENABLE_REALTIME_CHAT_SAVE: # Save message in the database Chats.upsert_message_to_chat_by_id_and_message_id( @@ -1320,22 +1330,22 @@ async def process_chat_response( } ) except Exception as e: - done = "data: [DONE]" in line if done: - # Clean up the last text block - if content_blocks[-1]["type"] == "text": - content_blocks[-1]["content"] = content_blocks[-1][ - "content" - ].strip() - - if not content_blocks[-1]["content"]: - content_blocks.pop() pass else: log.debug("Error: ", e) continue + # Clean up the last text block + if content_blocks[-1]["type"] == "text": + content_blocks[-1]["content"] = content_blocks[-1][ + "content" + ].strip() + + if not content_blocks[-1]["content"]: + content_blocks.pop() + if response.background: await response.background() @@ -1459,7 +1469,7 @@ async def process_chat_response( metadata["chat_id"], metadata["message_id"], { - "content": content, + "content": serialize_content_blocks(content_blocks), }, )