chore: 更新版本号至2.1.8

- 新增CLOUDFLARE_IMGBED_UPLOAD_FOLDER环境变量配置
- 用户可通过该配置项指定图片在CloudFlare图床中的上传路径

.env.example

@@ -43,6 +43,7 @@ SMMS_SECRET_TOKEN=XXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
 PICGO_API_KEY=xxxx
 CLOUDFLARE_IMGBED_URL=https://xxxxxxx.pages.dev/upload
 CLOUDFLARE_IMGBED_AUTH_CODE=xxxxxxxxx
+CLOUDFLARE_IMGBED_UPLOAD_FOLDER=
 ##########################################################################
 #########################stream_optimizer 相关配置########################
 STREAM_OPTIMIZER_ENABLED=false

Dockerfile

@@ -15,6 +15,7 @@ ENV TOOLS_CODE_EXECUTION_ENABLED=false
 ENV IMAGE_MODELS='["gemini-2.0-flash-exp"]'
 ENV SEARCH_MODELS='["gemini-2.0-flash-exp","gemini-2.0-pro-exp"]'
 ENV URL_NORMALIZATION_ENABLED=false
+ENV CLOUDFLARE_IMGBED_UPLOAD_FOLDER=""
 
 # Expose port
 EXPOSE 8000

README.md

@@ -2,6 +2,12 @@
 
 # Gemini Balance - Gemini API Proxy and Load Balancer
 
+<p align="center">
+  <a href="https://trendshift.io/repositories/13692" target="_blank">
+    <img src="https://trendshift.io/api/badge/repositories/13692" alt="snailyp%2Fgemini-balance | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
+  </a>
+</p>
+
 > ⚠️ This project is licensed under the CC BY-NC 4.0 (Attribution-NonCommercial) license. Any form of commercial resale service is prohibited. See the LICENSE file for details.
 
 > I have never sold this service on any platform. If you encounter someone selling this service, they are definitely a reseller. Please be careful not to be deceived.
