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()