1<?php 2 3/** @noinspection SqlResolve */ 4 5namespace dokuwiki\plugin\aichat\Storage; 6 7use dokuwiki\plugin\aichat\AIChat; 8use dokuwiki\plugin\aichat\Chunk; 9use dokuwiki\plugin\sqlite\SQLiteDB; 10use KMeans\Cluster; 11use KMeans\Space; 12 13/** 14 * Implements the storage backend using a SQLite database 15 * 16 * Note: all embeddings are stored and returned as normalized vectors 17 */ 18class SQLiteStorage extends AbstractStorage 19{ 20 /** @var float minimum similarity to consider a chunk a match */ 21 final public const SIMILARITY_THRESHOLD = 0.75; 22 23 /** @var int Number of documents to randomly sample to create the clusters */ 24 final public const SAMPLE_SIZE = 2000; 25 /** @var int The average size of each cluster */ 26 final public const CLUSTER_SIZE = 400; 27 28 /** @var SQLiteDB */ 29 protected $db; 30 31 protected $useLanguageClusters = false; 32 33 /** 34 * Initializes the database connection and registers our custom function 35 * 36 * @throws \Exception 37 */ 38 public function __construct() 39 { 40 $this->db = new SQLiteDB('aichat', DOKU_PLUGIN . 'aichat/db/'); 41 $this->db->getPdo()->sqliteCreateFunction('COSIM', $this->sqliteCosineSimilarityCallback(...), 2); 42 43 $helper = plugin_load('helper', 'aichat'); 44 $this->useLanguageClusters = $helper->getConf('preferUIlanguage') >= AIChat::LANG_UI_LIMITED; 45 } 46 47 /** @inheritdoc */ 48 public function getChunk($chunkID) 49 { 50 $record = $this->db->queryRecord('SELECT * FROM embeddings WHERE id = ?', [$chunkID]); 51 if (!$record) return null; 52 53 return new Chunk( 54 $record['page'], 55 $record['id'], 56 $record['chunk'], 57 json_decode((string) $record['embedding'], true, 512, JSON_THROW_ON_ERROR), 58 $record['lang'], 59 $record['created'] 60 ); 61 } 62 63 /** @inheritdoc */ 64 public function startCreation($clear = false) 65 { 66 if ($clear) { 67 /** @noinspection SqlWithoutWhere */ 68 $this->db->exec('DELETE FROM embeddings'); 69 } 70 } 71 72 /** @inheritdoc */ 73 public function reusePageChunks($page, $firstChunkID) 74 { 75 // no-op 76 } 77 78 /** @inheritdoc */ 79 public function deletePageChunks($page, $firstChunkID) 80 { 81 $this->db->exec('DELETE FROM embeddings WHERE page = ?', [$page]); 82 } 83 84 /** @inheritdoc */ 85 public function addPageChunks($chunks) 86 { 87 foreach ($chunks as $chunk) { 88 $this->db->saveRecord('embeddings', [ 89 'page' => $chunk->getPage(), 90 'id' => $chunk->getId(), 91 'chunk' => $chunk->getText(), 92 'embedding' => json_encode($chunk->getEmbedding(), JSON_THROW_ON_ERROR), 93 'created' => $chunk->getCreated(), 94 'lang' => $chunk->getLanguage(), 95 ]); 96 } 97 } 98 99 /** @inheritdoc */ 100 public function finalizeCreation() 101 { 102 if (!$this->hasClusters()) { 103 $this->createClusters(); 104 } 105 $this->setChunkClusters(); 106 107 $this->db->exec('VACUUM'); 108 } 109 110 /** @inheritdoc */ 111 public function runMaintenance() 112 { 113 $this->createClusters(); 114 $this->setChunkClusters(); 115 } 116 117 /** @inheritdoc */ 118 public function getPageChunks($page, $firstChunkID) 119 { 120 $result = $this->db->queryAll( 121 'SELECT * FROM embeddings WHERE page = ?', 122 [$page] 123 ); 124 $chunks = []; 125 foreach ($result as $record) { 126 $chunks[] = new Chunk( 127 $record['page'], 128 $record['id'], 129 $record['chunk'], 130 json_decode((string) $record['embedding'], true, 512, JSON_THROW_ON_ERROR), 131 $record['lang'], 132 $record['created'] 133 ); 134 } 135 return $chunks; 136 } 137 138 /** @inheritdoc */ 139 public function getSimilarChunks($vector, $lang = '', $limit = 4) 140 { 141 $cluster = $this->getCluster($vector, $lang); 142 if ($this->logger) $this->logger->info( 143 'Using cluster {cluster} for similarity search', 144 ['cluster' => $cluster] 145 ); 146 147 $result = $this->db->queryAll( 148 'SELECT *, COSIM(?, embedding) AS similarity 149 FROM embeddings 150 WHERE cluster = ? 151 AND GETACCESSLEVEL(page) > 0 152 AND similarity > CAST(? AS FLOAT) 153 ORDER BY similarity DESC 154 LIMIT ?', 155 [json_encode($vector, JSON_THROW_ON_ERROR), $cluster, self::SIMILARITY_THRESHOLD, $limit] 156 ); 157 $chunks = []; 158 foreach ($result as $record) { 159 $chunks[] = new Chunk( 160 $record['page'], 161 $record['id'], 162 $record['chunk'], 163 json_decode((string) $record['embedding'], true, 512, JSON_THROW_ON_ERROR), 164 $record['lang'], 165 $record['created'], 166 $record['similarity'] 167 ); 168 } 169 return $chunks; 170 } 171 172 /** @inheritdoc */ 173 public function statistics() 174 { 175 $items = $this->db->queryValue('SELECT COUNT(*) FROM embeddings'); 176 $size = $this->db->queryValue( 177 'SELECT page_count * page_size as size FROM pragma_page_count(), pragma_page_size()' 178 ); 179 $query = "SELECT cluster || ' ' || lang, COUNT(*) || ' chunks' as cnt 180 FROM embeddings 181 GROUP BY cluster 182 ORDER BY cluster"; 183 $clusters = $this->db->queryKeyValueList($query); 184 185 return [ 186 'storage type' => 'SQLite', 187 'chunks' => $items, 188 'db size' => filesize_h($size), 189 'clusters' => $clusters, 190 ]; 191 } 192 193 /** 194 * Method registered as SQLite callback to calculate the cosine similarity 195 * 196 * @param string $query JSON encoded vector array 197 * @param string $embedding JSON encoded vector array 198 * @return float 199 */ 200 public function sqliteCosineSimilarityCallback($query, $embedding) 201 { 202 return (float)$this->cosineSimilarity(json_decode($query, null, 512, JSON_THROW_ON_ERROR), json_decode($embedding, null, 512, JSON_THROW_ON_ERROR)); 203 } 204 205 /** 206 * Calculate the cosine similarity between two vectors 207 * 208 * Actually just calculating the dot product of the two vectors, since they are normalized 209 * 210 * @param float[] $queryVector The normalized vector of the search phrase 211 * @param float[] $embedding The normalized vector of the chunk 212 * @return float 213 */ 214 protected function cosineSimilarity($queryVector, $embedding) 215 { 216 $dotProduct = 0; 217 foreach ($queryVector as $key => $value) { 218 $dotProduct += $value * $embedding[$key]; 219 } 220 return $dotProduct; 221 } 222 223 /** 224 * Create new clusters based on random chunks 225 * 226 * @return void 227 */ 228 protected function createClusters() 229 { 230 if ($this->useLanguageClusters) { 231 $result = $this->db->queryAll('SELECT DISTINCT lang FROM embeddings'); 232 $langs = array_column($result, 'lang'); 233 foreach ($langs as $lang) { 234 $this->createLanguageClusters($lang); 235 } 236 } else { 237 $this->createLanguageClusters(''); 238 } 239 } 240 241 /** 242 * Create new clusters based on random chunks for the given Language 243 * 244 * @param string $lang The language to cluster, empty when all languages go into the same cluster 245 * @noinspection SqlWithoutWhere 246 */ 247 protected function createLanguageClusters($lang) 248 { 249 if ($lang != '') { 250 $where = 'WHERE lang = ' . $this->db->getPdo()->quote($lang); 251 } else { 252 $where = ''; 253 } 254 255 if ($this->logger) $this->logger->info('Creating new {lang} clusters...', ['lang' => $lang]); 256 $this->db->getPdo()->beginTransaction(); 257 try { 258 // clean up old cluster data 259 $query = "DELETE FROM clusters $where"; 260 $this->db->exec($query); 261 $query = "UPDATE embeddings SET cluster = NULL $where"; 262 $this->db->exec($query); 263 264 // get a random selection of chunks 265 $query = "SELECT id, embedding FROM embeddings $where ORDER BY RANDOM() LIMIT ?"; 266 $result = $this->db->queryAll($query, [self::SAMPLE_SIZE]); 267 if (!$result) return; // no data to cluster 268 $dimensions = count(json_decode((string) $result[0]['embedding'], true, 512, JSON_THROW_ON_ERROR)); 269 270 // how many clusters? 271 if (count($result) < self::CLUSTER_SIZE * 3) { 272 // there would be less than 3 clusters, so just use one 273 $clustercount = 1; 274 } else { 275 // get the number of all chunks, to calculate the number of clusters 276 $query = "SELECT COUNT(*) FROM embeddings $where"; 277 $total = $this->db->queryValue($query); 278 $clustercount = ceil($total / self::CLUSTER_SIZE); 279 } 280 if ($this->logger) $this->logger->info('Creating {clusters} clusters', ['clusters' => $clustercount]); 281 282 // cluster them using kmeans 283 $space = new Space($dimensions); 284 foreach ($result as $record) { 285 $space->addPoint(json_decode((string) $record['embedding'], true, 512, JSON_THROW_ON_ERROR)); 286 } 287 $clusters = $space->solve($clustercount, function ($space, $clusters) { 288 static $iterations = 0; 289 ++$iterations; 290 if ($this->logger) { 291 $clustercounts = implode(',', array_map('count', $clusters)); 292 $this->logger->info('Iteration {iteration}: [{clusters}]', [ 293 'iteration' => $iterations, 'clusters' => $clustercounts 294 ]); 295 } 296 }, Cluster::INIT_KMEANS_PLUS_PLUS); 297 298 // store the clusters 299 foreach ($clusters as $cluster) { 300 /** @var Cluster $cluster */ 301 $centroid = $cluster->getCoordinates(); 302 $query = 'INSERT INTO clusters (lang, centroid) VALUES (?, ?)'; 303 $this->db->exec($query, [$lang, json_encode($centroid, JSON_THROW_ON_ERROR)]); 304 } 305 306 $this->db->getPdo()->commit(); 307 if ($this->logger) $this->logger->success('Created {clusters} clusters', ['clusters' => count($clusters)]); 308 } catch (\Exception $e) { 309 $this->db->getPdo()->rollBack(); 310 throw new \RuntimeException('Clustering failed: ' . $e->getMessage(), 0, $e); 311 } 312 } 313 314 /** 315 * Assign the nearest cluster for all chunks that don't have one 316 * 317 * @return void 318 */ 319 protected function setChunkClusters() 320 { 321 if ($this->logger) $this->logger->info('Assigning clusters to chunks...'); 322 $query = 'SELECT id, embedding, lang FROM embeddings WHERE cluster IS NULL'; 323 $handle = $this->db->query($query); 324 325 while ($record = $handle->fetch(\PDO::FETCH_ASSOC)) { 326 $vector = json_decode((string) $record['embedding'], true, 512, JSON_THROW_ON_ERROR); 327 $cluster = $this->getCluster($vector, $this->useLanguageClusters ? $record['lang'] : ''); 328 $query = 'UPDATE embeddings SET cluster = ? WHERE id = ?'; 329 $this->db->exec($query, [$cluster, $record['id']]); 330 if ($this->logger) $this->logger->success( 331 'Chunk {id} assigned to cluster {cluster}', 332 ['id' => $record['id'], 'cluster' => $cluster] 333 ); 334 } 335 $handle->closeCursor(); 336 } 337 338 /** 339 * Get the nearest cluster for the given vector 340 * 341 * @param float[] $vector 342 * @return int|null 343 */ 344 protected function getCluster($vector, $lang) 345 { 346 if ($lang != '') { 347 $where = 'WHERE lang = ' . $this->db->getPdo()->quote($lang); 348 } else { 349 $where = ''; 350 } 351 352 $query = "SELECT cluster, centroid 353 FROM clusters 354 $where 355 ORDER BY COSIM(centroid, ?) DESC 356 LIMIT 1"; 357 358 $result = $this->db->queryRecord($query, [json_encode($vector, JSON_THROW_ON_ERROR)]); 359 if (!$result) return null; 360 return $result['cluster']; 361 } 362 363 /** 364 * Check if clustering has been done before 365 * @return bool 366 */ 367 protected function hasClusters() 368 { 369 $query = 'SELECT COUNT(*) FROM clusters'; 370 return $this->db->queryValue($query) > 0; 371 } 372 373 /** 374 * Writes TSV files for visualizing with http://projector.tensorflow.org/ 375 * 376 * @param string $vectorfile path to the file with the vectors 377 * @param string $metafile path to the file with the metadata 378 * @return void 379 */ 380 public function dumpTSV($vectorfile, $metafile) 381 { 382 $query = 'SELECT * FROM embeddings'; 383 $handle = $this->db->query($query); 384 385 $header = implode("\t", ['id', 'page', 'created']); 386 file_put_contents($metafile, $header . "\n", FILE_APPEND); 387 388 while ($row = $handle->fetch(\PDO::FETCH_ASSOC)) { 389 $vector = json_decode((string) $row['embedding'], true, 512, JSON_THROW_ON_ERROR); 390 $vector = implode("\t", $vector); 391 392 $meta = implode("\t", [$row['id'], $row['page'], $row['created']]); 393 394 file_put_contents($vectorfile, $vector . "\n", FILE_APPEND); 395 file_put_contents($metafile, $meta . "\n", FILE_APPEND); 396 } 397 } 398} 399