From df035d07569ddb324914f78085022f218aab4816 Mon Sep 17 00:00:00 2001 From: Keith Randall Date: Fri, 7 Jun 2024 16:50:00 -0700 Subject: [PATCH 1/2] test that streaming and final answer contain the same text. This should already be the case, just testing it. --- go/plugins/googleai/googleai_test.go | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/go/plugins/googleai/googleai_test.go b/go/plugins/googleai/googleai_test.go index 57df79c4c9..cc91693b16 100644 --- a/go/plugins/googleai/googleai_test.go +++ b/go/plugins/googleai/googleai_test.go @@ -141,7 +141,7 @@ func TestLive(t *testing.T) { out := "" parts := 0 g := googleai.Model(generativeModel) - _, err = ai.Generate(ctx, g, req, func(ctx context.Context, c *ai.GenerateResponseChunk) error { + final, err := ai.Generate(ctx, g, req, func(ctx context.Context, c *ai.GenerateResponseChunk) error { parts++ out += c.Content[0].Text return nil @@ -149,6 +149,15 @@ func TestLive(t *testing.T) { if err != nil { t.Fatal(err) } + out2 := "" + for _, c := range final.Candidates { + for _, p := range c.Message.Content { + out2 += p.Text + } + } + if out != out2 { + t.Errorf("streaming and final should contain the same text.\nstreaming:%s\nfinal:%s", out, out2) + } const want = "Golden" if !strings.Contains(out, want) { t.Errorf("got %q, expecting it to contain %q", out, want) From 83686278132c7df8f34d008df4a2a573275730ff Mon Sep 17 00:00:00 2001 From: Keith Randall Date: Mon, 10 Jun 2024 13:29:17 -0700 Subject: [PATCH 2/2] Just use the one candidate to check streaming results --- go/plugins/googleai/googleai_test.go | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/go/plugins/googleai/googleai_test.go b/go/plugins/googleai/googleai_test.go index cc91693b16..9b68654751 100644 --- a/go/plugins/googleai/googleai_test.go +++ b/go/plugins/googleai/googleai_test.go @@ -150,10 +150,8 @@ func TestLive(t *testing.T) { t.Fatal(err) } out2 := "" - for _, c := range final.Candidates { - for _, p := range c.Message.Content { - out2 += p.Text - } + for _, p := range final.Candidates[0].Message.Content { + out2 += p.Text } if out != out2 { t.Errorf("streaming and final should contain the same text.\nstreaming:%s\nfinal:%s", out, out2)