feat: support aws claude

relay/channel/aws/adaptor.go

@@ -0,0 +1,79 @@
+package aws
+
+import (
+	"errors"
+	"github.com/gin-gonic/gin"
+	"io"
+	"net/http"
+	"one-api/dto"
+	"one-api/relay/channel/claude"
+	relaycommon "one-api/relay/common"
+	"strings"
+)
+
+const (
+	RequestModeCompletion = 1
+	RequestModeMessage    = 2
+)
+
+type Adaptor struct {
+	RequestMode int
+}
+
+func (a *Adaptor) Init(info *relaycommon.RelayInfo, request dto.GeneralOpenAIRequest) {
+	if strings.HasPrefix(info.UpstreamModelName, "claude-3") {
+		a.RequestMode = RequestModeMessage
+	} else {
+		a.RequestMode = RequestModeCompletion
+	}
+}
+
+func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
+	return "", nil
+}
+
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
+	return nil
+}
+
+func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *dto.GeneralOpenAIRequest) (any, error) {
+	if request == nil {
+		return nil, errors.New("request is nil")
+	}
+
+	var claudeReq *claude.ClaudeRequest
+	var err error
+	if a.RequestMode == RequestModeCompletion {
+		claudeReq = claude.RequestOpenAI2ClaudeComplete(*request)
+	} else {
+		claudeReq, err = claude.RequestOpenAI2ClaudeMessage(*request)
+	}
+	c.Set("request_model", request.Model)
+	c.Set("converted_request", claudeReq)
+	return claudeReq, err
+}
+
+func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
+	return nil, nil
+}
+
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
+	if info.IsStream {
+		err, usage = awsStreamHandler(c, info, a.RequestMode)
+	} else {
+		err, usage = awsHandler(c, info, a.RequestMode)
+	}
+	return
+}
+
+func (a *Adaptor) GetModelList() (models []string) {
+	for n := range awsModelIDMap {
+		models = append(models, n)
+	}
+
+	return
+}
+
+func (a *Adaptor) GetChannelName() string {
+	return ChannelName
+}

relay/channel/aws/constants.go

@@ -0,0 +1,12 @@
+package aws
+
+var awsModelIDMap = map[string]string{
+	"claude-instant-1.2":       "anthropic.claude-instant-v1",
+	"claude-2.0":               "anthropic.claude-v2",
+	"claude-2.1":               "anthropic.claude-v2:1",
+	"claude-3-sonnet-20240229": "anthropic.claude-3-sonnet-20240229-v1:0",
+	"claude-3-opus-20240229":   "anthropic.claude-3-opus-20240229-v1:0",
+	"claude-3-haiku-20240307":  "anthropic.claude-3-haiku-20240307-v1:0",
+}
+
+var ChannelName = "aws"

relay/channel/aws/dto.go

@@ -0,0 +1,14 @@
+package aws
+
+import "one-api/relay/channel/claude"
+
+type AwsClaudeRequest struct {
+	// AnthropicVersion should be "bedrock-2023-05-31"
+	AnthropicVersion string                 `json:"anthropic_version"`
+	Messages         []claude.ClaudeMessage `json:"messages"`
+	MaxTokens        int                    `json:"max_tokens,omitempty"`
+	Temperature      float64                `json:"temperature,omitempty"`
+	TopP             float64                `json:"top_p,omitempty"`
+	TopK             int                    `json:"top_k,omitempty"`
+	StopSequences    []string               `json:"stop_sequences,omitempty"`
+}

relay/channel/aws/relay-aws.go

