Skip to content

phsl / new-api

Go to a project
Commit bd585d78 by Calcium-Ion Committed by GitHub
Options

fix(aws): cancel Bedrock requests on client disconnect (#6589) 

* fix(aws): cancel Bedrock requests on client disconnect

* fix(billing): log effective usage billing path
parent cfaba1dd
Showing with 61 additions and 28 deletions
relay/channel/aws/relay-aws.go
...@@ -40,11 +40,24 @@ func getAwsErrorStatusCode(err error) int { ...@@ -40,11 +40,24 @@ func getAwsErrorStatusCode(err error) int {
return http.StatusInternalServerError return http.StatusInternalServerError
} }
func newAwsInvokeContext() (context.Context, context.CancelFunc) { func newAwsInvokeContext(parent context.Context) (context.Context, context.CancelFunc) {
if common.RelayTimeout <= 0 { if common.RelayTimeout <= 0 {
return context.Background(), func() {} return context.WithCancel(parent)
} }
return context.WithTimeout(context.Background(), time.Duration(common.RelayTimeout)*time.Second) 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) { func newAwsClient(c *gin.Context, info *relaycommon.RelayInfo) (*bedrockruntime.Client, error) {
...@@ -215,13 +228,13 @@ func getAwsModelID(requestModel string) string { ...@@ -215,13 +228,13 @@ func getAwsModelID(requestModel string) string {
func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
ctx, cancel := newAwsInvokeContext() requestContext := c.Request.Context()
ctx, cancel := newAwsInvokeContext(requestContext)
defer cancel() defer cancel()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput)) awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
if err != nil { if err != nil {
statusCode := getAwsErrorStatusCode(err) return newAwsInvokeError(requestContext, err, "InvokeModel"), nil
return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil
} }
claudeInfo := &claude.ClaudeResponseInfo{ claudeInfo := &claude.ClaudeResponseInfo{
...@@ -245,13 +258,13 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types ...@@ -245,13 +258,13 @@ func awsHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types
} }
func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
ctx, cancel := newAwsInvokeContext() requestContext := c.Request.Context()
ctx, cancel := newAwsInvokeContext(requestContext)
defer cancel() defer cancel()
awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput)) awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput))
if err != nil { if err != nil {
statusCode := getAwsErrorStatusCode(err) return newAwsInvokeError(requestContext, err, "InvokeModelWithResponseStream"), nil
return types.NewOpenAIError(errors.Wrap(err, "InvokeModelWithResponseStream"), types.ErrorCodeAwsInvokeError, statusCode), nil
} }
stream := awsResp.GetStream() stream := awsResp.GetStream()
defer stream.Close() defer stream.Close()
...@@ -264,7 +277,20 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ( ...@@ -264,7 +277,20 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (
Usage: &dto.Usage{}, Usage: &dto.Usage{},
} }
for event := range stream.Events() { 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) { switch v := event.(type) {
case *bedrockruntimeTypes.ResponseStreamMemberChunk: case *bedrockruntimeTypes.ResponseStreamMemberChunk:
info.SetFirstResponseTime() info.SetFirstResponseTime()
...@@ -280,7 +306,9 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ( ...@@ -280,7 +306,9 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (
return types.NewError(errors.New("nil or unknown response type"), types.ErrorCodeInvalidRequest), nil return types.NewError(errors.New("nil or unknown response type"), types.ErrorCodeInvalidRequest), nil
} }
} }
}
_ = stream.Close()
claude.HandleStreamFinalResponse(c, info, claudeInfo) claude.HandleStreamFinalResponse(c, info, claudeInfo)
return nil, claudeInfo.Usage return nil, claudeInfo.Usage
} }
...@@ -288,13 +316,13 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) ( ...@@ -288,13 +316,13 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (
// Nova模型处理函数 // Nova模型处理函数
func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) { func handleNovaRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (*types.NewAPIError, *dto.Usage) {
ctx, cancel := newAwsInvokeContext() requestContext := c.Request.Context()
ctx, cancel := newAwsInvokeContext(requestContext)
defer cancel() defer cancel()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput)) awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
if err != nil { if err != nil {
statusCode := getAwsErrorStatusCode(err) return newAwsInvokeError(requestContext, err, "InvokeModel"), nil
return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil
} }
// 解析Nova响应 // 解析Nova响应
......
service/billing_usage.go
...@@ -25,36 +25,32 @@ func effectiveBillingUsage(usage *dto.Usage) *dto.Usage { ...@@ -25,36 +25,32 @@ func effectiveBillingUsage(usage *dto.Usage) *dto.Usage {
} }
func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string { func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string {
effectiveUsage, ok := usageFromBillingUsage(usage)
if !ok {
if isLocalCountTokens { if isLocalCountTokens {
return usageBillingPathLocal return usageBillingPathLocal
} }
if usage == nil || usage.BillingUsage == nil {
return usageBillingPathUpstream return usageBillingPathUpstream
} }
source := strings.TrimSpace(usage.BillingUsage.Source)
semantic := strings.TrimSpace(usage.BillingUsage.Semantic) switch effectiveUsage.UsageSemantic {
if strings.EqualFold(source, dto.BillingUsageSourceOAIChat) || case dto.BillingUsageSemanticOpenAI:
strings.EqualFold(source, dto.BillingUsageSourceOAIResponses) ||
strings.EqualFold(semantic, dto.BillingUsageSemanticOpenAI) {
if usage.BillingUsage.Estimated { if usage.BillingUsage.Estimated {
return usageBillingPathOpenAIEstimated return usageBillingPathOpenAIEstimated
} }
return usageBillingPathOpenAI return usageBillingPathOpenAI
} case dto.BillingUsageSemanticAnthropic:
if strings.EqualFold(source, dto.BillingUsageSourceClaudeMessages) ||
strings.EqualFold(semantic, dto.BillingUsageSemanticAnthropic) {
if usage.BillingUsage.Estimated { if usage.BillingUsage.Estimated {
return usageBillingPathAnthropicEstimated return usageBillingPathAnthropicEstimated
} }
return usageBillingPathAnthropic return usageBillingPathAnthropic
} case dto.BillingUsageSemanticGemini:
if strings.EqualFold(source, dto.BillingUsageSourceGeminiChat) ||
strings.EqualFold(semantic, dto.BillingUsageSemanticGemini) {
if usage.BillingUsage.Estimated { if usage.BillingUsage.Estimated {
return usageBillingPathGeminiEstimated return usageBillingPathGeminiEstimated
} }
return usageBillingPathGemini return usageBillingPathGemini
} }
return usageBillingPathUpstream return usageBillingPathUpstream
} }
......
service/text_quota_test.go
...@@ -284,9 +284,18 @@ func TestCalculateTextQuotaSummaryUsesOpenAIBillingUsageBeforeTopLevelUsage(t *t ...@@ -284,9 +284,18 @@ func TestCalculateTextQuotaSummaryUsesOpenAIBillingUsageBeforeTopLevelUsage(t *t
} }
func TestUsageBillingPathForLog(t *testing.T) { func TestUsageBillingPathForLog(t *testing.T) {
require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, &dto.Usage{ require.Equal(t, usageBillingPathAnthropic, usageBillingPathForLog(true, &dto.Usage{
BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}),
})) }))
invalidBillingUsage := &dto.Usage{
PromptTokens: 1,
BillingUsage: &dto.BillingUsage{
Source: dto.BillingUsageSourceClaudeMessages,
Semantic: dto.BillingUsageSemanticAnthropic,
},
}
require.Equal(t, usageBillingPathLocal, usageBillingPathForLog(true, invalidBillingUsage))
require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, invalidBillingUsage))
require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, &dto.Usage{})) require.Equal(t, usageBillingPathUpstream, usageBillingPathForLog(false, &dto.Usage{}))
require.Equal(t, usageBillingPathOpenAI, usageBillingPathForLog(false, &dto.Usage{ require.Equal(t, usageBillingPathOpenAI, usageBillingPathForLog(false, &dto.Usage{
BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: 1}), BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: 1}),
...@@ -297,7 +306,7 @@ func TestUsageBillingPathForLog(t *testing.T) { ...@@ -297,7 +306,7 @@ func TestUsageBillingPathForLog(t *testing.T) {
require.Equal(t, usageBillingPathGemini, usageBillingPathForLog(false, &dto.Usage{ require.Equal(t, usageBillingPathGemini, usageBillingPathForLog(false, &dto.Usage{
BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{PromptTokenCount: 1}), BillingUsage: dto.NewGeminiChatBillingUsage(&dto.GeminiUsageMetadata{PromptTokenCount: 1}),
})) }))
require.Equal(t, usageBillingPathGeminiEstimated, usageBillingPathForLog(false, &dto.Usage{ require.Equal(t, usageBillingPathGeminiEstimated, usageBillingPathForLog(true, &dto.Usage{
BillingUsage: dto.NewEstimatedGeminiChatBillingUsage(&dto.Usage{PromptTokens: 1}), BillingUsage: dto.NewEstimatedGeminiChatBillingUsage(&dto.Usage{PromptTokens: 1}),
})) }))
} }
...@@ -306,7 +315,7 @@ func TestAppendUsageBillingPathForLogWritesAdminInfo(t *testing.T) { ...@@ -306,7 +315,7 @@ func TestAppendUsageBillingPathForLogWritesAdminInfo(t *testing.T) {
other := map[string]interface{}{ other := map[string]interface{}{
"admin_info": map[string]interface{}{}, "admin_info": map[string]interface{}{},
} }
appendUsageBillingPathForLog(other, false, &dto.Usage{ appendUsageBillingPathForLog(other, true, &dto.Usage{
BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}), BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{InputTokens: 1}),
}) })
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment