From 8c9c098e255d35071d7d0423fef52a06feed8a3f Mon Sep 17 00:00:00 2001
From: ikechan8370 <geyinchibuaa@gmail.com>
Date: Sat, 25 Feb 2023 11:38:34 +0800
Subject: [PATCH] feat: support dalle api

---
 apps/draw.js    | 55 +++++++++++++++++++++++++++++++++++++++++++++++++
 apps/help.js    |  5 +++++
 package.json    |  1 +
 utils/common.js |  2 +-
 utils/dalle.js  | 19 +++++++++++++++++
 5 files changed, 81 insertions(+), 1 deletion(-)
 create mode 100644 apps/draw.js
 create mode 100644 utils/dalle.js

diff --git a/apps/draw.js b/apps/draw.js
new file mode 100644
index 0000000..3341ad0
--- /dev/null
+++ b/apps/draw.js
@@ -0,0 +1,55 @@
+import plugin from '../../../lib/plugins/plugin.js'
+import { segment } from 'oicq'
+import { createImage } from '../utils/dalle.js'
+import { makeForwardMsg} from "../utils/common.js";
+import _ from 'lodash'
+
+export class dalle extends plugin {
+  constructor (e) {
+    super({
+      name: 'ChatGPT-Plugin Dalle 绘图',
+      dsc: 'ChatGPT-Plugin基于OpenAI Dalle的绘图插件',
+      event: 'message',
+      priority: 500,
+      rule: [
+        {
+          reg: '#(chatgpt|ChatGPT|dalle|Dalle)(绘图|画图)',
+          fnc: 'draw'
+        }
+      ]
+    })
+  }
+
+  async draw (e) {
+    let splits = _.split(e.msg, '图', 2)
+    if (splits.length < 2) {
+      this.reply('请带上绘图要求')
+      return false
+    }
+    let rules = _.split(splits[1], '/')
+    let [prompt = '', num = '1', size = '512x512'] = rules.slice(0, 3)
+    if (['256x256', '512x512', '1024x1024'].indexOf(size) === -1) {
+      this.reply('大小不符合要求，必须是256x256/512x512/1024x1024中的一个')
+      return false
+    }
+    let priceMap = {
+      '1024x1024': 0.02,
+      '512x512': 0.018,
+      '256x256': 0.016
+    }
+    num = parseInt(num, 10)
+    await this.reply(`正在为您绘制大小为${size}的${num}张图片，预计消耗${priceMap[size] * num}美元余额，请稍候……`)
+    try {
+      let images = (await createImage(prompt, num, size)).map(image => segment.image(`base64://${image}`))
+      if (images.length > 1) {
+        this.reply(await makeForwardMsg(e, images, prompt))
+      } else {
+        this.reply(images[0], true)
+      }
+    } catch (err) {
+      this.reply(`绘图失败: err`, true)
+    }
+
+
+  }
+}
diff --git a/apps/help.js b/apps/help.js
index 438e6c3..5f692ff 100644
--- a/apps/help.js
+++ b/apps/help.js
@@ -24,6 +24,11 @@ let helpData = [
         icon: 'text',
         title: '#chatgpt文本模式',
         desc: '机器人以文本形式回答，默认选项'
+      },
+      {
+        icon: 'text',
+        title: '#chatgpt画图+prompt(/张数/图片大小)',
+        desc: '调用OpenAI Dalle API进行绘图，需要有API key并消耗余额。图片大小只能是256x256/512x512/1024x1024中的一个.默认为1张、512x512'
       }
 
     ]
diff --git a/package.json b/package.json
index 12b265d..0b7d793 100644
--- a/package.json
+++ b/package.json
@@ -8,6 +8,7 @@
     "delay": "^5.0.0",
     "keyv-file": "^0.2.0",
     "node-fetch": "^3.3.0",
+    "openai": "^3.1.0",
     "puppeteer-extra": "^3.3.4",
     "puppeteer-extra-plugin-recaptcha": "^3.6.6",
     "puppeteer-extra-plugin-stealth": "^2.11.1",
diff --git a/utils/common.js b/utils/common.js
index 110cc10..0a0de4c 100644
--- a/utils/common.js
+++ b/utils/common.js
@@ -20,7 +20,7 @@ export function escapeHtml (str) {
   return str.replace(/[&<>"'/]/g, (match) => htmlEntities[match])
 }
 
-export function randomString(length = 5) {
+export function randomString (length = 5) {
   let str = ''
   for (let i = 0; i < length; i++) {
     str += lodash.random(36).toString(36)
diff --git a/utils/dalle.js b/utils/dalle.js
new file mode 100644
index 0000000..2076049
--- /dev/null
+++ b/utils/dalle.js
@@ -0,0 +1,19 @@
+import { Configuration, OpenAIApi } from 'openai'
+import { Config } from './config.js'
+
+export async function createImage (prompt, n = 1, size = '512x512') {
+  const configuration = new Configuration({
+    apiKey: Config.apiKey
+  })
+  const openai = new OpenAIApi(configuration)
+  if (Config.debug) {
+    logger.info({ prompt, n, size })
+  }
+  const response = await openai.createImage({
+    prompt,
+    n,
+    size,
+    response_format: 'b64_json'
+  })
+  return response.data.data?.map(pic => pic.b64_json)
+}