feat: support aws bedrockruntime claude3 (#1328)

* feat: support aws bedrockruntime claude3

closes #622, closes #749, closes #1300

* fix: convert to aws claude model id

* fix: Update AWS adapter to handle stream completions and calculate usage metrics

Based on the file summaries provided, here are the important bullet points for the commit message:

- Add functionality to handle stream completion events from AWS in the relay/adaptor/aws/main.go file
- Marshall AWS response to OpenAI format and calculate usage metrics in the same file
- Implement a custom render function for streaming events in the same file
- Improve error handling for JSON unmarshalling and marshalling errors in the same file

* fix: Implement AWS handler with usage tracking and error handling

- Implemented streaming response handling for AWS handler
- Set response content type to text/event-stream
- Added error handling for failed marshaling/unmarshaling
- Updated return values to include `relaymodel.ErrorWithStatusCode` and `relaymodel.Usage`
- Improved error handling and response formatting for AWS adaptor

* fix: Refactor AWS Adapter for Improved Model Mapping and Error Handling

* Refactor AWS adapter to improve model management
  - Replace hardcoded model list in `adapter.go` with a function to get models from `awsModelIDMap`
  - Update `GetModelList` function to return model list directly
  - Add `GetChannelName` function to get channel name from `Adaptor` object
* Improve error handling and code organization in main.go
  - Replace switch statement with a map to map AWS model IDs to OpenAI model IDs
  - Return an error if the model is not found in the map
  - Use a single return statement instead of wrapping multiple return statements in the `awsModelID` function
  - Add a new error message for when the model is not found in the map in the `Handler` function

* fix: bug fix

* chore: change variable name & package

* chore: change variable name

* perf: update config related code

---------

Co-authored-by: JustSong <songquanpeng@foxmail.com>

relay/adaptor/anthropic/main.go

@@ -91,7 +91,7 @@ func ConvertRequest(textRequest model.GeneralOpenAIRequest) *Request {
 }
 
 // https://docs.anthropic.com/claude/reference/messages-streaming
-func streamResponseClaude2OpenAI(claudeResponse *StreamResponse) (*openai.ChatCompletionsStreamResponse, *Response) {
+func StreamResponseClaude2OpenAI(claudeResponse *StreamResponse) (*openai.ChatCompletionsStreamResponse, *Response) {
 	var response *Response
 	var responseText string
 	var stopReason string
@@ -129,7 +129,7 @@ func streamResponseClaude2OpenAI(claudeResponse *StreamResponse) (*openai.ChatCo
 	return &openaiResponse, response
 }
 
-func responseClaude2OpenAI(claudeResponse *Response) *openai.TextResponse {
+func ResponseClaude2OpenAI(claudeResponse *Response) *openai.TextResponse {
 	var responseText string
 	if len(claudeResponse.Content) > 0 {
 		responseText = claudeResponse.Content[0].Text
@@ -199,7 +199,7 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
 				logger.SysError("error unmarshalling stream response: " + err.Error())
 				return true
 			}
-			response, meta := streamResponseClaude2OpenAI(&claudeResponse)
+			response, meta := StreamResponseClaude2OpenAI(&claudeResponse)
 			if meta != nil {
 				usage.PromptTokens += meta.Usage.InputTokens
 				usage.CompletionTokens += meta.Usage.OutputTokens
@@ -254,7 +254,7 @@ func Handler(c *gin.Context, resp *http.Response, promptTokens int, modelName st
 			StatusCode: resp.StatusCode,
 		}, nil
 	}
-	fullTextResponse := responseClaude2OpenAI(&claudeResponse)
+	fullTextResponse := ResponseClaude2OpenAI(&claudeResponse)
 	fullTextResponse.Model = modelName
 	usage := model.Usage{
 		PromptTokens:     claudeResponse.Usage.InputTokens,

relay/adaptor/aws/adapter.go

@@ -0,0 +1,74 @@
+package aws
+
+import (
+	"github.com/songquanpeng/one-api/common/ctxkey"
+	"io"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/relay/adaptor"
+	"github.com/songquanpeng/one-api/relay/adaptor/anthropic"
+	"github.com/songquanpeng/one-api/relay/meta"
+	"github.com/songquanpeng/one-api/relay/model"
+)
+
+var _ adaptor.Adaptor = new(Adaptor)
+
+type Adaptor struct {
+}
+
+func (a *Adaptor) Init(meta *meta.Meta) {
+
+}
+
+func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
+	return "", nil
+}
+
+func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, meta *meta.Meta) error {
+	return nil
+}
+
+func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *model.GeneralOpenAIRequest) (any, error) {
+	if request == nil {
+		return nil, errors.New("request is nil")
+	}
+
+	claudeReq := anthropic.ConvertRequest(*request)
+	c.Set(ctxkey.RequestModel, request.Model)
+	c.Set(ctxkey.ConvertedRequest, claudeReq)
+	return claudeReq, nil
+}
+
+func (a *Adaptor) ConvertImageRequest(request *model.ImageRequest) (any, error) {
+	if request == nil {
+		return nil, errors.New("request is nil")
+	}
+	return request, nil
+}
+
+func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Reader) (*http.Response, error) {
+	return nil, nil
+}
+
+func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, meta *meta.Meta) (usage *model.Usage, err *model.ErrorWithStatusCode) {
+	if meta.IsStream {
+		err, usage = StreamHandler(c, resp)
+	} else {
+		err, usage = Handler(c, resp, meta.PromptTokens, meta.ActualModelName)
+	}
+	return
+}
+
+func (a *Adaptor) GetModelList() (models []string) {
+	for n := range awsModelIDMap {
+		models = append(models, n)
+	}
+
+	return
+}
+
+func (a *Adaptor) GetChannelName() string {
+	return "aws"
+}

relay/adaptor/aws/main.go

@@ -0,0 +1,215 @@
+// Package aws provides the AWS adaptor for the relay service.
+package aws
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"github.com/songquanpeng/one-api/common/config"
+	"github.com/songquanpeng/one-api/common/ctxkey"
+	"io"
+	"net/http"
+
+	"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"
+	"github.com/gin-gonic/gin"
+	"github.com/jinzhu/copier"
+	"github.com/pkg/errors"
+	"github.com/songquanpeng/one-api/common"
+	"github.com/songquanpeng/one-api/common/helper"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/relay/adaptor/anthropic"
+	relaymodel "github.com/songquanpeng/one-api/relay/model"
+)
+
+func newAwsClient(c *gin.Context) (*bedrockruntime.Client, error) {
+	ak := c.GetString(config.KeyAK)
+	sk := c.GetString(config.KeySK)
+	region := c.GetString(config.KeyRegion)
+	client := bedrockruntime.New(bedrockruntime.Options{
+		Region:      region,
+		Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(ak, sk, "")),
+	})
+
+	return client, nil
+}
+
+func wrapErr(err error) *relaymodel.ErrorWithStatusCode {
+	return &relaymodel.ErrorWithStatusCode{
+		StatusCode: http.StatusInternalServerError,
+		Error: relaymodel.Error{
+			Message: fmt.Sprintf("%s", err.Error()),
+		},
+	}
+}
+
+// https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
+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",
+}
+
+func awsModelID(requestModel string) (string, error) {
+	if awsModelID, ok := awsModelIDMap[requestModel]; ok {
+		return awsModelID, nil
+	}
+
+	return "", errors.Errorf("model %s not found", requestModel)
+}
+
+func Handler(c *gin.Context, resp *http.Response, promptTokens int, modelName string) (*relaymodel.ErrorWithStatusCode, *relaymodel.Usage) {
+	awsCli, err := newAwsClient(c)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "newAwsClient")), nil
+	}
+
+	awsModelId, err := awsModelID(c.GetString(ctxkey.RequestModel))
+	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(ctxkey.ConvertedRequest)
+	if !ok {
+		return wrapErr(errors.New("request not found")), nil
+	}
+	claudeReq := claudeReq_.(*anthropic.Request)
+	awsClaudeReq := &Request{
+		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(anthropic.Response)
+	err = json.Unmarshal(awsResp.Body, claudeResponse)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "unmarshal response")), nil
+	}
+
+	openaiResp := anthropic.ResponseClaude2OpenAI(claudeResponse)
+	openaiResp.Model = modelName
+	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 StreamHandler(c *gin.Context, resp *http.Response) (*relaymodel.ErrorWithStatusCode, *relaymodel.Usage) {
+	createdTime := helper.GetTimestamp()
+	awsCli, err := newAwsClient(c)
+	if err != nil {
+		return wrapErr(errors.Wrap(err, "newAwsClient")), nil
+	}
+
+	awsModelId, err := awsModelID(c.GetString(ctxkey.RequestModel))
+	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(ctxkey.ConvertedRequest)
+	if !ok {
+		return wrapErr(errors.New("request not found")), nil
+	}
+	claudeReq := claudeReq_.(*anthropic.Request)
+
+	awsClaudeReq := &Request{
+		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
+	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(anthropic.StreamResponse)
+			err := json.NewDecoder(bytes.NewReader(v.Value.Bytes)).Decode(claudeResp)
+			if err != nil {
+				logger.SysError("error unmarshalling stream response: " + err.Error())
+				return false
+			}
+
+			response, meta := anthropic.StreamResponseClaude2OpenAI(claudeResp)
+			if meta != nil {
+				usage.PromptTokens += meta.Usage.InputTokens
+				usage.CompletionTokens += meta.Usage.OutputTokens
+				id = fmt.Sprintf("chatcmpl-%s", meta.Id)
+				return true
+			}
+			if response == nil {
+				return true
+			}
+			response.Id = id
+			response.Model = c.GetString(ctxkey.OriginalModel)
+			response.Created = createdTime
+			jsonStr, err := json.Marshal(response)
+			if err != nil {
+				logger.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/adaptor/aws/model.go

@@ -0,0 +1,17 @@
+package aws
+
+import "github.com/songquanpeng/one-api/relay/adaptor/anthropic"
+
+// Request is the request to AWS Claude
+//
+// https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages.html
+type Request struct {
+	// AnthropicVersion should be "bedrock-2023-05-31"
+	AnthropicVersion string              `json:"anthropic_version"`
+	Messages         []anthropic.Message `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"`
+}