xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-09 02:03:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge commit '1c2654320e5b6268b13b6efca40ce37a523d032b'

Browse Source
This commit is contained in:
Laisky.Cai
2024-04-28 01:53:31 +00:00
parent e607a44d53 1c2654320e
commit ca23dd9a97
162 changed files with 1351 additions and 731 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
relay/adaptor/aws/adapter.go
View File

@@ -4,11 +4,14 @@ import (
"io"
"net/http"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/relay/adaptor"
"github.com/Laisky/one-api/relay/adaptor/anthropic"
"github.com/Laisky/one-api/relay/meta"
"github.com/Laisky/one-api/relay/model"
"github.com/songquanpeng/one-api/common/ctxkey"
"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"
"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/gin-gonic/gin"
"github.com/pkg/errors"
)
@@ -16,10 +19,16 @@ import (
var _ adaptor.Adaptor = new(Adaptor)
type Adaptor struct {
meta *meta.Meta
awsClient *bedrockruntime.Client
}
func (a *Adaptor) Init(meta *meta.Meta) {
a.meta = meta
a.awsClient = bedrockruntime.New(bedrockruntime.Options{
Region: meta.Config.Region,
Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(meta.Config.AK, meta.Config.SK, "")),
})
}
func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
@@ -54,9 +63,9 @@ func (a *Adaptor) DoRequest(c *gin.Context, meta *meta.Meta, requestBody io.Read
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)
err, usage = StreamHandler(c, a.awsClient)
} else {
err, usage = Handler(c, resp, meta.PromptTokens, meta.ActualModelName)
err, usage = Handler(c, a.awsClient, meta.ActualModelName)
}
return
}
@@ -65,7 +74,6 @@ func (a *Adaptor) GetModelList() (models []string) {
for n := range awsModelIDMap {
models = append(models, n)
}
return
}

39
relay/adaptor/aws/main.go
View File

@@ -8,14 +8,13 @@ import (
"io"
"net/http"
"github.com/Laisky/one-api/common"
"github.com/Laisky/one-api/common/ctxkey"
"github.com/Laisky/one-api/common/helper"
"github.com/Laisky/one-api/common/logger"
"github.com/Laisky/one-api/relay/adaptor/anthropic"
relaymodel "github.com/Laisky/one-api/relay/model"
"github.com/songquanpeng/one-api/common"
"github.com/songquanpeng/one-api/common/ctxkey"
"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"
"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"
@@ -23,18 +22,6 @@ import (
"github.com/pkg/errors"
)
func newAwsClient(c *gin.Context) (*bedrockruntime.Client, error) {
ak := c.GetString(ctxkey.ConfigAK)
sk := c.GetString(ctxkey.ConfigSK)
region := c.GetString(ctxkey.ConfigRegion)
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,
@@ -62,12 +49,7 @@ func awsModelID(requestModel string) (string, error) {
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
}
func Handler(c *gin.Context, awsCli *bedrockruntime.Client, modelName string) (*relaymodel.ErrorWithStatusCode, *relaymodel.Usage) {
awsModelId, err := awsModelID(c.GetString(ctxkey.RequestModel))
if err != nil {
return wrapErr(errors.Wrap(err, "awsModelID")), nil
@@ -120,13 +102,8 @@ func Handler(c *gin.Context, resp *http.Response, promptTokens int, modelName st
return nil, &usage
}
func StreamHandler(c *gin.Context, resp *http.Response) (*relaymodel.ErrorWithStatusCode, *relaymodel.Usage) {
func StreamHandler(c *gin.Context, awsCli *bedrockruntime.Client) (*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

2
relay/adaptor/aws/model.go
View File

@@ -1,6 +1,6 @@
package aws
import "github.com/Laisky/one-api/relay/adaptor/anthropic"
import "github.com/songquanpeng/one-api/relay/adaptor/anthropic"
// Request is the request to AWS Claude
//