369 lines
12 KiB
Python
369 lines
12 KiB
Python
"""
|
||
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]
|