Skip to content

Commit

Permalink
feat(vertexai): add Candidate.FunctionCalls accessor (#10149)
Browse files Browse the repository at this point in the history
Convenience accessor to align with the Google AI SDK
  • Loading branch information
eliben committed May 13, 2024
1 parent 1f9ecc4 commit 6c76a67
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
4 changes: 4 additions & 0 deletions vertexai/genai/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
14 changes: 14 additions & 0 deletions vertexai/genai/content.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}

0 comments on commit 6c76a67

Please sign in to comment.