Rhinox13/chatapi

1from __future__ import annotations
2
3import re
4from typing import Literal, TypedDict
5
6
7class ThinkingPart(TypedDict):
8    type: Literal["thinking", "answer"]
9    text: str
10
11
12_THINK_RE = re.compile(r"<think(?:\s[^>]*)?>\s*([\s\S]*?)\s*</think>", re.IGNORECASE)
13
14
15def split_thinking_parts(text: str) -> list[ThinkingPart]:
16    if not text:
17        return []
18
19    parts: list[ThinkingPart] = []
20    last_index = 0
21    for match in _THINK_RE.finditer(text):
22        before = text[last_index:match.start()].strip()
23        if before:
24            parts.append({"type": "answer", "text": before})
25        thinking = str(match.group(1) or "").strip()
26        if thinking:
27            parts.append({"type": "thinking", "text": thinking})
28        last_index = match.end()
29
30    after = text[last_index:].strip()
31    if after:
32        parts.append({"type": "answer", "text": after})
33
34    if not parts and text.strip():
35        parts.append({"type": "answer", "text": text.strip()})
36    return parts
37
38
39def has_thinking(text: str) -> bool:
40    return any(part["type"] == "thinking" for part in split_thinking_parts(text))
41
42
43def thinking_text(text: str) -> str:
44    return "\n\n".join(
45        part["text"] for part in split_thinking_parts(text) if part["type"] == "thinking"
46    ).strip()
47
48
49def answer_text(text: str) -> str:
50    return "\n\n".join(
51        part["text"] for part in split_thinking_parts(text) if part["type"] == "answer"
52    ).strip()
53
54
55def compose_thinking_text(parts: list[dict[str, str]]) -> str:
56    chunks: list[str] = []
57    for part in parts:
58        part_type = str(part.get("type") or "")
59        text = str(part.get("text") or "")
60        if not text:
61            continue
62        if part_type == "thinking":
63            chunks.append(f"<think>{text}</think>")
64        else:
65            chunks.append(text)
66    return "".join(chunks).strip()
67
68
69class ThinkingStreamParser:
70    def __init__(self) -> None:
71        self._buffer = ""
72        self._in_thinking = False
73
74    def feed(self, chunk: str) -> list[ThinkingPart]:
75        if not chunk:
76            return []
77        self._buffer += chunk
78        return self._drain(allow_incomplete_answer=True)
79
80    def flush(self) -> list[ThinkingPart]:
81        return self._drain(allow_incomplete_answer=True, flush=True)
82
83    def _drain(self, *, allow_incomplete_answer: bool, flush: bool = False) -> list[ThinkingPart]:
84        emitted: list[ThinkingPart] = []
85        while self._buffer:
86            lower = self._buffer.lower()
87            if self._in_thinking:
88                close_index = lower.find("</think>")
89                if close_index < 0:
90                    if flush:
91                        text = self._buffer.strip()
92                        if text:
93                            emitted.append({"type": "thinking", "text": text})
94                        self._buffer = ""
95                        self._in_thinking = False
96                    break
97                thinking = self._buffer[:close_index].strip()
98                if thinking:
99                    emitted.append({"type": "thinking", "text": thinking})
100                self._buffer = self._buffer[close_index + len("</think>"):].lstrip()
101                self._in_thinking = False
102                continue
103
104            open_index = lower.find("<think>")
105            if open_index < 0:
106                if allow_incomplete_answer or flush:
107                    text = self._buffer
108                    if not flush:
109                        keep = _partial_think_prefix_len(text)
110                        if keep > 0:
111                            emit_text = text[:-keep]
112                            self._buffer = text[-keep:]
113                        else:
114                            emit_text = text
115                            self._buffer = ""
116                    else:
117                        emit_text = text
118                        self._buffer = ""
119                    if emit_text.strip():
120                        emitted.append({"type": "answer", "text": emit_text})
121                break
122
123            before = self._buffer[:open_index]
124            if before.strip():
125                emitted.append({"type": "answer", "text": before})
126            self._buffer = self._buffer[open_index + len("<think>"):]
127            self._in_thinking = True
128        return emitted
129
130
131def _partial_think_prefix_len(text: str) -> int:
132    marker = "<think>"
133    lower = text.lower()
134    max_len = min(len(marker) - 1, len(lower))
135    for size in range(max_len, 0, -1):
136        if marker.startswith(lower[-size:]):
137            return size
138    return 0
139