mirror of https://github.com/Mai-with-u/MaiBot.git
490 lines
20 KiB
Python
490 lines
20 KiB
Python
"""
|
||
豆包图片生成插件
|
||
|
||
基于火山引擎豆包模型的AI图片生成插件。
|
||
|
||
功能特性：
|
||
- 智能LLM判定：根据聊天内容智能判断是否需要生成图片
|
||
- 高质量图片生成：使用豆包Seed Dream模型生成图片
|
||
- 结果缓存：避免重复生成相同内容的图片
|
||
- 配置验证：自动验证和修复配置文件
|
||
- 参数验证：完整的输入参数验证和错误处理
|
||
- 多尺寸支持：支持多种图片尺寸生成
|
||
|
||
包含组件：
|
||
- 图片生成Action - 根据描述使用火山引擎API生成图片
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import urllib.request
|
||
import urllib.error
|
||
import base64
|
||
import traceback
|
||
from typing import List, Tuple, Type, Optional
|
||
|
||
# 导入新插件系统
|
||
from src.plugin_system.base.base_plugin import BasePlugin
|
||
from src.plugin_system.base.base_plugin import register_plugin
|
||
from src.plugin_system.base.base_action import BaseAction
|
||
from src.plugin_system.base.component_types import ComponentInfo, ActionActivationType, ChatMode
|
||
from src.plugin_system.base.config_types import ConfigField
|
||
from src.common.logger import get_logger
|
||
|
||
logger = get_logger("doubao_pic_plugin")
|
||
|
||
|
||
# ===== Action组件 =====
|
||
|
||
|
||
class DoubaoImageGenerationAction(BaseAction):
|
||
"""豆包图片生成Action - 根据描述使用火山引擎API生成图片"""
|
||
|
||
# 激活设置
|
||
focus_activation_type = ActionActivationType.LLM_JUDGE # Focus模式使用LLM判定，精确理解需求
|
||
normal_activation_type = ActionActivationType.KEYWORD # Normal模式使用关键词激活，快速响应
|
||
mode_enable = ChatMode.ALL
|
||
parallel_action = True
|
||
|
||
# 动作基本信息
|
||
action_name = "doubao_image_generation"
|
||
action_description = (
|
||
"可以根据特定的描述，生成并发送一张图片，如果没提供描述，就根据聊天内容生成,你可以立刻画好，不用等待"
|
||
)
|
||
|
||
# 关键词设置（用于Normal模式）
|
||
activation_keywords = ["画", "绘制", "生成图片", "画图", "draw", "paint", "图片生成"]
|
||
keyword_case_sensitive = False
|
||
|
||
# LLM判定提示词（用于Focus模式）
|
||
llm_judge_prompt = """
|
||
判定是否需要使用图片生成动作的条件：
|
||
1. 用户明确要求画图、生成图片或创作图像
|
||
2. 用户描述了想要看到的画面或场景
|
||
3. 对话中提到需要视觉化展示某些概念
|
||
4. 用户想要创意图片或艺术作品
|
||
|
||
适合使用的情况：
|
||
- "画一张..."、"画个..."、"生成图片"
|
||
- "我想看看...的样子"
|
||
- "能画出...吗"
|
||
- "创作一幅..."
|
||
|
||
绝对不要使用的情况：
|
||
1. 纯文字聊天和问答
|
||
2. 只是提到"图片"、"画"等词但不是要求生成
|
||
3. 谈论已存在的图片或照片
|
||
4. 技术讨论中提到绘图概念但无生成需求
|
||
5. 用户明确表示不需要图片时
|
||
"""
|
||
|
||
# 动作参数定义
|
||
action_parameters = {
|
||
"description": "图片描述，输入你想要生成并发送的图片的描述，必填",
|
||
"size": "图片尺寸，例如 '1024x1024' (可选, 默认从配置或 '1024x1024')",
|
||
}
|
||
|
||
# 动作使用场景
|
||
action_require = [
|
||
"当有人让你画东西时使用，你可以立刻画好，不用等待",
|
||
"当有人要求你生成并发送一张图片时使用",
|
||
"当有人让你画一张图时使用",
|
||
]
|
||
|
||
# 关联类型
|
||
associated_types = ["image", "text"]
|
||
|
||
# 简单的请求缓存，避免短时间内重复请求
|
||
_request_cache = {}
|
||
_cache_max_size = 10
|
||
|
||
async def execute(self) -> Tuple[bool, Optional[str]]:
|
||
"""执行图片生成动作"""
|
||
logger.info(f"{self.log_prefix} 执行豆包图片生成动作")
|
||
|
||
# 配置验证
|
||
http_base_url = self.api.get_config("api.base_url")
|
||
http_api_key = self.api.get_config("api.volcano_generate_api_key")
|
||
|
||
if not (http_base_url and http_api_key):
|
||
error_msg = "抱歉，图片生成功能所需的HTTP配置（如API地址或密钥）不完整，无法提供服务。"
|
||
await self.send_text(error_msg)
|
||
logger.error(f"{self.log_prefix} HTTP调用配置缺失: base_url 或 volcano_generate_api_key.")
|
||
return False, "HTTP配置不完整"
|
||
|
||
# API密钥验证
|
||
if http_api_key == "YOUR_DOUBAO_API_KEY_HERE":
|
||
error_msg = "图片生成功能尚未配置，请设置正确的API密钥。"
|
||
await self.send_text(error_msg)
|
||
logger.error(f"{self.log_prefix} API密钥未配置")
|
||
return False, "API密钥未配置"
|
||
|
||
# 参数验证
|
||
description = self.action_data.get("description")
|
||
if not description or not description.strip():
|
||
logger.warning(f"{self.log_prefix} 图片描述为空，无法生成图片。")
|
||
await self.send_text("你需要告诉我想要画什么样的图片哦~ 比如说'画一只可爱的小猫'")
|
||
return False, "图片描述为空"
|
||
|
||
# 清理和验证描述
|
||
description = description.strip()
|
||
if len(description) > 1000: # 限制描述长度
|
||
description = description[:1000]
|
||
logger.info(f"{self.log_prefix} 图片描述过长，已截断")
|
||
|
||
# 获取配置
|
||
default_model = self.api.get_config("generation.default_model", "doubao-seedream-3-0-t2i-250415")
|
||
image_size = self.action_data.get("size", self.api.get_config("generation.default_size", "1024x1024"))
|
||
|
||
# 验证图片尺寸格式
|
||
if not self._validate_image_size(image_size):
|
||
logger.warning(f"{self.log_prefix} 无效的图片尺寸: {image_size}，使用默认值")
|
||
image_size = "1024x1024"
|
||
|
||
# 检查缓存
|
||
cache_key = self._get_cache_key(description, default_model, image_size)
|
||
if cache_key in self._request_cache:
|
||
cached_result = self._request_cache[cache_key]
|
||
logger.info(f"{self.log_prefix} 使用缓存的图片结果")
|
||
await self.send_text("我之前画过类似的图片，用之前的结果~")
|
||
|
||
# 直接发送缓存的结果
|
||
send_success = await self._send_image(cached_result)
|
||
if send_success:
|
||
await self.send_text("图片已发送！")
|
||
return True, "图片已发送(缓存)"
|
||
else:
|
||
# 缓存失败，清除这个缓存项并继续正常流程
|
||
del self._request_cache[cache_key]
|
||
|
||
# 获取其他配置参数
|
||
guidance_scale_val = self._get_guidance_scale()
|
||
seed_val = self._get_seed()
|
||
watermark_val = self._get_watermark()
|
||
|
||
await self.send_text(
|
||
f"收到！正在为您生成关于 '{description}' 的图片，请稍候...（模型: {default_model}, 尺寸: {image_size}）"
|
||
)
|
||
|
||
try:
|
||
success, result = await asyncio.to_thread(
|
||
self._make_http_image_request,
|
||
prompt=description,
|
||
model=default_model,
|
||
size=image_size,
|
||
seed=seed_val,
|
||
guidance_scale=guidance_scale_val,
|
||
watermark=watermark_val,
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"{self.log_prefix} (HTTP) 异步请求执行失败: {e!r}", exc_info=True)
|
||
traceback.print_exc()
|
||
success = False
|
||
result = f"图片生成服务遇到意外问题: {str(e)[:100]}"
|
||
|
||
if success:
|
||
image_url = result
|
||
# print(f"image_url: {image_url}")
|
||
# print(f"result: {result}")
|
||
logger.info(f"{self.log_prefix} 图片URL获取成功: {image_url[:70]}... 下载并编码.")
|
||
|
||
try:
|
||
encode_success, encode_result = await asyncio.to_thread(self._download_and_encode_base64, image_url)
|
||
except Exception as e:
|
||
logger.error(f"{self.log_prefix} (B64) 异步下载/编码失败: {e!r}", exc_info=True)
|
||
traceback.print_exc()
|
||
encode_success = False
|
||
encode_result = f"图片下载或编码时发生内部错误: {str(e)[:100]}"
|
||
|
||
if encode_success:
|
||
base64_image_string = encode_result
|
||
send_success = await self._send_image(base64_image_string)
|
||
if send_success:
|
||
# 缓存成功的结果
|
||
self._request_cache[cache_key] = base64_image_string
|
||
self._cleanup_cache()
|
||
|
||
await self.send_message_by_expressor("图片已发送！")
|
||
return True, "图片已成功生成并发送"
|
||
else:
|
||
print(f"send_success: {send_success}")
|
||
await self.send_message_by_expressor("图片已处理为Base64，但发送失败了。")
|
||
return False, "图片发送失败 (Base64)"
|
||
else:
|
||
await self.send_message_by_expressor(f"获取到图片URL，但在处理图片时失败了：{encode_result}")
|
||
return False, f"图片处理失败(Base64): {encode_result}"
|
||
else:
|
||
error_message = result
|
||
await self.send_message_by_expressor(f"哎呀，生成图片时遇到问题：{error_message}")
|
||
return False, f"图片生成失败: {error_message}"
|
||
|
||
def _get_guidance_scale(self) -> float:
|
||
"""获取guidance_scale配置值"""
|
||
guidance_scale_input = self.api.get_config("generation.default_guidance_scale", 2.5)
|
||
try:
|
||
return float(guidance_scale_input)
|
||
except (ValueError, TypeError):
|
||
logger.warning(f"{self.log_prefix} default_guidance_scale 值无效，使用默认值 2.5")
|
||
return 2.5
|
||
|
||
def _get_seed(self) -> int:
|
||
"""获取seed配置值"""
|
||
seed_config_value = self.api.get_config("generation.default_seed")
|
||
if seed_config_value is not None:
|
||
try:
|
||
return int(seed_config_value)
|
||
except (ValueError, TypeError):
|
||
logger.warning(f"{self.log_prefix} default_seed 值无效，使用默认值 42")
|
||
return 42
|
||
|
||
def _get_watermark(self) -> bool:
|
||
"""获取watermark配置值"""
|
||
watermark_source = self.api.get_config("generation.default_watermark", True)
|
||
if isinstance(watermark_source, bool):
|
||
return watermark_source
|
||
elif isinstance(watermark_source, str):
|
||
return watermark_source.lower() == "true"
|
||
else:
|
||
logger.warning(f"{self.log_prefix} default_watermark 值无效，使用默认值 True")
|
||
return True
|
||
|
||
async def _send_image(self, base64_image: str) -> bool:
|
||
"""发送图片"""
|
||
try:
|
||
# 使用聊天流信息确定发送目标
|
||
chat_stream = self.api.get_service("chat_stream")
|
||
if not chat_stream:
|
||
logger.error(f"{self.log_prefix} 没有可用的聊天流发送图片")
|
||
return False
|
||
|
||
if chat_stream.group_info:
|
||
# 群聊
|
||
return await self.api.send_message_to_target(
|
||
message_type="image",
|
||
content=base64_image,
|
||
platform=chat_stream.platform,
|
||
target_id=str(chat_stream.group_info.group_id),
|
||
is_group=True,
|
||
display_message="发送生成的图片",
|
||
)
|
||
else:
|
||
# 私聊
|
||
return await self.api.send_message_to_target(
|
||
message_type="image",
|
||
content=base64_image,
|
||
platform=chat_stream.platform,
|
||
target_id=str(chat_stream.user_info.user_id),
|
||
is_group=False,
|
||
display_message="发送生成的图片",
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"{self.log_prefix} 发送图片时出错: {e}")
|
||
return False
|
||
|
||
@classmethod
|
||
def _get_cache_key(cls, description: str, model: str, size: str) -> str:
|
||
"""生成缓存键"""
|
||
return f"{description[:100]}|{model}|{size}"
|
||
|
||
@classmethod
|
||
def _cleanup_cache(cls):
|
||
"""清理缓存，保持大小在限制内"""
|
||
if len(cls._request_cache) > cls._cache_max_size:
|
||
keys_to_remove = list(cls._request_cache.keys())[: -cls._cache_max_size // 2]
|
||
for key in keys_to_remove:
|
||
del cls._request_cache[key]
|
||
|
||
def _validate_image_size(self, image_size: str) -> bool:
|
||
"""验证图片尺寸格式"""
|
||
try:
|
||
width, height = map(int, image_size.split("x"))
|
||
return 100 <= width <= 10000 and 100 <= height <= 10000
|
||
except (ValueError, TypeError):
|
||
return False
|
||
|
||
def _download_and_encode_base64(self, image_url: str) -> Tuple[bool, str]:
|
||
"""下载图片并将其编码为Base64字符串"""
|
||
logger.info(f"{self.log_prefix} (B64) 下载并编码图片: {image_url[:70]}...")
|
||
try:
|
||
with urllib.request.urlopen(image_url, timeout=30) as response:
|
||
if response.status == 200:
|
||
image_bytes = response.read()
|
||
base64_encoded_image = base64.b64encode(image_bytes).decode("utf-8")
|
||
logger.info(f"{self.log_prefix} (B64) 图片下载编码完成. Base64长度: {len(base64_encoded_image)}")
|
||
return True, base64_encoded_image
|
||
else:
|
||
error_msg = f"下载图片失败 (状态: {response.status})"
|
||
logger.error(f"{self.log_prefix} (B64) {error_msg} URL: {image_url}")
|
||
return False, error_msg
|
||
except Exception as e:
|
||
logger.error(f"{self.log_prefix} (B64) 下载或编码时错误: {e!r}", exc_info=True)
|
||
traceback.print_exc()
|
||
return False, f"下载或编码图片时发生错误: {str(e)[:100]}"
|
||
|
||
def _make_http_image_request(
|
||
self, prompt: str, model: str, size: str, seed: int, guidance_scale: float, watermark: bool
|
||
) -> Tuple[bool, str]:
|
||
"""发送HTTP请求生成图片"""
|
||
base_url = self.api.get_config("api.base_url")
|
||
generate_api_key = self.api.get_config("api.volcano_generate_api_key")
|
||
|
||
endpoint = f"{base_url.rstrip('/')}/images/generations"
|
||
|
||
payload_dict = {
|
||
"model": model,
|
||
"prompt": prompt,
|
||
"response_format": "url",
|
||
"size": size,
|
||
"guidance_scale": guidance_scale,
|
||
"watermark": watermark,
|
||
"seed": seed,
|
||
"api-key": generate_api_key,
|
||
}
|
||
|
||
data = json.dumps(payload_dict).encode("utf-8")
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"Accept": "application/json",
|
||
"Authorization": f"Bearer {generate_api_key}",
|
||
}
|
||
|
||
logger.info(f"{self.log_prefix} (HTTP) 发起图片请求: {model}, Prompt: {prompt[:30]}... To: {endpoint}")
|
||
|
||
req = urllib.request.Request(endpoint, data=data, headers=headers, method="POST")
|
||
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=60) as response:
|
||
response_status = response.status
|
||
response_body_bytes = response.read()
|
||
response_body_str = response_body_bytes.decode("utf-8")
|
||
|
||
logger.info(f"{self.log_prefix} (HTTP) 响应: {response_status}. Preview: {response_body_str[:150]}...")
|
||
|
||
if 200 <= response_status < 300:
|
||
response_data = json.loads(response_body_str)
|
||
image_url = None
|
||
if (
|
||
isinstance(response_data.get("data"), list)
|
||
and response_data["data"]
|
||
and isinstance(response_data["data"][0], dict)
|
||
):
|
||
image_url = response_data["data"][0].get("url")
|
||
elif response_data.get("url"):
|
||
image_url = response_data.get("url")
|
||
|
||
if image_url:
|
||
logger.info(f"{self.log_prefix} (HTTP) 图片生成成功，URL: {image_url[:70]}...")
|
||
return True, image_url
|
||
else:
|
||
logger.error(f"{self.log_prefix} (HTTP) API成功但无图片URL")
|
||
return False, "图片生成API响应成功但未找到图片URL"
|
||
else:
|
||
logger.error(f"{self.log_prefix} (HTTP) API请求失败. 状态: {response.status}")
|
||
return False, f"图片API请求失败(状态码 {response.status})"
|
||
except Exception as e:
|
||
logger.error(f"{self.log_prefix} (HTTP) 图片生成时意外错误: {e!r}", exc_info=True)
|
||
traceback.print_exc()
|
||
return False, f"图片生成HTTP请求时发生意外错误: {str(e)[:100]}"
|
||
|
||
|
||
# ===== 插件主类 =====
|
||
|
||
|
||
@register_plugin
|
||
class DoubaoImagePlugin(BasePlugin):
|
||
"""豆包图片生成插件
|
||
|
||
基于火山引擎豆包模型的AI图片生成插件：
|
||
- 图片生成Action：根据描述使用火山引擎API生成图片
|
||
"""
|
||
|
||
# 插件基本信息
|
||
plugin_name = "doubao_pic_plugin"
|
||
plugin_description = "基于火山引擎豆包模型的AI图片生成插件"
|
||
plugin_version = "2.0.0"
|
||
plugin_author = "MaiBot开发团队"
|
||
enable_plugin = True
|
||
config_file_name = "config.toml"
|
||
|
||
# 配置节描述
|
||
config_section_descriptions = {
|
||
"plugin": "插件基本信息配置",
|
||
"api": "API相关配置，包含火山引擎API的访问信息",
|
||
"generation": "图片生成参数配置，控制生成图片的各种参数",
|
||
"cache": "结果缓存配置",
|
||
"components": "组件启用配置"
|
||
}
|
||
|
||
# 配置Schema定义
|
||
config_schema = {
|
||
"plugin": {
|
||
"name": ConfigField(type=str, default="doubao_pic_plugin", description="插件名称", required=True),
|
||
"version": ConfigField(type=str, default="2.0.0", description="插件版本号"),
|
||
"enabled": ConfigField(type=bool, default=True, description="是否启用插件"),
|
||
"description": ConfigField(
|
||
type=str,
|
||
default="基于火山引擎豆包模型的AI图片生成插件",
|
||
description="插件描述",
|
||
required=True
|
||
)
|
||
},
|
||
"api": {
|
||
"base_url": ConfigField(
|
||
type=str,
|
||
default="https://ark.cn-beijing.volces.com/api/v3",
|
||
description="API基础URL",
|
||
example="https://api.example.com/v1"
|
||
),
|
||
"volcano_generate_api_key": ConfigField(
|
||
type=str,
|
||
default="YOUR_DOUBAO_API_KEY_HERE",
|
||
description="火山引擎豆包API密钥",
|
||
required=True
|
||
)
|
||
},
|
||
"generation": {
|
||
"default_model": ConfigField(
|
||
type=str,
|
||
default="doubao-seedream-3-0-t2i-250415",
|
||
description="默认使用的文生图模型",
|
||
choices=["doubao-seedream-3-0-t2i-250415", "doubao-seedream-2-0-t2i"]
|
||
),
|
||
"default_size": ConfigField(
|
||
type=str,
|
||
default="1024x1024",
|
||
description="默认图片尺寸",
|
||
example="1024x1024",
|
||
choices=["1024x1024", "1024x1280", "1280x1024", "1024x1536", "1536x1024"]
|
||
),
|
||
"default_watermark": ConfigField(type=bool, default=True, description="是否默认添加水印"),
|
||
"default_guidance_scale": ConfigField(
|
||
type=float,
|
||
default=2.5,
|
||
description="模型指导强度，影响图片与提示的关联性",
|
||
example="2.0"
|
||
),
|
||
"default_seed": ConfigField(type=int, default=42, description="随机种子，用于复现图片")
|
||
},
|
||
"cache": {
|
||
"enabled": ConfigField(type=bool, default=True, description="是否启用请求缓存"),
|
||
"max_size": ConfigField(type=int, default=10, description="最大缓存数量")
|
||
},
|
||
"components": {
|
||
"enable_image_generation": ConfigField(type=bool, default=True, description="是否启用图片生成Action")
|
||
}
|
||
}
|
||
|
||
def get_plugin_components(self) -> List[Tuple[ComponentInfo, Type]]:
|
||
"""返回插件包含的组件列表"""
|
||
|
||
# 从配置获取组件启用状态
|
||
enable_image_generation = self.get_config("components.enable_image_generation", True)
|
||
|
||
components = []
|
||
|
||
# 添加图片生成Action
|
||
if enable_image_generation:
|
||
components.append((DoubaoImageGenerationAction.get_action_info(), DoubaoImageGenerationAction))
|
||
|
||
return components
|