diff --git a/src/callback_commands/cancel.ts b/src/callback_commands/cancel.ts new file mode 100644 index 0000000..69e2eb2 --- /dev/null +++ b/src/callback_commands/cancel.ts @@ -0,0 +1,22 @@ +import {CallbackCommand} from "../base/callback-command"; + +export class Cancel extends CallbackCommand { + + text = "❌ Отменить"; + data = null; + + constructor(text?: string, data?: string) { + super(); + + this.text = text ?? this.text; + this.data = data ?? this.data; + } + + static withData(data?: string): Cancel { + return new Cancel(null, data); + } + + async execute(): Promise { + return Promise.resolve(); + } +} \ No newline at end of file diff --git a/src/callback_commands/ollama-cancel.ts b/src/callback_commands/ollama-cancel.ts new file mode 100644 index 0000000..90717bf --- /dev/null +++ b/src/callback_commands/ollama-cancel.ts @@ -0,0 +1,32 @@ +import {CallbackCommand} from "../base/callback-command"; +import {CallbackQuery} from "typescript-telegram-bot-api"; +import {abortOllamaRequest, bot, getOllamaRequest} from "../index"; +import {logError} from "../util/utils"; + +export class OllamaCancel extends CallbackCommand { + + data = "/cancel_ollama"; + text = "Cancel Ollama generation"; + + async execute(query: CallbackQuery): Promise { + const chatId = query.message.chat.id; + const fromId = query.from.id; + const messageId = query.message.message_id; + + const uuid = query.data.split(" ")[1]; + if (!uuid) return; + + const request = getOllamaRequest(uuid); + if (!request) return; + if (request.fromId !== fromId) return; + + const aborted = abortOllamaRequest(uuid); + console.log(`aborted request ${uuid}:`, aborted); + + await bot.editMessageReplyMarkup({ + chat_id: chatId, + message_id: messageId, + reply_markup: {inline_keyboard: []} + }).catch(logError); + } +} \ No newline at end of file diff --git a/src/commands/ollama-chat.ts b/src/commands/ollama-chat.ts index 512d766..888440b 100644 --- a/src/commands/ollama-chat.ts +++ b/src/commands/ollama-chat.ts @@ -1,6 +1,6 @@ import {ChatCommand} from "../base/chat-command"; import {Message} from "typescript-telegram-bot-api"; -import {bot, ollama} from "../index"; +import {abortOllamaRequest, bot, getOllamaRequest, ollama, ollamaRequests} from "../index"; import { collectReplyChainText, editMessageText, @@ -16,6 +16,8 @@ import {MessageStore} from "../common/message-store"; import axios from "axios"; import * as fs from "node:fs"; import path from "node:path"; +import {Cancel} from "../callback_commands/cancel"; +import {OllamaCancel} from "../callback_commands/ollama-cancel"; export class OllamaChat extends ChatCommand { regexp = /^\/ollama\s([^]+)/; @@ -73,13 +75,18 @@ export class OllamaChat extends ChatCommand { const startTime = Date.now(); try { + let isOver: boolean = false; + const uuid = crypto.randomUUID(); + const cancelMarkup = {inline_keyboard: [[Cancel.withData(new OllamaCancel().data + " " + uuid).asButton()]]}; + waitMessage = await bot.sendMessage({ chat_id: chatId, text: maxSize !== null ? `🔍 Внимательно изучаю изображение...\n🤓 ${maxSize.width}x${maxSize.height}px` : Environment.waitText, reply_parameters: { chat_id: chatId, message_id: msg.message_id - } + }, + reply_markup: cancelMarkup }); const stream = await ollama.chat({ @@ -90,6 +97,8 @@ export class OllamaChat extends ChatCommand { messages: chatMessages }); + ollamaRequests.push({uuid: uuid, stream: stream, done: false, fromId: msg.from.id, chatId: msg.chat.id}); + let currentText = ""; let shouldBreak = false; @@ -97,7 +106,13 @@ export class OllamaChat extends ChatCommand { intervalMs: 4500, getText: () => currentText, editFn: async (text) => { - await editMessageText(chatId, waitMessage.message_id, escapeMarkdownV2Text(text), "Markdown"); + await editMessageText( + chatId, + waitMessage.message_id, + escapeMarkdownV2Text(text), + "Markdown", + isOver ? {inline_keyboard: []} : cancelMarkup + ).catch(logError); }, onStop: async () => { } @@ -105,16 +120,22 @@ export class OllamaChat extends ChatCommand { try { for await (const chunk of stream) { - const content = chunk.message.content; - currentText += content; + if (!getOllamaRequest(uuid).done) { + const content = chunk.message.content; + currentText += content; - const length = currentText.length; - if (length > 4096) { - currentText = currentText.slice(0, 4093) + "..."; + const length = currentText.length; + if (length > 4096) { + currentText = currentText.slice(0, 4093) + "..."; + shouldBreak = true; + } + } else { shouldBreak = true; } if (shouldBreak || chunk.done) { + isOver = true; + console.log("messageText", currentText); console.log("length", length); @@ -124,7 +145,7 @@ export class OllamaChat extends ChatCommand { console.log("ended", true); } - stream.abort(); + console.log(`aborted request ${uuid}:`, abortOllamaRequest(uuid)); const diff = Math.abs(Date.now() - startTime) / 1000; @@ -140,13 +161,15 @@ export class OllamaChat extends ChatCommand { } } } finally { + console.log(`aborted request ${uuid}:`, abortOllamaRequest(uuid)); await editor.tick(); await editor.stop(); } } catch (error) { + if (error.message === "This operation was aborted") return; + console.error(error); await replyToMessage(waitMessage, `Произошла ошибка!\n${error.toString()}`).catch(logError); } - return Promise.resolve(); } } \ No newline at end of file diff --git a/src/index.ts b/src/index.ts index dfc9f26..248ee02 100644 --- a/src/index.ts +++ b/src/index.ts @@ -6,6 +6,7 @@ import { checkRequirements, executeChatCommand, extractTextMessage, + findAndExecuteCallbackCommand, initSystemSpecs, logError, randomValue, @@ -53,6 +54,9 @@ import {MessageDao} from "./db/message-dao"; import {DatabaseManager} from "./db/database-manager"; import {UserDao} from "./db/user-dao"; import {UserStore} from "./common/user-store"; +import {OllamaRequest} from "./model/ollama-request"; +import {CallbackCommand} from "./base/callback-command"; +import {OllamaCancel} from "./callback_commands/ollama-cancel"; process.setUncaughtExceptionCaptureCallback(console.error); @@ -70,6 +74,33 @@ export const ollama = new Ollama({ headers: {"Authorization": `Bearer ${Environment.OLLAMA_API_KEY}`} }); +export const ollamaRequests: OllamaRequest[] = []; + +export function getOllamaRequest(uuid: string): OllamaRequest | null { + return ollamaRequests.find(r => r.uuid === uuid); +} + +export function updateOllamaRequest(uuid: string, request: OllamaRequest) { + const index = ollamaRequests.findIndex(r => r.uuid === uuid); + if (index >= 0) { + ollamaRequests[index] = request; + } +} + +export function abortOllamaRequest(uuid: string): boolean { + const request = getOllamaRequest(uuid); + if (!request || request.done) return false; + + try { + request.stream.abort(); + updateOllamaRequest(uuid, {...request, done: true}); + return true; + } catch (e) { + console.error(e); + return false; + } +} + export const googleAi = new GoogleGenAI({apiKey: Environment.GEMINI_API_KEY}); export let systemInfoText: string = ""; @@ -113,6 +144,10 @@ export const chatCommands: ChatCommand[] = [ new Leave(), ]; +export const callbackCommands: CallbackCommand[] = [ + new OllamaCancel() +]; + if (Environment.OLLAMA_ADDRESS && Environment.OLLAMA_MODEL && Environment.SYSTEM_PROMPT) { chatCommands.push(new OllamaChat(), new OllamaPrompt(), new OllamaKill()); } @@ -257,4 +292,9 @@ bot.on("inline_query", async (query) => { } }); +bot.on("callback_query", async (query) => { + console.log(query); + await findAndExecuteCallbackCommand(callbackCommands, query); +}); + main().catch(console.error); \ No newline at end of file diff --git a/src/model/ollama-request.ts b/src/model/ollama-request.ts new file mode 100644 index 0000000..279d920 --- /dev/null +++ b/src/model/ollama-request.ts @@ -0,0 +1,7 @@ +export type OllamaRequest = { + uuid: string; + stream: any; + done: boolean; + fromId: number; + chatId: number; +} \ No newline at end of file diff --git a/src/util/utils.ts b/src/util/utils.ts index fe52627..b5fd69e 100644 --- a/src/util/utils.ts +++ b/src/util/utils.ts @@ -127,20 +127,21 @@ export async function findAndExecuteCallbackCommand(commands: CallbackCommand[], const fromId = query.from.id; const data = query.data || ""; - const command = searchCallbackCommand(commands, data); - if (!command) return false; + const cmd = searchCallbackCommand(commands, data); + if (!cmd) return false; - const requirements = command.requirements; + // TODO: 15/01/2026, Danil Nikolaev: reimplement + const requirements = cmd.requirements; if (requirements) { if (requirements.isRequiresBotAdmin() && !Environment.ADMIN_IDS.has(fromId)) { - console.log(`${command.data}: adminId is bad: ${fromId}`); + console.log(`${cmd.data}: adminId is bad: ${fromId}`); return false; } } - await command.execute(query); - await command.answerCallbackQuery(query); - await command.afterExecute(query); + await cmd.execute(query); + await cmd.answerCallbackQuery(query); + await cmd.afterExecute(query); return true; }