File size: 1,461 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
package claude

import (
	"encoding/json"
	"fmt"
	"strings"

	"ds2api/internal/util"
)

func (s *claudeStreamRuntime) send(event string, v any) {
	b, _ := json.Marshal(v)
	_, _ = s.w.Write([]byte("event: "))
	_, _ = s.w.Write([]byte(event))
	_, _ = s.w.Write([]byte("\n"))
	_, _ = s.w.Write([]byte("data: "))
	_, _ = s.w.Write(b)
	_, _ = s.w.Write([]byte("\n\n"))
	if s.canFlush {
		_ = s.rc.Flush()
	}
}

func (s *claudeStreamRuntime) sendError(message string) {
	msg := strings.TrimSpace(message)
	if msg == "" {
		msg = "upstream stream error"
	}
	s.send("error", map[string]any{
		"type": "error",
		"error": map[string]any{
			"type":    "api_error",
			"message": msg,
			"code":    "internal_error",
			"param":   nil,
		},
	})
}

func (s *claudeStreamRuntime) sendPing() {
	s.send("ping", map[string]any{"type": "ping"})
}

func (s *claudeStreamRuntime) sendMessageStart() {
	inputTokens := countClaudeInputTokensFromText(s.promptTokenText, s.model)
	if inputTokens == 0 {
		inputTokens = util.CountPromptTokens(fmt.Sprintf("%v", s.messages), s.model)
	}
	s.send("message_start", map[string]any{
		"type": "message_start",
		"message": map[string]any{
			"id":            s.messageID,
			"type":          "message",
			"role":          "assistant",
			"model":         s.model,
			"content":       []any{},
			"stop_reason":   nil,
			"stop_sequence": nil,
			"usage":         map[string]any{"input_tokens": inputTokens, "output_tokens": 0},
		},
	})
}