xref: /plugin/aichat/Storage/SQLiteStorage.php (revision 7ebc78955c65af90e7ee0afbd07adc15271113ba) 
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    public const SIMILARITY_THRESHOLD = 0.75;
22
23    /** @var int Number of documents to randomly sample to create the clusters */
24    public const SAMPLE_SIZE = 2000;
25    /** @var int The average size of each cluster */
26    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($record['embedding'], true),
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()),
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($record['embedding'], true),
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), $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($record['embedding'], true),
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 FROM embeddings GROUP BY cluster ORDER BY cluster";
180        $clusters = $this->db->queryKeyValueList($query);
181
182        return [
183            'storage type' => 'SQLite',
184            'chunks' => $items,
185            'db size' => filesize_h($size),
186            'clusters' => $clusters,
187        ];
188    }
189
190    /**
191     * Method registered as SQLite callback to calculate the cosine similarity
192     *
193     * @param string $query JSON encoded vector array
194     * @param string $embedding JSON encoded vector array
195     * @return float
196     */
197    public function sqliteCosineSimilarityCallback($query, $embedding)
198    {
199        return (float)$this->cosineSimilarity(json_decode($query), json_decode($embedding));
200    }
201
202    /**
203     * Calculate the cosine similarity between two vectors
204     *
205     * Actually just calculating the dot product of the two vectors, since they are normalized
206     *
207     * @param float[] $queryVector The normalized vector of the search phrase
208     * @param float[] $embedding The normalized vector of the chunk
209     * @return float
210     */
211    protected function cosineSimilarity($queryVector, $embedding)
212    {
213        $dotProduct = 0;
214        foreach ($queryVector as $key => $value) {
215            $dotProduct += $value * $embedding[$key];
216        }
217        return $dotProduct;
218    }
219
220    /**
221     * Create new clusters based on random chunks
222     *
223     * @return void
224     */
225    protected function createClusters()
226    {
227        if ($this->useLanguageClusters) {
228            $result = $this->db->queryAll('SELECT DISTINCT lang FROM embeddings');
229            $langs = array_column($result, 'lang');
230            foreach ($langs as $lang) {
231                $this->createLanguageClusters($lang);
232            }
233        } else {
234            $this->createLanguageClusters('');
235        }
236    }
237
238    /**
239     * Create new clusters based on random chunks for the given Language
240     *
241     * @param string $lang The language to cluster, empty when all languages go into the same cluster
242     * @noinspection SqlWithoutWhere
243     */
244    protected function createLanguageClusters($lang)
245    {
246        if ($lang != '') {
247            $where = 'WHERE lang = ' . $this->db->getPdo()->quote($lang);
248        } else {
249            $where = '';
250        }
251
252        if ($this->logger) $this->logger->info('Creating new {lang} clusters...', ['lang' => $lang]);
253        $this->db->getPdo()->beginTransaction();
254        try {
255            // clean up old cluster data
256            $query = "DELETE FROM clusters $where";
257            $this->db->exec($query);
258            $query = "UPDATE embeddings SET cluster = NULL $where";
259            $this->db->exec($query);
260
261            // get a random selection of chunks
262            $query = "SELECT id, embedding FROM embeddings $where ORDER BY RANDOM() LIMIT ?";
263            $result = $this->db->queryAll($query, [self::SAMPLE_SIZE]);
264            if (!$result) return; // no data to cluster
265            $dimensions = count(json_decode($result[0]['embedding'], true));
266
267            // how many clusters?
268            if (count($result) < self::CLUSTER_SIZE * 3) {
269                // there would be less than 3 clusters, so just use one
270                $clustercount = 1;
271            } else {
272                // get the number of all chunks, to calculate the number of clusters
273                $query = "SELECT COUNT(*) FROM embeddings $where";
274                $total = $this->db->queryValue($query);
275                $clustercount = ceil($total / self::CLUSTER_SIZE);
276            }
277            if ($this->logger) $this->logger->info('Creating {clusters} clusters', ['clusters' => $clustercount]);
278
279            // cluster them using kmeans
280            $space = new Space($dimensions);
281            foreach ($result as $record) {
282                $space->addPoint(json_decode($record['embedding'], true));
283            }
284            $clusters = $space->solve($clustercount, function ($space, $clusters) {
285                static $iterations = 0;
286                ++$iterations;
287                if ($this->logger) {
288                    $clustercounts = implode(',', array_map('count', $clusters));
289                    $this->logger->info('Iteration {iteration}: [{clusters}]', [
290                        'iteration' => $iterations, 'clusters' => $clustercounts
291                    ]);
292                }
293            }, Cluster::INIT_KMEANS_PLUS_PLUS);
294
295            // store the clusters
296            foreach ($clusters as $cluster) {
297                /** @var Cluster $cluster */
298                $centroid = $cluster->getCoordinates();
299                $query = 'INSERT INTO clusters (lang, centroid) VALUES (?, ?)';
300                $this->db->exec($query, [$lang, json_encode($centroid)]);
301            }
302
303            $this->db->getPdo()->commit();
304            if ($this->logger) $this->logger->success('Created {clusters} clusters', ['clusters' => count($clusters)]);
305        } catch (\Exception $e) {
306            $this->db->getPdo()->rollBack();
307            throw new \RuntimeException('Clustering failed: ' . $e->getMessage(), 0, $e);
308        }
309    }
310
311    /**
312     * Assign the nearest cluster for all chunks that don't have one
313     *
314     * @return void
315     */
316    protected function setChunkClusters()
317    {
318        if ($this->logger) $this->logger->info('Assigning clusters to chunks...');
319        $query = 'SELECT id, embedding, lang FROM embeddings WHERE cluster IS NULL';
320        $handle = $this->db->query($query);
321
322        while ($record = $handle->fetch(\PDO::FETCH_ASSOC)) {
323            $vector = json_decode($record['embedding'], true);
324            $cluster = $this->getCluster($vector, $this->useLanguageClusters ? $record['lang'] : '');
325            $query = 'UPDATE embeddings SET cluster = ? WHERE id = ?';
326            $this->db->exec($query, [$cluster, $record['id']]);
327            if ($this->logger) $this->logger->success(
328                'Chunk {id} assigned to cluster {cluster}',
329                ['id' => $record['id'], 'cluster' => $cluster]
330            );
331        }
332        $handle->closeCursor();
333    }
334
335    /**
336     * Get the nearest cluster for the given vector
337     *
338     * @param float[] $vector
339     * @return int|null
340     */
341    protected function getCluster($vector, $lang)
342    {
343        if ($lang != '') {
344            $where = 'WHERE lang = ' . $this->db->getPdo()->quote($lang);
345        } else {
346            $where = '';
347        }
348
349        $query = "SELECT cluster, centroid
350                    FROM clusters
351                   $where
352                ORDER BY COSIM(centroid, ?) DESC
353                   LIMIT 1";
354
355        $result = $this->db->queryRecord($query, [json_encode($vector)]);
356        if (!$result) return null;
357        return $result['cluster'];
358    }
359
360    /**
361     * Check if clustering has been done before
362     * @return bool
363     */
364    protected function hasClusters()
365    {
366        $query = 'SELECT COUNT(*) FROM clusters';
367        return $this->db->queryValue($query) > 0;
368    }
369
370    /**
371     * Writes TSV files for visualizing with http://projector.tensorflow.org/
372     *
373     * @param string $vectorfile path to the file with the vectors
374     * @param string $metafile path to the file with the metadata
375     * @return void
376     */
377    public function dumpTSV($vectorfile, $metafile)
378    {
379        $query = 'SELECT * FROM embeddings';
380        $handle = $this->db->query($query);
381
382        $header = implode("\t", ['id', 'page', 'created']);
383        file_put_contents($metafile, $header . "\n", FILE_APPEND);
384
385        while ($row = $handle->fetch(\PDO::FETCH_ASSOC)) {
386            $vector = json_decode($row['embedding'], true);
387            $vector = implode("\t", $vector);
388
389            $meta = implode("\t", [$row['id'], $row['page'], $row['created']]);
390
391            file_put_contents($vectorfile, $vector . "\n", FILE_APPEND);
392            file_put_contents($metafile, $meta . "\n", FILE_APPEND);
393        }
394    }
395}
396