From 71bd84bab718f0891ed15928dc89c5ff88e98613 Mon Sep 17 00:00:00 2001
From: ZhangWei-KUMO <kanseefoil@gmail.com>
Date: Sun, 15 Dec 2024 20:41:40 +0800
Subject: [PATCH] add ui

---
 db/agent.js    |  2 +-
 util/gemini.js | 16 ++++++++++------
 util/handle.js | 29 +++++++++++++++++++----------
 3 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/db/agent.js b/db/agent.js
index 2dd3c27..072591e 100644
--- a/db/agent.js
+++ b/db/agent.js
@@ -28,7 +28,7 @@ const db = new sqlite3.Database(dbFile, (err) => {
                         console.error("Error checking agentConfig table:", err.message);
                     } else if (row.count === 0) {
                         const defaultEmailConfig = {
-                            model: 'gemini-1.5-pro',
+                            model: 'gemini-1.5-flash',
                             prompt: '你是一个机器人',
                             tasks: '其他',
                             news: 0,
diff --git a/util/gemini.js b/util/gemini.js
index 5c2fda0..abce1b1 100644
--- a/util/gemini.js
+++ b/util/gemini.js
@@ -1,18 +1,22 @@
 /* eslint-disable no-undef */
 import {GoogleGenerativeAI} from '@google/generative-ai';
 import 'dotenv/config'
+import {getAgentConfig} from '../db/agent.js'
+import { getConfig } from '../db/config.js';
+let agent = await getAgentConfig();
+let config = await getConfig();
+const genAI = new GoogleGenerativeAI(config.geminiApiKey || configprocess.env.GEMINI_API_KEY);
+const model = genAI.getGenerativeModel({ model:agent.model || "gemini-1.5-flash" });
 
-const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY);
-const model = genAI.getGenerativeModel({ model: "gemini-1.5-flash" });
-
-export const recgonizeImage = async (prompt) => {
+export const recgonizeImage = async (buffer,question) => {
     const image = {
         inlineData: {
-            data: Buffer.from(fs.readFileSync("cookie.jpg")).toString("base64"),
+            // data: Buffer.from(fs.readFileSync("cookie.jpg")).toString("base64"),
+            data: buffer.toString("base64"),
             mimeType: "image/png",
         }};
 
-    const result = await model.generateContent([prompt, image]);
+    const result = await model.generateContent([question, image]);
     return result.response.text()
 }
 
diff --git a/util/handle.js b/util/handle.js
index 8f1cd06..f11b99f 100644
--- a/util/handle.js
+++ b/util/handle.js
@@ -8,6 +8,7 @@ import fs from 'fs';
 import process from 'process';
 import schedule from 'node-schedule';
 import {getNews} from './group.js'
+import {recgonizeImage} from './gemini.js'
 
 export const singleChat = async (talkerId,listenerId,text) => {
     if(talkerId!==listenerId && talkerId!=='weixin' && text!==''){
@@ -260,14 +261,22 @@ export const handlePush = async (rooms) => {
       });
     }
 }
-
-export const handleImage = async (message,talkerId) => {
-    try{
-        // const imageFileBox = await message.toFileBox();
-        // let image = await imageFileBox.toFile();
-        await sendMessage(talkerId, "好的，我先看下");
-    }catch(e){
-        await sendMessage(talkerId, e)
-        return
+// 处理图片，该函数的触发条件是指用户在当前文本信息之前发送了一张图片
+export const handleImage = async (message, talkerId,question) => {
+  try{
+    const imageFileBox = await message.toFileBox();
+    let image = await imageFileBox.toBuffer();
+    let mimeType = imageFileBox.mimeType
+    if (!mimeType) {
+       await sendMessage(talkerId, '无法获取图片 mimeType，请检查图片格式');
+       return
     }
-}
\ No newline at end of file
+    let res = await recgonizeImage(image, question)
+    await sendMessage(talkerId, res);
+}catch(e){
+    // 修改这里，提取错误信息
+    let errorMessage = `图片处理出错: ${e.message || String(e)}`;
+    await sendMessage(talkerId, errorMessage);
+    return
+}
+} 
\ No newline at end of file