File size: 2,742 Bytes
8d3471e | 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 | package shared
import (
"strings"
"ds2api/internal/sse"
)
type StreamAccumulator struct {
ThinkingEnabled bool
SearchEnabled bool
StripReferenceMarkers bool
RawThinking strings.Builder
Thinking strings.Builder
ToolDetectionThinking strings.Builder
RawText strings.Builder
Text strings.Builder
}
type StreamPartDelta struct {
Type string
RawText string
VisibleText string
CitationOnly bool
}
type StreamAccumulatorResult struct {
ContentSeen bool
Parts []StreamPartDelta
}
func (a *StreamAccumulator) Apply(parsed sse.LineResult) StreamAccumulatorResult {
out := StreamAccumulatorResult{}
for _, p := range parsed.ToolDetectionThinkingParts {
trimmed := sse.TrimContinuationOverlapFromBuilder(&a.ToolDetectionThinking, p.Text)
if trimmed != "" {
a.ToolDetectionThinking.WriteString(trimmed)
}
}
for _, p := range parsed.Parts {
if p.Type == "thinking" {
delta := a.applyThinkingPart(p.Text)
if delta.RawText != "" {
out.ContentSeen = true
}
if delta.RawText != "" || delta.VisibleText != "" {
out.Parts = append(out.Parts, delta)
}
continue
}
delta := a.applyTextPart(p.Text)
if delta.RawText != "" {
out.ContentSeen = true
}
if delta.RawText != "" || delta.VisibleText != "" || delta.CitationOnly {
out.Parts = append(out.Parts, delta)
}
}
return out
}
func (a *StreamAccumulator) applyThinkingPart(text string) StreamPartDelta {
rawTrimmed := sse.TrimContinuationOverlapFromBuilder(&a.RawThinking, text)
if rawTrimmed != "" {
a.RawThinking.WriteString(rawTrimmed)
}
delta := StreamPartDelta{Type: "thinking", RawText: rawTrimmed}
if !a.ThinkingEnabled || rawTrimmed == "" {
return delta
}
cleanedText := CleanVisibleOutput(rawTrimmed, a.StripReferenceMarkers)
if cleanedText == "" {
return delta
}
trimmed := sse.TrimContinuationOverlapFromBuilder(&a.Thinking, cleanedText)
if trimmed == "" {
return delta
}
a.Thinking.WriteString(trimmed)
delta.VisibleText = trimmed
return delta
}
func (a *StreamAccumulator) applyTextPart(text string) StreamPartDelta {
rawTrimmed := sse.TrimContinuationOverlapFromBuilder(&a.RawText, text)
if rawTrimmed == "" {
return StreamPartDelta{Type: "text"}
}
a.RawText.WriteString(rawTrimmed)
delta := StreamPartDelta{Type: "text", RawText: rawTrimmed}
if a.SearchEnabled && sse.IsCitation(rawTrimmed) {
delta.CitationOnly = true
return delta
}
cleanedText := CleanVisibleOutput(rawTrimmed, a.StripReferenceMarkers)
trimmed := sse.TrimContinuationOverlapFromBuilder(&a.Text, cleanedText)
if trimmed == "" {
return delta
}
a.Text.WriteString(trimmed)
delta.VisibleText = trimmed
return delta
}
|