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,