118 lines
4.5 KiB
Python
118 lines
4.5 KiB
Python
"""
|
|
术语表管理器 (Glossary Manager)
|
|
|
|
负责从书籍内容中提取采样文本，调用 LLM 生成术语表，并管理术语表的持久化。
|
|
"""
|
|
|
|
import json
|
|
import random
|
|
from pathlib import Path
|
|
from typing import Dict, List, Any
|
|
from loguru import logger
|
|
from .manifest_manager import ManifestManager
|
|
from .llm_client import OpenRouterClient
|
|
|
|
class GlossaryManager:
|
|
def __init__(self, config: Dict, llm_client: OpenRouterClient):
|
|
self.config = config
|
|
self.llm_client = llm_client
|
|
self.glossary_path = Path("cache/glossary.json")
|
|
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 Exception:
|
|
logger.warning("未找到 config/prompts.json，使用默认 Prompt")
|
|
return {}
|
|
|
|
def extract_samples(self, manifest: ManifestManager, sample_size: int = 3000) -> str:
|
|
"""
|
|
从 Manifest 中提取采样文本。
|
|
策略：
|
|
1. 优先提取前言/绪论 (通常在文件的前部)。
|
|
2. 随机抽取中间段落。
|
|
"""
|
|
all_items = manifest.get_items()
|
|
if not all_items:
|
|
return ""
|
|
|
|
# 1. 提取开头部分 (Preface/Intro) - 假设在前 50 个段落中
|
|
intro_sample = [item.clean_text for item in all_items[:50] if len(item.clean_text) > 50]
|
|
|
|
# 2. 随机提取正文
|
|
body_items = [item for item in all_items[50:] if len(item.clean_text) > 50]
|
|
random_sample = []
|
|
if body_items:
|
|
# 随机取 10 个片段
|
|
sample_count = min(10, len(body_items))
|
|
random_items = random.sample(body_items, sample_count)
|
|
random_sample = [item.clean_text for item in random_items]
|
|
|
|
# 组合并截断
|
|
full_text = "\n\n".join(intro_sample + random_sample)
|
|
if len(full_text) > sample_size:
|
|
full_text = full_text[:sample_size] + "..."
|
|
|
|
return full_text
|
|
|
|
async def generate_glossary(self, manifest: ManifestManager) -> Dict[str, str]:
|
|
"""
|
|
生成术语表。
|
|
"""
|
|
# 1. 采样
|
|
sample_text = self.extract_samples(manifest)
|
|
if not sample_text:
|
|
logger.warning("采样文本为空，跳过术语表生成")
|
|
return {}
|
|
|
|
logger.info(f"提取了 {len(sample_text)} 字符的采样文本，正在生成术语表...")
|
|
|
|
# 2. 构建 Prompt
|
|
prompt_cfg = self.prompts.get("glossary_extraction", {})
|
|
system_prompt = prompt_cfg.get("system", "Analyze the text and extract named entities.")
|
|
user_template = prompt_cfg.get("user_template", "Text:\n{{content}}")
|
|
user_prompt = user_template.replace("{{content}}", sample_text)
|
|
|
|
# 3. 调用 LLM (使用 smart 模型)
|
|
# 注意：这里需要 LLMClient 支持直接传入 system/user prompt，而不是封装好的 translate 接口
|
|
# 我们稍后会扩展 LLMClient
|
|
try:
|
|
response = await self.llm_client.raw_chat_completion(
|
|
system_prompt,
|
|
user_prompt,
|
|
model_type="smart"
|
|
)
|
|
|
|
# 4. 解析 JSON
|
|
# 简单的 JSON 提取逻辑 (处理可能的 markdown code block)
|
|
json_str = response.strip()
|
|
if "```json" in json_str:
|
|
json_str = json_str.split("```json")[1].split("```")[0].strip()
|
|
elif "```" in json_str:
|
|
json_str = json_str.split("```")[1].split("```")[0].strip()
|
|
|
|
glossary = json.loads(json_str)
|
|
self.save_glossary(glossary)
|
|
return glossary
|
|
|
|
except Exception as e:
|
|
logger.error(f"术语表生成失败: {e}")
|
|
return {}
|
|
|
|
def save_glossary(self, glossary: Dict[str, str]):
|
|
self.glossary_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(self.glossary_path, "w", encoding="utf-8") as f:
|
|
json.dump(glossary, f, ensure_ascii=False, indent=2)
|
|
logger.info(f"术语表已保存至: {self.glossary_path}")
|
|
|
|
def load_glossary(self) -> Dict[str, str]:
|
|
if self.glossary_path.exists():
|
|
try:
|
|
with open(self.glossary_path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except:
|
|
pass
|
|
return {}
|