""" NanoBanana Tool - 图像生成 通用图像生成工具，可以接受自然语言描述和/或图像输入，生成新的图像。 支持通过 OpenRouter 调用 Gemini 2.5 Flash Image 模型。 """ import base64 import json import mimetypes import os import re from pathlib import Path from typing import Optional, Dict, Any, List, Tuple import httpx from dotenv import load_dotenv from agent.tools import tool, ToolResult OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1" DEFAULT_TIMEOUT = 120.0 DEFAULT_IMAGE_PROMPT = "根据输入生成图像。" DEFAULT_IMAGE_MODEL_CANDIDATES = [ "google/gemini-2.5-flash-image", ] def _resolve_api_key() -> Optional[str]: """优先读取环境变量，缺失时尝试从 .env 加载。""" api_key = os.getenv("OPENROUTER_API_KEY") or os.getenv("OPEN_ROUTER_API_KEY") if api_key: return api_key load_dotenv() return os.getenv("OPENROUTER_API_KEY") or os.getenv("OPEN_ROUTER_API_KEY") def _image_to_data_url(image_path: Path) -> str: """将图片文件编码为 data URL。""" mime_type = mimetypes.guess_type(str(image_path))[0] or "application/octet-stream" raw = image_path.read_bytes() b64 = base64.b64encode(raw).decode("utf-8") return f"data:{mime_type};base64,{b64}" def _safe_json_parse(content: str) -> Dict[str, Any]: """尽量从模型文本中提取 JSON。""" try: return json.loads(content) except json.JSONDecodeError: start = content.find("{") end = content.rfind("}") if start != -1 and end != -1 and end > start: candidate = content[start:end + 1] return json.loads(candidate) raise def _extract_data_url_images(message: Dict[str, Any]) -> List[Tuple[str, str]]: """ 从 OpenRouter 响应消息中提取 data URL 图片。 Returns: List[(mime_type, base64_data)] """ extracted: List[Tuple[str, str]] = [] # 官方文档中的主要位置：message.images[] for img in message.get("images", []) or []: if not isinstance(img, dict): continue if img.get("type") != "image_url": continue data_url = ((img.get("image_url") or {}).get("url") or "").strip() if not data_url.startswith("data:"): continue m = re.match(r"^data:([^;]+);base64,(.+)$", data_url, flags=re.DOTALL) if not m: continue extracted.append((m.group(1), m.group(2))) # 兼容某些模型可能把 image_url 放在 content 数组中 content = message.get("content") if isinstance(content, list): for part in content: if not isinstance(part, dict): continue if part.get("type") != "image_url": continue data_url = ((part.get("image_url") or {}).get("url") or "").strip() if not data_url.startswith("data:"): continue m = re.match(r"^data:([^;]+);base64,(.+)$", data_url, flags=re.DOTALL) if not m: continue extracted.append((m.group(1), m.group(2))) return extracted def _extract_image_refs(choice: Dict[str, Any], message: Dict[str, Any]) -> List[Dict[str, str]]: """ 尝试从不同响应格式中提取图片引用。 返回格式: - {"kind": "data_url", "value": "data:image/png;base64,..."} - {"kind": "base64", "value": "...", "mime_type": "image/png"} - {"kind": "url", "value": "https://..."} """ refs: List[Dict[str, str]] = [] # 1) 标准 message.images for img in message.get("images", []) or []: if not isinstance(img, dict): continue # image_url 结构 data_url = ((img.get("image_url") or {}).get("url") or "").strip() if data_url.startswith("data:"): refs.append({"kind": "data_url", "value": data_url}) continue if data_url.startswith("http"): refs.append({"kind": "url", "value": data_url}) continue # 兼容 base64 字段 b64 = (img.get("b64_json") or img.get("base64") or "").strip() if b64: refs.append({"kind": "base64", "value": b64, "mime_type": img.get("mime_type", "image/png")}) # 2) 某些格式可能在 choice.images for img in choice.get("images", []) or []: if not isinstance(img, dict): continue data_url = ((img.get("image_url") or {}).get("url") or "").strip() if data_url.startswith("data:"): refs.append({"kind": "data_url", "value": data_url}) continue if data_url.startswith("http"): refs.append({"kind": "url", "value": data_url}) continue b64 = (img.get("b64_json") or img.get("base64") or "").strip() if b64: refs.append({"kind": "base64", "value": b64, "mime_type": img.get("mime_type", "image/png")}) # 3) content 数组里的 image_url content = message.get("content") if isinstance(content, list): for part in content: if not isinstance(part, dict): continue if part.get("type") != "image_url": continue url = ((part.get("image_url") or {}).get("url") or "").strip() if url.startswith("data:"): refs.append({"kind": "data_url", "value": url}) elif url.startswith("http"): refs.append({"kind": "url", "value": url}) # 4) 极端兼容：文本中可能出现 data:image 或 http 图片 URL if isinstance(content, str): # data URL for m in re.finditer(r"(data:image\/[a-zA-Z0-9.+-]+;base64,[A-Za-z0-9+/=]+)", content): refs.append({"kind": "data_url", "value": m.group(1)}) # http(s) 图片链接 for m in re.finditer(r"(https?://\S+\.(?:png|jpg|jpeg|webp))", content, flags=re.IGNORECASE): refs.append({"kind": "url", "value": m.group(1)}) return refs def _mime_to_ext(mime_type: str) -> str: """MIME 类型映射到扩展名。""" mapping = { "image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp", } return mapping.get(mime_type.lower(), ".png") def _normalize_model_id(model_id: str) -> str: """ 规范化常见误写模型 ID，减少无效重试。 """ if not model_id: return model_id m = model_id.strip() # 常见误写：gemini/gemini-xxx -> google/gemini-xxx if m.startswith("gemini/"): m = "google/" + m.split("/", 1)[1] # 常见顺序误写：preview-image -> image if "gemini-2.5-flash-preview-image" in m: m = m.replace("gemini-2.5-flash-preview-image", "gemini-2.5-flash-image") # 兼容旧 ID 到当前可用 ID if "gemini-2.5-flash-image-preview" in m: m = m.replace("gemini-2.5-flash-image-preview", "gemini-2.5-flash-image") return m @tool(description="通用图像生成工具，可以接受自然语言描述和/或图像输入，生成新的图像") async def nanobanana( image_path: str = "", image_paths: Optional[List[str]] = None, prompt: Optional[str] = None, model: Optional[str] = None, max_tokens: int = 1200, image_output_path: Optional[str] = None, ) -> ToolResult: """ 通用图像生成工具，可以接受自然语言描述和/或图像输入，生成新的图像。 Args: image_path: 输入图片路径（单图模式，可选） image_paths: 输入图片路径列表（多图模式，可选） prompt: 自定义生成描述（可选，默认使用通用prompt） model: OpenRouter 模型名（可选，默认使用 gemini-2.5-flash-image） max_tokens: 最大输出 token image_output_path: 生成图片保存路径（可选） Returns: ToolResult: 包含生成的图片路径 """ raw_paths: List[str] = [] if image_paths: raw_paths.extend(image_paths) if image_path: raw_paths.append(image_path) # 图像输入是可选的，但如果提供了就需要验证 input_paths: List[Path] = [] if raw_paths: # 去重并检查路径 unique_raw: List[str] = [] seen = set() for p in raw_paths: if p and p not in seen: unique_raw.append(p) seen.add(p) input_paths = [Path(p) for p in unique_raw] invalid = [str(p) for p in input_paths if (not p.exists() or not p.is_file())] if invalid: return ToolResult( title="NanoBanana 生成失败", output="", error=f"以下图片不存在或不可读: {invalid}", ) api_key = _resolve_api_key() if not api_key: return ToolResult( title="NanoBanana 生成失败", output="", error="未找到 OpenRouter API Key，请设置 OPENROUTER_API_KEY 或 OPEN_ROUTER_API_KEY", ) user_prompt = prompt or DEFAULT_IMAGE_PROMPT # 编码图像（如果有） image_data_urls = [] if input_paths: try: image_data_urls = [_image_to_data_url(p) for p in input_paths] except Exception as e: return ToolResult( title="NanoBanana 生成失败", output="", error=f"图片编码失败: {e}", ) user_content: List[Dict[str, Any]] = [{"type": "text", "text": user_prompt}] for u in image_data_urls: user_content.append({"type": "image_url", "image_url": {"url": u}}) payload: Dict[str, Any] = { "messages": [ { "role": "system", "content": "你是图像生成助手。请根据用户的描述和/或输入图像生成新的图像。", }, { "role": "user", "content": user_content, }, ], "temperature": 0.2, "max_tokens": max_tokens, "modalities": ["image", "text"], } headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "HTTP-Referer": "https://local-agent", "X-Title": "Agent NanoBanana Tool", } endpoint = f"{OPENROUTER_BASE_URL}/chat/completions" # 自动尝试多个可用模型，减少 404/invalid model 影响 candidates: List[str] = [] if model: candidates.append(_normalize_model_id(model)) if env_model := os.getenv("NANOBANANA_IMAGE_MODEL"): candidates.append(_normalize_model_id(env_model)) candidates.extend([_normalize_model_id(x) for x in DEFAULT_IMAGE_MODEL_CANDIDATES]) # 去重并保持顺序 dedup: List[str] = [] seen = set() for m in candidates: if m and m not in seen: dedup.append(m) seen.add(m) candidates = dedup data: Optional[Dict[str, Any]] = None used_model: Optional[str] = None errors: List[Dict[str, Any]] = [] for cand in candidates: modality_attempts: List[Optional[List[str]]] = [["image", "text"], ["image"], None] for mods in modality_attempts: trial_payload = dict(payload) trial_payload["model"] = cand if mods is None: trial_payload.pop("modalities", None) else: trial_payload["modalities"] = mods try: async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client: resp = await client.post(endpoint, json=trial_payload, headers=headers) resp.raise_for_status() data = resp.json() used_model = cand break except httpx.HTTPStatusError as e: errors.append({ "model": cand, "modalities": mods, "status_code": e.response.status_code, "body": e.response.text[:600], }) continue except Exception as e: errors.append({ "model": cand, "modalities": mods, "status_code": None, "body": str(e)[:600], }) continue if data is not None: break if data is None: return ToolResult( title="NanoBanana 生成失败", output=json.dumps({"attempted_models": candidates, "errors": errors}, ensure_ascii=False, indent=2), long_term_memory="All candidate models failed for this request", metadata={"attempted_models": candidates, "errors": errors}, ) chosen_model = used_model or candidates[0] choices = data.get("choices") or [] message = choices[0].get("message", {}) if choices else {} # 提取生成的图像 refs = _extract_image_refs(choices[0] if choices else {}, message) if not refs: content = message.get("content") preview = "" if isinstance(content, str): preview = content[:500] elif isinstance(content, list): preview = json.dumps(content[:3], ensure_ascii=False)[:500] return ToolResult( title="NanoBanana 生成失败", output=json.dumps(data, ensure_ascii=False, indent=2), error="模型未返回可解析图片（未在 message.images/choice.images/content 中发现图片）", metadata={ "model": chosen_model, "choice_keys": list((choices[0] if choices else {}).keys()), "message_keys": list(message.keys()) if isinstance(message, dict) else [], "content_preview": preview, }, ) output_paths: List[str] = [] if image_output_path: base_path = Path(image_output_path) else: if len(input_paths) > 1: base_path = input_paths[0].parent / "set_generated.png" else: base_path = input_paths[0].parent / f"{input_paths[0].stem}_generated.png" base_path.parent.mkdir(parents=True, exist_ok=True) for idx, ref in enumerate(refs): kind = ref.get("kind", "") mime_type = "image/png" raw_bytes: Optional[bytes] = None if kind == "data_url": m = re.match(r"^data:([^;]+);base64,(.+)$", ref.get("value", ""), flags=re.DOTALL) if not m: continue mime_type = m.group(1) raw_bytes = base64.b64decode(m.group(2)) elif kind == "base64": mime_type = ref.get("mime_type", "image/png") raw_bytes = base64.b64decode(ref.get("value", "")) elif kind == "url": url = ref.get("value", "") try: with httpx.Client(timeout=DEFAULT_TIMEOUT) as client: r = client.get(url) r.raise_for_status() raw_bytes = r.content mime_type = r.headers.get("content-type", "image/png").split(";")[0] except Exception: continue else: continue if not raw_bytes: continue ext = _mime_to_ext(mime_type) if len(refs) == 1: target = base_path if target.suffix.lower() not in [".png", ".jpg", ".jpeg", ".webp"]: target = target.with_suffix(ext) else: stem = base_path.stem target = base_path.with_name(f"{stem}_{idx+1}{ext}") try: target.write_bytes(raw_bytes) output_paths.append(str(target)) except Exception as e: return ToolResult( title="NanoBanana 生成失败", output="", error=f"写入生成图片失败: {e}", metadata={"model": chosen_model}, ) if not output_paths: return ToolResult( title="NanoBanana 生成失败", output=json.dumps(data, ensure_ascii=False, indent=2), error="检测到图片引用但写入失败（可能是无效 base64 或 URL 不可访问）", metadata={"model": chosen_model, "ref_count": len(refs)}, ) usage = data.get("usage", {}) prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens", 0) completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens", 0) summary = { "model": chosen_model, "input_images": [str(p) for p in input_paths], "input_count": len(input_paths), "generated_images": output_paths, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, } return ToolResult( title="NanoBanana 图片生成完成", output=json.dumps({"summary": summary}, ensure_ascii=False, indent=2), long_term_memory=f"Generated {len(output_paths)} image(s) from {len(input_paths)} input image(s) using {chosen_model}", attachments=output_paths, metadata=summary, tool_usage={ "model": chosen_model, "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, } )