From dff590bf9a2b1466ad5da29649facf41618c12cf Mon Sep 17 00:00:00 2001 From: Andrey Nering Date: Tue, 9 Sep 2025 17:53:20 -0300 Subject: [PATCH] feat(gemini): fixes for streaming + thinking --- google/google.go | 124 +++++++++++++++--- .../TestTool/google-gemini-2.5-flash.yaml | 12 +- .../TestTool/google-gemini-2.5-pro.yaml | 12 +- 3 files changed, 117 insertions(+), 31 deletions(-) diff --git a/google/google.go b/google/google.go index eef7f43d34b1abb5a5d28371916048a25cac39c3..558125a7892a00807e83896bf2781abaf8dc2ef1 100644 --- a/google/google.go +++ b/google/google.go @@ -431,9 +431,13 @@ func (g *languageModel) Stream(ctx context.Context, call ai.Call) (ai.StreamResp var currentContent string var toolCalls []ai.ToolCallContent var isActiveText bool + var isActiveReasoning bool + var blockCounter int + var currentTextBlockID string + var currentReasoningBlockID string var usage ai.Usage + var lastFinishReason ai.FinishReason - // Stream the response for resp, err := range chat.SendMessageStream(ctx, depointerSlice(lastMessage.Parts)...) { if err != nil { yield(ai.StreamPart{ @@ -449,30 +453,91 @@ func (g *languageModel) Stream(ctx context.Context, call ai.Call) (ai.StreamResp case part.Text != "": delta := part.Text if delta != "" { - if !isActiveText { - isActiveText = true + // Check if this is a reasoning/thought part + if part.Thought { + // End any active text block before starting reasoning + if isActiveText { + isActiveText = false + if !yield(ai.StreamPart{ + Type: ai.StreamPartTypeTextEnd, + ID: currentTextBlockID, + }) { + return + } + } + + // Start new reasoning block if not already active + if !isActiveReasoning { + isActiveReasoning = true + currentReasoningBlockID = fmt.Sprintf("%d", blockCounter) + blockCounter++ + if !yield(ai.StreamPart{ + Type: ai.StreamPartTypeReasoningStart, + ID: currentReasoningBlockID, + }) { + return + } + } + if !yield(ai.StreamPart{ - Type: ai.StreamPartTypeTextStart, - ID: "0", + Type: ai.StreamPartTypeReasoningDelta, + ID: currentReasoningBlockID, + Delta: delta, }) { return } + } else { + // Regular text part + // End any active reasoning block before starting text + if isActiveReasoning { + isActiveReasoning = false + if !yield(ai.StreamPart{ + Type: ai.StreamPartTypeReasoningEnd, + ID: currentReasoningBlockID, + }) { + return + } + } + + // Start new text block if not already active + if !isActiveText { + isActiveText = true + currentTextBlockID = fmt.Sprintf("%d", blockCounter) + blockCounter++ + if !yield(ai.StreamPart{ + Type: ai.StreamPartTypeTextStart, + ID: currentTextBlockID, + }) { + return + } + } + + if !yield(ai.StreamPart{ + Type: ai.StreamPartTypeTextDelta, + ID: currentTextBlockID, + Delta: delta, + }) { + return + } + currentContent += delta } - if !yield(ai.StreamPart{ - Type: ai.StreamPartTypeTextDelta, - ID: "0", - Delta: delta, - }) { - return - } - currentContent += delta } case part.FunctionCall != nil: + // End any active text or reasoning blocks if isActiveText { isActiveText = false if !yield(ai.StreamPart{ Type: ai.StreamPartTypeTextEnd, - ID: "0", + ID: currentTextBlockID, + }) { + return + } + } + if isActiveReasoning { + isActiveReasoning = false + if !yield(ai.StreamPart{ + Type: ai.StreamPartTypeReasoningEnd, + ID: currentReasoningBlockID, }) { return } @@ -535,20 +600,35 @@ func (g *languageModel) Stream(ctx context.Context, call ai.Call) (ai.StreamResp if resp.UsageMetadata != nil { usage = mapUsage(resp.UsageMetadata) } + + if len(resp.Candidates) > 0 && resp.Candidates[0].FinishReason != "" { + lastFinishReason = mapFinishReason(resp.Candidates[0].FinishReason) + } } + // Close any open blocks before finishing if isActiveText { if !yield(ai.StreamPart{ Type: ai.StreamPartTypeTextEnd, - ID: "0", + ID: currentTextBlockID, + }) { + return + } + } + if isActiveReasoning { + if !yield(ai.StreamPart{ + Type: ai.StreamPartTypeReasoningEnd, + ID: currentReasoningBlockID, }) { return } } - finishReason := ai.FinishReasonStop + finishReason := lastFinishReason if len(toolCalls) > 0 { finishReason = ai.FinishReasonToolCalls + } else if finishReason == "" { + finishReason = ai.FinishReasonStop } yield(ai.StreamPart{ @@ -720,21 +800,27 @@ func mapResponse(response *genai.GenerateContentResponse, warnings []ai.CallWarn for _, part := range candidate.Content.Parts { switch { case part.Text != "": - content = append(content, ai.TextContent{Text: part.Text}) + if part.Thought { + content = append(content, ai.ReasoningContent{Text: part.Text}) + } else { + content = append(content, ai.TextContent{Text: part.Text}) + } case part.FunctionCall != nil: input, err := json.Marshal(part.FunctionCall.Args) if err != nil { return nil, err } + toolCallID := cmp.Or(part.FunctionCall.ID, part.FunctionCall.Name, uuid.NewString()) content = append(content, ai.ToolCallContent{ - ToolCallID: part.FunctionCall.ID, + ToolCallID: toolCallID, ToolName: part.FunctionCall.Name, Input: string(input), ProviderExecuted: false, }) hasToolCalls = true default: - return nil, fmt.Errorf("not implemented part type") + // Silently skip unknown part types instead of erroring + // This allows for forward compatibility with new part types } } diff --git a/providertests/testdata/TestTool/google-gemini-2.5-flash.yaml b/providertests/testdata/TestTool/google-gemini-2.5-flash.yaml index 9d9b8efbecb56fda34f8cafcc64ed3f3f0719db2..ba50fa977c18d902abdea70e640d898f1fb22e3b 100644 --- a/providertests/testdata/TestTool/google-gemini-2.5-flash.yaml +++ b/providertests/testdata/TestTool/google-gemini-2.5-flash.yaml @@ -22,21 +22,21 @@ interactions: proto_minor: 0 content_length: -1 uncompressed: true - body: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"functionCall\": {\n \"name\": \"weather\",\n \"args\": {\n \"location\": \"Florence\"\n }\n },\n \"thoughtSignature\": \"CpUCAVSoXO7jOuZmcU+HAEV/vBr3c+MDukNpwF7G8D2V1FB5+SS3o4GLSA4PqaT+0qwv6xtHan6ExItRzcphhhqVD4JUK13scMvwfy7r2r3KjbyFLq/JqQjMY5KHJXxuYu5GKuwKyYJHhHTjP0p62ORdPaUZg3umBcdZM3nD9bMneU6zIPBB4l/t0s69c/+0nooCeYI+r9s7NpI2AvqQNKjaiIeJrC1W1Qw1cUAB7L/l1ZJjqiZ5CiBTiW4WJ4pxhP9WT3vwqqv9SFi1FpvycBiLCUmLkH/rHmHHoq2ExxRUANcgqXiF0q3beAP+W1iudaGHkz9023CrN6YR5oQj6ABqYTYVDYTE8qG95MWYhWk6suiv9SD/MA==\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 52,\n \"candidatesTokenCount\": 13,\n \"totalTokenCount\": 121,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 52\n }\n ],\n \"thoughtsTokenCount\": 56\n },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"_ki7aPmbEbLhz7IP8vvqoQk\"\n}\n" + body: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"functionCall\": {\n \"name\": \"weather\",\n \"args\": {\n \"location\": \"Florence\"\n }\n },\n \"thoughtSignature\": \"CuoBAdHtim8Vp9NlNShmGo0RDhWEoSVPAdPgu2WPmt3EVCgITsoxv35IRyOao5Xl6E7MOMlrQQeQu/agp4wJ2bkeEHEZbPf7Iosdec5nLrwjXwf0WRN/LFik08aPKh3VeKN1F4wcqc6eSYEHQgQdLtJ6Ke0DrXCQWUuPgykk5QR4o12ZHZrtW3Suj0xizUOWf7/vMoqRoJEHVVHL7qznMEH1OMCbwRMuR6pwCkSgthshvEvx6kpk4Nt1Ymf9M8cCQuQeM4LvdPpT0UteZ6ZLqwbaYF4OPistAL19QI5lelFaMDoGZx5Ian/aiH9j\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 52,\n \"candidatesTokenCount\": 13,\n \"totalTokenCount\": 111,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 52\n }\n ],\n \"thoughtsTokenCount\": 46\n },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"-ZPAaOeBIMmtz7IP6MLh4Aw\"\n}\n" headers: Content-Type: - application/json; charset=UTF-8 status: 200 OK code: 200 - duration: 1.063079458s + duration: 1.84992975s - id: 1 request: proto: HTTP/1.1 proto_major: 1 proto_minor: 1 - content_length: 670 + content_length: 700 host: generativelanguage.googleapis.com - body: "{\"contents\":[{\"parts\":[{\"text\":\"What's the weather in Florence?\"}],\"role\":\"user\"},{\"parts\":[{\"functionCall\":{\"args\":{\"location\":\"Florence\"},\"name\":\"weather\"}}],\"role\":\"model\"},{\"parts\":[{\"functionResponse\":{\"name\":\"weather\",\"response\":{\"result\":\"40 C\"}}}],\"role\":\"user\"}],\"generationConfig\":{},\"systemInstruction\":{\"parts\":[{\"text\":\"You are a helpful assistant\"}],\"role\":\"user\"},\"toolConfig\":{\"functionCallingConfig\":{\"mode\":\"AUTO\"}},\"tools\":[{\"functionDeclarations\":[{\"description\":\"Get weather information for a location\",\"name\":\"weather\",\"parameters\":{\"properties\":{\"location\":{\"description\":\"the city\",\"type\":\"STRING\"}},\"required\":[\"location\"],\"type\":\"OBJECT\"}}]}]}\n" + body: "{\"contents\":[{\"parts\":[{\"text\":\"What's the weather in Florence?\"}],\"role\":\"user\"},{\"parts\":[{\"functionCall\":{\"args\":{\"location\":\"Florence\"},\"id\":\"weather\",\"name\":\"weather\"}}],\"role\":\"model\"},{\"parts\":[{\"functionResponse\":{\"id\":\"weather\",\"name\":\"weather\",\"response\":{\"result\":\"40 C\"}}}],\"role\":\"user\"}],\"generationConfig\":{},\"systemInstruction\":{\"parts\":[{\"text\":\"You are a helpful assistant\"}],\"role\":\"user\"},\"toolConfig\":{\"functionCallingConfig\":{\"mode\":\"AUTO\"}},\"tools\":[{\"functionDeclarations\":[{\"description\":\"Get weather information for a location\",\"name\":\"weather\",\"parameters\":{\"properties\":{\"location\":{\"description\":\"the city\",\"type\":\"STRING\"}},\"required\":[\"location\"],\"type\":\"OBJECT\"}}]}]}\n" headers: Content-Type: - application/json @@ -50,10 +50,10 @@ interactions: proto_minor: 0 content_length: -1 uncompressed: true - body: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \"The weather in Florence is 40 C.\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 80,\n \"candidatesTokenCount\": 10,\n \"totalTokenCount\": 90,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 80\n }\n ]\n },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"_ki7aN-APaaHz7IPrL60iQE\"\n}\n" + body: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \"The weather in Florence is 40 C.\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 80,\n \"candidatesTokenCount\": 10,\n \"totalTokenCount\": 90,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 80\n }\n ]\n },\n \"modelVersion\": \"gemini-2.5-flash\",\n \"responseId\": \"-pPAaOmKDa7Vz7IPi7fmgQ0\"\n}\n" headers: Content-Type: - application/json; charset=UTF-8 status: 200 OK code: 200 - duration: 678.656417ms + duration: 739.748334ms diff --git a/providertests/testdata/TestTool/google-gemini-2.5-pro.yaml b/providertests/testdata/TestTool/google-gemini-2.5-pro.yaml index 70f8bf364924ac2d1e69daeae3db02afc543958e..6671cdaa35c06eb6d5f3f424b9eefad59ea2a596 100644 --- a/providertests/testdata/TestTool/google-gemini-2.5-pro.yaml +++ b/providertests/testdata/TestTool/google-gemini-2.5-pro.yaml @@ -22,21 +22,21 @@ interactions: proto_minor: 0 content_length: -1 uncompressed: true - body: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"functionCall\": {\n \"name\": \"weather\",\n \"args\": {\n \"location\": \"Florence\"\n }\n },\n \"thoughtSignature\": \"CqQEAdHtim+OnMkt9f1IrDaI9fiX/zDMQre89tEPHMT7hUY8LCgMhigmIoih3mEKqTCnk7F2prHKcg6UnerobTbNBTy4r6nFCvwRAs2DZkeZ9uK+xGqYRPGTEZaFS5oUlHoANUH+bSHA1kmwgjxy3c8VrZTwnLWyQAZw6pabJqrRJe/G5ADjNX/lMCDqM/rNRiqRaZqstv0K8QjwD4pz9rAgwP7BeTu9moPXzUwv9lY+7RP1J39DQnPjlRA++5h2F3yjIgO7NKiJp+aXCKPlCYb+o6uPskH1tlahXQdnNydrFhtHXMe7+5G76fXuR+vAxz9McjLbploeLqFZgPekHWOcY6EU0lcrm/39YILOy3iyAiu42ZbTfoOVdCsyttcfJPjyZ8U75fyIX/Z+45ky8sE1Nw7+cv4Frm7sOH63lSu3uA0ml3xkZSkJ3nZKtMSRCUBJOLfndMy2gAzeo2QNeLuUPksvg318ehZP48oTnDipo+Zr5575ycekNDLueQn3a7tPUw+uckudUzQFLzJiqs01qq0+1hhxGeHUHe19UcWBTZ/AbrJKNOhUA7TSU2tbza/bf7xGBvDpwxmlDPzvkiYH83nBPUUZUo+jf3FOn+lwI2PzOlJjN0GaGyIp/GY2tPrToElqZnfzsgxp/cBb7ZlbZIwP5Sw3uepfu3Y3RhFaZk4n9yotZWWITfyQFHRQT/fmkPyuYi2zPn3R+EEf9DL5J5oHCeE=\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 52,\n \"candidatesTokenCount\": 13,\n \"totalTokenCount\": 186,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 52\n }\n ],\n \"thoughtsTokenCount\": 121\n },\n \"modelVersion\": \"gemini-2.5-pro\",\n \"responseId\": \"aIvAaM70OfuNmtkPjZGvoAM\"\n}\n" + body: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"functionCall\": {\n \"name\": \"weather\",\n \"args\": {\n \"location\": \"Florence\"\n }\n },\n \"thoughtSignature\": \"CpkEAdHtim+EtnrVkWr/1MyuKGUtmnac2ihex9vY7ws2/Y2MPtaieXMdvYOE0YwAA3oV6VELDoHoz4++dx3qGkTwnxMachj5W+EuyQ0FaXtPP/OrMZDEmGVjhQZ0ADj9Q/8BJ/J6KonretAC3baPhvUCCmnQy4/5UPoVkcmQJkusUABUnCQoXQMv1nIEoQZDa3OWw3kZzHrIlk5/1/vaWqEVPtQ76EHOeQic3C3waGzDtmbqV4f9Ygy2ImLNfB8DlPq3JBfG/gkGaxA6RiOW3yRoRBdIxpsaPmzEyt1AZpNzJR0Jv6HuzjM1bv4tcjcTQwXZagjkFC/TkJ8wwDhhG0L78rCFguZZJhBVULYWXiBpIkWzhemreSHE5q8X5nzNxV2BIddYXEMWjy31IimBPYzmH32ithFjpDuW2T5I/4MYhCB8BlBw++ii5xbIL/WfFDzJwRXaWNx2/cmcT9Un+AruUfHFDUGCFeOCqq/1daPk1qQkyZY6d3T7dx9dt1vKICqXSETtoZMfuBdre931vOmklm7oRrWROKIz/tIBXU4NvWf9wT1dDWegusxLMl76UHSb2//ArXtvj1FbY2loas7VpxSW0zIeRmsGQhY5/hskKk917ppNKNvUSqX7OlHP9rW+/OIuxQpGj3t8PNtoqGQsO1FDrYhNq/M/56jVt6rkkYzRffno0tvDiJa2fySVFo7F0m2XNgGyavbQ\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 52,\n \"candidatesTokenCount\": 13,\n \"totalTokenCount\": 186,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 52\n }\n ],\n \"thoughtsTokenCount\": 121\n },\n \"modelVersion\": \"gemini-2.5-pro\",\n \"responseId\": \"_JPAaPK5ML_Vz7IP-4KM-QQ\"\n}\n" headers: Content-Type: - application/json; charset=UTF-8 status: 200 OK code: 200 - duration: 2.122576125s + duration: 2.469309625s - id: 1 request: proto: HTTP/1.1 proto_major: 1 proto_minor: 1 - content_length: 670 + content_length: 700 host: generativelanguage.googleapis.com - body: "{\"contents\":[{\"parts\":[{\"text\":\"What's the weather in Florence?\"}],\"role\":\"user\"},{\"parts\":[{\"functionCall\":{\"args\":{\"location\":\"Florence\"},\"name\":\"weather\"}}],\"role\":\"model\"},{\"parts\":[{\"functionResponse\":{\"name\":\"weather\",\"response\":{\"result\":\"40 C\"}}}],\"role\":\"user\"}],\"generationConfig\":{},\"systemInstruction\":{\"parts\":[{\"text\":\"You are a helpful assistant\"}],\"role\":\"user\"},\"toolConfig\":{\"functionCallingConfig\":{\"mode\":\"AUTO\"}},\"tools\":[{\"functionDeclarations\":[{\"description\":\"Get weather information for a location\",\"name\":\"weather\",\"parameters\":{\"properties\":{\"location\":{\"description\":\"the city\",\"type\":\"STRING\"}},\"required\":[\"location\"],\"type\":\"OBJECT\"}}]}]}\n" + body: "{\"contents\":[{\"parts\":[{\"text\":\"What's the weather in Florence?\"}],\"role\":\"user\"},{\"parts\":[{\"functionCall\":{\"args\":{\"location\":\"Florence\"},\"id\":\"weather\",\"name\":\"weather\"}}],\"role\":\"model\"},{\"parts\":[{\"functionResponse\":{\"id\":\"weather\",\"name\":\"weather\",\"response\":{\"result\":\"40 C\"}}}],\"role\":\"user\"}],\"generationConfig\":{},\"systemInstruction\":{\"parts\":[{\"text\":\"You are a helpful assistant\"}],\"role\":\"user\"},\"toolConfig\":{\"functionCallingConfig\":{\"mode\":\"AUTO\"}},\"tools\":[{\"functionDeclarations\":[{\"description\":\"Get weather information for a location\",\"name\":\"weather\",\"parameters\":{\"properties\":{\"location\":{\"description\":\"the city\",\"type\":\"STRING\"}},\"required\":[\"location\"],\"type\":\"OBJECT\"}}]}]}\n" headers: Content-Type: - application/json @@ -50,10 +50,10 @@ interactions: proto_minor: 0 content_length: -1 uncompressed: true - body: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \"The weather in Florence is 40 C. \\n\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 80,\n \"candidatesTokenCount\": 10,\n \"totalTokenCount\": 90,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 80\n }\n ]\n },\n \"modelVersion\": \"gemini-2.5-pro\",\n \"responseId\": \"aovAaM-PLLT4qtsPs9jkuQQ\"\n}\n" + body: "{\n \"candidates\": [\n {\n \"content\": {\n \"parts\": [\n {\n \"text\": \"The weather in Florence is 40 C. \\n\"\n }\n ],\n \"role\": \"model\"\n },\n \"finishReason\": \"STOP\",\n \"index\": 0\n }\n ],\n \"usageMetadata\": {\n \"promptTokenCount\": 80,\n \"candidatesTokenCount\": 10,\n \"totalTokenCount\": 90,\n \"promptTokensDetails\": [\n {\n \"modality\": \"TEXT\",\n \"tokenCount\": 80\n }\n ]\n },\n \"modelVersion\": \"gemini-2.5-pro\",\n \"responseId\": \"_pPAaJuZF-itz7IP-rWq8Aw\"\n}\n" headers: Content-Type: - application/json; charset=UTF-8 status: 200 OK code: 200 - duration: 1.854787917s + duration: 1.571857334s