Initial commit: Comfyui_o1key v1.10.0
This commit is contained in:
o1key
@@ -0,0 +1,368 @@
|
||||
"""
|
||||
API 客户端基类
|
||||
提供通用的 HTTP 请求、响应解析和错误处理功能
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
|
||||
class BaseAPIClient(ABC):
|
||||
"""
|
||||
API 客户端抽象基类
|
||||
|
||||
子类需要实现以下方法：
|
||||
- get_endpoint(): 获取 API 端点
|
||||
- build_request_body(): 构建请求体
|
||||
- parse_response(): 解析响应
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
max_request_size: int = 20 * 1024 * 1024
|
||||
):
|
||||
"""
|
||||
初始化客户端
|
||||
|
||||
Args:
|
||||
base_url: API 基础 URL
|
||||
api_key: API 密钥
|
||||
max_request_size: 最大请求体大小（字节），默认 20MB
|
||||
"""
|
||||
self.base_url = base_url
|
||||
self.api_key = api_key
|
||||
self.max_request_size = max_request_size
|
||||
|
||||
@abstractmethod
|
||||
def get_endpoint(self, **kwargs) -> str:
|
||||
"""
|
||||
获取 API 端点路径
|
||||
|
||||
Args:
|
||||
**kwargs: 额外参数（如模型名、分辨率等）
|
||||
|
||||
Returns:
|
||||
端点路径字符串
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def build_request_body(self, **kwargs) -> Dict[str, Any]:
|
||||
"""
|
||||
构建 API 请求体
|
||||
|
||||
Args:
|
||||
**kwargs: 请求参数
|
||||
|
||||
Returns:
|
||||
请求体字典
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def parse_response(self, response: Dict[str, Any]) -> Any:
|
||||
"""
|
||||
解析 API 响应
|
||||
|
||||
Args:
|
||||
response: API 响应字典
|
||||
|
||||
Returns:
|
||||
解析后的结果
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_headers(self, use_bearer_token: bool = False) -> Dict[str, str]:
|
||||
"""
|
||||
获取请求头
|
||||
|
||||
Args:
|
||||
use_bearer_token: 是否使用 Bearer Token 认证（默认为 False）
|
||||
|
||||
Returns:
|
||||
请求头字典
|
||||
"""
|
||||
if use_bearer_token:
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"x-goog-api-key": self.api_key,
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
def check_request_size(self, request_body: Dict[str, Any]) -> None:
|
||||
"""
|
||||
检查请求体大小是否超过限制
|
||||
|
||||
Args:
|
||||
request_body: 请求体字典
|
||||
|
||||
Raises:
|
||||
ValueError: 如果请求体超过限制
|
||||
"""
|
||||
request_json = json.dumps(request_body)
|
||||
request_size = len(request_json.encode('utf-8'))
|
||||
|
||||
if request_size > self.max_request_size:
|
||||
size_mb = request_size / 1024 / 1024
|
||||
limit_mb = self.max_request_size / 1024 / 1024
|
||||
raise ValueError(
|
||||
f"请求体大小 {size_mb:.2f}MB 超过限制 {limit_mb:.0f}MB，"
|
||||
"请降低分辨率或减少图片数量"
|
||||
)
|
||||
|
||||
async def request_async(
|
||||
self,
|
||||
endpoint: str,
|
||||
request_body: Dict[str, Any],
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
use_bearer_token: bool = False,
|
||||
timeout: Optional[int] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
发送异步 HTTP 请求
|
||||
|
||||
Args:
|
||||
endpoint: API 端点
|
||||
request_body: 请求体
|
||||
session: aiohttp 会话（可选）
|
||||
use_bearer_token: 是否使用 Bearer Token 认证
|
||||
|
||||
Returns:
|
||||
响应 JSON
|
||||
|
||||
Raises:
|
||||
RuntimeError: 请求失败时
|
||||
"""
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token)
|
||||
|
||||
# 检查请求大小
|
||||
self.check_request_size(request_body)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
# 设置超时
|
||||
timeout_obj = aiohttp.ClientTimeout(total=timeout) if timeout else None
|
||||
async with session.post(url, json=request_body, headers=headers, timeout=timeout_obj) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
|
||||
# 针对常见错误状态码提供友好提示
|
||||
if response.status == 504:
|
||||
raise RuntimeError(
|
||||
f"API 请求超时 (504 Gateway Timeout)\n"
|
||||
f"原因：服务器响应超时或该端点暂时不可用\n"
|
||||
f"建议：\n"
|
||||
f" - 尝试使用其他模型\n"
|
||||
f" - 稍后重试\n"
|
||||
f" - 降低分辨率或减少输入图像数量\n"
|
||||
f"详细错误: {error_text[:200]}"
|
||||
)
|
||||
elif response.status == 503:
|
||||
raise RuntimeError(
|
||||
f"服务暂时不可用 (503 Service Unavailable)\n"
|
||||
f"原因：模型服务过载或维护中\n"
|
||||
f"建议：\n"
|
||||
f" - 稍后重试\n"
|
||||
f" - 尝试使用其他模型"
|
||||
)
|
||||
elif response.status == 429:
|
||||
raise RuntimeError(
|
||||
f"请求频率超限 (429 Too Many Requests)\n"
|
||||
f"原因：API 配额用尽或请求过于频繁\n"
|
||||
f"建议：\n"
|
||||
f" - 等待一段时间后重试\n"
|
||||
f" - 检查 API 配额是否充足"
|
||||
)
|
||||
elif response.status == 404:
|
||||
raise RuntimeError(
|
||||
f"端点不存在 (404 Not Found)\n"
|
||||
f"原因：API 端点路径错误或模型不存在\n"
|
||||
f"建议：\n"
|
||||
f" - 检查模型名称是否正确\n"
|
||||
f" - 使用其他可用模型"
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"API 请求失败 (状态码: {response.status}): {error_text}"
|
||||
)
|
||||
|
||||
return await response.json()
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
async def request_get_async(
|
||||
self,
|
||||
endpoint: str,
|
||||
session: Optional[aiohttp.ClientSession] = None,
|
||||
use_bearer_token: bool = True,
|
||||
timeout: Optional[int] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
发送异步 HTTP GET 请求
|
||||
|
||||
Args:
|
||||
endpoint: API 端点
|
||||
session: aiohttp 会话（可选）
|
||||
use_bearer_token: 是否使用 Bearer Token 认证（默认为 True）
|
||||
|
||||
Returns:
|
||||
响应 JSON
|
||||
|
||||
Raises:
|
||||
RuntimeError: 请求失败时
|
||||
"""
|
||||
url = f"{self.base_url}{endpoint}"
|
||||
headers = self.get_headers(use_bearer_token)
|
||||
|
||||
close_session = False
|
||||
if session is None:
|
||||
session = aiohttp.ClientSession()
|
||||
close_session = True
|
||||
|
||||
try:
|
||||
# 设置超时
|
||||
timeout_obj = aiohttp.ClientTimeout(total=timeout) if timeout else None
|
||||
async with session.get(url, headers=headers, timeout=timeout_obj) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
|
||||
# 针对常见错误状态码提供友好提示
|
||||
if response.status == 504:
|
||||
raise RuntimeError(
|
||||
f"API 请求超时 (504 Gateway Timeout)\n"
|
||||
f"原因：服务器响应超时或该端点暂时不可用\n"
|
||||
f"建议：稍后重试"
|
||||
)
|
||||
elif response.status == 503:
|
||||
raise RuntimeError(
|
||||
f"服务暂时不可用 (503 Service Unavailable)\n"
|
||||
f"原因：服务过载或维护中\n"
|
||||
f"建议：稍后重试"
|
||||
)
|
||||
elif response.status == 429:
|
||||
raise RuntimeError(
|
||||
f"请求频率超限 (429 Too Many Requests)\n"
|
||||
f"原因：API 配额用尽或请求过于频繁\n"
|
||||
f"建议：等待一段时间后重试"
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"API 请求失败 (状态码: {response.status}): {error_text}"
|
||||
)
|
||||
|
||||
return await response.json()
|
||||
|
||||
finally:
|
||||
if close_session:
|
||||
await session.close()
|
||||
|
||||
async def batch_request_async(
|
||||
self,
|
||||
requests: List[Dict[str, Any]],
|
||||
progress_callback: Optional[Callable[[int, int], None]] = None
|
||||
) -> List[Any]:
|
||||
"""
|
||||
批量并发请求
|
||||
|
||||
Args:
|
||||
requests: 请求列表，每个元素包含 endpoint 和 request_body
|
||||
progress_callback: 进度回调函数 (current, total)
|
||||
|
||||
Returns:
|
||||
响应结果列表
|
||||
"""
|
||||
results = []
|
||||
completed = 0
|
||||
total = len(requests)
|
||||
|
||||
# 创建无限制的连接器
|
||||
connector = aiohttp.TCPConnector(limit=0, limit_per_host=0)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector) as session:
|
||||
tasks = []
|
||||
|
||||
for req in requests:
|
||||
task = self.request_async(
|
||||
endpoint=req['endpoint'],
|
||||
request_body=req['request_body'],
|
||||
session=session
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# 并发执行
|
||||
responses = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
for i, resp in enumerate(responses):
|
||||
if isinstance(resp, Exception):
|
||||
print(f"⚠️ 第 {i+1} 个请求失败: {str(resp)}")
|
||||
continue
|
||||
|
||||
try:
|
||||
parsed = self.parse_response(resp)
|
||||
results.append(parsed)
|
||||
completed += 1
|
||||
|
||||
if progress_callback:
|
||||
progress_callback(completed, total)
|
||||
|
||||
except Exception as e:
|
||||
print(f"⚠️ 第 {i+1} 个响应解析失败: {str(e)}")
|
||||
|
||||
return results
|
||||
|
||||
def run_async_in_thread(self, coro) -> Any:
|
||||
"""
|
||||
在独立线程中运行异步代码（用于 ComfyUI 同步接口）
|
||||
|
||||
Args:
|
||||
coro: 协程对象
|
||||
|
||||
Returns:
|
||||
协程执行结果
|
||||
"""
|
||||
result_container = []
|
||||
error_container = []
|
||||
|
||||
def run_in_thread():
|
||||
try:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
try:
|
||||
result = loop.run_until_complete(coro)
|
||||
result_container.append(result)
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
except Exception as e:
|
||||
error_container.append(e)
|
||||
|
||||
thread = threading.Thread(target=run_in_thread)
|
||||
thread.start()
|
||||
thread.join()
|
||||
|
||||
if error_container:
|
||||
raise error_container[0]
|
||||
|
||||
if not result_container:
|
||||
raise RuntimeError("异步任务未返回结果")
|
||||
|
||||
return result_container[0]
|
||||
Reference in New Issue
Block a user