重构前
This commit is contained in:
@@ -0,0 +1,188 @@
|
||||
package aws
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/claude"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
|
||||
"github.com/pkg/errors"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type ClientMode int
|
||||
|
||||
const (
|
||||
ClientModeApiKey ClientMode = iota + 1
|
||||
ClientModeAKSK
|
||||
)
|
||||
|
||||
type Adaptor struct {
|
||||
ClientMode ClientMode
|
||||
AwsClient *bedrockruntime.Client
|
||||
AwsModelId string
|
||||
AwsReq any
|
||||
IsNova bool
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertGeminiRequest(*gin.Context, *relaycommon.RelayInfo, *dto.GeminiChatRequest) (any, error) {
|
||||
//TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertClaudeRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.ClaudeRequest) (any, error) {
|
||||
for i, message := range request.Messages {
|
||||
updated := false
|
||||
if !message.IsStringContent() {
|
||||
content, err := message.ParseContent()
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to parse message content")
|
||||
}
|
||||
for i2, mediaMessage := range content {
|
||||
if mediaMessage.Source != nil {
|
||||
if mediaMessage.Source.Type == "url" {
|
||||
// 使用统一的文件服务获取图片数据
|
||||
source := types.NewURLFileSource(mediaMessage.Source.Url)
|
||||
base64Data, mimeType, err := service.GetBase64Data(c, source, "formatting image for Claude")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get file base64 from url failed: %s", err.Error())
|
||||
}
|
||||
mediaMessage.Source.MediaType = mimeType
|
||||
mediaMessage.Source.Data = base64Data
|
||||
mediaMessage.Source.Url = ""
|
||||
mediaMessage.Source.Type = "base64"
|
||||
content[i2] = mediaMessage
|
||||
updated = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if updated {
|
||||
message.SetContent(content)
|
||||
}
|
||||
}
|
||||
if updated {
|
||||
request.Messages[i] = message
|
||||
}
|
||||
}
|
||||
return request, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertAudioRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.AudioRequest) (io.Reader, error) {
|
||||
//TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertImageRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.ImageRequest) (any, error) {
|
||||
//TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (a *Adaptor) Init(info *relaycommon.RelayInfo) {
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
|
||||
if info.ChannelOtherSettings.AwsKeyType == dto.AwsKeyTypeApiKey {
|
||||
awsModelId := getAwsModelID(info.UpstreamModelName)
|
||||
a.ClientMode = ClientModeApiKey
|
||||
awsSecret := strings.Split(info.ApiKey, "|")
|
||||
if len(awsSecret) != 2 {
|
||||
return "", errors.New("invalid aws api key, should be in format of <api-key>|<region>")
|
||||
}
|
||||
return fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com/model/%s/converse", awsModelId, awsSecret[1]), nil
|
||||
} else {
|
||||
a.ClientMode = ClientModeAKSK
|
||||
return "", nil
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
|
||||
claude.CommonClaudeHeadersOperation(c, req, info)
|
||||
if a.ClientMode == ClientModeApiKey {
|
||||
req.Set("Authorization", "Bearer "+info.ApiKey)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertOpenAIRequest(c *gin.Context, info *relaycommon.RelayInfo, request *dto.GeneralOpenAIRequest) (any, error) {
|
||||
if request == nil {
|
||||
return nil, errors.New("request is nil")
|
||||
}
|
||||
// 检查是否为Nova模型
|
||||
if isNovaModel(request.Model) {
|
||||
novaReq := convertToNovaRequest(request)
|
||||
a.IsNova = true
|
||||
return novaReq, nil
|
||||
}
|
||||
|
||||
// 原有的Claude模型处理逻辑
|
||||
result, err := service.ConvertRequest(c, info, types.RelayFormatClaude, request)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "failed to convert openai request to claude request")
|
||||
}
|
||||
claudeReq, ok := result.Value.(*dto.ClaudeRequest)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("expected Anthropic Messages request, got %T", result.Value)
|
||||
}
|
||||
info.UpstreamModelName = claudeReq.Model
|
||||
return claudeReq, err
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.EmbeddingRequest) (any, error) {
|
||||
//TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) {
|
||||
// TODO implement me
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) {
|
||||
if a.ClientMode == ClientModeApiKey {
|
||||
return channel.DoApiRequest(a, c, info, requestBody)
|
||||
} else {
|
||||
return doAwsClientRequest(c, info, a, requestBody)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
|
||||
if a.ClientMode == ClientModeApiKey {
|
||||
claudeAdaptor := claude.Adaptor{}
|
||||
usage, err = claudeAdaptor.DoResponse(c, resp, info)
|
||||
} else {
|
||||
if a.IsNova {
|
||||
err, usage = handleNovaRequest(c, info, a)
|
||||
} else {
|
||||
if info.IsStream {
|
||||
err, usage = awsStreamHandler(c, info, a)
|
||||
} else {
|
||||
err, usage = awsHandler(c, info, a)
|
||||
}
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetModelList() (models []string) {
|
||||
for n := range awsModelIDMap {
|
||||
models = append(models, n)
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetChannelName() string {
|
||||
return ChannelName
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
package aws
|
||||
|
||||
import "strings"
|
||||
|
||||
var awsModelIDMap = map[string]string{
|
||||
"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",
|
||||
"claude-3-5-sonnet-20240620": "anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"claude-3-5-sonnet-20241022": "anthropic.claude-3-5-sonnet-20241022-v2:0",
|
||||
"claude-3-5-haiku-20241022": "anthropic.claude-3-5-haiku-20241022-v1:0",
|
||||
"claude-3-7-sonnet-20250219": "anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
"claude-sonnet-4-20250514": "anthropic.claude-sonnet-4-20250514-v1:0",
|
||||
"claude-opus-4-20250514": "anthropic.claude-opus-4-20250514-v1:0",
|
||||
"claude-opus-4-1-20250805": "anthropic.claude-opus-4-1-20250805-v1:0",
|
||||
"claude-sonnet-4-5-20250929": "anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"claude-sonnet-4-6": "anthropic.claude-sonnet-4-6",
|
||||
"claude-haiku-4-5-20251001": "anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"claude-opus-4-5-20251101": "anthropic.claude-opus-4-5-20251101-v1:0",
|
||||
"claude-opus-4-6": "anthropic.claude-opus-4-6-v1",
|
||||
"claude-opus-4-7": "anthropic.claude-opus-4-7",
|
||||
"claude-opus-4-8": "anthropic.claude-opus-4-8",
|
||||
// Nova models
|
||||
"nova-micro-v1:0": "amazon.nova-micro-v1:0",
|
||||
"nova-lite-v1:0": "amazon.nova-lite-v1:0",
|
||||
"nova-pro-v1:0": "amazon.nova-pro-v1:0",
|
||||
"nova-premier-v1:0": "amazon.nova-premier-v1:0",
|
||||
"nova-canvas-v1:0": "amazon.nova-canvas-v1:0",
|
||||
"nova-reel-v1:0": "amazon.nova-reel-v1:0",
|
||||
"nova-reel-v1:1": "amazon.nova-reel-v1:1",
|
||||
"nova-sonic-v1:0": "amazon.nova-sonic-v1:0",
|
||||
}
|
||||
|
||||
var awsModelCanCrossRegionMap = map[string]map[string]bool{
|
||||
"anthropic.claude-3-sonnet-20240229-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"ap": true,
|
||||
},
|
||||
"anthropic.claude-3-opus-20240229-v1:0": {
|
||||
"us": true,
|
||||
},
|
||||
"anthropic.claude-3-haiku-20240307-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"ap": true,
|
||||
},
|
||||
"anthropic.claude-3-5-sonnet-20240620-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"ap": true,
|
||||
},
|
||||
"anthropic.claude-3-5-sonnet-20241022-v2:0": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
},
|
||||
"anthropic.claude-3-5-haiku-20241022-v1:0": {
|
||||
"us": true,
|
||||
},
|
||||
"anthropic.claude-3-7-sonnet-20250219-v1:0": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
"anthropic.claude-sonnet-4-20250514-v1:0": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
"anthropic.claude-opus-4-20250514-v1:0": {
|
||||
"us": true,
|
||||
},
|
||||
"anthropic.claude-opus-4-1-20250805-v1:0": {
|
||||
"us": true,
|
||||
},
|
||||
"anthropic.claude-sonnet-4-5-20250929-v1:0": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
"anthropic.claude-sonnet-4-6": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
"anthropic.claude-opus-4-5-20251101-v1:0": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
"anthropic.claude-opus-4-6-v1": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
"anthropic.claude-opus-4-7": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
"anthropic.claude-opus-4-8": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
"anthropic.claude-haiku-4-5-20251001-v1:0": {
|
||||
"us": true,
|
||||
"ap": true,
|
||||
"eu": true,
|
||||
},
|
||||
// Nova models - all support three major regions
|
||||
"amazon.nova-micro-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"apac": true,
|
||||
},
|
||||
"amazon.nova-lite-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"apac": true,
|
||||
},
|
||||
"amazon.nova-pro-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"apac": true,
|
||||
},
|
||||
"amazon.nova-premier-v1:0": {
|
||||
"us": true,
|
||||
},
|
||||
"amazon.nova-canvas-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"apac": true,
|
||||
},
|
||||
"amazon.nova-reel-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"apac": true,
|
||||
},
|
||||
"amazon.nova-reel-v1:1": {
|
||||
"us": true,
|
||||
},
|
||||
"amazon.nova-sonic-v1:0": {
|
||||
"us": true,
|
||||
"eu": true,
|
||||
"apac": true,
|
||||
},
|
||||
}
|
||||
|
||||
var awsRegionCrossModelPrefixMap = map[string]string{
|
||||
"us": "us",
|
||||
"eu": "eu",
|
||||
"ap": "apac",
|
||||
}
|
||||
|
||||
var ChannelName = "aws"
|
||||
|
||||
// 判断是否为Nova模型
|
||||
func isNovaModel(modelId string) bool {
|
||||
return strings.Contains(modelId, "nova-")
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package aws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
)
|
||||
|
||||
type AwsClaudeRequest struct {
|
||||
// AnthropicVersion should be "bedrock-2023-05-31"
|
||||
AnthropicVersion string `json:"anthropic_version"`
|
||||
AnthropicBeta json.RawMessage `json:"anthropic_beta,omitempty"`
|
||||
System any `json:"system,omitempty"`
|
||||
Messages []dto.ClaudeMessage `json:"messages"`
|
||||
MaxTokens *uint `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"`
|
||||
Tools any `json:"tools,omitempty"`
|
||||
ToolChoice any `json:"tool_choice,omitempty"`
|
||||
ContextManagement json.RawMessage `json:"context_management,omitempty"`
|
||||
Thinking *dto.Thinking `json:"thinking,omitempty"`
|
||||
OutputConfig json.RawMessage `json:"output_config,omitempty"`
|
||||
//Metadata json.RawMessage `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
func formatRequest(requestBody io.Reader, requestHeader http.Header) (*AwsClaudeRequest, error) {
|
||||
var awsClaudeRequest AwsClaudeRequest
|
||||
err := common.DecodeJson(requestBody, &awsClaudeRequest)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
awsClaudeRequest.AnthropicVersion = "bedrock-2023-05-31"
|
||||
|
||||
// check header anthropic-beta
|
||||
anthropicBetaValues := requestHeader.Get("anthropic-beta")
|
||||
if len(anthropicBetaValues) > 0 {
|
||||
var tempArray []string
|
||||
tempArray = strings.Split(anthropicBetaValues, ",")
|
||||
if len(tempArray) > 0 {
|
||||
betaJson, err := json.Marshal(tempArray)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
awsClaudeRequest.AnthropicBeta = betaJson
|
||||
}
|
||||
}
|
||||
logger.LogJson(context.Background(), "json", awsClaudeRequest)
|
||||
return &awsClaudeRequest, nil
|
||||
}
|
||||
|
||||
// NovaMessage Nova模型使用messages-v1格式
|
||||
type NovaMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content []NovaContent `json:"content"`
|
||||
}
|
||||
|
||||
type NovaContent struct {
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
type NovaRequest struct {
|
||||
SchemaVersion string `json:"schemaVersion"` // 请求版本，例如 "1.0"
|
||||
Messages []NovaMessage `json:"messages"` // 对话消息列表
|
||||
InferenceConfig *NovaInferenceConfig `json:"inferenceConfig,omitempty"` // 推理配置，可选
|
||||
}
|
||||
|
||||
type NovaInferenceConfig struct {
|
||||
MaxTokens int `json:"maxTokens,omitempty"` // 最大生成的 token 数
|
||||
Temperature float64 `json:"temperature,omitempty"` // 随机性 (默认 0.7, 范围 0-1)
|
||||
TopP float64 `json:"topP,omitempty"` // nucleus sampling (默认 0.9, 范围 0-1)
|
||||
TopK int `json:"topK,omitempty"` // 限制候选 token 数 (默认 50, 范围 0-128)
|
||||
StopSequences []string `json:"stopSequences,omitempty"` // 停止生成的序列
|
||||
}
|
||||
|
||||
// 转换OpenAI请求为Nova格式
|
||||
func convertToNovaRequest(req *dto.GeneralOpenAIRequest) *NovaRequest {
|
||||
novaMessages := make([]NovaMessage, len(req.Messages))
|
||||
for i, msg := range req.Messages {
|
||||
novaMessages[i] = NovaMessage{
|
||||
Role: msg.Role,
|
||||
Content: []NovaContent{{Text: msg.StringContent()}},
|
||||
}
|
||||
}
|
||||
|
||||
novaReq := &NovaRequest{
|
||||
SchemaVersion: "messages-v1",
|
||||
Messages: novaMessages,
|
||||
}
|
||||
|
||||
// 设置推理配置
|
||||
if (req.MaxTokens != nil && *req.MaxTokens != 0) || (req.Temperature != nil && *req.Temperature != 0) || (req.TopP != nil && *req.TopP != 0) || (req.TopK != nil && *req.TopK != 0) || req.Stop != nil {
|
||||
novaReq.InferenceConfig = &NovaInferenceConfig{}
|
||||
if req.MaxTokens != nil && *req.MaxTokens != 0 {
|
||||
novaReq.InferenceConfig.MaxTokens = int(*req.MaxTokens)
|
||||
}
|
||||
if req.Temperature != nil && *req.Temperature != 0 {
|
||||
novaReq.InferenceConfig.Temperature = *req.Temperature
|
||||
}
|
||||
if req.TopP != nil && *req.TopP != 0 {
|
||||
novaReq.InferenceConfig.TopP = *req.TopP
|
||||
}
|
||||
if req.TopK != nil && *req.TopK != 0 {
|
||||
novaReq.InferenceConfig.TopK = *req.TopK
|
||||
}
|
||||
if req.Stop != nil {
|
||||
if stopSequences := parseStopSequences(req.Stop); len(stopSequences) > 0 {
|
||||
novaReq.InferenceConfig.StopSequences = stopSequences
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return novaReq
|
||||
}
|
||||
|
||||
// parseStopSequences 解析停止序列，支持字符串或字符串数组
|
||||
func parseStopSequences(stop any) []string {
|
||||
if stop == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch v := stop.(type) {
|
||||
case string:
|
||||
if v != "" {
|
||||
return []string{v}
|
||||
}
|
||||
case []string:
|
||||
return v
|
||||
case []interface{}:
|
||||
var sequences []string
|
||||
for _, item := range v {
|
||||
if str, ok := item.(string); ok && str != "" {
|
||||
sequences = append(sequences, str)
|
||||
}
|
||||
}
|
||||
return sequences
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,371 @@
|
||||
package aws
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/claude"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/pkg/errors"
|
||||
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"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"
|
||||
bedrockruntimeTypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
|
||||
"github.com/aws/smithy-go/auth/bearer"
|
||||
)
|
||||
|
||||
// getAwsErrorStatusCode extracts HTTP status code from AWS SDK error
|
||||
func getAwsErrorStatusCode(err error) int {
|
||||
// Check for HTTP response error which contains status code
|
||||
var httpErr interface{ HTTPStatusCode() int }
|
||||
if errors.As(err, &httpErr) {
|
||||
return httpErr.HTTPStatusCode()
|
||||
}
|
||||
// Default to 500 if we can't determine the status code
|
||||
return http.StatusInternalServerError
|
||||
}
|
||||
|
||||
func newAwsInvokeContext(parent context.Context) (context.Context, context.CancelFunc) {
|
||||
if common.RelayTimeout <= 0 {
|
||||
return context.WithCancel(parent)
|
||||
}
|
||||
return context.WithTimeout(parent, time.Duration(common.RelayTimeout)*time.Second)
|
||||
}
|
||||
|
||||
func newAwsInvokeError(requestContext context.Context, err error, operation string) *types.NewAPIError {
|
||||
options := make([]types.NewAPIErrorOptions, 0, 1)
|
||||
if requestContext.Err() != nil {
|
||||
options = append(options, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
return types.NewOpenAIError(
|
||||
errors.Wrap(err, operation),
|
||||
types.ErrorCodeAwsInvokeError,
|
||||
getAwsErrorStatusCode(err),
|
||||
options...,
|
||||
)
|
||||
}
|
||||
|
||||
func newAwsClient(c *gin.Context, info *relaycommon.RelayInfo) (*bedrockruntime.Client, error) {
|
||||
httpClient, err := service.GetHttpClientWithProxySettings(info.ChannelSetting.Proxy, info.ChannelSetting)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("new proxy http client failed: %w", err)
|
||||
}
|
||||
|
||||
awsSecret := strings.Split(info.ApiKey, "|")
|
||||
var client *bedrockruntime.Client
|
||||
switch len(awsSecret) {
|
||||
case 2:
|
||||
apiKey := awsSecret[0]
|
||||
region := awsSecret[1]
|
||||
client = bedrockruntime.New(bedrockruntime.Options{
|
||||
Region: region,
|
||||
BearerAuthTokenProvider: bearer.StaticTokenProvider{Token: bearer.Token{Value: apiKey}},
|
||||
HTTPClient: httpClient,
|
||||
})
|
||||
case 3:
|
||||
ak := awsSecret[0]
|
||||
sk := awsSecret[1]
|
||||
region := awsSecret[2]
|
||||
client = bedrockruntime.New(bedrockruntime.Options{
|
||||
Region: region,
|
||||
Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(ak, sk, "")),
|
||||
HTTPClient: httpClient,
|
||||
})
|
||||
default:
|
||||
return nil, errors.New("invalid aws secret key")
|
||||
}
|
||||
|
||||
return client, nil
|
||||
}
|
||||
|
||||
func doAwsClientRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor, requestBody io.Reader) (any, error) {
|
||||
awsCli, err := newAwsClient(c, info)
|
||||
if err != nil {
|
||||
return nil, types.NewError(err, types.ErrorCodeChannelAwsClientError)
|
||||
}
|
||||
a.AwsClient = awsCli
|
||||
|
||||
// 获取对应的AWS模型ID
|
||||
awsModelId := getAwsModelID(info.UpstreamModelName)
|
||||
|
||||
awsRegionPrefix := getAwsRegionPrefix(awsCli.Options().Region)
|
||||
canCrossRegion := awsModelCanCrossRegion(awsModelId, awsRegionPrefix)
|
||||
if canCrossRegion {
|
||||
awsModelId = awsModelCrossRegion(awsModelId, awsRegionPrefix)
|
||||
}
|
||||
|
||||
// init empty request.header
|
||||
requestHeader := http.Header{}
|
||||
a.SetupRequestHeader(c, &requestHeader, info)
|
||||
headerOverride, err := channel.ResolveHeaderOverride(info, c)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for key, value := range headerOverride {
|
||||
requestHeader.Set(key, value)
|
||||
}
|
||||
|
||||
if isNovaModel(awsModelId) {
|
||||
var novaReq *NovaRequest
|
||||
err = common.DecodeJson(requestBody, &novaReq)
|
||||
if err != nil {
|
||||
return nil, types.NewError(errors.Wrap(err, "decode nova request fail"), types.ErrorCodeBadRequestBody)
|
||||
}
|
||||
|
||||
// 使用InvokeModel API，但使用Nova格式的请求体
|
||||
awsReq := &bedrockruntime.InvokeModelInput{
|
||||
ModelId: aws.String(awsModelId),
|
||||
Accept: aws.String("application/json"),
|
||||
ContentType: aws.String("application/json"),
|
||||
}
|
||||
|
||||
reqBody, err := common.Marshal(novaReq)
|
||||
if err != nil {
|
||||
return nil, types.NewError(errors.Wrap(err, "marshal nova request"), types.ErrorCodeBadResponseBody)
|
||||
}
|
||||
awsReq.Body = reqBody
|
||||
a.AwsReq = awsReq
|
||||
return nil, nil
|
||||
} else {
|
||||
awsClaudeReq, err := formatRequest(requestBody, requestHeader)
|
||||
if err != nil {
|
||||
return nil, types.NewError(errors.Wrap(err, "format aws request fail"), types.ErrorCodeBadRequestBody)
|
||||
}
|
||||
|
||||
if info.IsStream {
|
||||
awsReq := &bedrockruntime.InvokeModelWithResponseStreamInput{
|
||||
ModelId: aws.String(awsModelId),
|
||||
Accept: aws.String("application/json"),
|
||||
ContentType: aws.String("application/json"),
|
||||
}
|
||||
awsReq.Body, err = buildAwsRequestBody(c, info, awsClaudeReq)
|
||||
if err != nil {
|
||||
return nil, types.NewError(errors.Wrap(err, "marshal aws request fail"), types.ErrorCodeBadRequestBody)
|
||||
}
|
||||
a.AwsReq = awsReq
|
||||
return nil, nil
|
||||
} else {
|
||||
awsReq := &bedrockruntime.InvokeModelInput{
|
||||
ModelId: aws.String(awsModelId),
|
||||
Accept: aws.String("application/json"),
|
||||
ContentType: aws.String("application/json"),
|
||||
}
|
||||
awsReq.Body, err = buildAwsRequestBody(c, info, awsClaudeReq)
|
||||
if err != nil {
|
||||
return nil, types.NewError(errors.Wrap(err, "marshal aws request fail"), types.ErrorCodeBadRequestBody)
|
||||
}
|
||||
a.AwsReq = awsReq
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// buildAwsRequestBody prepares the payload for AWS requests, applying passthrough rules when enabled.
|
||||
func buildAwsRequestBody(c *gin.Context, info *relaycommon.RelayInfo, awsClaudeReq any) ([]byte, error) {
|
||||
if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled {
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "get request body for pass-through fail")
|
||||
}
|
||||
body, err := storage.Bytes()
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "get request body bytes fail")
|
||||
}
|
||||
var data map[string]interface{}
|
||||
if err := common.Unmarshal(body, &data); err != nil {
|
||||
return nil, errors.Wrap(err, "pass-through unmarshal request body fail")
|
||||
}
|
||||
delete(data, "model")
|
||||
delete(data, "stream")
|
||||
return common.Marshal(data)
|
||||
}
|
||||
return common.Marshal(awsClaudeReq)
|
||||
}
|
||||
|
||||
func getAwsRegionPrefix(awsRegionId string) string {
|
||||
parts := strings.Split(awsRegionId, "-")
|
||||
regionPrefix := ""
|
||||
if len(parts) > 0 {
|
||||
regionPrefix = parts[0]
|
||||
}
|
||||
return regionPrefix
|
||||
}
|
||||
|
||||
func awsModelCanCrossRegion(awsModelId, awsRegionPrefix string) bool {
|
||||
regionSet, exists := awsModelCanCrossRegionMap[awsModelId]
|
||||
return exists && regionSet[awsRegionPrefix]
|
||||
}
|
||||
|
||||
func awsModelCrossRegion(awsModelId, awsRegionPrefix string) string {
|
||||
modelPrefix, find := awsRegionCrossModelPrefixMap[awsRegionPrefix]
|
||||
if !find {
|
||||
return awsModelId
|
||||
}
|
||||
return modelPrefix + "." + awsModelId
|
||||
}
|
||||
|
||||
func getAwsModelID(requestModel string) string {
|
||||
if awsModelIDName, ok := awsModelIDMap[requestModel]; ok {
|
||||
return awsModelIDName
|
||||
}
|
||||
return requestModel
|
||||
}
|
||||
|
||||
func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
|
||||
|
||||
requestContext := c.Request.Context()
|
||||
ctx, cancel := newAwsInvokeContext(requestContext)
|
||||
defer cancel()
|
||||
|
||||
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
|
||||
if err != nil {
|
||||
return newAwsInvokeError(requestContext, err, "InvokeModel"), nil
|
||||
}
|
||||
|
||||
claudeInfo := &claude.ClaudeResponseInfo{
|
||||
ResponseId: helper.GetResponseID(c),
|
||||
Created: common.GetTimestamp(),
|
||||
Model: info.UpstreamModelName,
|
||||
ResponseText: strings.Builder{},
|
||||
Usage: &dto.Usage{},
|
||||
}
|
||||
|
||||
// 复制上游 Content-Type 到客户端响应头
|
||||
if awsResp.ContentType != nil && *awsResp.ContentType != "" {
|
||||
c.Writer.Header().Set("Content-Type", *awsResp.ContentType)
|
||||
}
|
||||
|
||||
handlerErr := claude.HandleClaudeResponseData(c, info, claudeInfo, nil, awsResp.Body)
|
||||
if handlerErr != nil {
|
||||
return handlerErr, nil
|
||||
}
|
||||
return nil, claudeInfo.Usage
|
||||
}
|
||||
|
||||
func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
|
||||
requestContext := c.Request.Context()
|
||||
ctx, cancel := newAwsInvokeContext(requestContext)
|
||||
defer cancel()
|
||||
|
||||
awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput))
|
||||
if err != nil {
|
||||
return newAwsInvokeError(requestContext, err, "InvokeModelWithResponseStream"), nil
|
||||
}
|
||||
stream := awsResp.GetStream()
|
||||
defer stream.Close()
|
||||
|
||||
claudeInfo := &claude.ClaudeResponseInfo{
|
||||
ResponseId: helper.GetResponseID(c),
|
||||
Created: common.GetTimestamp(),
|
||||
Model: info.UpstreamModelName,
|
||||
ResponseText: strings.Builder{},
|
||||
Usage: &dto.Usage{},
|
||||
}
|
||||
|
||||
events := stream.Events()
|
||||
streamLoop:
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
break streamLoop
|
||||
case event, ok := <-events:
|
||||
if !ok {
|
||||
break streamLoop
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
break streamLoop
|
||||
}
|
||||
|
||||
switch v := event.(type) {
|
||||
case *bedrockruntimeTypes.ResponseStreamMemberChunk:
|
||||
info.SetFirstResponseTime()
|
||||
respErr := claude.HandleStreamResponseData(c, info, claudeInfo, string(v.Value.Bytes))
|
||||
if respErr != nil {
|
||||
return respErr, nil
|
||||
}
|
||||
case *bedrockruntimeTypes.UnknownUnionMember:
|
||||
fmt.Println("unknown tag:", v.Tag)
|
||||
return types.NewError(errors.New("unknown response type"), types.ErrorCodeInvalidRequest), nil
|
||||
default:
|
||||
fmt.Println("union is nil or unknown type")
|
||||
return types.NewError(errors.New("nil or unknown response type"), types.ErrorCodeInvalidRequest), nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_ = stream.Close()
|
||||
claude.HandleStreamFinalResponse(c, info, claudeInfo)
|
||||
return nil, claudeInfo.Usage
|
||||
}
|
||||
|
||||
// Nova模型处理函数
|
||||
func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
|
||||
|
||||
requestContext := c.Request.Context()
|
||||
ctx, cancel := newAwsInvokeContext(requestContext)
|
||||
defer cancel()
|
||||
|
||||
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
|
||||
if err != nil {
|
||||
return newAwsInvokeError(requestContext, err, "InvokeModel"), nil
|
||||
}
|
||||
|
||||
// 解析Nova响应
|
||||
var novaResp struct {
|
||||
Output struct {
|
||||
Message struct {
|
||||
Content []struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"output"`
|
||||
Usage struct {
|
||||
InputTokens int `json:"inputTokens"`
|
||||
OutputTokens int `json:"outputTokens"`
|
||||
TotalTokens int `json:"totalTokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
if err := json.Unmarshal(awsResp.Body, &novaResp); err != nil {
|
||||
return types.NewError(errors.Wrap(err, "unmarshal nova response"), types.ErrorCodeBadResponseBody), nil
|
||||
}
|
||||
|
||||
// 构造OpenAI格式响应
|
||||
response := dto.OpenAITextResponse{
|
||||
Id: helper.GetResponseID(c),
|
||||
Object: "chat.completion",
|
||||
Created: common.GetTimestamp(),
|
||||
Model: info.UpstreamModelName,
|
||||
Choices: []dto.OpenAITextResponseChoice{{
|
||||
Index: 0,
|
||||
Message: dto.Message{
|
||||
Role: "assistant",
|
||||
Content: novaResp.Output.Message.Content[0].Text,
|
||||
},
|
||||
FinishReason: "stop",
|
||||
}},
|
||||
Usage: dto.Usage{
|
||||
PromptTokens: novaResp.Usage.InputTokens,
|
||||
CompletionTokens: novaResp.Usage.OutputTokens,
|
||||
TotalTokens: novaResp.Usage.TotalTokens,
|
||||
},
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, response)
|
||||
return nil, &response.Usage
|
||||
}
|
||||
@@ -0,0 +1,458 @@
|
||||
package aws
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
relaytypes "github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
|
||||
"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream/eventstreamapi"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
const awsTestModel = "anthropic.claude-3-5-sonnet-20240620-v1:0"
|
||||
|
||||
type awsHTTPClientFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f awsHTTPClientFunc) Do(request *http.Request) (*http.Response, error) {
|
||||
return f(request)
|
||||
}
|
||||
|
||||
type awsNotifyingResponseWriter struct {
|
||||
*httptest.ResponseRecorder
|
||||
notifyOn []byte
|
||||
notified chan int
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func newAwsNotifyingResponseWriter(notifyOn string) *awsNotifyingResponseWriter {
|
||||
return &awsNotifyingResponseWriter{
|
||||
ResponseRecorder: httptest.NewRecorder(),
|
||||
notifyOn: []byte(notifyOn),
|
||||
notified: make(chan int, 1),
|
||||
}
|
||||
}
|
||||
|
||||
func (w *awsNotifyingResponseWriter) Write(data []byte) (int, error) {
|
||||
return w.ResponseRecorder.Write(data)
|
||||
}
|
||||
|
||||
func (w *awsNotifyingResponseWriter) Flush() {
|
||||
w.ResponseRecorder.Flush()
|
||||
if bytes.Contains(w.Body.Bytes(), w.notifyOn) {
|
||||
w.once.Do(func() {
|
||||
w.notified <- w.Body.Len()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func newAwsTestClient(httpClient bedrockruntime.HTTPClient) *bedrockruntime.Client {
|
||||
return bedrockruntime.New(bedrockruntime.Options{
|
||||
Region: "us-east-1",
|
||||
BaseEndpoint: aws.String("https://bedrock.test"),
|
||||
Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider(
|
||||
"access-key", "secret-key", "",
|
||||
)),
|
||||
HTTPClient: httpClient,
|
||||
Retryer: aws.NopRetryer{},
|
||||
})
|
||||
}
|
||||
|
||||
func newAwsTestContext(writer http.ResponseWriter, requestContext context.Context) *gin.Context {
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(requestContext)
|
||||
return c
|
||||
}
|
||||
|
||||
func newAwsTestRelayInfo() *relaycommon.RelayInfo {
|
||||
return &relaycommon.RelayInfo{
|
||||
StartTime: time.Now(),
|
||||
IsStream: true,
|
||||
OriginModelName: awsTestModel,
|
||||
RelayFormat: relaytypes.RelayFormatOpenAI,
|
||||
ShouldIncludeUsage: true,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
UpstreamModelName: awsTestModel,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newAwsInvokeModelInput() *bedrockruntime.InvokeModelInput {
|
||||
return &bedrockruntime.InvokeModelInput{
|
||||
ModelId: aws.String(awsTestModel),
|
||||
Body: []byte(`{}`),
|
||||
Accept: aws.String("application/json"),
|
||||
ContentType: aws.String("application/json"),
|
||||
}
|
||||
}
|
||||
|
||||
func newAwsStreamInput() *bedrockruntime.InvokeModelWithResponseStreamInput {
|
||||
return &bedrockruntime.InvokeModelWithResponseStreamInput{
|
||||
ModelId: aws.String(awsTestModel),
|
||||
Body: []byte(`{}`),
|
||||
Accept: aws.String("application/json"),
|
||||
ContentType: aws.String("application/json"),
|
||||
}
|
||||
}
|
||||
|
||||
func writeAwsStreamEvent(writer io.Writer, data string) error {
|
||||
payload, err := common.Marshal(struct {
|
||||
Bytes []byte `json:"bytes"`
|
||||
}{Bytes: []byte(data)})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return eventstream.NewEncoder().Encode(writer, eventstream.Message{
|
||||
Headers: eventstream.Headers{
|
||||
{Name: eventstreamapi.MessageTypeHeader, Value: eventstream.StringValue(eventstreamapi.EventMessageType)},
|
||||
{Name: eventstreamapi.EventTypeHeader, Value: eventstream.StringValue("chunk")},
|
||||
{Name: eventstreamapi.ContentTypeHeader, Value: eventstream.StringValue("application/json")},
|
||||
},
|
||||
Payload: payload,
|
||||
})
|
||||
}
|
||||
|
||||
func newAwsStreamResponse(request *http.Request, body io.ReadCloser) *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Status: "200 OK",
|
||||
Header: http.Header{
|
||||
"Content-Type": []string{"application/vnd.amazon.eventstream"},
|
||||
"X-Amzn-Bedrock-Content-Type": []string{"application/json"},
|
||||
},
|
||||
Body: body,
|
||||
Request: request,
|
||||
}
|
||||
}
|
||||
|
||||
func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "claude-3-5-sonnet-20240620",
|
||||
IsStream: false,
|
||||
UseRuntimeHeadersOverride: true,
|
||||
RuntimeHeadersOverride: map[string]any{
|
||||
"anthropic-beta": "computer-use-2025-01-24",
|
||||
},
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ApiKey: "access-key|secret-key|us-east-1",
|
||||
UpstreamModelName: "claude-3-5-sonnet-20240620",
|
||||
},
|
||||
}
|
||||
|
||||
requestBody := bytes.NewBufferString(`{"messages":[{"role":"user","content":"hello"}],"max_tokens":128}`)
|
||||
adaptor := &Adaptor{}
|
||||
|
||||
_, err := doAwsClientRequest(ctx, info, adaptor, requestBody)
|
||||
require.NoError(t, err)
|
||||
|
||||
awsReq, ok := adaptor.AwsReq.(*bedrockruntime.InvokeModelInput)
|
||||
require.True(t, ok)
|
||||
|
||||
var payload map[string]any
|
||||
require.NoError(t, common.Unmarshal(awsReq.Body, &payload))
|
||||
|
||||
anthropicBeta, exists := payload["anthropic_beta"]
|
||||
require.True(t, exists)
|
||||
|
||||
values, ok := anthropicBeta.([]any)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, []any{"computer-use-2025-01-24"}, values)
|
||||
}
|
||||
|
||||
func TestNewAwsInvokeContextInheritsParent(t *testing.T) {
|
||||
originalRelayTimeout := common.RelayTimeout
|
||||
t.Cleanup(func() {
|
||||
common.RelayTimeout = originalRelayTimeout
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
relayTimeout int
|
||||
wantDeadline bool
|
||||
}{
|
||||
{name: "without relay timeout", relayTimeout: 0, wantDeadline: false},
|
||||
{name: "with relay timeout", relayTimeout: 30, wantDeadline: true},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
common.RelayTimeout = test.relayTimeout
|
||||
parent, cancelParent := context.WithCancel(context.Background())
|
||||
invokeContext, cancelInvoke := newAwsInvokeContext(parent)
|
||||
defer cancelInvoke()
|
||||
|
||||
_, hasDeadline := invokeContext.Deadline()
|
||||
assert.Equal(t, test.wantDeadline, hasDeadline)
|
||||
|
||||
cancelParent()
|
||||
require.ErrorIs(t, invokeContext.Err(), context.Canceled)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewAwsInvokeErrorSkipsRetryOnlyForClientCancellation(t *testing.T) {
|
||||
canceledContext, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
requestContext context.Context
|
||||
err error
|
||||
wantSkipRetry bool
|
||||
}{
|
||||
{
|
||||
name: "client context canceled",
|
||||
requestContext: canceledContext,
|
||||
err: context.Canceled,
|
||||
wantSkipRetry: true,
|
||||
},
|
||||
{
|
||||
name: "relay timeout with live client context",
|
||||
requestContext: context.Background(),
|
||||
err: context.DeadlineExceeded,
|
||||
wantSkipRetry: false,
|
||||
},
|
||||
{
|
||||
name: "upstream error with live client context",
|
||||
requestContext: context.Background(),
|
||||
err: errors.New("upstream failed"),
|
||||
wantSkipRetry: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
err := newAwsInvokeError(test.requestContext, test.err, "InvokeModel")
|
||||
assert.Equal(t, test.wantSkipRetry, relaytypes.IsSkipRetryError(err))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAwsHandlersCancelSdkRequestAndSkipRetry(t *testing.T) {
|
||||
originalRelayTimeout := common.RelayTimeout
|
||||
common.RelayTimeout = 0
|
||||
t.Cleanup(func() {
|
||||
common.RelayTimeout = originalRelayTimeout
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
request any
|
||||
handle func(*gin.Context, *relaycommon.RelayInfo, *Adaptor) (*relaytypes.NewAPIError, *dto.Usage)
|
||||
}{
|
||||
{name: "non-stream", request: newAwsInvokeModelInput(), handle: awsHandler},
|
||||
{name: "stream", request: newAwsStreamInput(), handle: awsStreamHandler},
|
||||
{name: "nova", request: newAwsInvokeModelInput(), handle: handleNovaRequest},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
requestContext, cancelRequest := context.WithCancel(context.Background())
|
||||
t.Cleanup(cancelRequest)
|
||||
|
||||
upstreamContexts := make(chan context.Context, 1)
|
||||
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
|
||||
upstreamContexts <- request.Context()
|
||||
<-request.Context().Done()
|
||||
return nil, request.Context().Err()
|
||||
}))
|
||||
adaptor := &Adaptor{AwsClient: client, AwsReq: test.request}
|
||||
c := newAwsTestContext(httptest.NewRecorder(), requestContext)
|
||||
info := newAwsTestRelayInfo()
|
||||
|
||||
type handlerResult struct {
|
||||
err *relaytypes.NewAPIError
|
||||
usage *dto.Usage
|
||||
}
|
||||
results := make(chan handlerResult, 1)
|
||||
go func() {
|
||||
err, usage := test.handle(c, info, adaptor)
|
||||
results <- handlerResult{err: err, usage: usage}
|
||||
}()
|
||||
|
||||
var upstreamContext context.Context
|
||||
select {
|
||||
case upstreamContext = <-upstreamContexts:
|
||||
case result := <-results:
|
||||
t.Fatalf("handler returned before issuing AWS request: %v", result.err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("AWS request did not start")
|
||||
}
|
||||
|
||||
cancelRequest()
|
||||
|
||||
var result handlerResult
|
||||
select {
|
||||
case result = <-results:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("handler did not stop after client cancellation")
|
||||
}
|
||||
|
||||
require.ErrorIs(t, upstreamContext.Err(), context.Canceled)
|
||||
require.NotNil(t, result.err)
|
||||
assert.True(t, relaytypes.IsSkipRetryError(result.err))
|
||||
assert.Nil(t, result.usage)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAwsStreamHandlerUsesFinalUpstreamUsage(t *testing.T) {
|
||||
originalRelayTimeout := common.RelayTimeout
|
||||
common.RelayTimeout = 0
|
||||
t.Cleanup(func() {
|
||||
common.RelayTimeout = originalRelayTimeout
|
||||
})
|
||||
|
||||
events := []string{
|
||||
`{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`,
|
||||
`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
|
||||
`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`,
|
||||
`{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`,
|
||||
`{"type":"message_stop"}`,
|
||||
}
|
||||
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
|
||||
var body bytes.Buffer
|
||||
for _, event := range events {
|
||||
if err := writeAwsStreamEvent(&body, event); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return newAwsStreamResponse(request, io.NopCloser(bytes.NewReader(body.Bytes()))), nil
|
||||
}))
|
||||
adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()}
|
||||
recorder := httptest.NewRecorder()
|
||||
c := newAwsTestContext(recorder, context.Background())
|
||||
|
||||
handlerErr, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor)
|
||||
|
||||
require.Nil(t, handlerErr)
|
||||
require.NotNil(t, usage)
|
||||
require.NotNil(t, usage.BillingUsage)
|
||||
require.NotNil(t, usage.BillingUsage.ClaudeUsage)
|
||||
assert.Equal(t, 100, usage.BillingUsage.ClaudeUsage.InputTokens)
|
||||
assert.Equal(t, 423, usage.BillingUsage.ClaudeUsage.OutputTokens)
|
||||
assert.Contains(t, recorder.Body.String(), "[DONE]")
|
||||
}
|
||||
|
||||
func TestAwsStreamHandlerStopsAtClientCancellationAndKeepsPartialBillingUsage(t *testing.T) {
|
||||
originalRelayTimeout := common.RelayTimeout
|
||||
common.RelayTimeout = 0
|
||||
t.Cleanup(func() {
|
||||
common.RelayTimeout = originalRelayTimeout
|
||||
})
|
||||
|
||||
requestContext, cancelRequest := context.WithCancel(context.Background())
|
||||
t.Cleanup(cancelRequest)
|
||||
releaseFinal := make(chan struct{})
|
||||
var releaseFinalOnce sync.Once
|
||||
release := func() {
|
||||
releaseFinalOnce.Do(func() {
|
||||
close(releaseFinal)
|
||||
})
|
||||
}
|
||||
t.Cleanup(release)
|
||||
|
||||
producerResults := make(chan error, 1)
|
||||
upstreamContexts := make(chan context.Context, 1)
|
||||
client := newAwsTestClient(awsHTTPClientFunc(func(request *http.Request) (*http.Response, error) {
|
||||
upstreamContexts <- request.Context()
|
||||
reader, writer := io.Pipe()
|
||||
go func() {
|
||||
defer writer.Close()
|
||||
initialEvents := []string{
|
||||
`{"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":100,"output_tokens":1}}}`,
|
||||
`{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
|
||||
`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}`,
|
||||
}
|
||||
for _, event := range initialEvents {
|
||||
if err := writeAwsStreamEvent(writer, event); err != nil {
|
||||
producerResults <- err
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
<-releaseFinal
|
||||
producerResults <- writeAwsStreamEvent(writer, `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":423}}`)
|
||||
}()
|
||||
return newAwsStreamResponse(request, reader), nil
|
||||
}))
|
||||
|
||||
responseWriter := newAwsNotifyingResponseWriter("partial")
|
||||
c := newAwsTestContext(responseWriter, requestContext)
|
||||
adaptor := &Adaptor{AwsClient: client, AwsReq: newAwsStreamInput()}
|
||||
|
||||
type handlerResult struct {
|
||||
err *relaytypes.NewAPIError
|
||||
usage *dto.Usage
|
||||
}
|
||||
results := make(chan handlerResult, 1)
|
||||
go func() {
|
||||
err, usage := awsStreamHandler(c, newAwsTestRelayInfo(), adaptor)
|
||||
results <- handlerResult{err: err, usage: usage}
|
||||
}()
|
||||
|
||||
var upstreamContext context.Context
|
||||
select {
|
||||
case upstreamContext = <-upstreamContexts:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("AWS stream request did not start")
|
||||
}
|
||||
|
||||
var bodyLengthBeforeCancel int
|
||||
select {
|
||||
case bodyLengthBeforeCancel = <-responseWriter.notified:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("partial response was not written")
|
||||
}
|
||||
cancelRequest()
|
||||
|
||||
var result handlerResult
|
||||
select {
|
||||
case result = <-results:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("stream handler did not stop after client cancellation")
|
||||
}
|
||||
|
||||
require.ErrorIs(t, upstreamContext.Err(), context.Canceled)
|
||||
require.Nil(t, result.err)
|
||||
require.NotNil(t, result.usage)
|
||||
require.NotNil(t, result.usage.BillingUsage)
|
||||
require.NotNil(t, result.usage.BillingUsage.ClaudeUsage)
|
||||
assert.Equal(t, dto.BillingUsageSourceClaudeMessages, result.usage.BillingUsage.Source)
|
||||
assert.Equal(t, dto.BillingUsageSemanticAnthropic, result.usage.BillingUsage.Semantic)
|
||||
assert.Equal(t, 100, result.usage.BillingUsage.ClaudeUsage.InputTokens)
|
||||
assert.Equal(t, 1, result.usage.BillingUsage.ClaudeUsage.OutputTokens)
|
||||
assert.Equal(t, bodyLengthBeforeCancel, responseWriter.Body.Len())
|
||||
assert.NotContains(t, responseWriter.Body.String(), "[DONE]")
|
||||
|
||||
release()
|
||||
select {
|
||||
case producerErr := <-producerResults:
|
||||
require.Error(t, producerErr)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("upstream producer did not observe the closed stream")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user