Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat (core): add extractReasoningMiddleware #4541

Merged
merged 14 commits into from
Jan 27, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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' }),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we show how it would be used and accessed too?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added a note, I don't want to duplicate everywhere

Copy link
Contributor

@nicoalbanese nicoalbanese Jan 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

heard but I mean more for how you access reasoning - eg. is it await result.reasoning?

});
```

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
Loading