diff --git a/dto/claude.go b/dto/claude.go index 0c552b4b2815..97c7bca28015 100644 --- a/dto/claude.go +++ b/dto/claude.go @@ -22,6 +22,7 @@ type ClaudeMediaMessage struct { Source *ClaudeMessageSource `json:"source,omitempty"` Usage *ClaudeUsage `json:"usage,omitempty"` StopReason *string `json:"stop_reason,omitempty"` + StopSequence *string `json:"stop_sequence,omitempty"` PartialJson *string `json:"partial_json,omitempty"` Role string `json:"role,omitempty"` Thinking *string `json:"thinking,omitempty"` diff --git a/relay/channel/claude/adaptor.go b/relay/channel/claude/adaptor.go index b8e4a0366dd7..033620cfd453 100644 --- a/relay/channel/claude/adaptor.go +++ b/relay/channel/claude/adaptor.go @@ -124,6 +124,11 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request 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 { + clientStream := info.Request != nil && info.Request.IsStream(c) + if !clientStream { + info.IsStream = false + return ClaudeBufferedStreamHandler(c, resp, info) + } return ClaudeStreamHandler(c, resp, info) } else { return ClaudeHandler(c, resp, info) diff --git a/relay/channel/claude/buffered_stream.go b/relay/channel/claude/buffered_stream.go new file mode 100644 index 000000000000..828e1a3e24dd --- /dev/null +++ b/relay/channel/claude/buffered_stream.go @@ -0,0 +1,316 @@ +package claude + +import ( + "bufio" + "fmt" + "net/http" + "sort" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" +) + +type claudeBufferedStreamAccumulator struct { + response dto.ClaudeResponse + stopSequence *string + blocks map[int]*dto.ClaudeMediaMessage + toolInputs map[int]*strings.Builder + started bool + stopped bool +} + +type claudeBufferedResponse struct { + *dto.ClaudeResponse + StopSequence *string `json:"stop_sequence"` +} + +func newClaudeBufferedStreamAccumulator() *claudeBufferedStreamAccumulator { + return &claudeBufferedStreamAccumulator{ + response: dto.ClaudeResponse{ + Type: "message", + Role: "assistant", + }, + blocks: make(map[int]*dto.ClaudeMediaMessage), + toolInputs: make(map[int]*strings.Builder), + } +} + +func cloneClaudeUsage(usage *dto.ClaudeUsage) *dto.ClaudeUsage { + if usage == nil { + return nil + } + cloned := *usage + if usage.CacheCreation != nil { + cacheCreation := *usage.CacheCreation + cloned.CacheCreation = &cacheCreation + } + if usage.ServerToolUse != nil { + serverToolUse := *usage.ServerToolUse + cloned.ServerToolUse = &serverToolUse + } + cloned.BillingUsage = dto.CloneBillingUsage(usage.BillingUsage) + return &cloned +} + +func mergeClaudeUsage(target **dto.ClaudeUsage, incoming *dto.ClaudeUsage) { + if incoming == nil { + return + } + if *target == nil { + *target = cloneClaudeUsage(incoming) + return + } + + usage := *target + if incoming.InputTokens > 0 { + usage.InputTokens = incoming.InputTokens + } + if incoming.CacheCreationInputTokens > 0 { + usage.CacheCreationInputTokens = incoming.CacheCreationInputTokens + } + if incoming.CacheReadInputTokens > 0 { + usage.CacheReadInputTokens = incoming.CacheReadInputTokens + } + if incoming.OutputTokens > 0 { + usage.OutputTokens = incoming.OutputTokens + } + if incoming.ClaudeCacheCreation5mTokens > 0 { + usage.ClaudeCacheCreation5mTokens = incoming.ClaudeCacheCreation5mTokens + } + if incoming.ClaudeCacheCreation1hTokens > 0 { + usage.ClaudeCacheCreation1hTokens = incoming.ClaudeCacheCreation1hTokens + } + if incoming.CacheCreation != nil { + if usage.CacheCreation == nil { + cacheCreation := *incoming.CacheCreation + usage.CacheCreation = &cacheCreation + } else { + if incoming.CacheCreation.Ephemeral5mInputTokens > 0 { + usage.CacheCreation.Ephemeral5mInputTokens = incoming.CacheCreation.Ephemeral5mInputTokens + } + if incoming.CacheCreation.Ephemeral1hInputTokens > 0 { + usage.CacheCreation.Ephemeral1hInputTokens = incoming.CacheCreation.Ephemeral1hInputTokens + } + } + } + if incoming.ServerToolUse != nil { + serverToolUse := *incoming.ServerToolUse + usage.ServerToolUse = &serverToolUse + } + if incoming.BillingUsage != nil { + usage.BillingUsage = dto.CloneBillingUsage(incoming.BillingUsage) + } +} + +func (a *claudeBufferedStreamAccumulator) block(index int) *dto.ClaudeMediaMessage { + block, ok := a.blocks[index] + if ok { + return block + } + block = &dto.ClaudeMediaMessage{} + a.blocks[index] = block + return block +} + +func (a *claudeBufferedStreamAccumulator) finalizeToolInput(index int) error { + builder, ok := a.toolInputs[index] + if !ok { + return nil + } + delete(a.toolInputs, index) + + input := make(map[string]interface{}) + partialJSON := strings.TrimSpace(builder.String()) + if partialJSON != "" { + if err := common.Unmarshal([]byte(partialJSON), &input); err != nil { + return fmt.Errorf("invalid Claude tool input at content block %d: %w", index, err) + } + } + a.block(index).Input = input + return nil +} + +func (a *claudeBufferedStreamAccumulator) process(event *dto.ClaudeResponse) error { + if event == nil { + return nil + } + + switch event.Type { + case "message_start": + if event.Message == nil { + return fmt.Errorf("Claude message_start is missing message") + } + a.started = true + a.response.Id = event.Message.Id + a.response.Model = event.Message.Model + if event.Message.Type != "" { + a.response.Type = event.Message.Type + } + if event.Message.Role != "" { + a.response.Role = event.Message.Role + } + if event.Message.StopReason != nil { + a.response.StopReason = *event.Message.StopReason + } + a.stopSequence = event.Message.StopSequence + mergeClaudeUsage(&a.response.Usage, event.Message.Usage) + case "content_block_start": + if event.ContentBlock == nil { + return fmt.Errorf("Claude content_block_start is missing content_block") + } + contentBlock := *event.ContentBlock + a.blocks[event.GetIndex()] = &contentBlock + case "content_block_delta": + if event.Delta == nil { + return fmt.Errorf("Claude content_block_delta is missing delta") + } + index := event.GetIndex() + block := a.block(index) + switch event.Delta.Type { + case "text_delta": + block.Type = "text" + if event.Delta.Text != nil { + text := block.GetText() + *event.Delta.Text + block.Text = &text + } + case "thinking_delta": + block.Type = "thinking" + if event.Delta.Thinking != nil { + thinking := "" + if block.Thinking != nil { + thinking = *block.Thinking + } + thinking += *event.Delta.Thinking + block.Thinking = &thinking + } + case "signature_delta": + block.Signature += event.Delta.Signature + case "input_json_delta": + if event.Delta.PartialJson != nil { + builder, ok := a.toolInputs[index] + if !ok { + builder = &strings.Builder{} + a.toolInputs[index] = builder + } + builder.WriteString(*event.Delta.PartialJson) + } + } + case "content_block_stop": + return a.finalizeToolInput(event.GetIndex()) + case "message_delta": + if event.Delta != nil { + if event.Delta.StopReason != nil { + a.response.StopReason = *event.Delta.StopReason + } + if event.Delta.StopSequence != nil { + a.stopSequence = event.Delta.StopSequence + } + } + mergeClaudeUsage(&a.response.Usage, event.Usage) + case "message_stop": + a.stopped = true + return nil + case "ping": + return nil + } + return nil +} + +func (a *claudeBufferedStreamAccumulator) finalResponse(model string) (*claudeBufferedResponse, error) { + if !a.started { + return nil, fmt.Errorf("Claude stream ended without message_start") + } + if !a.stopped { + return nil, fmt.Errorf("Claude stream ended without message_stop") + } + for index := range a.toolInputs { + if err := a.finalizeToolInput(index); err != nil { + return nil, err + } + } + + indices := make([]int, 0, len(a.blocks)) + for index := range a.blocks { + indices = append(indices, index) + } + sort.Ints(indices) + a.response.Content = make([]dto.ClaudeMediaMessage, 0, len(indices)) + for _, index := range indices { + a.response.Content = append(a.response.Content, *a.blocks[index]) + } + + if a.response.Model == "" { + a.response.Model = model + } + if a.response.Usage == nil { + a.response.Usage = &dto.ClaudeUsage{} + } + return &claudeBufferedResponse{ + ClaudeResponse: &a.response, + StopSequence: a.stopSequence, + }, nil +} + +func ClaudeBufferedStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.NewAPIError) { + if resp == nil || resp.Body == nil { + return nil, types.NewError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse) + } + defer service.CloseResponseBodyGracefully(resp) + + accumulator := newClaudeBufferedStreamAccumulator() + scanner := helper.NewStreamScanner(resp.Body) + scanner.Split(bufio.ScanLines) + for scanner.Scan() { + line := scanner.Text() + if len(line) < 5 || line[:5] != "data:" { + continue + } + data := strings.TrimSpace(line[5:]) + if data == "" || data == "[DONE]" { + continue + } + + info.SetFirstResponseTime() + info.ReceivedResponseCount++ + + var event dto.ClaudeResponse + if err := common.UnmarshalJsonStr(data, &event); err != nil { + return nil, types.NewError(err, types.ErrorCodeBadResponseBody) + } + if claudeError := event.GetClaudeError(); claudeError != nil && claudeError.Type != "" { + return nil, types.WithClaudeError(*claudeError, http.StatusInternalServerError) + } + if err := accumulator.process(&event); err != nil { + return nil, types.NewError(err, types.ErrorCodeBadResponseBody) + } + if event.Type == "message_stop" { + break + } + } + if err := scanner.Err(); err != nil { + return nil, types.NewError(err, types.ErrorCodeBadResponse) + } + + response, err := accumulator.finalResponse(info.UpstreamModelName) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeBadResponseBody) + } + responseBody, err := common.Marshal(response) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeJsonMarshalFailed) + } + + resp.Header.Set("Content-Type", "application/json; charset=utf-8") + claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{}} + if handleErr := HandleClaudeResponseData(c, info, claudeInfo, resp, responseBody); handleErr != nil { + return nil, handleErr + } + return claudeInfo.Usage, nil +} diff --git a/relay/channel/claude/buffered_stream_test.go b/relay/channel/claude/buffered_stream_test.go new file mode 100644 index 000000000000..0d4770f2b0ff --- /dev/null +++ b/relay/channel/claude/buffered_stream_test.go @@ -0,0 +1,170 @@ +package claude + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/gin-gonic/gin" +) + +func TestAdaptorDoResponseBuffersClaudeStreamForNonStreamClient(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + clientStream := false + info := &relaycommon.RelayInfo{ + Request: &dto.ClaudeRequest{ + Stream: &clientStream, + }, + RelayFormat: types.RelayFormatClaude, + IsStream: true, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "claude-test", + }, + } + sse := strings.Join([]string{ + `event: message_start`, + `data: {"type":"message_start","message":{"id":"msg_test","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":11,"cache_creation_input_tokens":2,"cache_read_input_tokens":3,"output_tokens":1}}}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"consider"}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sig"}}`, + ``, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":0}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"OK"}}`, + ``, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":1}`, + ``, + `event: content_block_start`, + `data: {"type":"content_block_start","index":2,"content_block":{"type":"tool_use","id":"tool_1","name":"lookup","input":{}}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"{\"city\":"}}`, + ``, + `event: content_block_delta`, + `data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"\"Paris\"}"}}`, + ``, + `event: content_block_stop`, + `data: {"type":"content_block_stop","index":2}`, + ``, + `event: message_delta`, + `data: {"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":"tool-end"},"usage":{"output_tokens":9}}`, + ``, + `event: message_stop`, + `data: {"type":"message_stop"}`, + ``, + }, "\n") + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + }, + Body: io.NopCloser(strings.NewReader(sse)), + } + + usageValue, apiErr := (&Adaptor{}).DoResponse(c, resp, info) + require.Nil(t, apiErr) + usage, ok := usageValue.(*dto.Usage) + require.True(t, ok) + assert.Equal(t, 11, usage.PromptTokens) + assert.Equal(t, 9, usage.CompletionTokens) + assert.Equal(t, "application/json; charset=utf-8", recorder.Header().Get("Content-Type")) + assert.False(t, info.IsStream) + + var response dto.ClaudeResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + var responseObject map[string]interface{} + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &responseObject)) + stopSequence, exists := responseObject["stop_sequence"] + assert.True(t, exists) + assert.Equal(t, "tool-end", stopSequence) + assert.Equal(t, "msg_test", response.Id) + assert.Equal(t, "message", response.Type) + assert.Equal(t, "assistant", response.Role) + assert.Equal(t, "claude-test", response.Model) + assert.Equal(t, "tool_use", response.StopReason) + require.Len(t, response.Content, 3) + assert.Equal(t, "thinking", response.Content[0].Type) + require.NotNil(t, response.Content[0].Thinking) + assert.Equal(t, "consider", *response.Content[0].Thinking) + assert.Equal(t, "sig", response.Content[0].Signature) + assert.Equal(t, "text", response.Content[1].Type) + assert.Equal(t, "OK", response.Content[1].GetText()) + assert.Equal(t, "tool_use", response.Content[2].Type) + assert.Equal(t, "tool_1", response.Content[2].Id) + assert.Equal(t, "lookup", response.Content[2].Name) + assert.Equal(t, map[string]interface{}{"city": "Paris"}, response.Content[2].Input) + require.NotNil(t, response.Usage) + assert.Equal(t, 11, response.Usage.InputTokens) + assert.Equal(t, 2, response.Usage.CacheCreationInputTokens) + assert.Equal(t, 3, response.Usage.CacheReadInputTokens) + assert.Equal(t, 9, response.Usage.OutputTokens) +} + +func TestAdaptorDoResponseKeepsClaudeStreamForStreamClient(t *testing.T) { + gin.SetMode(gin.TestMode) + previousStreamingTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { + constant.StreamingTimeout = previousStreamingTimeout + }) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + clientStream := true + info := &relaycommon.RelayInfo{ + Request: &dto.ClaudeRequest{ + Stream: &clientStream, + }, + RelayFormat: types.RelayFormatClaude, + IsStream: true, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "claude-test", + }, + } + sse := strings.Join([]string{ + `data: {"type":"message_start","message":{"id":"msg_stream","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":1,"output_tokens":1}}}`, + `data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}`, + `data: {"type":"message_stop"}`, + ``, + }, "\n") + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + }, + Body: io.NopCloser(strings.NewReader(sse)), + } + + _, apiErr := (&Adaptor{}).DoResponse(c, resp, info) + require.Nil(t, apiErr) + assert.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type")) + assert.True(t, info.IsStream) + assert.Contains(t, recorder.Body.String(), `"type":"message_start"`) +}