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,21 +1283,28 @@ async def process_chat_response(
) )
if DETECT_REASONING: if DETECT_REASONING:
content, content_blocks = tag_content_handler( content, content_blocks, _ = (
"reasoning", tag_content_handler(
reasoning_tags, "reasoning",
content, reasoning_tags,
content_blocks, content,
content_blocks,
)
) )
if DETECT_CODE_INTERPRETER: if DETECT_CODE_INTERPRETER:
content, content_blocks = tag_content_handler( content, content_blocks, end = (
"code_interpreter", tag_content_handler(
code_interpreter_tags, "code_interpreter",
content, code_interpreter_tags,
content_blocks, content,
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
Chats.upsert_message_to_chat_by_id_and_message_id( Chats.upsert_message_to_chat_by_id_and_message_id(
@ -1320,22 +1330,22 @@ 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:
# 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 pass
else: else:
log.debug("Error: ", e) log.debug("Error: ", e)
continue 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: 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),
}, },
) )