add postgres db, use settings and user data, lots of cleanup and logic fixes, bug fixes, better error handling, update docs and docker

src/commands/ai.ts

@@ -38,6 +38,9 @@ import { languageCode } from "../utils/language-code"
 import axios from "axios"
 import { rateLimiter } from "../utils/rate-limiter"
 import { logger } from "../utils/log"
+import { ensureUserInDb } from "../utils/ensure-user"
+import * as schema from '../db/schema'
+import type { NodePgDatabase } from "drizzle-orm/node-postgres"
 
 const spamwatchMiddleware = spamwatchMiddlewareModule(isOnSpamWatch)
 export const flash_model = process.env.flashModel || "gemma3:4b"
@@ -45,6 +48,94 @@ export const thinking_model = process.env.thinkingModel || "qwen3:4b"
 
 type TextContext = Context & { message: Message.TextMessage }
 
+type User = typeof schema.usersTable.$inferSelect
+
+interface ModelInfo {
+  name: string;
+  label: string;
+  descriptionEn: string;
+  descriptionPt: string;
+  models: Array<{
+    name: string;
+    label: string;
+    parameterSize: string;
+  }>;
+}
+
+interface OllamaResponse {
+  response: string;
+}
+
+export const models: ModelInfo[] = [
+  {
+    name: 'gemma3n',
+    label: 'Gemma3n',
+    descriptionEn: 'Gemma3n is a family of open, light on-device models for general tasks.',
+    descriptionPt: 'Gemma3n é uma família de modelos abertos, leves e para dispositivos locais, para tarefas gerais.',
+    models: [
+      { name: 'gemma3n:e2b', label: 'Gemma3n e2b', parameterSize: '2B' },
+      { name: 'gemma3n:e4b', label: 'Gemma3n e4b', parameterSize: '4B' },
+    ]
+  },
+  {
+    name: 'gemma3-abliterated',
+    label: 'Gemma3 Uncensored',
+    descriptionEn: 'Gemma3-abliterated is a family of open, uncensored models for general tasks.',
+    descriptionPt: 'Gemma3-abliterated é uma família de modelos abertos, não censurados, para tarefas gerais.',
+    models: [
+      { name: 'huihui_ai/gemma3-abliterated:1b', label: 'Gemma3-abliterated 1B', parameterSize: '1b' },
+      { name: 'huihui_ai/gemma3-abliterated:4b', label: 'Gemma3-abliterated 4B', parameterSize: '4b' },
+    ]
+  },
+  {
+    name: 'qwen3',
+    label: 'Qwen3',
+    descriptionEn: 'Qwen3 is a multilingual reasoning model series.',
+    descriptionPt: 'Qwen3 é uma série de modelos multilingues.',
+    models: [
+      { name: 'qwen3:4b', label: 'Qwen3 4B', parameterSize: '4B' },
+    ]
+  },
+  {
+    name: 'deepseek',
+    label: 'DeepSeek',
+    descriptionEn: 'DeepSeek is a research model for reasoning tasks.',
+    descriptionPt: 'DeepSeek é um modelo de pesquisa para tarefas de raciocínio.',
+    models: [
+      { name: 'deepseek-r1:1.5b', label: 'DeepSeek 1.5B', parameterSize: '1.5B' },
+      { name: 'huihui_ai/deepseek-r1-abliterated:1.5b', label: 'DeepSeek Uncensored 1.5B', parameterSize: '1.5B' },
+    ]
+  }
+];
+
+const enSystemPrompt = `You are a plaintext-only, helpful assistant called {botName}.
+Current Date/Time (UTC): {date}
+
+---
+
+Respond to the user's message:
+{message}`
+
+const ptSystemPrompt = `Você é um assistente de texto puro e útil chamado {botName}.
+Data/Hora atual (UTC): {date}
+
+---
+
+Responda à mensagem do usuário:
+{message}`
+
+async function usingSystemPrompt(ctx: TextContext, db: NodePgDatabase<typeof schema>, botName: string): Promise<string> {
+  const user = await db.query.usersTable.findMany({ where: (fields, { eq }) => eq(fields.telegramId, String(ctx.from!.id)), limit: 1 });
+  if (user.length === 0) await ensureUserInDb(ctx, db);
+  const userData = user[0];
+  const lang = userData?.languageCode || "en";
+  const utcDate = new Date().toISOString();
+  const prompt = lang === "pt"
+    ? ptSystemPrompt.replace("{botName}", botName).replace("{date}", utcDate).replace("{message}", ctx.message.text)
+    : enSystemPrompt.replace("{botName}", botName).replace("{date}", utcDate).replace("{message}", ctx.message.text);
+  return prompt;
+}
+
 export function sanitizeForJson(text: string): string {
   return text
     .replace(/\\/g, '\\\\')
@@ -69,23 +160,50 @@ export async function preChecks() {
     }
     checked++;
   }
