Release v0.13: Typography V3 (Safe mode) & Endnote Extraction Fix
This commit is contained in:
谭凯
@@ -0,0 +1,231 @@
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Optional
|
||||
from openai import AsyncOpenAI
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
from loguru import logger
|
||||
|
||||
from src.common.data_model import ManifestEntry
|
||||
from src.common.utils import setup_logger, ensure_directory
|
||||
|
||||
logger = setup_logger("llm_client")
|
||||
|
||||
# Default chunk save directory (can be overridden)
|
||||
DEFAULT_CHUNK_DIR = Path("tmp/chunks")
|
||||
|
||||
|
||||
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 with short ID strategy."""
|
||||
|
||||
def __init__(self, api_key: str, base_url: str, model: str = "gpt-3.5-turbo",
|
||||
requests_per_minute: int = 60, concurrent_requests: int = 5,
|
||||
extra_headers: Dict = None, chunk_dir: Path = None):
|
||||
|
||||
# Configure proxy client to avoid SOCKS issues and ensure connectivity
|
||||
import httpx
|
||||
import os
|
||||
|
||||
# Prefer HTTP proxy if available to avoid missing socksio support
|
||||
proxy_url = os.environ.get("http_proxy") or os.environ.get("https_proxy")
|
||||
http_client = httpx.AsyncClient(
|
||||
proxy=proxy_url,
|
||||
timeout=60.0,
|
||||
follow_redirects=True
|
||||
) if proxy_url else None
|
||||
|
||||
self.client = AsyncOpenAI(
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
default_headers=extra_headers,
|
||||
http_client=http_client
|
||||
)
|
||||
self.model = model
|
||||
|
||||
self.rate_limiter = RateLimiter(requests_per_minute, concurrent_requests)
|
||||
self.prompts = self._load_prompts()
|
||||
self._chunk_counter = 0
|
||||
|
||||
# Chunk directory for debug output
|
||||
self.chunk_dir = chunk_dir or DEFAULT_CHUNK_DIR
|
||||
ensure_directory(self.chunk_dir)
|
||||
|
||||
def _load_prompts(self) -> Dict:
|
||||
try:
|
||||
with open("config/prompts.json", "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load config/prompts.json: {e}")
|
||||
return {}
|
||||
|
||||
async def translate_chunk(self, items: List[ManifestEntry], glossary: Dict = None,
|
||||
instruction: str = None, mode: str = "bilingual") -> Dict[str, str]:
|
||||
"""
|
||||
Translate a chunk of items using short ID strategy.
|
||||
Returns: Dict[entry_id, translated_text]
|
||||
"""
|
||||
if not items: return {}
|
||||
|
||||
# Build prompt with short IDs
|
||||
id_map, prompt = self._build_prompt_with_short_ids(items)
|
||||
|
||||
try:
|
||||
# Build System Prompt
|
||||
base_sys_prompt = self.prompts.get("translation", {}).get("system",
|
||||
"You are a professional English to Chinese translator.")
|
||||
|
||||
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}"
|
||||
|
||||
# Short ID format instructions
|
||||
base_sys_prompt += """
|
||||
|
||||
Output Format:
|
||||
- Each line MUST start with #N: (keep this ID exactly as given)
|
||||
- Preserve any φXφ or φ/Xφ placeholders EXACTLY as-is
|
||||
- Only output translations, no explanations
|
||||
- Match the number of output lines to input lines"""
|
||||
|
||||
# Save chunk before translation
|
||||
chunk_id = self._save_chunk("before", prompt, base_sys_prompt)
|
||||
|
||||
logger.debug(f"Sending request to LLM (Chunk: {chunk_id}, Items: {len(items)})")
|
||||
|
||||
raw_response = await self._make_request(base_sys_prompt, prompt)
|
||||
|
||||
# Save chunk after translation
|
||||
self._save_chunk("after", raw_response, base_sys_prompt, chunk_id)
|
||||
|
||||
# Parse with short ID mapping
|
||||
results = self._parse_short_id_response(raw_response, id_map)
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Translation failed: {e}")
|
||||
return {item.entry_id: f"[Error - {str(e)}]" for item in items}
|
||||
|
||||
def _build_prompt_with_short_ids(self, items: List[ManifestEntry]) -> tuple:
|
||||
"""
|
||||
Build prompt with short IDs (#1, #2, ...).
|
||||
Returns: (id_map, prompt_text)
|
||||
"""
|
||||
id_map = {} # short_id -> entry_id
|
||||
lines = []
|
||||
|
||||
for i, item in enumerate(items, 1):
|
||||
short_id = f"#{i}"
|
||||
id_map[short_id] = item.entry_id
|
||||
|
||||
# Clean text (remove extra whitespace)
|
||||
text = re.sub(r'\s+', ' ', item.original_text).strip()
|
||||
lines.append(f"{short_id}: {text}")
|
||||
|
||||
return id_map, "\n".join(lines)
|
||||
|
||||
def _parse_short_id_response(self, response: str, id_map: Dict[str, str]) -> Dict[str, str]:
|
||||
"""
|
||||
Parse response with short ID format.
|
||||
Returns: Dict[entry_id, translated_text]
|
||||
"""
|
||||
results = {}
|
||||
|
||||
for line in response.split("\n"):
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
# Match #N: pattern
|
||||
# Handles both ASCII colon (:) and full-width colon (：)
|
||||
match = re.match(r'^#(\d+)[:：]\s*(.+)$', line)
|
||||
if match:
|
||||
short_id = f"#{match.group(1)}"
|
||||
content = match.group(2).strip()
|
||||
|
||||
# Recursively strip repeated IDs (e.g. "#12: #12: Text")
|
||||
while True:
|
||||
sub_match = re.match(r'^#(\d+)[:：]\s*(.+)$', content)
|
||||
if sub_match:
|
||||
# Check if the inner ID matches the outer ID, or just strip it anyway
|
||||
# Usually LLM repeats the same ID.
|
||||
content = sub_match.group(2).strip()
|
||||
else:
|
||||
break
|
||||
|
||||
if short_id in id_map:
|
||||
full_id = id_map[short_id]
|
||||
results[full_id] = content
|
||||
else:
|
||||
logger.warning(f"Unknown short ID in response: {short_id}")
|
||||
|
||||
return results
|
||||
|
||||
def _save_chunk(self, stage: str, content: str, system_prompt: str = None,
|
||||
chunk_id: str = None) -> str:
|
||||
"""Save chunk to tmp directory for debugging."""
|
||||
if chunk_id is None:
|
||||
self._chunk_counter += 1
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
chunk_id = f"{timestamp}_{self._chunk_counter:04d}"
|
||||
|
||||
filename = self.chunk_dir / f"chunk_{chunk_id}_{stage}.txt"
|
||||
|
||||
with open(filename, "w", encoding="utf-8") as f:
|
||||
if system_prompt and stage == "before":
|
||||
f.write("=== SYSTEM PROMPT ===\n")
|
||||
f.write(system_prompt)
|
||||
f.write("\n\n=== USER PROMPT ===\n")
|
||||
f.write(content)
|
||||
|
||||
logger.debug(f"Saved chunk: {filename}")
|
||||
return chunk_id
|
||||
|
||||
async def raw_chat_completion(self, system_prompt: str, user_prompt: str) -> str:
|
||||
"""Generic chat completion."""
|
||||
return await self._make_request(system_prompt, user_prompt)
|
||||
|
||||
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
|
||||
async def _make_request(self, system_prompt: str, user_prompt: str) -> str:
|
||||
await self.rate_limiter.acquire()
|
||||
try:
|
||||
resp = await self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt}
|
||||
],
|
||||
temperature=0.3,
|
||||
)
|
||||
return resp.choices[0].message.content.strip()
|
||||
finally:
|
||||
self.rate_limiter.release()
|
||||
|
||||
async def close(self):
|
||||
await self.client.close()
|
||||
Reference in New Issue
Block a user