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                'customprompt' => $this->getConf('customprompt'),
206            ]);
207        } else {
208            $prompt = $this->getPrompt('noanswer', [
209                'question' => $question,
210            ]);
211            $history = [];
212        }
213
214        $messages = $this->prepareMessages(
215            $this->getChatModel(),
216            $prompt,
217            $history,
218            $this->getConf('chatHistory')
219        );
220        $answer = $this->getChatModel()->getAnswer($messages);
221
222        return [
223            'question' => $question,
224            'contextQuestion' => $contextQuestion,
225            'answer' => $answer,
226            'sources' => $similar,
227        ];
228    }
229
230    /**
231     * Rephrase a question into a standalone question based on the chat history
232     *
233     * @param string $question The original user question
234     * @param array[] $history The chat history [[user, ai], [user, ai], ...]
235     * @return string The rephrased question
236     * @throws Exception
237     */
238    public function rephraseChatQuestion($question, $history)
239    {
240        $prompt = $this->getPrompt('rephrase', [
241            'question' => $question,
242        ]);
243        $messages = $this->prepareMessages(
244            $this->getRephraseModel(),
245            $prompt,
246            $history,
247            $this->getConf('rephraseHistory')
248        );
249        return $this->getRephraseModel()->getAnswer($messages);
250    }
251
252    /**
253     * Prepare the messages for the AI
254     *
255     * @param ChatInterface $model The used model
256     * @param string $promptedQuestion The user question embedded in a prompt
257     * @param array[] $history The chat history [[user, ai], [user, ai], ...]
258     * @param int $historySize The maximum number of messages to use from the history
259     * @return array An OpenAI compatible array of messages
260     */
261    protected function prepareMessages(
262        ChatInterface $model,
263        string $promptedQuestion,
264        array $history,
265        int $historySize
266    ): array {
267        // calculate the space for context
268        $remainingContext = $model->getMaxInputTokenLength();
269        $remainingContext -= $this->countTokens($promptedQuestion);
270        $safetyMargin = $remainingContext * 0.05; // 5% safety margin
271        $remainingContext -= $safetyMargin;
272        // FIXME we may want to also have an upper limit for the history and not always use the full context
273
274        $messages = $this->historyMessages($history, $remainingContext, $historySize);
275        $messages[] = [
276            'role' => 'user',
277            'content' => $promptedQuestion
278        ];
279        return $messages;
280    }
281
282    /**
283     * Create an array of OpenAI compatible messages from the given history
284     *
285     * Only as many messages are used as fit into the token limit
286     *
287     * @param array[] $history The chat history [[user, ai], [user, ai], ...]
288     * @param int $tokenLimit The maximum number of tokens to use
289     * @param int $sizeLimit The maximum number of messages to use
290     * @return array
291     */
292    protected function historyMessages(array $history, int $tokenLimit, int $sizeLimit): array
293    {
294        $remainingContext = $tokenLimit;
295
296        $messages = [];
297        $history = array_reverse($history);
298        $history = array_slice($history, 0, $sizeLimit);
299        foreach ($history as $row) {
300            $length = $this->countTokens($row[0] . $row[1]);
301            if ($length > $remainingContext) {
302                break;
303            }
304            $remainingContext -= $length;
305
306            $messages[] = [
307                'role' => 'assistant',
308                'content' => $row[1]
309            ];
310            $messages[] = [
311                'role' => 'user',
312                'content' => $row[0]
313            ];
314        }
315        return array_reverse($messages);
316    }
317
318    /**
319     * Get an aproximation of the token count for the given text
320     *
321     * @param $text
322     * @return int
323     */
324    protected function countTokens($text)
325    {
326        return count($this->getEmbeddings()->getTokenEncoder()->encode($text));
327    }
328
329    /**
330     * Load the given prompt template and fill in the variables
331     *
332     * @param string $type
333     * @param string[] $vars
334     * @return string
335     */
336    protected function getPrompt($type, $vars = [])
337    {
338        $template = file_get_contents($this->localFN($type, 'prompt'));
339        $vars['language'] = $this->getLanguagePrompt();
340
341        $replace = [];
342        foreach ($vars as $key => $val) {
343            $replace['{{' . strtoupper($key) . '}}'] = $val;
344        }
345
346        return strtr($template, $replace);
347    }
348
349    /**
350     * Construct the prompt to define the answer language
351     *
352     * @return string
353     */
354    protected function getLanguagePrompt()
355    {
356        global $conf;
357        $isoLangnames = include(__DIR__ . '/lang/languages.php');
358
359        $currentLang = $isoLangnames[$conf['lang']] ?? 'English';
360
361        if ($this->getConf('preferUIlanguage') > AIChat::LANG_AUTO_ALL) {
362            if (isset($isoLangnames[$conf['lang']])) {
363                $languagePrompt = 'Always answer in ' . $isoLangnames[$conf['lang']] . '.';
364                return $languagePrompt;
365            }
366        }
367
368        $languagePrompt = 'Always answer in the user\'s language. ' .
369            "If you are unsure about the language, speak $currentLang.";
370        return $languagePrompt;
371    }
372
373    /**
374     * Should sources be limited to current language?
375     *
376     * @return string The current language code or empty string
377     */
378    public function getLanguageLimit()
379    {
380        if ($this->getConf('preferUIlanguage') >= AIChat::LANG_UI_LIMITED) {
381            global $conf;
382            return $conf['lang'];
383        } else {
384            return '';
385        }
386    }
387
388    /**
389     * Store info about the last run
390     *
391     * @param array $data
392     * @return void
393     */
394    public function setRunData(array $data)
395    {
396        file_put_contents($this->runDataFile, json_encode($data, JSON_PRETTY_PRINT));
397    }
398
399    /**
400     * Get info about the last run
401     *
402     * @return array
403     */
404    public function getRunData()
405    {
406        if (!file_exists($this->runDataFile)) {
407            return [];
408        }
409        return json_decode(file_get_contents($this->runDataFile), true);
410    }
411}
412