feat(infra/llm): Anthropic client with stream + count_tokens via official SDK

2026-05-04 09:29:35 +08:00
parent 9c8e4175c4
commit a1ba157d29
4 changed files with 250 additions and 0 deletions
@@ -0,0 +1,80 @@
+package llm
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+const anthropicStreamFixture = `event: message_start
+data: {"type":"message_start","message":{"id":"m1","content":[],"model":"claude-opus","role":"assistant","stop_reason":null,"stop_sequence":null,"type":"message","usage":{"input_tokens":10,"output_tokens":0}}}
+
+event: content_block_start
+data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
+
+event: content_block_delta
+data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" world"}}
+
+event: content_block_stop
+data: {"type":"content_block_stop","index":0}
+
+event: message_delta
+data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}
+
+event: message_stop
+data: {"type":"message_stop"}
+
+`
+
+func TestAnthropic_Stream_TextDeltaAndDone(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		require.Equal(t, "/v1/messages", r.URL.Path)
+		w.Header().Set("Content-Type", "text/event-stream")
+		w.WriteHeader(200)
+		_, _ = w.Write([]byte(anthropicStreamFixture))
+	}))
+	defer srv.Close()
+
+	c, err := NewAnthropic("test-key", srv.URL)
+	require.NoError(t, err)
+
+	req := Request{
+		Messages:        []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hi"}}}},
+		MaxOutputTokens: 100,
+	}
+	ch, err := c.Stream(context.Background(), "claude-opus-4-7", req)
+	require.NoError(t, err)
+
+	var deltas []string
+	var sawUsage, sawDone bool
+	var promptTokens, completionTokens int
+	for ev := range ch {
+		switch ev.Type {
+		case EventTextDelta:
+			deltas = append(deltas, ev.Text)
+		case EventUsage:
+			sawUsage = true
+			promptTokens = ev.Usage.PromptTokens
+			completionTokens = ev.Usage.CompletionTokens
+		case EventDone:
+			sawDone = true
+		}
+	}
+	require.Equal(t, "Hello world", strings.Join(deltas, ""))
+	require.True(t, sawUsage)
+	require.Equal(t, 10, promptTokens)
+	require.Equal(t, 2, completionTokens)
+	require.True(t, sawDone)
+}
+
+func TestAnthropic_Provider(t *testing.T) {
+	c, _ := NewAnthropic("k", "http://localhost")
+	require.Equal(t, "anthropic", c.Provider())
+}