xref: /plugin/aichat/Embeddings.php (revision 30b9cbc7090f3bc7eca85292f4ff98a1cf513f8f) 
1<?php
2
3namespace dokuwiki\plugin\aichat;
4
5use dokuwiki\Extension\PluginInterface;
6use dokuwiki\plugin\aichat\Model\AbstractModel;
7use dokuwiki\plugin\aichat\Storage\AbstractStorage;
8use dokuwiki\Search\Indexer;
9use splitbrain\phpcli\CLI;
10use TikToken\Encoder;
11use Vanderlee\Sentence\Sentence;
12
13/**
14 * Manage the embeddings index
15 *
16 * Pages are split into chunks of 1000 tokens each. For each chunk the embedding vector is fetched from
17 * OpenAI and stored in the Storage backend.
18 */
19class Embeddings
20{
21    /** @var int maximum overlap between chunks in tokens */
22    final public const MAX_OVERLAP_LEN = 200;
23
24    /** @var AbstractModel */
25    protected $model;
26    /** @var CLI|null */
27    protected $logger;
28    /** @var Encoder */
29    protected $tokenEncoder;
30
31    /** @var AbstractStorage */
32    protected $storage;
33
34    /** @var array remember sentences when chunking */
35    private $sentenceQueue = [];
36
37    public function __construct(AbstractModel $model, AbstractStorage $storage)
38    {
39        $this->model = $model;
40        $this->storage = $storage;
41    }
42
43    /**
44     * Access storage
45     *
46     * @return AbstractStorage
47     */
48    public function getStorage()
49    {
50        return $this->storage;
51    }
52
53    /**
54     * Add a logger instance
55     *
56     * @return void
57     */
58    public function setLogger(CLI $logger)
59    {
60        $this->logger = $logger;
61    }
62
63    /**
64     * Get the token encoder instance
65     *
66     * @return Encoder
67     */
68    public function getTokenEncoder()
69    {
70        if (!$this->tokenEncoder instanceof Encoder) {
71            $this->tokenEncoder = new Encoder();
72        }
73        return $this->tokenEncoder;
74    }
75
76    /**
77     * Update the embeddings storage
78     *
79     * @param string $skipRE Regular expression to filter out pages (full RE with delimiters)
80     * @param bool $clear Should any existing storage be cleared before updating?
81     * @return void
82     * @throws \Exception
83     */
84    public function createNewIndex($skipRE = '', $clear = false)
85    {
86        $indexer = new Indexer();
87        $pages = $indexer->getPages();
88
89        $this->storage->startCreation($clear);
90        foreach ($pages as $pid => $page) {
91            $chunkID = $pid * 100; // chunk IDs start at page ID * 100
92
93            if (
94                !page_exists($page) ||
95                isHiddenPage($page) ||
96                filesize(wikiFN($page)) < 150 || // skip very small pages
97                ($skipRE && preg_match($skipRE, (string) $page))
98            ) {
99                // this page should not be in the index (anymore)
100                $this->storage->deletePageChunks($page, $chunkID);
101                continue;
102            }
103
104            $firstChunk = $this->storage->getChunk($chunkID);
105            if ($firstChunk && @filemtime(wikiFN($page)) < $firstChunk->getCreated()) {
106                // page is older than the chunks we have, reuse the existing chunks
107                $this->storage->reusePageChunks($page, $chunkID);
108                if ($this->logger instanceof CLI) $this->logger->info("Reusing chunks for $page");
109            } else {
110                // page is newer than the chunks we have, create new chunks
111                $this->storage->deletePageChunks($page, $chunkID);
112                $this->storage->addPageChunks($this->createPageChunks($page, $chunkID));
113            }
114        }
115        $this->storage->finalizeCreation();
116    }
117
118    /**
119     * Split the given page, fetch embedding vectors and return Chunks
120     *
121     * Will use the text renderer plugin if available to get the rendered text.
122     * Otherwise the raw wiki text is used.
123     *
124     * @param string $page Name of the page to split
125     * @param int $firstChunkID The ID of the first chunk of this page
126     * @return Chunk[] A list of chunks created for this page
127     * @throws \Exception
128     */
129    protected function createPageChunks($page, $firstChunkID)
130    {
131        $chunkList = [];
132
133        $textRenderer = plugin_load('renderer', 'text');
134        if ($textRenderer instanceof PluginInterface) {
135            global $ID;
136            $ID = $page;
137            $text = p_cached_output(wikiFN($page), 'text', $page);
138        } else {
139            $text = rawWiki($page);
140        }
141
142        $parts = $this->splitIntoChunks($text);
143        foreach ($parts as $part) {
144            if (trim((string) $part) == '') continue; // skip empty chunks
145
146            try {
147                $embedding = $this->model->getEmbedding($part);
148            } catch (\Exception $e) {
149                if ($this->logger instanceof CLI) {
150                    $this->logger->error(
151                        'Failed to get embedding for chunk of page {page}: {msg}',
152                        ['page' => $page, 'msg' => $e->getMessage()]
153                    );
154                }
155                continue;
156            }
157            $chunkList[] = new Chunk($page, $firstChunkID, $part, $embedding);
158            $firstChunkID++;
159        }
160        if ($this->logger instanceof CLI) {
161            if ($chunkList !== []) {
162                $this->logger->success(
163                    '{id} split into {count} chunks',
164                    ['id' => $page, 'count' => count($chunkList)]
165                );
166            } else {
167                $this->logger->warning('{id} could not be split into chunks', ['id' => $page]);
168            }
169        }
170        return $chunkList;
171    }
172
173    /**
174     * Do a nearest neighbor search for chunks similar to the given question
175     *
176     * Returns only chunks the current user is allowed to read, may return an empty result.
177     * The number of returned chunks depends on the MAX_CONTEXT_LEN setting.
178     *
179     * @param string $query The question
180     * @param string $lang Limit results to this language
181     * @return Chunk[]
182     * @throws \Exception
183     */
184    public function getSimilarChunks($query, $lang = '')
185    {
186        global $auth;
187        $vector = $this->model->getEmbedding($query);
188
189        $fetch = ceil(
190            ($this->model->getMaxContextTokenLength() / $this->model->getMaxEmbeddingTokenLength())
191            * 1.5 // fetch a few more than needed, since not all chunks are maximum length
192        );
193
194        $time = microtime(true);
195        $chunks = $this->storage->getSimilarChunks($vector, $lang, $fetch);
196        if ($this->logger instanceof CLI) {
197            $this->logger->info(
198                'Fetched {count} similar chunks from store in {time} seconds',
199                ['count' => count($chunks), 'time' => round(microtime(true) - $time, 2)]
200            );
201        }
202
203        $size = 0;
204        $result = [];
205        foreach ($chunks as $chunk) {
206            // filter out chunks the user is not allowed to read
207            if ($auth && auth_quickaclcheck($chunk->getPage()) < AUTH_READ) continue;
208
209            $chunkSize = count($this->getTokenEncoder()->encode($chunk->getText()));
210            if ($size + $chunkSize > $this->model->getMaxContextTokenLength()) break; // we have enough
211
212            $result[] = $chunk;
213            $size += $chunkSize;
214        }
215        return $result;
216    }
217
218
219    /**
220     * @param $text
221     * @return array
222     * @throws \Exception
223     * @todo support splitting too long sentences
224     */
225    public function splitIntoChunks($text)
226    {
227        $sentenceSplitter = new Sentence();
228        $tiktok = $this->getTokenEncoder();
229
230        $chunks = [];
231        $sentences = $sentenceSplitter->split($text);
232
233        $chunklen = 0;
234        $chunk = '';
235        while ($sentence = array_shift($sentences)) {
236            $slen = count($tiktok->encode($sentence));
237            if ($slen > $this->model->getMaxEmbeddingTokenLength()) {
238                // sentence is too long, we need to split it further
239                if ($this->logger instanceof CLI) $this->logger->warning(
240                    'Sentence too long, splitting not implemented yet'
241                );
242                continue;
243            }
244
245            if ($chunklen + $slen < $this->model->getMaxEmbeddingTokenLength()) {
246                // add to current chunk
247                $chunk .= $sentence;
248                $chunklen += $slen;
249                // remember sentence for overlap check
250                $this->rememberSentence($sentence);
251            } else {
252                // add current chunk to result
253                $chunks[] = $chunk;
254
255                // start new chunk with remembered sentences
256                $chunk = implode(' ', $this->sentenceQueue);
257                $chunk .= $sentence;
258                $chunklen = count($tiktok->encode($chunk));
259            }
260        }
261        $chunks[] = $chunk;
262
263        return $chunks;
264    }
265
266    /**
267     * Add a sentence to the queue of remembered sentences
268     *
269     * @param string $sentence
270     * @return void
271     */
272    protected function rememberSentence($sentence)
273    {
274        // add sentence to queue
275        $this->sentenceQueue[] = $sentence;
276
277        // remove oldest sentences from queue until we are below the max overlap
278        $encoder = $this->getTokenEncoder();
279        while (count($encoder->encode(implode(' ', $this->sentenceQueue))) > self::MAX_OVERLAP_LEN) {
280            array_shift($this->sentenceQueue);
281        }
282    }
283}
284