From c30659e526fff60b6f350971eb976e25bb3017ae Mon Sep 17 00:00:00 2001
From: tetsuya-ki <64536338+tetsuya-ki@users.noreply.github.com>
Date: Fri, 8 Mar 2024 23:25:17 +0900
Subject: [PATCH] =?UTF-8?q?enhance:=20aichat=E6=A9=9F=E8=83=BD(=E7=94=BB?=
 =?UTF-8?q?=E5=83=8F=E5=AF=BE=E5=BF=9C)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/modules/aichat/index.ts | 75 +++++++++++++++++++++++++++++++------
 src/utils/url2base64.ts     | 16 ++++++++
 2 files changed, 79 insertions(+), 12 deletions(-)
 create mode 100644 src/utils/url2base64.ts

diff --git a/src/modules/aichat/index.ts b/src/modules/aichat/index.ts
index 5c6dfbc..deb4466 100644
--- a/src/modules/aichat/index.ts
+++ b/src/modules/aichat/index.ts
@@ -3,6 +3,7 @@ import Module from '@/module.js';
 import serifs from '@/serifs.js';
 import Message from '@/message.js';
 import config from '@/config.js';
+import urlToBase64 from '@/utils/url2base64.js';
 import got from 'got';
 
 type AiChat = {
@@ -11,7 +12,12 @@ type AiChat = {
 	api: string;
 	key: string;
 };
+type Base64Image = {
+	type: string;
+	base64: string;
+};
 const GEMINI_API = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent';
+const GEMINI_VISION_API = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-pro-vision:generateContent';
 
 export default class extends Module {
 	public readonly name = 'aichat';
@@ -24,21 +30,34 @@ export default class extends Module {
 	}
 
 	@bindThis
-	private async genTextByGemini(aiChat: AiChat) {
+	private async genTextByGemini(aiChat: AiChat, image:Base64Image|null) {
 		this.log('Generate Text By Gemini...');
-		var options = {
+		let parts: ({ text: string; inline_data?: undefined; } | { inline_data: { mime_type: string; data: string; }; text?: undefined; })[];
+		if (image === null) {
+			// 画像がない場合、メッセージのみで問い合わせ
+			parts = [{text: aiChat.prompt + aiChat.question}];
+		} else {
+			// 画像が存在する場合、画像を添付して問い合わせ
+			parts = [
+				{ text: aiChat.prompt + aiChat.question },
+				{
+					inline_data: {
+						mime_type: image.type,
+						data: image.base64,
+					},
+				},
+			];
+		}
+		let options = {
 			url: aiChat.api,
 			searchParams: {
 				key: aiChat.key,
 			},
 			json: {
-				contents: [{
-					parts:[{
-						text: aiChat.prompt + aiChat.question
-					}]
-				}]
+				contents: {parts: parts}
 			},
 		};
+		this.log(JSON.stringify(options));
 		let res_data:any = null;
 		try {
 			res_data = await got.post(options,
@@ -60,7 +79,35 @@ export default class extends Module {
 		} catch (err: unknown) {
 			this.log('Error By Call Gemini');
 			if (err instanceof Error) {
-				this.log(`${err.name}\n${err.message}`);
+				this.log(`${err.name}\n${err.message}\n${err.stack}`);
+			}
+		}
+		return null;
+	}
+
+	@bindThis
+	private async note2base64Image(notesId: string) {
+		const noteData = await this.ai.api('notes/show', { noteId: notesId });
+		let fileType: string | undefined,thumbnailUrl: string | undefined;
+		if (noteData !== null && noteData.hasOwnProperty('files')) {
+			if (noteData.files.length > 0) {
+				if (noteData.files[0].hasOwnProperty('type')) {
+					fileType = noteData.files[0].type;
+				}
+				if (noteData.files[0].hasOwnProperty('thumbnailUrl')) {
+					thumbnailUrl = noteData.files[0].thumbnailUrl;
+				}
+			}
+			if (fileType !== undefined && thumbnailUrl !== undefined) {
+				try {
+					const image = await urlToBase64(thumbnailUrl);
+					const base64Image:Base64Image = {type: fileType, base64: image};
+					return base64Image;
+				} catch (err: unknown) {
+					if (err instanceof Error) {
+						this.log(`${err.name}\n${err.message}\n${err.stack}`);
+					}
+				}
 			}
 		}
 		return null;
@@ -88,8 +135,8 @@ export default class extends Module {
 							.replace(kigo + type, '')
 							.trim();
 
-		let text;
-		let prompt = '';
+		let text:string, aiChat:AiChat;
+		let prompt:string = '';
 		if (config.prompt) {
 			prompt = config.prompt;
 		}
@@ -99,13 +146,17 @@ export default class extends Module {
 					msg.reply(serifs.aichat.nothing(type));
 					return false;
 				}
-				const aiChat = {
+				const base64Image:Base64Image|null = await this.note2base64Image(msg.id);
+				aiChat = {
 					question: question,
 					prompt: prompt,
 					api: GEMINI_API,
 					key: config.geminiProApiKey
 				};
-				text = await this.genTextByGemini(aiChat);
+				if (base64Image !== null) {
+					aiChat.api = GEMINI_VISION_API;
+				}
+				text = await this.genTextByGemini(aiChat, base64Image);
 				break;
 			default:
 				msg.reply(serifs.aichat.nothing(type));
diff --git a/src/utils/url2base64.ts b/src/utils/url2base64.ts
new file mode 100644
index 0000000..7350a66
--- /dev/null
+++ b/src/utils/url2base64.ts
@@ -0,0 +1,16 @@
+import log from '@/utils/log.js';
+import got from 'got';
+
+export default async function(url: string): Promise<string> {
+	try {
+			const buffer = await got(url).buffer();
+			const base64Image = buffer.toString('base64');
+			return base64Image;
+	} catch (err: unknown) {
+		log('Error in url2base64');
+		if (err instanceof Error) {
+			log(`${err.name}\n${err.message}\n${err.stack}`);
+		}
+		throw err;
+	}
+}