242 lines
9.4 KiB
Python
242 lines
9.4 KiB
Python
"""
|
||
LLM Client Module - Generic OpenAI Compatible
|
||
|
||
Features:
|
||
1. Fully configurable via config.json (base_url, headers).
|
||
2. Mode-aware prompt building (bilingual vs chinese).
|
||
3. Format repair capability for chinese mode.
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import re
|
||
from openai import AsyncOpenAI
|
||
from typing import List, Dict, Optional, Any
|
||
from loguru import logger
|
||
import time
|
||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||
from .manifest_manager import ManifestItem
|
||
|
||
|
||
class RateLimiter:
|
||
"""Rate limiter for concurrency and 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
|
||
self._lock = asyncio.Lock()
|
||
|
||
async def acquire(self):
|
||
await self.semaphore.acquire()
|
||
async with self._lock:
|
||
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 LLMClient:
|
||
"""Generic OpenAI-compatible API Client."""
|
||
|
||
def __init__(self, config: Dict):
|
||
self.config = config
|
||
llm_config = config["llm"]
|
||
|
||
api_key = llm_config.get("api_key")
|
||
base_url = llm_config.get("base_url")
|
||
extra_headers = llm_config.get("extra_headers", {})
|
||
|
||
if not api_key:
|
||
raise ValueError("API Key is missing in config")
|
||
|
||
self.client = AsyncOpenAI(
|
||
base_url=base_url,
|
||
api_key=api_key,
|
||
default_headers=extra_headers
|
||
)
|
||
|
||
self.models = llm_config.get("models", {"fast": "gpt-3.5-turbo", "smart": "gpt-4"})
|
||
|
||
self.rate_limiter = RateLimiter(
|
||
llm_config["rate_limits"]["requests_per_minute"],
|
||
llm_config["rate_limits"]["concurrent_requests"]
|
||
)
|
||
self.prompts = self._load_prompts()
|
||
|
||
def _load_prompts(self) -> Dict:
|
||
try:
|
||
with open("config/prompts.json", "r", encoding="utf-8") as f:
|
||
return json.load(f)
|
||
except:
|
||
return {}
|
||
|
||
async def translate_chunk(self, items: List[ManifestItem], glossary: Dict = None,
|
||
instruction: str = None, model_type: str = "fast",
|
||
mode: str = "bilingual") -> Dict[str, str]:
|
||
"""
|
||
Translate a chunk of items.
|
||
|
||
Args:
|
||
items: List of ManifestItem to translate
|
||
glossary: Term dictionary
|
||
instruction: Style guide
|
||
model_type: "fast" or "smart"
|
||
mode: "bilingual" or "chinese"
|
||
"""
|
||
if not items: return {}
|
||
|
||
model = self.models.get(model_type, self.models.get("fast"))
|
||
prompt = self._build_prompt(items, mode)
|
||
|
||
try:
|
||
# Build System Prompt
|
||
base_sys_prompt = self.prompts.get("translation", {}).get("system", "You are a professional translator.")
|
||
|
||
# 中文模式：添加占位符保护指令
|
||
if mode == "chinese":
|
||
base_sys_prompt += """
|
||
|
||
Placeholder Instructions (CRITICAL):
|
||
1. Text contains PAIRED placeholders: φNφ (start) and φ/Nφ (end), like HTML tags.
|
||
2. Example: "φ1φTable Talkφ/1φ" means italic text, translate as "φ1φ桌谈φ/1φ"
|
||
3. Single placeholders φNφ without φ/Nφ are inline elements (footnotes, formulas) - keep them in place.
|
||
4. RULES:
|
||
- DO NOT create new placeholder numbers that don't exist in the original
|
||
- DO NOT remove or modify existing placeholders
|
||
- Keep placeholders in the SAME relative position in your translation
|
||
- If word order changes, keep placeholders with their associated text
|
||
5. Each line starts with paragraph ID (p_xxxxx). Preserve them.
|
||
"""
|
||
|
||
|
||
if instruction:
|
||
base_sys_prompt += f"\n\nBook Style Guide:\n{instruction}"
|
||
|
||
if glossary:
|
||
glossary_text = "\n".join([f"{k} -> {v}" for k, v in glossary.items()])
|
||
base_sys_prompt += f"\n\nTerminology:\n{glossary_text}"
|
||
|
||
# Strict formatting instructions
|
||
base_sys_prompt += "\n\nRequirements:\n1. Each line MUST start with ID (p_xxxxx).\n2. DO NOT modify IDs.\n3. Return only translations."
|
||
|
||
raw_response = await self._make_request(model, base_sys_prompt, prompt)
|
||
|
||
if not raw_response:
|
||
return {item.global_id: f"[Error - Empty Response]" for item in items}
|
||
|
||
return self._simple_parse(raw_response, items, mode)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Translation failed ({model}): {e}")
|
||
return {item.global_id: f"[Error - {str(e)}]" for item in items}
|
||
|
||
async def repair_format(self, original_text: str, broken_translation: str) -> str:
|
||
"""
|
||
修复翻译格式：将占位符正确插入到译文中。
|
||
"""
|
||
model = self.models.get("fast")
|
||
|
||
system_prompt = "You are a format repair assistant. Your ONLY job is to insert placeholders into the translation."
|
||
user_prompt = f"""
|
||
Original Text (with placeholders):
|
||
{original_text}
|
||
|
||
Translation (placeholders missing/incorrect):
|
||
{broken_translation}
|
||
|
||
Task:
|
||
Please rewrite the Translation to include ALL placeholders (φcXXXXXφ) from the Original Text.
|
||
1. DO NOT translate again. Keep the meaning of the Translation.
|
||
2. Place φcXXXXXφ tags exactly where they correspond to the original format (bold, italic, links).
|
||
3. Output ONLY the fixed translation.
|
||
"""
|
||
try:
|
||
return await self._make_request(model, system_prompt, user_prompt)
|
||
except Exception as e:
|
||
logger.error(f"Format repair failed: {e}")
|
||
return broken_translation
|
||
|
||
async def raw_chat_completion(self, system_prompt: str, user_prompt: str, model_type: str = "smart") -> str:
|
||
"""Generic chat completion (for Profiler)."""
|
||
model = self.models.get(model_type, self.models.get("smart"))
|
||
return await self._make_request(model, system_prompt, user_prompt)
|
||
|
||
def _build_prompt(self, items: List[ManifestItem], mode: str = "bilingual") -> str:
|
||
"""构建翻译提示词"""
|
||
lines = []
|
||
for item in items:
|
||
if mode == "chinese":
|
||
# 中文模式：使用带占位符的文本和段落类型
|
||
text = item.text_with_placeholders if item.text_with_placeholders else item.clean_text
|
||
p_type = getattr(item, 'paragraph_type', 'body').upper()
|
||
lines.append(f"{item.global_id} [{p_type}] {text}")
|
||
else:
|
||
# 双语模式：使用纯文本
|
||
lines.append(f"{item.global_id} {item.clean_text}")
|
||
return "\n".join(lines)
|
||
|
||
def _simple_parse(self, response: str, items: List[ManifestItem], mode: str = "bilingual") -> Dict[str, str]:
|
||
"""解析 LLM 响应"""
|
||
results = {}
|
||
for i, item in enumerate(items):
|
||
current_id = item.global_id
|
||
start_idx = response.find(current_id)
|
||
if start_idx == -1: continue
|
||
|
||
end_idx = len(response)
|
||
if i + 1 < len(items):
|
||
next_id = items[i+1].global_id
|
||
next_found = response.find(next_id, start_idx + len(current_id))
|
||
if next_found != -1:
|
||
end_idx = next_found
|
||
|
||
content = response[start_idx:end_idx].strip()
|
||
clean_content = content[len(current_id):].strip()
|
||
clean_content = clean_content.lstrip(":： \t")
|
||
|
||
# 移除类型标记 (如 [BODY])
|
||
if mode == "chinese":
|
||
clean_content = re.sub(r'^\[[A-Z]+\]\s*', '', clean_content)
|
||
|
||
if clean_content:
|
||
results[current_id] = clean_content
|
||
|
||
# Fallback: 逐行解析
|
||
if len(results) < len(items):
|
||
for line in response.split("\n"):
|
||
line = line.strip()
|
||
for item in items:
|
||
if item.global_id not in results and line.startswith(item.global_id):
|
||
res = line[len(item.global_id):].strip().lstrip(":： ")
|
||
if mode == "chinese":
|
||
res = re.sub(r'^\[[A-Z]+\]\s*', '', res)
|
||
if res: results[item.global_id] = res
|
||
return results
|
||
|
||
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
|
||
async def _make_request(self, model: str, system_prompt: str, user_prompt: str) -> str:
|
||
await self.rate_limiter.acquire()
|
||
try:
|
||
resp = await self.client.chat.completions.create(
|
||
model=model,
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_prompt}
|
||
],
|
||
temperature=self.config['translation'].get('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()
|
||
|
||
# Alias for backward compatibility
|
||
OpenRouterClient = LLMClient
|