diff --git a/dto/gemini.go b/dto/gemini.go index 686be06f..63c99b86 100644 --- a/dto/gemini.go +++ b/dto/gemini.go @@ -121,6 +121,11 @@ func (r *GeminiChatRequest) IsStream(c *gin.Context) bool { if c.Query("alt") == "sse" { return true } + // Native Gemini API uses URL action to indicate streaming: + // /v1beta/models/{model}:streamGenerateContent + if strings.Contains(c.Request.URL.Path, "streamGenerateContent") { + return true + } return false } diff --git a/dto/gemini_isstream_test.go b/dto/gemini_isstream_test.go new file mode 100644 index 00000000..ecd461b9 --- /dev/null +++ b/dto/gemini_isstream_test.go @@ -0,0 +1,73 @@ +package dto + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" +) + +func TestGeminiChatRequest_IsStream(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + path string + query string + expected bool + }{ + { + name: "streamGenerateContent without alt=sse", + path: "/v1beta/models/gemini-2.0-flash:streamGenerateContent", + query: "key=sk-xxx", + expected: true, + }, + { + name: "streamGenerateContent with alt=sse", + path: "/v1beta/models/gemini-2.0-flash:streamGenerateContent", + query: "alt=sse&key=sk-xxx", + expected: true, + }, + { + name: "generateContent without alt=sse", + path: "/v1beta/models/gemini-2.0-flash:generateContent", + query: "key=sk-xxx", + expected: false, + }, + { + name: "generateContent with alt=sse", + path: "/v1beta/models/gemini-2.0-flash:generateContent", + query: "alt=sse", + expected: true, + }, + { + name: "GenerateContent capitalized", + path: "/v1beta/models/gemini-2.0-flash:GenerateContent", + query: "key=sk-xxx", + expected: false, + }, + { + name: "embedding path", + path: "/v1beta/models/gemini-2.0-flash:embedContent", + query: "", + expected: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + url := tt.path + if tt.query != "" { + url += "?" + tt.query + } + c.Request, _ = http.NewRequest("POST", url, nil) + + req := &GeminiChatRequest{} + assert.Equal(t, tt.expected, req.IsStream(c)) + }) + } +}