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 {
return http.StatusInternalServerError
}
func newAwsInvokeContext() (context.Context, context.CancelFunc) {
func newAwsInvokeContext(parent context.Context) (context.Context, context.CancelFunc) {
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) {
......@@ -215,13 +228,13 @@ func getAwsModelID(requestModel string) string {
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()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
if err != nil {
statusCode := getAwsErrorStatusCode(err)
return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil
return newAwsInvokeError(requestContext, err, "InvokeModel"), nil
}
claudeInfo := &claude.ClaudeResponseInfo{
......@@ -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) {
ctx, cancel := newAwsInvokeContext()
requestContext := c.Request.Context()
ctx, cancel := newAwsInvokeContext(requestContext)
defer cancel()
awsResp, err := a.AwsClient.InvokeModelWithResponseStream(ctx, a.AwsReq.(*bedrockruntime.InvokeModelWithResponseStreamInput))
if err != nil {
statusCode := getAwsErrorStatusCode(err)
return types.NewOpenAIError(errors.Wrap(err, "InvokeModelWithResponseStream"), types.ErrorCodeAwsInvokeError, statusCode), nil
return newAwsInvokeError(requestContext, err, "InvokeModelWithResponseStream"), nil
}
stream := awsResp.GetStream()
defer stream.Close()
......@@ -264,7 +277,20 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (
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) {
case *bedrockruntimeTypes.ResponseStreamMemberChunk:
info.SetFirstResponseTime()
......@@ -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
}
}
}
_ = stream.Close()
claude.HandleStreamFinalResponse(c, info, claudeInfo)
return nil, claudeInfo.Usage
}
......@@ -288,13 +316,13 @@ func awsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor) (
// Nova模型处理函数
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()
awsResp, err := a.AwsClient.InvokeModel(ctx, a.AwsReq.(*bedrockruntime.InvokeModelInput))
if err != nil {
statusCode := getAwsErrorStatusCode(err)
return types.NewOpenAIError(errors.Wrap(err, "InvokeModel"), types.ErrorCodeAwsInvokeError, statusCode), nil
return newAwsInvokeError(requestContext, err, "InvokeModel"), nil
}
// 解析Nova响应
......
service/billing_usage.go
......@@ -25,36 +25,32 @@ func effectiveBillingUsage(usage *dto.Usage) *dto.Usage {
}
func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string {
effectiveUsage, ok := usageFromBillingUsage(usage)
if !ok {
if isLocalCountTokens {
return usageBillingPathLocal
}
if usage == nil || usage.BillingUsage == nil {
return usageBillingPathUpstream
}
source := strings.TrimSpace(usage.BillingUsage.Source)
semantic := strings.TrimSpace(usage.BillingUsage.Semantic)
if strings.EqualFold(source, dto.BillingUsageSourceOAIChat) ||
strings.EqualFold(source, dto.BillingUsageSourceOAIResponses) ||
strings.EqualFold(semantic, dto.BillingUsageSemanticOpenAI) {
switch effectiveUsage.UsageSemantic {
case dto.BillingUsageSemanticOpenAI:
if usage.BillingUsage.Estimated {
return usageBillingPathOpenAIEstimated
}
return usageBillingPathOpenAI
}
if strings.EqualFold(source, dto.BillingUsageSourceClaudeMessages) ||
strings.EqualFold(semantic, dto.BillingUsageSemanticAnthropic) {
case dto.BillingUsageSemanticAnthropic:
if usage.BillingUsage.Estimated {
return usageBillingPathAnthropicEstimated
}
return usageBillingPathAnthropic
}
if strings.EqualFold(source, dto.BillingUsageSourceGeminiChat) ||
strings.EqualFold(semantic, dto.BillingUsageSemanticGemini) {
case dto.BillingUsageSemanticGemini:
if usage.BillingUsage.Estimated {
return usageBillingPathGeminiEstimated
}
return usageBillingPathGemini
}
return usageBillingPathUpstream
}
......
service/text_quota_test.go
......@@ -284,9 +284,18 @@ func TestCalculateTextQuotaSummaryUsesOpenAIBillingUsageBeforeTopLevelUsage(t *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}),
}))
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, usageBillingPathOpenAI, usageBillingPathForLog(false, &dto.Usage{
BillingUsage: dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: 1}),
......@@ -297,7 +306,7 @@ func TestUsageBillingPathForLog(t *testing.T) {
require.Equal(t, usageBillingPathGemini, usageBillingPathForLog(false, &dto.Usage{
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}),
}))
}
......@@ -306,7 +315,7 @@ func TestAppendUsageBillingPathForLogWritesAdminInfo(t *testing.T) {
other := 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}),
})
......
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