From 4d4e222d20ee24bfc551175bed092893d64955fc Mon Sep 17 00:00:00 2001
From: Shigma <shigma10826@gmail.com>
Date: Sun, 4 Feb 2024 00:25:37 +0800
Subject: [PATCH] fix: escape for prompt and uc, fix #219

---
 package.json        |  2 +-
 src/config.ts       |  6 ++++--
 src/index.ts        |  8 +++-----
 tests/index.spec.ts | 14 ++++++++++++++
 4 files changed, 22 insertions(+), 8 deletions(-)
 create mode 100644 tests/index.spec.ts

diff --git a/package.json b/package.json
index d16d0c6..6ad5d30 100644
--- a/package.json
+++ b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "koishi-plugin-novelai",
   "description": "Generate images by diffusion models",
-  "version": "1.20.0",
+  "version": "1.20.1",
   "main": "lib/index.js",
   "typings": "lib/index.d.ts",
   "files": [
diff --git a/src/config.ts b/src/config.ts
index e4f0a7d..41fc906 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -387,7 +387,7 @@ export function parseForbidden(input: string) {
 
 const backslash = /@@__BACKSLASH__@@/g
 
-export function parseInput(session: Session, input: string, config: Config, override: boolean, addDefault: boolean): string[] {
+export function parseInput(session: Session, input: string, config: Config, override: boolean): string[] {
   if (!input) {
     return [
       null,
@@ -401,6 +401,8 @@ export function parseInput(session: Session, input: string, config: Config, over
     .replace(/，/g, ',')
     .replace(/（/g, '(')
     .replace(/）/g, ')')
+    .replace(/《/g, '<')
+    .replace(/》/g, '>')
 
   if (config.type === 'sd-webui') {
     input = input
@@ -470,7 +472,7 @@ export function parseInput(session: Session, input: string, config: Config, over
   if (!override) {
     appendToList(positive, session.resolve(config.basePrompt))
     appendToList(negative, session.resolve(config.negativePrompt))
-    if (addDefault) appendToList(positive, session.resolve(config.defaultPrompt))
+    if (config.defaultPromptSw) appendToList(positive, session.resolve(config.defaultPrompt))
   }
 
   return [null, positive.join(', '), negative.join(', ')]
diff --git a/src/index.ts b/src/index.ts
index 5f45c59..0e89a6f 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -188,9 +188,7 @@ export function apply(ctx: Context, config: Config) {
         }
       }
 
-      const [errPath, prompt, uc] = parseInput(
-        session, input, config, options.override, config.defaultPromptSw,
-      )
+      const [errPath, prompt, uc] = parseInput(session, input, config, options.override)
       if (errPath) return session.text(errPath)
 
       let token: string
@@ -486,9 +484,9 @@ export function apply(ctx: Context, config: Config) {
             }
           }
           result.children.push(h('message', attrs, lines.join('\n')))
-          result.children.push(h('message', attrs, `prompt = ${finalPrompt}`))
+          result.children.push(h('message', attrs, `prompt = ${h.escape(finalPrompt)}`))
           if (output === 'verbose') {
-            result.children.push(h('message', attrs, `undesired = ${uc}`))
+            result.children.push(h('message', attrs, `undesired = ${h.escape(uc)}`))
           }
           result.children.push(h('message', attrs, h.image(dataUrl)))
           return result
diff --git a/tests/index.spec.ts b/tests/index.spec.ts
new file mode 100644
index 0000000..0546715
--- /dev/null
+++ b/tests/index.spec.ts
@@ -0,0 +1,14 @@
+import { describe, test } from 'node:test'
+import * as novelai from '../src'
+import { Context } from 'koishi'
+import mock from '@koishijs/plugin-mock'
+
+describe('koishi-plugin-novelai', () => {
+  test('parse input', () => {
+    const ctx = new Context()
+    ctx.plugin(mock)
+    const session = ctx.bots[0].session({})
+    const fork = ctx.plugin(novelai)
+    console.log(novelai.parseInput(session, '<lora:skr2:1>,1girl', fork.config, false))
+  })
+})