Skip to content

Commit

Permalink
anthropic: Add support for beta header (tmc#967)
Browse files Browse the repository at this point in the history
  • Loading branch information
tmc authored Jul 25, 2024
1 parent 4fc43db commit 1975058
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 3 deletions.
1 change: 1 addition & 0 deletions llms/anthropic/anthropicllm.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ func newClient(opts ...Option) (*anthropicclient.Client, error) {
return anthropicclient.New(options.token, options.model, options.baseURL,
anthropicclient.WithHTTPClient(options.httpClient),
anthropicclient.WithLegacyTextCompletionsAPI(options.useLegacyTextCompletionsAPI),
anthropicclient.WithAnthropicBetaHeader(options.anthropicBetaHeader),
)
}

Expand Down
14 changes: 14 additions & 0 deletions llms/anthropic/anthropicllm_option.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,13 +8,20 @@ const (
tokenEnvVarName = "ANTHROPIC_API_KEY" //nolint:gosec
)

// MaxTokensAnthropicSonnet35 is the header value for specifying the maximum number of tokens
// when using the Anthropic Sonnet 3.5 model.
const MaxTokensAnthropicSonnet35 = "max-tokens-3-5-sonnet-2024-07-15" //nolint:gosec // This is not a sensitive value.

type options struct {
token string
model string
baseURL string
httpClient anthropicclient.Doer

useLegacyTextCompletionsAPI bool

// If supplied, the 'anthropic-beta' header will be added to the request with the given value.
anthropicBetaHeader string
}

type Option func(*options)
Expand Down Expand Up @@ -56,3 +63,10 @@ func WithLegacyTextCompletionsAPI() Option {
opts.useLegacyTextCompletionsAPI = true
}
}

// WithAnthropicBetaHeader adds the Anthropic Beta header to support extended options.
func WithAnthropicBetaHeader(value string) Option {
return func(opts *options) {
opts.anthropicBetaHeader = value
}
}
21 changes: 18 additions & 3 deletions llms/anthropic/internal/anthropicclient/anthropicclient.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@ type Client struct {

httpClient Doer

anthropicBetaHeader string

// UseLegacyTextCompletionsAPI is a flag to use the legacy text completions API.
UseLegacyTextCompletionsAPI bool
}
Expand Down Expand Up @@ -56,6 +58,14 @@ func WithLegacyTextCompletionsAPI(val bool) Option {
}
}

// WithAnthropicBetaHeader sets the anthropic-beta header.
func WithAnthropicBetaHeader(val string) Option {
return func(opts *Client) error {
opts.anthropicBetaHeader = val
return nil
}
}

// New returns a new Anthropic client.
func New(token string, model string, baseURL string, opts ...Option) (*Client, error) {
c := &Client{
Expand Down Expand Up @@ -149,9 +159,14 @@ func (c *Client) CreateMessage(ctx context.Context, r *MessageRequest) (*Message

func (c *Client) setHeaders(req *http.Request) {
req.Header.Set("Content-Type", "application/json")
req.Header.Set("x-api-key", c.token)
// TODO: expose version as a option/parameter
req.Header.Set("anthropic-version", "2023-06-01")
req.Header.Set("x-api-key", c.token) //nolint:canonicalheader

// This is necessary as per https://docs.anthropic.com/en/api/versioning
// If this changes frequently enough we should expose it as an option..
req.Header.Set("anthropic-version", "2023-06-01") // nolint:canonicalheader
if c.anthropicBetaHeader != "" {
req.Header.Set("anthropic-beta", c.anthropicBetaHeader) // nolint:canonicalheader
}
}

func (c *Client) do(ctx context.Context, path string, payloadBytes []byte) (*http.Response, error) {
Expand Down

0 comments on commit 1975058

Please sign in to comment.