84 lines
3.4 KiB
TypeScript
84 lines
3.4 KiB
TypeScript
import type { TextStreamPart, ToolSet } from 'ai';
|
|
|
|
export function transformCerebrasReasoningStream<TOOLS extends ToolSet>(): (options: {
|
|
tools: TOOLS;
|
|
stopStream: () => void;
|
|
}) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>> {
|
|
|
|
return (_opts) => {
|
|
let isThinking = false;
|
|
let currentReasoningId: string | null = null;
|
|
|
|
let bufferedTextStart: TextStreamPart<TOOLS> | null = null;
|
|
let hasEmittedTextStart = false;
|
|
|
|
return new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
|
|
transform(chunk, controller) {
|
|
if (chunk.type === 'text-start') {
|
|
bufferedTextStart = chunk;
|
|
return;
|
|
}
|
|
|
|
if (chunk.type === 'text-delta') {
|
|
let text = chunk.text;
|
|
|
|
if (text.includes('<think>')) {
|
|
isThinking = true;
|
|
currentReasoningId = crypto.randomUUID();
|
|
|
|
const [before, after] = text.split('<think>');
|
|
|
|
if (before && before.trim().length > 0) {
|
|
if (bufferedTextStart && !hasEmittedTextStart) {
|
|
controller.enqueue(bufferedTextStart);
|
|
hasEmittedTextStart = true;
|
|
}
|
|
controller.enqueue({ type: 'text-delta', text: before, id: chunk.id });
|
|
}
|
|
|
|
controller.enqueue({ type: 'reasoning-start', id: currentReasoningId });
|
|
|
|
if (after) {
|
|
controller.enqueue({ type: 'reasoning-delta', text: after, id: currentReasoningId });
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (text.includes('</think>')) {
|
|
isThinking = false;
|
|
const [before, after] = text.split('</think>');
|
|
|
|
if (before && currentReasoningId !== null) {
|
|
controller.enqueue({ type: 'reasoning-delta', text: before, id: currentReasoningId });
|
|
}
|
|
|
|
if (currentReasoningId !== null) {
|
|
controller.enqueue({ type: 'reasoning-end', id: currentReasoningId });
|
|
}
|
|
|
|
if (after && after.length > 0) {
|
|
if (bufferedTextStart && !hasEmittedTextStart) {
|
|
controller.enqueue(bufferedTextStart);
|
|
hasEmittedTextStart = true;
|
|
}
|
|
controller.enqueue({ type: 'text-delta', text: after, id: chunk.id });
|
|
}
|
|
return;
|
|
}
|
|
|
|
if (isThinking && currentReasoningId !== null) {
|
|
controller.enqueue({ type: 'reasoning-delta', text: text, id: currentReasoningId });
|
|
} else {
|
|
if (bufferedTextStart && !hasEmittedTextStart) {
|
|
controller.enqueue(bufferedTextStart);
|
|
hasEmittedTextStart = true;
|
|
}
|
|
controller.enqueue(chunk);
|
|
}
|
|
} else {
|
|
controller.enqueue(chunk);
|
|
}
|
|
},
|
|
});
|
|
}
|
|
} |