-
Notifications
You must be signed in to change notification settings - Fork 10
Expand file tree
/
Copy pathsqliteVecClient.ts
More file actions
374 lines (330 loc) · 9.66 KB
/
Copy pathsqliteVecClient.ts
File metadata and controls
374 lines (330 loc) · 9.66 KB
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
import Database from "better-sqlite3"
import * as sqlite_vec from "sqlite-vec"
import { createHash } from "crypto"
import path from "path"
import fs from "fs"
import type { VectorStore, VectorStoreSearchResult } from "./interface.js"
const DEFAULT_MAX_SEARCH_RESULTS = 50
const DEFAULT_SEARCH_MIN_SCORE = 0.4
/**
* SQLite-vec based vector store for local vector search.
* Stores vectors in a local SQLite database using the sqlite-vec extension.
*
* Advantages:
* - No external services required
* - Portable (single .db file)
* - Can be committed with the project
*
* Best for:
* - Small to medium projects (<100k chunks)
* - Local development
* - Offline usage
*/
export class SqliteVecClient implements VectorStore {
private db: Database.Database
private vectorSize: number
private readonly dbPath: string
private readonly workspacePath: string
private readonly tableName = "code_vectors"
constructor(
workspacePath: string,
vectorSize: number,
dbPath?: string,
) {
this.workspacePath = workspacePath
this.vectorSize = vectorSize
// Default to .codebase/vectors.db
if (dbPath) {
this.dbPath = dbPath
} else {
const codebaseDir = path.join(workspacePath, ".codebase")
if (!fs.existsSync(codebaseDir)) {
fs.mkdirSync(codebaseDir, { recursive: true })
}
this.dbPath = path.join(codebaseDir, "vectors.db")
}
// Initialize database
this.db = new Database(this.dbPath)
this.db.pragma("journal_mode = WAL")
// Load sqlite-vec extension
sqlite_vec.load(this.db)
}
/**
* Initialize the vector store.
* Creates the virtual table if it doesn't exist.
* Returns initialization result with created and cleanup flags.
* SQLite never does cleanup (no remote state to validate).
*/
async initialize(): Promise<{ created: boolean; didCleanup: boolean }> {
let created = false
try {
// Check if table exists
const tableExists = this.db
.prepare(
`SELECT name FROM sqlite_master WHERE type='table' AND name=?`
)
.get(this.tableName)
if (!tableExists) {
// Create virtual table using vec0
this.db.exec(`
CREATE VIRTUAL TABLE ${this.tableName} USING vec0(
id TEXT PRIMARY KEY,
embedding float[${this.vectorSize}],
file_path TEXT NOT NULL,
code_chunk TEXT NOT NULL,
start_line INTEGER NOT NULL,
end_line INTEGER NOT NULL,
segment_hash TEXT
)
`)
// Note: Cannot create indexes on virtual tables
// Filtering will be done in the WHERE clause
created = true
} else {
// Verify vector dimension matches
const tableInfo = this.db.prepare(`PRAGMA table_info(${this.tableName})`).all() as any[]
const embeddingColumn = tableInfo.find((col: any) => col.name === "embedding")
if (embeddingColumn) {
// Extract dimension from type like "float[4096]"
const match = embeddingColumn.type.match(/float\[(\d+)\]/)
if (match) {
const existingDimension = parseInt(match[1], 10)
if (existingDimension !== this.vectorSize) {
console.warn(
`[SqliteVecClient] Table ${this.tableName} has dimension ${existingDimension}, expected ${this.vectorSize}. Recreating.`
)
await this.resetCollection()
created = true
}
}
}
}
return { created, didCleanup: false }
} catch (error: any) {
console.error(`[SqliteVecClient] Failed to initialize database:`, error.message)
throw new Error(`Failed to initialize SQLite vector store: ${error.message}`)
}
}
/**
* Reset the collection by dropping and recreating the table.
*/
async resetCollection(): Promise<void> {
try {
this.db.exec(`DROP TABLE IF EXISTS ${this.tableName}`)
await this.initialize()
} catch (error: any) {
console.error(`[SqliteVecClient] Failed to reset collection:`, error.message)
throw error
}
}
/**
* Upsert points into the vector store.
* Uses INSERT OR REPLACE to handle both inserts and updates.
*/
async upsertPoints(
points: Array<{
id: string
vector: number[]
payload: Record<string, any>
}>,
): Promise<void> {
if (points.length === 0) {
return
}
try {
const stmt = this.db.prepare(`
INSERT OR REPLACE INTO ${this.tableName}
(id, embedding, file_path, code_chunk, start_line, end_line, segment_hash)
VALUES (?, ?, ?, ?, CAST(? AS INTEGER), CAST(? AS INTEGER), ?)
`)
const insertMany = this.db.transaction((points: any[]) => {
for (const point of points) {
// Convert vector to JSON string for storage
const vectorJson = JSON.stringify(point.vector)
stmt.run(
point.id,
vectorJson,
point.payload.filePath || "",
point.payload.codeChunk || "",
point.payload.startLine || 0,
point.payload.endLine || 0,
point.payload.segmentHash || null,
)
}
})
insertMany(points)
} catch (error: any) {
console.error("[SqliteVecClient] Failed to upsert points:", error.message)
throw error
}
}
/**
* Search for similar vectors using KNN.
*
* @param queryVector - The query vector to search for
* @param directoryPrefix - Optional directory prefix to filter results
* @param minScore - Minimum similarity score (0-1, higher is better)
* @param maxResults - Maximum number of results to return
*/
async search(
queryVector: number[],
directoryPrefix?: string,
minScore?: number,
maxResults?: number,
): Promise<VectorStoreSearchResult[]> {
try {
const limit = maxResults ?? DEFAULT_MAX_SEARCH_RESULTS
const scoreThreshold = minScore ?? DEFAULT_SEARCH_MIN_SCORE
// Convert query vector to JSON for sqlite-vec
const queryVectorJson = JSON.stringify(queryVector)
let query = `
SELECT
id,
file_path,
code_chunk,
start_line,
end_line,
segment_hash,
distance
FROM ${this.tableName}
WHERE embedding MATCH ?
`
const params: any[] = [queryVectorJson]
// Add directory prefix filter if provided
if (directoryPrefix && directoryPrefix !== "." && directoryPrefix !== "./") {
const normalizedPrefix = path.posix.normalize(directoryPrefix.replace(/\\/g, "/"))
const cleanedPrefix = normalizedPrefix.startsWith("./")
? normalizedPrefix.slice(2)
: normalizedPrefix
if (cleanedPrefix) {
query += ` AND file_path LIKE ?`
params.push(`${cleanedPrefix}%`)
}
}
query += ` ORDER BY distance LIMIT ?`
params.push(limit)
const stmt = this.db.prepare(query)
const rows = stmt.all(...params) as any[]
// Convert distance to similarity score (1 - distance for cosine)
// Filter by score threshold
const results: VectorStoreSearchResult[] = rows
.map((row) => {
// sqlite-vec returns distance, convert to similarity score
// For cosine distance: similarity = 1 - distance
const score = 1 - row.distance
return {
id: row.id,
score,
payload: {
filePath: row.file_path,
codeChunk: row.code_chunk,
startLine: row.start_line,
endLine: row.end_line,
segmentHash: row.segment_hash,
},
}
})
.filter((result) => result.score >= scoreThreshold)
return results
} catch (error: any) {
console.error("[SqliteVecClient] Failed to search:", error.message)
throw error
}
}
/**
* Delete points by file path.
*/
async deletePointsByFilePath(filePath: string): Promise<void> {
await this.deletePointsByMultipleFilePaths([filePath])
}
/**
* Delete points by multiple file paths.
*/
async deletePointsByMultipleFilePaths(filePaths: string[]): Promise<void> {
if (filePaths.length === 0) {
return
}
try {
const workspaceRoot = this.workspacePath
const normalizedPaths = filePaths.map((filePath) => {
const relativePath = path.isAbsolute(filePath)
? path.relative(workspaceRoot, filePath)
: filePath
return path.normalize(relativePath)
})
const placeholders = normalizedPaths.map(() => "?").join(",")
const stmt = this.db.prepare(`
DELETE FROM ${this.tableName}
WHERE file_path IN (${placeholders})
`)
stmt.run(...normalizedPaths)
} catch (error: any) {
console.error("[SqliteVecClient] Failed to delete points:", error.message)
}
}
/**
* Delete the entire collection (drop table).
*/
async deleteCollection(): Promise<void> {
try {
this.db.exec(`DROP TABLE IF EXISTS ${this.tableName}`)
} catch (error: any) {
console.error("[SqliteVecClient] Failed to delete collection:", error.message)
throw error
}
}
/**
* Clear all points from the collection.
*/
async clearCollection(): Promise<void> {
try {
this.db.exec(`DELETE FROM ${this.tableName}`)
} catch (error: any) {
console.error("[SqliteVecClient] Failed to clear collection:", error.message)
throw error
}
}
/**
* Check if the collection exists.
*/
async collectionExists(): Promise<boolean> {
const result = this.db
.prepare(
`SELECT name FROM sqlite_master WHERE type='table' AND name=?`
)
.get(this.tableName)
return result !== undefined
}
/**
* Ensure vector dimension matches.
* If not, recreate the collection.
*/
async ensureVectorDimension(actualDimension: number): Promise<void> {
if (actualDimension === this.vectorSize) {
return
}
console.warn(
`[SqliteVecClient] Adjusting vector dimension from ${this.vectorSize} to ${actualDimension}.`
)
this.vectorSize = actualDimension
await this.resetCollection()
}
/**
* Close the database connection.
*/
close(): void {
this.db.close()
}
}
export function generatePointId(segmentHash: string): string {
const hex = segmentHash.replace(/-/g, "").toLowerCase()
const base = hex.slice(0, 32).padEnd(32, "0")
const sections = [
base.slice(0, 8),
base.slice(8, 12),
base.slice(12, 16),
base.slice(16, 20),
base.slice(20, 32),
]
return sections.join("-")
}