156 lines
6.0 KiB
Python
156 lines
6.0 KiB
Python
"""
|
||
LLM 客户端模块 (LLM Client Module) - 简化 ID 锚点匹配版
|
||
|
||
该模块负责发送翻译请求，并使用简化后的 ID 作为锚点解析 LLM 的响应。
|
||
"""
|
||
|
||
import asyncio
|
||
from openai import AsyncOpenAI
|
||
from typing import List, Dict, Optional, Any
|
||
from loguru import logger
|
||
import time
|
||
import re
|
||
from .manifest_manager import ManifestItem
|
||
|
||
|
||
class RateLimiter:
|
||
"""并发与 RPM 速率限制器。"""
|
||
def __init__(self, requests_per_minute: int, concurrent_requests: int):
|
||
self.semaphore = asyncio.Semaphore(concurrent_requests)
|
||
self.min_interval = 60.0 / requests_per_minute if requests_per_minute > 0 else 0
|
||
self.last_request_time = 0
|
||
|
||
async def acquire(self):
|
||
await self.semaphore.acquire()
|
||
current_time = time.time()
|
||
wait_time = self.min_interval - (current_time - self.last_request_time)
|
||
if wait_time > 0:
|
||
await asyncio.sleep(wait_time)
|
||
self.last_request_time = time.time()
|
||
|
||
def release(self):
|
||
self.semaphore.release()
|
||
|
||
|
||
class OpenRouterClient:
|
||
"""基于简化 ID 锚点匹配逻辑的 LLM 客户端。"""
|
||
|
||
def __init__(self, config: Dict):
|
||
self.config = config
|
||
or_config = config["openrouter"]
|
||
api_key = or_config.get("api_key")
|
||
if not api_key or api_key == "YOUR_OPENROUTER_API_KEY":
|
||
raise ValueError("请设置有效的 OpenRouter API Key")
|
||
|
||
self.client = AsyncOpenAI(
|
||
base_url=or_config["base_url"],
|
||
api_key=api_key,
|
||
default_headers={"HTTP-Referer": "https://github.com/epub-translator", "X-Title": "EPUB Translator"}
|
||
)
|
||
self.models = or_config["models"]
|
||
self.rate_limiter = RateLimiter(
|
||
or_config["rate_limits"]["requests_per_minute"],
|
||
or_config["rate_limits"]["concurrent_requests"]
|
||
)
|
||
|
||
async def translate_chunk(self, items: List[ManifestItem], model_type: str = "production") -> Dict[str, str]:
|
||
"""
|
||
翻译一个段落块。
|
||
"""
|
||
if not items: return {}
|
||
|
||
prompt = self._build_prompt(items)
|
||
model = self.models.get(model_type, self.models["production"])
|
||
|
||
try:
|
||
raw_response = await self._make_request(prompt, model)
|
||
if not raw_response:
|
||
return {item.global_id: f"[翻译失败 - API空响应]" for item in items}
|
||
|
||
return self._parse_with_anchors(raw_response, items)
|
||
|
||
except Exception as e:
|
||
logger.error(f"翻译请求异常: {e}")
|
||
return {item.global_id: f"[翻译失败 - {str(e)}]" for item in items}
|
||
|
||
def _build_prompt(self, items: List[ManifestItem]) -> str:
|
||
lines = ["请将以下编号的英文段落翻译成中文。每段翻译前必须带上原编号，格式为 p_xxxxx。", "原文：", ""]
|
||
for item in items:
|
||
lines.append(f"{item.global_id} {item.clean_text}")
|
||
lines.extend(["", "要求：只返回翻译，不要解释。必须保留原编号。"])
|
||
return "\n".join(lines)
|
||
|
||
def _parse_with_anchors(self, response: str, items: List[ManifestItem]) -> Dict[str, str]:
|
||
results = {}
|
||
|
||
for i, item in enumerate(items):
|
||
current_id = item.global_id
|
||
|
||
# 使用简单的字符串拼接构造正则
|
||
start_pat = r"\b" + current_id + r"\b"
|
||
start_match = re.search(start_pat, response)
|
||
if not start_match:
|
||
continue
|
||
|
||
start_pos = start_match.start()
|
||
end_pos = len(response)
|
||
|
||
if i + 1 < len(items):
|
||
next_id = items[i+1].global_id
|
||
next_pat = r"\b" + next_id + r"\b"
|
||
next_match = re.search(next_pat, response[start_pos + 1:])
|
||
if next_match:
|
||
end_pos = start_pos + 1 + next_match.start()
|
||
|
||
segment = response[start_pos:end_pos].strip()
|
||
|
||
# 清洗 ID
|
||
# 注意：这里的正则使用了更稳健的拼接方式
|
||
id_clean_pat = r"^(\[?" + current_id + r"\]?[:：\s]*)+"
|
||
cleaned = re.sub(id_clean_pat, "", segment).strip()
|
||
|
||
if cleaned:
|
||
results[current_id] = cleaned
|
||
|
||
if len(results) < len(items):
|
||
expected_ids = {item.global_id for item in items}
|
||
missing_ids = expected_ids - set(results.keys())
|
||
logger.warning(f"锚点匹配缺失 {len(missing_ids)} 个。尝试降级解析...")
|
||
fallback_results = self._fallback_parse(response, items)
|
||
for pid, trans in fallback_results.items():
|
||
if pid not in results:
|
||
results[pid] = trans
|
||
|
||
return results
|
||
|
||
def _fallback_parse(self, response: str, items: List[ManifestItem]) -> Dict[str, str]:
|
||
results = {}
|
||
lines = [l.strip() for l in response.split("\n") if l.strip()]
|
||
for line in lines:
|
||
for item in items:
|
||
if item.global_id in line:
|
||
# 同样的稳健拼接
|
||
fallback_clean_pat = r"^\(?" + item.global_id + r"\)?[:：\s]*"
|
||
clean_line = re.sub(fallback_clean_pat, "", line).strip()
|
||
if clean_line:
|
||
results[item.global_id] = clean_line
|
||
return results
|
||
|
||
async def _make_request(self, prompt: str, model: str) -> str:
|
||
await self.rate_limiter.acquire()
|
||
try:
|
||
resp = await self.client.chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": "你是一位专业的翻译。请严格按编号输出翻译。"},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
temperature=0.2,
|
||
max_tokens=8000
|
||
)
|
||
return resp.choices[0].message.content.strip()
|
||
finally:
|
||
self.rate_limiter.release()
|
||
|
||
async def close(self):
|
||
await self.client.close() |