232 lines
8.6 KiB
Python
232 lines
8.6 KiB
Python
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()
|