mirror of
https://github.com/snailyp/gemini-balance.git
synced 2026-05-16 07:17:36 +08:00
67dd1af583
这次提交重构了整个应用的异常处理机制，保证了处理方式的一致性，还能提供更详细的错误信息。 主要改动包括： - 修改了 `ApiClient`，现在抛出的异常会同时包含状态码和消息。这样上游服务就能传递准确的 HTTP 错误响应啦。 - 更新了所有服务层（`gemini`、`openai`、`vertex`、`embedding`），现在会捕获这些结构化的异常，不再从字符串里解析错误消息了。 - 增强了路由级别的错误处理，特别是针对流式端点，能正确捕获初始化错误，并返回结构化的 JSON 错误响应，而不是格式错误的 SSE 事件。 - 在所有 API 路由中添加了 `allowed_token` 的日志记录，方便追踪和调试授权问题。 - 还有一些常规的代码清理，比如调整了 import 顺序和格式化代码，提高了可读性和可维护性。
394 lines
16 KiB
Python
394 lines
16 KiB
Python
# app/services/chat/api_client.py
|
|
|
|
import random
|
|
from abc import ABC, abstractmethod
|
|
from typing import Any, AsyncGenerator, Dict, Optional
|
|
|
|
import httpx
|
|
|
|
from app.config.config import settings
|
|
from app.core.constants import DEFAULT_TIMEOUT
|
|
from app.log.logger import get_api_client_logger
|
|
|
|
logger = get_api_client_logger()
|
|
|
|
|
|
class ApiClient(ABC):
|
|
"""API客户端基类"""
|
|
|
|
@abstractmethod
|
|
async def generate_content(
|
|
self, payload: Dict[str, Any], model: str, api_key: str
|
|
) -> Dict[str, Any]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
async def stream_generate_content(
|
|
self, payload: Dict[str, Any], model: str, api_key: str
|
|
) -> AsyncGenerator[str, None]:
|
|
pass
|
|
|
|
|
|
class GeminiApiClient(ApiClient):
|
|
"""Gemini API客户端"""
|
|
|
|
def __init__(self, base_url: str, timeout: int = DEFAULT_TIMEOUT):
|
|
self.base_url = base_url
|
|
self.timeout = timeout
|
|
|
|
def _get_real_model(self, model: str) -> str:
|
|
if model.endswith("-search"):
|
|
model = model[:-7]
|
|
if model.endswith("-image"):
|
|
model = model[:-6]
|
|
if model.endswith("-non-thinking"):
|
|
model = model[:-13]
|
|
if "-search" in model and "-non-thinking" in model:
|
|
model = model[:-20]
|
|
return model
|
|
|
|
def _prepare_headers(self) -> Dict[str, str]:
|
|
headers = {}
|
|
if settings.CUSTOM_HEADERS:
|
|
headers.update(settings.CUSTOM_HEADERS)
|
|
logger.info(f"Using custom headers: {settings.CUSTOM_HEADERS}")
|
|
return headers
|
|
|
|
async def get_models(self, api_key: str) -> Optional[Dict[str, Any]]:
|
|
"""获取可用的 Gemini 模型列表"""
|
|
timeout = httpx.Timeout(timeout=5)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for getting models: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers()
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/models?key={api_key}&pageSize=1000"
|
|
try:
|
|
response = await client.get(url, headers=headers)
|
|
response.raise_for_status()
|
|
return response.json()
|
|
except httpx.HTTPStatusError as e:
|
|
logger.error(f"获取模型列表失败: {e.response.status_code}")
|
|
logger.error(e.response.text)
|
|
return None
|
|
except httpx.RequestError as e:
|
|
logger.error(f"请求模型列表失败: {e}")
|
|
return None
|
|
|
|
async def generate_content(
|
|
self, payload: Dict[str, Any], model: str, api_key: str
|
|
) -> Dict[str, Any]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
model = self._get_real_model(model)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for getting models: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers()
|
|
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/models/{model}:generateContent?key={api_key}"
|
|
|
|
try:
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
logger.error(
|
|
f"API call failed - Status: {response.status_code}, Content: {error_content}"
|
|
)
|
|
raise Exception(response.status_code, error_content)
|
|
|
|
response_data = response.json()
|
|
|
|
# 检查响应结构的基本信息
|
|
if not response_data.get("candidates"):
|
|
logger.warning("No candidates found in API response")
|
|
|
|
return response_data
|
|
|
|
except httpx.TimeoutException as e:
|
|
logger.error(f"Request timeout: {e}")
|
|
raise Exception(500, f"Request timeout: {e}")
|
|
except httpx.RequestError as e:
|
|
logger.error(f"Request error: {e}")
|
|
raise Exception(500, f"Request error: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Unexpected error: {e}")
|
|
raise Exception(500, f"Unexpected error: {e}")
|
|
|
|
async def stream_generate_content(
|
|
self, payload: Dict[str, Any], model: str, api_key: str
|
|
) -> AsyncGenerator[str, None]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
model = self._get_real_model(model)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for getting models: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers()
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse&key={api_key}"
|
|
async with client.stream(
|
|
method="POST", url=url, json=payload, headers=headers
|
|
) as response:
|
|
if response.status_code != 200:
|
|
error_content = await response.aread()
|
|
error_msg = error_content.decode("utf-8")
|
|
raise Exception(response.status_code, error_msg)
|
|
async for line in response.aiter_lines():
|
|
yield line
|
|
|
|
async def count_tokens(
|
|
self, payload: Dict[str, Any], model: str, api_key: str
|
|
) -> Dict[str, Any]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
model = self._get_real_model(model)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for counting tokens: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers()
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/models/{model}:countTokens?key={api_key}"
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
raise Exception(response.status_code, error_content)
|
|
return response.json()
|
|
|
|
async def embed_content(
|
|
self, payload: Dict[str, Any], model: str, api_key: str
|
|
) -> Dict[str, Any]:
|
|
"""单一嵌入内容生成"""
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
model = self._get_real_model(model)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for embedding: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers()
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/models/{model}:embedContent?key={api_key}"
|
|
|
|
try:
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
logger.error(
|
|
f"Embedding API call failed - Status: {response.status_code}, Content: {error_content}"
|
|
)
|
|
raise Exception(response.status_code, error_content)
|
|
|
|
return response.json()
|
|
|
|
except httpx.TimeoutException as e:
|
|
logger.error(f"Embedding request timeout: {e}")
|
|
raise Exception(500, f"Request timeout: {e}")
|
|
except httpx.RequestError as e:
|
|
logger.error(f"Embedding request error: {e}")
|
|
raise Exception(500, f"Request error: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Unexpected embedding error: {e}")
|
|
raise Exception(500, f"Unexpected embedding error: {e}")
|
|
|
|
async def batch_embed_contents(
|
|
self, payload: Dict[str, Any], model: str, api_key: str
|
|
) -> Dict[str, Any]:
|
|
"""批量嵌入内容生成"""
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
model = self._get_real_model(model)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for batch embedding: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers()
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/models/{model}:batchEmbedContents?key={api_key}"
|
|
|
|
try:
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
logger.error(
|
|
f"Batch embedding API call failed - Status: {response.status_code}, Content: {error_content}"
|
|
)
|
|
raise Exception(response.status_code, error_content)
|
|
|
|
return response.json()
|
|
|
|
except httpx.TimeoutException as e:
|
|
logger.error(f"Batch embedding request timeout: {e}")
|
|
raise Exception(500, f"Request timeout: {e}")
|
|
except httpx.RequestError as e:
|
|
logger.error(f"Batch embedding request error: {e}")
|
|
raise Exception(500, f"Request error: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Unexpected batch embedding error: {e}")
|
|
raise Exception(500, f"Unexpected batch embedding error: {e}")
|
|
|
|
|
|
class OpenaiApiClient(ApiClient):
|
|
"""OpenAI API客户端"""
|
|
|
|
def __init__(self, base_url: str, timeout: int = DEFAULT_TIMEOUT):
|
|
self.base_url = base_url
|
|
self.timeout = timeout
|
|
|
|
def _prepare_headers(self, api_key: str) -> Dict[str, str]:
|
|
headers = {"Authorization": f"Bearer {api_key}"}
|
|
if settings.CUSTOM_HEADERS:
|
|
headers.update(settings.CUSTOM_HEADERS)
|
|
logger.info(f"Using custom headers: {settings.CUSTOM_HEADERS}")
|
|
return headers
|
|
|
|
async def get_models(self, api_key: str) -> Dict[str, Any]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for getting models: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers(api_key)
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/openai/models"
|
|
response = await client.get(url, headers=headers)
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
raise Exception(response.status_code, error_content)
|
|
return response.json()
|
|
|
|
async def generate_content(
|
|
self, payload: Dict[str, Any], api_key: str
|
|
) -> Dict[str, Any]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
logger.info(
|
|
f"settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: {settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY}"
|
|
)
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for getting models: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers(api_key)
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/openai/chat/completions"
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
raise Exception(response.status_code, error_content)
|
|
return response.json()
|
|
|
|
async def stream_generate_content(
|
|
self, payload: Dict[str, Any], api_key: str
|
|
) -> AsyncGenerator[str, None]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for getting models: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers(api_key)
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/openai/chat/completions"
|
|
async with client.stream(
|
|
method="POST", url=url, json=payload, headers=headers
|
|
) as response:
|
|
if response.status_code != 200:
|
|
error_content = await response.aread()
|
|
error_msg = error_content.decode("utf-8")
|
|
raise Exception(response.status_code, error_msg)
|
|
async for line in response.aiter_lines():
|
|
yield line
|
|
|
|
async def create_embeddings(
|
|
self, input: str, model: str, api_key: str
|
|
) -> Dict[str, Any]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for getting models: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers(api_key)
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/openai/embeddings"
|
|
payload = {
|
|
"input": input,
|
|
"model": model,
|
|
}
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
raise Exception(response.status_code, error_content)
|
|
return response.json()
|
|
|
|
async def generate_images(
|
|
self, payload: Dict[str, Any], api_key: str
|
|
) -> Dict[str, Any]:
|
|
timeout = httpx.Timeout(self.timeout, read=self.timeout)
|
|
|
|
proxy_to_use = None
|
|
if settings.PROXIES:
|
|
if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
|
|
proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
|
|
else:
|
|
proxy_to_use = random.choice(settings.PROXIES)
|
|
logger.info(f"Using proxy for getting models: {proxy_to_use}")
|
|
|
|
headers = self._prepare_headers(api_key)
|
|
async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
|
|
url = f"{self.base_url}/openai/images/generations"
|
|
response = await client.post(url, json=payload, headers=headers)
|
|
if response.status_code != 200:
|
|
error_content = response.text
|
|
raise Exception(response.status_code, error_content)
|
|
return response.json()
|