| 1 | import fs from 'fs' |
| 2 | import path from 'path' |
| 3 | import { createRequire } from 'module' |
| 4 | import log from 'electron-log/main.js' |
| 5 | import { Jieba } from '@node-rs/jieba' |
| 6 | |
| 7 | type JiebaDictModule = { |
| 8 | dict: Uint8Array |
| 9 | } |
| 10 | |
| 11 | export type ReferenceDocumentSnippet = { |
| 12 | sourcePath: string |
| 13 | headingPath: string[] |
| 14 | text: string |
| 15 | startLine: number |
| 16 | endLine: number |
| 17 | score: number |
| 18 | } |
| 19 | |
| 20 | type ReferenceChunk = { |
| 21 | id: string |
| 22 | sourcePath: string |
| 23 | absolutePath: string |
| 24 | headingPath: string[] |
| 25 | text: string |
| 26 | normalizedText: string |
| 27 | normalizedHeading: string |
| 28 | startLine: number |
| 29 | endLine: number |
| 30 | } |
| 31 | |
| 32 | type SearchInput = { |
| 33 | pageId: string |
| 34 | pageTitle: string |
| 35 | pageOutline: string |
| 36 | userMessage: string |
| 37 | } |
| 38 | |
| 39 | type ReferenceDocumentRetriever = { |
| 40 | search: (input: SearchInput) => ReferenceDocumentSnippet[] |
| 41 | } |
| 42 | |
| 43 | let jiebaInstance: Jieba | null = null |
| 44 | |
| 45 | const getJieba = (): Jieba => { |
| 46 | if (!jiebaInstance) { |
| 47 | const require = createRequire(import.meta.url) |
| 48 | const { dict } = require('@node-rs/jieba/dict.js') as JiebaDictModule |
| 49 | jiebaInstance = Jieba.withDict(dict) |
| 50 | } |
| 51 | return jiebaInstance |
| 52 | } |
| 53 | |
| 54 | const MAX_CHUNK_CHARS = 1200 |
| 55 | const MIN_CHUNK_CHARS = 300 |
| 56 | const MAX_SNIPPETS = 5 |
| 57 | const MAX_INJECTED_CHARS = 6000 |
| 58 | |
| 59 | const GENERIC_TERMS = new Set([ |
| 60 | '背景', |
| 61 | '目标', |
| 62 | '价值', |
| 63 | '方案', |
| 64 | '介绍', |
| 65 | '说明', |
| 66 | '分析', |
| 67 | '概述', |
| 68 | '总结', |
| 69 | '规划', |
| 70 | '设计', |
| 71 | '能力', |
| 72 | '功能', |
| 73 | '核心', |
| 74 | '重点', |
| 75 | '整体', |
| 76 | '业务', |
| 77 | '内容', |
| 78 | '页面', |
| 79 | '展示' |
| 80 | ]) |
| 81 | |
| 82 | const normalizeForSearch = (value: string): string => |
| 83 | value |
| 84 | .toLowerCase() |
| 85 | .replace(/\s+/g, '') |
| 86 | .replace(/[，。！？；：、“”‘’（）【】《》,.!?;:"'()[\]<>]/g, '') |
| 87 | .trim() |
| 88 | |
| 89 | const cleanText = (value: string): string => value.replace(/\s+/g, ' ').trim() |
| 90 | |
| 91 | const isMostlyChinese = (value: string): boolean => /[\u4e00-\u9fff]/.test(value) |
| 92 | |
| 93 | const isGenericTerm = (value: string): boolean => { |
| 94 | const normalized = normalizeForSearch(value) |
| 95 | if (!normalized) return true |
| 96 | if (GENERIC_TERMS.has(normalized)) return true |
| 97 | if (normalized.length <= 1) return true |
| 98 | return false |
| 99 | } |
| 100 | |
| 101 | const splitQueryPhrases = (input: string): string[] => { |
| 102 | const phrases = input |
| 103 | .split(/[\n\r。！？；;：:，,、|/\\()[\]【】《》<>]+/g) |
| 104 | .map((item) => item.replace(/^\s*(?:第?\d+[.、）)]|[-*•]+)\s*/u, '').trim()) |
| 105 | .filter((item) => item.length > 0) |
| 106 | .filter((item) => item.length >= 3 && item.length <= 40) |
| 107 | .filter((item) => !isGenericTerm(item)) |
| 108 | return Array.from(new Set(phrases)).slice(0, 8) |
| 109 | } |
| 110 | |
| 111 | const extractStructuredTokens = (input: string): string[] => { |
| 112 | const tokenMatches = input.match( |
| 113 | /(?:[Pp][0-9]+)|(?:\d{4}[-/.年]\d{1,2}(?:[-/.月]\d{1,2}日?)?)|(?:\d+(?:\.\d+)?%)|(?:\d+(?:\.\d+)?(?:万|亿|元|万元|亿元)?)|(?:[A-Za-z][A-Za-z0-9_-]{1,})/g |
| 114 | ) |
| 115 | return Array.from(new Set(tokenMatches || [])).filter((item) => !isGenericTerm(item)) |
| 116 | } |
| 117 | |
| 118 | const extractWordQueries = (input: string): string[] => { |
| 119 | const words = getJieba().cutForSearch(input, false) |
| 120 | const structured = extractStructuredTokens(input) |
| 121 | const candidates = [...words, ...structured] |
| 122 | .map((item) => item.trim()) |
| 123 | .filter((item) => item.length > 0) |
| 124 | .filter((item) => { |
| 125 | if (isGenericTerm(item)) return false |
| 126 | if (isMostlyChinese(item)) return item.length >= 2 && item.length <= 12 |
| 127 | return item.length >= 2 && item.length <= 32 |
| 128 | }) |
| 129 | return Array.from(new Set(candidates)).slice(0, 20) |
| 130 | } |
| 131 | |
| 132 | const createQueries = (input: SearchInput): { phraseQueries: string[]; wordQueries: string[] } => { |
| 133 | const source = [input.pageTitle, input.pageOutline, input.userMessage].filter(Boolean).join('\n') |
| 134 | const phraseQueries = splitQueryPhrases(source) |
| 135 | const wordQueries = extractWordQueries(source).filter( |
| 136 | (word) => !phraseQueries.some((phrase) => normalizeForSearch(phrase) === normalizeForSearch(word)) |
| 137 | ) |
| 138 | return { phraseQueries, wordQueries } |
| 139 | } |
| 140 | |
| 141 | const splitLongText = (text: string): string[] => { |
| 142 | const normalized = text.trim() |
| 143 | if (normalized.length <= MAX_CHUNK_CHARS) return [normalized] |
| 144 | const parts = normalized |
| 145 | .split(/(?<=[。！？；;.!?])\s*/u) |
| 146 | .map((item) => item.trim()) |
| 147 | .filter(Boolean) |
| 148 | const chunks: string[] = [] |
| 149 | let current = '' |
| 150 | for (const part of parts.length > 0 ? parts : [normalized]) { |
| 151 | if (current && `${current}${part}`.length > MAX_CHUNK_CHARS) { |
| 152 | chunks.push(current.trim()) |
| 153 | current = '' |
| 154 | } |
| 155 | current = current ? `${current}${part}` : part |
| 156 | } |
| 157 | if (current.trim()) chunks.push(current.trim()) |
| 158 | return chunks.length > 0 ? chunks : [normalized.slice(0, MAX_CHUNK_CHARS)] |
| 159 | } |
| 160 | |
| 161 | const createChunk = (args: { |
| 162 | sourcePath: string |
| 163 | absolutePath: string |
| 164 | headingPath: string[] |
| 165 | text: string |
| 166 | startLine: number |
| 167 | endLine: number |
| 168 | index: number |
| 169 | }): ReferenceChunk => { |
| 170 | const text = cleanText(args.text) |
| 171 | const headingPath = args.headingPath.filter(Boolean) |
| 172 | return { |
| 173 | id: `${args.sourcePath}:${args.startLine}-${args.endLine}:${args.index}`, |
| 174 | sourcePath: args.sourcePath, |
| 175 | absolutePath: args.absolutePath, |
| 176 | headingPath, |
| 177 | text, |
| 178 | normalizedText: normalizeForSearch(text), |
| 179 | normalizedHeading: normalizeForSearch(headingPath.join(' ')), |
| 180 | startLine: args.startLine, |
| 181 | endLine: args.endLine |
| 182 | } |
| 183 | } |
| 184 | |
| 185 | const chunkDocument = (args: { |
| 186 | sourcePath: string |
| 187 | absolutePath: string |
| 188 | content: string |
| 189 | }): ReferenceChunk[] => { |
| 190 | const lines = args.content.replace(/\r\n/g, '\n').replace(/\r/g, '\n').split('\n') |
| 191 | const chunks: ReferenceChunk[] = [] |
| 192 | const headingPath: string[] = [] |
| 193 | let buffer: string[] = [] |
| 194 | let bufferStartLine = 1 |
| 195 | let chunkIndex = 0 |
| 196 | |
| 197 | const pushBuffer = (endLine: number) => { |
| 198 | const text = buffer.join('\n').trim() |
| 199 | if (!text) { |
| 200 | buffer = [] |
| 201 | return |
| 202 | } |
| 203 | for (const part of splitLongText(text)) { |
| 204 | chunks.push( |
| 205 | createChunk({ |
| 206 | sourcePath: args.sourcePath, |
| 207 | absolutePath: args.absolutePath, |
| 208 | headingPath, |
| 209 | text: part, |
| 210 | startLine: bufferStartLine, |
| 211 | endLine, |
| 212 | index: chunkIndex |
| 213 | }) |
| 214 | ) |
| 215 | chunkIndex += 1 |
| 216 | } |
| 217 | buffer = [] |
| 218 | } |
| 219 | |
| 220 | const appendLine = (line: string, lineNumber: number) => { |
| 221 | if (buffer.length === 0) bufferStartLine = lineNumber |
| 222 | buffer.push(line) |
| 223 | const currentLength = buffer.join('\n').length |
| 224 | if (currentLength >= MAX_CHUNK_CHARS) { |
| 225 | pushBuffer(lineNumber) |
| 226 | } |
| 227 | } |
| 228 | |
| 229 | for (let index = 0; index < lines.length; index += 1) { |
| 230 | const line = lines[index] |
| 231 | const lineNumber = index + 1 |
| 232 | const headingMatch = line.match(/^(#{1,6})\s+(.+?)\s*#*$/) |
| 233 | if (headingMatch) { |
| 234 | pushBuffer(lineNumber - 1) |
| 235 | const level = headingMatch[1].length |
| 236 | headingPath.splice(level - 1) |
| 237 | headingPath[level - 1] = headingMatch[2].trim() |
| 238 | continue |
| 239 | } |
| 240 | if (!line.trim()) { |
| 241 | pushBuffer(lineNumber - 1) |
| 242 | continue |
| 243 | } |
| 244 | appendLine(line, lineNumber) |
| 245 | } |
| 246 | pushBuffer(lines.length) |
| 247 | |
| 248 | const merged: ReferenceChunk[] = [] |
| 249 | for (const chunk of chunks) { |
| 250 | const prev = merged[merged.length - 1] |
| 251 | if ( |
| 252 | prev && |
| 253 | prev.text.length < MIN_CHUNK_CHARS && |
| 254 | chunk.text.length < MIN_CHUNK_CHARS && |
| 255 | prev.sourcePath === chunk.sourcePath && |
| 256 | prev.headingPath.join('/') === chunk.headingPath.join('/') |
| 257 | ) { |
| 258 | const text = `${prev.text}\n${chunk.text}`.trim() |
| 259 | merged[merged.length - 1] = { |
| 260 | ...prev, |
| 261 | text, |
| 262 | normalizedText: normalizeForSearch(text), |
| 263 | endLine: chunk.endLine |
| 264 | } |
| 265 | continue |
| 266 | } |
| 267 | merged.push(chunk) |
| 268 | } |
| 269 | return merged |
| 270 | } |
| 271 | |
| 272 | const resolveSourcePath = (projectDir: string, sourceDocumentPath: string): string | null => { |
| 273 | if (!sourceDocumentPath.startsWith('/docs/')) return null |
| 274 | const absolutePath = path.resolve(projectDir, sourceDocumentPath.replace(/^\/+/, '')) |
| 275 | const relativeToProject = path.relative(projectDir, absolutePath) |
| 276 | if (relativeToProject.startsWith('..') || path.isAbsolute(relativeToProject)) return null |
| 277 | return absolutePath |
| 278 | } |
| 279 | |
| 280 | const scoreChunk = ( |
| 281 | chunk: ReferenceChunk, |
| 282 | queries: { phraseQueries: string[]; wordQueries: string[] } |
| 283 | ): number => { |
| 284 | let score = 0 |
| 285 | const seen = new Set<string>() |
| 286 | for (const phrase of queries.phraseQueries) { |
| 287 | const normalized = normalizeForSearch(phrase) |
| 288 | if (!normalized || seen.has(`phrase:${normalized}`)) continue |
| 289 | seen.add(`phrase:${normalized}`) |
| 290 | if (chunk.normalizedText.includes(normalized)) score += 2 |
| 291 | if (chunk.normalizedHeading.includes(normalized)) score += 2 |
| 292 | } |
| 293 | for (const word of queries.wordQueries) { |
| 294 | const normalized = normalizeForSearch(word) |
| 295 | if (!normalized || seen.has(`word:${normalized}`)) continue |
| 296 | seen.add(`word:${normalized}`) |
| 297 | if (chunk.normalizedText.includes(normalized)) score += 1 |
| 298 | if (chunk.normalizedHeading.includes(normalized)) score += 2 |
| 299 | } |
| 300 | return score |
| 301 | } |
| 302 | |
| 303 | const selectSnippets = ( |
| 304 | chunks: ReferenceChunk[], |
| 305 | queries: { phraseQueries: string[]; wordQueries: string[] } |
| 306 | ): { snippets: ReferenceDocumentSnippet[]; matchedChunkCount: number } => { |
| 307 | const scored = chunks |
| 308 | .map((chunk) => ({ chunk, score: scoreChunk(chunk, queries) })) |
| 309 | .filter((item) => item.score > 0) |
| 310 | .sort((a, b) => b.score - a.score || a.chunk.startLine - b.chunk.startLine) |
| 311 | |
| 312 | const snippets: ReferenceDocumentSnippet[] = [] |
| 313 | let injectedChars = 0 |
| 314 | for (const item of scored) { |
| 315 | if (snippets.length >= MAX_SNIPPETS) break |
| 316 | if (injectedChars >= MAX_INJECTED_CHARS) break |
| 317 | const text = |
| 318 | item.chunk.text.length > 1400 ? `${item.chunk.text.slice(0, 1400).trimEnd()}...` : item.chunk.text |
| 319 | snippets.push({ |
| 320 | sourcePath: item.chunk.sourcePath, |
| 321 | headingPath: item.chunk.headingPath, |
| 322 | text, |
| 323 | startLine: item.chunk.startLine, |
| 324 | endLine: item.chunk.endLine, |
| 325 | score: item.score |
| 326 | }) |
| 327 | injectedChars += text.length |
| 328 | } |
| 329 | return { snippets, matchedChunkCount: scored.length } |
| 330 | } |
| 331 | |
| 332 | export const formatReferenceDocumentSnippets = ( |
| 333 | snippets: ReferenceDocumentSnippet[] |
| 334 | ): string => { |
| 335 | if (snippets.length === 0) return '' |
| 336 | return [ |
| 337 | '参考文档检索片段（程序侧根据当前页标题和大纲预检索，优先使用）：', |
| 338 | '', |
| 339 | ...snippets.flatMap((snippet, index) => [ |
| 340 | `[片段 ${index + 1}] ${snippet.sourcePath}#L${snippet.startLine}-L${snippet.endLine}`, |
| 341 | snippet.headingPath.length > 0 ? `标题路径：${snippet.headingPath.join(' / ')}` : '', |
| 342 | `内容：${snippet.text}`, |
| 343 | '' |
| 344 | ]) |
| 345 | ] |
| 346 | .filter((line) => line !== '') |
| 347 | .join('\n') |
| 348 | } |
| 349 | |
| 350 | export const createReferenceDocumentRetriever = async (args: { |
| 351 | sessionId: string |
| 352 | projectDir: string |
| 353 | sourceDocumentPaths?: string[] |
| 354 | }): Promise<ReferenceDocumentRetriever | null> => { |
| 355 | const sourceDocumentPaths = (args.sourceDocumentPaths || []).filter(Boolean) |
| 356 | if (sourceDocumentPaths.length === 0) return null |
| 357 | |
| 358 | const chunkCache = new Map<string, ReferenceChunk[]>() |
| 359 | for (const sourcePath of sourceDocumentPaths) { |
| 360 | const absolutePath = resolveSourcePath(args.projectDir, sourcePath) |
| 361 | if (!absolutePath || !fs.existsSync(absolutePath)) { |
| 362 | log.warn('[referenceDocument:grep] source missing', { |
| 363 | sessionId: args.sessionId, |
| 364 | sourcePath |
| 365 | }) |
| 366 | continue |
| 367 | } |
| 368 | try { |
| 369 | const content = await fs.promises.readFile(absolutePath, 'utf-8') |
| 370 | const chunks = chunkDocument({ sourcePath, absolutePath, content }) |
| 371 | chunkCache.set(sourcePath, chunks) |
| 372 | log.info('[referenceDocument:grep] chunked', { |
| 373 | sessionId: args.sessionId, |
| 374 | sourcePath, |
| 375 | chunkCount: chunks.length, |
| 376 | characterCount: content.length |
| 377 | }) |
| 378 | } catch (error) { |
| 379 | log.warn('[referenceDocument:grep] read failed', { |
| 380 | sessionId: args.sessionId, |
| 381 | sourcePath, |
| 382 | message: error instanceof Error ? error.message : String(error) |
| 383 | }) |
| 384 | } |
| 385 | } |
| 386 | |
| 387 | const chunks = Array.from(chunkCache.values()).flat() |
| 388 | if (chunks.length === 0) return null |
| 389 | |
| 390 | return { |
| 391 | search: (input: SearchInput): ReferenceDocumentSnippet[] => { |
| 392 | const queries = createQueries(input) |
| 393 | const { snippets, matchedChunkCount } = selectSnippets(chunks, queries) |
| 394 | log.info('[referenceDocument:grep] search', { |
| 395 | sessionId: args.sessionId, |
| 396 | pageId: input.pageId, |
| 397 | sourceDocumentPaths, |
| 398 | phraseQueryCount: queries.phraseQueries.length, |
| 399 | wordQueryCount: queries.wordQueries.length, |
| 400 | chunkCount: chunks.length, |
| 401 | matchedChunkCount, |
| 402 | injectedChunkCount: snippets.length |
| 403 | }) |
| 404 | return snippets |
| 405 | } |
| 406 | } |
| 407 | } |
| 408 |