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,137 @@
+package llm
+
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+
+	"github.com/anthropics/anthropic-sdk-go"
+	"github.com/anthropics/anthropic-sdk-go/option"
+)
+
+// Compile-time assertion: AnthropicClient must satisfy LLMClient.
+var _ LLMClient = (*AnthropicClient)(nil)
+
+// AnthropicClient 通过官方 SDK 调用 Anthropic Messages API。
+// baseURL 可改写以走 LiteLLM/OpenRouter 等代理。
+type AnthropicClient struct {
+	sdk anthropic.Client
+}
+
+// NewAnthropic 构造客户端。baseURL 可空（用 SDK 默认）。
+func NewAnthropic(apiKey, baseURL string) (*AnthropicClient, error) {
+	opts := []option.RequestOption{option.WithAPIKey(apiKey)}
+	if baseURL != "" {
+		opts = append(opts, option.WithBaseURL(baseURL))
+	}
+	c := anthropic.NewClient(opts...)
+	return &AnthropicClient{sdk: c}, nil
+}
+
+func (c *AnthropicClient) Provider() string { return "anthropic" }
+
+// toSDKMessages 把中立 Message 转 SDK 类型。
+// image/document 的 Data 字段为原始字节，此处进行 base64 编码。
+func toSDKMessages(msgs []Message) []anthropic.MessageParam {
+	out := make([]anthropic.MessageParam, 0, len(msgs))
+	for _, m := range msgs {
+		blocks := make([]anthropic.ContentBlockParamUnion, 0, len(m.Blocks))
+		for _, b := range m.Blocks {
+			switch b.Type {
+			case "text":
+				blocks = append(blocks, anthropic.NewTextBlock(b.Text))
+			case "image":
+				// NewImageBlockBase64 期望 base64 字符串，Data 字段存原始字节
+				encoded := base64.StdEncoding.EncodeToString(b.Data)
+				blocks = append(blocks, anthropic.NewImageBlockBase64(b.MimeType, encoded))
+			case "document":
+				// NewDocumentBlock 泛型支持 Base64PDFSourceParam，Data 字段存原始字节
+				encoded := base64.StdEncoding.EncodeToString(b.Data)
+				blocks = append(blocks, anthropic.NewDocumentBlock(anthropic.Base64PDFSourceParam{
+					Data: encoded,
+				}))
+			}
+		}
+		role := anthropic.MessageParamRoleUser
+		if m.Role == RoleAssistant {
+			role = anthropic.MessageParamRoleAssistant
+		}
+		out = append(out, anthropic.MessageParam{Role: role, Content: blocks})
+	}
+	return out
+}
+
+func (c *AnthropicClient) Stream(ctx context.Context, modelID string, req Request) (<-chan StreamEvent, error) {
+	params := anthropic.MessageNewParams{
+		Model:     anthropic.Model(modelID),
+		MaxTokens: int64(req.MaxOutputTokens),
+		Messages:  toSDKMessages(req.Messages),
+	}
+	if req.SystemPrompt != "" {
+		params.System = []anthropic.TextBlockParam{{Text: req.SystemPrompt}}
+	}
+	if req.Reasoning {
+		params.Thinking = anthropic.ThinkingConfigParamOfEnabled(int64(req.MaxOutputTokens / 2))
+	}
+
+	stream := c.sdk.Messages.NewStreaming(ctx, params)
+	out := make(chan StreamEvent, 16)
+
+	go func() {
+		defer close(out)
+		var inputTokens int
+		for stream.Next() {
+			event := stream.Current()
+			// SDK v1.38: AsAny() returns unexported anyMessageStreamEvent which cannot
+			// be used in a type switch. Use event.Type string + typed As*() methods.
+			switch event.Type {
+			case "message_start":
+				ev := event.AsMessageStart()
+				inputTokens = int(ev.Message.Usage.InputTokens)
+			case "content_block_delta":
+				ev := event.AsContentBlockDelta()
+				switch ev.Delta.Type {
+				case "text_delta":
+					if text := ev.Delta.AsTextDelta().Text; text != "" {
+						out <- StreamEvent{Type: EventTextDelta, Text: text}
+					}
+				case "thinking_delta":
+					if thinking := ev.Delta.AsThinkingDelta().Thinking; thinking != "" {
+						out <- StreamEvent{Type: EventThinkingDelta, Text: thinking}
+					}
+				}
+			case "message_delta":
+				ev := event.AsMessageDelta()
+				out <- StreamEvent{Type: EventUsage, Usage: &Usage{
+					PromptTokens:     inputTokens,
+					CompletionTokens: int(ev.Usage.OutputTokens),
+				}}
+			case "message_stop":
+				out <- StreamEvent{Type: EventDone, Reason: "end_turn"}
+			}
+		}
+		if err := stream.Err(); err != nil {
+			out <- StreamEvent{Type: EventError, Err: fmt.Errorf("anthropic stream: %w", err)}
+		}
+	}()
+	return out, nil
+}
+
+func (c *AnthropicClient) CountTokens(ctx context.Context, modelID string, req Request) (int, error) {
+	params := anthropic.MessageCountTokensParams{
+		Model:    anthropic.Model(modelID),
+		Messages: toSDKMessages(req.Messages),
+	}
+	if req.SystemPrompt != "" {
+		// MessageCountTokensParams.System 是 MessageCountTokensParamsSystemUnion，
+		// 不同于 MessageNewParams.System ([]TextBlockParam)
+		params.System = anthropic.MessageCountTokensParamsSystemUnion{
+			OfTextBlockArray: []anthropic.TextBlockParam{{Text: req.SystemPrompt}},
+		}
+	}
+	res, err := c.sdk.Messages.CountTokens(ctx, params)
+	if err != nil {
+		return 0, fmt.Errorf("anthropic count_tokens: %w", err)
+	}
+	return int(res.InputTokens), nil
+}