| 1 | using System.Buffers.Binary; |
| 2 | using HybridCodebaseIndex.Core.Embeddings; |
| 3 | using Microsoft.Data.Sqlite; |
| 4 | |
| 5 | namespace HybridCodebaseIndex.Core; |
| 6 | |
| 7 | internal static partial class SqliteFtsIndex |
| 8 | { |
| 9 | private static void EnsureVectorsTable(SqliteConnection conn) |
| 10 | { |
| 11 | Exec(conn, """ |
| 12 | CREATE TABLE IF NOT EXISTS vectors( |
| 13 | chunk_rowid INTEGER PRIMARY KEY, |
| 14 | dim INTEGER NOT NULL, |
| 15 | norm REAL NOT NULL, |
| 16 | vec BLOB NOT NULL, |
| 17 | updated_at_utc_ticks INTEGER NOT NULL |
| 18 | ); |
| 19 | """); |
| 20 | } |
| 21 | |
| 22 | internal static Task<(int vectorsUpserted, string? err)> ReindexVectorsAsync( |
| 23 | string workspaceRoot, |
| 24 | string dbPath, |
| 25 | CancellationToken cancellationToken) |
| 26 | => Task.Run(() => ReindexVectors(workspaceRoot, dbPath, cancellationToken), cancellationToken); |
| 27 | |
| 28 | private static (int vectorsUpserted, string? err) ReindexVectors( |
| 29 | string workspaceRoot, |
| 30 | string dbPath, |
| 31 | CancellationToken cancellationToken) |
| 32 | { |
| 33 | workspaceRoot = Path.GetFullPath(workspaceRoot.TrimEnd(Path.DirectorySeparatorChar)); |
| 34 | if (!File.Exists(dbPath)) |
| 35 | return (0, "Index database not found; run codebase_index_reindex."); |
| 36 | |
| 37 | using var conn = new SqliteConnection($"Data Source={dbPath};Mode=ReadWrite"); |
| 38 | conn.Open(); |
| 39 | |
| 40 | EnsureMetaTable(conn); |
| 41 | EnsureChunksTable(conn); |
| 42 | EnsureFileStateTable(conn); |
| 43 | EnsureVectorsTable(conn); |
| 44 | |
| 45 | var settings = IndexSettings.TryLoadFromIndexDirectory(Path.GetDirectoryName(dbPath)!); |
| 46 | if (!settings.SemanticEnabled) |
| 47 | return (0, "semantic_enabled=false in settings; enable it to build vectors."); |
| 48 | |
| 49 | var allowedExt = settings.GetEffectiveVecExtensionsSet(); |
| 50 | |
| 51 | var provider = EmbeddingProviderFactory.Create(settings, Path.GetDirectoryName(dbPath)); |
| 52 | var dim = provider.Dimension; |
| 53 | |
| 54 | var indexDir = Path.GetDirectoryName(dbPath)!; |
| 55 | var sqliteVecAvailable = SqliteVecInterop.TryEnableAndLoad(conn, settings, indexDir, out _) |
| 56 | && SqliteVecInterop.TryEnsureVecChunksTable(conn, dim, out _); |
| 57 | |
| 58 | // Vec может наследовать FTS effective extensions, но допускает override через semantic.vec_*. |
| 59 | PruneVectorsAndVecChunksOutsideAllowedExtensions(conn, allowedExt, sqliteVecAvailable); |
| 60 | |
| 61 | if (allowedExt.Count == 0) |
| 62 | { |
| 63 | UpsertMeta(conn, "vec_dim", dim.ToString(System.Globalization.CultureInfo.InvariantCulture)); |
| 64 | UpsertMeta(conn, "vec_indexed_at", DateTime.UtcNow.ToString("O", System.Globalization.CultureInfo.InvariantCulture)); |
| 65 | return (0, null); |
| 66 | } |
| 67 | |
| 68 | using var select = conn.CreateCommand(); |
| 69 | select.CommandText = "SELECT rowid, substr(body, 1, 4000), extension FROM chunks;"; |
| 70 | |
| 71 | using var upsert = conn.CreateCommand(); |
| 72 | upsert.CommandText = """ |
| 73 | INSERT INTO vectors(chunk_rowid, dim, norm, vec, updated_at_utc_ticks) |
| 74 | VALUES($id, $dim, $norm, $vec, $t) |
| 75 | ON CONFLICT(chunk_rowid) DO UPDATE SET |
| 76 | dim=excluded.dim, |
| 77 | norm=excluded.norm, |
| 78 | vec=excluded.vec, |
| 79 | updated_at_utc_ticks=excluded.updated_at_utc_ticks; |
| 80 | """; |
| 81 | |
| 82 | var nowTicks = DateTime.UtcNow.Ticks; |
| 83 | var count = 0; |
| 84 | |
| 85 | using var r = select.ExecuteReader(); |
| 86 | while (r.Read()) |
| 87 | { |
| 88 | cancellationToken.ThrowIfCancellationRequested(); |
| 89 | |
| 90 | var id = r.GetInt64(0); |
| 91 | var text = r.IsDBNull(1) ? "" : r.GetString(1); |
| 92 | var extStored = r.IsDBNull(2) ? "" : r.GetString(2); |
| 93 | if (!allowedExt.Contains(extStored)) |
| 94 | continue; |
| 95 | |
| 96 | var emb = provider.EmbedAsync(text, cancellationToken).GetAwaiter().GetResult(); |
| 97 | if (emb.Length != dim) |
| 98 | continue; |
| 99 | |
| 100 | var (blob, norm) = SerializeVector(emb); |
| 101 | |
| 102 | upsert.Parameters.Clear(); |
| 103 | upsert.Parameters.AddWithValue("$id", id); |
| 104 | upsert.Parameters.AddWithValue("$dim", dim); |
| 105 | upsert.Parameters.AddWithValue("$norm", norm); |
| 106 | upsert.Parameters.AddWithValue("$vec", blob); |
| 107 | upsert.Parameters.AddWithValue("$t", nowTicks); |
| 108 | upsert.ExecuteNonQuery(); |
| 109 | |
| 110 | if (sqliteVecAvailable) |
| 111 | _ = SqliteVecInterop.TryUpsertVector(conn, id, emb, out _); |
| 112 | |
| 113 | count++; |
| 114 | } |
| 115 | |
| 116 | UpsertMeta(conn, "vec_dim", dim.ToString(System.Globalization.CultureInfo.InvariantCulture)); |
| 117 | UpsertMeta(conn, "vec_indexed_at", DateTime.UtcNow.ToString("O", System.Globalization.CultureInfo.InvariantCulture)); |
| 118 | return (count, null); |
| 119 | } |
| 120 | |
| 121 | private static void PruneVectorsAndVecChunksOutsideAllowedExtensions( |
| 122 | SqliteConnection conn, |
| 123 | HashSet<string> allowedExtensions, |
| 124 | bool sqliteVecActive) |
| 125 | { |
| 126 | if (allowedExtensions.Count == 0) |
| 127 | { |
| 128 | Exec(conn, "DELETE FROM vectors;"); |
| 129 | if (sqliteVecActive && TableExists(conn, "vec_chunks")) |
| 130 | Exec(conn, "DELETE FROM vec_chunks;"); |
| 131 | return; |
| 132 | } |
| 133 | |
| 134 | var pruneIds = new List<long>(); |
| 135 | using (var q = conn.CreateCommand()) |
| 136 | { |
| 137 | q.CommandText = "SELECT rowid, extension FROM chunks;"; |
| 138 | using var r = q.ExecuteReader(); |
| 139 | while (r.Read()) |
| 140 | { |
| 141 | var id = r.GetInt64(0); |
| 142 | var ext = r.IsDBNull(1) ? "" : r.GetString(1); |
| 143 | if (!allowedExtensions.Contains(ext)) |
| 144 | pruneIds.Add(id); |
| 145 | } |
| 146 | } |
| 147 | |
| 148 | DeleteVectorsAndVecChunksForRowIds(conn, pruneIds, sqliteVecActive); |
| 149 | |
| 150 | Exec(conn, """ |
| 151 | DELETE FROM vectors WHERE NOT EXISTS ( |
| 152 | SELECT 1 FROM chunks c WHERE c.rowid = vectors.chunk_rowid |
| 153 | ); |
| 154 | """); |
| 155 | |
| 156 | if (sqliteVecActive && TableExists(conn, "vec_chunks")) |
| 157 | { |
| 158 | Exec(conn, """ |
| 159 | DELETE FROM vec_chunks WHERE NOT EXISTS ( |
| 160 | SELECT 1 FROM chunks c WHERE c.rowid = vec_chunks.rowid |
| 161 | ); |
| 162 | """); |
| 163 | } |
| 164 | } |
| 165 | |
| 166 | private static void DeleteVectorsAndVecChunksForRowIds(SqliteConnection conn, List<long> rowIds, bool sqliteVecActive) |
| 167 | { |
| 168 | if (rowIds.Count == 0) |
| 169 | return; |
| 170 | |
| 171 | const int Batch = 400; |
| 172 | for (var offset = 0; offset < rowIds.Count; offset += Batch) |
| 173 | { |
| 174 | var slice = rowIds.Skip(offset).Take(Batch).ToArray(); |
| 175 | if (slice.Length == 0) |
| 176 | break; |
| 177 | |
| 178 | var placeholders = string.Join(",", slice.Select(static (_, i) => "$p" + i)); |
| 179 | using var dv = conn.CreateCommand(); |
| 180 | dv.CommandText = $"DELETE FROM vectors WHERE chunk_rowid IN ({placeholders});"; |
| 181 | for (var i = 0; i < slice.Length; i++) |
| 182 | dv.Parameters.AddWithValue($"$p{i}", slice[i]); |
| 183 | dv.ExecuteNonQuery(); |
| 184 | |
| 185 | if (sqliteVecActive && TableExists(conn, "vec_chunks")) |
| 186 | { |
| 187 | using var dvc = conn.CreateCommand(); |
| 188 | dvc.CommandText = $"DELETE FROM vec_chunks WHERE rowid IN ({placeholders});"; |
| 189 | for (var i = 0; i < slice.Length; i++) |
| 190 | dvc.Parameters.AddWithValue($"$p{i}", slice[i]); |
| 191 | dvc.ExecuteNonQuery(); |
| 192 | } |
| 193 | } |
| 194 | } |
| 195 | |
| 196 | private static (byte[] blob, double norm) SerializeVector(float[] v) |
| 197 | { |
| 198 | double sum = 0; |
| 199 | foreach (var x in v) |
| 200 | sum += x * x; |
| 201 | var norm = Math.Sqrt(sum); |
| 202 | |
| 203 | var blob = new byte[v.Length * 4]; |
| 204 | for (var i = 0; i < v.Length; i++) |
| 205 | BinaryPrimitives.WriteSingleLittleEndian(blob.AsSpan(i * 4, 4), v[i]); |
| 206 | |
| 207 | return (blob, norm); |
| 208 | } |
| 209 | } |
| 210 | |
| 211 | |