genai: support constrained decoding (#82)

Let the user determine how the model should respond with FunctionCalls.

Author: jba

genai/client.go

@@ -105,6 +105,7 @@ type GenerativeModel struct {
 	GenerationConfig
 	SafetySettings []*SafetySetting
 	Tools          []*Tool
+	ToolConfig     *ToolConfig
 }
 
 // GenerativeModel creates a new instance of the named generative model.
@@ -168,6 +169,7 @@ func (m *GenerativeModel) newGenerateContentRequest(contents ...*Content) *pb.Ge
 		Contents:         support.TransformSlice(contents, (*Content).toProto),
 		SafetySettings:   support.TransformSlice(m.SafetySettings, (*SafetySetting).toProto),
 		Tools:            support.TransformSlice(m.Tools, (*Tool).toProto),
+		ToolConfig:       m.ToolConfig.toProto(),
 		GenerationConfig: m.GenerationConfig.toProto(),
 	}
 }

genai/client_test.go

@@ -297,7 +297,7 @@ func TestLive(t *testing.T) {
 	})
 	t.Run("tools", func(t *testing.T) {
 
-		weatherChat := func(t *testing.T, s *Schema) {
+		weatherChat := func(t *testing.T, s *Schema, fcm FunctionCallingMode) {
 			weatherTool := &Tool{
 				FunctionDeclarations: []*FunctionDeclaration{{
 					Name:        "CurrentWeather",
@@ -308,12 +308,23 @@ func TestLive(t *testing.T) {
 			model := client.GenerativeModel(*modelName)
 			model.SetTemperature(0)
 			model.Tools = []*Tool{weatherTool}
+			model.ToolConfig = &ToolConfig{
+				FunctionCallingConfig: &FunctionCallingConfig{
+					Mode: fcm,
+				},
+			}
 			session := model.StartChat()
 			res, err := session.SendMessage(ctx, Text("What is the weather like in New York?"))
 			if err != nil {
 				t.Fatal(err)
 			}
 			funcalls := res.Candidates[0].FunctionCalls()
+			if fcm == FunctionCallingNone {
+				if len(funcalls) != 0 {
+					t.Fatalf("got %d FunctionCalls, want 0", len(funcalls))
+				}
+				return
+			}
 			if len(funcalls) != 1 {
 				t.Fatalf("got %d FunctionCalls, want 1", len(funcalls))
 			}
@@ -339,22 +350,25 @@ func TestLive(t *testing.T) {
 			}
 			checkMatch(t, responseString(res), "(it's|it is|weather) .*cold")
 		}
-
-		t.Run("direct", func(t *testing.T) {
-			weatherChat(t, &Schema{
-				Type: TypeObject,
-				Properties: map[string]*Schema{
-					"location": {
-						Type:        TypeString,
-						Description: "The city and state, e.g. San Francisco, CA",
-					},
-					"unit": {
-						Type: TypeString,
-						Enum: []string{"celsius", "fahrenheit"},
-					},
+		schema := &Schema{
+			Type: TypeObject,
+			Properties: map[string]*Schema{
+				"location": {
+					Type:        TypeString,
+					Description: "The city and state, e.g. San Francisco, CA",
 				},
-				Required: []string{"location"},
-			})
+				"unit": {
+					Type: TypeString,
+					Enum: []string{"celsius", "fahrenheit"},
+				},
+			},
+			Required: []string{"location"},
+		}
+		t.Run("direct", func(t *testing.T) {
+			weatherChat(t, schema, FunctionCallingAuto)
+		})
+		t.Run("none", func(t *testing.T) {
+			weatherChat(t, schema, FunctionCallingNone)
 		})
 	})
 }

genai/config.yaml

@@ -117,6 +117,7 @@ types:
 
     # Types for function calling
     Tool:
+    ToolConfig:
     FunctionDeclaration:
     FunctionCall:
     FunctionResponse:
@@ -126,6 +127,15 @@ types:
       protoPrefix: Type_
       veneerPrefix: ''
 
+    FunctionCallingConfig:
+      doc: 'holds configuration for function calling.'
+
+    FunctionCallingConfig_Mode:
+      name: FunctionCallingMode
+      protoPrefix: FunctionCallingConfig
+      veneerPrefix: FunctionCalling
+      valueNames:
+        FunctionCallingConfig_MODE_UNSPECIFIED: FunctionCallingUnspecified
 
 
 

genai/example_test.go

@@ -394,6 +394,40 @@ func ExampleTool() {
 	printResponse(res)
 }
 
+func ExampleToolConifg() {
+	// This example shows how to affect how the model uses the tools provided to it.
+	// By setting the ToolConfig, you can disable function calling.
+
+	// Assume we have created a Model and have set its Tools field with some functions.
+	// See the Example for Tool for details.
+	var model *genai.GenerativeModel
+
+	// By default, the model will use the functions in its responses if it thinks they are
+	// relevant, by returning FunctionCall parts.
+	// Here we set the model's ToolConfig to disable function calling completely.
+	model.ToolConfig = &genai.ToolConfig{
+		FunctionCallingConfig: &genai.FunctionCallingConfig{
+			Mode: genai.FunctionCallingNone,
+		},
+	}
+
+	// Subsequent calls to ChatSession.SendMessage will not result in FunctionCall responses.
+	session := model.StartChat()
+	res, err := session.SendMessage(context.Background(), genai.Text("What is the weather like in New York?"))
+	if err != nil {
+		log.Fatal(err)
+	}
+	for _, part := range res.Candidates[0].Content.Parts {
+		if _, ok := part.(genai.FunctionCall); ok {
+			log.Fatal("did not expect FunctionCall")
+		}
+	}
+
+	// It is also possible to force a function call by using FunctionCallingAny
+	// instead of FunctionCallingNone. See the documentation for FunctionCallingMode
+	// for details.
+}
+
 func printResponse(resp *genai.GenerateContentResponse) {
 	for _, cand := range resp.Candidates {
 		if cand.Content != nil {

genai/generativelanguagepb_veneer.gen.go

@@ -400,6 +400,74 @@ func (FunctionCall) fromProto(p *pb.FunctionCall) *FunctionCall {
 	}
 }
 
+// FunctionCallingConfig holds configuration for function calling.
+type FunctionCallingConfig struct {
+	// Optional. Specifies the mode in which function calling should execute. If
+	// unspecified, the default value will be set to AUTO.
+	Mode FunctionCallingMode
+	// Optional. A set of function names that, when provided, limits the functions
+	// the model will call.
+	//
+	// This should only be set when the Mode is ANY. Function names
+	// should match [FunctionDeclaration.name]. With mode set to ANY, model will
+	// predict a function call from the set of function names provided.
+	AllowedFunctionNames []string
+}
+
+func (v *FunctionCallingConfig) toProto() *pb.FunctionCallingConfig {
+	if v == nil {
+		return nil
+	}
+	return &pb.FunctionCallingConfig{
+		Mode:                 pb.FunctionCallingConfig_Mode(v.Mode),
+		AllowedFunctionNames: v.AllowedFunctionNames,
+	}
+}
+
+func (FunctionCallingConfig) fromProto(p *pb.FunctionCallingConfig) *FunctionCallingConfig {
+	if p == nil {
+		return nil
+	}
+	return &FunctionCallingConfig{
+		Mode:                 FunctionCallingMode(p.Mode),
+		AllowedFunctionNames: p.AllowedFunctionNames,
+	}
+}
+
+// FunctionCallingMode is defines the execution behavior for function calling by defining the
+// execution mode.
+type FunctionCallingMode int32
+
+const (
+	// FunctionCallingUnspecified means unspecified function calling mode. This value should not be used.
+	FunctionCallingUnspecified FunctionCallingMode = 0
+	// FunctionCallingAuto means default model behavior, model decides to predict either a function call
+	// or a natural language repspose.
+	FunctionCallingAuto FunctionCallingMode = 1
+	// FunctionCallingAny means model is constrained to always predicting a function call only.
+	// If "allowed_function_names" are set, the predicted function call will be
+	// limited to any one of "allowed_function_names", else the predicted
+	// function call will be any one of the provided "function_declarations".
+	FunctionCallingAny FunctionCallingMode = 2
+	// FunctionCallingNone means model will not predict any function call. Model behavior is same as when
+	// not passing any function declarations.
+	FunctionCallingNone FunctionCallingMode = 3
+)
+
+var namesForFunctionCallingMode = map[FunctionCallingMode]string{
+	FunctionCallingUnspecified: "FunctionCallingUnspecified",
+	FunctionCallingAuto:        "FunctionCallingAuto",
+	FunctionCallingAny:         "FunctionCallingAny",
+	FunctionCallingNone:        "FunctionCallingNone",
+}
+
+func (v FunctionCallingMode) String() string {
+	if n, ok := namesForFunctionCallingMode[v]; ok {
+		return n
+	}
+	return fmt.Sprintf("FunctionCallingMode(%d)", v)
+}
+
 // FunctionDeclaration is structured representation of a function declaration as defined by the
 // [OpenAPI 3.03 specification](https://spec.openapis.org/oas/v3.0.3). Included
 // in this declaration are the function name and parameters. This
@@ -517,7 +585,7 @@ func (GenerateContentResponse) fromProto(p *pb.GenerateContentResponse) *Generat
 type GenerationConfig struct {
 	// Optional. Number of generated responses to return.
 	//
-	// This value must be between [1, 8], inclusive. If unset, this will default
+	// Currently, this value can only be set to 1. If unset, this will default
 	// to 1.
 	CandidateCount *int32
 	// Optional. The set of character sequences (up to 5) that will stop output
@@ -527,17 +595,15 @@ type GenerationConfig struct {
 	StopSequences []string
 	// Optional. The maximum number of tokens to include in a candidate.
 	//
-	// If unset, this will default to output_token_limit specified in the `Model`
-	// specification.
+	// Note: The default value varies by model, see the `Model.output_token_limit`
+	// attribute of the `Model` returned from the `getModel` function.
 	MaxOutputTokens *int32
 	// Optional. Controls the randomness of the output.
+	//
 	// Note: The default value varies by model, see the `Model.temperature`
-	// attribute of the `Model` returned the `getModel` function.
+	// attribute of the `Model` returned from the `getModel` function.
 	//
-	// Values can range from [0.0,1.0],
-	// inclusive. A value closer to 1.0 will produce responses that are more
-	// varied and creative, while a value closer to 0.0 will typically result in
-	// more straightforward responses from the model.
+	// Values can range from [0.0, infinity).
 	Temperature *float32
 	// Optional. The maximum cumulative probability of tokens to consider when
 	// sampling.
@@ -550,17 +616,16 @@ type GenerationConfig struct {
 	// of tokens based on the cumulative probability.
 	//
 	// Note: The default value varies by model, see the `Model.top_p`
-	// attribute of the `Model` returned the `getModel` function.
+	// attribute of the `Model` returned from the `getModel` function.
 	TopP *float32
 	// Optional. The maximum number of tokens to consider when sampling.
 	//
 	// The model uses combined Top-k and nucleus sampling.
 	//
 	// Top-k sampling considers the set of `top_k` most probable tokens.
-	// Defaults to 40.
 	//
 	// Note: The default value varies by model, see the `Model.top_k`
-	// attribute of the `Model` returned the `getModel` function.
+	// attribute of the `Model` returned from the `getModel` function.
 	TopK *int32
 }
 
@@ -634,9 +699,9 @@ const (
 	HarmCategoryUnspecified HarmCategory = 0
 	// HarmCategoryDerogatory means negative or harmful comments targeting identity and/or protected attribute.
 	HarmCategoryDerogatory HarmCategory = 1
-	// HarmCategoryToxicity means content that is rude, disrepspectful, or profane.
+	// HarmCategoryToxicity means content that is rude, disrespectful, or profane.
 	HarmCategoryToxicity HarmCategory = 2
-	// HarmCategoryViolence means describes scenarios depictng violence against an individual or group, or
+	// HarmCategoryViolence means describes scenarios depicting violence against an individual or group, or
 	// general descriptions of gore.
 	HarmCategoryViolence HarmCategory = 3
 	// HarmCategorySexual means contains references to sexual acts or other lewd content.
@@ -1044,6 +1109,31 @@ func (Tool) fromProto(p *pb.Tool) *Tool {
 	}
 }
 
+// ToolConfig is the Tool configuration containing parameters for specifying `Tool` use
+// in the request.
+type ToolConfig struct {
+	// Optional. Function calling config.
+	FunctionCallingConfig *FunctionCallingConfig
+}
+
+func (v *ToolConfig) toProto() *pb.ToolConfig {
+	if v == nil {
+		return nil
+	}
+	return &pb.ToolConfig{
+		FunctionCallingConfig: v.FunctionCallingConfig.toProto(),
+	}
+}
+
+func (ToolConfig) fromProto(p *pb.ToolConfig) *ToolConfig {
+	if p == nil {
+		return nil
+	}
+	return &ToolConfig{
+		FunctionCallingConfig: (FunctionCallingConfig{}).fromProto(p.FunctionCallingConfig),
+	}
+}
+
 // Type contains the list of OpenAPI data types as defined by
 // https://spec.openapis.org/oas/v3.0.3#data-types
 type Type int32

go.mod

@@ -3,31 +3,40 @@ module github.com/google/generative-ai-go
 go 1.21
 
 require (
-	cloud.google.com/go/ai v0.3.0
-	google.golang.org/api v0.149.0
+	cloud.google.com/go/ai v0.3.5-0.20240409161017-ce55ad694f21
+	google.golang.org/api v0.172.0
 	google.golang.org/protobuf v1.33.0
 )
 
 require (
-	cloud.google.com/go v0.110.8 // indirect
-	cloud.google.com/go/compute v1.23.1 // indirect
+	cloud.google.com/go v0.112.1 // indirect
+	cloud.google.com/go/compute v1.24.0 // indirect
 	cloud.google.com/go/compute/metadata v0.2.3 // indirect
-	cloud.google.com/go/longrunning v0.5.2 // indirect
+	cloud.google.com/go/longrunning v0.5.6 // indirect
+	github.com/felixge/httpsnoop v1.0.4 // indirect
+	github.com/go-logr/logr v1.4.1 // indirect
+	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
-	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/s2a-go v0.1.7 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
-	github.com/googleapis/gax-go/v2 v2.12.0 // indirect
+	github.com/googleapis/gax-go/v2 v2.12.3 // indirect
 	go.opencensus.io v0.24.0 // indirect
-	golang.org/x/crypto v0.17.0 // indirect
-	golang.org/x/net v0.17.0 // indirect
-	golang.org/x/oauth2 v0.13.0 // indirect
-	golang.org/x/sync v0.4.0 // indirect
-	golang.org/x/sys v0.15.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.49.0 // indirect
+	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.49.0 // indirect
+	go.opentelemetry.io/otel v1.24.0 // indirect
+	go.opentelemetry.io/otel/metric v1.24.0 // indirect
+	go.opentelemetry.io/otel/trace v1.24.0 // indirect
+	golang.org/x/crypto v0.21.0 // indirect
+	golang.org/x/net v0.22.0 // indirect
+	golang.org/x/oauth2 v0.18.0 // indirect
+	golang.org/x/sync v0.6.0 // indirect
+	golang.org/x/sys v0.18.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
-	google.golang.org/appengine v1.6.7 // indirect
-	google.golang.org/genproto v0.0.0-20231016165738-49dd2c1f3d0b // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20231016165738-49dd2c1f3d0b // indirect
-	google.golang.org/genproto/googleapis/rpc v0.0.0-20231016165738-49dd2c1f3d0b // indirect
-	google.golang.org/grpc v1.59.0 // indirect
+	golang.org/x/time v0.5.0 // indirect
+	google.golang.org/appengine v1.6.8 // indirect
+	google.golang.org/genproto v0.0.0-20240213162025-012b6fc9bca9 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20240401170217-c3f982113cda // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20240325203815-454cdb8f5daa // indirect
+	google.golang.org/grpc v1.62.1 // indirect
 )