@@ -211,6 +217,7 @@ If you want to run the source code directly locally for development or testing,
 | `PICGO_API_KEY`                | Optional, API Key for [PicoGo](https://www.picgo.net/) image hosting        | `your-picogo-apikey`                                                                                                                                                                                                                         |
 | `CLOUDFLARE_IMGBED_URL`        | Optional, [CloudFlare](https://github.com/MarSeventh/CloudFlare-ImgBed) image hosting upload address | `https://xxxxxxx.pages.dev/upload`                                                                                                                                                                                           |
 | `CLOUDFLARE_IMGBED_AUTH_CODE`  | Optional, authentication key for CloudFlare image hosting                   | `your-cloudflare-imgber-auth-code`                                                                                                                                                                                                           |
+| `CLOUDFLARE_IMGBED_UPLOAD_FOLDER` | Optional, upload folder path for CloudFlare image hosting               | `""`                                                                                                                                                                                                                                          |
 | **Stream Optimizer Related**   |                                                                             |                                                                                                                                                                                                                                          |
 | `STREAM_OPTIMIZER_ENABLED`     | Optional, whether to enable stream output optimization                      | `false`                                                                                                                                                                                                                                  |
 | `STREAM_MIN_DELAY`             | Optional, minimum delay for stream output                                   | `0.016`                                                                                                                                                                                                                                  |

README_ZH.md

@@ -1,5 +1,11 @@
 # Gemini Balance - Gemini API 代理和负载均衡器
 
+<p align="center">
+  <a href="https://trendshift.io/repositories/13692" target="_blank">
+    <img src="https://trendshift.io/api/badge/repositories/13692" alt="snailyp%2Fgemini-balance | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
+  </a>
+</p>
+
 > ⚠️ 本项目采用 CC BY-NC 4.0（署名-非商业性使用）协议，禁止任何形式的商业倒卖服务，详见 LICENSE 文件。
 
 > 本人从未在各个平台售卖服务，如有遇到售卖此服务者，那一定是倒卖狗，大家切记不要上当受骗。
@@ -204,6 +210,7 @@ app/
 | `PICGO_API_KEY`              | 可选，[PicoGo](https://www.picgo.net/)图床的API Key                                      | `your-picogo-apikey`                                  |
 | `CLOUDFLARE_IMGBED_URL`      | 可选，[CloudFlare](https://github.com/MarSeventh/CloudFlare-ImgBed) 图床上传地址                                  | `https://xxxxxxx.pages.dev/upload`                    |
 | `CLOUDFLARE_IMGBED_AUTH_CODE`| 可选，CloudFlare图床的鉴权key                                  | `your-cloudflare-imgber-auth-code`                    |
+| `CLOUDFLARE_IMGBED_UPLOAD_FOLDER`| 可选，CloudFlare图床的上传文件夹路径                        | `""`                                                  |
 | **流式优化器相关**           |                                                          |                                                       |
 | `STREAM_OPTIMIZER_ENABLED`   | 可选，是否启用流式输出优化                                     | `false`                                               |
 | `STREAM_MIN_DELAY`           | 可选，流式输出最小延迟                                         | `0.016`                                               |

VERSION

@@ -1 +1 @@
-2.1.7
+2.1.8

app/config/config.py

@@ -90,6 +90,7 @@ class Settings(BaseSettings):
     PICGO_API_KEY: str = ""
     CLOUDFLARE_IMGBED_URL: str = ""
     CLOUDFLARE_IMGBED_AUTH_CODE: str = ""
+    CLOUDFLARE_IMGBED_UPLOAD_FOLDER: str = ""
 
     # 流式输出优化器配置
     STREAM_OPTIMIZER_ENABLED: bool = False

app/handler/response_handler.py

@@ -238,6 +238,7 @@ def _extract_image_data(part: dict) -> str:
             provider=settings.UPLOAD_PROVIDER,
             base_url=settings.CLOUDFLARE_IMGBED_URL,
             auth_code=settings.CLOUDFLARE_IMGBED_AUTH_CODE,
+            upload_folder=settings.CLOUDFLARE_IMGBED_UPLOAD_FOLDER,
         )
     current_date = time.strftime("%Y/%m/%d")
     filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"

app/router/gemini_routes.py

@@ -151,6 +151,35 @@ async def stream_generate_content(
         return StreamingResponse(response_stream, media_type="text/event-stream")
 
 
+@router.post("/models/{model_name}:countTokens")
+@router_v1beta.post("/models/{model_name}:countTokens")
+@RetryHandler(key_arg="api_key")
+async def count_tokens(
+    model_name: str,
+    request: GeminiRequest,
+    _=Depends(security_service.verify_key_or_goog_api_key),
+    api_key: str = Depends(get_next_working_key),
+    key_manager: KeyManager = Depends(get_key_manager),
+    chat_service: GeminiChatService = Depends(get_chat_service)
+):
+    """处理 Gemini token 计数请求。"""
+    operation_name = "gemini_count_tokens"
+    async with handle_route_errors(logger, operation_name, failure_message="Token counting failed"):
+        logger.info(f"Handling Gemini token count request for model: {model_name}")
+        logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
+        logger.info(f"Using API key: {api_key}")
+
+        if not await model_service.check_model_support(model_name):
+            raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported")
+
+        response = await chat_service.count_tokens(
+            model=model_name,
+            request=request,
+            api_key=api_key
+        )
+        return response
+
+
 @router.post("/reset-all-fail-counts")
 async def reset_all_key_fail_counts(key_type: str = None, key_manager: KeyManager = Depends(get_key_manager)):
     """批量重置Gemini API密钥的失败计数，可选择性地仅重置有效或无效密钥"""

app/service/chat/gemini_chat_service.py

@@ -195,6 +195,54 @@ class GeminiChatService:
                 request_time=request_datetime
             )
 
+    async def count_tokens(
+        self, model: str, request: GeminiRequest, api_key: str
+    ) -> Dict[str, Any]:
+        """计算token数量"""
+        # countTokens API只需要contents
+        payload = {"contents": request.model_dump().get("contents", [])}
+        start_time = time.perf_counter()
+        request_datetime = datetime.datetime.now()
+        is_success = False
+        status_code = None
+        response = None
+
+        try:
+            response = await self.api_client.count_tokens(payload, model, api_key)
+            is_success = True
+            status_code = 200
+            return response
+        except Exception as e:
+            is_success = False
+            error_log_msg = str(e)
+            logger.error(f"Count tokens API call failed with error: {error_log_msg}")
+            match = re.search(r"status code (\d+)", error_log_msg)
+            if match:
+                status_code = int(match.group(1))
+            else:
+                status_code = 500
+
+            await add_error_log(
+                gemini_key=api_key,
+                model_name=model,
+                error_type="gemini-count-tokens",
+                error_log=error_log_msg,
+                error_code=status_code,
+                request_msg=payload
+            )
+            raise e
+        finally:
+            end_time = time.perf_counter()
+            latency_ms = int((end_time - start_time) * 1000)
+            await add_request_log(
+                model_name=model,
+                api_key=api_key,
+                is_success=is_success,
+                status_code=status_code,
+                latency_ms=latency_ms,
+                request_time=request_datetime
+            )
+
     async def stream_generate_content(
         self, model: str, request: GeminiRequest, api_key: str
     ) -> AsyncGenerator[str, None]:

app/service/client/api_client.py

@@ -21,6 +21,10 @@ class ApiClient(ABC):
     async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]:
         pass
 
+    @abstractmethod
+    async def count_tokens(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
+        pass
+
 
 class GeminiApiClient(ApiClient):
     """Gemini API客户端"""
@@ -108,6 +112,26 @@ class GeminiApiClient(ApiClient):
                 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}")
+
+        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)
+            if response.status_code != 200:
+                error_content = response.text
+                raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
+            return response.json()
+
 
 class OpenaiApiClient(ApiClient):
     """OpenAI API客户端"""

app/service/image/image_create_service.py

@@ -121,6 +121,7 @@ class ImageCreateService:
                             provider=settings.UPLOAD_PROVIDER,
                             base_url=settings.CLOUDFLARE_IMGBED_URL,
                             auth_code=settings.CLOUDFLARE_IMGBED_AUTH_CODE,
+                            upload_folder=settings.CLOUDFLARE_IMGBED_UPLOAD_FOLDER,
                         )
                     else:
                         raise ValueError(

app/templates/config_editor.html

@@ -1602,6 +1602,23 @@ endblock %} {% block head_extra_styles %}
           />
           <small class="text-gray-500 mt-1 block">Cloudflare图床的认证码</small>
         </div>
