feat(vertexai): add Candidate.FunctionCalls accessor (#10149)

eliben · web-flow · commit 6c76a67af1b6 · 2024-05-13T22:38:22.000Z
Convenience accessor to align with the Google AI SDK
diff --git a/vertexai/genai/client_test.go b/vertexai/genai/client_test.go
@@ -257,6 +257,10 @@ func TestLive(t *testing.T) {
 				t.Fatal(err)
 			}
 			part := res.Candidates[0].Content.Parts[0]
+			funcalls := res.Candidates[0].FunctionCalls()
+			if len(funcalls) != 1 {
+				t.Fatalf("got %d FunctionCalls, want 1", len(funcalls))
+			}
 			funcall, ok := part.(FunctionCall)
 			if !ok {
 				t.Fatalf("want FunctionCall, got %T", part)
diff --git a/vertexai/genai/content.go b/vertexai/genai/content.go
@@ -134,3 +134,17 @@ func (c *GenerationConfig) SetTopP(x float32) { c.TopP = &x }
 
 // SetTopK sets the TopK field.
 func (c *GenerationConfig) SetTopK(x int32) { c.TopK = &x }
+
+// FunctionCalls return all the FunctionCall parts in the candidate.
+func (c *Candidate) FunctionCalls() []FunctionCall {
+	if c.Content == nil {
+		return nil
+	}
+	var fcs []FunctionCall
+	for _, p := range c.Content.Parts {
+		if fc, ok := p.(FunctionCall); ok {
+			fcs = append(fcs, fc)
+		}
+	}
+	return fcs
+}
