prompt: Result add Tokens field (#82)

sunshineplan · Jul 10, 2024 · e9d14b7 · e9d14b7
1 parent 9aba1bf
commit e9d14b7
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/prompt/prompt.go b/prompt/prompt.go
@@ -119,6 +119,7 @@ type Result struct {
 	Index  int
 	Prompt string
 	Result []string
+	Tokens int
 	Error  error
 }
 
@@ -133,9 +134,9 @@ func (prompt *Prompt) Execute(ai ai.AI, input []string, prefix string) (<-chan *
 		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, 0, err}
 			} else {
-				c <- &Result{i, p, resp.Results(), nil}
+				c <- &Result{i, p, resp.Results(), resp.TokenCount().Total, nil}
 			}
 		}))
 		close(c)