You've already forked agentic-coding-workflow
feat(infra/llm): Anthropic client with stream + count_tokens via official SDK
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
Reference in New Issue
Block a user