-  console.log(`[✨ AI] Pre-checks passed [${checked}/${envs.length}]\n`)
+
+  const ollamaApi = process.env.ollamaApi
+  if (!ollamaApi) {
+    console.error("[✨ AI | !] ❌ ollamaApi not set!")
+    return false
+  }
+  let ollamaOk = false
+  for (let i = 0; i < 10; i++) {
+    try {
+      const res = await axios.get(ollamaApi, { timeout: 2000 })
+      if (res && res.data && typeof res.data === 'object' && 'ollama' in res.data) {
+        ollamaOk = true
+        break
+      }
+      if (res && res.status === 200) {
+        ollamaOk = true
+        break
+      }
+    } catch (err) {
+      await new Promise(resolve => setTimeout(resolve, 1000))
+    }
+  }
+  if (!ollamaOk) {
+    console.error("[✨ AI | !] ❌ Ollama API is not responding at ", ollamaApi)
+    return false
+  }
+  checked++;
+  console.log(`[✨  AI] Pre-checks passed [${checked}/${envs.length + 1}]`)
   return true
 }
 
-function isAxiosError(error: unknown): error is { response?: { data?: { error?: string }, status?: number }, request?: unknown, message?: string } {
+function isAxiosError(error: unknown): error is { response?: { data?: { error?: string }, status?: number, statusText?: string }, request?: unknown, message?: string } {
   return typeof error === 'object' && error !== null && (
     'response' in error || 'request' in error || 'message' in error
-  )
+  );
 }
 
 function extractAxiosErrorMessage(error: unknown): string {
   if (isAxiosError(error)) {
-    const err = error as Record<string, unknown>;
+    const err = error as { response?: { data?: { error?: string }, status?: number, statusText?: string }, request?: unknown, message?: string };
     if (err.response && typeof err.response === 'object') {
-      const resp = err.response as Record<string, unknown>;
+      const resp = err.response;
       if (resp.data && typeof resp.data === 'object' && 'error' in resp.data) {
-        return String((resp.data as Record<string, unknown>).error);
+        return String(resp.data.error);
       }
       if ('status' in resp && 'statusText' in resp) {
         return `HTTP ${resp.status}: ${resp.statusText}`;
@@ -102,71 +220,68 @@ function extractAxiosErrorMessage(error: unknown): string {
   return 'An unexpected error occurred.';
 }
 
-async function getResponse(prompt: string, ctx: TextContext, replyGenerating: Message, model: string) {
-  const Strings = getStrings(languageCode(ctx))
-
+async function getResponse(prompt: string, ctx: TextContext, replyGenerating: Message, model: string, aiTemperature: number): Promise<{ success: boolean; response?: string; error?: string }> {
+  const Strings = getStrings(languageCode(ctx));
   if (!ctx.chat) {
     return {
       success: false,
       error: Strings.unexpectedErr.replace("{error}", "No chat found"),
-    }
+    };
   }
-
   try {
-    const aiResponse = await axios.post(
+    const aiResponse = await axios.post<unknown>(
       `${process.env.ollamaApi}/api/generate`,
       {
         model,
         prompt,
         stream: true,
+        options: {
+          temperature: aiTemperature
+        }
       },
       {
         responseType: "stream",
       }
-    )
-
-    let fullResponse = ""
-    let thoughts = ""
-    let lastUpdate = Date.now()
-
-    const stream = aiResponse.data
+    );
+    let fullResponse = "";
+    let thoughts = "";
+    let lastUpdate = Date.now();
+    const stream: NodeJS.ReadableStream = aiResponse.data as any;
     for await (const chunk of stream) {
-      const lines = chunk.toString().split('\n')
+      const lines = chunk.toString().split('\n');
       for (const line of lines) {
-        if (!line.trim()) continue
-        let ln
+        if (!line.trim()) continue;
+        let ln: OllamaResponse;
         try {
-          ln = JSON.parse(line)
+          ln = JSON.parse(line);
         } catch (e) {
-          console.error("[✨ AI | !] Error parsing chunk:", e)
-          continue
+          console.error("[✨ AI | !] Error parsing chunk:", e);
+          continue;
         }
-
-        if (model === thinking_model) {
+        if (model === thinking_model && ln.response) {
           if (ln.response.includes('<think>')) {
-            const thinkMatch = ln.response.match(/<think>([\s\S]*?)<\/think>/)
+            const thinkMatch = ln.response.match(/<think>([\s\S]*?)<\/think>/);
             if (thinkMatch && thinkMatch[1].trim().length > 0) {
-              logger.logThinking(ctx.chat.id, replyGenerating.message_id, true)
+              logger.logThinking(ctx.chat.id, replyGenerating.message_id, true);
             } else if (!thinkMatch) {
-              logger.logThinking(ctx.chat.id, replyGenerating.message_id, true)
+              logger.logThinking(ctx.chat.id, replyGenerating.message_id, true);
             }
           } else if (ln.response.includes('</think>')) {
-            logger.logThinking(ctx.chat.id, replyGenerating.message_id, false)
+            logger.logThinking(ctx.chat.id, replyGenerating.message_id, false);
           }
         }
-
-        const now = Date.now()
+        const now = Date.now();
         if (ln.response) {
           if (model === thinking_model) {
-            let patchedThoughts = ln.response
-            const thinkTagRx = /<think>([\s\S]*?)<\/think>/g
-            patchedThoughts = patchedThoughts.replace(thinkTagRx, (match, p1) => p1.trim().length > 0 ? '`Thinking...`' + p1 + '`Finished thinking`' : '')
-            patchedThoughts = patchedThoughts.replace(/<think>/g, '`Thinking...`')
-            patchedThoughts = patchedThoughts.replace(/<\/think>/g, '`Finished thinking`')
-            thoughts += patchedThoughts
-            fullResponse += patchedThoughts
+            let patchedThoughts = ln.response;
+            const thinkTagRx = /<think>([\s\S]*?)<\/think>/g;
+            patchedThoughts = patchedThoughts.replace(thinkTagRx, (match, p1) => p1.trim().length > 0 ? '`Thinking...`' + p1 + '`Finished thinking`' : '');
+            patchedThoughts = patchedThoughts.replace(/<think>/g, '`Thinking...`');
+            patchedThoughts = patchedThoughts.replace(/<\/think>/g, '`Finished thinking`');
+            thoughts += patchedThoughts;
+            fullResponse += patchedThoughts;
           } else {
-            fullResponse += ln.response
+            fullResponse += ln.response;
           }
           if (now - lastUpdate >= 1000) {
             await rateLimiter.editMessageWithRetry(
@@ -175,67 +290,104 @@ async function getResponse(prompt: string, ctx: TextContext, replyGenerating: Me
               replyGenerating.message_id,
               thoughts,
               { parse_mode: 'Markdown' }
-            )
-            lastUpdate = now
+            );
+            lastUpdate = now;
           }
         }
       }
     }
-
     return {
       success: true,
       response: fullResponse,
-    }
+    };
   } catch (error: unknown) {
-    const errorMsg = extractAxiosErrorMessage(error)
-    console.error("[✨ AI | !] Error:", errorMsg)
-
-    // model not found or 404
+    const errorMsg = extractAxiosErrorMessage(error);
+    console.error("[✨ AI | !] Error:", errorMsg);
     if (isAxiosError(error) && error.response && typeof error.response === 'object') {
-      const resp = error.response as Record<string, unknown>;
-      const errData = resp.data && typeof resp.data === 'object' && 'error' in resp.data ? (resp.data as Record<string, unknown>).error : undefined;
+      const resp = error.response as { data?: { error?: string }, status?: number };
+      const errData = resp.data && typeof resp.data === 'object' && 'error' in resp.data ? (resp.data as { error?: string }).error : undefined;
       const errStatus = 'status' in resp ? resp.status : undefined;
       if ((typeof errData === 'string' && errData.includes(`model '${model}' not found`)) || errStatus === 404) {
-        ctx.telegram.editMessageText(
-          ctx.chat.id,
+        await ctx.telegram.editMessageText(
+          ctx.chat!.id,
           replyGenerating.message_id,
           undefined,
-          `🔄 *Pulling ${model} from Ollama...*\n\nThis may take a few minutes...`,
+          Strings.ai.pulling.replace("{model}", model),
           { parse_mode: 'Markdown' }
-        )
-        console.log(`[✨ AI | i] Pulling ${model} from ollama...`)
+        );
+        console.log(`[✨ AI | i] Pulling ${model} from ollama...`);
         try {
           await axios.post(
             `${process.env.ollamaApi}/api/pull`,
             {
               model,
               stream: false,
-              timeout: process.env.ollamaApiTimeout || 10000,
+              timeout: Number(process.env.ollamaApiTimeout) || 10000,
             }
-          )
+          );
         } catch (e: unknown) {
-          const pullMsg = extractAxiosErrorMessage(e)
-          console.error("[✨ AI | !] Pull error:", pullMsg)
+          const pullMsg = extractAxiosErrorMessage(e);
+          console.error("[✨ AI | !] Pull error:", pullMsg);
           return {
             success: false,
             error: `❌ Something went wrong while pulling ${model}: ${pullMsg}`,
-          }
+          };
         }
-        console.log(`[✨ AI | i] ${model} pulled successfully`)
+        console.log(`[✨ AI | i] ${model} pulled successfully`);
         return {
           success: true,
           response: `✅ Pulled ${model} successfully, please retry the command.`,
-        }
+        };
       }
     }
     return {
       success: false,
       error: errorMsg,
-    }
+    };
   }
 }
 
