feat: add vertex claude channel

relay/channel/vertex-claude/adaptor.go

@@ -0,0 +1,83 @@
+package vertex_claude
+
+import (
+	"errors"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"io"
+	"net/http"
+	"one-api/dto"
+	"one-api/relay/channel"
+	relaycommon "one-api/relay/common"
+	"strings"
+)
+
+const (
+	// LOCATION europe-west1 or us-east5
+	LOCATION = "us-east5"
+)
+
+type Adaptor struct {
+}
+
+func (a *Adaptor) Init(info *relaycommon.RelayInfo, request dto.GeneralOpenAIRequest) {
+}
+
+func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
+	parts := strings.SplitN(info.ApiKey, "|", 2)
+	if len(parts) != 2 {
+		return "", fmt.Errorf("invalid api key: %s", info.ApiKey)
+	}
+	projectId := strings.TrimSpace(parts[0])
+	model, err := getRedirectModel(info.UpstreamModelName)
+	if err != nil {
+		return "", err
+	}
+	return fmt.Sprintf("https://%s-aiplatform.googleapis.com/v1/projects/%s/locations/%s/publishers/anthropic/models/%s:streamRawPredict", LOCATION, projectId, LOCATION, model), nil
+}
+
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+	channel.SetupApiRequestHeader(info, c, req)
+	parts := strings.SplitN(info.ApiKey, "|", 2)
+	if len(parts) != 2 {
+		return fmt.Errorf("invalid api key: %s", info.ApiKey)
+	}
+	json := strings.TrimSpace(parts[1])
+	accessToken, err := getAccessToken(json)
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Authorization", "Bearer "+accessToken)
+	return nil
+}
+
+func (a *Adaptor) ConvertRequest(c *gin.Context, _ int, request *dto.GeneralOpenAIRequest) (any, error) {
+	if request == nil {
+		return nil, errors.New("request is nil")
+	}
+	return requestOpenAI2VertexClaude(*request)
+}
+
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+	return channel.DoApiRequest(a, c, info, requestBody)
+}
+
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+	if info.IsStream {
+		err, usage = vertexClaudeStreamHandler(c, resp)
+	} else {
+		err, usage = vertexClaudeHandler(c, resp)
+	}
+	return
+}
+
+func (a *Adaptor) GetModelList() (models []string) {
+	for n := range modelIdMap {
+		models = append(models, n)
+	}
+	return
+}
+
+func (a *Adaptor) GetChannelName() string {
+	return ChannelName
+}

relay/channel/vertex-claude/constants.go

@@ -0,0 +1,7 @@
+package vertex_claude
+
+var modelIdMap = map[string]string{
+	"claude-3-5-sonnet-20240620": "claude-3-5-sonnet@20240620",
+}
+
+var ChannelName = "vertex-claude"

relay/channel/vertex-claude/dto.go

@@ -0,0 +1,12 @@
+package vertex_claude
+
+import "one-api/relay/channel/claude"
+
+type VertexClaudeRequest struct {
+	// vertex-2023-10-16
+	AnthropicVersion string                 `json:"anthropic_version"`
+	System           string                 `json:"system,omitempty"`
+	Messages         []claude.ClaudeMessage `json:"messages"`
+	MaxTokens        int                    `json:"max_tokens,omitempty"`
+	Stream           bool                   `json:"stream,omitempty"`
+}

relay/channel/vertex-claude/relay-vertex-claude.go

