CrossRelay/internal/anthropic/stream.go

package anthropic

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"net/http"

	"gitea.rutua.cn/kate522/CrossRelay/internal/chatcmpl"
)

// StreamConverter synthesizes Anthropic SSE events from Chat Completions
// stream chunks.
//
// Content blocks are tracked by their Anthropic index: one text block
// (opened on the first text delta) and one tool_use block per upstream
// tool-call index. Deltas always carry the block index, so interleaved
// tool-call arguments from the upstream stream are attributed to the
// right block. All open blocks are closed in Finish, followed by
// message_delta and message_stop.
type StreamConverter struct {
	w       http.ResponseWriter
	flusher http.Flusher
	model   string
	msgID   string

	started   bool
	nextIndex int

	textOpen  bool
	textIndex int

	tools     map[int]*toolBlock
	toolOrder []int

	inputEstimate int
	outputChars   int
	usage         *chatcmpl.Usage
	finishReason  string
}

type toolBlock struct {
	index   int
	id      string
	name    string
	started bool
	pending string
}

// NewStreamConverter builds a converter writing to w. model is echoed in
// message_start; inputTokens seeds the usage shown to the client.
func NewStreamConverter(w http.ResponseWriter, model string, inputTokens int) *StreamConverter {
	return &StreamConverter{
		w:             w,
		flusher:       flusherOf(w),
		model:         model,
		msgID:         "msg_" + randHex(16),
		tools:         map[int]*toolBlock{},
		inputEstimate: inputTokens,
	}
}

// WriteStart sends the SSE headers and the message_start event.
func (s *StreamConverter) WriteStart() error {
	h := s.w.Header()
	h.Set("Content-Type", "text/event-stream; charset=utf-8")
	h.Set("Cache-Control", "no-cache")
	h.Set("Connection", "keep-alive")
	h.Set("X-Accel-Buffering", "no")
	s.w.WriteHeader(http.StatusOK)
	s.started = true
	return writeEvent(s.w, s.flusher, "message_start", map[string]any{
		"type": "message_start",
		"message": map[string]any{
			"id":            s.msgID,
			"type":          "message",
			"role":          "assistant",
			"content":       []any{},
			"model":         s.model,
			"stop_reason":   nil,
			"stop_sequence": nil,
			"usage":         map[string]any{"input_tokens": s.inputEstimate, "output_tokens": 0},
		},
	})
}

// OnChunk folds one upstream chunk into the Anthropic event stream.
func (s *StreamConverter) OnChunk(chunk *chatcmpl.StreamChunk) error {
	if chunk == nil {
		return nil
	}
	if chunk.Usage != nil {
		u := *chunk.Usage
		s.usage = &u
	}
	for _, choice := range chunk.Choices {
		if choice.Delta.Content != nil && *choice.Delta.Content != "" {
			if err := s.writeText(*choice.Delta.Content); err != nil {
				return err
			}
		}
		for _, tc := range choice.Delta.ToolCalls {
			if err := s.toolDelta(tc); err != nil {
				return err
			}
		}
		if choice.FinishReason != nil && *choice.FinishReason != "" {
			s.finishReason = *choice.FinishReason
		}
	}
	return nil
}

func (s *StreamConverter) writeText(text string) error {
	if !s.textOpen {
		s.textIndex = s.nextIndex
		s.nextIndex++
		s.textOpen = true
		if err := writeEvent(s.w, s.flusher, "content_block_start", map[string]any{
			"type":          "content_block_start",
			"index":         s.textIndex,
			"content_block": map[string]any{"type": "text", "text": ""},
		}); err != nil {
			return err
		}
	}
	s.outputChars += len(text)
	return writeEvent(s.w, s.flusher, "content_block_delta", map[string]any{
		"type":  "content_block_delta",
		"index": s.textIndex,
		"delta": map[string]any{"type": "text_delta", "text": text},
	})
}

func (s *StreamConverter) toolDelta(tc chatcmpl.ToolCall) error {
	st := s.tools[tc.Index]
	if st == nil {
		st = &toolBlock{}
		s.tools[tc.Index] = st
		s.toolOrder = append(s.toolOrder, tc.Index)
	}
	if tc.ID != "" && st.id == "" {
		st.id = tc.ID
	}
	if tc.Function.Name != "" && st.name == "" {
		st.name = tc.Function.Name
	}
	if tc.Function.Arguments != "" {
		st.pending += tc.Function.Arguments
	}
	// Open the block once we know a name or have arguments to flush, so
	// the client never sees an unnamed tool_use block.
	if !st.started && (st.name != "" || st.pending != "") {
		st.index = s.nextIndex
		s.nextIndex++
		st.started = true
		if st.id == "" {
			st.id = fmt.Sprintf("toolu_%s_%d", s.msgID[len("msg_"):], tc.Index)
		}
		if err := writeEvent(s.w, s.flusher, "content_block_start", map[string]any{
			"type":  "content_block_start",
			"index": st.index,
			"content_block": map[string]any{
				"type":  "tool_use",
				"id":    st.id,
				"name":  st.name,
				"input": map[string]any{},
			},
		}); err != nil {
			return err
		}
		s.outputChars += len(st.name)
	}
	if st.started && st.pending != "" {
		s.outputChars += len(st.pending)
		pending := st.pending
		st.pending = ""
		return writeEvent(s.w, s.flusher, "content_block_delta", map[string]any{
			"type":  "content_block_delta",
			"index": st.index,
			"delta": map[string]any{"type": "input_json_delta", "partial_json": pending},
		})
	}
	return nil
}

// Finish closes every open block and ends the message.
func (s *StreamConverter) Finish() error {
	if s.textOpen {
		if err := writeEvent(s.w, s.flusher, "content_block_stop", map[string]any{
			"type":  "content_block_stop",
			"index": s.textIndex,
		}); err != nil {
			return err
		}
		s.textOpen = false
	}
	for _, idx := range s.toolOrder {
		st := s.tools[idx]
		if st == nil || !st.started {
			continue
		}
		if err := writeEvent(s.w, s.flusher, "content_block_stop", map[string]any{
			"type":  "content_block_stop",
			"index": st.index,
		}); err != nil {
			return err
		}
	}
	outputTokens := 0
	if s.usage != nil {
		outputTokens = s.usage.CompletionTokens
	} else {
		outputTokens = (s.outputChars + 3) / 4
	}
	usage := map[string]any{"output_tokens": outputTokens}
	if s.usage != nil {
		usage["input_tokens"] = s.usage.PromptTokens
	}
	if err := writeEvent(s.w, s.flusher, "message_delta", map[string]any{
		"type":  "message_delta",
		"delta": map[string]any{"stop_reason": MapStopReason(s.finishReason), "stop_sequence": nil},
		"usage": usage,
	}); err != nil {
		return err
	}
	return writeEvent(s.w, s.flusher, "message_stop", map[string]any{"type": "message_stop"})
}

// WriteError terminates the stream with an Anthropic error event.
func (s *StreamConverter) WriteError(errType, message string) {
	_ = writeEvent(s.w, s.flusher, "error", map[string]any{
		"type":  "error",
		"error": map[string]any{"type": errType, "message": message},
	})
}

func randHex(n int) string {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return fmt.Sprintf("%x", n)
	}
	return hex.EncodeToString(b)
}