File size: 4,654 Bytes
4fc4790
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import Foundation

struct AssistantTextSegment: Identifiable {
    enum Kind {
        case thinking
        case response
    }

    let id = UUID()
    let kind: Kind
    let text: String
}

enum AssistantTextParser {
    static func segments(from raw: String) -> [AssistantTextSegment] {
        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else { return [] }
        guard raw.contains("<") else {
            return [AssistantTextSegment(kind: .response, text: trimmed)]
        }

        var segments: [AssistantTextSegment] = []
        var cursor = raw.startIndex
        var currentKind: AssistantTextSegment.Kind = .response
        var matchedTag = false

        while let match = self.nextTag(in: raw, from: cursor) {
            matchedTag = true
            if match.range.lowerBound > cursor {
                self.appendSegment(kind: currentKind, text: raw[cursor..<match.range.lowerBound], to: &segments)
            }

            guard let tagEnd = raw.range(of: ">", range: match.range.upperBound..<raw.endIndex) else {
                cursor = raw.endIndex
                break
            }

            let isSelfClosing = self.isSelfClosingTag(in: raw, tagEnd: tagEnd)
            cursor = tagEnd.upperBound
            if isSelfClosing { continue }

            if match.closing {
                currentKind = .response
            } else {
                currentKind = match.kind == .think ? .thinking : .response
            }
        }

        if cursor < raw.endIndex {
            self.appendSegment(kind: currentKind, text: raw[cursor..<raw.endIndex], to: &segments)
        }

        guard matchedTag else {
            return [AssistantTextSegment(kind: .response, text: trimmed)]
        }

        return segments
    }

    static func hasVisibleContent(in raw: String) -> Bool {
        !self.segments(from: raw).isEmpty
    }

    private enum TagKind {
        case think
        case final
    }

    private struct TagMatch {
        let kind: TagKind
        let closing: Bool
        let range: Range<String.Index>
    }

    private static func nextTag(in text: String, from start: String.Index) -> TagMatch? {
        let candidates: [TagMatch] = [
            self.findTagStart(tag: "think", closing: false, in: text, from: start).map {
                TagMatch(kind: .think, closing: false, range: $0)
            },
            self.findTagStart(tag: "think", closing: true, in: text, from: start).map {
                TagMatch(kind: .think, closing: true, range: $0)
            },
            self.findTagStart(tag: "final", closing: false, in: text, from: start).map {
                TagMatch(kind: .final, closing: false, range: $0)
            },
            self.findTagStart(tag: "final", closing: true, in: text, from: start).map {
                TagMatch(kind: .final, closing: true, range: $0)
            },
        ].compactMap(\.self)

        return candidates.min { $0.range.lowerBound < $1.range.lowerBound }
    }

    private static func findTagStart(
        tag: String,
        closing: Bool,
        in text: String,
        from start: String.Index) -> Range<String.Index>?
    {
        let token = closing ? "</\(tag)" : "<\(tag)"
        var searchRange = start..<text.endIndex
        while let range = text.range(
            of: token,
            options: [.caseInsensitive, .diacriticInsensitive],
            range: searchRange)
        {
            let boundaryIndex = range.upperBound
            guard boundaryIndex < text.endIndex else { return range }
            let boundary = text[boundaryIndex]
            let isBoundary = boundary == ">" || boundary.isWhitespace || (!closing && boundary == "/")
            if isBoundary {
                return range
            }
            searchRange = boundaryIndex..<text.endIndex
        }
        return nil
    }

    private static func isSelfClosingTag(in text: String, tagEnd: Range<String.Index>) -> Bool {
        var cursor = tagEnd.lowerBound
        while cursor > text.startIndex {
            cursor = text.index(before: cursor)
            let char = text[cursor]
            if char.isWhitespace { continue }
            return char == "/"
        }
        return false
    }

    private static func appendSegment(
        kind: AssistantTextSegment.Kind,
        text: Substring,
        to segments: inout [AssistantTextSegment])
    {
        let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else { return }
        segments.append(AssistantTextSegment(kind: kind, text: trimmed))
    }
}