prompt: Add JobList (#78)

go.mod

@@ -5,7 +5,7 @@ go 1.22
 require (
 	github.com/google/generative-ai-go v0.16.0
 	github.com/sashabaranov/go-openai v1.26.2
-	github.com/sunshineplan/utils v0.1.71
+	github.com/sunshineplan/workers v1.0.1
 	golang.org/x/time v0.5.0
 	google.golang.org/api v0.187.0
 )
@@ -26,6 +26,7 @@ require (
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.3.2 // indirect
 	github.com/googleapis/gax-go/v2 v2.12.5 // indirect
+	github.com/sunshineplan/utils v0.1.71 // indirect
 	go.opencensus.io v0.24.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.51.0 // indirect
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.51.0 // indirect

go.sum

@@ -79,6 +79,8 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 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=
 go.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0=
 go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo=
 go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.51.0 h1:A3SayB3rNyt+1S6qpI9mHPkeHTZbD7XILEqWnYZb2l0=

prompt/prompt.go

@@ -8,7 +8,7 @@ import (
 	"time"
 
 	"github.com/sunshineplan/ai"
-	"github.com/sunshineplan/utils/workers"
+	"github.com/sunshineplan/workers"
 )
 
 const (
@@ -50,7 +50,7 @@ type Prompt struct {
 	n      int
 
 	d       time.Duration
-	workers int
+	workers int64
 }
 
 func New(prompt string) *Prompt {
@@ -79,7 +79,7 @@ func (prompt *Prompt) SetAITimeout(d time.Duration) *Prompt {
 	return prompt
 }
 
-func (prompt *Prompt) SetWorkers(n int) *Prompt {
+func (prompt *Prompt) SetWorkers(n int64) *Prompt {
 	prompt.workers = n
 	return prompt
 }
@@ -122,31 +122,57 @@ type Result struct {
 	Error  error
 }
 
-func (prompt *Prompt) Execute(ai ai.AI, input []string, prefix string) (<-chan Result, int, error) {
+func (prompt *Prompt) Execute(ai ai.AI, input []string, prefix string) (<-chan *Result, int, error) {
 	prompts, err := prompt.Prompts(input, prefix)
 	if err != nil {
 		return nil, 0, err
 	}
 	n := len(prompts)
-	c := make(chan Result, n)
+	c := make(chan *Result, n)
 	go func() {
-		workers.RunSlice(prompt.workers, prompts, func(i int, p string) {
-			var ctx context.Context
-			var cancel context.CancelFunc
-			if prompt.d != 0 {
-				ctx, cancel = context.WithTimeout(context.Background(), prompt.d)
-			} else {
-				ctx, cancel = context.WithCancel(context.Background())
-			}
-			defer cancel()
-			resp, err := ai.Chat(ctx, p)
+		workers.NewWorkers(prompt.workers).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, err}
+				c <- &Result{i, p, nil, err}
 			} else {
-				c <- Result{i, p, resp.Results(), nil}
+				c <- &Result{i, p, resp.Results(), nil}
 			}
-		})
+		}))
 		close(c)
 	}()
 	return c, n, nil
 }
+
+func (prompt *Prompt) JobList(ai ai.AI, input []string, prefix string, c chan<- *Result) (*workers.JobList[*Result], error) {
+	prompts, err := prompt.Prompts(input, prefix)
+	if err != nil {
+		return nil, err
+	}
+	jobList := workers.NewJobList(workers.NewWorkers(prompt.workers), func(r *Result) {
+		resp, err := chat(ai, prompt.d, r.Prompt)
+		if err != nil {
+			r.Result = nil
+			r.Error = err
+		} else {
+			r.Result = resp.Results()
+			r.Error = nil
+		}
+		c <- r
+	})
+	for i, p := range prompts {
+		jobList.PushBack(&Result{Index: i, Prompt: p})
+	}
+	return jobList, nil
+}
+
+func chat(ai ai.AI, d time.Duration, p string) (ai.ChatResponse, error) {
+	var ctx context.Context
+	var cancel context.CancelFunc
+	if d > 0 {
+		ctx, cancel = context.WithTimeout(context.Background(), d)
+	} else {
+		ctx, cancel = context.WithCancel(context.Background())
+	}
+	defer cancel()
+	return ai.Chat(ctx, p)
+}