feat (ai/core): add continuationSteps to generateText (#3104)

vercel · Sep 24, 2024 · e6c7e98 · e6c7e98
1 parent 3926476
commit e6c7e98
Show file tree

Hide file tree

Showing 7 changed files with 460 additions and 27 deletions.
diff --git a/.changeset/metal-stingrays-prove.md b/.changeset/metal-stingrays-prove.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/core): add continuationSteps to generateText
diff --git a/content/docs/03-ai-sdk-core/05-generating-text.mdx b/content/docs/03-ai-sdk-core/05-generating-text.mdx
@@ -48,6 +48,34 @@ The result object of `generateText` contains several promises that resolve when
 - `result.finishReason`: The reason the model finished generating text.
 - `result.usage`: The usage of the model during text generation.
 
+### Generating Long Text
+
+Most language models have an output limit that is much shorter than their context window.
+This means that you cannot generate long text in one go,
+but it is possible to add responses back to the input and continue generating
+to create longer text.
+
+`generateText` supports such continuations for long text generation using the experimental `continuationSteps` setting:
+
+```tsx
+import { openai } from '@ai-sdk/openai';
+import { generateText } from 'ai';
+
+const {
+  text, // combined text
+  usage, // combined usage of all steps
+} = await generateText({
+  model: openai('gpt-4o'), // 4096 output tokens
+  maxSteps: 5, // enable multi-step calls
+  experimental_continuationSteps: true,
+  prompt:
+    'Write a book about Roman history, ' +
+    'from the founding of the city of Rome ' +
+    'to the fall of the Western Roman Empire. ' +
+    'Each chapter MUST HAVE at least 1000 words.',
+});
+```
+
 ## `streamText`
 
 Depending on your model and prompt, it can take a large language model (LLM) up to a minute to finish generating it's response. This delay can be unacceptable for interactive use cases such as chatbots or real-time applications, where users expect immediate responses.

diff --git a/content/docs/07-reference/ai-sdk-core/01-generate-text.mdx b/content/docs/07-reference/ai-sdk-core/01-generate-text.mdx
@@ -353,6 +353,12 @@ To see `generateText` in action, check out [these examples](#examples).
       description:
         'Maximum number of sequential LLM calls (steps), e.g. when you use tool calls. A maximum number is required to prevent infinite loops in the case of misconfigured tools. By default, it is set to 1.',
     },
