chore(vertexai): use unary RPC for GenerateContent (#9391)

googleapis · Feb 8, 2024 · 1879551 · 1879551
1 parent e8bd85b
commit 1879551
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 14 deletions.
diff --git a/vertexai/genai/client.go b/vertexai/genai/client.go
@@ -116,20 +116,12 @@ func (m *GenerativeModel) GenerateContentStream(ctx context.Context, parts ...Pa
 }
 
 func (m *GenerativeModel) generateContent(ctx context.Context, req *pb.GenerateContentRequest) (*GenerateContentResponse, error) {
-	streamClient, err := m.c.c.StreamGenerateContent(ctx, req)
-	iter := &GenerateContentResponseIterator{
-		sc:  streamClient,
-		err: err,
-	}
-	for {
-		_, err := iter.Next()
-		if err == iterator.Done {
-			return iter.merged, nil
-		}
-		if err != nil {
-			return nil, err
-		}
+	res, err := m.c.c.GenerateContent(ctx, req)
+
+	if err != nil {
+		return nil, err
 	}
+	return protoToResponse(res)
 }
 
 func (m *GenerativeModel) newGenerateContentRequest(contents ...*Content) *pb.GenerateContentRequest {

diff --git a/vertexai/genai/client_test.go b/vertexai/genai/client_test.go
@@ -259,7 +259,7 @@ func TestLive(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		checkMatch(t, responseString(res), "(it's}|weather) .*cold")
+		checkMatch(t, responseString(res), "(it's|it is|weather) .*cold")
 	})
 }