xref: /plugin/aichat/cli.php (revision f6ef2e505783ac17f756e44bf15c66238362377a) 
1<?php
2
3use dokuwiki\Extension\CLIPlugin;
4use dokuwiki\plugin\aichat\Chunk;
5use splitbrain\phpcli\Colors;
6use splitbrain\phpcli\Options;
7
8
9/**
10 * DokuWiki Plugin aichat (CLI Component)
11 *
12 * @license GPL 2 http://www.gnu.org/licenses/gpl-2.0.html
13 * @author  Andreas Gohr <gohr@cosmocode.de>
14 */
15class cli_plugin_aichat extends CLIPlugin
16{
17    /** @var helper_plugin_aichat */
18    protected $helper;
19
20    public function __construct($autocatch = true)
21    {
22        parent::__construct($autocatch);
23        $this->helper = plugin_load('helper', 'aichat');
24        $this->helper->getEmbeddings()->setLogger($this);
25    }
26
27    /** @inheritDoc */
28    protected function setup(Options $options)
29    {
30        $options->useCompactHelp();
31
32        $options->setHelp(
33            'Manage and query the AI chatbot data. Please note that calls to your LLM provider will be made. ' .
34            'This may incur costs.'
35        );
36
37        $options->registerCommand(
38            'embed',
39            'Create embeddings for all pages. This skips pages that already have embeddings'
40        );
41        $options->registerOption(
42            'clear',
43            'Clear all existing embeddings before creating new ones',
44            'c', false, 'embed'
45        );
46
47        $options->registerCommand('similar', 'Search for similar pages');
48        $options->registerArgument('query', 'Look up chunks similar to this query', true, 'similar');
49
50        $options->registerCommand('ask', 'Ask a question');
51        $options->registerArgument('question', 'The question to ask', true, 'ask');
52
53        $options->registerCommand('chat', 'Start an interactive chat session');
54
55        $options->registerCommand('split', 'Split a page into chunks (for debugging)');
56        $options->registerArgument('page', 'The page to split', true, 'split');
57
58        $options->registerCommand('info', 'Get Info about the vector storage');
59    }
60
61    /** @inheritDoc */
62    protected function main(Options $options)
63    {
64        switch ($options->getCmd()) {
65
66            case 'embed':
67                $this->createEmbeddings($options->getOpt('clear'));
68                break;
69            case 'similar':
70                $this->similar($options->getArgs()[0]);
71                break;
72            case 'ask':
73                $this->ask($options->getArgs()[0]);
74                break;
75            case 'chat':
76                $this->chat();
77                break;
78            case 'split':
79                $this->split($options->getArgs()[0]);
80                break;
81            case 'info':
82                $this->showinfo();
83                break;
84            default:
85                echo $options->help();
86        }
87    }
88
89    /**
90     * @return void
91     */
92    protected function showinfo()
93    {
94        echo 'model: '. $this->getConf('model') . "\n";
95        $stats = $this->helper->getEmbeddings()->getStorage()->statistics();
96        foreach ($stats as $key => $value) {
97            echo $key . ': ' . $value . "\n";
98        }
99    }
100
101    /**
102     * Split the given page into chunks and print them
103     *
104     * @param string $page
105     * @return void
106     * @throws Exception
107     */
108    protected function split($page)
109    {
110        $text = rawWiki($page);
111        $chunks = $this->helper->getEmbeddings()->splitIntoChunks($text);
112        foreach ($chunks as $chunk) {
113            echo $chunk;
114            echo "\n";
115            $this->colors->ptln('--------------------------------', Colors::C_LIGHTPURPLE);
116        }
117        $this->success('Split into ' . count($chunks) . ' chunks');
118    }
119
120    /**
121     * Interactive Chat Session
122     *
123     * @return void
124     * @throws Exception
125     */
126    protected function chat()
127    {
128        $history = [];
129        while ($q = $this->readLine('Your Question')) {
130            $this->helper->getModel()->resetUsageStats();
131            $result = $this->helper->askChatQuestion($q, $history);
132            $this->colors->ptln("Interpretation: {$result['question']}", Colors::C_LIGHTPURPLE);
133            $history[] = [$result['question'], $result['answer']];
134            $this->printAnswer($result);
135        }
136    }
137
138    /**
139     * Handle a single, standalone question
140     *
141     * @param string $query
142     * @return void
143     * @throws Exception
144     */
145    protected function ask($query)
146    {
147        $result = $this->helper->askQuestion($query);
148        $this->printAnswer($result);
149    }
150
151    /**
152     * Get the pages that are similar to the query
153     *
154     * @param string $query
155     * @return void
156     */
157    protected function similar($query)
158    {
159        $sources = $this->helper->getEmbeddings()->getSimilarChunks($query);
160        $this->printSources($sources);
161    }
162
163    /**
164     * Recreate chunks and embeddings for all pages
165     *
166     * @return void
167     * @todo make skip regex configurable
168     */
169    protected function createEmbeddings($clear)
170    {
171        ini_set('memory_limit', -1); // we may need a lot of memory here
172        $this->helper->getEmbeddings()->createNewIndex('/(^|:)(playground|sandbox)(:|$)/', $clear);
173        $this->notice('Peak memory used: {memory}', ['memory' => filesize_h(memory_get_peak_usage(true))]);
174    }
175
176    /**
177     * Print the given detailed answer in a nice way
178     *
179     * @param array $answer
180     * @return void
181     */
182    protected function printAnswer($answer)
183    {
184        $this->colors->ptln($answer['answer'], Colors::C_LIGHTCYAN);
185        echo "\n";
186        $this->printSources($answer['sources']);
187        echo "\n";
188        $this->printUsage();
189    }
190
191    /**
192     * Print the given sources
193     *
194     * @param Chunk[] $sources
195     * @return void
196     */
197    protected function printSources($sources)
198    {
199        foreach ($sources as $source) {
200            /** @var Chunk $source */
201            $this->colors->ptln("\t" . $source->getPage() . ' ' . $source->getId(), Colors::C_LIGHTBLUE);
202        }
203    }
204
205    /**
206     * Print the usage statistics for OpenAI
207     *
208     * @return void
209     */
210    protected function printUsage()
211    {
212        $this->info(
213            'Made {requests} requests in {time}s to Model. Used {tokens} tokens for about ${cost}.',
214            $this->helper->getModel()->getUsageStats()
215        );
216    }
217
218    /**
219     * Interactively ask for a value from the user
220     *
221     * @param string $prompt
222     * @return string
223     */
224    protected function readLine($prompt)
225    {
226        $value = '';
227
228        while ($value === '') {
229            echo $prompt;
230            echo ': ';
231
232            $fh = fopen('php://stdin', 'r');
233            $value = trim(fgets($fh));
234            fclose($fh);
235        }
236
237        return $value;
238    }
239}
240
241