Upgrade workers package (#133)

sunshineplan · Oct 22, 2024 · 8519b89 · 8519b89
1 parent c0d4cd4
commit 8519b89
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 11 deletions.
diff --git a/go.mod b/go.mod
@@ -5,7 +5,7 @@ go 1.23
 require (
 	github.com/google/generative-ai-go v0.18.0
 	github.com/openai/openai-go v0.1.0-alpha.29
-	github.com/sunshineplan/workers v1.0.1
+	github.com/sunshineplan/workers v1.0.3
 	golang.org/x/time v0.7.0
 	google.golang.org/api v0.201.0
 )
@@ -25,7 +25,7 @@ require (
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
 	github.com/googleapis/gax-go/v2 v2.13.0 // indirect
-	github.com/sunshineplan/utils v0.1.71 // indirect
+	github.com/sunshineplan/utils v0.1.73 // indirect
 	github.com/tidwall/gjson v1.14.4 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/tidwall/pretty v1.2.1 // indirect

diff --git a/go.sum b/go.sum
@@ -77,10 +77,10 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/sunshineplan/utils v0.1.71 h1:FTe6qAdevePih+hLil1KvKhMAU2Bp+qHTtlO0wxoAYI=
-github.com/sunshineplan/utils v0.1.71/go.mod h1:QMe2vCEFq9VliPac8WfXndi5jP0ZrPHus8wvJlgBy5M=
-github.com/sunshineplan/workers v1.0.1 h1:vhDo4ZaerVu0QM1udRmWpvHY2U61/vTLRVwmKcYOxNw=
-github.com/sunshineplan/workers v1.0.1/go.mod h1:EjLgk5vd/2ig6vdWDAZy/fpkzdSm8IVDsARPQVSyQy0=
+github.com/sunshineplan/utils v0.1.73 h1:Kvi45QYzdvNqQKON5Q67M7U/35+RrOnX7K8sbZtHsgI=
+github.com/sunshineplan/utils v0.1.73/go.mod h1:61jHA9jn2CVlZzt9LbPz+MpOEiR6Jy5SNHHEcc/KQNk=
+github.com/sunshineplan/workers v1.0.3 h1:9NSB3nm9yxZFT/ZU4DaOXdWJ56Y2WXh238apdZUdA44=
+github.com/sunshineplan/workers v1.0.3/go.mod h1:wUlVWxvfwTiuULOw4atF17bgeL30gIMI6gE050reEVw=
 github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
 github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM=
 github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=

diff --git a/prompt/prompt.go b/prompt/prompt.go
@@ -115,11 +115,12 @@ type Result struct {
 	Error  error
 }
 
-func newWorkers(ai ai.AI) *workers.Workers {
+func limit(ai ai.AI) int {
+	var limit int
 	if rpm := ai.Limit(); rpm != math.MaxInt64 {
-		return workers.NewWorkers(rpm)
+		return int(rpm)
 	}
-	return workers.NewWorkers(0)
+	return limit
 }
 
 func (prompt *Prompt) Execute(ai ai.AI, input []string, prefix string) (<-chan *Result, int, error) {
@@ -130,7 +131,7 @@ func (prompt *Prompt) Execute(ai ai.AI, input []string, prefix string) (<-chan *
 	n := len(prompts)
 	c := make(chan *Result, n)
 	go func() {
-		newWorkers(ai).Run(context.Background(), workers.SliceJob(prompts, func(i int, p string) {
+		workers.Workers(limit(ai)).Run(context.Background(), workers.SliceJob(prompts, func(i int, p string) {
 			resp, err := chat(ai, prompt.d, p)
 			if err != nil {
 				c <- &Result{i, p, nil, 0, err}
@@ -149,7 +150,7 @@ func (prompt *Prompt) JobList(ctx context.Context, ai ai.AI, input []string, pre
 	if err != nil {
 		return nil, 0, err
 	}
-	jobList := workers.NewJobList(newWorkers(ai), func(r *Result) {
+	jobList := workers.NewJobList(limit(ai), func(r *Result) {
 		resp, err := chat(ai, prompt.d, r.Prompt)
 		if err != nil {
 			r.Result = nil