GithubBackup/gemini-balance
1
0
Fork 0
mirror of https://github.com/snailyp/gemini-balance.git synced 2026-09-04 23:19:17 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 添加自定义 Headers 功能

Browse Source 
- 在配置中添加 `CUSTOM_HEADERS` 选项，允许用户定义全局请求头。
- 更新 API 客户端，将自定义 `header` 应用于所有出站请求。
- 在配置页面上为 `CUSTOM_HEADERS` 添加了完整的前端编辑功能。
This commit is contained in:
snaily
2025-07-08 13:58:05 +08:00
parent 0a08913677
commit 231b69ecf8
4 changed files with 147 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/service/client/api_client.py
+27 -9
View File
@@ -44,6 +44,13 @@ class GeminiApiClient(ApiClient):
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)
@@ -56,10 +63,11 @@ class GeminiApiClient(ApiClient):
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)
response = await client.get(url, headers=headers)
response.raise_for_status()
return response.json()
except httpx.HTTPStatusError as e:
@@ -82,9 +90,10 @@ class GeminiApiClient(ApiClient):
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}"
response = await client.post(url, json=payload)
response = await client.post(url, json=payload, headers=headers)
if response.status_code != 200:
error_content = response.text
raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
@@ -102,9 +111,10 @@ class GeminiApiClient(ApiClient):
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) as response:
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")
@@ -124,9 +134,10 @@ class GeminiApiClient(ApiClient):
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)
response = await client.post(url, json=payload, headers=headers)
if response.status_code != 200:
error_content = response.text
raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
@@ -140,6 +151,13 @@ class OpenaiApiClient(ApiClient):
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)
@@ -151,9 +169,9 @@ class OpenaiApiClient(ApiClient):
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"
headers = {"Authorization": f"Bearer {api_key}"}
response = await client.get(url, headers=headers)
if response.status_code != 200:
error_content = response.text
@@ -171,9 +189,9 @@ class OpenaiApiClient(ApiClient):
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"
headers = {"Authorization": f"Bearer {api_key}"}
response = await client.post(url, json=payload, headers=headers)
if response.status_code != 200:
error_content = response.text
@@ -190,9 +208,9 @@ class OpenaiApiClient(ApiClient):
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"
headers = {"Authorization": f"Bearer {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()
@@ -212,9 +230,9 @@ class OpenaiApiClient(ApiClient):
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"
headers = {"Authorization": f"Bearer {api_key}"}
payload = {
"input": input,
"model": model,
@@ -236,9 +254,9 @@ class OpenaiApiClient(ApiClient):
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"
headers = {"Authorization": f"Bearer {api_key}"}
response = await client.post(url, json=payload, headers=headers)
if response.status_code != 200:
error_content = response.text