feat: publish sanitized custom.97 source

relaykit/relayconvert/response_registry_test.go

@@ -1,6 +1,7 @@
 package relayconvert
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/QuantumNous/new-api/relaykit/dto"
@@ -699,3 +700,104 @@ func textRegistryResponsesResponse() *dto.OpenAIResponsesResponse {
 func respPtr[T any](value T) *T {
 	return &value
 }
+
+func TestConvertResponseToResponsesRestoresRecordedCustomTools(t *testing.T) {
+	info := &convmeta.Values{ResponsesTools: &convmeta.ResponsesToolState{CustomToolNames: map[string]struct{}{"exec": {}}}}
+	chatMessage := dto.Message{Role: "assistant"}
+	chatMessage.SetToolCalls([]dto.ToolCallRequest{{ID: "call_exec", Type: "function", Function: dto.FunctionRequest{Name: "exec", Arguments: `{"input":"ls"}`}}})
+	geminiExecCall := `{"candidates":[{"finishReason":"STOP","content":{"role":"model","parts":[{"functionCall":{"name":"exec","args":{"input":"ls"}}}]}}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}`
+	tests := []struct {
+		name     string
+		from     types.RelayFormat
+		response any
+		stream   []any
+	}{
+		{
+			name: "chat",
+			from: types.RelayFormatOpenAI,
+			response: &dto.OpenAITextResponse{
+				Id:      "chatcmpl_1",
+				Model:   "gpt-test",
+				Choices: []dto.OpenAITextResponseChoice{{Message: chatMessage, FinishReason: "tool_calls"}},
+			},
+			stream: []any{
+				chatStreamChunk(`{"id":"chatcmpl_1","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_exec","type":"function","function":{"name":"exec","arguments":"{\"input\":"}}]}}]}`),
+				chatStreamChunk(`{"id":"chatcmpl_1","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"ls\"}"}}]},"finish_reason":"tool_calls"}]}`),
+			},
+		},
+		{
+			name: "claude",
+			from: types.RelayFormatClaude,
+			response: &dto.ClaudeResponse{
+				Id:         "msg_1",
+				Type:       "message",
+				Role:       "assistant",
+				Model:      "claude-test",
+				StopReason: "tool_use",
+				Content:    []dto.ClaudeMediaMessage{{Type: "tool_use", Id: "toolu_exec", Name: "exec", Input: map[string]any{"input": "ls"}}},
+			},
+			stream: []any{
+				claudeStreamChunk(`{"type":"message_start","message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[],"usage":{"input_tokens":4,"output_tokens":0}}}`),
+				claudeStreamChunk(`{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_exec","name":"exec","input":{}}}`),
+				claudeStreamChunk(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"input\":"}}`),
+				claudeStreamChunk(`{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"ls\"}"}}`),
+				claudeStreamChunk(`{"type":"content_block_stop","index":0}`),
+				claudeStreamChunk(`{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":2}}`),
+				claudeStreamChunk(`{"type":"message_stop"}`),
+			},
+		},
+		{
+			name:     "gemini",
+			from:     types.RelayFormatGemini,
+			response: geminiStreamChunk(geminiExecCall),
+			stream:   []any{geminiStreamChunk(geminiExecCall)},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := ConvertResponse(nil, info, types.RelayFormatOpenAIResponses, tt.response)
+			require.NoError(t, err)
+			responses, ok := result.Value.(*dto.OpenAIResponsesResponse)
+			require.True(t, ok)
+			require.Len(t, responses.Output, 1)
+			assert.Equal(t, "custom_tool_call", responses.Output[0].Type)
+			assert.Equal(t, "exec", responses.Output[0].Name)
+			assert.Equal(t, `"ls"`, string(responses.Output[0].Input))
+
+			state, err := NewResponseStreamState(tt.from, types.RelayFormatOpenAIResponses, ResponseStreamOptions{ID: "resp_1", Model: "model-test"})
+			require.NoError(t, err)
+			var results []ResponseResult
+			for _, chunk := range tt.stream {
+				chunkResults, err := ConvertStreamResponseChunk(nil, info, state, chunk)
+				require.NoError(t, err)
+				results = append(results, chunkResults...)
+			}
+			finals, err := FinalizeStreamResponse(nil, info, state)
+			require.NoError(t, err)
+			results = append(results, finals...)
+
+			var toolEvents []string
+			for _, result := range results {
+				event, ok := result.Value.(ChatToResponsesStreamEvent)
+				require.True(t, ok)
+				switch {
+				case event.Payload.Item != nil && event.Payload.Item.Type != "message":
+					toolEvents = append(toolEvents, event.Type+" "+event.Payload.Item.Type+" "+string(event.Payload.Item.Input))
+				case event.Type == "response.custom_tool_call_input.delta":
+					toolEvents = append(toolEvents, event.Type+" "+event.Payload.Delta)
+				case event.Type == "response.custom_tool_call_input.done":
+					require.NotNil(t, event.Payload.Input)
+					toolEvents = append(toolEvents, event.Type+" "+*event.Payload.Input)
+				case strings.Contains(event.Type, "function_call_arguments"):
+					toolEvents = append(toolEvents, event.Type)
+				}
+			}
+			assert.Equal(t, []string{
+				`response.output_item.added custom_tool_call ""`,
+				"response.custom_tool_call_input.delta ls",
+				"response.custom_tool_call_input.done ls",
+				`response.output_item.done custom_tool_call "ls"`,
+			}, toolEvents)
+		})
+	}
+}