1<?php
2
3use dokuwiki\plugin\aichat\AbstractCLI;
4use splitbrain\phpcli\Colors;
5use splitbrain\phpcli\Options;
6
7/**
8 * DokuWiki Plugin aichat (CLI Component)
9 *
10 * @license GPL 2 http://www.gnu.org/licenses/gpl-2.0.html
11 * @author  Andreas Gohr <gohr@cosmocode.de>
12 */
13class cli_plugin_aichat_simulate extends AbstractCLI
14{
15    /** @inheritDoc */
16    public function getInfo()
17    {
18        $info = parent::getInfo();
19        $info['desc'] = 'Run a prepared chat session against multiple LLM models';
20        return $info;
21    }
22
23    /** @inheritDoc */
24    protected function setup(Options $options)
25    {
26        parent::setup($options);
27
28        $options->setHelp('Run a prepared chat session against multiple models');
29        $options->registerArgument('input', 'A file with the chat questions. Each question separated by two newlines');
30        $options->registerArgument('output', 'Where to write the result CSV to');
31
32        $options->registerOption(
33            'filter',
34            'Use only models matching this case-insensitive regex (no delimiters)',
35            'f',
36            'regex'
37        );
38    }
39
40    /** @inheritDoc */
41    protected function main(Options $options)
42    {
43        parent::main($options);
44
45        [$input, $output] = $options->getArgs();
46        $questions = $this->readInputFile($input);
47        $outFH = @fopen($output, 'w');
48        if (!$outFH) throw new \Exception("Could not open $output for writing");
49
50        $models = $this->helper->factory->getModels(true, 'chat');
51
52        $results = [];
53        foreach ($models as $name => $info) {
54            if ($options->getOpt('filter') && !preg_match('/' . $options->getOpt('filter') . '/i', $name)) {
55                continue;
56            }
57            $this->success("Running on $name...");
58            $results[$name] = $this->simulate($questions, $info);
59        }
60
61        foreach ($this->records2rows($results) as $row) {
62            fputcsv($outFH, $row);
63        }
64        fclose($outFH);
65        $this->success("Results written to $output");
66    }
67
68    protected function simulate($questions, $model)
69    {
70        // override models
71        $this->helper->factory->chatModel = $model['instance'];
72        $this->helper->factory->rephraseModel = clone $model['instance'];
73
74        $records = [];
75
76        $history = [];
77        foreach ($questions as $q) {
78            $this->helper->getChatModel()->resetUsageStats();
79            $this->helper->getRephraseModel()->resetUsageStats();
80            $this->helper->getEmbeddingModel()->resetUsageStats();
81
82            $this->colors->ptln($q, Colors::C_LIGHTPURPLE);
83            try {
84                $result = $this->helper->askChatQuestion($q, $history);
85                $history[] = [$result['question'], $result['answer']];
86                $this->colors->ptln($result['question'], Colors::C_LIGHTBLUE);
87            } catch (Exception $e) {
88                $this->error($e->getMessage());
89                $this->debug($e->getTraceAsString());
90                $result = ['question' => $q, 'answer' => "ERROR\n" . $e->getMessage(), 'sources' => []];
91            }
92
93            $record = [
94                'question' => $q,
95                'rephrased' => $result['contextQuestion'],
96                'answer' => $result['answer'],
97                'source.list' => implode("\n", $result['sources']),
98                'source.time' => $this->helper->getEmbeddings()->timeSpent,
99                ...$this->flattenStats('stats.embedding', $this->helper->getEmbeddingModel()->getUsageStats()),
100                ...$this->flattenStats('stats.rephrase', $this->helper->getRephraseModel()->getUsageStats()),
101                ...$this->flattenStats('stats.chat', $this->helper->getChatModel()->getUsageStats()),
102            ];
103            $records[] = $record;
104            $this->colors->ptln($result['answer'], Colors::C_LIGHTCYAN);
105        }
106
107        return $records;
108    }
109
110    /**
111     * Reformat the result array into a CSV friendly array
112     */
113    protected function records2rows(array $result): array
114    {
115        $rowkeys = [
116            'question' => ['question', 'stats.embedding.cost', 'stats.embedding.time'],
117            'rephrased' => ['rephrased', 'stats.rephrase.cost', 'stats.rephrase.time'],
118            'sources' => ['source.list', '', 'source.time'],
119            'answer' => ['answer', 'stats.chat.cost', 'stats.chat.time'],
120        ];
121
122        $models = array_keys($result);
123        $numberOfRecords = count($result[$models[0]]);
124        $rows = [];
125
126        // write headers
127        $row = [];
128        $row[] = 'type';
129        foreach ($models as $model) {
130            $row[] = $model;
131            $row[] = 'Cost USD';
132            $row[] = 'Time s';
133        }
134        $rows[] = $row;
135
136        // write rows
137        for ($i = 0; $i < $numberOfRecords; $i++) {
138            foreach ($rowkeys as $type => $keys) {
139                $row = [];
140                $row[] = $type;
141                foreach ($models as $model) {
142                    foreach ($keys as $key) {
143                        if ($key) {
144                            $row[] = $result[$model][$i][$key];
145                        } else {
146                            $row[] = '';
147                        }
148                    }
149                }
150                $rows[] = $row;
151            }
152        }
153
154
155        return $rows;
156    }
157
158
159    /**
160     * Prefix each key in the given stats array to be merged with a larger array
161     *
162     * @param string $prefix
163     * @param array $stats
164     * @return array
165     */
166    protected function flattenStats(string $prefix, array $stats)
167    {
168        $result = [];
169        foreach ($stats as $key => $value) {
170            $result["$prefix.$key"] = $value;
171        }
172        return $result;
173    }
174
175    /**
176     * @param string $file
177     * @return array
178     * @throws Exception
179     */
180    protected function readInputFile(string $file): array
181    {
182        if (!file_exists($file)) throw new \Exception("File not found: $file");
183        $lines = file_get_contents($file);
184        $questions = explode("\n\n", $lines);
185        $questions = array_map('trim', $questions);
186        return $questions;
187    }
188}
189