1// Copyright 2025 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package genai
16
17import (
18 "context"
19 "errors"
20 "iter"
21)
22
23// ErrPageDone is the error returned by an iterator's Next method when no more pages are available.
24var ErrPageDone = errors.New("no more pages")
25
26// Page represents a page of results from a paginated API call.
27// It contains a slice of items and information about the next page.
28type Page[T any] struct {
29 Name string // The name of the resource.
30 Items []*T // The items in the current page.
31 NextPageToken string // The token to use to retrieve the next page of results.
32
33 config map[string]any // The configuration used for the API call.
34 listFunc func(ctx context.Context, config map[string]any) ([]*T, string, error) // The function used to retrieve the next page.
35}
36
37func newPage[T any](ctx context.Context, name string, config map[string]any, listFunc func(ctx context.Context, config map[string]any) ([]*T, string, error)) (Page[T], error) {
38 p := Page[T]{
39 Name: name,
40 config: config,
41 listFunc: listFunc,
42 }
43 items, nextPageToken, err := listFunc(ctx, config)
44 if err != nil {
45 return p, err
46 }
47 p.Items = items
48 p.NextPageToken = nextPageToken
49 return p, nil
50}
51
52// all returns an iterator that yields all items across all pages of results.
53//
54// The iterator retrieves each page sequentially and yields each item within
55// the page. If an error occurs during retrieval, the iterator will stop
56// and the error will be returned as the second value in the next call to Next().
57// A genai.PageDone error indicates that all pages have been processed.
58func (p Page[T]) all(ctx context.Context) iter.Seq2[*T, error] {
59 return func(yield func(*T, error) bool) {
60 for {
61 for _, item := range p.Items {
62 if !yield(item, nil) {
63 return
64 }
65 }
66 var err error
67 p, err = p.Next(ctx)
68 if err == ErrPageDone {
69 return
70 }
71 if err != nil {
72 yield(nil, err)
73 return
74 }
75 }
76 }
77}
78
79// Next retrieves the next page of results.
80//
81// If there are no more pages, PageDone is returned. Otherwise,
82// a new Page struct containing the next set of results is returned.
83// Any other errors encountered during retrieval will also be returned.
84func (p Page[T]) Next(ctx context.Context) (Page[T], error) {
85 if p.NextPageToken == "" {
86 return p, ErrPageDone
87 }
88 c := make(map[string]any)
89 for k, v := range p.config {
90 c[k] = v
91 }
92 c["PageToken"] = p.NextPageToken
93
94 return newPage[T](ctx, p.Name, c, p.listFunc)
95}