@@ -0,0 +1,254 @@
+package vertex_claude
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/pkg/errors"
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/google"
+	"io"
+	"net/http"
+	"one-api/common"
+	"one-api/dto"
+	relaymodel "one-api/dto"
+	"one-api/relay/channel/claude"
+	"one-api/service"
+	"strings"
+	"sync"
+	"time"
+)
+
+var accessTokenMap sync.Map
+
+func getAccessToken(json string) (string, error) {
+	data, ok := accessTokenMap.Load(json)
+	if ok {
+		token := data.(oauth2.Token)
+		if time.Now().Before(token.Expiry) {
+			return token.AccessToken, nil
+		}
+	}
+	creds, err := google.CredentialsFromJSON(context.Background(), []byte(json), "https://www.googleapis.com/auth/cloud-platform")
+	if err != nil {
+		return "", err
+	}
+	token, err := creds.TokenSource.Token()
+	if err != nil {
+		return "", err
+	}
+	accessTokenMap.Store(json, *token)
+	return token.AccessToken, nil
+}
+
+func getRedirectModel(requestModel string) (string, error) {
+	if model, ok := modelIdMap[requestModel]; ok {
+		return model, nil
+	}
+	return "", errors.Errorf("model %s not found", requestModel)
+}
+
+func requestOpenAI2VertexClaude(request dto.GeneralOpenAIRequest) (*VertexClaudeRequest, error) {
+	vertexClaudeRequest := VertexClaudeRequest{
+		AnthropicVersion: "vertex-2023-10-16",
+		Stream:           request.Stream,
+	}
+	if vertexClaudeRequest.MaxTokens == 0 {
+		vertexClaudeRequest.MaxTokens = 4096
+	}
+	formatMessages := make([]dto.Message, 0)
+	var lastMessage *dto.Message
+	for i, message := range request.Messages {
+		if message.Role == "" {
+			request.Messages[i].Role = "user"
+		}
+		fmtMessage := dto.Message{
+			Role:    message.Role,
+			Content: message.Content,
+		}
+		if lastMessage != nil && lastMessage.Role == message.Role {
+			if lastMessage.IsStringContent() && message.IsStringContent() {
+				content, _ := json.Marshal(strings.Trim(fmt.Sprintf("%s %s", lastMessage.StringContent(), message.StringContent()), "\""))
+				fmtMessage.Content = content
+				// delete last message
+				formatMessages = formatMessages[:len(formatMessages)-1]
+			}
+		}
+		if fmtMessage.Content == nil {
+			content, _ := json.Marshal("...")
+			fmtMessage.Content = content
+		}
+		formatMessages = append(formatMessages, fmtMessage)
+		lastMessage = &request.Messages[i]
+	}
+	claudeMessages := make([]claude.ClaudeMessage, 0)
+	for _, message := range formatMessages {
+		if message.Role == "system" {
+			if message.IsStringContent() {
+				vertexClaudeRequest.System = message.StringContent()
+			} else {
+				contents := message.ParseContent()
+				content := ""
+				for _, ctx := range contents {
+					if ctx.Type == "text" {
+						content += ctx.Text
+					}
+				}
+				vertexClaudeRequest.System = content
+			}
+		} else {
+			claudeMessage := claude.ClaudeMessage{
+				Role: message.Role,
+			}
+			if message.IsStringContent() {
+				claudeMessage.Content = message.StringContent()
+			} else {
+				claudeMediaMessages := make([]claude.ClaudeMediaMessage, 0)
+				for _, mediaMessage := range message.ParseContent() {
+					claudeMediaMessage := claude.ClaudeMediaMessage{
+						Type: mediaMessage.Type,
+					}
+					if mediaMessage.Type == "text" {
+						claudeMediaMessage.Text = mediaMessage.Text
+					} else {
+						imageUrl := mediaMessage.ImageUrl.(dto.MessageImageUrl)
+						claudeMediaMessage.Type = "image"
+						claudeMediaMessage.Source = &claude.ClaudeMessageSource{
+							Type: "base64",
+						}
+						// 判断是否是url
+						if strings.HasPrefix(imageUrl.Url, "http") {
+							// 是url，获取图片的类型和base64编码的数据
+							mimeType, data, _ := common.GetImageFromUrl(imageUrl.Url)
+							claudeMediaMessage.Source.MediaType = mimeType
+							claudeMediaMessage.Source.Data = data
+						} else {
+							_, format, base64String, err := common.DecodeBase64ImageData(imageUrl.Url)
+							if err != nil {
+								return nil, err
+							}
+							claudeMediaMessage.Source.MediaType = "image/" + format
+							claudeMediaMessage.Source.Data = base64String
+						}
+					}
+					claudeMediaMessages = append(claudeMediaMessages, claudeMediaMessage)
+				}
+				claudeMessage.Content = claudeMediaMessages
+			}
+			claudeMessages = append(claudeMessages, claudeMessage)
+		}
+	}
+	vertexClaudeRequest.Messages = claudeMessages
+	return &vertexClaudeRequest, nil
+}
+
+func vertexClaudeHandler(c *gin.Context, resp *http.Response) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) {
+	var claudeResponse claude.ClaudeResponse
+	responseBody, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = resp.Body.Close()
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
+	err = json.Unmarshal(responseBody, &claudeResponse)
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
+	}
+	openaiResp := claude.ResponseClaude2OpenAI(claude.RequestModeMessage, &claudeResponse)
+	usage := relaymodel.Usage{
+		PromptTokens:     claudeResponse.Usage.InputTokens,
+		CompletionTokens: claudeResponse.Usage.OutputTokens,
+		TotalTokens:      claudeResponse.Usage.InputTokens + claudeResponse.Usage.OutputTokens,
+	}
+	openaiResp.Usage = usage
+	c.JSON(http.StatusOK, openaiResp)
+	return nil, &usage
+}
+
+func vertexClaudeStreamHandler(c *gin.Context, resp *http.Response) (*relaymodel.OpenAIErrorWithStatusCode, *relaymodel.Usage) {
+	scanner := bufio.NewScanner(resp.Body)
+	scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
+		if atEOF && len(data) == 0 {
+			return 0, nil, nil
+		}
+		if i := strings.Index(string(data), "\n"); i >= 0 {
+			return i + 1, data[0:i], nil
+		}
+		if atEOF {
+			return len(data), data, nil
+		}
+		return 0, nil, nil
+	})
+	dataChan := make(chan string)
+	stopChan := make(chan bool)
+	go func() {
+		for scanner.Scan() {
+			data := scanner.Text()
+			if len(data) < 5 { // ignore blank line or wrong format
+				continue
+			}
+			if data[:5] != "data:" {
+				continue
+			}
+			data = data[5:]
+			dataChan <- data
+		}
+		stopChan <- true
+	}()
+	var id string
+	var model string
+	createdTime := common.GetTimestamp()
+	var usage relaymodel.Usage
+	service.SetEventStreamHeaders(c)
+	c.Stream(func(w io.Writer) bool {
+		select {
+		case data := <-dataChan:
+			claudeResp := new(claude.ClaudeResponse)
+			err := json.Unmarshal([]byte(data), &claudeResp)
+			if err != nil {
+				common.SysError("error unmarshalling stream response: " + err.Error())
+				return true
+			}
+			response, claudeUsage := claude.StreamResponseClaude2OpenAI(claude.RequestModeMessage, claudeResp)
+
+			if claudeUsage != nil {
+				usage.PromptTokens += claudeUsage.InputTokens
+				usage.CompletionTokens += claudeUsage.OutputTokens
+			}
+
+			if response == nil {
+				return true
+			}
+
+			if response.Id != "" {
+				id = response.Id
+			}
+			if response.Model != "" {
+				model = response.Model
+			}
+			response.Created = createdTime
+			response.Id = id
+			response.Model = model
+
+			jsonStr, err := json.Marshal(response)
+			if err != nil {
+				common.SysError("error marshalling stream response: " + err.Error())
+				return true
+			}
+			c.Render(-1, common.CustomEvent{Data: "data: " + string(jsonStr)})
+			return true
+		case <-stopChan:
+			c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
+			return false
+		}
+	})
+	err := resp.Body.Close()
+	if err != nil {
+		return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
+	}
+	return nil, &usage
+}

