重构前
This commit is contained in:
@@ -0,0 +1,288 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
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/relayconvert"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func stopReasonClaude2OpenAI(reason string) string {
|
||||
return relayconvert.StopReasonClaudeToOpenAI(reason)
|
||||
}
|
||||
|
||||
func maybeMarkClaudeRefusal(c *gin.Context, stopReason string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
if strings.EqualFold(stopReason, "refusal") {
|
||||
common.SetContextKey(c, constant.ContextKeyAdminRejectReason, "claude_stop_reason=refusal")
|
||||
}
|
||||
}
|
||||
|
||||
func StreamResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.ChatCompletionsStreamResponse {
|
||||
return relayconvert.StreamResponseClaude2OpenAI(claudeResponse)
|
||||
}
|
||||
|
||||
func ResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.OpenAITextResponse {
|
||||
return relayconvert.ResponseClaude2OpenAI(claudeResponse)
|
||||
}
|
||||
|
||||
type ClaudeResponseInfo = relayconvert.ClaudeResponseInfo
|
||||
|
||||
func cacheCreationTokensForOpenAIUsage(usage *dto.Usage) int {
|
||||
if usage == nil {
|
||||
return 0
|
||||
}
|
||||
openAIUsage := relayconvert.UsageFromClaudeUsage(usage)
|
||||
if openAIUsage == nil {
|
||||
return 0
|
||||
}
|
||||
return openAIUsage.PromptTokens - usage.PromptTokens - usage.PromptTokensDetails.CachedTokens
|
||||
}
|
||||
|
||||
func buildOpenAIStyleUsageFromClaudeUsage(usage *dto.Usage) dto.Usage {
|
||||
mapped := relayconvert.UsageFromClaudeUsage(usage)
|
||||
if mapped == nil {
|
||||
return dto.Usage{}
|
||||
}
|
||||
return *mapped
|
||||
}
|
||||
|
||||
func buildMessageDeltaPatchUsage(claudeResponse *dto.ClaudeResponse, claudeInfo *ClaudeResponseInfo) *dto.ClaudeUsage {
|
||||
return relayconvert.BuildMessageDeltaPatchUsage(claudeResponse, claudeInfo)
|
||||
}
|
||||
|
||||
func shouldSkipClaudeMessageDeltaUsagePatch(info *relaycommon.RelayInfo) bool {
|
||||
if model_setting.GetGlobalSettings().PassThroughRequestEnabled {
|
||||
return true
|
||||
}
|
||||
if info == nil {
|
||||
return false
|
||||
}
|
||||
return info.ChannelSetting.PassThroughBodyEnabled
|
||||
}
|
||||
|
||||
func patchClaudeMessageDeltaUsageData(data string, usage *dto.ClaudeUsage) string {
|
||||
return relayconvert.PatchClaudeMessageDeltaUsageData(data, usage)
|
||||
}
|
||||
|
||||
func FormatClaudeResponseInfo(claudeResponse *dto.ClaudeResponse, oaiResponse *dto.ChatCompletionsStreamResponse, claudeInfo *ClaudeResponseInfo) bool {
|
||||
return relayconvert.FormatClaudeResponseInfo(claudeResponse, oaiResponse, claudeInfo)
|
||||
}
|
||||
|
||||
func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo, data string) *types.NewAPIError {
|
||||
var claudeResponse dto.ClaudeResponse
|
||||
err := common.UnmarshalJsonStr(data, &claudeResponse)
|
||||
if err != nil {
|
||||
common.SysLog("error unmarshalling stream response: " + err.Error())
|
||||
return types.NewError(err, types.ErrorCodeBadResponseBody)
|
||||
}
|
||||
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
|
||||
return types.WithClaudeError(*claudeError, http.StatusInternalServerError)
|
||||
}
|
||||
if claudeResponse.StopReason != "" {
|
||||
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
|
||||
}
|
||||
if claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil {
|
||||
maybeMarkClaudeRefusal(c, *claudeResponse.Delta.StopReason)
|
||||
}
|
||||
if info.RelayFormat == types.RelayFormatClaude {
|
||||
FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo)
|
||||
|
||||
if claudeResponse.Type == "message_start" {
|
||||
// message_start, 获取usage
|
||||
if claudeResponse.Message != nil {
|
||||
info.UpstreamModelName = claudeResponse.Message.Model
|
||||
}
|
||||
} else if claudeResponse.Type == "message_delta" {
|
||||
// 确保 message_delta 的 usage 包含完整的 input_tokens 和 cache 相关字段
|
||||
// 解决 AWS Bedrock 等上游返回的 message_delta 缺少这些字段的问题
|
||||
if !shouldSkipClaudeMessageDeltaUsagePatch(info) {
|
||||
data = patchClaudeMessageDeltaUsageData(data, buildMessageDeltaPatchUsage(&claudeResponse, claudeInfo))
|
||||
}
|
||||
}
|
||||
countClaudeStreamBillableTools(c, info, &claudeResponse)
|
||||
helper.ClaudeChunkData(c, claudeResponse, data)
|
||||
} else if info.RelayFormat == types.RelayFormatOpenAI {
|
||||
response := StreamResponseClaude2OpenAI(&claudeResponse)
|
||||
|
||||
if !FormatClaudeResponseInfo(&claudeResponse, response, claudeInfo) {
|
||||
return nil
|
||||
}
|
||||
|
||||
countClaudeStreamBillableTools(c, info, &claudeResponse)
|
||||
|
||||
err = helper.ObjectData(c, response)
|
||||
if err != nil {
|
||||
logger.LogError(c, "send_stream_response_failed: "+err.Error())
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func countClaudeStreamBillableTools(c *gin.Context, info *relaycommon.RelayInfo, claudeResponse *dto.ClaudeResponse) {
|
||||
if claudeResponse == nil {
|
||||
return
|
||||
}
|
||||
if claudeResponse.Type == "content_block_start" &&
|
||||
claudeResponse.ContentBlock != nil &&
|
||||
claudeResponse.ContentBlock.Type == "tool_use" {
|
||||
info.CountBillableToolCall(dto.BuildInCallToolUse, claudeResponse.ContentBlock.Name)
|
||||
}
|
||||
if claudeResponse.Type == "message_delta" &&
|
||||
claudeResponse.Usage != nil &&
|
||||
claudeResponse.Usage.ServerToolUse != nil &&
|
||||
claudeResponse.Usage.ServerToolUse.WebSearchRequests > 0 {
|
||||
c.Set("claude_web_search_requests", claudeResponse.Usage.ServerToolUse.WebSearchRequests)
|
||||
}
|
||||
}
|
||||
|
||||
func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo) {
|
||||
if claudeInfo.Usage.PromptTokens == 0 {
|
||||
//上游出错
|
||||
}
|
||||
if claudeInfo.Usage.CompletionTokens == 0 || !claudeInfo.Done {
|
||||
if common.DebugEnabled {
|
||||
common.SysLog("claude response usage is not complete, maybe upstream error")
|
||||
}
|
||||
// 只补缺失字段，不整份覆盖——保留 message_start 已拿到的 cache 字段
|
||||
fallback := service.ResponseText2Usage(c, claudeInfo.ResponseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens())
|
||||
if claudeInfo.Usage.CompletionTokens == 0 ||
|
||||
(!claudeInfo.Done && fallback.CompletionTokens > claudeInfo.Usage.CompletionTokens) {
|
||||
claudeInfo.Usage.CompletionTokens = fallback.CompletionTokens
|
||||
}
|
||||
if claudeInfo.Usage.PromptTokens == 0 {
|
||||
claudeInfo.Usage.PromptTokens = fallback.PromptTokens
|
||||
}
|
||||
claudeInfo.Usage.TotalTokens = claudeInfo.Usage.PromptTokens + claudeInfo.Usage.CompletionTokens
|
||||
}
|
||||
if claudeInfo.Usage != nil {
|
||||
claudeInfo.Usage.UsageSemantic = "anthropic"
|
||||
}
|
||||
if claudeInfo.Usage != nil && claudeInfo.Usage.BillingUsage == nil {
|
||||
claudeInfo.Usage.BillingUsage = dto.NewClaudeMessagesBillingUsage(buildMessageDeltaPatchUsage(nil, claudeInfo))
|
||||
}
|
||||
|
||||
if info.RelayFormat == types.RelayFormatClaude {
|
||||
//
|
||||
} else if info.RelayFormat == types.RelayFormatOpenAI {
|
||||
if info.ShouldIncludeUsage {
|
||||
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(claudeInfo.Usage)
|
||||
response := helper.GenerateFinalUsageResponse(claudeInfo.ResponseId, claudeInfo.Created, info.UpstreamModelName, openAIUsage)
|
||||
err := helper.ObjectData(c, response)
|
||||
if err != nil {
|
||||
common.SysLog("send final response failed: " + err.Error())
|
||||
}
|
||||
}
|
||||
helper.Done(c)
|
||||
}
|
||||
}
|
||||
|
||||
func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
|
||||
claudeInfo := &ClaudeResponseInfo{
|
||||
ResponseId: helper.GetResponseID(c),
|
||||
Created: common.GetTimestamp(),
|
||||
Model: info.UpstreamModelName,
|
||||
ResponseText: strings.Builder{},
|
||||
Usage: &dto.Usage{},
|
||||
}
|
||||
var err *types.NewAPIError
|
||||
helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) {
|
||||
err = HandleStreamResponseData(c, info, claudeInfo, data)
|
||||
if err != nil {
|
||||
sr.Stop(err)
|
||||
}
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
HandleStreamFinalResponse(c, info, claudeInfo)
|
||||
return claudeInfo.Usage, nil
|
||||
}
|
||||
|
||||
func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo, httpResp *http.Response, data []byte) *types.NewAPIError {
|
||||
var claudeResponse dto.ClaudeResponse
|
||||
err := common.Unmarshal(data, &claudeResponse)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeBadResponseBody)
|
||||
}
|
||||
if claudeError := claudeResponse.GetClaudeError(); claudeError != nil && claudeError.Type != "" {
|
||||
return types.WithClaudeError(*claudeError, http.StatusInternalServerError)
|
||||
}
|
||||
maybeMarkClaudeRefusal(c, claudeResponse.StopReason)
|
||||
if claudeInfo.Usage == nil {
|
||||
claudeInfo.Usage = &dto.Usage{}
|
||||
}
|
||||
if claudeResponse.Usage != nil {
|
||||
claudeInfo.Usage.PromptTokens = claudeResponse.Usage.InputTokens
|
||||
claudeInfo.Usage.CompletionTokens = claudeResponse.Usage.OutputTokens
|
||||
claudeInfo.Usage.TotalTokens = claudeResponse.Usage.InputTokens + claudeResponse.Usage.OutputTokens
|
||||
claudeInfo.Usage.UsageSemantic = "anthropic"
|
||||
claudeInfo.Usage.BillingUsage = dto.NewClaudeMessagesBillingUsage(claudeResponse.Usage)
|
||||
claudeInfo.Usage.PromptTokensDetails.CachedTokens = claudeResponse.Usage.CacheReadInputTokens
|
||||
claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens = claudeResponse.Usage.CacheCreationInputTokens
|
||||
claudeInfo.Usage.ClaudeCacheCreation5mTokens = claudeResponse.Usage.GetCacheCreation5mTokens()
|
||||
claudeInfo.Usage.ClaudeCacheCreation1hTokens = claudeResponse.Usage.GetCacheCreation1hTokens()
|
||||
}
|
||||
var responseData []byte
|
||||
switch info.RelayFormat {
|
||||
case types.RelayFormatOpenAI:
|
||||
openaiResponse := ResponseClaude2OpenAI(&claudeResponse)
|
||||
openaiResponse.Usage = buildOpenAIStyleUsageFromClaudeUsage(claudeInfo.Usage)
|
||||
responseData, err = common.Marshal(openaiResponse)
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeBadResponseBody)
|
||||
}
|
||||
case types.RelayFormatClaude:
|
||||
responseData = data
|
||||
}
|
||||
|
||||
if claudeResponse.Usage != nil && claudeResponse.Usage.ServerToolUse != nil && claudeResponse.Usage.ServerToolUse.WebSearchRequests > 0 {
|
||||
c.Set("claude_web_search_requests", claudeResponse.Usage.ServerToolUse.WebSearchRequests)
|
||||
}
|
||||
|
||||
for _, block := range claudeResponse.Content {
|
||||
if block.Type == "tool_use" {
|
||||
info.CountBillableToolCall(dto.BuildInCallToolUse, block.Name)
|
||||
}
|
||||
}
|
||||
|
||||
service.IOCopyBytesGracefully(c, httpResp, responseData)
|
||||
return nil
|
||||
}
|
||||
|
||||
func ClaudeHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) {
|
||||
defer service.CloseResponseBodyGracefully(resp)
|
||||
|
||||
claudeInfo := &ClaudeResponseInfo{
|
||||
ResponseId: helper.GetResponseID(c),
|
||||
Created: common.GetTimestamp(),
|
||||
Model: info.UpstreamModelName,
|
||||
ResponseText: strings.Builder{},
|
||||
Usage: &dto.Usage{},
|
||||
}
|
||||
responseBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, types.NewError(err, types.ErrorCodeBadResponseBody)
|
||||
}
|
||||
logger.LogDebug(c, "responseBody: %s", responseBody)
|
||||
handleErr := HandleClaudeResponseData(c, info, claudeInfo, resp, responseBody)
|
||||
if handleErr != nil {
|
||||
return nil, handleErr
|
||||
}
|
||||
return claudeInfo.Usage, nil
|
||||
}
|
||||
Reference in New Issue
Block a user