Spaces:
Sleeping
Sleeping
File size: 3,289 Bytes
9f0bed6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | /**
* Persistent vector store backed by a JSON file.
* - In-memory cache: file is read once at startup; all ops hit memory first.
* - Atomic saves: write to .tmp then rename, so a crash never corrupts the store.
* - Cosine similarity for nearest-neighbour search.
*/
import fs from 'fs'
import path from 'path'
import { fileURLToPath } from 'url'
const __dirname = path.dirname(fileURLToPath(import.meta.url))
const DATA_DIR = process.env.DATA_DIR || path.join(__dirname, '..', 'vector_db')
const DB_PATH = path.join(DATA_DIR, 'store.json')
// In-memory store — null means "not loaded yet"
let _store = null
function getStore() {
if (_store === null) {
if (!fs.existsSync(DB_PATH)) {
_store = []
} else {
try {
_store = JSON.parse(fs.readFileSync(DB_PATH, 'utf8'))
} catch (err) {
console.error('[VectorStore] store.json is corrupted — reinitialising empty store.', err.message)
_store = []
}
}
}
return _store
}
function persist() {
// Atomic write: tmp file → rename
fs.mkdirSync(DATA_DIR, { recursive: true })
const tmp = DB_PATH + '.tmp'
fs.writeFileSync(tmp, JSON.stringify(_store), 'utf8')
fs.renameSync(tmp, DB_PATH)
}
// ── Cosine similarity ────────────────────────────────────────────────────────
function cosineSimilarity(a, b) {
if (!Array.isArray(a) || !Array.isArray(b) || a.length !== b.length || a.length === 0) return 0
let dot = 0, na = 0, nb = 0
for (let i = 0; i < a.length; i++) {
dot += a[i] * b[i]
na += a[i] * a[i]
nb += b[i] * b[i]
}
return na === 0 || nb === 0 ? 0 : dot / (Math.sqrt(na) * Math.sqrt(nb))
}
// ── Public API ───────────────────────────────────────────────────────────────
export function storeChunks(chunks, userId) {
const store = getStore()
store.push(...chunks.map(c => ({ ...c, userId })))
persist()
}
export function searchSimilar(queryEmbedding, userId, topK = 5) {
const store = getStore().filter(c => c.userId === userId)
if (!store.length) return []
const scored = store.map((c) => ({
text: c.text,
documentName: c.documentName,
pageNumber: c.pageNumber,
score: cosineSimilarity(queryEmbedding, c.embedding),
}))
scored.sort((a, b) => b.score - a.score)
return scored.slice(0, topK)
}
export function deleteDocumentChunks(documentName, userId) {
_store = getStore().filter(c => !(c.documentName === documentName && c.userId === userId))
persist()
}
export function listDocuments(userId) {
const store = getStore().filter(c => c.userId === userId)
return [...new Set(store.map(c => c.documentName))]
}
export function getDocumentStats(documentName, userId) {
const count = getStore().filter(c => c.documentName === documentName && c.userId === userId).length
return { chunks: count }
}
export function clearUserData(userId) {
_store = getStore().filter(c => c.userId !== userId)
persist()
}
export function getUserChunkCount(userId) {
return getStore().filter(c => c.userId === userId).length
}
|