增加rpm,tpm限制

This commit is contained in:
xunbu
2025-12-17 21:53:15 +08:00
parent 579f0b8a9c
commit c5658afcd9
11 changed files with 48 additions and 21 deletions

View File

@@ -122,7 +122,6 @@ class RateLimiter:
# Check RPM
if self.rpm and len(self.request_timestamps) >= self.rpm:
# 取最早的一条记录,计算还需要等待多久才能腾出位置
earliest = self.request_timestamps[0]
wait_time = max(wait_time, 60 - (now - earliest))
@@ -130,14 +129,10 @@ class RateLimiter:
if self.tpm:
current_tokens = sum(t[1] for t in self.token_timestamps)
if current_tokens + tokens > self.tpm:
# 稍微复杂点需要移除足够多的旧token才能放入新token
# 这里做一个简化估算:如果超限,等到最早的记录过期
if self.token_timestamps:
earliest = self.token_timestamps[0][0]
wait_time = max(wait_time, 60 - (now - earliest))
else:
# 这种情况理论上不应该发生除非单次请求超过了TPM上限
# 如果单次超过上限强制等待1秒防止死循环并允许通过(或者抛异常,这里选择允许)
pass
return wait_time
@@ -156,14 +151,19 @@ class RateLimiter:
return
while True:
# print(f"[RateLimiter-Async] 准备获取锁...")
with self.lock:
print(f"[RateLimiter-Async] 已加锁 (Checking)")
wait_time = self._check_and_get_wait_time(tokens)
if wait_time <= 0:
self._record_usage(tokens)
print(f"[RateLimiter-Async] 释放锁 (成功获取配额)")
return
# 释放锁后等待,避免阻塞其他协程/线程的检查
# 添加一点点缓冲时间,避免刚唤醒时毫秒级误差导致再次等待
print(f"[RateLimiter-Async] 释放锁 (需等待 {wait_time:.2f}s)")
# 释放锁后等待
await asyncio.sleep(wait_time + 0.1)
def acquire_sync(self, tokens: int = 0):
@@ -172,12 +172,18 @@ class RateLimiter:
return
while True:
# print(f"[RateLimiter-Sync] 准备获取锁...")
with self.lock:
print(f"[RateLimiter-Sync] 已加锁 (Checking)")
wait_time = self._check_and_get_wait_time(tokens)
if wait_time <= 0:
self._record_usage(tokens)
print(f"[RateLimiter-Sync] 释放锁 (成功获取配额)")
return
print(f"[RateLimiter-Sync] 释放锁 (需等待 {wait_time:.2f}s)")
time.sleep(wait_time + 0.1)

View File

@@ -40,7 +40,9 @@ class AssTranslator(AiTranslator):
glossary_dict=config.glossary_dict,
retry=config.retry,
system_proxy_enable=config.system_proxy_enable,
force_json=config.force_json
force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.insert_mode = config.insert_mode

View File

@@ -129,7 +129,9 @@ class DocxTranslator(AiTranslator):
api_key=config.api_key, model_id=config.model_id, temperature=config.temperature,
thinking=config.thinking, concurrent=config.concurrent, timeout=config.timeout,
logger=self.logger, glossary_dict=config.glossary_dict, retry=config.retry,
system_proxy_enable=config.system_proxy_enable, force_json=config.force_json
system_proxy_enable=config.system_proxy_enable, force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.insert_mode = config.insert_mode

View File

@@ -51,7 +51,9 @@ class EpubTranslator(AiTranslator):
glossary_dict=config.glossary_dict,
retry=config.retry,
system_proxy_enable=config.system_proxy_enable,
force_json=config.force_json
force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.insert_mode = config.insert_mode

View File

@@ -69,7 +69,9 @@ class HtmlTranslator(AiTranslator):
glossary_dict=config.glossary_dict,
retry=config.retry,
system_proxy_enable=config.system_proxy_enable,
force_json=config.force_json
force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.insert_mode = config.insert_mode

View File

@@ -37,7 +37,9 @@ class JsonTranslator(AiTranslator):
glossary_dict=config.glossary_dict,
retry=config.retry,
system_proxy_enable=config.system_proxy_enable,
force_json=config.force_json
force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.json_paths = config.json_paths

View File

@@ -40,7 +40,10 @@ class MDTranslator(AiTranslator):
logger=self.logger,
glossary_dict=config.glossary_dict,
retry=config.retry,
system_proxy_enable=config.system_proxy_enable)
system_proxy_enable=config.system_proxy_enable,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = MDTranslateAgent(agent_config)
def translate(self, document: MarkdownDocument) -> Self:

View File

@@ -3,14 +3,14 @@
import asyncio
from dataclasses import dataclass
from io import BytesIO
from typing import Self, Literal, List, Dict, Any, Tuple, Optional
from typing import Self, Literal, List, Dict, Any, Tuple
from pptx import Presentation
from pptx.enum.dml import MSO_COLOR_TYPE
from pptx.enum.shapes import MSO_SHAPE_TYPE
from pptx.enum.text import MSO_AUTO_SIZE
from pptx.enum.dml import MSO_COLOR_TYPE
from pptx.text.text import _Paragraph, TextFrame
from pptx.oxml.ns import qn
from pptx.text.text import _Paragraph, TextFrame
from docutranslate.agents.segments_agent import SegmentsTranslateAgentConfig, SegmentsTranslateAgent
from docutranslate.ir.document import Document
@@ -46,7 +46,9 @@ class PPTXTranslator(AiTranslator):
api_key=config.api_key, model_id=config.model_id, temperature=config.temperature,
thinking=config.thinking, concurrent=config.concurrent, timeout=config.timeout,
logger=self.logger, glossary_dict=config.glossary_dict, retry=config.retry,
system_proxy_enable=config.system_proxy_enable, force_json=config.force_json
system_proxy_enable=config.system_proxy_enable, force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.insert_mode = config.insert_mode
@@ -335,4 +337,4 @@ class PPTXTranslator(AiTranslator):
translated = await self.translate_agent.send_segments_async(originals,
self.chunk_size) if self.translate_agent else originals
document.content = await asyncio.to_thread(self._after_translate, prs, elements, translated, originals)
return self
return self

View File

@@ -42,7 +42,9 @@ class SrtTranslator(AiTranslator):
glossary_dict=config.glossary_dict,
retry=config.retry,
system_proxy_enable=config.system_proxy_enable,
force_json=config.force_json
force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.insert_mode = config.insert_mode

View File

@@ -74,7 +74,9 @@ class TXTTranslator(AiTranslator):
glossary_dict=config.glossary_dict,
retry=config.retry,
system_proxy_enable=config.system_proxy_enable,
force_json=config.force_json
force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.insert_mode = config.insert_mode

View File

@@ -43,7 +43,9 @@ class XlsxTranslator(AiTranslator):
glossary_dict=config.glossary_dict,
retry=config.retry,
system_proxy_enable=config.system_proxy_enable,
force_json=config.force_json
force_json=config.force_json,
rpm=config.rpm,
tpm=config.tpm
)
self.translate_agent = SegmentsTranslateAgent(agent_config)
self.insert_mode = config.insert_mode