Unverified Commit 0b50c525 authored by Bowen Liang's avatar Bowen Liang Committed by GitHub

feat: support error correction and border size in qrcode tool (#2731)

parent 8ba38e8e
......@@ -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()
......@@ -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
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment