diff --git a/modules/html_generator.py b/modules/html_generator.py index 6844c244..8777acf7 100644 --- a/modules/html_generator.py +++ b/modules/html_generator.py @@ -116,29 +116,56 @@ def extract_thinking_block(string): THINK_START_TAG = "<think>" THINK_END_TAG = "</think>" - # Look for think tag + # Look for think tag first start_pos = string.find(THINK_START_TAG) end_pos = string.find(THINK_END_TAG) - # Return if neither tag is in string - if start_pos == -1 and end_pos == -1: - return None, string + # If think tags found, use existing logic + if start_pos != -1 or end_pos != -1: + # handle missing start or end tags + if start_pos == -1: + thought_start = 0 + else: + thought_start = start_pos + len(THINK_START_TAG) + if end_pos == -1: + thought_end = len(string) + content_start = len(string) + else: + thought_end = end_pos + content_start = end_pos + len(THINK_END_TAG) + thinking_content = string[thought_start:thought_end] + remaining_content = string[content_start:] + return thinking_content, remaining_content - # handle missing start or end tags - if start_pos == -1: - thought_start = 0 - else: - thought_start = start_pos + len(THINK_START_TAG) - if end_pos == -1: - thought_end = len(string) - content_start = len(string) - else: - thought_end = end_pos - content_start = end_pos + len(THINK_END_TAG) + # If think tags not found, try alternative format + ALT_START = "<|channel|>analysis<|message|>" + ALT_END = "<|end|>" + ALT_CONTENT_START = "<|start|>assistant<|channel|>final<|message|>" - thinking_content = string[thought_start:thought_end] - remaining_content = string[content_start:] - return thinking_content, remaining_content + alt_start_pos = string.find(ALT_START) + alt_end_pos = string.find(ALT_END) + alt_content_pos = string.find(ALT_CONTENT_START) + + # Check if start tag or end tag is found + if alt_start_pos != -1 or alt_end_pos != -1: + if alt_start_pos == -1: + thought_start = 0 + else: + thought_start = alt_start_pos + len(ALT_START) + + if alt_end_pos == -1: + thought_end = len(string) + content_start = len(string) + else: + thought_end = alt_end_pos + content_start = alt_content_pos + len(ALT_CONTENT_START) if alt_content_pos != -1 else len(string) + + thinking_content = string[thought_start:thought_end] + remaining_content = string[content_start:] + return thinking_content, remaining_content + + # Return if neither format is found + return None, string @functools.lru_cache(maxsize=None)