Initial commit: Comfyui_o1key v1.10.0

clients/gemini_flash_client.py

@@ -0,0 +1,273 @@
+"""
+Gemini Flash API 客户端
+用于调用 Gemini 3 Flash 模型进行多模态文本生成
+"""
+
+import asyncio
+from typing import Any, Dict, List, Optional
+
+import aiohttp
+
+from ..utils.config import get_api_key_or_raise
+from ..models_config import get_flash_model_endpoint, get_enabled_flash_models
+from .base_client import BaseAPIClient
+
+
+# API 基础配置
+API_BASE_URL = "https://api.o1key.com"
+
+
+class GeminiFlashClient(BaseAPIClient):
+    """
+    Gemini Flash API 客户端
+    用于调用 Gemini 3 Flash 模型进行多模态文本生成
+    
+    特点：
+    - 支持图片和视频输入
+    - 支持系统指令
+    - 支持不同思考深度
+    """
+    
+    def __init__(self, api_key: Optional[str] = None):
+        """
+        初始化客户端
+        
+        Args:
+            api_key: API 密钥，如果为 None 则从配置文件或环境变量读取
+        """
+        if api_key is None:
+            api_key = get_api_key_or_raise("O1KEY_API_KEY")
+        
+        super().__init__(
+            base_url=API_BASE_URL,
+            api_key=api_key,
+            max_request_size=20 * 1024 * 1024  # 20MB
+        )
+    
+    def get_endpoint(
+        self, 
+        model: str = "gemini-3-flash-preview",
+        thinking_depth: str = "不思考", 
+        **kwargs
+    ) -> str:
+        """
+        根据模型和思考深度获取 API 端点
+        
+        Args:
+            model: 模型名称
+            thinking_depth: 思考深度 ("不思考" 或 "高")
+        
+        Returns:
+            API 端点路径
+        """
+        endpoint = get_flash_model_endpoint(model, thinking_depth)
+        
+        if endpoint is None:
+            # 回退到默认端点
+            default_models = get_enabled_flash_models()
+            if default_models:
+                endpoint = get_flash_model_endpoint(default_models[0], thinking_depth)
+        
+        if endpoint is None:
+            raise ValueError(f"无法获取模型 '{model}' 的端点 (思考深度: {thinking_depth})")
+        
+        return endpoint
+    
+    def build_request_body(
+        self,
+        prompt: str = "",
+        system_instruction: Optional[str] = None,
+        image_data: Optional[List[Dict[str, str]]] = None,
+        video_data: Optional[Dict[str, str]] = None,
+        **kwargs
+    ) -> Dict[str, Any]:
+        """
+        构建 API 请求体
+        
+        Args:
+            prompt: 用户提示词
+            system_instruction: 系统指令（可选）
+            image_data: 图片数据列表，每个元素包含 mime_type 和 data
+            video_data: 视频数据，包含 mime_type 和 data
+        
+        Returns:
+            请求体字典
+        """
+        parts = []
+        
+        # 添加文本部分
+        if prompt:
+            parts.append({"text": prompt})
+        
+        # 添加图片部分（如果有）
+        if image_data:
+            for img in image_data:
+                parts.append({
+                    "inline_data": {
+                        "mime_type": img["mime_type"],
+                        "data": img["data"]
+                    }
+                })
+        
+        # 添加视频部分（如果有）
+        if video_data:
+            parts.append({
+                "inline_data": {
+                    "mime_type": video_data["mime_type"],
+                    "data": video_data["data"]
+                }
+            })
+        
+        # 构建请求体
+        request_body = {
+            "contents": [
+                {
+                    "parts": parts
+                }
+            ]
+        }
+        
+        # 添加系统指令（如果有）
+        if system_instruction and system_instruction.strip():
+            request_body["system_instruction"] = {
+                "parts": [
+                    {"text": system_instruction}
+                ]
+            }
+        
+        return request_body
+    
+    def parse_response(self, response: Dict[str, Any]) -> str:
+        """
+        解析 API 响应，提取生成的文本
+        
+        Args:
+            response: API 响应字典
+        
+        Returns:
+            生成的文本内容
+        
+        Raises:
+            RuntimeError: 解析失败或 API 拒绝时
+        """
+        # 检查 candidatesTokenCount
+        usage_metadata = response.get("usageMetadata", {})
+        candidates_token_count = usage_metadata.get("candidatesTokenCount", -1)
+        
+        if candidates_token_count == 0:
+            raise RuntimeError(
+                "内容审核拒绝 - candidatesTokenCount = 0\n\n"
+                "原因：提示词或输入内容包含不适当内容\n"
+                "建议：检查并调整输入内容"
+            )
+        
+        # 检查 finishReason
+        candidates = response.get("candidates", [])
+        if candidates:
+            for candidate in candidates:
+                finish_reason = candidate.get("finishReason", "")
+                
+                if finish_reason and finish_reason not in ["STOP", "MAX_TOKENS"]:
+                    reason_messages = {
+                        "PROHIBITED_CONTENT": "违禁内容拒绝",
+                        "SAFETY": "安全过滤器拒绝",
+                        "RECITATION": "版权问题"
+                    }
+                    error_title = reason_messages.get(finish_reason, f"生成异常 ({finish_reason})")
+                    raise RuntimeError(f"{error_title}\n建议：调整输入内容后重试")
+        
+        # 提取文本内容
+        text_parts = []
+        
+        for candidate in candidates:
+            content = candidate.get("content", {})
+            parts = content.get("parts", [])
+            
+            for part in parts:
+                if "text" in part:
+                    text_parts.append(part["text"])
+        
+        if not text_parts:
+            raise RuntimeError("API 响应中未找到生成的文本")
+        
+        # 合并所有文本部分
+        return "\n".join(text_parts)
+    
+    async def generate_async(
+        self,
+        prompt: str,
+        model: str = "gemini-3-flash-preview",
+        thinking_depth: str = "不思考",
+        system_instruction: Optional[str] = None,
+        image_data: Optional[List[Dict[str, str]]] = None,
+        video_data: Optional[Dict[str, str]] = None,
+        session: Optional[aiohttp.ClientSession] = None
+    ) -> str:
+        """
+        异步生成文本
+        
+        Args:
+            prompt: 用户提示词
+            model: 模型名称
+            thinking_depth: 思考深度
+            system_instruction: 系统指令
+            image_data: 图片数据列表
+            video_data: 视频数据
+            session: aiohttp 会话
+        
+        Returns:
+            生成的文本内容
+        """
+        endpoint = self.get_endpoint(model=model, thinking_depth=thinking_depth)
+        request_body = self.build_request_body(
+            prompt=prompt,
+            system_instruction=system_instruction,
+            image_data=image_data,
+            video_data=video_data
+        )
+        
+        # 根据是否有视频设置超时（视频处理需要更长时间）
+        timeout = 300 if video_data else 180
+        
+        response = await self.request_async(
+            endpoint, 
+            request_body, 
+            session, 
+            timeout=timeout
+        )
+        
+        return self.parse_response(response)
+    
+    def generate_sync(
+        self,
+        prompt: str,
+        model: str = "gemini-3-flash-preview",
+        thinking_depth: str = "不思考",
+        system_instruction: Optional[str] = None,
+        image_data: Optional[List[Dict[str, str]]] = None,
+        video_data: Optional[Dict[str, str]] = None
+    ) -> str:
+        """
+        同步生成文本（用于 ComfyUI 节点）
+        
+        Args:
+            prompt: 用户提示词
+            model: 模型名称
+            thinking_depth: 思考深度
+            system_instruction: 系统指令
+            image_data: 图片数据列表
+            video_data: 视频数据
+        
+        Returns:
+            生成的文本内容
+        """
+        coro = self.generate_async(
+            prompt=prompt,
+            model=model,
+            thinking_depth=thinking_depth,
+            system_instruction=system_instruction,
+            image_data=image_data,
+            video_data=video_data
+        )
+        
+        return self.run_async_in_thread(coro)