diff --git a/translate_anthropic.go b/translate_anthropic.go index 313fd4b..4935b9d 100644 --- a/translate_anthropic.go +++ b/translate_anthropic.go @@ -693,10 +693,8 @@ type anthropicStreamState struct { // index also counts any preceding non-tool_use blocks (a tool_use // block right after a text block is content-block index 1, but it is // still the first — index 0 — tool call). - toolOrdinals map[int]int - // contentTail caches the per-stream suffix contentChunk appends after - // the delta text; nil until first use and reset when id changes. - contentTail []byte + toolOrdinals map[int]int + content contentChunker id string model string u usage @@ -709,7 +707,9 @@ type anthropicStreamState struct { // created the response timestamp every emitted chunk carries — the // adapter fixes both once, before the first upstream event arrives. func newAnthropicStreamState(model string, created int64) *anthropicStreamState { - return &anthropicStreamState{model: model, created: created, toolOrdinals: map[int]int{}} + st := &anthropicStreamState{model: model, created: created, toolOrdinals: map[int]int{}} + st.content.set(chatCompletionIDPrefix, model, created) + return st } // usage returns the prompt/completion token counts captured so far, from @@ -745,37 +745,45 @@ func (st *anthropicStreamState) chunk(delta map[string]any, finishReason *string return b } -// contentChunkHead frames a text-delta chunk with -// the same alphabetical key order json.Marshal gives chunk's map, so the -// bytes are identical to chunk(map[string]any{"content": text}, nil). +// contentChunkHead frames a text-delta chunk with the same alphabetical +// key order json.Marshal gives a map, so the bytes are identical to the +// chunk method's output for map[string]any{"content": text}. const contentChunkHead = `{"choices":[{"delta":{"content":` -// contentChunk builds a text-delta chunk without allocating the envelope -// maps chunk needs for arbitrary deltas. -func (st *anthropicStreamState) contentChunk(text string) []byte { +// contentChunker builds text-delta chunks without allocating the envelope +// maps a state's chunk method needs for arbitrary deltas. +type contentChunker struct { + // tail is the per-stream suffix after the delta text, rebuilt by set. + tail []byte +} + +// set fixes the envelope fields every later chunk carries; call it again +// whenever the stream id changes. +func (c *contentChunker) set(id, model string, created int64) { + idb, ierr := json.Marshal(id) + mb, merr := json.Marshal(model) + if ierr != nil || merr != nil { + panic("llmgateway: stream chunk failed to marshal id or model") + } + t := append([]byte(nil), `},"finish_reason":null,"index":0}],"created":`...) + t = strconv.AppendInt(t, created, 10) + t = append(t, `,"id":`...) + t = append(t, idb...) + t = append(t, `,"model":`...) + t = append(t, mb...) + t = append(t, `,"object":"chat.completion.chunk"}`...) + c.tail = t +} + +func (c *contentChunker) build(text string) []byte { tb, err := json.Marshal(text) if err != nil { - panic(fmt.Sprintf("llmgateway: anthropic stream chunk failed to marshal: %v", err)) + panic(fmt.Sprintf("llmgateway: stream chunk failed to marshal text: %v", err)) } - if st.contentTail == nil { - idb, ierr := json.Marshal(chatCompletionIDPrefix + st.id) - mb, merr := json.Marshal(st.model) - if ierr != nil || merr != nil { - panic("llmgateway: anthropic stream chunk failed to marshal id or model") - } - t := append([]byte(nil), `},"finish_reason":null,"index":0}],"created":`...) - t = strconv.AppendInt(t, st.created, 10) - t = append(t, `,"id":`...) - t = append(t, idb...) - t = append(t, `,"model":`...) - t = append(t, mb...) - t = append(t, `,"object":"chat.completion.chunk"}`...) - st.contentTail = t - } - out := make([]byte, 0, len(contentChunkHead)+len(tb)+len(st.contentTail)) + out := make([]byte, 0, len(contentChunkHead)+len(tb)+len(c.tail)) out = append(out, contentChunkHead...) out = append(out, tb...) - return append(out, st.contentTail...) + return append(out, c.tail...) } // usageChunk builds the final usage-only "chat.completion.chunk" payload — @@ -891,7 +899,7 @@ func (st *anthropicStreamState) translate(ev sseEvent) ([][]byte, error) { return nil, fmt.Errorf("%w: decode anthropic message_start event: %w", errUpstream, err) } st.id = p.Message.ID - st.contentTail = nil + st.content.set(chatCompletionIDPrefix+st.id, st.model, st.created) st.u.prompt = p.Message.Usage.totalInputTokens() return [][]byte{st.chunk(map[string]any{"role": "assistant"}, nil)}, nil @@ -924,7 +932,7 @@ func (st *anthropicStreamState) translate(ev sseEvent) ([][]byte, error) { } switch p.Delta.Type { case "text_delta": - return [][]byte{st.contentChunk(p.Delta.Text)}, nil + return [][]byte{st.content.build(p.Delta.Text)}, nil case "input_json_delta": ordinal, ok := st.toolOrdinals[p.Index] if !ok { diff --git a/translate_anthropic_test.go b/translate_anthropic_test.go index 13280c4..d630812 100644 --- a/translate_anthropic_test.go +++ b/translate_anthropic_test.go @@ -1129,16 +1129,23 @@ func TestOpenAIRequestFromAnthropic_NoTools_OmitsToolsAndToolChoiceFields(t *tes } } -func TestAnthropicContentChunkMatchesChunk(t *testing.T) { +func TestContentChunkerMatchesChunk(t *testing.T) { texts := []string{"", "Hello", "a \"quoted\" &", "line\nbreak\ttab", "unicode é 日本 \u2028\u2029", "bad\xffutf8"} for _, id := range []string{"", "msg_01", "we\"ird"} { - st := newAnthropicStreamState("gw-model", 1734000000) - st.id = id + ast := newAnthropicStreamState("gw-model", 1734000000) + ast.id = id + ast.content.set(chatCompletionIDPrefix+id, ast.model, ast.created) + gst := newGeminiStreamState("gw-model", 1734000000) + gst.id = id + gst.content.set(chatCompletionIDPrefix+id, gst.model, gst.created) for _, text := range texts { - got := st.contentChunk(text) - want := st.chunk(map[string]any{"content": text}, nil) - if !bytes.Equal(got, want) { - t.Errorf("id=%q text=%q:\n got %s\nwant %s", id, text, got, want) + gotA := ast.content.build(text) + if want := ast.chunk(map[string]any{"content": text}, nil); !bytes.Equal(gotA, want) { + t.Errorf("anthropic id=%q text=%q:\n got %s\nwant %s", id, text, gotA, want) + } + gotG := gst.content.build(text) + if want := gst.chunk(map[string]any{"content": text}, nil); !bytes.Equal(gotG, want) { + t.Errorf("gemini id=%q text=%q:\n got %s\nwant %s", id, text, gotG, want) } } } diff --git a/translate_gemini.go b/translate_gemini.go index 48279ad..4bbe95e 100644 --- a/translate_gemini.go +++ b/translate_gemini.go @@ -661,6 +661,7 @@ type geminiStreamState struct { // (a) requires. nextToolOrdinal int startedRole bool + content contentChunker // sawFunctionCall latches true the first time any chunk in this stream // carries a functionCall part, and stays true for the rest of the // stream. Gemini's finishReason is "STOP" whether the turn ended in @@ -681,7 +682,9 @@ type geminiStreamState struct { // shares one stable, non-empty id instead of drifting between an empty // string and whatever a later chunk happens to report. func newGeminiStreamState(model string, created int64) *geminiStreamState { - return &geminiStreamState{model: model, created: created, id: geminiResponseIDOrFallback("", created)} + st := &geminiStreamState{model: model, created: created, id: geminiResponseIDOrFallback("", created)} + st.content.set(chatCompletionIDPrefix+st.id, model, created) + return st } // usage returns the prompt/completion token counts captured so far, from @@ -783,6 +786,7 @@ func (st *geminiStreamState) translate(ev sseEvent) ([][]byte, error) { isFirstChunk := !st.startedRole if isFirstChunk && resp.ResponseID != "" { st.id = resp.ResponseID + st.content.set(chatCompletionIDPrefix+st.id, st.model, st.created) } if resp.UsageMetadata != nil { @@ -832,7 +836,7 @@ func (st *geminiStreamState) translate(ev sseEvent) ([][]byte, error) { st.nextToolOrdinal++ chunks = append(chunks, st.chunk(delta, nil)) case p.Text != "": - chunks = append(chunks, st.chunk(map[string]any{"content": p.Text}, nil)) + chunks = append(chunks, st.content.build(p.Text)) } }