aiData/Util/BananaClient.py

import asyncio
import logging
import os
import re
import sys
import uuid
from Config.Config import *
import aiohttp

# 配置日志
# 获取项目根目录
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

from Util.ObsUtil import ObsUploader

# 配置日志
logger = logging.getLogger(__name__)

class BananaClient:
    """
    通用 Banana 图片生成客户端，支持 TuZi 和 GPTNB 供应商
    """
    
    PROVIDERS = {
        "TuZi": {
            "api_key": TUZI_API_KEY,
            "base_url": TUZI_BASE_URL,
            "default_model": "gemini-3-pro-image-preview",
            "models": {
                "gemini-3-pro-image-preview": "gemini-3-pro-image-preview",
                "gemini-3-pro-image-preview-2k": "gemini-3-pro-image-preview-2k",
                "gemini-3-pro-image-preview-4k": "gemini-3-pro-image-preview-4k"
            }
        },
        "GPTNB": {
            "api_key": GPTNB_API_KEY,
            "base_url": GPTNB_BASE_URL,
            "default_model": "nano-banana-2",
            "models": {
                "nano-banana-2": "nano-banana-2",
                "nano-banana-2-4k": "nano-banana-2-4k",
                "nano-banana-2-hd": "nano-banana-2-hd"
            }
        }
    }
    def __init__(self, provider=None, api_key=None, base_url=None, model=None, response_format="url", timeout=120):
        """
        初始化 BananaClient
        
        Args:
            provider: 供应商名称，"TuZi" 或 "GPTNB"。如果为None则从配置读取BANANA_MODEL
            api_key: API Key，如果不提供则使用配置中的默认值
            base_url: Base URL，如果不提供则使用配置中的默认值
            model: 模型名称，如果不提供则使用供应商的默认模型
            response_format: 响应格式，默认为 "url"
            timeout: 超时时间，默认为 120 秒
        """
        if provider is None:
            # 尝试从全局配置获取
            try:
                provider = BANANA_MODEL
            except NameError:
                provider = None

            if provider:
                logger.info(f"从配置中读取 Banana 提供商: {provider}")
            else:
                provider = "TuZi"
                logger.warning("未在配置中找到 BANANA_MODEL，默认使用 TuZi")

        if provider not in self.PROVIDERS:
            raise ValueError(f"不支持的供应商: {provider}。支持的供应商: {list(self.PROVIDERS.keys())}")
            
        self.provider_name = provider
        provider_config = self.PROVIDERS[provider]
        
        self.api_key = api_key or provider_config["api_key"]
        self.base_url = base_url or provider_config["base_url"]
        self.model = model or provider_config["default_model"]
        self.available_models = provider_config["models"]
        
        self.response_format = response_format
        self.timeout = timeout
        self.obs_uploader = ObsUploader()
        
        # 统一 provider 标识用于 OBS 路径等 (全小写)
        self.provider_code = provider.lower()
        
        logger.info(f"初始化 BananaClient - 供应商: {self.provider_name}, base_url: {self.base_url}, 模型: {self.model}")

    async def prepare_image_data(self, image_paths):
        """
        准备图片数据，将本地文件上传到OBS，URL直接使用，包装为正确的API调用格式
        
        参数:
            image_paths: 单个图片路径字符串或图片路径列表，可以是本地文件路径或URL
        
        返回:
            单个图片的API调用格式数据或图片API调用格式数据列表
        """

        async def process_single_image(image_path):
            """处理单个图片路径，支持本地文件和URL"""
            image_path = image_path.strip()
            logger.info(f"开始处理图片: {image_path}")
            
            try:
                # 检查是否是URL
                if re.match(r'^https?://', image_path):
                    # 是URL，直接使用
                    logger.info(f"图片是URL类型，直接使用: {image_path}")
                    return {
                        "type": "image_url",
                        "image_url": {"url": image_path}
                    }
                else:
                    # 是本地文件，上传到OBS
                    logger.info(f"图片是本地文件，开始上传到OBS: {image_path}")
                    if not os.path.exists(image_path):
                        raise FileNotFoundError(f"本地图片文件不存在: {image_path}")
                    if not os.path.isfile(image_path):
                        raise IsADirectoryError(f"指定路径不是文件: {image_path}")
                    
                    file_size = os.path.getsize(image_path)
                    logger.info(f"本地文件存在，大小: {file_size} bytes")
                    
                    # 生成GUID作为文件名
                    file_guid = str(uuid.uuid4())
                    
                    # 从文件名获取扩展名
                    _, ext = os.path.splitext(image_path)
                    ext = ext.lower() if ext else '.jpg'  # 默认使用jpg
                    
                    # 构建OBS key，使用临时文件前缀
                    obs_key = f"{OBS_TMP_PREFIX}/{file_guid}{ext}"
                    logger.debug(f"构建OBS key: {obs_key}")
                    
                    # 上传到OBS
                    logger.info(f"正在上传到OBS: {obs_key}")
                    success, result = self.obs_uploader.upload_file(
                        file_path=image_path,
                        object_key=obs_key,
                        bucket_name=OBS_BUCKET
                    )
                    
                    if success:
                        # 构建CDN访问URL
                        cdn_url = f"https://{CDN_DOMAIN}/{obs_key}"
                        logger.info(f"✅ 图片上传到OBS成功: {cdn_url}")
                        logger.info(f"   OBS Key: {obs_key}")
                        
                        return {
                            "type": "image_url",
                            "image_url": {"url": cdn_url}
                        }
                    else:
                        logger.error(f"上传到OBS失败: {result}")
                        raise Exception(f"上传到OBS失败: {result}")

            except Exception as e:
                logger.error(f"处理单个图片时出错: {image_path} - {str(e)}")
                raise

        try:
            logger.info(f"准备图片数据: {image_paths}")
            # 如果是单个图片路径字符串
            if isinstance(image_paths, str):
                return await process_single_image(image_paths)
            # 如果是图片路径列表
            elif isinstance(image_paths, list):
                result = []
                for i, image_path in enumerate(image_paths):
                    logger.info(f"处理列表中的图片 {i+1}/{len(image_paths)}: {image_path}")
                    result.append(await process_single_image(image_path))
                logger.info(f"成功处理 {len(result)} 张图片")
                return result
            else:
                raise TypeError(f"image_paths参数类型错误，应为str或list，实际为{type(image_paths)}")
        except Exception as e:
            logger.error(f"准备图片数据时出错: {e}")
            raise

    async def generate_image(self, prompt, image_paths=None, n=1, size="1024x1024"):
        """
        异步调用 API 生成图片（支持文字生图和图片+文字改图）
        
        参数:
            prompt: 图片描述
            image_paths: 图片路径或URL，可以是单个字符串或列表（用于改图功能）
            n: 生成图片数量，默认为 1
            size: 图片尺寸，默认为 "1024x1024"
            
        返回:
            response_json: API 响应的 JSON 数据
        """
        
        # 确保使用有效的模型
        current_model = self.model
        if current_model not in self.available_models:
             logger.warning(f"警告: 模型 {current_model} 可能不在 {self.provider_name} 的支持列表中: {list(self.available_models.keys())}")
        
        logger.info(f"开始生成图片 - 提供商: {self.provider_name}, 模型: {current_model}, 尺寸: {size}, 数量: {n}")
        logger.info(f"提示词: {prompt}")
        if image_paths:
            logger.info(f"图片输入: {image_paths}")
        
        # 构建API调用消息
        if image_paths:
            # 处理图片数据
            image_data = await self.prepare_image_data(image_paths)
            
            # 构建payload，支持改图功能
            image_urls = []
            if isinstance(image_data, list):
                for item in image_data:
                    if "image_url" in item and "url" in item["image_url"]:
                        image_urls.append(item["image_url"]["url"])
            else:
                if "image_url" in image_data and "url" in image_data["image_url"]:
                    image_urls.append(image_data["image_url"]["url"])
            
            # 如果只有一个URL，直接传递字符串；如果有多个，传递列表
            final_image_input = image_urls[0] if len(image_urls) == 1 else image_urls

            payload = {
                "prompt": prompt,
                "n": n,
                "model": current_model,
                "size": size,
                "response_format": self.response_format,
                "image": final_image_input,  # 支持单个URL或URL列表
                "timeout": self.timeout
            }
            logger.info(f"构建改图请求，使用模型: {current_model}，尺寸: {size}")
        else:
            # 仅文字生图
            payload = {
                "prompt": prompt,
                "n": n,
                "model": current_model,
                "size": size,
                "response_format": self.response_format,
                "timeout": self.timeout
            }
            logger.info(f"构建文字生图请求，使用模型: {current_model}，尺寸: {size}")
            
        endpoint = f"{self.base_url}/v1/images/generations"
        
        headers = {
            'Content-Type': 'application/json',
            "Authorization": f"Bearer {self.api_key}"
        }
        logger.debug(f"请求头: {headers}")
        logger.debug(f"请求payload: {payload}")
        
        try:
            logger.info(f"正在异步调用{self.provider_name} API...")
            logger.info(f"使用端点: {endpoint}")
            
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    endpoint,
                    json=payload,
                    headers=headers,
                    timeout=self.timeout
                ) as response:
                    response.raise_for_status()  # 检查HTTP响应状态
                    response_json = await response.json()
            
            logger.info("API调用完成！")
            logger.debug(f"API响应: {response_json}")  # 使用debug级别记录完整响应
            return response_json
            
        except Exception as e:
            logger.error(f"API调用失败: {str(e)}")
            raise

    async def get_image_urls(self, response_json):
        """
        异步获取图片 URL 列表
        
        参数:
            response_json: API 响应的 JSON 数据
            
        返回:
            image_urls: 图片 URL 列表
        """
        if response_json is None:
            logger.info("未提供响应数据")
            return []
        
        # 检查数据结构并提取图片地址
        image_urls = []
        
        # 处理不同的响应格式
        if "data" in response_json:
            data = response_json["data"]
            
            # 格式1: data是字符串URL
            if isinstance(data, str):
                image_urls.append(data)
            
            # 格式2: data是包含url字段的对象列表
            elif isinstance(data, list):
                for item in data:
                    if isinstance(item, dict):
                        # 检查是否有url字段
                        if "url" in item:
                            image_urls.append(item["url"])
                        # 检查是否有嵌套的数据结构
                        elif "data" in item:
                            if isinstance(item["data"], str):
                                image_urls.append(item["data"])
                            elif isinstance(item["data"], list):
                                for nested_item in item["data"]:
                                    if isinstance(nested_item, dict) and "url" in nested_item:
                                        image_urls.append(nested_item["url"])
                                    elif isinstance(nested_item, str):
                                        image_urls.append(nested_item)
                    # 格式3: item本身是字符串URL
                    elif isinstance(item, str):
                        image_urls.append(item)
            
            # 格式4: data是一个对象，包含url字段
            elif isinstance(data, dict):
                if "url" in data:
                    image_urls.append(data["url"])
                # 检查是否有嵌套的数据结构
                elif "data" in data:
                    if isinstance(data["data"], str):
                        image_urls.append(data["data"])
                    elif isinstance(data["data"], list):
                        for item in data["data"]:
                            if isinstance(item, dict) and "url" in item:
                                image_urls.append(item["url"])
                            elif isinstance(item, str):
                                image_urls.append(item)
                                
        # Fallback for TuZi provider which might return 'url' at top level or nested differently
        if not image_urls and "url" in response_json:
             image_urls.append(response_json["url"])
             
        # Log if no URLs found but response exists
        if not image_urls:
             logger.warning(f"Failed to extract image URLs from response: {response_json}")

        return image_urls
    
    async def print_image_urls(self, response_json):
        """
        异步打印图片 URL
        
        参数:
            response_json: API 响应的 JSON 数据
        """
        logger.info("\n正在解析图片地址...")
        image_urls = await self.get_image_urls(response_json)
        
        if image_urls:
            for i, url in enumerate(image_urls):
                logger.info(f"图片 {i+1} URL: {url}")
        else:
            logger.info("未找到有效的图片 URL")

    async def download_and_upload_to_obs(self, response_json, overlay_logo_path=None):
        """
        异步从临时图片URL下载图片并直接上传到OBS（不落盘）
        
        参数:
            response_json: API 响应的 JSON 数据
            overlay_logo_path: 水印Logo图片路径（可选）
            
        返回:
            uploaded_urls: 上传到OBS的图片URL列表
        """
        logger.info("\n正在下载图片并上传到OBS...")
        
        # 获取图片URL列表
        image_urls = await self.get_image_urls(response_json)
        uploaded_urls = []
        
        if not image_urls:
            logger.info("未找到有效的图片 URL")
            return uploaded_urls
        
        logger.info(f"找到 {len(image_urls)} 个图片URL，开始下载并上传到OBS")
        
        for i, image_url in enumerate(image_urls):
            try:
                logger.info(f"正在处理图片 {i+1}/{len(image_urls)}: {image_url}")
                
                # 生成GUID作为文件名
                file_guid = str(uuid.uuid4())
                
                # 从URL获取文件扩展名
                if image_url.lower().endswith('.png'):
                    ext = 'png'
                elif image_url.lower().endswith('.jpg') or image_url.lower().endswith('.jpeg'):
                    ext = 'jpg'
                elif image_url.lower().endswith('.webp'):
                    ext = 'webp'
                else:
                    ext = 'jpg'  # 默认使用jpg
                
                # 构建OBS key
                obs_key = f"{OBS_CLOUD_PREFIX}/{self.provider_code.capitalize()}/{file_guid}.{ext}"
                logger.debug(f"构建OBS key: {obs_key}")
                
                # 使用异步方式下载图片
                logger.info(f"正在下载图片: {image_url}")
                async with aiohttp.ClientSession() as session:
                    async with session.get(image_url, timeout=120) as response:
                        response.raise_for_status()
                        image_data = await response.read()
                logger.debug(f"成功下载图片，大小: {len(image_data)} 字节")
                
                # 添加水印逻辑
                if overlay_logo_path and os.path.exists(overlay_logo_path):
                    try:
                        from PIL import Image
                        import io
                        
                        logger.info(f"正在添加水印: {overlay_logo_path}")
                        
                        # 打开主图
                        img = Image.open(io.BytesIO(image_data))
                        
                        # 打开Logo
                        logo = Image.open(overlay_logo_path)
                        
                        # 计算Logo大小 (宽度为原图的20%)
                        target_logo_width = int(img.width * 0.2)
                        aspect_ratio = logo.height / logo.width
                        target_logo_height = int(target_logo_width * aspect_ratio)
                        
                        # 调整Logo大小
                        logo = logo.resize((target_logo_width, target_logo_height), Image.Resampling.LANCZOS)
                        
                        # 计算位置 (左上角，3% padding)
                        padding = int(img.width * 0.03)
                        position = (padding, padding)
                        
                        # 粘贴Logo (处理透明度)
                        if logo.mode in ('RGBA', 'LA') or (logo.mode == 'P' and 'transparency' in logo.info):
                            img.paste(logo, position, logo)
                        else:
                            img.paste(logo, position)
                        
                        # 保存回bytes
                        output_buffer = io.BytesIO()
                        # 保持原格式，如果无法识别则使用PNG
                        save_format = img.format if img.format else 'PNG'
                        img.save(output_buffer, format=save_format)
                        image_data = output_buffer.getvalue()
                        
                        logger.info("水印添加成功")
                        
                    except Exception as e:
                        logger.error(f"添加水印失败: {e}")
                        # 失败后继续上传原图
                
                # 直接上传到OBS（不落盘）
                
                # 直接上传到OBS（不落盘）
                logger.info(f"正在上传到OBS: {obs_key}")
                success, result = self.obs_uploader.upload_base64_image(
                    object_key=obs_key,
                    base64_data=image_data,
                    bucket_name=OBS_BUCKET
                )
                
                if success:
                    # 构建OBS访问URL
                    obs_url = f"https://{OBS_BUCKET}.{OBS_SERVER}/{obs_key}"
                    # 构建CDN访问URL
                    cdn_url = f"https://{CDN_DOMAIN}/{obs_key}"
                    logger.info(f"✅ 图片 {i+1} 上传成功: {cdn_url}")
                    logger.info(f"   OBS Key: {obs_key}")
                    uploaded_urls.append(cdn_url)
                else:
                    logger.error(f"上传到OBS失败: {result}")
            
            except Exception as e:
                logger.error(f"处理图片 {image_url} 失败: {str(e)}")
                import traceback
                logger.error(f"详细错误堆栈: {traceback.format_exc()}")
        
        logger.info(f"图片处理完成，成功上传 {len(uploaded_urls)}/{len(image_urls)} 个图片到OBS")
        return uploaded_urls


