Jack1808/Claude_Code

1"""Think tag parser for extracting reasoning content from responses."""2 3from collections.abc import Iterator4from dataclasses import dataclass5from enum import Enum6 7 8class ContentType(Enum):9    """Type of content chunk."""10 11    TEXT = "text"12    THINKING = "thinking"13 14 15@dataclass16class ContentChunk:17    """A chunk of parsed content."""18 19    type: ContentType20    content: str21 22 23class ThinkTagParser:24    """25    Streaming parser for <think>...</think> tags.26 27    Handles partial tags at chunk boundaries by buffering.28    """29 30    OPEN_TAG = "<think>"31    CLOSE_TAG = "</think>"32    OPEN_TAG_LEN = 733    CLOSE_TAG_LEN = 834 35    def __init__(self):36        self._buffer: str = ""37        self._in_think_tag: bool = False38 39    @property40    def in_think_mode(self) -> bool:41        """Whether currently inside a think tag."""42        return self._in_think_tag43 44    def feed(self, content: str) -> Iterator[ContentChunk]:45        """46        Feed content and yield parsed chunks.47 48        Handles partial tags by buffering content near potential tag boundaries.49        Uses an iterative loop instead of mutual recursion to avoid stack overflow50        on inputs with many consecutive think tags.51        """52        self._buffer += content53 54        while self._buffer:55            prev_len = len(self._buffer)56            if not self._in_think_tag:57                chunk = self._parse_outside_think()58            else:59                chunk = self._parse_inside_think()60 61            if chunk:62                yield chunk63            elif len(self._buffer) == prev_len:64                # No progress: waiting for more data65                break66 67    def _parse_outside_think(self) -> ContentChunk | None:68        """Parse content outside think tags."""69        think_start = self._buffer.find(self.OPEN_TAG)70        orphan_close = self._buffer.find(self.CLOSE_TAG)71 72        # Handle orphan </think> - strip it (Step Fun AI sends reasoning via73        # reasoning_content but may leak closing tags in content)74        if orphan_close != -1 and (think_start == -1 or orphan_close < think_start):75            pre_orphan = self._buffer[:orphan_close]76            self._buffer = self._buffer[orphan_close + self.CLOSE_TAG_LEN :]77            if pre_orphan:78                return ContentChunk(ContentType.TEXT, pre_orphan)79            # Buffer shrunk; the feed() loop will continue parsing80            return None81 82        if think_start == -1:83            # No tag found - check for partial tag at end84            # We buffer any trailing '<' and subsequent characters that could be part of <think> or </think>85            last_bracket = self._buffer.rfind("<")86            if last_bracket != -1:87                potential_tag = self._buffer[last_bracket:]88                tag_len = len(potential_tag)89                # Check if could be partial <think> or </think>90                if (91                    tag_len < self.OPEN_TAG_LEN92                    and self.OPEN_TAG.startswith(potential_tag)93                ) or (94                    tag_len < self.CLOSE_TAG_LEN95                    and self.CLOSE_TAG.startswith(potential_tag)96                ):97                    emit = self._buffer[:last_bracket]98                    self._buffer = self._buffer[last_bracket:]99                    if emit:100                        return ContentChunk(ContentType.TEXT, emit)101                    return None102 103            # No partial tag found or it's irrelevant104            emit = self._buffer105            self._buffer = ""106            if emit:107                return ContentChunk(ContentType.TEXT, emit)108            return None109        else:110            # Found <think> tag111            pre_think = self._buffer[:think_start]112            self._buffer = self._buffer[think_start + self.OPEN_TAG_LEN :]113            self._in_think_tag = True114            if pre_think:115                return ContentChunk(ContentType.TEXT, pre_think)116            # Buffer shrunk (consumed <think>); the feed() loop will continue117            # parsing inside the think tag on the next iteration118            return None119 120    def _parse_inside_think(self) -> ContentChunk | None:121        """Parse content inside think tags."""122        think_end = self._buffer.find(self.CLOSE_TAG)123 124        if think_end == -1:125            # No closing tag - check for partial at end126            last_bracket = self._buffer.rfind("<")127            if (128                last_bracket != -1129                and len(self._buffer) - last_bracket < self.CLOSE_TAG_LEN130            ):131                # Check if the partial string could be the start of </think>132                potential_tag = self._buffer[last_bracket:]133                if self.CLOSE_TAG.startswith(potential_tag):134                    emit = self._buffer[:last_bracket]135                    self._buffer = self._buffer[last_bracket:]136                    if emit:137                        return ContentChunk(ContentType.THINKING, emit)138                    return None139 140            emit = self._buffer141            self._buffer = ""142            if emit:143                return ContentChunk(ContentType.THINKING, emit)144            return None145        else:146            # Found </think> tag147            thinking_content = self._buffer[:think_end]148            self._buffer = self._buffer[think_end + self.CLOSE_TAG_LEN :]149            self._in_think_tag = False150            if thinking_content:151                return ContentChunk(ContentType.THINKING, thinking_content)152            # Buffer shrunk (consumed </think>); the feed() loop will continue153            # parsing outside the think tag on the next iteration154            return None155 156    def flush(self) -> ContentChunk | None:157        """Flush any remaining buffered content."""158        if self._buffer:159            chunk_type = (160                ContentType.THINKING if self._in_think_tag else ContentType.TEXT161            )162            content = self._buffer163            self._buffer = ""164            return ContentChunk(chunk_type, content)165        return None166