feat: 算力引擎 Go 后端核心(new-api)
- 统一 OpenAI 兼容 /v1 中继 + 渠道/额度/令牌/流水 - 各领域包:relay 模型网关、model 数据层、controller 管理 API
This commit is contained in:
@@ -0,0 +1,139 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"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/setting/model_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type Adaptor struct {
|
||||
}
|
||||
|
||||
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) {
|
||||
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) {
|
||||
requestURL := fmt.Sprintf("%s/v1/messages", info.ChannelBaseUrl)
|
||||
if !shouldAppendClaudeBetaQuery(info) {
|
||||
return requestURL, nil
|
||||
}
|
||||
|
||||
parsedURL, err := url.Parse(requestURL)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
query := parsedURL.Query()
|
||||
query.Set("beta", "true")
|
||||
parsedURL.RawQuery = query.Encode()
|
||||
return parsedURL.String(), nil
|
||||
}
|
||||
|
||||
func shouldAppendClaudeBetaQuery(info *relaycommon.RelayInfo) bool {
|
||||
if info == nil {
|
||||
return false
|
||||
}
|
||||
if info.IsClaudeBetaQuery {
|
||||
return true
|
||||
}
|
||||
if info.ChannelOtherSettings.ClaudeBetaQuery {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func CommonClaudeHeadersOperation(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) {
|
||||
// common headers operation
|
||||
anthropicBeta := c.Request.Header.Get("anthropic-beta")
|
||||
if anthropicBeta != "" {
|
||||
req.Set("anthropic-beta", anthropicBeta)
|
||||
}
|
||||
model_setting.GetClaudeSettings().WriteHeaders(info.OriginModelName, req)
|
||||
}
|
||||
|
||||
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *relaycommon.RelayInfo) error {
|
||||
channel.SetupApiRequestHeader(info, c, req)
|
||||
req.Set("x-api-key", info.ApiKey)
|
||||
anthropicVersion := c.Request.Header.Get("anthropic-version")
|
||||
if anthropicVersion == "" {
|
||||
anthropicVersion = "2023-06-01"
|
||||
}
|
||||
req.Set("anthropic-version", anthropicVersion)
|
||||
CommonClaudeHeadersOperation(c, req, info)
|
||||
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")
|
||||
}
|
||||
result, err := relayconvert.ConvertRequest(c, info, types.RelayFormatClaude, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result.Value, nil
|
||||
}
|
||||
|
||||
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) {
|
||||
return channel.DoApiRequest(a, c, info, requestBody)
|
||||
}
|
||||
|
||||
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) {
|
||||
info.FinalRequestRelayFormat = types.RelayFormatClaude
|
||||
if info.IsStream {
|
||||
return ClaudeStreamHandler(c, resp, info)
|
||||
} else {
|
||||
return ClaudeHandler(c, resp, info)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetModelList() []string {
|
||||
return ModelList
|
||||
}
|
||||
|
||||
func (a *Adaptor) GetChannelName() string {
|
||||
return ChannelName
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package claude
|
||||
|
||||
var ModelList = []string{
|
||||
"claude-3-sonnet-20240229",
|
||||
"claude-3-opus-20240229",
|
||||
"claude-3-haiku-20240307",
|
||||
"claude-3-5-haiku-20241022",
|
||||
"claude-haiku-4-5-20251001",
|
||||
"claude-3-5-sonnet-20240620",
|
||||
"claude-3-5-sonnet-20241022",
|
||||
"claude-3-7-sonnet-20250219",
|
||||
"claude-3-7-sonnet-20250219-thinking",
|
||||
"claude-sonnet-4-20250514",
|
||||
"claude-sonnet-4-20250514-thinking",
|
||||
"claude-opus-4-20250514",
|
||||
"claude-opus-4-20250514-thinking",
|
||||
"claude-opus-4-1-20250805",
|
||||
"claude-opus-4-1-20250805-thinking",
|
||||
"claude-sonnet-4-5-20250929",
|
||||
"claude-sonnet-4-5-20250929-thinking",
|
||||
"claude-opus-4-5-20251101",
|
||||
"claude-opus-4-5-20251101-thinking",
|
||||
"claude-opus-4-6",
|
||||
"claude-opus-4-6-max",
|
||||
"claude-opus-4-6-high",
|
||||
"claude-opus-4-6-medium",
|
||||
"claude-opus-4-6-low",
|
||||
"claude-sonnet-4-6",
|
||||
"claude-opus-4-7",
|
||||
"claude-opus-4-7-max",
|
||||
"claude-opus-4-7-xhigh",
|
||||
"claude-opus-4-7-high",
|
||||
"claude-opus-4-7-medium",
|
||||
"claude-opus-4-7-low",
|
||||
"claude-opus-4-7-thinking",
|
||||
"claude-opus-4-8",
|
||||
"claude-opus-4-8-max",
|
||||
"claude-opus-4-8-xhigh",
|
||||
"claude-opus-4-8-high",
|
||||
"claude-opus-4-8-medium",
|
||||
"claude-opus-4-8-low",
|
||||
"claude-opus-4-8-thinking",
|
||||
}
|
||||
|
||||
var ChannelName = "claude"
|
||||
@@ -0,0 +1,95 @@
|
||||
package claude
|
||||
|
||||
//
|
||||
//type ClaudeMetadata struct {
|
||||
// UserId string `json:"user_id"`
|
||||
//}
|
||||
//
|
||||
//type ClaudeMediaMessage struct {
|
||||
// Type string `json:"type"`
|
||||
// Text string `json:"text,omitempty"`
|
||||
// Source *ClaudeMessageSource `json:"source,omitempty"`
|
||||
// Usage *ClaudeUsage `json:"usage,omitempty"`
|
||||
// StopReason *string `json:"stop_reason,omitempty"`
|
||||
// PartialJson string `json:"partial_json,omitempty"`
|
||||
// Thinking string `json:"thinking,omitempty"`
|
||||
// Signature string `json:"signature,omitempty"`
|
||||
// Delta string `json:"delta,omitempty"`
|
||||
// // tool_calls
|
||||
// Id string `json:"id,omitempty"`
|
||||
// Name string `json:"name,omitempty"`
|
||||
// Input any `json:"input,omitempty"`
|
||||
// Content string `json:"content,omitempty"`
|
||||
// ToolUseId string `json:"tool_use_id,omitempty"`
|
||||
//}
|
||||
//
|
||||
//type ClaudeMessageSource struct {
|
||||
// Type string `json:"type"`
|
||||
// MediaType string `json:"media_type"`
|
||||
// Data string `json:"data"`
|
||||
//}
|
||||
//
|
||||
//type ClaudeMessage struct {
|
||||
// Role string `json:"role"`
|
||||
// Content any `json:"content"`
|
||||
//}
|
||||
//
|
||||
//type Tool struct {
|
||||
// Name string `json:"name"`
|
||||
// Description string `json:"description,omitempty"`
|
||||
// InputSchema map[string]interface{} `json:"input_schema"`
|
||||
//}
|
||||
//
|
||||
//type InputSchema struct {
|
||||
// Type string `json:"type"`
|
||||
// Properties any `json:"properties,omitempty"`
|
||||
// Required any `json:"required,omitempty"`
|
||||
//}
|
||||
//
|
||||
//type ClaudeRequest struct {
|
||||
// Model string `json:"model"`
|
||||
// Prompt string `json:"prompt,omitempty"`
|
||||
// System string `json:"system,omitempty"`
|
||||
// Messages []ClaudeMessage `json:"messages,omitempty"`
|
||||
// MaxTokens uint `json:"max_tokens,omitempty"`
|
||||
// MaxTokensToSample uint `json:"max_tokens_to_sample,omitempty"`
|
||||
// StopSequences []string `json:"stop_sequences,omitempty"`
|
||||
// Temperature *float64 `json:"temperature,omitempty"`
|
||||
// TopP float64 `json:"top_p,omitempty"`
|
||||
// TopK int `json:"top_k,omitempty"`
|
||||
// //ClaudeMetadata `json:"metadata,omitempty"`
|
||||
// Stream bool `json:"stream,omitempty"`
|
||||
// Tools any `json:"tools,omitempty"`
|
||||
// ToolChoice any `json:"tool_choice,omitempty"`
|
||||
// Thinking *Thinking `json:"thinking,omitempty"`
|
||||
//}
|
||||
//
|
||||
//type Thinking struct {
|
||||
// Type string `json:"type"`
|
||||
// BudgetTokens int `json:"budget_tokens"`
|
||||
//}
|
||||
//
|
||||
//type ClaudeError struct {
|
||||
// Type string `json:"type"`
|
||||
// Message string `json:"message"`
|
||||
//}
|
||||
//
|
||||
//type ClaudeResponse struct {
|
||||
// Id string `json:"id"`
|
||||
// Type string `json:"type"`
|
||||
// Content []ClaudeMediaMessage `json:"content"`
|
||||
// Completion string `json:"completion"`
|
||||
// StopReason string `json:"stop_reason"`
|
||||
// Model string `json:"model"`
|
||||
// Error ClaudeError `json:"error"`
|
||||
// Usage ClaudeUsage `json:"usage"`
|
||||
// Index int `json:"index"` // stream only
|
||||
// ContentBlock *ClaudeMediaMessage `json:"content_block"`
|
||||
// Delta *ClaudeMediaMessage `json:"delta"` // stream only
|
||||
// Message *ClaudeResponse `json:"message"` // stream only: message_start
|
||||
//}
|
||||
//
|
||||
//type ClaudeUsage struct {
|
||||
// InputTokens int `json:"input_tokens"`
|
||||
// OutputTokens int `json:"output_tokens"`
|
||||
//}
|
||||
@@ -0,0 +1,129 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestPatchClaudeMessageDeltaUsageDataPreserveUnknownFields(t *testing.T) {
|
||||
originalData := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":53},"vendor_meta":{"trace_id":"trace_001"}}`
|
||||
usage := &dto.ClaudeUsage{
|
||||
InputTokens: 100,
|
||||
CacheReadInputTokens: 30,
|
||||
CacheCreationInputTokens: 50,
|
||||
}
|
||||
|
||||
patchedData := patchClaudeMessageDeltaUsageData(originalData, usage)
|
||||
|
||||
require.Equal(t, "message_delta", gjson.Get(patchedData, "type").String())
|
||||
require.Equal(t, "end_turn", gjson.Get(patchedData, "delta.stop_reason").String())
|
||||
require.Equal(t, "trace_001", gjson.Get(patchedData, "vendor_meta.trace_id").String())
|
||||
require.EqualValues(t, 53, gjson.Get(patchedData, "usage.output_tokens").Int())
|
||||
require.EqualValues(t, 100, gjson.Get(patchedData, "usage.input_tokens").Int())
|
||||
require.EqualValues(t, 30, gjson.Get(patchedData, "usage.cache_read_input_tokens").Int())
|
||||
require.EqualValues(t, 50, gjson.Get(patchedData, "usage.cache_creation_input_tokens").Int())
|
||||
}
|
||||
|
||||
func TestPatchClaudeMessageDeltaUsageDataZeroValueChecks(t *testing.T) {
|
||||
originalData := `{"type":"message_delta","usage":{"output_tokens":53,"input_tokens":9,"cache_read_input_tokens":0}}`
|
||||
usage := &dto.ClaudeUsage{
|
||||
InputTokens: 100,
|
||||
CacheReadInputTokens: 30,
|
||||
CacheCreationInputTokens: 0,
|
||||
}
|
||||
|
||||
patchedData := patchClaudeMessageDeltaUsageData(originalData, usage)
|
||||
|
||||
require.EqualValues(t, 9, gjson.Get(patchedData, "usage.input_tokens").Int())
|
||||
require.EqualValues(t, 30, gjson.Get(patchedData, "usage.cache_read_input_tokens").Int())
|
||||
assert.False(t, gjson.Get(patchedData, "usage.cache_creation_input_tokens").Exists())
|
||||
}
|
||||
|
||||
func TestShouldSkipClaudeMessageDeltaUsagePatch(t *testing.T) {
|
||||
originGlobalPassThrough := model_setting.GetGlobalSettings().PassThroughRequestEnabled
|
||||
t.Cleanup(func() {
|
||||
model_setting.GetGlobalSettings().PassThroughRequestEnabled = originGlobalPassThrough
|
||||
})
|
||||
|
||||
model_setting.GetGlobalSettings().PassThroughRequestEnabled = true
|
||||
assert.True(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{}))
|
||||
|
||||
model_setting.GetGlobalSettings().PassThroughRequestEnabled = false
|
||||
assert.True(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelSetting: dto.ChannelSettings{PassThroughBodyEnabled: true}},
|
||||
}))
|
||||
assert.False(t, shouldSkipClaudeMessageDeltaUsagePatch(&relaycommon.RelayInfo{
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelSetting: dto.ChannelSettings{PassThroughBodyEnabled: false}},
|
||||
}))
|
||||
}
|
||||
|
||||
func TestBuildMessageDeltaPatchUsage(t *testing.T) {
|
||||
t.Run("merge missing fields from claudeInfo", func(t *testing.T) {
|
||||
claudeResponse := &dto.ClaudeResponse{Usage: &dto.ClaudeUsage{OutputTokens: 53}}
|
||||
claudeInfo := &ClaudeResponseInfo{
|
||||
Usage: &dto.Usage{
|
||||
PromptTokens: 100,
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedTokens: 30,
|
||||
CachedCreationTokens: 50,
|
||||
},
|
||||
ClaudeCacheCreation5mTokens: 10,
|
||||
ClaudeCacheCreation1hTokens: 20,
|
||||
},
|
||||
}
|
||||
|
||||
usage := buildMessageDeltaPatchUsage(claudeResponse, claudeInfo)
|
||||
require.NotNil(t, usage)
|
||||
require.EqualValues(t, 100, usage.InputTokens)
|
||||
require.EqualValues(t, 30, usage.CacheReadInputTokens)
|
||||
require.EqualValues(t, 50, usage.CacheCreationInputTokens)
|
||||
require.EqualValues(t, 53, usage.OutputTokens)
|
||||
require.NotNil(t, usage.CacheCreation)
|
||||
require.EqualValues(t, 30, usage.CacheCreation.Ephemeral5mInputTokens)
|
||||
require.EqualValues(t, 20, usage.CacheCreation.Ephemeral1hInputTokens)
|
||||
})
|
||||
|
||||
t.Run("keep upstream non-zero values", func(t *testing.T) {
|
||||
claudeResponse := &dto.ClaudeResponse{Usage: &dto.ClaudeUsage{
|
||||
InputTokens: 9,
|
||||
CacheReadInputTokens: 7,
|
||||
CacheCreationInputTokens: 6,
|
||||
}}
|
||||
claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{
|
||||
PromptTokens: 100,
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedTokens: 30,
|
||||
CachedCreationTokens: 50,
|
||||
},
|
||||
}}
|
||||
|
||||
usage := buildMessageDeltaPatchUsage(claudeResponse, claudeInfo)
|
||||
require.EqualValues(t, 9, usage.InputTokens)
|
||||
require.EqualValues(t, 7, usage.CacheReadInputTokens)
|
||||
require.EqualValues(t, 6, usage.CacheCreationInputTokens)
|
||||
})
|
||||
|
||||
t.Run("default aggregate cache creation to 5m when split missing", func(t *testing.T) {
|
||||
claudeResponse := &dto.ClaudeResponse{Usage: &dto.ClaudeUsage{
|
||||
OutputTokens: 53,
|
||||
CacheCreationInputTokens: 50,
|
||||
}}
|
||||
claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedCreationTokens: 50,
|
||||
},
|
||||
}}
|
||||
|
||||
usage := buildMessageDeltaPatchUsage(claudeResponse, claudeInfo)
|
||||
require.NotNil(t, usage)
|
||||
require.NotNil(t, usage.CacheCreation)
|
||||
require.EqualValues(t, 50, usage.CacheCreation.Ephemeral5mInputTokens)
|
||||
require.EqualValues(t, 0, usage.CacheCreation.Ephemeral1hInputTokens)
|
||||
})
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,376 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/relayconvert"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func commonPointer[T any](value T) *T {
|
||||
return &value
|
||||
}
|
||||
|
||||
func TestResponseOpenAI2ClaudeToolUseInputIsObject(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
args string
|
||||
want map[string]interface{}
|
||||
}{
|
||||
{name: "object", args: `{"q":"x"}`, want: map[string]interface{}{"q": "x"}},
|
||||
{name: "empty", args: "", want: map[string]interface{}{}},
|
||||
{name: "invalid", args: "{", want: map[string]interface{}{}},
|
||||
{name: "null", args: "null", want: map[string]interface{}{}},
|
||||
{name: "array", args: `["x"]`, want: map[string]interface{}{}},
|
||||
{name: "string", args: `"x"`, want: map[string]interface{}{}},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
msg := dto.Message{Role: "assistant"}
|
||||
msg.SetToolCalls([]dto.ToolCallRequest{
|
||||
{
|
||||
ID: "call_1",
|
||||
Type: "function",
|
||||
Function: dto.FunctionRequest{
|
||||
Name: "lookup",
|
||||
Arguments: tt.args,
|
||||
},
|
||||
},
|
||||
})
|
||||
resp := relayconvert.ResponseOpenAI2Claude(&dto.OpenAITextResponse{
|
||||
Id: "chatcmpl_1",
|
||||
Model: "gpt-test",
|
||||
Choices: []dto.OpenAITextResponseChoice{
|
||||
{Message: msg, FinishReason: "tool_calls"},
|
||||
},
|
||||
}, nil)
|
||||
|
||||
require.Len(t, resp.Content, 1)
|
||||
assert.Equal(t, "tool_use", resp.Content[0].Type)
|
||||
assert.Equal(t, tt.want, resp.Content[0].Input)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) {
|
||||
claudeInfo := &ClaudeResponseInfo{
|
||||
Usage: &dto.Usage{},
|
||||
}
|
||||
claudeResponse := &dto.ClaudeResponse{
|
||||
Type: "message_start",
|
||||
Message: &dto.ClaudeMediaMessage{
|
||||
Id: "msg_123",
|
||||
Model: "claude-3-5-sonnet",
|
||||
Usage: &dto.ClaudeUsage{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 1,
|
||||
CacheCreationInputTokens: 50,
|
||||
CacheReadInputTokens: 30,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
|
||||
if !ok {
|
||||
t.Fatal("expected true")
|
||||
}
|
||||
if claudeInfo.Usage.PromptTokens != 100 {
|
||||
t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
|
||||
}
|
||||
if claudeInfo.Usage.PromptTokensDetails.CachedTokens != 30 {
|
||||
t.Errorf("CachedTokens = %d, want 30", claudeInfo.Usage.PromptTokensDetails.CachedTokens)
|
||||
}
|
||||
if claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens != 50 {
|
||||
t.Errorf("CachedCreationTokens = %d, want 50", claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens)
|
||||
}
|
||||
if claudeInfo.ResponseId != "msg_123" {
|
||||
t.Errorf("ResponseId = %s, want msg_123", claudeInfo.ResponseId)
|
||||
}
|
||||
if claudeInfo.Model != "claude-3-5-sonnet" {
|
||||
t.Errorf("Model = %s, want claude-3-5-sonnet", claudeInfo.Model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatClaudeResponseInfo_MessageDelta_FullUsage(t *testing.T) {
|
||||
// message_start 先积累 usage
|
||||
claudeInfo := &ClaudeResponseInfo{
|
||||
Usage: &dto.Usage{
|
||||
PromptTokens: 100,
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedTokens: 30,
|
||||
CachedCreationTokens: 50,
|
||||
},
|
||||
CompletionTokens: 1,
|
||||
},
|
||||
}
|
||||
|
||||
// message_delta 带完整 usage(原生 Anthropic 场景)
|
||||
claudeResponse := &dto.ClaudeResponse{
|
||||
Type: "message_delta",
|
||||
Usage: &dto.ClaudeUsage{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 200,
|
||||
CacheCreationInputTokens: 50,
|
||||
CacheReadInputTokens: 30,
|
||||
},
|
||||
}
|
||||
|
||||
ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
|
||||
if !ok {
|
||||
t.Fatal("expected true")
|
||||
}
|
||||
if claudeInfo.Usage.PromptTokens != 100 {
|
||||
t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
|
||||
}
|
||||
if claudeInfo.Usage.CompletionTokens != 200 {
|
||||
t.Errorf("CompletionTokens = %d, want 200", claudeInfo.Usage.CompletionTokens)
|
||||
}
|
||||
if claudeInfo.Usage.TotalTokens != 300 {
|
||||
t.Errorf("TotalTokens = %d, want 300", claudeInfo.Usage.TotalTokens)
|
||||
}
|
||||
if !claudeInfo.Done {
|
||||
t.Error("expected Done = true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatClaudeResponseInfo_MessageDelta_OnlyOutputTokens(t *testing.T) {
|
||||
// 模拟 Bedrock: message_start 已积累 usage
|
||||
claudeInfo := &ClaudeResponseInfo{
|
||||
Usage: &dto.Usage{
|
||||
PromptTokens: 100,
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedTokens: 30,
|
||||
CachedCreationTokens: 50,
|
||||
},
|
||||
CompletionTokens: 1,
|
||||
ClaudeCacheCreation5mTokens: 10,
|
||||
ClaudeCacheCreation1hTokens: 20,
|
||||
},
|
||||
}
|
||||
|
||||
// Bedrock 的 message_delta 只有 output_tokens,缺少 input_tokens 和 cache 字段
|
||||
claudeResponse := &dto.ClaudeResponse{
|
||||
Type: "message_delta",
|
||||
Usage: &dto.ClaudeUsage{
|
||||
OutputTokens: 200,
|
||||
// InputTokens, CacheCreationInputTokens, CacheReadInputTokens 都是 0
|
||||
},
|
||||
}
|
||||
|
||||
ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
|
||||
if !ok {
|
||||
t.Fatal("expected true")
|
||||
}
|
||||
// PromptTokens 应保持 message_start 的值(因为 message_delta 的 InputTokens=0,不更新)
|
||||
if claudeInfo.Usage.PromptTokens != 100 {
|
||||
t.Errorf("PromptTokens = %d, want 100", claudeInfo.Usage.PromptTokens)
|
||||
}
|
||||
if claudeInfo.Usage.CompletionTokens != 200 {
|
||||
t.Errorf("CompletionTokens = %d, want 200", claudeInfo.Usage.CompletionTokens)
|
||||
}
|
||||
if claudeInfo.Usage.TotalTokens != 300 {
|
||||
t.Errorf("TotalTokens = %d, want 300", claudeInfo.Usage.TotalTokens)
|
||||
}
|
||||
// cache 字段应保持 message_start 的值
|
||||
if claudeInfo.Usage.PromptTokensDetails.CachedTokens != 30 {
|
||||
t.Errorf("CachedTokens = %d, want 30", claudeInfo.Usage.PromptTokensDetails.CachedTokens)
|
||||
}
|
||||
if claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens != 50 {
|
||||
t.Errorf("CachedCreationTokens = %d, want 50", claudeInfo.Usage.PromptTokensDetails.CachedCreationTokens)
|
||||
}
|
||||
if claudeInfo.Usage.ClaudeCacheCreation5mTokens != 10 {
|
||||
t.Errorf("ClaudeCacheCreation5mTokens = %d, want 10", claudeInfo.Usage.ClaudeCacheCreation5mTokens)
|
||||
}
|
||||
if claudeInfo.Usage.ClaudeCacheCreation1hTokens != 20 {
|
||||
t.Errorf("ClaudeCacheCreation1hTokens = %d, want 20", claudeInfo.Usage.ClaudeCacheCreation1hTokens)
|
||||
}
|
||||
if !claudeInfo.Done {
|
||||
t.Error("expected Done = true")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatClaudeResponseInfo_NilClaudeInfo(t *testing.T) {
|
||||
claudeResponse := &dto.ClaudeResponse{Type: "message_start"}
|
||||
ok := FormatClaudeResponseInfo(claudeResponse, nil, nil)
|
||||
if ok {
|
||||
t.Error("expected false for nil claudeInfo")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFormatClaudeResponseInfo_ContentBlockDelta(t *testing.T) {
|
||||
text := "hello"
|
||||
claudeInfo := &ClaudeResponseInfo{
|
||||
Usage: &dto.Usage{},
|
||||
ResponseText: strings.Builder{},
|
||||
}
|
||||
claudeResponse := &dto.ClaudeResponse{
|
||||
Type: "content_block_delta",
|
||||
Delta: &dto.ClaudeMediaMessage{
|
||||
Text: &text,
|
||||
},
|
||||
}
|
||||
|
||||
ok := FormatClaudeResponseInfo(claudeResponse, nil, claudeInfo)
|
||||
if !ok {
|
||||
t.Fatal("expected true")
|
||||
}
|
||||
if claudeInfo.ResponseText.String() != "hello" {
|
||||
t.Errorf("ResponseText = %q, want %q", claudeInfo.ResponseText.String(), "hello")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAIStyleUsageFromClaudeUsage(t *testing.T) {
|
||||
usage := &dto.Usage{
|
||||
PromptTokens: 100,
|
||||
CompletionTokens: 20,
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedTokens: 30,
|
||||
CachedCreationTokens: 50,
|
||||
},
|
||||
ClaudeCacheCreation5mTokens: 10,
|
||||
ClaudeCacheCreation1hTokens: 20,
|
||||
UsageSemantic: "anthropic",
|
||||
}
|
||||
|
||||
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(usage)
|
||||
|
||||
if openAIUsage.PromptTokens != 180 {
|
||||
t.Fatalf("PromptTokens = %d, want 180", openAIUsage.PromptTokens)
|
||||
}
|
||||
if openAIUsage.InputTokens != 180 {
|
||||
t.Fatalf("InputTokens = %d, want 180", openAIUsage.InputTokens)
|
||||
}
|
||||
if openAIUsage.TotalTokens != 200 {
|
||||
t.Fatalf("TotalTokens = %d, want 200", openAIUsage.TotalTokens)
|
||||
}
|
||||
if openAIUsage.UsageSemantic != "openai" {
|
||||
t.Fatalf("UsageSemantic = %s, want openai", openAIUsage.UsageSemantic)
|
||||
}
|
||||
if openAIUsage.UsageSource != "anthropic" {
|
||||
t.Fatalf("UsageSource = %s, want anthropic", openAIUsage.UsageSource)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAIStyleUsageFromClaudeUsagePreservesCacheCreationRemainder(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
cachedCreationTokens int
|
||||
cacheCreationTokens5m int
|
||||
cacheCreationTokens1h int
|
||||
expectedTotalInputToken int
|
||||
}{
|
||||
{
|
||||
name: "prefers aggregate when it includes remainder",
|
||||
cachedCreationTokens: 50,
|
||||
cacheCreationTokens5m: 10,
|
||||
cacheCreationTokens1h: 20,
|
||||
expectedTotalInputToken: 180,
|
||||
},
|
||||
{
|
||||
name: "falls back to split tokens when aggregate missing",
|
||||
cachedCreationTokens: 0,
|
||||
cacheCreationTokens5m: 10,
|
||||
cacheCreationTokens1h: 20,
|
||||
expectedTotalInputToken: 160,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
usage := &dto.Usage{
|
||||
PromptTokens: 100,
|
||||
CompletionTokens: 20,
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedTokens: 30,
|
||||
CachedCreationTokens: tt.cachedCreationTokens,
|
||||
},
|
||||
ClaudeCacheCreation5mTokens: tt.cacheCreationTokens5m,
|
||||
ClaudeCacheCreation1hTokens: tt.cacheCreationTokens1h,
|
||||
UsageSemantic: "anthropic",
|
||||
}
|
||||
|
||||
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(usage)
|
||||
|
||||
if openAIUsage.PromptTokens != tt.expectedTotalInputToken {
|
||||
t.Fatalf("PromptTokens = %d, want %d", openAIUsage.PromptTokens, tt.expectedTotalInputToken)
|
||||
}
|
||||
if openAIUsage.InputTokens != tt.expectedTotalInputToken {
|
||||
t.Fatalf("InputTokens = %d, want %d", openAIUsage.InputTokens, tt.expectedTotalInputToken)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpenAIStyleUsageFromClaudeUsageDefaultsAggregateCacheCreationTo5m(t *testing.T) {
|
||||
usage := &dto.Usage{
|
||||
PromptTokens: 100,
|
||||
CompletionTokens: 20,
|
||||
PromptTokensDetails: dto.InputTokenDetails{
|
||||
CachedTokens: 30,
|
||||
CachedCreationTokens: 50,
|
||||
},
|
||||
UsageSemantic: "anthropic",
|
||||
}
|
||||
|
||||
openAIUsage := buildOpenAIStyleUsageFromClaudeUsage(usage)
|
||||
|
||||
require.Equal(t, 50, openAIUsage.ClaudeCacheCreation5mTokens)
|
||||
require.Equal(t, 0, openAIUsage.ClaudeCacheCreation1hTokens)
|
||||
}
|
||||
|
||||
func TestOpenAIChatRequestToClaudeMessages_ClaudeOpus48HighUsesAdaptiveThinking(t *testing.T) {
|
||||
request := dto.GeneralOpenAIRequest{
|
||||
Model: "claude-opus-4-8-high",
|
||||
Temperature: commonPointer(0.7),
|
||||
TopP: commonPointer(0.9),
|
||||
TopK: commonPointer(40),
|
||||
Messages: []dto.Message{
|
||||
{
|
||||
Role: "user",
|
||||
Content: "hello",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
claudeRequest, err := relayconvert.OpenAIChatRequestToClaudeMessages(nil, &relaycommon.RelayInfo{}, request)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "claude-opus-4-8", claudeRequest.Model)
|
||||
require.NotNil(t, claudeRequest.Thinking)
|
||||
require.Equal(t, "adaptive", claudeRequest.Thinking.Type)
|
||||
require.Equal(t, "summarized", claudeRequest.Thinking.Display)
|
||||
require.JSONEq(t, `{"effort":"high"}`, string(claudeRequest.OutputConfig))
|
||||
require.Nil(t, claudeRequest.Temperature)
|
||||
require.Nil(t, claudeRequest.TopP)
|
||||
require.Nil(t, claudeRequest.TopK)
|
||||
}
|
||||
|
||||
func TestOpenAIChatRequestToClaudeMessages_ClaudeOpus48ThinkingUsesAdaptiveHighEffort(t *testing.T) {
|
||||
request := dto.GeneralOpenAIRequest{
|
||||
Model: "claude-opus-4-8-thinking",
|
||||
Temperature: commonPointer(0.7),
|
||||
TopP: commonPointer(0.9),
|
||||
TopK: commonPointer(40),
|
||||
Messages: []dto.Message{
|
||||
{
|
||||
Role: "user",
|
||||
Content: "hello",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
claudeRequest, err := relayconvert.OpenAIChatRequestToClaudeMessages(nil, &relaycommon.RelayInfo{}, request)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "claude-opus-4-8", claudeRequest.Model)
|
||||
require.NotNil(t, claudeRequest.Thinking)
|
||||
require.Equal(t, "adaptive", claudeRequest.Thinking.Type)
|
||||
require.Equal(t, "summarized", claudeRequest.Thinking.Display)
|
||||
require.JSONEq(t, `{"effort":"high"}`, string(claudeRequest.OutputConfig))
|
||||
require.Nil(t, claudeRequest.Temperature)
|
||||
require.Nil(t, claudeRequest.TopP)
|
||||
require.Nil(t, claudeRequest.TopK)
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
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/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestHandleClaudeResponseDataCountsToolUse(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
operation_setting.SetToolPriceForTest("lookup_fn", 3.0)
|
||||
t.Cleanup(func() {
|
||||
operation_setting.DeleteToolPriceForTest("lookup_fn")
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
info := &relaycommon.RelayInfo{
|
||||
OriginModelName: "claude-3-7-sonnet",
|
||||
RelayFormat: types.RelayFormatClaude,
|
||||
}
|
||||
claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{}}
|
||||
|
||||
data := []byte(`{
|
||||
"type":"message",
|
||||
"content":[
|
||||
{"type":"text","text":"hi"},
|
||||
{"type":"tool_use","id":"tu1","name":"lookup_fn","input":{}},
|
||||
{"type":"server_tool_use","id":"stu1","name":"web_search","input":{}}
|
||||
],
|
||||
"usage":{"input_tokens":1,"output_tokens":1}
|
||||
}`)
|
||||
|
||||
err := HandleClaudeResponseData(c, info, claudeInfo, nil, data)
|
||||
require.Nil(t, err)
|
||||
require.NotNil(t, info.ResponsesUsageInfo)
|
||||
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, "lookup_fn")
|
||||
assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["lookup_fn"].CallCount)
|
||||
assert.NotContains(t, info.ResponsesUsageInfo.BuiltInTools, "web_search")
|
||||
}
|
||||
|
||||
func TestCountClaudeStreamBillableToolsSetsWebSearchRequests(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
info := &relaycommon.RelayInfo{OriginModelName: "claude-3-7-sonnet"}
|
||||
|
||||
countClaudeStreamBillableTools(c, info, &dto.ClaudeResponse{
|
||||
Type: "message_delta",
|
||||
Usage: &dto.ClaudeUsage{
|
||||
ServerToolUse: &dto.ClaudeServerToolUse{WebSearchRequests: 3},
|
||||
},
|
||||
})
|
||||
assert.Equal(t, 3, c.GetInt("claude_web_search_requests"))
|
||||
|
||||
operation_setting.SetToolPriceForTest("stream_fn", 2.0)
|
||||
t.Cleanup(func() {
|
||||
operation_setting.DeleteToolPriceForTest("stream_fn")
|
||||
})
|
||||
countClaudeStreamBillableTools(c, info, &dto.ClaudeResponse{
|
||||
Type: "content_block_start",
|
||||
ContentBlock: &dto.ClaudeMediaMessage{
|
||||
Type: "tool_use",
|
||||
Name: "stream_fn",
|
||||
},
|
||||
})
|
||||
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, "stream_fn")
|
||||
assert.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["stream_fn"].CallCount)
|
||||
}
|
||||
Reference in New Issue
Block a user