You've already forked agentic-coding-workflow
138 lines
4.5 KiB
Go
138 lines
4.5 KiB
Go
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
|
|
}
|