1<?php
2
3use dokuwiki\Extension\CLIPlugin;
4use dokuwiki\Extension\Plugin;
5use dokuwiki\plugin\aichat\AIChat;
6use dokuwiki\plugin\aichat\Chunk;
7use dokuwiki\plugin\aichat\Embeddings;
8use dokuwiki\plugin\aichat\Model\ChatInterface;
9use dokuwiki\plugin\aichat\Model\EmbeddingInterface;
10use dokuwiki\plugin\aichat\ModelFactory;
11use dokuwiki\plugin\aichat\Storage\AbstractStorage;
12
13/**
14 * DokuWiki Plugin aichat (Helper Component)
15 *
16 * @license GPL 2 http://www.gnu.org/licenses/gpl-2.0.html
17 * @author  Andreas Gohr <gohr@cosmocode.de>
18 */
19class helper_plugin_aichat extends Plugin
20{
21    /** @var ModelFactory */
22    public $factory;
23
24    /** @var CLIPlugin $logger */
25    protected $logger;
26
27    /** @var Embeddings */
28    protected $embeddings;
29    /** @var AbstractStorage */
30    protected $storage;
31
32    /** @var array where to store meta data on the last run */
33    protected $runDataFile;
34
35
36    /**
37     * Constructor. Initializes vendor autoloader
38     */
39    public function __construct()
40    {
41        require_once __DIR__ . '/vendor/autoload.php'; // FIXME obsolete from Kaos onwards
42        global $conf;
43        $this->runDataFile = $conf['metadir'] . '/aichat__run.json';
44        $this->loadConfig();
45        $this->factory = new ModelFactory($this->conf);
46    }
47
48    /**
49     * Use the given CLI plugin for logging
50     *
51     * @param CLIPlugin $logger
52     * @return void
53     */
54    public function setLogger($logger)
55    {
56        $this->logger = $logger;
57    }
58
59    /**
60     * Update the configuration
61     *
62     * @param array $config
63     * @return void
64     */
65    public function updateConfig(array $config)
66    {
67        $this->conf = array_merge($this->conf, $config);
68        $this->factory->updateConfig($config);
69    }
70
71    /**
72     * Check if the current user is allowed to use the plugin (if it has been restricted)
73     *
74     * @return bool
75     */
76    public function userMayAccess()
77    {
78        global $auth;
79        global $USERINFO;
80        global $INPUT;
81
82        if (!$auth) return true;
83        if (!$this->getConf('restrict')) return true;
84        if (!isset($USERINFO)) return false;
85
86        return auth_isMember($this->getConf('restrict'), $INPUT->server->str('REMOTE_USER'), $USERINFO['grps']);
87    }
88
89    /**
90     * Access the Chat Model
91     *
92     * @return ChatInterface
93     */
94    public function getChatModel()
95    {
96        return $this->factory->getChatModel();
97    }
98
99    /**
100     * @return ChatInterface
101     */
102    public function getRephraseModel()
103    {
104        return $this->factory->getRephraseModel();
105    }
106
107    /**
108     * Access the Embedding Model
109     *
110     * @return EmbeddingInterface
111     */
112    public function getEmbeddingModel()
113    {
114        return $this->factory->getEmbeddingModel();
115    }
116
117    /**
118     * Access the Embeddings interface
119     *
120     * @return Embeddings
121     */
122    public function getEmbeddings()
123    {
124        if ($this->embeddings instanceof Embeddings) {
125            return $this->embeddings;
126        }
127
128        $this->embeddings = new Embeddings(
129            $this->getChatModel(),
130            $this->getEmbeddingModel(),
131            $this->getStorage(),
132            $this->conf
133        );
134        if ($this->logger) {
135            $this->embeddings->setLogger($this->logger);
136        }
137
138        return $this->embeddings;
139    }
140
141    /**
142     * Access the Storage interface
143     *
144     * @return AbstractStorage
145     */
146    public function getStorage()
147    {
148        if ($this->storage instanceof AbstractStorage) {
149            return $this->storage;
150        }
151
152        $class = '\\dokuwiki\\plugin\\aichat\\Storage\\' . $this->getConf('storage') . 'Storage';
153        $this->storage = new $class($this->conf);
154
155        if ($this->logger) {
156            $this->storage->setLogger($this->logger);
157        }
158
159        return $this->storage;
160    }
161
162    /**
163     * Ask a question with a chat history
164     *
165     * @param string $question
166     * @param array[] $history The chat history [[user, ai], [user, ai], ...]
167     * @return array ['question' => $question, 'answer' => $answer, 'sources' => $sources]
168     * @throws Exception
169     */
170    public function askChatQuestion($question, $history = [])
171    {
172        if ($history && $this->getConf('rephraseHistory') > 0) {
173            $contextQuestion = $this->rephraseChatQuestion($question, $history);
174
175            // Only use the rephrased question if it has more history than the chat history provides
176            if ($this->getConf('rephraseHistory') > $this->getConf('chatHistory')) {
177                $question = $contextQuestion;
178            }
179        } else {
180            $contextQuestion = $question;
181        }
182        return $this->askQuestion($question, $history, $contextQuestion);
183    }
184
185    /**
186     * Ask a single standalone question
187     *
188     * @param string $question The question to ask
189     * @param array $history [user, ai] of the previous question
190     * @param string $contextQuestion The question to use for context search
191     * @return array ['question' => $question, 'answer' => $answer, 'sources' => $sources]
192     * @throws Exception
193     */
194    public function askQuestion($question, $history = [], $contextQuestion = '')
195    {
196        $similar = $this->getEmbeddings()->getSimilarChunks($contextQuestion ?: $question, $this->getLanguageLimit());
197        if ($similar) {
198            $context = implode(
199                "\n",
200                array_map(static fn(Chunk $chunk) => "\n```\n" . $chunk->getText() . "\n```\n", $similar)
201            );
202            $prompt = $this->getPrompt('question', [
203                'context' => $context,
204                'question' => $question,
205            ]);
206        } else {
207            $prompt = $this->getPrompt('noanswer', [
208                'question' => $question,
209            ]);
210            $history = [];
211        }
212
213        $messages = $this->prepareMessages(
214            $this->getChatModel(),
215            $prompt,
216            $history,
217            $this->getConf('chatHistory')
218        );
219        $answer = $this->getChatModel()->getAnswer($messages);
220
221        return [
222            'question' => $question,
223            'contextQuestion' => $contextQuestion,
224            'answer' => $answer,
225            'sources' => $similar,
226        ];
227    }
228
229    /**
230     * Rephrase a question into a standalone question based on the chat history
231     *
232     * @param string $question The original user question
233     * @param array[] $history The chat history [[user, ai], [user, ai], ...]
234     * @return string The rephrased question
235     * @throws Exception
236     */
237    public function rephraseChatQuestion($question, $history)
238    {
239        $prompt = $this->getPrompt('rephrase', [
240            'question' => $question,
241        ]);
242        $messages = $this->prepareMessages(
243            $this->getRephraseModel(),
244            $prompt,
245            $history,
246            $this->getConf('rephraseHistory')
247        );
248        return $this->getRephraseModel()->getAnswer($messages);
249    }
250
251    /**
252     * Prepare the messages for the AI
253     *
254     * @param ChatInterface $model The used model
255     * @param string $promptedQuestion The user question embedded in a prompt
256     * @param array[] $history The chat history [[user, ai], [user, ai], ...]
257     * @param int $historySize The maximum number of messages to use from the history
258     * @return array An OpenAI compatible array of messages
259     */
260    protected function prepareMessages(
261        ChatInterface $model,
262        string $promptedQuestion,
263        array $history,
264        int $historySize
265    ): array {
266        // calculate the space for context
267        $remainingContext = $model->getMaxInputTokenLength();
268        $remainingContext -= $this->countTokens($promptedQuestion);
269        $safetyMargin = $remainingContext * 0.05; // 5% safety margin
270        $remainingContext -= $safetyMargin;
271        // FIXME we may want to also have an upper limit for the history and not always use the full context
272
273        $messages = $this->historyMessages($history, $remainingContext, $historySize);
274        $messages[] = [
275            'role' => 'user',
276            'content' => $promptedQuestion
277        ];
278        return $messages;
279    }
280
281    /**
282     * Create an array of OpenAI compatible messages from the given history
283     *
284     * Only as many messages are used as fit into the token limit
285     *
286     * @param array[] $history The chat history [[user, ai], [user, ai], ...]
287     * @param int $tokenLimit The maximum number of tokens to use
288     * @param int $sizeLimit The maximum number of messages to use
289     * @return array
290     */
291    protected function historyMessages(array $history, int $tokenLimit, int $sizeLimit): array
292    {
293        $remainingContext = $tokenLimit;
294
295        $messages = [];
296        $history = array_reverse($history);
297        $history = array_slice($history, 0, $sizeLimit);
298        foreach ($history as $row) {
299            $length = $this->countTokens($row[0] . $row[1]);
300            if ($length > $remainingContext) {
301                break;
302            }
303            $remainingContext -= $length;
304
305            $messages[] = [
306                'role' => 'assistant',
307                'content' => $row[1]
308            ];
309            $messages[] = [
310                'role' => 'user',
311                'content' => $row[0]
312            ];
313        }
314        return array_reverse($messages);
315    }
316
317    /**
318     * Get an aproximation of the token count for the given text
319     *
320     * @param $text
321     * @return int
322     */
323    protected function countTokens($text)
324    {
325        return count($this->getEmbeddings()->getTokenEncoder()->encode($text));
326    }
327
328    /**
329     * Load the given prompt template and fill in the variables
330     *
331     * @param string $type
332     * @param string[] $vars
333     * @return string
334     */
335    protected function getPrompt($type, $vars = [])
336    {
337        $template = file_get_contents($this->localFN($type, 'prompt'));
338        $vars['language'] = $this->getLanguagePrompt();
339
340        $replace = [];
341        foreach ($vars as $key => $val) {
342            $replace['{{' . strtoupper($key) . '}}'] = $val;
343        }
344
345        return strtr($template, $replace);
346    }
347
348    /**
349     * Construct the prompt to define the answer language
350     *
351     * @return string
352     */
353    protected function getLanguagePrompt()
354    {
355        global $conf;
356        $isoLangnames = include(__DIR__ . '/lang/languages.php');
357
358        $currentLang = $isoLangnames[$conf['lang']] ?? 'English';
359
360        if ($this->getConf('preferUIlanguage') > AIChat::LANG_AUTO_ALL) {
361            if (isset($isoLangnames[$conf['lang']])) {
362                $languagePrompt = 'Always answer in ' . $isoLangnames[$conf['lang']] . '.';
363                return $languagePrompt;
364            }
365        }
366
367        $languagePrompt = 'Always answer in the user\'s language. ' .
368            "If you are unsure about the language, speak $currentLang.";
369        return $languagePrompt;
370    }
371
372    /**
373     * Should sources be limited to current language?
374     *
375     * @return string The current language code or empty string
376     */
377    public function getLanguageLimit()
378    {
379        if ($this->getConf('preferUIlanguage') >= AIChat::LANG_UI_LIMITED) {
380            global $conf;
381            return $conf['lang'];
382        } else {
383            return '';
384        }
385    }
386
387    /**
388     * Store info about the last run
389     *
390     * @param array $data
391     * @return void
392     */
393    public function setRunData(array $data)
394    {
395        file_put_contents($this->runDataFile, json_encode($data, JSON_PRETTY_PRINT));
396    }
397
398    /**
399     * Get info about the last run
400     *
401     * @return array
402     */
403    public function getRunData()
404    {
405        if (!file_exists($this->runDataFile)) {
406            return [];
407        }
408        return json_decode(file_get_contents($this->runDataFile), true);
409    }
410}
411