From 0b50c525cf94b4f3a38a25077445733f6dbf69bf Mon Sep 17 00:00:00 2001
From: Bowen Liang <liangbowen@gf.com.cn>
Date: Thu, 7 Mar 2024 20:54:14 +0800
Subject: [PATCH] feat: support error correction and border size in qrcode tool
 (#2731)

---
 .../builtin/qrcode/tools/qrcode_generator.py  | 48 +++++++++++++---
 .../qrcode/tools/qrcode_generator.yaml        | 56 ++++++++++++++++++-
 2 files changed, 94 insertions(+), 10 deletions(-)

diff --git a/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.py b/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.py
index a86f17a999..8db5b1f8e8 100644
--- a/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.py
+++ b/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.py
@@ -2,14 +2,23 @@ import io
 import logging
 from typing import Any, Union
 
-import qrcode
+from qrcode.constants import ERROR_CORRECT_H, ERROR_CORRECT_L, ERROR_CORRECT_M, ERROR_CORRECT_Q
+from qrcode.image.base import BaseImage
 from qrcode.image.pure import PyPNGImage
+from qrcode.main import QRCode
 
 from core.tools.entities.tool_entities import ToolInvokeMessage
 from core.tools.tool.builtin_tool import BuiltinTool
 
 
 class QRCodeGeneratorTool(BuiltinTool):
+    error_correction_levels = {
+        'L': ERROR_CORRECT_L,  # <=7%
+        'M': ERROR_CORRECT_M,  # <=15%
+        'Q': ERROR_CORRECT_Q,  # <=25%
+        'H': ERROR_CORRECT_H,  # <=30%
+    }
+
     def _invoke(self,
                 user_id: str,
                 tool_parameters: dict[str, Any],
@@ -17,19 +26,44 @@ class QRCodeGeneratorTool(BuiltinTool):
         """
             invoke tools
         """
-        # get expression
+        # get text content
         content = tool_parameters.get('content', '')
         if not content:
             return self.create_text_message('Invalid parameter content')
 
+        # get border size
+        border = tool_parameters.get('border', 0)
+        if border < 0 or border > 100:
+            return self.create_text_message('Invalid parameter border')
+
+        # get error_correction
+        error_correction = tool_parameters.get('error_correction', '')
+        if error_correction not in self.error_correction_levels.keys():
+            return self.create_text_message('Invalid parameter error_correction')
+
         try:
-            img = qrcode.make(data=content, image_factory=PyPNGImage)
-            byte_stream = io.BytesIO()
-            img.save(byte_stream)
-            byte_array = byte_stream.getvalue()
-            return self.create_blob_message(blob=byte_array,
+            image = self._generate_qrcode(content, border, error_correction)
+            image_bytes = self._image_to_byte_array(image)
+            return self.create_blob_message(blob=image_bytes,
                                             meta={'mime_type': 'image/png'},
                                             save_as=self.VARIABLE_KEY.IMAGE.value)
         except Exception:
             logging.exception(f'Failed to generate QR code for content: {content}')
             return self.create_text_message('Failed to generate QR code')
+
+    def _generate_qrcode(self, content: str, border: int, error_correction: str) -> BaseImage:
+        qr = QRCode(
+            image_factory=PyPNGImage,
+            error_correction=self.error_correction_levels.get(error_correction),
+            border=border,
+        )
+        qr.add_data(data=content)
+        qr.make(fit=True)
+        img = qr.make_image()
+        return img
+
+    @staticmethod
+    def _image_to_byte_array(image: BaseImage) -> bytes:
+        byte_stream = io.BytesIO()
+        image.save(byte_stream)
+        return byte_stream.getvalue()
diff --git a/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.yaml b/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.yaml
index ca562ac094..8c8b8c449a 100644
--- a/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.yaml
+++ b/api/core/tools/provider/builtin/qrcode/tools/qrcode_generator.yaml
@@ -2,9 +2,9 @@ identity:
   name: qrcode_generator
   author: Bowen Liang
   label:
-    en_US: QR Code Generator
-    zh_Hans: 二维码生成器
-    pt_BR: QR Code Generator
+    en_US: Generate QR Code
+    zh_Hans: 生成二维码
+    pt_BR: Generate QR Code
 description:
   human:
     en_US: A tool for generating QR code image
@@ -24,3 +24,53 @@ parameters:
       zh_Hans: 二维码文本内容
       pt_BR: 二维码文本内容
     form: llm
+  - name: error_correction
+    type: select
+    required: true
+    default: M
+    label:
+      en_US: Error Correction
+      zh_Hans: 容错等级
+      pt_BR: Error Correction
+    human_description:
+      en_US: Error Correction in L, M, Q or H, from low to high, the bigger size of generated QR code with the better error correction effect
+      zh_Hans: 容错等级，可设置为低、中、偏高或高，从低到高，生成的二维码越大且容错效果越好
+      pt_BR: Error Correction in L, M, Q or H, from low to high, the bigger size of generated QR code with the better error correction effect
+    options:
+      - value: L
+        label:
+          en_US: Low
+          zh_Hans: 低
+          pt_BR: Low
+      - value: M
+        label:
+          en_US: Medium
+          zh_Hans: 中
+          pt_BR: Medium
+      - value: Q
+        label:
+          en_US: Quartile
+          zh_Hans: 偏高
+          pt_BR: Quartile
+      - value: H
+        label:
+          en_US: High
+          zh_Hans: 高
+          pt_BR: High
+    form: form
+  - name: border
+    type: number
+    required: true
+    default: 2
+    min: 0
+    max: 100
+    label:
+      en_US: border size
+      zh_Hans: 边框粗细
+      pt_BR: border size
+    human_description:
+      en_US: border size（default to 2）
+      zh_Hans: 边框粗细的格数（默认为2）
+      pt_BR: border size（default to 2）
+    llm: border size, default to 2
+    form: form