ds2api/internal/sse/consumer.go

package sse

import (
	"net/http"
	"strings"

	"ds2api/internal/deepseek"
)

// CollectResult holds the aggregated text and thinking content from a
// DeepSeek SSE stream, consumed to completion (non-streaming use case).
type CollectResult struct {
	Text          string
	Thinking      string
	PromptTokens  int
	OutputTokens  int
	ContentFilter bool
}

// CollectStream fully consumes a DeepSeek SSE response and separates
// thinking content from text content. This replaces the duplicated
// stream-collection logic in openai.handleNonStream, claude.collectDeepSeek,
// and admin.testAccount.
//
// The caller is responsible for closing resp.Body unless closeBody is true.
func CollectStream(resp *http.Response, thinkingEnabled bool, closeBody bool) CollectResult {
	if closeBody {
		defer func() { _ = resp.Body.Close() }()
	}
	text := strings.Builder{}
	thinking := strings.Builder{}
	promptTokens := 0
	outputTokens := 0
	contentFilter := false
	currentType := "text"
	if thinkingEnabled {
		currentType = "thinking"
	}
	_ = deepseek.ScanSSELines(resp, func(line []byte) bool {
		result := ParseDeepSeekContentLine(line, thinkingEnabled, currentType)
		currentType = result.NextType
		if !result.Parsed {
			return true
		}
		if result.PromptTokens > 0 {
			promptTokens = result.PromptTokens
		}
		if result.OutputTokens > 0 {
			outputTokens = result.OutputTokens
		}
		if result.Stop {
			if result.ContentFilter {
				contentFilter = true
			}
			return false
		}
		for _, p := range result.Parts {
			if p.Type == "thinking" {
				trimmed := TrimContinuationOverlap(thinking.String(), p.Text)
				thinking.WriteString(trimmed)
			} else {
				trimmed := TrimContinuationOverlap(text.String(), p.Text)
				text.WriteString(trimmed)
			}
		}
		return true
	})
	return CollectResult{
		Text:          text.String(),
		Thinking:      thinking.String(),
		PromptTokens:  promptTokens,
		OutputTokens:  outputTokens,
		ContentFilter: contentFilter,
	}
}