relay/constant/api_type.go

@@ -21,6 +21,7 @@ const (
 	APITypeAws
 	APITypeCohere
 	APITypeScholarAI
+	APITypeVertexClaude
 	APITypeDummy // this one is only for count, do not add any channel after this
 )
 
@@ -59,6 +60,8 @@ func ChannelType2APIType(channelType int) (int, bool) {
 		apiType = APITypeCohere
 	case common.ChannelTypeScholarAI:
 		apiType = APITypeScholarAI
+	case common.ChannelTypeVertexClaude:
+		apiType = APITypeVertexClaude
 	}
 	if apiType == -1 {
 		return APITypeOpenAI, false

relay/relay_adaptor.go

@@ -14,6 +14,7 @@ import (
 	"one-api/relay/channel/perplexity"
 	"one-api/relay/channel/scholarai"
 	"one-api/relay/channel/tencent"
+	vertex_claude "one-api/relay/channel/vertex-claude"
 	"one-api/relay/channel/xunfei"
 	"one-api/relay/channel/zhipu"
 	"one-api/relay/channel/zhipu_4v"
@@ -54,6 +55,8 @@ func GetAdaptor(apiType int) channel.Adaptor {
 		return &cohere.Adaptor{}
 	case constant.APITypeScholarAI:
 		return &scholarai.Adaptor{}
+	case constant.APITypeVertexClaude:
+		return &vertex_claude.Adaptor{}
 	}
 	return nil
 }