async def test_provider(provider_name):
    """测试指定供应商"""
    print(f"\n{'='*20} 测试 {provider_name} 供应商 {'='*20}")
    
    try:
        # 创建客户端实例
        client = BananaClient(provider=provider_name)
        
        # 测试1: 文字生图
        logger.info("\n---------------- 文字生图测试 ----------------")
        text_prompt = "一只可爱的猫咪坐在月亮上，星空背景，卡通风格"
        logger.info(f"提示词: {text_prompt}")
        
        response = await client.generate_image(
            prompt=text_prompt,
            n=1,
            size="1024x1024"
        )
        
        if response:
            logger.info("文字生图成功！")
            await client.print_image_urls(response)
            # 可选：上传到OBS
            # await client.download_and_upload_to_obs(response)
        else:
            logger.error("文字生图失败！")
            
    except Exception as e:
        logger.error(f"{provider_name} 测试失败: {e}")
        import traceback
        logger.error(traceback.format_exc())

async def main():
    """
    异步主函数，用于测试 BananaClient 的功能
    """
    # 配置控制台日志输出
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[logging.StreamHandler(sys.stdout)]
    )
    
    logger.info("开始 BananaClient 整合测试...")
    
    # 测试 TuZi (默认)
    await test_provider("TuZi")
    
    # 测试 GPTNB
    #await test_provider("GPTNB")
    
    logger.info("\n" + "="*50)
    logger.info("测试全部完成！")
    logger.info("="*50)


if __name__ == "__main__":
    asyncio.run(main())