Spaces:
Paused
Paused
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))
}
}
|