enh: code interpreter

This commit is contained in:
Timothy Jaeryang Baek 2025-02-03 00:24:09 -08:00
parent 58d2fd9ac1
commit ca4b839e6d

View File

@ -1124,6 +1124,8 @@ async def process_chat_response(
return content return content
def tag_content_handler(content_type, tags, content, content_blocks): def tag_content_handler(content_type, tags, content, content_blocks):
end_flag = False
def extract_attributes(tag_content): def extract_attributes(tag_content):
"""Extract attributes from a tag if they exist.""" """Extract attributes from a tag if they exist."""
attributes = {} attributes = {}
@ -1173,6 +1175,7 @@ async def process_chat_response(
end_tag_pattern, "", block_content end_tag_pattern, "", block_content
).strip() ).strip()
if block_content: if block_content:
end_flag = True
content_blocks[-1]["content"] = block_content content_blocks[-1]["content"] = block_content
content_blocks[-1]["ended_at"] = time.time() content_blocks[-1]["ended_at"] = time.time()
content_blocks[-1]["duration"] = int( content_blocks[-1]["duration"] = int(
@ -1196,7 +1199,7 @@ async def process_chat_response(
else: else:
# Remove the block if content is empty # Remove the block if content is empty
content_blocks.pop() content_blocks.pop()
return content, content_blocks return content, content_blocks, end_flag
message = Chats.get_message_by_id_and_message_id( message = Chats.get_message_by_id_and_message_id(
metadata["chat_id"], metadata["message_id"] metadata["chat_id"], metadata["message_id"]
@ -1280,20 +1283,27 @@ async def process_chat_response(
) )
if DETECT_REASONING: if DETECT_REASONING:
content, content_blocks = tag_content_handler( content, content_blocks, _ = (
tag_content_handler(
"reasoning", "reasoning",
reasoning_tags, reasoning_tags,
content, content,
content_blocks, content_blocks,
) )
)
if DETECT_CODE_INTERPRETER: if DETECT_CODE_INTERPRETER:
content, content_blocks = tag_content_handler( content, content_blocks, end = (
tag_content_handler(
"code_interpreter", "code_interpreter",
code_interpreter_tags, code_interpreter_tags,
content, content,
content_blocks, content_blocks,
) )
)
if end:
break
if ENABLE_REALTIME_CHAT_SAVE: if ENABLE_REALTIME_CHAT_SAVE:
# Save message in the database # Save message in the database
@ -1320,9 +1330,13 @@ async def process_chat_response(
} }
) )
except Exception as e: except Exception as e:
done = "data: [DONE]" in line done = "data: [DONE]" in line
if done: if done:
pass
else:
log.debug("Error: ", e)
continue
# Clean up the last text block # Clean up the last text block
if content_blocks[-1]["type"] == "text": if content_blocks[-1]["type"] == "text":
content_blocks[-1]["content"] = content_blocks[-1][ content_blocks[-1]["content"] = content_blocks[-1][
@ -1331,10 +1345,6 @@ async def process_chat_response(
if not content_blocks[-1]["content"]: if not content_blocks[-1]["content"]:
content_blocks.pop() content_blocks.pop()
pass
else:
log.debug("Error: ", e)
continue
if response.background: if response.background:
await response.background() await response.background()
@ -1459,7 +1469,7 @@ async def process_chat_response(
metadata["chat_id"], metadata["chat_id"],
metadata["message_id"], metadata["message_id"],
{ {
"content": content, "content": serialize_content_blocks(content_blocks),
}, },
) )