Skip to content

Commit

Permalink
feat (core): add extractReasoningMiddleware (#4541)
Browse files Browse the repository at this point in the history
  • Loading branch information
lgrammel authored Jan 27, 2025
1 parent 3e65901 commit 92f5f36
Show file tree
Hide file tree
Showing 11 changed files with 616 additions and 0 deletions.
5 changes: 5 additions & 0 deletions .changeset/itchy-pumpkins-punch.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'ai': patch
---

feat (core): add extractReasoningMiddleware
23 changes: 23 additions & 0 deletions content/docs/03-ai-sdk-core/45-middleware.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,29 @@ const result = streamText({
});
```

## Built-in Middleware

### Extract Reasoning

Some providers and models expose reasoning information in the generated text using special tags,
e.g. <think> and </think>.

The `extractReasoningMiddleware` function can be used to extract this reasoning information and expose it as a `reasoning` property on the result.

```ts
import {
experimental_wrapLanguageModel as wrapLanguageModel,
extractReasoningMiddleware,
} from 'ai';

const model = wrapLanguageModel({
model: yourModel,
middleware: extractReasoningMiddleware({ tagName: 'think' }),
});
```

You can then use that enhanced model in functions like `generateText` and `streamText`.

## Implementing Language Model Middleware

<Note>
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
---
title: extractReasoningMiddleware
description: Middleware that extracts XML-tagged reasoning sections from generated text
---

# `extractReasoningMiddleware()`

`extractReasoningMiddleware` is a middleware function that extracts XML-tagged reasoning sections from generated text and exposes them separately from the main text content. This is particularly useful when you want to separate an AI model's reasoning process from its final output.

```ts
import { extractReasoningMiddleware } from 'ai';

const middleware = extractReasoningMiddleware({
tagName: 'reasoning',
separator: '\n',
});
```

## Import

<Snippet
text={`import { extractReasoningMiddleware } from "ai"`}
prompt={false}
/>

## API Signature

### Parameters

<PropertiesTable
content={[
{
name: 'tagName',
type: 'string',
isOptional: false,
description:
'The name of the XML tag to extract reasoning from (without angle brackets)',
},
{
name: 'separator',
type: 'string',
isOptional: true,
description:
'The separator to use between reasoning and text sections. Defaults to "\\n"',
},
]}
/>

### Returns

Returns a middleware object that:

- Processes both streaming and non-streaming responses
- Extracts content between specified XML tags as reasoning
- Removes the XML tags and reasoning from the main text
- Adds a `reasoning` property to the result containing the extracted content
- Maintains proper separation between text sections using the specified separator

### Type Parameters

The middleware works with the `LanguageModelV1StreamPart` type for streaming responses.
11 changes: 11 additions & 0 deletions content/docs/07-reference/01-ai-sdk-core/index.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,17 @@ It also contains the following helper functions:
'Creates a ReadableStream that emits values with configurable delays.',
href: '/docs/reference/ai-sdk-core/simulate-readable-stream',
},
{
title: 'wrapLanguageModel()',
description: 'Wraps a language model with middleware.',
href: '/docs/reference/ai-sdk-core/wrap-language-model',
},
{
title: 'extractReasoningMiddleware()',
description:
'Extracts reasoning from the generated text and exposes it as a `reasoning` property on the result.',
href: '/docs/reference/ai-sdk-core/extract-reasoning-middleware',
},
{
title: 'smoothStream()',
description: 'Smooths text streaming output.',
Expand Down
45 changes: 45 additions & 0 deletions examples/ai-core/src/stream-text/groq-reasoning-fullstream.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
import { groq } from '@ai-sdk/groq';
import {
experimental_wrapLanguageModel,
extractReasoningMiddleware,
streamText,
} from 'ai';
import 'dotenv/config';

async function main() {
const result = streamText({
model: experimental_wrapLanguageModel({
model: groq('deepseek-r1-distill-llama-70b'),
middleware: extractReasoningMiddleware({ tagName: 'think' }),
}),
prompt: 'Invent a new holiday and describe its traditions.',
});

let enteredReasoning = false;
let enteredText = false;
for await (const part of result.fullStream) {
if (part.type === 'reasoning') {
if (!enteredReasoning) {
enteredReasoning = true;
console.log('\nSTREAMING REASONING:\n');
}
process.stdout.write(part.textDelta);
} else if (part.type === 'text-delta') {
if (!enteredText) {
enteredText = true;
console.log('\nSTREAMING TEXT:\n');
}
process.stdout.write(part.textDelta);
}
}

console.log();
console.log('\nFINAL REASONING:\n', await result.reasoning);
console.log('\nFINAL TEXT:\n', await result.text);

console.log();
console.log('Token usage:', await result.usage);
console.log('Finish reason:', await result.finishReason);
}

main().catch(console.error);
Loading

0 comments on commit 92f5f36

Please sign in to comment.