diff --git a/packages/pagination/pagination.go b/packages/pagination/pagination.go index a6b5c164f..e71ce8ca4 100644 --- a/packages/pagination/pagination.go +++ b/packages/pagination/pagination.go @@ -328,10 +328,6 @@ func (r *NextCursorPage[T]) UnmarshalJSON(data []byte) error { // there is no next page, this function will return a 'nil' for the page value, but // will not return an error func (r *NextCursorPage[T]) GetNextPage() (res *NextCursorPage[T], err error) { - if len(r.Data) == 0 { - return nil, nil - } - if r.JSON.HasMore.Valid() && r.HasMore == false { return nil, nil } @@ -364,11 +360,12 @@ func (r *NextCursorPage[T]) SetPageConfig(cfg *requestconfig.RequestConfig, res } type NextCursorPageAutoPager[T any] struct { - page *NextCursorPage[T] - cur T - idx int - run int - err error + page *NextCursorPage[T] + cur T + idx int + run int + err error + seenCursors *map[string]struct{} paramObj } @@ -380,20 +377,33 @@ func NewNextCursorPageAutoPager[T any](page *NextCursorPage[T], err error) *Next } func (r *NextCursorPageAutoPager[T]) Next() bool { - if r.page == nil || len(r.page.Data) == 0 { - return false - } - if r.idx >= len(r.page.Data) { + for { + if r.page == nil { + return false + } + if r.idx < len(r.page.Data) { + r.cur = r.page.Data[r.idx] + r.run += 1 + r.idx += 1 + return true + } r.idx = 0 + next := r.page.Next + if next != "" { + if r.seenCursors == nil { + seenCursors := make(map[string]struct{}) + r.seenCursors = &seenCursors + } + if _, seen := (*r.seenCursors)[next]; seen { + return false + } + (*r.seenCursors)[next] = struct{}{} + } r.page, r.err = r.page.GetNextPage() - if r.err != nil || r.page == nil || len(r.page.Data) == 0 { + if r.err != nil || r.page == nil { return false } } - r.cur = r.page.Data[r.idx] - r.run += 1 - r.idx += 1 - return true } func (r *NextCursorPageAutoPager[T]) Current() T { diff --git a/pagination_next_cursor_empty_test.go b/pagination_next_cursor_empty_test.go new file mode 100644 index 000000000..a6cc5e13f --- /dev/null +++ b/pagination_next_cursor_empty_test.go @@ -0,0 +1,97 @@ +package openai_test + +import ( + "context" + "errors" + "io" + "net/http" + "reflect" + "strings" + "testing" + + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/option" + "github.com/openai/openai-go/v3/packages/pagination" +) + +func TestNextCursorPaginationFollowsEmptyPages(t *testing.T) { + tests := []struct { + name string + pages map[string]string + want []string + }{ + { + name: "empty first page", + pages: map[string]string{ + "": `{"data":[],"has_more":true,"next":"cursor-1"}`, + "cursor-1": `{"data":[{"id":"group-1","created_at":1,"group_type":"group","is_scim_managed":false,"name":"one"}],"has_more":false,"next":null}`, + }, + want: []string{"group-1"}, + }, + { + name: "repeated cursor stops", + pages: map[string]string{ + "": `{"data":[],"has_more":true,"next":"cursor-1"}`, + "cursor-1": `{"data":[],"has_more":true,"next":"cursor-1"}`, + }, + want: []string{}, + }, + { + name: "empty intermediate page", + pages: map[string]string{ + "": `{"data":[{"id":"group-1","created_at":1,"group_type":"group","is_scim_managed":false,"name":"one"}],"has_more":true,"next":"cursor-1"}`, + "cursor-1": `{"data":[],"has_more":true,"next":"cursor-2"}`, + "cursor-2": `{"data":[{"id":"group-2","created_at":2,"group_type":"group","is_scim_managed":false,"name":"two"}],"has_more":false,"next":null}`, + }, + want: []string{"group-1", "group-2"}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + calls := 0 + client := openai.NewClient( + option.WithBaseURL("https://example.com/v1"), + option.WithAdminAPIKey("test-admin-key"), + option.WithMaxRetries(0), + option.WithHTTPClient(paginationHTTPDoerFunc(func(req *http.Request) (*http.Response, error) { + calls++ + cursor := req.URL.Query().Get("after") + body, ok := test.pages[cursor] + if !ok { + return nil, errors.New("unexpected pagination cursor") + } + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"application/json"}}, + Body: io.NopCloser(strings.NewReader(body)), + Request: req, + }, nil + })), + ) + + pager := client.Admin.Organization.Groups.ListAutoPaging( + context.Background(), openai.AdminOrganizationGroupListParams{}, + ) + var got []string + for pager.Next() { + got = append(got, pager.Current().ID) + } + if err := pager.Err(); err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(got, test.want) { + t.Fatalf("group IDs = %v, want %v", got, test.want) + } + if calls != len(test.pages) { + t.Fatalf("HTTP calls = %d, want %d", calls, len(test.pages)) + } + }) + } +} + +func TestNextCursorPageAutoPagerRemainsComparable(t *testing.T) { + var left pagination.NextCursorPageAutoPager[int] + var right pagination.NextCursorPageAutoPager[int] + _ = left == right +}