116 lines
4.9 KiB
TypeScript
116 lines
4.9 KiB
TypeScript
import { ToolSet, type TextStreamPart, type ToolCallPart } from 'ai';
|
|
|
|
export function createLongcatTransformer<TOOLS extends ToolSet>(): (options: {
|
|
tools: TOOLS;
|
|
stopStream: () => void;
|
|
}) => TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>> {
|
|
let buffer = '';
|
|
let hasToolCallInStep = false;
|
|
let lastChunkId: string | undefined;
|
|
let lastChunkType: 'text' | 'reasoning' | undefined;
|
|
let step = 0;
|
|
|
|
return (_opts) => {
|
|
return new TransformStream<TextStreamPart<TOOLS>, TextStreamPart<TOOLS>>({
|
|
transform(chunk, controller) {
|
|
if (chunk.type === 'finish-step' || chunk.type === 'finish') {
|
|
step++;
|
|
|
|
if (hasToolCallInStep) {
|
|
// We clone the chunk and overwrite the finishReason.
|
|
// This tricks the SDK into thinking the model requested a tool natively.
|
|
const modifiedChunk = {
|
|
...chunk,
|
|
finishReason: 'tool-calls' as const,
|
|
};
|
|
|
|
// Reset for the next potential step
|
|
if (chunk.type === 'finish-step') {
|
|
hasToolCallInStep = false;
|
|
}
|
|
|
|
controller.enqueue(modifiedChunk);
|
|
return;
|
|
}
|
|
}
|
|
|
|
if (chunk.type === 'text-start' || chunk.type === 'reasoning-start') {
|
|
lastChunkId = chunk.id;
|
|
lastChunkType = chunk.type.split('-')[1] as 'text' | 'reasoning';
|
|
}
|
|
|
|
// We only care about text chunks
|
|
if (chunk.type !== 'text-delta' && chunk.type !== 'reasoning-delta') {
|
|
controller.enqueue(chunk);
|
|
return;
|
|
}
|
|
|
|
buffer += chunk.text;
|
|
|
|
// Check if we have a full tool call in the buffer
|
|
const pattern = /<longcat_tool_call>([\s\S]*?)<\/longcat_tool_call>/g;
|
|
let lastIndex = 0;
|
|
let match;
|
|
|
|
while ((match = pattern.exec(buffer)) !== null) {
|
|
console.log("longcat tool call found at index", match.index);
|
|
|
|
// 1. Enqueue any text that appeared BEFORE the tool call
|
|
const textBefore = buffer.substring(lastIndex, match.index);
|
|
if (textBefore) {
|
|
controller.enqueue({ type: chunk.type, text: textBefore, id: lastChunkId ?? chunk.type.includes('reasoning') ? `reasoning-${step}` : `text-${step}` });
|
|
}
|
|
|
|
// 2. Parse the XML content
|
|
const content = match[1]!.trim();
|
|
const toolNameMatch = content.match(/^([^\s<]+)/);
|
|
|
|
if (toolNameMatch) {
|
|
hasToolCallInStep = true;
|
|
|
|
const toolName = toolNameMatch[1];
|
|
const args: Record<string, any> = {};
|
|
const argRegex = /<longcat_arg_key>(.*?)<\/longcat_arg_key>\s*<longcat_arg_value>(.*?)<\/longcat_arg_value>/gs;
|
|
|
|
let argMatch;
|
|
while ((argMatch = argRegex.exec(content)) !== null) {
|
|
args[argMatch[1]!.trim()] = argMatch[2]!.trim();
|
|
}
|
|
|
|
// 3. EMIT A TOOL CALL PART
|
|
// This is the "magic" - the SDK will see this and act as if the LLM
|
|
// called a native tool.
|
|
const toolCallId = `lc-${Date.now()}-${Math.random().toString(36).substr(2, 5)}`;
|
|
controller.enqueue({
|
|
type: 'tool-call',
|
|
// @ts-ignore
|
|
id: toolCallId,
|
|
toolCallId,
|
|
toolName,
|
|
input: args,
|
|
dynamic: true,
|
|
});
|
|
}
|
|
|
|
lastIndex = pattern.lastIndex;
|
|
}
|
|
|
|
// Keep the remaining buffer (unclosed tags) for the next chunk
|
|
buffer = buffer.substring(lastIndex);
|
|
|
|
// If there's no open tag starting, we can flush the buffer as text
|
|
if (!buffer.includes('<longcat_tool_call>')) {
|
|
if (buffer) {
|
|
controller.enqueue({ type: chunk.type, text: buffer, id: lastChunkId ?? chunk.type.includes('reasoning') ? `reasoning-${step}` : `text-${step}` });
|
|
buffer = '';
|
|
}
|
|
}
|
|
},
|
|
flush(controller) {
|
|
if (buffer && lastChunkId && lastChunkType) {
|
|
controller.enqueue({ type: `${lastChunkType}-delta`, text: buffer, id: lastChunkId });
|
|
}
|
|
}
|
|
});
|
|
};
|
|
} |