@@ -0,0 +1,211 @@
+package aws
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/jinzhu/copier"
+	"github.com/pkg/errors"
+	"io"
+	"net/http"
+	"one-api/common"
+	relaymodel "one-api/dto"
+	"one-api/relay/channel/claude"
+	relaycommon "one-api/relay/common"
+	"strings"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/credentials"
+	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
+	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
+)
+
+func newAwsClient(c *gin.Context, info *relaycommon.RelayInfo) (*bedrockruntime.Client, error) {
+	awsSecret := strings.Split(info.ApiKey, "|")
+	if len(awsSecret) != 3 {
+		return nil, errors.New("invalid aws secret key")
+	}
+	ak := awsSecret[0]
+	sk := awsSecret[1]
+	region := awsSecret[2]
+	client := bedrockruntime.New(bedrockruntime.Options{
+		Region:      region,
+		Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(ak, sk, "")),
+	})
+
+	return client, nil
+}
+
+func wrapErr(err error) *relaymodel.OpenAIErrorWithStatusCode {
+	return &relaymodel.OpenAIErrorWithStatusCode{
+		StatusCode: http.StatusInternalServerError,
+		Error: relaymodel.OpenAIError{
+			Message: fmt.Sprintf("%s", err.Error()),
+		},
+	}
+}
+
+func awsModelID(requestModel string) (string, error) {
+	if awsModelID, ok := awsModelIDMap[requestModel]; ok {
+		return awsModelID, nil
+	}
+
+	return "", errors.Errorf("model %s not found", requestModel)
+}
+
+func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, requestMode int) (*relaymodel.OpenAIErrorWithStatusCode, *relaymodel.Usage) {
+	awsCli, err := newAwsClient(c, info)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "newAwsClient")), nil
+	}
+
+	awsModelId, err := awsModelID(c.GetString("request_model"))
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "awsModelID")), nil
+	}
+
+	awsReq := &bedrockruntime.InvokeModelInput{
+		ModelId:     aws.String(awsModelId),
+		Accept:      aws.String("application/json"),
+		ContentType: aws.String("application/json"),
+	}
+
+	claudeReq_, ok := c.Get("converted_request")
+	if !ok {
+		return wrapErr(errors.New("request not found")), nil
+	}
+	claudeReq := claudeReq_.(*claude.ClaudeRequest)
+	awsClaudeReq := &AwsClaudeRequest{
+		AnthropicVersion: "bedrock-2023-05-31",
+	}
+	if err = copier.Copy(awsClaudeReq, claudeReq); err != nil {
+		return wrapErr(errors.Wrap(err, "copy request")), nil
+	}
+
+	awsReq.Body, err = json.Marshal(awsClaudeReq)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "marshal request")), nil
+	}
+
+	awsResp, err := awsCli.InvokeModel(c.Request.Context(), awsReq)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "InvokeModel")), nil
+	}
+
+	claudeResponse := new(claude.ClaudeResponse)
+	err = json.Unmarshal(awsResp.Body, claudeResponse)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "unmarshal response")), nil
+	}
+
+	openaiResp := claude.ResponseClaude2OpenAI(requestMode, 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 awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, requestMode int) (*relaymodel.OpenAIErrorWithStatusCode, *relaymodel.Usage) {
+	awsCli, err := newAwsClient(c, info)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "newAwsClient")), nil
+	}
+
+	awsModelId, err := awsModelID(c.GetString("request_model"))
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "awsModelID")), nil
+	}
+
+	awsReq := &bedrockruntime.InvokeModelWithResponseStreamInput{
+		ModelId:     aws.String(awsModelId),
+		Accept:      aws.String("application/json"),
+		ContentType: aws.String("application/json"),
+	}
+
+	claudeReq_, ok := c.Get("converted_request")
+	if !ok {
+		return wrapErr(errors.New("request not found")), nil
+	}
+	claudeReq := claudeReq_.(*claude.ClaudeRequest)
+
+	awsClaudeReq := &AwsClaudeRequest{
+		AnthropicVersion: "bedrock-2023-05-31",
+	}
+	if err = copier.Copy(awsClaudeReq, claudeReq); err != nil {
+		return wrapErr(errors.Wrap(err, "copy request")), nil
+	}
+	awsReq.Body, err = json.Marshal(awsClaudeReq)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "marshal request")), nil
+	}
+
+	awsResp, err := awsCli.InvokeModelWithResponseStream(c.Request.Context(), awsReq)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "InvokeModelWithResponseStream")), nil
+	}
+	stream := awsResp.GetStream()
+	defer stream.Close()
+
+	c.Writer.Header().Set("Content-Type", "text/event-stream")
+	var usage relaymodel.Usage
+	var id string
+	var model string
+	c.Stream(func(w io.Writer) bool {
+		event, ok := <-stream.Events()
+		if !ok {
+			c.Render(-1, common.CustomEvent{Data: "data: [DONE]"})
+			return false
+		}
+
+		switch v := event.(type) {
+		case *types.ResponseStreamMemberChunk:
+			claudeResp := new(claude.ClaudeResponse)
+			err := json.NewDecoder(bytes.NewReader(v.Value.Bytes)).Decode(claudeResp)
+			if err != nil {
+				common.SysError("error unmarshalling stream response: " + err.Error())
+				return false
+			}
+
+			response, claudeUsage := claude.StreamResponseClaude2OpenAI(requestMode, 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.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 *types.UnknownUnionMember:
+			fmt.Println("unknown tag:", v.Tag)
+			return false
+		default:
+			fmt.Println("union is nil or unknown type")
+			return false
+		}
+	})
+
+	return nil, &usage
+}

