From 1975058648b5914fdd9dc53434c5b59f219e2b5c Mon Sep 17 00:00:00 2001
From: Travis Cline <travis.cline@gmail.com>
Date: Wed, 24 Jul 2024 21:14:51 -0700
Subject: [PATCH] anthropic: Add support for beta header (#967)

---
 llms/anthropic/anthropicllm.go                |  1 +
 llms/anthropic/anthropicllm_option.go         | 14 +++++++++++++
 .../anthropicclient/anthropicclient.go        | 21 ++++++++++++++++---
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/llms/anthropic/anthropicllm.go b/llms/anthropic/anthropicllm.go
index 8111b7f25..2d236ebe5 100644
--- a/llms/anthropic/anthropicllm.go
+++ b/llms/anthropic/anthropicllm.go
@@ -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),
 	)
 }
 
diff --git a/llms/anthropic/anthropicllm_option.go b/llms/anthropic/anthropicllm_option.go
index ea442d45b..75ae25ded 100644
--- a/llms/anthropic/anthropicllm_option.go
+++ b/llms/anthropic/anthropicllm_option.go
@@ -8,6 +8,10 @@ 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
@@ -15,6 +19,9 @@ type options struct {
 	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)
@@ -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
+	}
+}
diff --git a/llms/anthropic/internal/anthropicclient/anthropicclient.go b/llms/anthropic/internal/anthropicclient/anthropicclient.go
index a24b7e94c..cbe9fd9bf 100644
--- a/llms/anthropic/internal/anthropicclient/anthropicclient.go
+++ b/llms/anthropic/internal/anthropicclient/anthropicclient.go
@@ -27,6 +27,8 @@ type Client struct {
 
 	httpClient Doer
 
+	anthropicBetaHeader string
+
 	// UseLegacyTextCompletionsAPI is a flag to use the legacy text completions API.
 	UseLegacyTextCompletionsAPI bool
 }
@@ -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{
@@ -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) {