mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-18 05:47:38 +00:00
perf(anthropic): build text-delta stream chunks without envelope maps
This commit is contained in:
+40
-2
@@ -3,6 +3,7 @@ package traefikllmgateway
|
|||||||
import (
|
import (
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -692,7 +693,10 @@ type anthropicStreamState struct {
|
|||||||
// index also counts any preceding non-tool_use blocks (a tool_use
|
// 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
|
// block right after a text block is content-block index 1, but it is
|
||||||
// still the first — index 0 — tool call).
|
// still the first — index 0 — tool call).
|
||||||
toolOrdinals map[int]int
|
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
|
||||||
id string
|
id string
|
||||||
model string
|
model string
|
||||||
u usage
|
u usage
|
||||||
@@ -741,6 +745,39 @@ func (st *anthropicStreamState) chunk(delta map[string]any, finishReason *string
|
|||||||
return b
|
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).
|
||||||
|
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 {
|
||||||
|
tb, err := json.Marshal(text)
|
||||||
|
if err != nil {
|
||||||
|
panic(fmt.Sprintf("llmgateway: anthropic stream chunk failed to marshal: %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 = append(out, contentChunkHead...)
|
||||||
|
out = append(out, tb...)
|
||||||
|
return append(out, st.contentTail...)
|
||||||
|
}
|
||||||
|
|
||||||
// usageChunk builds the final usage-only "chat.completion.chunk" payload —
|
// usageChunk builds the final usage-only "chat.completion.chunk" payload —
|
||||||
// empty "choices" plus a populated "usage" — mirroring the shape
|
// empty "choices" plus a populated "usage" — mirroring the shape
|
||||||
// provider_openai.go's forwardStream relays from a native OpenAI upstream
|
// provider_openai.go's forwardStream relays from a native OpenAI upstream
|
||||||
@@ -854,6 +891,7 @@ func (st *anthropicStreamState) translate(ev sseEvent) ([][]byte, error) {
|
|||||||
return nil, fmt.Errorf("%w: decode anthropic message_start event: %w", errUpstream, err)
|
return nil, fmt.Errorf("%w: decode anthropic message_start event: %w", errUpstream, err)
|
||||||
}
|
}
|
||||||
st.id = p.Message.ID
|
st.id = p.Message.ID
|
||||||
|
st.contentTail = nil
|
||||||
st.u.prompt = p.Message.Usage.totalInputTokens()
|
st.u.prompt = p.Message.Usage.totalInputTokens()
|
||||||
return [][]byte{st.chunk(map[string]any{"role": "assistant"}, nil)}, nil
|
return [][]byte{st.chunk(map[string]any{"role": "assistant"}, nil)}, nil
|
||||||
|
|
||||||
@@ -886,7 +924,7 @@ func (st *anthropicStreamState) translate(ev sseEvent) ([][]byte, error) {
|
|||||||
}
|
}
|
||||||
switch p.Delta.Type {
|
switch p.Delta.Type {
|
||||||
case "text_delta":
|
case "text_delta":
|
||||||
return [][]byte{st.chunk(map[string]any{"content": p.Delta.Text}, nil)}, nil
|
return [][]byte{st.contentChunk(p.Delta.Text)}, nil
|
||||||
case "input_json_delta":
|
case "input_json_delta":
|
||||||
ordinal, ok := st.toolOrdinals[p.Index]
|
ordinal, ok := st.toolOrdinals[p.Index]
|
||||||
if !ok {
|
if !ok {
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package traefikllmgateway
|
package traefikllmgateway
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
@@ -1127,3 +1128,18 @@ func TestOpenAIRequestFromAnthropic_NoTools_OmitsToolsAndToolChoiceFields(t *tes
|
|||||||
t.Errorf(`out["tool_choice"] = %v, want the key absent when no tool_choice was requested`, out["tool_choice"])
|
t.Errorf(`out["tool_choice"] = %v, want the key absent when no tool_choice was requested`, out["tool_choice"])
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestAnthropicContentChunkMatchesChunk(t *testing.T) {
|
||||||
|
texts := []string{"", "Hello", "a \"quoted\" <b>&</b>", "line\nbreak\ttab", "unicode é 日本 \u2028\u2029", "bad\xffutf8"}
|
||||||
|
for _, id := range []string{"", "msg_01", "we\"ird"} {
|
||||||
|
st := newAnthropicStreamState("gw-model", 1734000000)
|
||||||
|
st.id = id
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user