542 lines
16 KiB
TypeScript
542 lines
16 KiB
TypeScript
// src/main/services/RAGService.ts
|
|
// Phase 13.2: 로컬 RAG — 문서 임베딩 + 코사인 유사도 검색 + LLM 컨텍스트 주입
|
|
// Ollama nomic-embed-text 모델, SQLite에 JSON 직렬화 벡터 저장.
|
|
|
|
import { EventEmitter } from 'events'
|
|
import path from 'path'
|
|
import { eq } from 'drizzle-orm'
|
|
import { getLogger } from './LoggerService'
|
|
import { getPremiumLLMService } from './PremiumLLMService'
|
|
import { configGet } from './ConfigService'
|
|
import { getDatabase } from '../db'
|
|
import { ragDocuments, ragChunks } from '../db/schema'
|
|
import { getMainWindow } from '../windows/WindowManager'
|
|
import { IPC_CHANNELS } from '@d3ro/core/ipc-channels'
|
|
import { D3ROError, ErrorCode } from '@d3ro/core/errors'
|
|
import type {
|
|
RAGDocument,
|
|
RAGQueryResult,
|
|
RAGState,
|
|
RAGStateInfo,
|
|
RAGIndexProgress,
|
|
} from '@d3ro/core/types'
|
|
|
|
const logger = getLogger('RAGService')
|
|
|
|
/** 임베딩 모델 */
|
|
const EMBED_MODEL = 'nomic-embed-text'
|
|
/** 청크 크기 (자) */
|
|
const CHUNK_SIZE = 500
|
|
/** 청크 오버랩 (자) */
|
|
const CHUNK_OVERLAP = 50
|
|
/** 검색 기본 topK */
|
|
const DEFAULT_TOP_K = 5
|
|
|
|
class RAGService extends EventEmitter {
|
|
private _state: RAGState = 'idle'
|
|
|
|
get state(): RAGState {
|
|
return this._state
|
|
}
|
|
|
|
getStateInfo(): RAGStateInfo {
|
|
const db = getDatabase()
|
|
const docs = db.select().from(ragDocuments).all()
|
|
const chunks = db.select().from(ragChunks).all()
|
|
return {
|
|
state: this._state,
|
|
documentCount: docs.length,
|
|
totalChunks: chunks.length,
|
|
}
|
|
}
|
|
|
|
getDocuments(): RAGDocument[] {
|
|
const db = getDatabase()
|
|
const rows = db.select().from(ragDocuments).all()
|
|
return rows.map((r) => ({
|
|
id: r.id,
|
|
fileName: r.fileName,
|
|
filePath: r.filePath,
|
|
fileType: r.fileType as RAGDocument['fileType'],
|
|
chunkCount: r.chunkCount,
|
|
indexed: r.indexed,
|
|
indexedAt: r.indexedAt,
|
|
addedAt: r.addedAt,
|
|
}))
|
|
}
|
|
|
|
/**
|
|
* 문서 추가 + 인덱싱 (청킹 → 임베딩 → DB 저장)
|
|
*/
|
|
async addDocument(filePath: string): Promise<RAGDocument> {
|
|
// 라이센스 체크
|
|
try {
|
|
const { getLicenseService } = await import('./LicenseService')
|
|
const { Feature } = await import('@d3ro/core/types')
|
|
const license = getLicenseService()
|
|
const access = license.canUse(Feature.LOCAL_RAG)
|
|
if (!access.allowed) {
|
|
license.promptUpgrade(Feature.LOCAL_RAG, 'tier_required')
|
|
throw new D3ROError(ErrorCode.FeatureNotAvailable, 'Pro+ required for Local RAG')
|
|
}
|
|
} catch (err) {
|
|
if (err instanceof D3ROError) throw err
|
|
}
|
|
|
|
const ext = path.extname(filePath).toLowerCase()
|
|
const supportedTypes: Record<string, RAGDocument['fileType']> = {
|
|
'.txt': 'txt',
|
|
'.md': 'md',
|
|
'.pdf': 'pdf',
|
|
'.docx': 'docx',
|
|
}
|
|
|
|
const fileType = supportedTypes[ext]
|
|
if (!fileType) {
|
|
throw new D3ROError(ErrorCode.RAGUnsupportedFormat, `Unsupported format: ${ext}. Supported: .txt, .md, .pdf, .docx`)
|
|
}
|
|
|
|
const fileName = path.basename(filePath)
|
|
const docId = crypto.randomUUID()
|
|
|
|
// 텍스트 추출
|
|
let content: string
|
|
try {
|
|
content = await this._extractText(filePath, fileType)
|
|
} catch (err) {
|
|
if (err instanceof D3ROError) throw err
|
|
throw new D3ROError(ErrorCode.RAGIndexingFailed, `Text extraction failed: ${err instanceof Error ? err.message : String(err)}`)
|
|
}
|
|
|
|
if (!content || content.trim().length < 20) {
|
|
throw new D3ROError(ErrorCode.RAGIndexingFailed, 'No readable text content found in document')
|
|
}
|
|
|
|
logger.info(`RAG text extracted: ${fileName} (${content.length} chars)`)
|
|
|
|
// 청킹
|
|
const chunks = this._chunkText(content)
|
|
if (chunks.length === 0) {
|
|
throw new D3ROError(ErrorCode.RAGIndexingFailed, 'Document produced no valid text chunks')
|
|
}
|
|
|
|
// DB에 문서 레코드 삽입
|
|
const db = getDatabase()
|
|
db.insert(ragDocuments).values({
|
|
id: docId,
|
|
fileName,
|
|
filePath,
|
|
fileType,
|
|
chunkCount: chunks.length,
|
|
indexed: false,
|
|
indexedAt: null,
|
|
addedAt: Date.now(),
|
|
}).run()
|
|
|
|
// 비동기 인덱싱 (임베딩 생성)
|
|
this._indexDocument(docId, fileName, chunks).catch((err) => {
|
|
logger.error(`Indexing failed for ${fileName}:`, err)
|
|
})
|
|
|
|
return {
|
|
id: docId,
|
|
fileName,
|
|
filePath,
|
|
fileType,
|
|
chunkCount: chunks.length,
|
|
indexed: false,
|
|
indexedAt: null,
|
|
addedAt: Date.now(),
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 문서 제거 (청크 포함)
|
|
*/
|
|
removeDocument(documentId: string): void {
|
|
const db = getDatabase()
|
|
db.delete(ragChunks).where(eq(ragChunks.documentId, documentId)).run()
|
|
db.delete(ragDocuments).where(eq(ragDocuments.id, documentId)).run()
|
|
logger.info(`RAG document removed: ${documentId}`)
|
|
}
|
|
|
|
/**
|
|
* 문서 재인덱싱
|
|
*/
|
|
async reindex(documentId: string): Promise<void> {
|
|
const db = getDatabase()
|
|
const rows = db.select().from(ragDocuments).where(eq(ragDocuments.id, documentId)).all()
|
|
if (rows.length === 0) {
|
|
throw new D3ROError(ErrorCode.RAGDocumentNotFound, 'Document not found')
|
|
}
|
|
const doc = rows[0]
|
|
|
|
// 기존 청크 삭제
|
|
db.delete(ragChunks).where(eq(ragChunks.documentId, documentId)).run()
|
|
|
|
// 텍스트 재추출 + 재인덱싱
|
|
const content = await this._extractText(doc.filePath, doc.fileType as RAGDocument['fileType'])
|
|
const chunks = this._chunkText(content)
|
|
|
|
db.update(ragDocuments)
|
|
.set({ chunkCount: chunks.length, indexed: false })
|
|
.where(eq(ragDocuments.id, documentId))
|
|
.run()
|
|
|
|
await this._indexDocument(documentId, doc.fileName, chunks)
|
|
}
|
|
|
|
/**
|
|
* 벡터 검색 + LLM 답변 생성
|
|
*/
|
|
async query(queryText: string, topK: number = DEFAULT_TOP_K): Promise<RAGQueryResult> {
|
|
this._state = 'querying'
|
|
|
|
try {
|
|
const db = getDatabase()
|
|
const allChunks = db.select().from(ragChunks).all()
|
|
if (allChunks.length === 0) {
|
|
throw new D3ROError(ErrorCode.RAGQueryFailed, 'No indexed chunks to query')
|
|
}
|
|
|
|
// 쿼리 임베딩
|
|
const queryEmbedding = await this._embed(queryText)
|
|
|
|
const allDocs = db.select().from(ragDocuments).all()
|
|
const docMap = new Map(allDocs.map((d) => [d.id, d.fileName]))
|
|
|
|
const scored = allChunks.map((chunk) => {
|
|
const embedding = JSON.parse(chunk.embedding) as number[]
|
|
const similarity = this._cosineSimilarity(queryEmbedding, embedding)
|
|
return {
|
|
documentId: chunk.documentId,
|
|
fileName: docMap.get(chunk.documentId) ?? 'unknown',
|
|
content: chunk.content,
|
|
similarity,
|
|
}
|
|
})
|
|
|
|
// 상위 topK
|
|
scored.sort((a, b) => b.similarity - a.similarity)
|
|
const topResults = scored.slice(0, topK)
|
|
|
|
// LLM에 컨텍스트 주입
|
|
const context = topResults
|
|
.map((r, i) => `[${i + 1}] (${r.fileName})\n${r.content}`)
|
|
.join('\n\n')
|
|
|
|
const systemPrompt = `You are a helpful assistant. Answer the user's question based on the following documents. If the documents don't contain relevant information, say so. Respond in the same language as the question.
|
|
|
|
Documents:
|
|
${context}`
|
|
|
|
const llmService = getPremiumLLMService()
|
|
const result = await llmService.generate(queryText, { systemPrompt })
|
|
|
|
return {
|
|
query: queryText,
|
|
results: topResults,
|
|
answer: result.text.trim(),
|
|
}
|
|
} finally {
|
|
this._state = 'idle'
|
|
}
|
|
}
|
|
|
|
// ── 내부 메서드 ──
|
|
|
|
private async _indexDocument(docId: string, fileName: string, chunks: string[]): Promise<void> {
|
|
this._state = 'indexing'
|
|
const db = getDatabase()
|
|
|
|
logger.info(`RAG indexing started: ${fileName} (${chunks.length} chunks)`)
|
|
|
|
// 초기 진행률 즉시 전송
|
|
this._sendToRenderer(IPC_CHANNELS.RAG.INDEX_PROGRESS, {
|
|
documentId: docId,
|
|
fileName,
|
|
currentChunk: 0,
|
|
totalChunks: chunks.length,
|
|
percent: 0,
|
|
} as RAGIndexProgress)
|
|
|
|
let successCount = 0
|
|
for (let i = 0; i < chunks.length; i++) {
|
|
try {
|
|
const embedding = await this._embed(chunks[i])
|
|
|
|
db.insert(ragChunks).values({
|
|
id: crypto.randomUUID(),
|
|
documentId: docId,
|
|
content: chunks[i],
|
|
embedding: JSON.stringify(embedding),
|
|
chunkIndex: i,
|
|
}).run()
|
|
|
|
successCount++
|
|
} catch (err) {
|
|
logger.warn(`RAG embedding failed for chunk ${i}/${chunks.length} of ${fileName}:`, err)
|
|
// 개별 청크 실패는 건너뛰고 계속 진행
|
|
}
|
|
|
|
const progress: RAGIndexProgress = {
|
|
documentId: docId,
|
|
fileName,
|
|
currentChunk: i + 1,
|
|
totalChunks: chunks.length,
|
|
percent: Math.round(((i + 1) / chunks.length) * 100),
|
|
}
|
|
this._sendToRenderer(IPC_CHANNELS.RAG.INDEX_PROGRESS, progress)
|
|
|
|
// 이벤트 루프 양보 (UI 블로킹 방지)
|
|
await new Promise((r) => setTimeout(r, 10))
|
|
}
|
|
|
|
// 인덱싱 완료 표시
|
|
db.update(ragDocuments)
|
|
.set({ indexed: true, indexedAt: Date.now(), chunkCount: successCount })
|
|
.where(eq(ragDocuments.id, docId))
|
|
.run()
|
|
|
|
this._state = 'idle'
|
|
this._sendToRenderer(IPC_CHANNELS.RAG.INDEX_COMPLETE, { documentId: docId, fileName })
|
|
logger.info(`RAG indexing complete: ${fileName} (${successCount}/${chunks.length} chunks embedded)`)
|
|
}
|
|
|
|
/**
|
|
* Ollama /api/embed 엔드포인트로 텍스트 임베딩
|
|
*/
|
|
private async _embed(text: string): Promise<number[]> {
|
|
const serverUrl = configGet('ollamaServerUrl')
|
|
|
|
try {
|
|
const response = await fetch(`${serverUrl}/api/embed`, {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
model: EMBED_MODEL,
|
|
input: text,
|
|
}),
|
|
signal: AbortSignal.timeout(30000),
|
|
})
|
|
|
|
if (!response.ok) {
|
|
throw new D3ROError(ErrorCode.RAGEmbeddingFailed, `Embed API error: ${response.status}`)
|
|
}
|
|
|
|
const data = (await response.json()) as { embeddings: number[][] }
|
|
if (!data.embeddings || data.embeddings.length === 0) {
|
|
throw new D3ROError(ErrorCode.RAGEmbeddingFailed, 'No embeddings returned')
|
|
}
|
|
|
|
return data.embeddings[0]
|
|
} catch (err) {
|
|
if (err instanceof D3ROError) throw err
|
|
throw new D3ROError(
|
|
ErrorCode.RAGEmbeddingFailed,
|
|
`Embedding failed: ${err instanceof Error ? err.message : String(err)}`,
|
|
)
|
|
}
|
|
}
|
|
|
|
private _chunkText(text: string): string[] {
|
|
// 텍스트 크기 제한 (500KB 초과 시 잘라냄)
|
|
const MAX_TEXT_LENGTH = 500000
|
|
const safeText = text.length > MAX_TEXT_LENGTH ? text.slice(0, MAX_TEXT_LENGTH) : text
|
|
|
|
const chunks: string[] = []
|
|
let start = 0
|
|
const MAX_CHUNKS = 2000
|
|
while (start < safeText.length && chunks.length < MAX_CHUNKS) {
|
|
const end = Math.min(start + CHUNK_SIZE, safeText.length)
|
|
const chunk = safeText.slice(start, end).trim()
|
|
if (chunk.length > 10) {
|
|
chunks.push(chunk)
|
|
}
|
|
start = end - CHUNK_OVERLAP
|
|
if (start >= safeText.length) break
|
|
}
|
|
return chunks
|
|
}
|
|
|
|
private async _extractText(filePath: string, fileType: string): Promise<string> {
|
|
const { promises: fsp } = await import('fs')
|
|
|
|
if (fileType === 'txt' || fileType === 'md') {
|
|
return fsp.readFile(filePath, 'utf-8')
|
|
}
|
|
|
|
if (fileType === 'pdf') {
|
|
try {
|
|
const buffer = await fsp.readFile(filePath)
|
|
const text = this._extractPdfText(buffer)
|
|
if (text.trim().length < 10) {
|
|
throw new Error('No readable text found in PDF')
|
|
}
|
|
return text.trim()
|
|
} catch (err) {
|
|
logger.error('PDF parsing failed:', err)
|
|
throw new D3ROError(ErrorCode.RAGUnsupportedFormat, 'PDF parsing failed. Ensure the PDF contains readable text.')
|
|
}
|
|
}
|
|
|
|
if (fileType === 'docx') {
|
|
try {
|
|
const buffer = await fsp.readFile(filePath)
|
|
// DOCX는 ZIP 내 XML — w:t 태그에서 텍스트 추출
|
|
const raw = buffer.toString('utf-8')
|
|
const matches = raw.match(/<w:t[^>]*>([^<]*)<\/w:t>/g)
|
|
if (matches) {
|
|
const text = matches
|
|
.map((m) => m.replace(/<[^>]+>/g, ''))
|
|
.join(' ')
|
|
return text.trim()
|
|
}
|
|
// XML 태그 제거 폴백
|
|
return raw.replace(/<[^>]+>/g, ' ').replace(/\s+/g, ' ').trim().slice(0, 500000)
|
|
} catch {
|
|
throw new D3ROError(ErrorCode.RAGUnsupportedFormat, 'DOCX parsing failed')
|
|
}
|
|
}
|
|
|
|
throw new D3ROError(ErrorCode.RAGUnsupportedFormat, `Unsupported: ${fileType}`)
|
|
}
|
|
|
|
/**
|
|
* PDF 바이너리에서 텍스트 추출 (순수 JS + zlib).
|
|
* FlateDecode 압축 해제 후 BT...ET 블록 내 Tj/TJ 파싱.
|
|
*/
|
|
private _extractPdfText(buffer: Buffer): string {
|
|
const zlib = require('zlib') as typeof import('zlib')
|
|
const raw = buffer.toString('binary')
|
|
const textParts: string[] = []
|
|
|
|
// 스트림 블록 추출 — 바이너리 오프셋 기반
|
|
const streamMarker = 'stream\r\n'
|
|
const streamMarker2 = 'stream\n'
|
|
const endMarker = 'endstream'
|
|
|
|
let pos = 0
|
|
while (pos < raw.length) {
|
|
let streamStart = raw.indexOf(streamMarker, pos)
|
|
let offset = streamMarker.length
|
|
if (streamStart === -1) {
|
|
streamStart = raw.indexOf(streamMarker2, pos)
|
|
offset = streamMarker2.length
|
|
}
|
|
if (streamStart === -1) break
|
|
|
|
const dataStart = streamStart + offset
|
|
const streamEnd = raw.indexOf(endMarker, dataStart)
|
|
if (streamEnd === -1) break
|
|
|
|
const streamData = Buffer.from(raw.slice(dataStart, streamEnd), 'binary')
|
|
pos = streamEnd + endMarker.length
|
|
|
|
// FlateDecode 해제 시도
|
|
let decoded: string
|
|
try {
|
|
const inflated = zlib.inflateSync(streamData)
|
|
decoded = inflated.toString('binary')
|
|
} catch {
|
|
// 압축 안 된 스트림
|
|
decoded = streamData.toString('binary')
|
|
}
|
|
|
|
// BT...ET 블록 파싱
|
|
this._extractTextFromStream(decoded, textParts)
|
|
}
|
|
|
|
// PDF 이스케이프 디코딩 + 정리
|
|
let text = textParts.join(' ')
|
|
text = text
|
|
.replace(/\\n/g, '\n')
|
|
.replace(/\\r/g, '\r')
|
|
.replace(/\\t/g, '\t')
|
|
.replace(/\\\(/g, '(')
|
|
.replace(/\\\)/g, ')')
|
|
.replace(/\\\\/g, '\\')
|
|
.replace(/[^\x20-\x7E\u00A0-\u00FF\u3000-\u9FFF\uAC00-\uD7AF\n]/g, ' ')
|
|
.replace(/\s+/g, ' ')
|
|
|
|
return text.trim()
|
|
}
|
|
|
|
private _extractTextFromStream(content: string, textParts: string[]): void {
|
|
const btRegex = /BT([\s\S]*?)ET/g
|
|
let btMatch: RegExpExecArray | null = null
|
|
while ((btMatch = btRegex.exec(content)) !== null) {
|
|
const block = btMatch[1]
|
|
|
|
// Tj: (text) Tj
|
|
const tjRegex = /\(([^)]*)\)\s*Tj/g
|
|
let tjMatch: RegExpExecArray | null = null
|
|
while ((tjMatch = tjRegex.exec(block)) !== null) {
|
|
if (tjMatch[1].trim()) textParts.push(tjMatch[1])
|
|
}
|
|
|
|
// TJ: [(text) num (text)] TJ
|
|
const tjArrayRegex = /\[(.*?)\]\s*TJ/g
|
|
let tjArrayMatch: RegExpExecArray | null = null
|
|
while ((tjArrayMatch = tjArrayRegex.exec(block)) !== null) {
|
|
const items = tjArrayMatch[1]
|
|
const itemRegex = /\(([^)]*)\)/g
|
|
let itemMatch: RegExpExecArray | null = null
|
|
while ((itemMatch = itemRegex.exec(items)) !== null) {
|
|
if (itemMatch[1].trim()) textParts.push(itemMatch[1])
|
|
}
|
|
}
|
|
|
|
// ' 연산자: (text) '
|
|
const quoteRegex = /\(([^)]*)\)\s*'/g
|
|
let quoteMatch: RegExpExecArray | null = null
|
|
while ((quoteMatch = quoteRegex.exec(block)) !== null) {
|
|
if (quoteMatch[1].trim()) textParts.push(quoteMatch[1])
|
|
}
|
|
}
|
|
}
|
|
|
|
private _cosineSimilarity(a: number[], b: number[]): number {
|
|
if (a.length !== b.length) return 0
|
|
let dotProduct = 0
|
|
let normA = 0
|
|
let normB = 0
|
|
for (let i = 0; i < a.length; i++) {
|
|
dotProduct += a[i] * b[i]
|
|
normA += a[i] * a[i]
|
|
normB += b[i] * b[i]
|
|
}
|
|
const denom = Math.sqrt(normA) * Math.sqrt(normB)
|
|
return denom === 0 ? 0 : dotProduct / denom
|
|
}
|
|
|
|
private _sendToRenderer(channel: string, data: unknown): void {
|
|
try {
|
|
const mainWindow = getMainWindow()
|
|
if (mainWindow && !mainWindow.isDestroyed()) {
|
|
mainWindow.webContents.send(channel, data)
|
|
}
|
|
} catch {
|
|
// ignore
|
|
}
|
|
}
|
|
|
|
dispose(): void {
|
|
this.removeAllListeners()
|
|
}
|
|
}
|
|
|
|
// ── 싱글톤 ──
|
|
let instance: RAGService | null = null
|
|
|
|
export function resetRAGServiceForTests(): void {
|
|
if (instance) instance.removeAllListeners()
|
|
instance = null
|
|
}
|
|
|
|
export function getRAGService(): RAGService {
|
|
if (!instance) {
|
|
instance = new RAGService()
|
|
}
|
|
return instance
|
|
}
|