relay/channel/claude/adaptor.go

@@ -53,9 +53,9 @@ func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *dto.Gen
 		return nil, errors.New("request is nil")
 	}
 	if a.RequestMode == RequestModeCompletion {
-		return requestOpenAI2ClaudeComplete(*request), nil
+		return RequestOpenAI2ClaudeComplete(*request), nil
 	} else {
-		return requestOpenAI2ClaudeMessage(*request)
+		return RequestOpenAI2ClaudeMessage(*request)
 	}
 }
 

relay/channel/claude/relay-claude.go

@@ -26,7 +26,7 @@ func stopReasonClaude2OpenAI(reason string) string {
 	}
 }
 
-func requestOpenAI2ClaudeComplete(textRequest dto.GeneralOpenAIRequest) *ClaudeRequest {
+func RequestOpenAI2ClaudeComplete(textRequest dto.GeneralOpenAIRequest) *ClaudeRequest {
 	claudeRequest := ClaudeRequest{
 		Model:             textRequest.Model,
 		Prompt:            "",
@@ -57,7 +57,7 @@ func requestOpenAI2ClaudeComplete(textRequest dto.GeneralOpenAIRequest) *ClaudeR
 	return &claudeRequest
 }
 
-func requestOpenAI2ClaudeMessage(textRequest dto.GeneralOpenAIRequest) (*ClaudeRequest, error) {
+func RequestOpenAI2ClaudeMessage(textRequest dto.GeneralOpenAIRequest) (*ClaudeRequest, error) {
 	claudeRequest := ClaudeRequest{
 		Model:         textRequest.Model,
 		MaxTokens:     textRequest.MaxTokens,
@@ -122,7 +122,7 @@ func requestOpenAI2ClaudeMessage(textRequest dto.GeneralOpenAIRequest) (*ClaudeR
 	return &claudeRequest, nil
 }
 
-func streamResponseClaude2OpenAI(reqMode int, claudeResponse *ClaudeResponse) (*dto.ChatCompletionsStreamResponse, *ClaudeUsage) {
+func StreamResponseClaude2OpenAI(reqMode int, claudeResponse *ClaudeResponse) (*dto.ChatCompletionsStreamResponse, *ClaudeUsage) {
 	var response dto.ChatCompletionsStreamResponse
 	var claudeUsage *ClaudeUsage
 	response.Object = "chat.completion.chunk"
@@ -149,6 +149,8 @@ func streamResponseClaude2OpenAI(reqMode int, claudeResponse *ClaudeResponse) (*
 				choice.FinishReason = &finishReason
 			}
 			claudeUsage = &claudeResponse.Usage
+		} else if claudeResponse.Type == "message_stop" {
+			return nil, nil
 		}
 	}
 	if claudeUsage == nil {
@@ -158,7 +160,7 @@ func streamResponseClaude2OpenAI(reqMode int, claudeResponse *ClaudeResponse) (*
 	return &response, claudeUsage
 }
 
-func responseClaude2OpenAI(reqMode int, claudeResponse *ClaudeResponse) *dto.OpenAITextResponse {
+func ResponseClaude2OpenAI(reqMode int, claudeResponse *ClaudeResponse) *dto.OpenAITextResponse {
 	choices := make([]dto.OpenAITextResponseChoice, 0)
 	fullTextResponse := dto.OpenAITextResponse{
 		Id:      fmt.Sprintf("chatcmpl-%s", common.GetUUID()),
@@ -242,7 +244,10 @@ func claudeStreamHandler(requestMode int, modelName string, promptTokens int, c
 				return true
 			}
 
-			response, claudeUsage := streamResponseClaude2OpenAI(requestMode, &claudeResponse)
+			response, claudeUsage := StreamResponseClaude2OpenAI(requestMode, &claudeResponse)
+			if response == nil {
+				return true
+			}
 			if requestMode == RequestModeCompletion {
 				responseText += claudeResponse.Completion
 				responseId = response.Id
@@ -317,7 +322,7 @@ func claudeHandler(requestMode int, c *gin.Context, resp *http.Response, promptT
 			StatusCode: resp.StatusCode,
 		}, nil
 	}
-	fullTextResponse := responseClaude2OpenAI(requestMode, &claudeResponse)
+	fullTextResponse := ResponseClaude2OpenAI(requestMode, &claudeResponse)
 	completionTokens, err, _ := service.CountTokenText(claudeResponse.Completion, model, false)
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "count_token_text_failed", http.StatusInternalServerError), nil

relay/constant/api_type.go

@@ -18,6 +18,7 @@ const (
 	APITypeZhipu_v4
 	APITypeOllama
 	APITypePerplexity
+	APITypeAws
 
 	APITypeDummy // this one is only for count, do not add any channel after this
 )
@@ -49,6 +50,8 @@ func ChannelType2APIType(channelType int) int {
 		apiType = APITypeOllama
 	case common.ChannelTypePerplexity:
 		apiType = APITypePerplexity
+	case common.ChannelTypeAws:
+		apiType = APITypeAws
 	}
 	return apiType
 }

relay/relay-text.go

@@ -159,14 +159,16 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode {
 	if err != nil {
 		return service.OpenAIErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
 	}
-	relayInfo.IsStream = relayInfo.IsStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
 
-	if resp.StatusCode != http.StatusOK {
-		returnPreConsumedQuota(c, relayInfo.TokenId, userQuota, preConsumedQuota)
-		openaiErr := service.RelayErrorHandler(resp)
-		// reset status code 重置状态码
-		service.ResetStatusCode(openaiErr, statusCodeMappingStr)
-		return openaiErr
+	if resp != nil {
+		relayInfo.IsStream = relayInfo.IsStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
+		if resp.StatusCode != http.StatusOK {
+			returnPreConsumedQuota(c, relayInfo.TokenId, userQuota, preConsumedQuota)
+			openaiErr := service.RelayErrorHandler(resp)
+			// reset status code 重置状态码
+			service.ResetStatusCode(openaiErr, statusCodeMappingStr)
+			return openaiErr
+		}
 	}
 
 	usage, openaiErr := adaptor.DoResponse(c, resp, relayInfo)

relay/relay_adaptor.go

@@ -3,6 +3,7 @@ package relay
 import (
 	"one-api/relay/channel"
 	"one-api/relay/channel/ali"
+	"one-api/relay/channel/aws"
 	"one-api/relay/channel/baidu"
 	"one-api/relay/channel/claude"
 	"one-api/relay/channel/gemini"
@@ -45,6 +46,8 @@ func GetAdaptor(apiType int) channel.Adaptor {
 		return &ollama.Adaptor{}
 	case constant.APITypePerplexity:
 		return &perplexity.Adaptor{}
+	case constant.APITypeAws:
+		return &aws.Adaptor{}
 	}
 	return nil
 }