genai: fix tool-using example and test (#163)

The model/API is sensitive to our weather example and emit 400s occasionally.

Reuse the movies example from the official docs (https://ai.google.dev/gemini-api/docs/function-calling) instead

Fixes #160

genai/client_test.go

@@ -38,6 +38,7 @@ import (
 const (
  defaultModel = "gemini-1.0-pro"
  modelForVision = "gemini-1.5-flash"
+ modelForTools = "gemini-1.5-pro-latest"
  imageFile = "personWorkingOnComputer.jpg"
 )
 
@@ -332,24 +333,24 @@ func TestLive(t *testing.T) {
  })
 
  t.Run("tools", func(t *testing.T) {
- weatherChat := func(t *testing.T, s *Schema, fcm FunctionCallingMode) {
- weatherTool := &Tool{
+ movieChat := func(t *testing.T, s *Schema, fcm FunctionCallingMode) {
+ movieTool := &Tool{
  FunctionDeclarations: []*FunctionDeclaration{{
- Name: "CurrentWeather",
- Description: "Get the current weather in a given location",
+ Name: "find_theaters",
+ Description: "find theaters based on location and optionally movie title which is currently playing in theaters",
  Parameters: s,
  }},
  }
- model := client.GenerativeModel(defaultModel)
+ model := client.GenerativeModel(modelForTools)
  model.SetTemperature(0)
- model.Tools = []*Tool{weatherTool}
+ model.Tools = []*Tool{movieTool}
  model.ToolConfig = &ToolConfig{
  FunctionCallingConfig: &FunctionCallingConfig{
  Mode: fcm,
  },
  }
  session := model.StartChat()
- res, err := session.SendMessage(ctx, Text("What is the weather like in New York?"))
+ res, err := session.SendMessage(ctx, Text("Which theaters in Mountain View show Barbie movie?"))
  if err != nil {
  t.Fatal(err)
  }
@@ -364,46 +365,46 @@ func TestLive(t *testing.T) {
  t.Fatalf("got %d FunctionCalls, want 1", len(funcalls))
  }
  funcall := funcalls[0]
- if g, w := funcall.Name, weatherTool.FunctionDeclarations[0].Name; g != w {
+ if g, w := funcall.Name, movieTool.FunctionDeclarations[0].Name; g != w {
  t.Errorf("FunctionCall.Name: got %q, want %q", g, w)
  }
  locArg, ok := funcall.Args["location"].(string)
  if !ok {
  t.Fatal(`funcall.Args["location"] is not a string`)
  }
- if c := "New York"; !strings.Contains(locArg, c) {
+ if c := "Mountain View"; !strings.Contains(locArg, c) {
  t.Errorf(`FunctionCall.Args["location"]: got %q, want string containing %q`, locArg, c)
  }
  res, err = session.SendMessage(ctx, FunctionResponse{
- Name: weatherTool.FunctionDeclarations[0].Name,
+ Name: movieTool.FunctionDeclarations[0].Name,
  Response: map[string]any{
- "weather_there": "cold",
+ "theater": "AMC16",
  },
  })
  if err != nil {
  t.Fatal(err)
  }
- checkMatch(t, responseString(res), "(it's|it is|weather) .*cold")
+ checkMatch(t, responseString(res), "AMC")
  }
  schema := &Schema{
  Type: TypeObject,
  Properties: map[string]*Schema{
  "location": {
  Type: TypeString,
- Description: "The city and state, e.g. San Francisco, CA",
+ Description: "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
  },
- "unit": {
- Type: TypeString,
- Enum: []string{"celsius", "fahrenheit"},
+ "title": {
+ Type:  TypeString,
+ Description: "Any movie title",
  },
  },
  Required: []string{"location"},
  }
  t.Run("direct", func(t *testing.T) {
- weatherChat(t, schema, FunctionCallingAuto)
+ movieChat(t, schema, FunctionCallingAuto)
  })
  t.Run("none", func(t *testing.T) {
- weatherChat(t, schema, FunctionCallingNone)
+ movieChat(t, schema, FunctionCallingNone)
  })
  })
 

genai/example_test.go

@@ -364,50 +364,36 @@ func ExampleTool() {
  }
  defer client.Close()
 
- currentWeather := func(city string) string {
- switch city {
- case "New York, NY":
- return "cold"
- case "Miami, FL":
- return "warm"
- default:
- return "unknown"
- }
- }
-
  // To use functions / tools, we have to first define a schema that describes
  // the function to the model. The schema is similar to OpenAPI 3.0.
- //
- // In this example, we create a single function that provides the model with
- // a weather forecast in a given location.
  schema := &genai.Schema{
  Type: genai.TypeObject,
  Properties: map[string]*genai.Schema{
  "location": {
  Type: genai.TypeString,
- Description: "The city and state, for example 'San Francisco, CA'. Both city and state are mandatory",
+ Description: "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
  },
- "unit": {
- Type: genai.TypeString,
- Enum: []string{"celsius", "fahrenheit"},
+ "title": {
+ Type:  genai.TypeString,
+ Description: "Any movie title",
  },
  },
  Required: []string{"location"},
  }
 
- weatherTool := &genai.Tool{
+ movieTool := &genai.Tool{
  FunctionDeclarations: []*genai.FunctionDeclaration{{
- Name: "CurrentWeather",
- Description: "Get the current weather in a given location",
+ Name: "find_theaters",
+ Description: "find theaters based on location and optionally movie title which is currently playing in theaters",
  Parameters: schema,
  }},
  }
 
- model := client.GenerativeModel("gemini-1.5-flash-latest")
+ model := client.GenerativeModel("gemini-1.5-pro-latest")
 
  // Before initiating a conversation, we tell the model which tools it has
  // at its disposal.
- model.Tools = []*genai.Tool{weatherTool}
+ model.Tools = []*genai.Tool{movieTool}
 
  // For using tools, the chat mode is useful because it provides the required
  // chat context. A model needs to have tools supplied to it in the chat
@@ -417,45 +403,35 @@ func ExampleTool() {
  //
  // 1. We send a question to the model
  // 2. The model recognizes that it needs to use a tool to answer the question,
- // an returns a FunctionCall response asking to use the CurrentWeather
- // tool.
+ // an returns a FunctionCall response asking to use the tool.
  // 3. We send a FunctionResponse message, simulating the return value of
- // CurrentWeather for the model's query.
+ // the tool for the model's query.
  // 4. The model provides its text answer in response to this message.
  session := model.StartChat()
 
- res, err := session.SendMessage(ctx, genai.Text("What is the weather like in New York?"))
+ res, err := session.SendMessage(ctx, genai.Text("Which theaters in Mountain View show Barbie movie?"))
  if err != nil {
- log.Fatal(err)
+ log.Fatalf("session.SendMessage: %v", err)
  }
 
  part := res.Candidates[0].Content.Parts[0]
  funcall, ok := part.(genai.FunctionCall)
- if !ok {
- log.Fatalf("expected FunctionCall: %v", part)
- }
-
- if funcall.Name != "CurrentWeather" {
- log.Fatalf("expected CurrentWeather: %v", funcall.Name)
+ if !ok || funcall.Name != "find_theaters" {
+ log.Fatalf("expected FunctionCall to find_theaters: %v", part)
  }
 
  // Expect the model to pass a proper string "location" argument to the tool.
- locArg, ok := funcall.Args["location"].(string)
- if !ok {
+ if _, ok := funcall.Args["location"].(string); !ok {
  log.Fatalf("expected string: %v", funcall.Args["location"])
  }
 
- weatherData := currentWeather(locArg)
+ // Provide the model with a hard-coded reply.
  res, err = session.SendMessage(ctx, genai.FunctionResponse{
- Name: weatherTool.FunctionDeclarations[0].Name,
+ Name: movieTool.FunctionDeclarations[0].Name,
  Response: map[string]any{
- "weather": weatherData,
+ "theater": "AMC16",
  },
  })
- if err != nil {
- log.Fatal(err)
- }
-
  printResponse(res)
 }