fix: implement improved headers for anthropic to support 8k outputs (s…

…ongquanpeng#1654)
mxdlzg · Oct 15, 2024 · 4e332e4 · 4e332e4
1 parent ae98274
commit 4e332e4
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 30 deletions.
diff --git a/common/config/config.go b/common/config/config.go
@@ -147,7 +147,6 @@ var InitialRootAccessToken = os.Getenv("INITIAL_ROOT_ACCESS_TOKEN")
 
 var GeminiVersion = env.String("GEMINI_VERSION", "v1")
 
-
 var OnlyOneLogFile = env.Bool("ONLY_ONE_LOG_FILE", false)
 
 var RelayProxy = env.String("RELAY_PROXY", "")

diff --git a/relay/adaptor/anthropic/adaptor.go b/relay/adaptor/anthropic/adaptor.go
@@ -3,12 +3,14 @@ package anthropic
 import (
 	"errors"
 	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
 	"github.com/gin-gonic/gin"
 	"github.com/songquanpeng/one-api/relay/adaptor"
 	"github.com/songquanpeng/one-api/relay/meta"
 	"github.com/songquanpeng/one-api/relay/model"
-	"io"
-	"net/http"
 )
 
 type Adaptor struct {
@@ -31,6 +33,13 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *me
 	}
 	req.Header.Set("anthropic-version", anthropicVersion)
 	req.Header.Set("anthropic-beta", "messages-2023-12-15")
+
+	// https://x.com/alexalbert__/status/1812921642143900036
+	// claude-3-5-sonnet can support 8k context
+	if strings.HasPrefix(meta.ActualModelName, "claude-3-5-sonnet") {
+		req.Header.Set("anthropic-beta", "max-tokens-3-5-sonnet-2024-07-15")
+	}
+
 	return nil
 }
 

diff --git a/relay/adaptor/vertexai/claude/model.go b/relay/adaptor/vertexai/claude/model.go
@@ -4,16 +4,16 @@ import "github.com/songquanpeng/one-api/relay/adaptor/anthropic"
 
 type Request struct {
 	// AnthropicVersion must be "vertex-2023-10-16"
-	AnthropicVersion string              `json:"anthropic_version"`
+	AnthropicVersion string `json:"anthropic_version"`
 	// Model            string              `json:"model"`
-	Messages         []anthropic.Message `json:"messages"`
-	System           string              `json:"system,omitempty"`
-	MaxTokens        int                 `json:"max_tokens,omitempty"`
-	StopSequences    []string            `json:"stop_sequences,omitempty"`
-	Stream           bool                `json:"stream,omitempty"`
-	Temperature      float64             `json:"temperature,omitempty"`
-	TopP             float64             `json:"top_p,omitempty"`
-	TopK             int                 `json:"top_k,omitempty"`
-	Tools            []anthropic.Tool    `json:"tools,omitempty"`
-	ToolChoice       any                 `json:"tool_choice,omitempty"`
+	Messages      []anthropic.Message `json:"messages"`
+	System        string              `json:"system,omitempty"`
+	MaxTokens     int                 `json:"max_tokens,omitempty"`
+	StopSequences []string            `json:"stop_sequences,omitempty"`
+	Stream        bool                `json:"stream,omitempty"`
+	Temperature   float64             `json:"temperature,omitempty"`
+	TopP          float64             `json:"top_p,omitempty"`
+	TopK          int                 `json:"top_k,omitempty"`
+	Tools         []anthropic.Tool    `json:"tools,omitempty"`
+	ToolChoice    any                 `json:"tool_choice,omitempty"`
 }
diff --git a/relay/adaptor/vertexai/registry.go b/relay/adaptor/vertexai/registry.go
@@ -32,7 +32,6 @@ func init() {
 	}
 }
 
-
 type innerAIAdapter interface {
 	ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error)
 	DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode)

diff --git a/relay/adaptor/vertexai/token.go b/relay/adaptor/vertexai/token.go
@@ -26,7 +26,6 @@ type ApplicationDefaultCredentials struct {
 	UniverseDomain          string `json:"universe_domain"`
 }
 
-
 var Cache = cache.New(50*time.Minute, 55*time.Minute)
 
 const defaultScope = "https://www.googleapis.com/auth/cloud-platform"

diff --git a/relay/channeltype/url.go b/relay/channeltype/url.go
@@ -43,7 +43,7 @@ var ChannelBaseURLs = []string{
 	"https://api.together.xyz",                  // 39
 	"https://ark.cn-beijing.volces.com",         // 40
 	"https://api.novita.ai/v3/openai",           // 41
-	"",										     // 42
+	"",                                          // 42
 }
 
 func init() {

diff --git a/relay/meta/relay_meta.go b/relay/meta/relay_meta.go
@@ -10,20 +10,22 @@ import (
 )
 
 type Meta struct {
-	Mode            int
-	ChannelType     int
-	ChannelId       int
-	TokenId         int
-	TokenName       string
-	UserId          int
-	Group           string
-	ModelMapping    map[string]string
-	BaseURL         string
-	APIKey          string
-	APIType         int
-	Config          model.ChannelConfig
-	IsStream        bool
+	Mode         int
+	ChannelType  int
+	ChannelId    int
+	TokenId      int
+	TokenName    string
+	UserId       int
+	Group        string
+	ModelMapping map[string]string
+	BaseURL      string
+	APIKey       string
+	APIType      int
+	Config       model.ChannelConfig
+	IsStream     bool
+	// OriginModelName is the model name from the raw user request
 	OriginModelName string
+	// ActualModelName is the model name after mapping
 	ActualModelName string
 	RequestURLPath  string
 	PromptTokens    int // only for DoResponse