""" NanoImage AI绘图插件 支持 OpenAI 格式的绘图 API，用户可自定义 URL、模型 ID、密钥 支持命令触发和 LLM 工具调用 """ import asyncio import os import tomllib import httpx import uuid import base64 import re from pathlib import Path from datetime import datetime from typing import List, Optional from urllib.parse import urlparse from loguru import logger from utils.plugin_base import PluginBase from utils.decorators import on_text_message from WechatHook import WechatHookClient class NanoImage(PluginBase): """NanoImage AI绘图插件""" description = "NanoImage AI绘图插件 - 支持 OpenAI 格式的绘图 API" author = "ShiHao" version = "1.0.0" def __init__(self): super().__init__() self.config = None self.images_dir = None async def async_init(self): """异步初始化""" config_path = Path(__file__).parent / "config.toml" logger.info(f"NanoImage 配置文件路径: {config_path}") with open(config_path, "rb") as f: self.config = tomllib.load(f) # 创建图片目录 self.images_dir = Path(__file__).parent / "images" self.images_dir.mkdir(exist_ok=True) logger.success(f"NanoImage AI插件初始化完成，base_url: {self.config['api']['base_url']}, 模型: {self.config['api']['model']}") async def generate_image(self, prompt: str) -> List[str]: """ 生成图像 Args: prompt: 提示词 Returns: 图片本地路径列表 """ api_config = self.config["api"] gen_config = self.config["generation"] max_retry = gen_config["max_retry_attempts"] simplify_on_fail = gen_config.get("simplify_prompt_on_fail", True) simplify_max_tags = gen_config.get("simplify_max_tags", 8) simplify_max_chars = gen_config.get("simplify_max_chars", 140) simplified_prompt = "" for attempt in range(max_retry): if attempt > 0: await asyncio.sleep(min(2 ** attempt, 10)) try: url = f"{api_config['base_url'].rstrip('/')}/v1/chat/completions" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {api_config['api_key']}" } current_prompt = prompt if simplified_prompt and attempt > 0: current_prompt = simplified_prompt # 添加明确的绘图指令前缀 full_prompt = f"请生成一张图片，并在回复中包含图片URL：{current_prompt}" stream_enabled = api_config.get("stream", True) payload = { "model": api_config["model"], "messages": [{"role": "user", "content": full_prompt}], "stream": stream_enabled } modalities = api_config.get("modalities") if isinstance(modalities, list) and modalities: payload["modalities"] = modalities logger.info(f"NanoImage请求: {url}, 模型: {api_config['model']}, 提示词长度: {len(current_prompt)} 字符") logger.debug(f"完整提示词: {current_prompt}") # 设置超时时间 max_timeout = min(api_config["timeout"], 600) timeout = httpx.Timeout( connect=10.0, read=max_timeout, write=10.0, pool=10.0 ) # 获取代理配置 proxy = await self._get_proxy() if proxy and not self._ensure_socksio(proxy): return [] client_kwargs = {"timeout": timeout, "trust_env": True} if proxy: client_kwargs["proxy"] = proxy async with httpx.AsyncClient(**client_kwargs) as client: async with client.stream("POST", url, json=payload, headers=headers) as response: logger.debug(f"收到响应状态码: {response.status_code}") if response.status_code == 200: content_type = (response.headers.get("content-type") or "").lower() logger.info(f"响应 Content-Type: {content_type}") # 放宽 SSE 检测：只要不是明确的 application/json 就当作流式处理 is_sse = "application/json" not in content_type # 处理流式响应（SSE） image_url = None image_base64 = None full_content = "" if is_sse: logger.info("使用流式响应处理模式") async for line in response.aiter_lines(): if not line: continue if line.startswith("event:") or line.startswith(":"): continue if line.startswith("data:"): data_str = line[5:].lstrip() if data_str == "[DONE]": break try: import json data = json.loads(data_str) url_candidate, b64_candidate = self._extract_image_from_payload(data) if url_candidate and not image_url: image_url = url_candidate logger.info(f"从 SSE JSON 提取到图片URL: {image_url}") if b64_candidate and not image_base64: image_base64 = b64_candidate logger.info("从 SSE JSON 提取到 base64 图片") if "choices" in data and data["choices"]: choice0 = data["choices"][0] if isinstance(data["choices"][0], dict) else {} delta = choice0.get("delta", {}) if isinstance(choice0, dict) else {} content = delta.get("content", "") reasoning = delta.get("reasoning_content", "") if content: full_content += content if reasoning: full_content += reasoning except Exception as e: # JSON 解析失败时，尝试从原始内容提取 if "http" in data_str or "data:image" in data_str: urls = re.findall(r'https?://[^\s\)\]"\']+', data_str) if urls and not image_url: image_url = urls[0].rstrip("'\"") logger.info(f"从 SSE 文本提取到图片URL: {image_url}") if "data:image" in data_str and not image_base64: image_base64 = data_str.strip() logger.info("从 SSE 文本提取到 base64 图片") else: logger.warning(f"解析响应数据失败: {e}") continue else: # 非流式（application/json）：某些网关即使传了 stream=true 也会返回完整 JSON raw = await response.aread() try: import json data = json.loads(raw.decode("utf-8", errors="ignore")) except Exception as e: logger.error(f"解析 JSON 响应失败: {type(e).__name__}: {e}") data = None if isinstance(data, dict): # 1) 标准 images endpoint 兼容：{"data":[{"url":...}|{"b64_json":...}]} items = data.get("data") if isinstance(items, list) and items: first = items[0] if isinstance(items[0], dict) else {} if isinstance(first, dict): b64_json = first.get("b64_json") if b64_json: image_base64 = b64_json logger.info("从 data[0].b64_json 提取到 base64 图片") else: u = first.get("url") or "" if isinstance(u, str) and u: image_url = u logger.info(f"从 data[0].url 提取到图片URL: {image_url}") # 2) chat.completion 兼容：choices[0].message.images[0].image_url.url if not image_url and not image_base64: try: choices = data.get("choices") or [] if choices: msg = (choices[0].get("message") or {}) if isinstance(choices[0], dict) else {} images = msg.get("images") or [] if isinstance(images, list) and images: img0 = images[0] if isinstance(images[0], dict) else {} if isinstance(img0, dict): img_data = ( (img0.get("image_url") or {}).get("url") if isinstance(img0.get("image_url"), dict) else img0.get("url") ) if isinstance(img_data, str) and img_data: if img_data.startswith("data:image"): image_base64 = img_data logger.info("从 message.images 提取到 base64 图片") elif img_data.startswith("http"): image_url = img_data logger.info(f"从 message.images 提取到图片URL: {image_url}") except Exception: pass if not image_url and not image_base64: url_candidate, b64_candidate = self._extract_image_from_payload(data) if url_candidate: image_url = url_candidate logger.info(f"从 JSON payload 提取到图片URL: {image_url}") if b64_candidate: image_base64 = b64_candidate logger.info("从 JSON payload 提取到 base64 图片") # 如果没有从流中提取到URL，尝试从完整内容中提取 if not image_url and not image_base64 and full_content: urls = re.findall(r'https?://[^\s\)\]"\']+', full_content) if urls: image_url = urls[0].rstrip("'\"") logger.info(f"从完整内容提取到图片URL: {image_url}") if not image_url and not image_base64: # 避免把 base64 打到日志里：只输出裁剪后的概要 if full_content: logger.error(f"未能提取到图片，完整响应(截断): {full_content[:500]}") else: # 非SSE时 full_content 可能为空，补充输出 content-type 便于定位 logger.error(f"未能提取到图片（content-type={content_type or 'unknown'}）") if simplify_on_fail and not simplified_prompt: simplified_prompt = self._simplify_prompt(prompt, simplify_max_tags, simplify_max_chars) if simplified_prompt and simplified_prompt != prompt: logger.warning(f"将使用简化提示词重试: {simplified_prompt}") # 处理 base64 图片 if image_base64: image_path = await self._save_base64_image(image_base64) if image_path: logger.success("成功生成图像 (base64)") return [image_path] else: logger.warning(f"base64图片保存失败，将重试 ({attempt + 1}/{max_retry})") continue # 处理 URL 图片 if image_url: image_path = await self._download_image(image_url) if image_path: logger.success("成功生成图像") return [image_path] else: logger.warning(f"图片下载失败，将重试 ({attempt + 1}/{max_retry})") continue elif response.status_code == 401: logger.error("API Key 认证失败") return [] else: error_text = await response.aread() logger.error(f"API请求失败: {response.status_code}, {error_text[:200]}") continue except asyncio.TimeoutError: logger.warning(f"请求超时，重试中... ({attempt + 1}/{max_retry})") continue except httpx.ReadTimeout: logger.warning(f"读取超时，重试中... ({attempt + 1}/{max_retry})") continue except Exception as e: import traceback logger.error(f"请求异常: {type(e).__name__}: {str(e)}") logger.error(f"异常详情:\n{traceback.format_exc()}") continue logger.error("图像生成失败") return [] def _decode_base64_image(self, image_base64: str) -> tuple[Optional[bytes], str, str]: """解析 base64 图片，返回 (bytes, ext, mime)。""" if not isinstance(image_base64, str) or not image_base64.strip(): return None, "png", "image/png" raw = image_base64.strip() payload = raw mime_type = "image/png" ext = "png" if raw.startswith("data:image"): try: header, payload = raw.split(",", 1) mime_match = re.search(r"^data:([^;]+);base64$", header, re.IGNORECASE) if mime_match: mime_type = mime_match.group(1).strip().lower() except Exception: payload = raw if "jpeg" in mime_type or "jpg" in mime_type: ext = "jpg" elif "webp" in mime_type: ext = "webp" elif "gif" in mime_type: ext = "gif" elif "png" in mime_type: ext = "png" try: image_bytes = base64.b64decode(payload) if not image_bytes: return None, ext, mime_type return image_bytes, ext, mime_type except Exception as e: logger.error(f"解析编辑图片 base64 失败: {e}") return None, ext, mime_type async def edit_image(self, prompt: str, image_base64: str) -> List[str]: """基于引用图调用 /v1/images/edits 进行改图。""" edits_config = self.config.get("edits", {}) if not edits_config.get("enabled", True): logger.warning("NanoImage 编辑能力已关闭，跳过 edits 调用") return [] api_config = self.config["api"] gen_config = self.config.get("generation", {}) max_retry = int(edits_config.get("max_retry_attempts", gen_config.get("max_retry_attempts", 2))) max_retry = max(max_retry, 1) endpoint = str(edits_config.get("endpoint", "/v1/images/edits") or "/v1/images/edits").strip() if not endpoint.startswith("/"): endpoint = f"/{endpoint}" url = f"{api_config['base_url'].rstrip('/')}{endpoint}" image_bytes, image_ext, image_mime = self._decode_base64_image(image_base64) if not image_bytes: logger.error("改图失败：引用图片解析失败") return [] size = str(edits_config.get("size", "") or "").strip() response_format = str(edits_config.get("response_format", "") or "").strip() image_count = int(edits_config.get("n", 1) or 1) allow_json_fallback = bool(edits_config.get("allow_json_fallback", True)) for attempt in range(max_retry): if attempt > 0: await asyncio.sleep(min(2 ** attempt, 10)) try: headers = { "Authorization": f"Bearer {api_config['api_key']}" } max_timeout = min(api_config["timeout"], 600) timeout = httpx.Timeout( connect=10.0, read=max_timeout, write=10.0, pool=10.0 ) proxy = await self._get_proxy() if proxy and not self._ensure_socksio(proxy): return [] client_kwargs = {"timeout": timeout, "trust_env": True} if proxy: client_kwargs["proxy"] = proxy form_data = { "model": api_config["model"], "prompt": prompt, "n": str(max(image_count, 1)), } if size: form_data["size"] = size if response_format: form_data["response_format"] = response_format logger.info(f"NanoImage 改图请求: {url}, 模型: {api_config['model']}, 提示词长度: {len(prompt)}") async with httpx.AsyncClient(**client_kwargs) as client: response = await client.post( url, data=form_data, files={"image": (f"edit_source.{image_ext}", image_bytes, image_mime)}, headers=headers, ) if response.status_code != 200 and allow_json_fallback and response.status_code in (400, 415, 422): json_payload = { "model": api_config["model"], "prompt": prompt, "n": max(image_count, 1), "image": image_base64, } if size: json_payload["size"] = size if response_format: json_payload["response_format"] = response_format logger.warning(f"edits multipart 失败({response.status_code})，尝试 JSON 回退") response = await client.post( url, json=json_payload, headers={ "Authorization": f"Bearer {api_config['api_key']}", "Content-Type": "application/json", }, ) logger.debug(f"NanoImage edits 状态码: {response.status_code}") if response.status_code != 200: err_text = (response.text or "")[:300] logger.error(f"edits 请求失败: {response.status_code}, {err_text}") continue payload = None try: payload = response.json() except Exception as e: logger.warning(f"edits 响应非 JSON: {e}") image_url = None result_base64 = None if isinstance(payload, list): payload = {"data": payload} if isinstance(payload, dict): image_url, result_base64 = self._extract_image_from_payload(payload) if not image_url and not result_base64: body_text = response.text or "" if body_text: urls = re.findall(r'https?://[^\s\)\]\"\']+', body_text) if urls: image_url = urls[0].rstrip("'\"") elif "data:image" in body_text: result_base64 = body_text.strip() if result_base64: image_path = await self._save_base64_image(result_base64) if image_path: logger.success("NanoImage 改图成功 (base64)") return [image_path] if image_url: image_path = await self._download_image(image_url) if image_path: logger.success("NanoImage 改图成功 (url)") return [image_path] logger.error("edits 响应中未解析到图片结果") except asyncio.TimeoutError: logger.warning(f"改图请求超时，重试中... ({attempt + 1}/{max_retry})") continue except httpx.ReadTimeout: logger.warning(f"改图读取超时，重试中... ({attempt + 1}/{max_retry})") continue except Exception as e: logger.error(f"改图请求异常: {type(e).__name__}: {e}") continue logger.error("NanoImage 改图失败") return [] def _extract_image_from_payload(self, data: dict) -> tuple[Optional[str], Optional[str]]: """从响应 payload 中提取图片 URL 或 base64""" if not isinstance(data, dict): return None, None image_url = None image_base64 = None items = data.get("data") if isinstance(items, list) and items: first = items[0] if isinstance(items[0], dict) else {} if isinstance(first, dict): b64 = first.get("b64_json") or first.get("image_base64") if isinstance(b64, str) and b64: image_base64 = b64 u = first.get("url") or first.get("image_url") if isinstance(u, dict): u = u.get("url") if isinstance(u, str) and u: image_url = u choices = data.get("choices") or [] for choice in choices: if not isinstance(choice, dict): continue for container in (choice.get("delta"), choice.get("message")): if not isinstance(container, dict): continue images = container.get("images") or [] if isinstance(images, list) and images: img0 = images[0] if isinstance(images[0], dict) else {} if isinstance(img0, dict): img_data = ( (img0.get("image_url") or {}).get("url") if isinstance(img0.get("image_url"), dict) else img0.get("url") ) if isinstance(img_data, str) and img_data: if img_data.startswith("data:image"): image_base64 = img_data else: image_url = img_data direct_url = container.get("image_url") or container.get("url") if isinstance(direct_url, dict): direct_url = direct_url.get("url") if isinstance(direct_url, str) and direct_url: if direct_url.startswith("data:image"): image_base64 = direct_url else: image_url = direct_url b64 = container.get("b64_json") or container.get("image_base64") if isinstance(b64, str) and b64: image_base64 = b64 content = container.get("content") if isinstance(content, str) and content: if "data:image" in content and not image_base64: image_base64 = content if "http" in content and not image_url: urls = re.findall(r'https?://[^\s\)\]"\']+', content) if urls: image_url = urls[0].rstrip("'\"") tool_calls = container.get("tool_calls") or [] if isinstance(tool_calls, list): for call in tool_calls: if not isinstance(call, dict): continue args = call.get("arguments") if not args: func = call.get("function") or {} if isinstance(func, dict): args = func.get("arguments") if isinstance(args, str): try: import json args = json.loads(args) except Exception: continue if isinstance(args, dict): for key in ("url", "image_url", "image", "b64_json", "image_base64"): val = args.get(key) if isinstance(val, dict): val = val.get("url") if isinstance(val, str) and val: if val.startswith("data:image"): image_base64 = val elif val.startswith("http"): image_url = val return image_url, image_base64 def _simplify_prompt(self, prompt: str, max_tags: int, max_chars: int) -> str: """简化提示词，减少失败概率""" if not isinstance(prompt, str): return "" parts = [p.strip() for p in prompt.split(",") if p.strip()] if len(parts) >= 2: return ", ".join(parts[:max_tags]) return prompt.strip()[:max_chars] async def _get_proxy(self) -> Optional[str]: """获取 AIChat 插件的代理配置""" try: aichat_config_path = Path(__file__).parent.parent / "AIChat" / "config.toml" if aichat_config_path.exists(): with open(aichat_config_path, "rb") as f: aichat_config = tomllib.load(f) proxy_config = aichat_config.get("proxy", {}) if proxy_config.get("enabled", False): proxy_type = proxy_config.get("type", "socks5") proxy_host = proxy_config.get("host", "127.0.0.1") proxy_port = proxy_config.get("port", 7890) proxy = f"{proxy_type}://{proxy_host}:{proxy_port}" logger.debug(f"使用代理: {proxy}") return proxy except Exception as e: logger.warning(f"读取代理配置失败: {e}") # 尝试读取环境变量代理（适配全局代理/系统代理） for key in ("HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy"): proxy = (os.environ.get(key) or "").strip() if proxy: logger.debug(f"使用环境变量代理: {key}={proxy}") return proxy return None async def _save_base64_image(self, base64_data: str) -> Optional[str]: """保存 base64 图片到本地""" try: # 去除 data:image/xxx;base64, 前缀 if base64_data.startswith("data:image"): # 提取格式和数据 header, data = base64_data.split(",", 1) # 从 header 中提取格式，如 data:image/jpeg;base64 if "jpeg" in header or "jpg" in header: ext = "jpg" elif "png" in header: ext = "png" elif "gif" in header: ext = "gif" elif "webp" in header: ext = "webp" else: ext = "jpg" else: data = base64_data ext = "jpg" # 解码 base64 image_bytes = base64.b64decode(data) # 生成文件名 ts = datetime.now().strftime("%Y%m%d_%H%M%S") uid = uuid.uuid4().hex[:8] file_path = self.images_dir / f"nano_{ts}_{uid}.{ext}" # 保存文件 with open(file_path, "wb") as f: f.write(image_bytes) logger.info(f"base64图片保存成功: {file_path}") return str(file_path) except Exception as e: logger.error(f"保存base64图片失败: {e}") import traceback logger.error(traceback.format_exc()) return None async def _download_image(self, url: str) -> Optional[str]: """下载图片到本地""" try: timeout = httpx.Timeout(connect=10.0, read=30.0, write=10.0, pool=10.0) proxy = await self._get_proxy() headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36", "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8", "Accept-Language": "zh-CN,zh;q=0.9", } if proxy and not self._ensure_socksio(proxy): return None client_kwargs = {"timeout": timeout, "trust_env": True} if proxy: client_kwargs["proxy"] = proxy async with httpx.AsyncClient(**client_kwargs) as client: response = await client.get(url, headers=headers, follow_redirects=True) response.raise_for_status() # 生成文件名 ts = datetime.now().strftime("%Y%m%d_%H%M%S") uid = uuid.uuid4().hex[:8] file_path = self.images_dir / f"nano_{ts}_{uid}.jpg" # 保存文件 with open(file_path, "wb") as f: f.write(response.content) logger.info(f"图片下载成功: {file_path}") return str(file_path) except Exception as e: logger.error(f"下载图片失败: {type(e).__name__}: {e!r}") return None def _ensure_socksio(self, proxy_url: str) -> bool: """检测 SOCKS 代理依赖""" try: scheme = urlparse(proxy_url).scheme.lower() except Exception: return True if scheme.startswith("socks"): try: import socksio # noqa: F401 except Exception: logger.error("检测到 SOCKS 代理，但未安装 socksio。请执行: pip install socksio") return False return True @on_text_message(priority=70) async def handle_message(self, bot: WechatHookClient, message: dict): """处理文本消息""" if not self.config["behavior"]["enable_command"]: return True content = message.get("Content", "").strip() from_wxid = message.get("FromWxid", "") is_group = message.get("IsGroup", False) # 检查群聊/私聊开关 if is_group and not self.config["behavior"]["enable_group"]: return True if not is_group and not self.config["behavior"]["enable_private"]: return True # 检查是否是绘图命令 keywords = self.config["behavior"]["command_keywords"] matched_keyword = None for keyword in keywords: if content.startswith(keyword + " ") or content == keyword: matched_keyword = keyword break if not matched_keyword: return True # 提取提示词 prompt = content[len(matched_keyword):].strip() if not prompt: await bot.send_text(from_wxid, f"❌ 请提供绘图提示词\n用法: {matched_keyword} <提示词>") return False logger.info(f"收到绘图请求: {prompt[:50]}...") try: # 生成图像 image_paths = await self.generate_image(prompt) if image_paths: # 直接发送图片 await bot.send_image(from_wxid, image_paths[0]) logger.success("绘图成功，已发送图片") else: await bot.send_text(from_wxid, "❌ 图像生成失败，请稍后重试") except Exception as e: logger.error(f"绘图处理失败: {e}") await bot.send_text(from_wxid, f"❌ 处理失败: {str(e)}") return False def get_llm_tools(self) -> List[dict]: """?? LLM ?????""" if not self.config["llm_tool"]["enabled"]: return [] return [{ "type": "function", "function": { "name": self.config["llm_tool"]["tool_name"], "description": ( str(self.config["llm_tool"]["tool_description"] or "").strip() or "????/???????????????????????????" ), "parameters": { "type": "object", "properties": { "prompt": { "type": "string", "description": "?????????????" }, "image_base64": { "type": "string", "description": "??????????????data:image/... ?? base64??????? /v1/images/edits?" } }, "required": ["prompt"], "additionalProperties": False } } }] async def execute_llm_tool(self, tool_name: str, arguments: dict, bot: WechatHookClient, from_wxid: str) -> dict: """?? LLM ?????""" expected_tool_name = self.config["llm_tool"]["tool_name"] if tool_name != expected_tool_name: return None try: prompt = str(arguments.get("prompt") or "").strip() image_base64 = arguments.get("image_base64") if not isinstance(image_base64, str): image_base64 = "" image_base64 = image_base64.strip() if not prompt: return {"success": False, "message": "???????"} image_paths: List[str] = [] if image_base64: logger.info(f"LLM??????: {prompt[:50]}...") image_paths = await self.edit_image(prompt, image_base64) if not image_paths and self.config.get("edits", {}).get("fallback_to_generation", False): logger.warning("?????????????") image_paths = await self.generate_image(prompt) else: logger.info(f"LLM??????: {prompt[:50]}...") image_paths = await self.generate_image(prompt) if image_paths: await bot.send_image(from_wxid, image_paths[0]) return { "success": True, "message": "????????", "images": [image_paths[0]] } return {"success": False, "message": "??????"} except Exception as e: logger.error(f"LLM??????: {e}") return {"success": False, "message": f"????: {str(e)}"}