diff --git a/internal/api/gemini.go b/internal/api/gemini.go index d8ef7c8..17b4776 100644 --- a/internal/api/gemini.go +++ b/internal/api/gemini.go @@ -828,46 +828,109 @@ func buildGeminiResponse(request aistudio.GenerateRequest, result generationResu func geminiOutputParts(result generationResult) []map[string]any { parts := make([]map[string]any, 0) + var pendingSignature string + for _, event := range result.events { switch event.Kind { case aistudio.EventText: - parts = append(parts, geminiSignedPart(geminiTextPart(event), event.ThoughtSignature)) + if len(parts) > 0 { + last := parts[len(parts)-1] + if lastText, ok := last["text"].(string); ok && last["thought"] != true && last["transcriptionMetadata"] == nil && event.Transcript == nil { + last["text"] = lastText + event.Text + if event.ThoughtSignature != "" { + last["thoughtSignature"] = event.ThoughtSignature + } + continue + } + } + part := geminiSignedPart(geminiTextPart(event), event.ThoughtSignature) + if event.ThoughtSignature == "" && pendingSignature != "" { + part["thoughtSignature"] = pendingSignature + pendingSignature = "" + } + parts = append(parts, part) case aistudio.EventReasoning: - parts = append(parts, geminiSignedPart(map[string]any{"text": event.Text, "thought": true}, event.ThoughtSignature)) + if len(parts) > 0 { + last := parts[len(parts)-1] + if lastText, ok := last["text"].(string); ok && last["thought"] == true { + last["text"] = lastText + event.Text + if event.ThoughtSignature != "" { + last["thoughtSignature"] = event.ThoughtSignature + } + continue + } + } + part := geminiSignedPart(map[string]any{"text": event.Text, "thought": true}, event.ThoughtSignature) + if event.ThoughtSignature == "" && pendingSignature != "" { + part["thoughtSignature"] = pendingSignature + pendingSignature = "" + } + parts = append(parts, part) case aistudio.EventToolCall: if event.ToolCall != nil { - parts = append(parts, geminiSignedPart(geminiFunctionCallPart(*event.ToolCall), event.ThoughtSignature)) + part := geminiSignedPart(geminiFunctionCallPart(*event.ToolCall), event.ThoughtSignature) + if event.ThoughtSignature == "" && pendingSignature != "" { + part["thoughtSignature"] = pendingSignature + pendingSignature = "" + } + parts = append(parts, part) } case aistudio.EventExecutableCode: if event.ExecutableCode != nil { - parts = append(parts, geminiSignedPart(map[string]any{"executableCode": map[string]any{ + part := geminiSignedPart(map[string]any{"executableCode": map[string]any{ "language": event.ExecutableCode.Language, "code": event.ExecutableCode.Code, - }}, event.ThoughtSignature)) + }}, event.ThoughtSignature) + if event.ThoughtSignature == "" && pendingSignature != "" { + part["thoughtSignature"] = pendingSignature + pendingSignature = "" + } + parts = append(parts, part) } case aistudio.EventCodeExecutionResult: if event.CodeExecutionResult != nil { - parts = append(parts, geminiSignedPart(map[string]any{ + part := geminiSignedPart(map[string]any{ "codeExecutionResult": geminiCodeExecutionResult(*event.CodeExecutionResult), - }, event.ThoughtSignature)) + }, event.ThoughtSignature) + if event.ThoughtSignature == "" && pendingSignature != "" { + part["thoughtSignature"] = pendingSignature + pendingSignature = "" + } + parts = append(parts, part) } case aistudio.EventMedia: if event.Media != nil { + var part map[string]any if len(event.Media.Data) > 0 { - parts = append(parts, geminiSignedPart(map[string]any{"inlineData": map[string]any{ + part = geminiSignedPart(map[string]any{"inlineData": map[string]any{ "mimeType": event.Media.MIME, "data": base64.StdEncoding.EncodeToString(event.Media.Data), - }}, event.ThoughtSignature)) + }}, event.ThoughtSignature) } else if event.Media.URL != "" { - parts = append(parts, geminiSignedPart(map[string]any{"fileData": map[string]any{ + part = geminiSignedPart(map[string]any{"fileData": map[string]any{ "mimeType": event.Media.MIME, "fileUri": event.Media.URL, "displayName": event.Media.Name, - }}, event.ThoughtSignature)) + }}, event.ThoughtSignature) + } + if part != nil { + if event.ThoughtSignature == "" && pendingSignature != "" { + part["thoughtSignature"] = pendingSignature + pendingSignature = "" + } + parts = append(parts, part) } } case aistudio.EventThoughtSignature: - if event.ThoughtSignature != "" { - parts = append(parts, map[string]any{"thoughtSignature": event.ThoughtSignature}) + if event.ThoughtSignature == "" { + continue + } + if len(parts) > 0 { + parts[len(parts)-1]["thoughtSignature"] = event.ThoughtSignature + } else { + pendingSignature = event.ThoughtSignature } } } + if len(parts) == 0 && pendingSignature != "" { + parts = append(parts, map[string]any{"thoughtSignature": pendingSignature}) + } return parts } diff --git a/internal/api/gemini_output_test.go b/internal/api/gemini_output_test.go new file mode 100644 index 0000000..28b41ae --- /dev/null +++ b/internal/api/gemini_output_test.go @@ -0,0 +1,138 @@ +package api + +import ( + "encoding/json" + "reflect" + "testing" + + "github.com/Mag1cFall/AIStudio2API/internal/aistudio" +) + +func TestGeminiOutputParts_MergeTextAndReasoning(t *testing.T) { + result := generationResult{ + events: []aistudio.Event{ + {Kind: aistudio.EventReasoning, Text: "I think "}, + {Kind: aistudio.EventReasoning, Text: "therefore "}, + {Kind: aistudio.EventReasoning, Text: "I am."}, + {Kind: aistudio.EventThoughtSignature, ThoughtSignature: "sig_abc123"}, + {Kind: aistudio.EventText, Text: "Hello "}, + {Kind: aistudio.EventText, Text: "world!"}, + }, + } + + parts := geminiOutputParts(result) + if len(parts) != 2 { + t.Fatalf("expected 2 parts, got %d: %+v", len(parts), parts) + } + + // First part: merged thought + expectedThought := map[string]any{ + "thought": true, + "text": "I think therefore I am.", + "thoughtSignature": "sig_abc123", + } + if !reflect.DeepEqual(parts[0], expectedThought) { + t.Errorf("parts[0] = %+v, want %+v", parts[0], expectedThought) + } + + // Second part: merged text + expectedText := map[string]any{ + "text": "Hello world!", + } + if !reflect.DeepEqual(parts[1], expectedText) { + t.Errorf("parts[1] = %+v, want %+v", parts[1], expectedText) + } +} + +func TestGeminiOutputParts_TextOnly(t *testing.T) { + result := generationResult{ + events: []aistudio.Event{ + {Kind: aistudio.EventText, Text: "Chunk 1 "}, + {Kind: aistudio.EventText, Text: "Chunk 2 "}, + {Kind: aistudio.EventText, Text: "Chunk 3"}, + }, + } + + parts := geminiOutputParts(result) + if len(parts) != 1 { + t.Fatalf("expected 1 part, got %d: %+v", len(parts), parts) + } + + expectedText := map[string]any{ + "text": "Chunk 1 Chunk 2 Chunk 3", + } + if !reflect.DeepEqual(parts[0], expectedText) { + t.Errorf("parts[0] = %+v, want %+v", parts[0], expectedText) + } +} + +func TestGeminiOutputParts_LeadingThoughtSignature(t *testing.T) { + result := generationResult{ + events: []aistudio.Event{ + {Kind: aistudio.EventThoughtSignature, ThoughtSignature: "sig_pre"}, + {Kind: aistudio.EventReasoning, Text: "Thinking step."}, + {Kind: aistudio.EventText, Text: "Final text."}, + }, + } + + parts := geminiOutputParts(result) + if len(parts) != 2 { + t.Fatalf("expected 2 parts, got %d: %+v", len(parts), parts) + } + + if parts[0]["thoughtSignature"] != "sig_pre" { + t.Errorf("expected thoughtSignature 'sig_pre' on parts[0], got %v", parts[0]["thoughtSignature"]) + } + if parts[0]["text"] != "Thinking step." { + t.Errorf("expected text 'Thinking step.', got %v", parts[0]["text"]) + } + if parts[1]["text"] != "Final text." { + t.Errorf("expected text 'Final text.', got %v", parts[1]["text"]) + } +} + +func TestGeminiOutputParts_StandaloneThoughtSignatureFallback(t *testing.T) { + result := generationResult{ + events: []aistudio.Event{ + {Kind: aistudio.EventThoughtSignature, ThoughtSignature: "sig_only"}, + }, + } + + parts := geminiOutputParts(result) + if len(parts) != 1 { + t.Fatalf("expected 1 part, got %d: %+v", len(parts), parts) + } + if parts[0]["thoughtSignature"] != "sig_only" { + t.Errorf("expected thoughtSignature 'sig_only', got %v", parts[0]["thoughtSignature"]) + } +} + +func TestGeminiOutputParts_WithToolCall(t *testing.T) { + call := aistudio.FunctionCall{ + ID: "call_1", + Name: "get_weather", + Arguments: json.RawMessage(`{"location":"Tokyo"}`), + } + result := generationResult{ + events: []aistudio.Event{ + {Kind: aistudio.EventReasoning, Text: "Need weather for Tokyo."}, + {Kind: aistudio.EventThoughtSignature, ThoughtSignature: "sig_call"}, + {Kind: aistudio.EventToolCall, ToolCall: &call}, + }, + } + + parts := geminiOutputParts(result) + if len(parts) != 2 { + t.Fatalf("expected 2 parts, got %d: %+v", len(parts), parts) + } + + if parts[0]["thought"] != true || parts[0]["text"] != "Need weather for Tokyo." { + t.Errorf("unexpected parts[0]: %+v", parts[0]) + } + if parts[0]["thoughtSignature"] != "sig_call" { + t.Errorf("expected sig_call on parts[0], got %v", parts[0]["thoughtSignature"]) + } + if parts[1]["functionCall"] == nil { + t.Errorf("expected functionCall on parts[1], got %+v", parts[1]) + } +}