From 66cb2acfdabc75fb6a755a263ee9bfa6649d45ed Mon Sep 17 00:00:00 2001 From: MikeRez0 <45215419+MikeRez0@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:47:19 -0700 Subject: [PATCH] fix(go/plugins/compat_oai): preserve request on streaming responses generateStream never set ModelResponse.Request, and convertChatCompletionToModelResponse defaulted it to an empty &ai.ModelRequest{}. Streaming responses therefore reported an empty message list, so ModelResponse.History() returned only the model reply and dropped the entire input conversation. The non-streaming path was unaffected because generateComplete already assigned resp.Request. Thread the originating request through generateStream and assign it after conversion, mirroring generateComplete. Drop the misleading default from the converter so it stays a pure conversion and callers own the request. Fixes #4683 Co-authored-by: Alex Pascal --- go/plugins/compat_oai/generate.go | 8 ++-- go/plugins/compat_oai/generate_test.go | 64 ++++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/go/plugins/compat_oai/generate.go b/go/plugins/compat_oai/generate.go index c4bfea4773..3333226aae 100644 --- a/go/plugins/compat_oai/generate.go +++ b/go/plugins/compat_oai/generate.go @@ -330,7 +330,7 @@ func (g *ModelGenerator) Generate(ctx context.Context, req *ai.ModelRequest, han } if handleChunk != nil { - return g.generateStream(ctx, handleChunk) + return g.generateStream(ctx, req, handleChunk) } return g.generateComplete(ctx, req) } @@ -393,7 +393,7 @@ func concatenateReasoningContent(parts []*ai.Part) string { } // generateStream generates a streaming model response -func (g *ModelGenerator) generateStream(ctx context.Context, handleChunk func(context.Context, *ai.ModelResponseChunk) error) (*ai.ModelResponse, error) { +func (g *ModelGenerator) generateStream(ctx context.Context, req *ai.ModelRequest, handleChunk func(context.Context, *ai.ModelResponseChunk) error) (*ai.ModelResponse, error) { stream := g.client.Chat.Completions.NewStreaming(ctx, *g.request) defer func() { _ = stream.Close() @@ -463,6 +463,7 @@ func (g *ModelGenerator) generateStream(ctx context.Context, handleChunk func(co resp.Message.Content..., ) } + resp.Request = req return resp, nil } @@ -525,8 +526,7 @@ func convertChatCompletionToModelResponse(completion *openai.ChatCompletion) (*a } resp := &ai.ModelResponse{ - Request: &ai.ModelRequest{}, - Usage: usage, + Usage: usage, Message: &ai.Message{ Role: ai.RoleModel, Content: make([]*ai.Part, 0), diff --git a/go/plugins/compat_oai/generate_test.go b/go/plugins/compat_oai/generate_test.go index 1074cb3c99..12e87d4096 100644 --- a/go/plugins/compat_oai/generate_test.go +++ b/go/plugins/compat_oai/generate_test.go @@ -15,11 +15,16 @@ package compat_oai import ( + "context" "encoding/json" + "io" + "net/http" + "net/http/httptest" "testing" "github.com/firebase/genkit/go/ai" "github.com/openai/openai-go" + "github.com/openai/openai-go/option" ) func TestConvertChatCompletionToModelResponseReasoningContent(t *testing.T) { @@ -65,6 +70,65 @@ func newGen() *ModelGenerator { return NewModelGenerator((*openai.Client)(nil), "test-model") } +// newStubGen returns a ModelGenerator pointed at a stub OpenAI-compatible +// endpoint that replies with body, so both generate paths can be exercised +// without a live API key. +func newStubGen(t *testing.T, contentType, body string) *ModelGenerator { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", contentType) + io.WriteString(w, body) + })) + t.Cleanup(srv.Close) + client := openai.NewClient(option.WithBaseURL(srv.URL), option.WithAPIKey("stub")) + return NewModelGenerator(&client, "test-model") +} + +// Regression test for #4683: the streaming path used to leave Request as an +// empty &ai.ModelRequest{}, so History() dropped every input message. +func TestGeneratePreservesRequest(t *testing.T) { + messages := []*ai.Message{ + ai.NewUserTextMessage("first user turn"), + ai.NewModelTextMessage("first model turn"), + ai.NewUserTextMessage("second user turn"), + } + + const streamBody = `data: {"id":"1","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{"role":"assistant","content":"answer"},"finish_reason":null}]} + +data: {"id":"1","object":"chat.completion.chunk","created":1,"model":"test-model","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]} + +data: [DONE] + +` + const completeBody = `{"id":"1","object":"chat.completion","created":1,"model":"test-model","choices":[{"index":0,"message":{"role":"assistant","content":"answer"},"finish_reason":"stop"}]}` + + for _, tc := range []struct { + name string + contentType string + body string + handleChunk func(context.Context, *ai.ModelResponseChunk) error + }{ + {"streaming", "text/event-stream", streamBody, func(context.Context, *ai.ModelResponseChunk) error { return nil }}, + {"complete", "application/json", completeBody, nil}, + } { + t.Run(tc.name, func(t *testing.T) { + req := &ai.ModelRequest{Messages: messages} + resp, err := newStubGen(t, tc.contentType, tc.body). + WithMessages(messages). + Generate(context.Background(), req, tc.handleChunk) + if err != nil { + t.Fatalf("Generate() error = %v", err) + } + if resp.Request != req { + t.Fatalf("Request = %#v, want the originating request", resp.Request) + } + if got := len(resp.History()); got != len(messages)+1 { + t.Errorf("len(History()) = %d, want %d (inputs plus model reply)", got, len(messages)+1) + } + }) + } +} + func TestConvertChatCompletionToModelResponseProviderFinishReasons(t *testing.T) { for finishReason, want := range map[string]ai.FinishReason{ "sensitive": ai.FinishReasonBlocked,