support for thinking models, ux improvements

src/callback_commands/ollama-cancel.ts

@@ -1,8 +1,12 @@
 import {CallbackCommand} from "../base/callback-command";
 import {CallbackQuery} from "typescript-telegram-bot-api";
 import {abortOllamaRequest, bot, getOllamaRequest} from "../index";
-import {logError} from "../util/utils";
 import {Environment} from "../common/environment";
+import {logError} from "../util/utils";
+import {MessageStore} from "../common/message-store";
+import {StoredMessage} from "../model/stored-message";
+
+const cancelledText = "```Ollama\n❌ Отменено```";
 
 export class OllamaCancel extends CallbackCommand {
 
@@ -18,16 +22,35 @@ export class OllamaCancel extends CallbackCommand {
         if (!uuid) return;
 
         const request = getOllamaRequest(uuid);
-        if (!request) return;
+        if (request) {
             if (request.fromId !== fromId && fromId !== Environment.CREATOR_ID) return;
 
             const aborted = abortOllamaRequest(uuid);
             console.log(`aborted request ${uuid}:`, aborted);
+        }
 
-        await bot.editMessageReplyMarkup({
+        let msg: StoredMessage | null = null;
+        try {
+            msg = await MessageStore.get(chatId, messageId);
+        } catch (e) {
+            logError(e);
+        }
+
+        let content: string | null = null;
+
+        if (msg?.text?.trim()?.length > 0) {
+            content = msg?.text.trim();
+            if (content.length + cancelledText.length > 4096) {
+                content = content.substring(0, 4096 - cancelledText.length - 2) + "\n";
+            }
+        }
+
+        await bot.editMessageText({
             chat_id: chatId,
             message_id: messageId,
-            reply_markup: {inline_keyboard: []}
+            text: `${content}${cancelledText}`,
+            parse_mode: "Markdown",
+            reply_markup: {inline_keyboard: []},
         }).catch(logError);
     }
 }

src/commands/ollama-chat.ts

@@ -57,7 +57,6 @@ export class OllamaChat extends ChatCommand {
             waitMessage = await bot.sendMessage({
                 chat_id: chatId,
                 text: Environment.waitText,
-                // text: maxSize !== null ? `🔍 Внимательно изучаю изображение...\n🤓 ${maxSize.width}x${maxSize.height}px` : Environment.waitText,
                 reply_parameters: {
                     chat_id: chatId,
                     message_id: msg.message_id
@@ -68,9 +67,9 @@ export class OllamaChat extends ChatCommand {
             const stream = await ollama.chat({
                 model: Environment.OLLAMA_MODEL,
                 stream: true,
-                think: false,
                 keep_alive: 300,
-                messages: chatMessages
+                think: true,
+                messages: chatMessages,
             });
 
             ollamaRequests.push({uuid: uuid, stream: stream, done: false, fromId: msg.from.id, chatId: msg.chat.id});
@@ -79,16 +78,25 @@ export class OllamaChat extends ChatCommand {
             let shouldBreak = false;
 
             const editor = startIntervalEditor({
+                uuid: uuid,
                 intervalMs: 4500,
                 getText: () => currentText,
                 editFn: async (text) => {
+                    try {
                         await editMessageText(
                             chatId,
                             waitMessage.message_id,
                             escapeMarkdownV2Text(text),
                             "Markdown",
                             isOver ? {inline_keyboard: []} : cancelMarkup
-                    ).catch(logError);
+                        );
+
+                        waitMessage.reply_to_message = msg;
+                        waitMessage.text = text;
+                        await MessageStore.put(waitMessage);
+                    } catch (e) {
+                        logError(e);
+                    }
                 },
                 onStop: async () => {
                 }
@@ -96,8 +104,43 @@ export class OllamaChat extends ChatCommand {
             await editor.tick();
 
             try {
+                let isThinking = false;
+
                 for await (const chunk of stream) {
-                    currentText += chunk.message.content;
+                    const content = chunk.message.content;
+
+                    if (content === "<think>" || chunk.message.thinking) {
+                        if (!isThinking) {
+                            await editMessageText(
+                                chatId,
+                                waitMessage.message_id,
+                                "🤔 Размышляю...",
+                                "Markdown",
+                                isOver ? {inline_keyboard: []} : cancelMarkup
+                            ).catch(logError);
+                            console.log("Thinking:\n");
+                        }
+
+                        isThinking = true;
+                    }
+
+                    if (chunk.message.thinking) {
+                        console.log(chunk.message.thinking);
+                    } else {
+                        console.log(chunk.message.content);
+                    }
+
+                    if (!isThinking) {
+                        currentText += content;
+                    }
+
+                    if (isThinking && !chunk.message.thinking) {
+                        currentText += content;
+                    }
+
+                    if (content === "</think>" || !chunk.message.thinking) {
+                        isThinking = false;
+                    }
 
                     if (currentText.length > 4096) {
                         currentText = currentText.slice(0, 4093) + "...";
@@ -130,8 +173,7 @@ export class OllamaChat extends ChatCommand {
                         waitMessage.reply_to_message = msg;
                         waitMessage.text = currentText;
                         await MessageStore.put(waitMessage);
-
+                        await oldReplyToMessage(waitMessage, `⏱️ ${diff}s`);
-                        await oldReplyToMessage(waitMessage, `⏱️ ${diff}s` /*+ (maxSize !== null ? `\n🤓 ${maxSize.width}x${maxSize.height}px` : "")*/);
                         break;
                     }
                 }

src/commands/ollama-get-model.ts

@@ -9,12 +9,12 @@ export class OllamaGetModel extends ChatCommand {
     description = "Ollama model info";
 
     async execute(msg: Message): Promise<void> {
-
+        try {
             const showResponse = await ollama.show({model: Environment.OLLAMA_MODEL});
 
             const caps = showResponse.capabilities;
 
-        const text = "```ollama\n" +
+            const text = "```Ollama\n" +
                 `model: ${Environment.OLLAMA_MODEL}\n\n` +
                 `vision: ${boolToEmoji(caps.includes("vision"))}\n` +
                 `thinking: ${boolToEmoji(caps.includes("thinking"))}\n` +
@@ -22,5 +22,9 @@ export class OllamaGetModel extends ChatCommand {
                 + "```";
 
             await replyToMessage({message: msg, text: text, parse_mode: "Markdown"}).catch(logError);
+        } catch (e) {
+            logError(e);
+            await replyToMessage({message: msg, text: e.toString()}).catch(logError);
+        }
     }
 }

src/commands/ollama-set-model.ts

@@ -4,6 +4,7 @@ import {Environment} from "../common/environment";
 import {logError, replyToMessage} from "../util/utils";
 import {Requirements} from "../base/requirements";
 import {Requirement} from "../base/requirement";
+import {ollama} from "../index";
 
 export class OllamaSetModel extends ChatCommand {
     argsMode = "required" as const;
@@ -15,11 +16,19 @@ export class OllamaSetModel extends ChatCommand {
 
     async execute(msg: Message, match?: RegExpExecArray | null): Promise<void> {
         const newModel = match?.[3];
+
+        try {
+            await ollama.show({model: newModel});
+
             Environment.setOllamaModel(newModel || Environment.OLLAMA_MODEL);
 
             const text = newModel ? `Выбрана модель "${newModel}"`
                 : `Модель не задана. Будет использоваться стандартная модель "${Environment.OLLAMA_MODEL}".`;
 
             await replyToMessage({message: msg, text: text}).catch(logError);
+        } catch (e) {
+            logError(e);
+            await replyToMessage({message: msg, text: e.toString()}).catch(logError);
+        }
     }
 }

src/util/utils.ts

@@ -4,7 +4,7 @@ import {CallbackCommand} from "../base/callback-command";
 import {CallbackQuery, InlineKeyboardMarkup, Message, ParseMode, PhotoSize, User} from "typescript-telegram-bot-api";
 import {Environment} from "../common/environment";
 import {TelegramError} from "typescript-telegram-bot-api/dist/errors";
-import {bot, botUser, messageDao, setSystemInfo} from "../index";
+import {bot, botUser, getOllamaRequest, messageDao, setSystemInfo} from "../index";
 import os from "os";
 import axios from "axios";
 import {MessagePart} from "../common/message-part";
@@ -746,6 +746,7 @@ function toHex(v: number) {
 }
 
 export function startIntervalEditor(params: {
+    uuid?: string;
     intervalMs: number;
     getText: () => string;
     editFn: (text: string) => Promise<void>;
@@ -755,7 +756,7 @@ export function startIntervalEditor(params: {
     let stopped = false;
 
     const tick = async () => {
-        if (stopped) return;
+        if (stopped || (params.uuid && getOllamaRequest(params.uuid)?.done)) return;
         const next = params.getText();
         if (!next || next === lastSent) return;