-export default (bot: Telegraf<Context>) => {
+async function handleAiReply(ctx: TextContext, db: NodePgDatabase<typeof schema>, model: string, prompt: string, replyGenerating: Message, aiTemperature: number) {
+  const Strings = getStrings(languageCode(ctx));
+  const aiResponse = await getResponse(prompt, ctx, replyGenerating, model, aiTemperature);
+  if (!aiResponse) return;
+  if (!ctx.chat) return;
+  if (aiResponse.success && aiResponse.response) {
+    const modelHeader = `🤖 *${model}*  |  🌡️ *${aiTemperature}*\n\n`;
+    await rateLimiter.editMessageWithRetry(
+      ctx,
+      ctx.chat.id,
+      replyGenerating.message_id,
+      modelHeader + aiResponse.response,
+      { parse_mode: 'Markdown' }
+    );
+    return;
+  }
+  const error = Strings.unexpectedErr.replace("{error}", aiResponse.error);
+  await rateLimiter.editMessageWithRetry(
+    ctx,
+    ctx.chat.id,
+    replyGenerating.message_id,
+    error,
+    { parse_mode: 'Markdown' }
+  );
+}
+
+async function getUserWithStringsAndModel(ctx: Context, db: NodePgDatabase<typeof schema>): Promise<{ user: User; Strings: ReturnType<typeof getStrings>; languageCode: string; customAiModel: string; aiTemperature: number }> {
+  const userArr = await db.query.usersTable.findMany({ where: (fields, { eq }) => eq(fields.telegramId, String(ctx.from!.id)), limit: 1 });
+  let user = userArr[0];
+  if (!user) {
+    await ensureUserInDb(ctx, db);
+    const newUserArr = await db.query.usersTable.findMany({ where: (fields, { eq }) => eq(fields.telegramId, String(ctx.from!.id)), limit: 1 });
+    user = newUserArr[0];
+    const Strings = getStrings(user.languageCode);
+    return { user, Strings, languageCode: user.languageCode, customAiModel: user.customAiModel, aiTemperature: user.aiTemperature };
+  }
+  const Strings = getStrings(user.languageCode);
+  return { user, Strings, languageCode: user.languageCode, customAiModel: user.customAiModel, aiTemperature: user.aiTemperature };
+}
+
+export default (bot: Telegraf<Context>, db: NodePgDatabase<typeof schema>) => {
   const botName = bot.botInfo?.first_name && bot.botInfo?.last_name ? `${bot.botInfo.first_name} ${bot.botInfo.last_name}` : "Kowalski"
 
   bot.command(["ask", "think"], spamwatchMiddleware, async (ctx) => {
@@ -244,65 +396,75 @@ export default (bot: Telegraf<Context>) => {
     const model = isAsk ? flash_model : thinking_model
     const textCtx = ctx as TextContext
     const reply_to_message_id = replyToMessageId(textCtx)
-    const Strings = getStrings(languageCode(textCtx))
+    const { Strings, aiTemperature } = await getUserWithStringsAndModel(textCtx, db)
     const message = textCtx.message.text
     const author = ("@" + ctx.from?.username) || ctx.from?.first_name
 
     logger.logCmdStart(author, model === flash_model ? "ask" : "think")
 
     if (!process.env.ollamaApi) {
-      await ctx.reply(Strings.aiDisabled, {
+      await ctx.reply(Strings.ai.disabled, {
         parse_mode: 'Markdown',
         ...({ reply_to_message_id })
       })
       return
     }
 
-    const replyGenerating = await ctx.reply(Strings.askGenerating.replace("{model}", model), {
+    const fixedMsg = message.replace(/^\/(ask|think)(@\w+)?\s*/, "").trim()
+    if (fixedMsg.length < 1) {
+      await ctx.reply(Strings.ai.askNoMessage, {
+        parse_mode: 'Markdown',
+        ...({ reply_to_message_id })
+      })
+      return
+    }
+
+    const replyGenerating = await ctx.reply(Strings.ai.askGenerating.replace("{model}", model), {
       parse_mode: 'Markdown',
       ...({ reply_to_message_id })
     })
 
-    const fixedMsg = message.replace(/\/(ask|think) /, "")
-    if (fixedMsg.length < 1) {
-      await ctx.reply(Strings.askNoMessage, {
+    logger.logPrompt(fixedMsg)
+
+    const prompt = sanitizeForJson(await usingSystemPrompt(textCtx, db, botName))
+    await handleAiReply(textCtx, db, model, prompt, replyGenerating, aiTemperature)
+  })
+
+  bot.command(["ai"], spamwatchMiddleware, async (ctx) => {
+    if (!ctx.message || !('text' in ctx.message)) return
+    const textCtx = ctx as TextContext
+    const reply_to_message_id = replyToMessageId(textCtx)
+    const { Strings, customAiModel, aiTemperature } = await getUserWithStringsAndModel(textCtx, db)
+    const message = textCtx.message.text
+    const author = ("@" + ctx.from?.username) || ctx.from?.first_name
+
+    logger.logCmdStart(author, "ask")
+
+    if (!process.env.ollamaApi) {
+      await ctx.reply(Strings.ai.disabled, {
         parse_mode: 'Markdown',
         ...({ reply_to_message_id })
       })
       return
     }
 
-    logger.logPrompt(fixedMsg)
-
-    const prompt = sanitizeForJson(
-`You are a plaintext-only, helpful assistant called ${botName}.
-Current Date/Time (UTC): ${new Date().toLocaleString()}
-
----
-
-Respond to the user's message:
-${fixedMsg}`)
-    const aiResponse = await getResponse(prompt, textCtx, replyGenerating, model)
-    if (!aiResponse) return
-
-    if (!ctx.chat) return
-    if (aiResponse.success && aiResponse.response) {
-      await rateLimiter.editMessageWithRetry(
-        ctx,
-        ctx.chat.id,
-        replyGenerating.message_id,
-        aiResponse.response,
-        { parse_mode: 'Markdown' }
-      )
+    const fixedMsg = message.replace(/^\/ai(@\w+)?\s*/, "").trim()
+    if (fixedMsg.length < 1) {
+      await ctx.reply(Strings.ai.askNoMessage, {
+        parse_mode: 'Markdown',
+        ...({ reply_to_message_id })
+      })
       return
     }
-    const error = Strings.unexpectedErr.replace("{error}", aiResponse.error)
-    await rateLimiter.editMessageWithRetry(
-      ctx,
-      ctx.chat.id,
-      replyGenerating.message_id,
-      error,
-      { parse_mode: 'Markdown' }
-    )
+
+    const replyGenerating = await ctx.reply(Strings.ai.askGenerating.replace("{model}", customAiModel), {
+      parse_mode: 'Markdown',
+      ...({ reply_to_message_id })
+    })
+
+    logger.logPrompt(fixedMsg)
+
+    const prompt = sanitizeForJson(await usingSystemPrompt(textCtx, db, botName))
+    await handleAiReply(textCtx, db, customAiModel, prompt, replyGenerating, aiTemperature)
   })
 }