53 lines
1.5 KiB
TypeScript
53 lines
1.5 KiB
TypeScript
import { parsePartialJson } from '@ai-sdk/ui-utils';
|
||
import { z } from 'zod';
|
||
|
||
export type DeepPartial<T> = T extends object
|
||
? T extends Array<any>
|
||
? T
|
||
: { [P in keyof T]?: DeepPartial<T[P]> }
|
||
: T;
|
||
|
||
|
||
export function removeJsonMarkdown(text: string) {
|
||
if (text.startsWith('```json')) {
|
||
text = text.slice(7);
|
||
} else if (text.startsWith('json')) {
|
||
text = text.slice(4);
|
||
} else if (text.startsWith('```')) {
|
||
text = text.slice(3);
|
||
}
|
||
if (text.endsWith('```')) {
|
||
text = text.slice(0, -3);
|
||
}
|
||
return text;
|
||
}
|
||
|
||
/**
|
||
* 解析流式的 JSON 数据
|
||
* @param textStream 字符串流
|
||
* @param schema zod schema 用于类型验证
|
||
* @param isValid 自定义验证函数，用于判断解析出的 JSON 是否有效
|
||
* @returns 异步生成器，yield 解析后的数据
|
||
*/
|
||
export async function* parseStreamingJson<T extends z.ZodType>(
|
||
textStream: AsyncIterable<string>,
|
||
schema: T,
|
||
isValid: (value: DeepPartial<z.infer<T>>) => boolean
|
||
): AsyncGenerator<DeepPartial<z.infer<T>>> {
|
||
let rawText = '';
|
||
let isParseSuccessful = false;
|
||
|
||
for await (const chunk of textStream) {
|
||
rawText = removeJsonMarkdown(rawText + chunk);
|
||
const parsed = parsePartialJson(rawText);
|
||
|
||
isParseSuccessful = parsed.state === 'repaired-parse' || parsed.state === 'successful-parse';
|
||
if (isParseSuccessful) {
|
||
yield parsed.value as DeepPartial<z.infer<T>>;
|
||
} else {
|
||
console.dir(parsed, { depth: null, colors: true });
|
||
}
|
||
}
|
||
|
||
return { isSuccessful: isParseSuccessful };
|
||
} |