fix: tag support

This commit is contained in:
Timothy Jaeryang Baek 2025-02-05 02:33:40 -08:00
parent c36c63f1e8
commit 90cd6f272f

View File

@ -1231,18 +1231,10 @@ async def process_chat_response(
content_blocks[-1]["content"] = content_blocks[-1][ content_blocks[-1]["content"] = content_blocks[-1][
"content" "content"
].replace(match.group(0), "") ].replace(match.group(0), "")
if not content_blocks[-1]["content"]: if not content_blocks[-1]["content"]:
content_blocks.pop() content_blocks.pop()
if not content_blocks:
# Append the new block
content_blocks.append(
{
"type": "text",
"content": "",
}
)
# Append the new block # Append the new block
content_blocks.append( content_blocks.append(
{ {
@ -1258,6 +1250,7 @@ async def process_chat_response(
tag = content_blocks[-1]["tag"] tag = content_blocks[-1]["tag"]
# Match end tag e.g., </tag> # Match end tag e.g., </tag>
end_tag_pattern = rf"</{tag}>" end_tag_pattern = rf"</{tag}>"
if re.search(end_tag_pattern, content): if re.search(end_tag_pattern, content):
block_content = content_blocks[-1]["content"] block_content = content_blocks[-1]["content"]
# Strip start and end tags from the content # Strip start and end tags from the content
@ -1265,9 +1258,23 @@ async def process_chat_response(
block_content = re.sub( block_content = re.sub(
start_tag_pattern, "", block_content start_tag_pattern, "", block_content
).strip() ).strip()
block_content = re.sub(
end_tag_pattern, "", block_content end_tag_regex = re.compile(end_tag_pattern, re.DOTALL)
).strip() split_content = end_tag_regex.split(block_content, maxsplit=1)
# Content inside the tag
block_content = (
split_content[0].strip() if split_content else ""
)
# Leftover content (everything after `</tag>`)
leftover_content = (
split_content[1].strip() if len(split_content) > 1 else ""
)
print(f"block_content: {block_content}")
print(f"leftover_content: {leftover_content}")
if block_content: if block_content:
end_flag = True end_flag = True
content_blocks[-1]["content"] = block_content content_blocks[-1]["content"] = block_content
@ -1280,19 +1287,31 @@ async def process_chat_response(
content_blocks.append( content_blocks.append(
{ {
"type": "text", "type": "text",
"content": "", "content": leftover_content,
} }
) )
# Clean processed content
content = re.sub(
rf"<{tag}(.*?)>(.|\n)*?</{tag}>",
"",
content,
flags=re.DOTALL,
)
else: else:
end_flag = True
# Remove the block if content is empty # Remove the block if content is empty
content_blocks.pop() content_blocks.pop()
if leftover_content:
content_blocks.append(
{
"type": "text",
"content": leftover_content,
}
)
# Clean processed content
content = re.sub(
rf"<{tag}(.*?)>(.|\n)*?</{tag}>",
"",
content,
flags=re.DOTALL,
)
return content, content_blocks, end_flag return content, content_blocks, end_flag
message = Chats.get_message_by_id_and_message_id( message = Chats.get_message_by_id_and_message_id(
@ -1358,6 +1377,7 @@ async def process_chat_response(
try: try:
data = json.loads(data) data = json.loads(data)
print(data)
if "selected_model_id" in data: if "selected_model_id" in data:
model_id = data["selected_model_id"] model_id = data["selected_model_id"]
@ -1412,6 +1432,15 @@ async def process_chat_response(
if value: if value:
content = f"{content}{value}" content = f"{content}{value}"
if not content_blocks:
content_blocks.append(
{
"type": "text",
"content": "",
}
)
content_blocks[-1]["content"] = ( content_blocks[-1]["content"] = (
content_blocks[-1]["content"] + value content_blocks[-1]["content"] + value
) )