mirror of
https://github.com/hellodigua/ChatLab.git
synced 2026-10-26 02:04:37 +08:00
fix(import): stabilize large and multi-file imports
This commit is contained in:
1 parent
4518a3aaa9
commit
c0942a72f0
6 files changed
+201
-17
No files matched your search
@@ -0,0 +1,83 @@
|
||||
import assert from 'node:assert/strict'
|
||||
import fs from 'node:fs'
|
||||
import os from 'node:os'
|
||||
import path from 'node:path'
|
||||
import test from 'node:test'
|
||||
import { MessageType } from '@openchatlab/shared-types'
|
||||
|
||||
import { findEntryFileInDirectory, parseFileWithFormat } from '../index'
|
||||
|
||||
function makeTempDir(): string {
|
||||
const baseDir = process.env.CHATLAB_TEST_TMPDIR ?? (fs.existsSync('/private/tmp') ? '/private/tmp' : os.tmpdir())
|
||||
return fs.mkdtempSync(path.join(baseDir, 'chatlab-jsonl-parser-'))
|
||||
}
|
||||
|
||||
function writeChatLabJsonl(filePath: string, messageCount: number): void {
|
||||
const lines = [
|
||||
JSON.stringify({
|
||||
_type: 'header',
|
||||
chatlab: { version: '0.0.2', exportedAt: 1711468800 },
|
||||
meta: { name: 'JSONL Test', platform: 'wechat', type: 'group', groupId: 'group-1' },
|
||||
}),
|
||||
JSON.stringify({ _type: 'member', platformId: 'member-1', accountName: 'Alice' }),
|
||||
]
|
||||
|
||||
for (let index = 0; index < messageCount; index++) {
|
||||
lines.push(
|
||||
JSON.stringify({
|
||||
_type: 'message',
|
||||
sender: 'member-1',
|
||||
accountName: 'Alice',
|
||||
timestamp: 1711468800 + index,
|
||||
type: MessageType.TEXT,
|
||||
content: `message-${index}`,
|
||||
})
|
||||
)
|
||||
}
|
||||
|
||||
fs.writeFileSync(filePath, `${lines.join('\n')}\n`, 'utf-8')
|
||||
}
|
||||
|
||||
test('ChatLab JSONL emits message batches and progress while reading', async () => {
|
||||
const root = makeTempDir()
|
||||
const filePath = path.join(root, 'large-chat.jsonl')
|
||||
writeChatLabJsonl(filePath, 12)
|
||||
|
||||
try {
|
||||
const batchSizes: number[] = []
|
||||
let progressCalls = 0
|
||||
let progressCallsAtFirstBatch: number | null = null
|
||||
|
||||
for await (const event of parseFileWithFormat('chatlab-jsonl', {
|
||||
filePath,
|
||||
batchSize: 5,
|
||||
onProgress: () => {
|
||||
progressCalls++
|
||||
},
|
||||
})) {
|
||||
if (event.type === 'messages') {
|
||||
progressCallsAtFirstBatch ??= progressCalls
|
||||
batchSizes.push(event.data.length)
|
||||
}
|
||||
}
|
||||
|
||||
assert.deepEqual(batchSizes, [5, 5, 2])
|
||||
assert.equal(progressCallsAtFirstBatch, 1)
|
||||
assert.equal(progressCalls, 4)
|
||||
} finally {
|
||||
fs.rmSync(root, { recursive: true, force: true })
|
||||
}
|
||||
})
|
||||
|
||||
test('directory entry detection accepts a top-level ChatLab JSONL file', () => {
|
||||
const root = makeTempDir()
|
||||
const filePath = path.join(root, 'chat.jsonl')
|
||||
fs.mkdirSync(path.join(root, 'media'))
|
||||
writeChatLabJsonl(filePath, 1)
|
||||
|
||||
try {
|
||||
assert.equal(findEntryFileInDirectory(root), filePath)
|
||||
} finally {
|
||||
fs.rmSync(root, { recursive: true, force: true })
|
||||
}
|
||||
})
|
||||
@@ -225,8 +225,10 @@ async function* parseChatLabJsonl(options: ParseOptions): AsyncGenerator<ParseEv
|
||||
})
|
||||
messagesProcessed++
|
||||
|
||||
// 达到批次大小,发送进度
|
||||
// Flush each batch immediately so large files do not retain every message or flood progress IPC.
|
||||
if (messageBatch.length >= batchSize) {
|
||||
yield { type: 'messages', data: messageBatch.splice(0) }
|
||||
|
||||
const progress = createProgress(
|
||||
'parsing',
|
||||
bytesRead,
|
||||
@@ -234,6 +236,7 @@ async function* parseChatLabJsonl(options: ParseOptions): AsyncGenerator<ParseEv
|
||||
messagesProcessed,
|
||||
`已处理 ${messagesProcessed.toLocaleString()} 条消息...`
|
||||
)
|
||||
yield { type: 'progress', data: progress }
|
||||
onProgress?.(progress)
|
||||
}
|
||||
break
|
||||
@@ -257,10 +260,9 @@ async function* parseChatLabJsonl(options: ParseOptions): AsyncGenerator<ParseEv
|
||||
yield { type: 'members', data: Array.from(memberMap.values()) }
|
||||
}
|
||||
|
||||
// 分批发送消息
|
||||
for (let i = 0; i < messageBatch.length; i += batchSize) {
|
||||
const batch = messageBatch.slice(i, i + batchSize)
|
||||
yield { type: 'messages', data: batch }
|
||||
// 发送不足一个批次的剩余消息
|
||||
if (messageBatch.length > 0) {
|
||||
yield { type: 'messages', data: messageBatch }
|
||||
}
|
||||
|
||||
// 完成
|
||||
|
||||
@@ -260,6 +260,13 @@ export async function parseFileInfo(
|
||||
* @returns path to the entry file, or null if none found
|
||||
*/
|
||||
export function findEntryFileInDirectory(dirPath: string, maxDepth = 1): string | null {
|
||||
const supportedExtensions = new Set(
|
||||
sniffer
|
||||
.getSupportedFormats()
|
||||
.flatMap((feature) => feature.extensions)
|
||||
.map((extension) => extension.toLowerCase())
|
||||
)
|
||||
|
||||
function walk(current: string, depth: number): string | null {
|
||||
if (depth > maxDepth) return null
|
||||
let entries: fs.Dirent[]
|
||||
@@ -270,7 +277,7 @@ export function findEntryFileInDirectory(dirPath: string, maxDepth = 1): string
|
||||
}
|
||||
for (const entry of entries) {
|
||||
const full = path.join(current, entry.name)
|
||||
if (entry.isFile() && entry.name.endsWith('.json')) {
|
||||
if (entry.isFile() && supportedExtensions.has(path.extname(entry.name).toLowerCase())) {
|
||||
const format = detectFormat(full)
|
||||
if (format) return full
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user