feat(generate-text): add onStepFinish() (#26)

* feat(generate-text): add onStepFinish() * chore(generate-text): update test
moeru-ai · Jan 15, 2025 · 2bc1733 · 2bc1733
1 parent c59ab2c
commit 2bc1733
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/packages/generate-text/src/index.ts b/packages/generate-text/src/index.ts
@@ -10,6 +10,7 @@ import {
 export interface GenerateTextOptions extends ChatOptions {
   /** @default 1 */
   maxSteps?: number
+  onStepFinish?: (step: StepResult) => Promise<void> | void
   /** @internal */
   steps?: StepResult[]
   /** if you want to enable stream, use `@xsai/stream-{text,object}` */
@@ -102,6 +103,9 @@ const rawGenerateText: RawGenerateText = async (options: GenerateTextOptions) =>
 
         steps.push(step)
 
+        if (options.onStepFinish)
+          await options.onStepFinish(step)
+
         return {
           finishReason,
           steps,
@@ -141,12 +145,17 @@ const rawGenerateText: RawGenerateText = async (options: GenerateTextOptions) =>
         })
       }
 
-      steps.push({
+      const step: StepResult = {
         text: message.content,
         toolCalls,
         toolResults,
         usage,
-      })
+      }
+
+      steps.push(step)
+
+      if (options.onStepFinish)
+        await options.onStepFinish(step)
 
       return async () => await rawGenerateText({
         ...options,

diff --git a/packages/generate-text/test/index.test.ts b/packages/generate-text/test/index.test.ts
@@ -1,10 +1,12 @@
 import { ollama } from '@xsai/providers'
 import { describe, expect, it } from 'vitest'
 
-import { generateText } from '../src'
+import { generateText, type StepResult } from '../src'
 
 describe('@xsai/generate-text', () => {
   it('basic', async () => {
+    let step: StepResult | undefined
+
     const { finishReason, steps, text, toolCalls, toolResults } = await generateText({
       ...ollama.chat('llama3.2'),
       messages: [
@@ -17,13 +19,18 @@ describe('@xsai/generate-text', () => {
           role: 'user',
         },
       ],
+      onStepFinish: (result) => {
+        step = result
+      },
     })
 
     expect(text).toStrictEqual('YES')
     expect(finishReason).toBe('stop')
     expect(toolCalls.length).toBe(0)
     expect(toolResults.length).toBe(0)
     expect(steps).toMatchSnapshot()
+
+    expect(steps[0]).toStrictEqual(step)
   })
 
   // TODO: error handling