genai: add sample for tokens_tools (#190)

We're requested to add a tool token count example even w/o automatic function calling. Just the function names are sufficient to demonstrate a change in the token count - we don't have to specify schemas
google · Jul 19, 2024 · ab21d2a · ab21d2a
1 parent 70f6989
commit ab21d2a
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 0 deletions.
diff --git a/genai/example_test.go b/genai/example_test.go
@@ -429,6 +429,44 @@ func ExampleGenerativeModel_CountTokens_textOnly() {
 
 }
 
+func ExampleGenerativeModel_CountTokens_tools() {
+	ctx := context.Background()
+	client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer client.Close()
+
+	model := client.GenerativeModel("gemini-1.5-flash-001")
+	prompt := "I have 57 cats, each owns 44 mittens, how many mittens is that in total?"
+
+	tokResp, err := model.CountTokens(ctx, genai.Text(prompt))
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Println("total_tokens:", tokResp.TotalTokens)
+	// ( total_tokens: 23 )
+
+	tools := []*genai.Tool{
+		&genai.Tool{FunctionDeclarations: []*genai.FunctionDeclaration{
+			{Name: "add"},
+			{Name: "subtract"},
+			{Name: "multiply"},
+			{Name: "divide"},
+		}}}
+
+	model.Tools = tools
+	tokResp, err = model.CountTokens(ctx, genai.Text(prompt))
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Println("total_tokens:", tokResp.TotalTokens)
+	// ( total_tokens: 99 )
+
+}
+
 func ExampleGenerativeModel_CountTokens_cachedContent() {
 	ctx := context.Background()
 	client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))

diff --git a/genai/internal/samples/docs-snippets_test.go b/genai/internal/samples/docs-snippets_test.go
@@ -442,6 +442,46 @@ func ExampleGenerativeModel_CountTokens_textOnly() {
 	// [END tokens_text_only]
 }
 
+func ExampleGenerativeModel_CountTokens_tools() {
+	ctx := context.Background()
+	client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer client.Close()
+
+	// [START tokens_tools]
+	model := client.GenerativeModel("gemini-1.5-flash-001")
+	prompt := "I have 57 cats, each owns 44 mittens, how many mittens is that in total?"
+
+	tokResp, err := model.CountTokens(ctx, genai.Text(prompt))
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Println("total_tokens:", tokResp.TotalTokens)
+	// ( total_tokens: 23 )
+
+	tools := []*genai.Tool{
+		&genai.Tool{FunctionDeclarations: []*genai.FunctionDeclaration{
+			{Name: "add"},
+			{Name: "subtract"},
+			{Name: "multiply"},
+			{Name: "divide"},
+		}}}
+
+	model.Tools = tools
+	tokResp, err = model.CountTokens(ctx, genai.Text(prompt))
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Println("total_tokens:", tokResp.TotalTokens)
+	// ( total_tokens: 99 )
+
+	// [END tokens_tools]
+}
+
 func ExampleGenerativeModel_CountTokens_cachedContent() {
 	ctx := context.Background()
 	client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))