+
+        <!-- Cloudflare上传文件夹 -->
+        <div class="mb-6 provider-config" data-provider="cloudflare_imgbed">
+          <label
+            for="CLOUDFLARE_IMGBED_UPLOAD_FOLDER"
+            class="block font-semibold mb-2 text-gray-700"
+            >Cloudflare上传文件夹</label
+          >
+          <input
+            type="text"
+            id="CLOUDFLARE_IMGBED_UPLOAD_FOLDER"
+            name="CLOUDFLARE_IMGBED_UPLOAD_FOLDER"
+            placeholder=""
+            class="w-full px-4 py-3 rounded-lg form-input-themed"
+          />
+          <small class="text-gray-500 mt-1 block">Cloudflare图床的上传文件夹路径（可选）</small>
+        </div>
       </div>
 
       <!-- 流式输出优化配置 -->

app/utils/uploader.py

@@ -261,18 +261,20 @@ class PicGoUploader(ImageUploader):
 
 class CloudFlareImgBedUploader(ImageUploader):
     """CloudFlare图床上传器"""
-    
-    def __init__(self, auth_code: str, api_url: str):
+
+    def __init__(self, auth_code: str, api_url: str, upload_folder: str = ""):
         """
         初始化CloudFlare图床上传器
         
         Args:
             auth_code: 认证码
             api_url: 上传API地址
+            upload_folder: 上传文件夹路径（可选）
         """
         self.auth_code = auth_code
         self.api_url = api_url
-        
+        self.upload_folder = upload_folder
+
     def upload(self, file: bytes, filename: str) -> UploadResponse:
         """
         上传图片到CloudFlare图床
@@ -288,12 +290,16 @@ class CloudFlareImgBedUploader(ImageUploader):
             UploadError: 上传失败时抛出异常
         """
         try:
-            # 准备请求URL（添加认证码参数，如果存在）
+            # 准备请求URL参数
+            params = []
+            if self.upload_folder:
+                params.append(f"uploadFolder={self.upload_folder}")
             if self.auth_code:
-                request_url = f"{self.api_url}?authCode={self.auth_code}&uploadNameType=origin"
-            else:
-                request_url = f"{self.api_url}?uploadNameType=origin"
-            
+                params.append(f"authCode={self.auth_code}")
+            params.append("uploadNameType=origin")
+
+            request_url = f"{self.api_url}?{'&'.join(params)}"
+
             # 准备文件数据
             files = {
                 "file": (filename, file)
@@ -388,6 +394,7 @@ class ImageUploaderFactory:
         elif provider == "cloudflare_imgbed":
             return CloudFlareImgBedUploader(
                 credentials["auth_code"],
-                credentials["base_url"]
+                credentials["base_url"],
+                credentials.get("upload_folder", ""),
             )
         raise ValueError(f"Unknown provider: {provider}")