genai: move source of truth of example_test and add code generator (#164

)

We'd like to automatically extract code snippets from our examples for
the official documentation; the documentation tooling extracts code
snippets delineated with special [START... and [END... markers.

However, we also want our examples to appear on the pkgsite without
these special comments.

Therefore, the source of truth for example_test is moved to an internal
directory from which snippets will be extracted; a tool processes this
file to remove these special comments before placing it in
genai/example_test.go

One snippet is added as an example - context window reporting for tokens

genai/example_test.go

@@ -1,3 +1,6 @@
+// DO NOT EDIT THIS FILE -- it is automatically generated.
+// See internal/cmd/gen-examples for details.
+
 // Copyright 2023 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -146,6 +149,24 @@ func ExampleGenerativeModel_GenerateContentStream() {
 	}
 }
 
+func ExampleGenerativeModel_CountTokens_contextWindow() {
+	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.0-pro-001")
+	info, err := model.Info(ctx)
+	if err != nil {
+		log.Fatal(err)
+	}
+	fmt.Printf("input_token_limit=%v\n", info.InputTokenLimit)
+	fmt.Printf("output_token_limit=%v\n", info.OutputTokenLimit)
+
+}
+
 func ExampleGenerativeModel_CountTokens() {
 	ctx := context.Background()
 	client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))

genai/internal/cmd/gen-examples/gen-examples.go

@@ -0,0 +1,79 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This code generator takes examples from the internal/snippets directory
+// and copies them to "official" examples in genai/example_test.go, while
+// removing snippet comments (between [START...] and [END...]) that are used
+// for website documentation purposes.
+//
+// Run from the root directory:
+//
+//   go run ./genai/internal/cmd/gen-examples/ < genai/internal/snippets/example_test.go > genai/example_test.go
+
+package main
+
+import (
+	"fmt"
+	"go/ast"
+	"go/format"
+	"go/parser"
+	"go/token"
+	"log"
+	"os"
+	"strings"
+)
+
+func main() {
+	fset := token.NewFileSet()
+	file, err := parser.ParseFile(fset, "src.go", os.Stdin, parser.ParseComments)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	for _, cgroup := range file.Comments {
+		sanitizeCommentGroup(cgroup)
+	}
+
+	fmt.Println(strings.TrimLeft(preamble, "\r\n"))
+	format.Node(os.Stdout, fset, file)
+}
+
+const preamble = `
+// DO NOT EDIT THIS FILE -- it is automatically generated.
+// See internal/cmd/gen-examples for details.
+`
+
+func printCommentGroup(cg *ast.CommentGroup) {
+	fmt.Printf("-- comment group %p\n", cg)
+	for _, c := range cg.List {
+		fmt.Println(c.Slash, c.Text)
+	}
+}
+
+// sanitizeCommentGroup removes comment blocks between [START... and [END...
+// (including these lines) - it modifies cg.
+func sanitizeCommentGroup(cg *ast.CommentGroup) {
+	var nl []*ast.Comment
+	exclude := false
+	for _, commentLine := range cg.List {
+		if strings.Contains(commentLine.Text, "[START") {
+			exclude = true
+		} else if strings.Contains(commentLine.Text, "[END") {
+			exclude = false
+		} else if !exclude {
+			nl = append(nl, commentLine)
+		}
+	}
+	cg.List = nl
+}