+    {
+      name: 'experimental_continuationSteps',
+      type: 'boolean',
+      isOptional: true,
+      description: 'Enable or disable continuation steps. Disabled by default.',
+    },
     {
       name: 'experimental_telemetry',
       type: 'TelemetrySettings',

diff --git a/examples/ai-core/src/generate-text/openai-multi-step-continuation.ts b/examples/ai-core/src/generate-text/openai-multi-step-continuation.ts
@@ -0,0 +1,23 @@
+import { openai } from '@ai-sdk/openai';
+import { generateText } from 'ai';
+import 'dotenv/config';
+
+async function main() {
+  const { text, usage, steps } = await generateText({
+    model: openai('gpt-4o'), // 4096 output tokens
+    maxSteps: 5,
+    experimental_continuationSteps: true,
+    prompt:
+      'Write a book about Roman history, ' +
+      'from the founding of the city of Rome ' +
+      'to the fall of the Western Roman Empire. ' +
+      'Each chapter MUST HAVE at least 1000 words.',
+  });
+
+  console.log(text);
+  console.log();
+  console.log('Usage:', usage);
+  console.log('# of steps:', steps.length);
+}
+
+main().catch(console.error);
diff --git a/packages/ai/core/generate-text/__snapshots__/generate-text.test.ts.snap b/packages/ai/core/generate-text/__snapshots__/generate-text.test.ts.snap
@@ -1,6 +1,6 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
-exports[`options.maxSteps > 2 steps > onStepFinish should be called for each step 1`] = `
+exports[`options.maxSteps > 2 steps: initial, tool-result > onStepFinish should be called for each step 1`] = `
 [
   {
     "experimental_providerMetadata": undefined,
@@ -65,7 +65,7 @@ exports[`options.maxSteps > 2 steps > onStepFinish should be called for each ste
 ]
 `;
 
-exports[`options.maxSteps > 2 steps > result.responseMessages should contain response messages from all steps 1`] = `
+exports[`options.maxSteps > 2 steps: initial, tool-result > result.responseMessages should contain response messages from all steps 1`] = `
 [
   {
     "content": [
@@ -107,7 +107,7 @@ exports[`options.maxSteps > 2 steps > result.responseMessages should contain res
 ]
 `;
 
-exports[`options.maxSteps > 2 steps > result.steps should contain all steps 1`] = `
+exports[`options.maxSteps > 2 steps: initial, tool-result > result.steps should contain all steps 1`] = `
 [
   {
     "experimental_providerMetadata": undefined,
@@ -172,6 +172,140 @@ exports[`options.maxSteps > 2 steps > result.steps should contain all steps 1`]
 ]
 `;
 
+exports[`options.maxSteps > 3 steps: initial, continuation, continuation > onStepFinish should be called for each step 1`] = `
+[
+  {
+    "experimental_providerMetadata": undefined,
+    "finishReason": "length",
+    "logprobs": undefined,
+    "response": {
+      "headers": undefined,
+      "id": "test-id-1-from-model",
+      "modelId": "test-response-model-id",
+      "timestamp": 1970-01-01T00:00:00.000Z,
+    },
+    "text": "part-1",
+    "toolCalls": [],
+    "toolResults": [],
+    "usage": {
+      "completionTokens": 20,
+      "promptTokens": 10,
+      "totalTokens": 30,
+    },
+    "warnings": undefined,
+  },
+  {
+    "experimental_providerMetadata": undefined,
+    "finishReason": "length",
+    "logprobs": undefined,
+    "response": {
+      "headers": {
+        "custom-response-header": "response-header-value",
+      },
+      "id": "test-id-2-from-model",
+      "modelId": "test-response-model-id",
+      "timestamp": 1970-01-01T00:00:10.000Z,
+    },
+    "text": "part-2",
+    "toolCalls": [],
+    "toolResults": [],
+    "usage": {
+      "completionTokens": 5,
+      "promptTokens": 30,
+      "totalTokens": 35,
+    },
+    "warnings": undefined,
+  },
+  {
+    "experimental_providerMetadata": undefined,
+    "finishReason": "stop",
+    "logprobs": undefined,
+    "response": {
+      "headers": undefined,
+      "id": "test-id-3-from-model",
+      "modelId": "test-response-model-id",
+      "timestamp": 1970-01-01T00:00:20.000Z,
+    },
+    "text": "part-3",
+    "toolCalls": [],
+    "toolResults": [],
+    "usage": {
+      "completionTokens": 2,
+      "promptTokens": 3,
+      "totalTokens": 5,
+    },
+    "warnings": undefined,
+  },
+]
+`;
+
+exports[`options.maxSteps > 3 steps: initial, continuation, continuation > result.steps should contain all steps 1`] = `
+[
+  {
+    "experimental_providerMetadata": undefined,
+    "finishReason": "length",
+    "logprobs": undefined,
+    "response": {
+      "headers": undefined,
+      "id": "test-id-1-from-model",
+      "modelId": "test-response-model-id",
+      "timestamp": 1970-01-01T00:00:00.000Z,
+    },
+    "text": "part-1",
+    "toolCalls": [],
+    "toolResults": [],
+    "usage": {
+      "completionTokens": 20,
+      "promptTokens": 10,
+      "totalTokens": 30,
+    },
+    "warnings": undefined,
+  },
+  {
+    "experimental_providerMetadata": undefined,
+    "finishReason": "length",
+    "logprobs": undefined,
+    "response": {
+      "headers": {
+        "custom-response-header": "response-header-value",
+      },
+      "id": "test-id-2-from-model",
+      "modelId": "test-response-model-id",
+      "timestamp": 1970-01-01T00:00:10.000Z,
+    },
+    "text": "part-2",
+    "toolCalls": [],
+    "toolResults": [],
+    "usage": {
+      "completionTokens": 5,
+      "promptTokens": 30,
+      "totalTokens": 35,
+    },
+    "warnings": undefined,
+  },
+  {
+    "experimental_providerMetadata": undefined,
+    "finishReason": "stop",
+    "logprobs": undefined,
+    "response": {
+      "headers": undefined,
+      "id": "test-id-3-from-model",
+      "modelId": "test-response-model-id",
+      "timestamp": 1970-01-01T00:00:20.000Z,
+    },
+    "text": "part-3",
+    "toolCalls": [],
+    "toolResults": [],
+    "usage": {
+      "completionTokens": 2,
+      "promptTokens": 3,
+      "totalTokens": 5,
+    },
+    "warnings": undefined,
+  },
+]
+`;
+
 exports[`result.responseMessages > should contain assistant response message and tool message when there are tool calls with results 1`] = `
 [
   {