275 lines
9.4 KiB
Python
275 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Entity Linker - 实体消歧辅助模块 (v5.4)
|
||
|
||
为 Data Agent 提供实体消歧的辅助功能：
|
||
- 置信度判断
|
||
- 别名索引管理 (通过 index.db aliases 表)
|
||
- 消歧结果记录
|
||
|
||
v5.1 变更（v5.4 沿用）:
|
||
- 别名存储从 state.json 迁移到 index.db aliases 表
|
||
- 使用 IndexManager 进行别名读写
|
||
- 移除对 state.json 的直接操作
|
||
"""
|
||
|
||
from typing import Dict, List, Optional, Tuple
|
||
from dataclasses import dataclass, field
|
||
|
||
from .config import get_config
|
||
from .index_manager import IndexManager
|
||
from .observability import safe_log_tool_call
|
||
|
||
|
||
@dataclass
|
||
class DisambiguationResult:
|
||
"""消歧结果"""
|
||
mention: str
|
||
entity_id: Optional[str]
|
||
confidence: float
|
||
candidates: List[str] = field(default_factory=list)
|
||
adopted: bool = False
|
||
warning: Optional[str] = None
|
||
|
||
|
||
class EntityLinker:
|
||
"""实体链接器 - 辅助 Data Agent 进行实体消歧 (v5.1 SQLite，v5.4 沿用)"""
|
||
|
||
def __init__(self, config=None):
|
||
self.config = config or get_config()
|
||
self._index_manager = IndexManager(self.config)
|
||
|
||
# ==================== 别名管理 (v5.1 SQLite，v5.4 沿用) ====================
|
||
|
||
def register_alias(self, entity_id: str, alias: str, entity_type: str = "角色") -> bool:
|
||
"""注册新别名（v5.1 引入：写入 index.db aliases 表）"""
|
||
if not alias or not entity_id:
|
||
return False
|
||
return self._index_manager.register_alias(alias, entity_id, entity_type)
|
||
|
||
def lookup_alias(self, mention: str, entity_type: str = None) -> Optional[str]:
|
||
"""查找别名对应的实体ID（返回第一个匹配，可选按类型过滤）"""
|
||
entries = self._index_manager.get_entities_by_alias(mention)
|
||
if not entries:
|
||
return None
|
||
|
||
if entity_type:
|
||
for entry in entries:
|
||
if entry.get("type") == entity_type:
|
||
return entry.get("id")
|
||
return None
|
||
else:
|
||
return entries[0].get("id") if entries else None
|
||
|
||
def lookup_alias_all(self, mention: str) -> List[Dict]:
|
||
"""查找别名对应的所有实体（一对多）"""
|
||
entries = self._index_manager.get_entities_by_alias(mention)
|
||
return [{"type": e.get("type"), "id": e.get("id")} for e in entries]
|
||
|
||
def get_all_aliases(self, entity_id: str, entity_type: str = None) -> List[str]:
|
||
"""获取实体的所有别名"""
|
||
return self._index_manager.get_entity_aliases(entity_id)
|
||
|
||
# ==================== 置信度判断 ====================
|
||
|
||
def evaluate_confidence(self, confidence: float) -> Tuple[str, bool, Optional[str]]:
|
||
"""
|
||
评估置信度，返回 (action, adopt, warning)
|
||
|
||
- action: "auto" | "warn" | "manual"
|
||
- adopt: 是否采用
|
||
- warning: 警告信息
|
||
"""
|
||
if confidence >= self.config.extraction_confidence_high:
|
||
return ("auto", True, None)
|
||
elif confidence >= self.config.extraction_confidence_medium:
|
||
return ("warn", True, f"中置信度匹配 (confidence: {confidence:.2f})")
|
||
else:
|
||
return ("manual", False, f"需人工确认 (confidence: {confidence:.2f})")
|
||
|
||
def process_uncertain(
|
||
self,
|
||
mention: str,
|
||
candidates: List[str],
|
||
suggested: str,
|
||
confidence: float,
|
||
context: str = ""
|
||
) -> DisambiguationResult:
|
||
"""
|
||
处理不确定的实体匹配
|
||
|
||
返回消歧结果，包含是否采用、警告信息等
|
||
"""
|
||
action, adopt, warning = self.evaluate_confidence(confidence)
|
||
|
||
result = DisambiguationResult(
|
||
mention=mention,
|
||
entity_id=suggested if adopt else None,
|
||
confidence=confidence,
|
||
candidates=candidates,
|
||
adopted=adopt,
|
||
warning=warning
|
||
)
|
||
|
||
return result
|
||
|
||
# ==================== 批量处理 ====================
|
||
|
||
def process_extraction_result(
|
||
self,
|
||
uncertain_items: List[Dict]
|
||
) -> Tuple[List[DisambiguationResult], List[str]]:
|
||
"""
|
||
处理 AI 提取结果中的 uncertain 项
|
||
|
||
返回 (results, warnings)
|
||
"""
|
||
results = []
|
||
warnings = []
|
||
|
||
for item in uncertain_items:
|
||
result = self.process_uncertain(
|
||
mention=item.get("mention", ""),
|
||
candidates=item.get("candidates", []),
|
||
suggested=item.get("suggested", ""),
|
||
confidence=item.get("confidence", 0.0),
|
||
context=item.get("context", "")
|
||
)
|
||
results.append(result)
|
||
|
||
if result.warning:
|
||
warnings.append(f"{result.mention} → {result.entity_id}: {result.warning}")
|
||
|
||
return results, warnings
|
||
|
||
def register_new_entities(
|
||
self,
|
||
new_entities: List[Dict]
|
||
) -> List[str]:
|
||
"""
|
||
注册新实体的别名 (v5.1 引入，v5.4 沿用)
|
||
|
||
返回注册的实体ID列表
|
||
"""
|
||
registered = []
|
||
|
||
for entity in new_entities:
|
||
entity_id = entity.get("suggested_id") or entity.get("id")
|
||
if not entity_id or entity_id == "NEW":
|
||
continue
|
||
|
||
entity_type = entity.get("type", "角色")
|
||
|
||
# 注册主名称
|
||
name = entity.get("name", "")
|
||
if name:
|
||
self.register_alias(entity_id, name, entity_type)
|
||
|
||
# 注册提及方式
|
||
for mention in entity.get("mentions", []):
|
||
if mention and mention != name:
|
||
self.register_alias(entity_id, mention, entity_type)
|
||
|
||
registered.append(entity_id)
|
||
|
||
return registered
|
||
|
||
|
||
# ==================== CLI 接口 ====================
|
||
|
||
def main():
|
||
import argparse
|
||
import sys
|
||
from .cli_output import print_success, print_error
|
||
from .cli_args import normalize_global_project_root
|
||
from .index_manager import IndexManager
|
||
|
||
parser = argparse.ArgumentParser(description="Entity Linker CLI (v5.4 SQLite)")
|
||
parser.add_argument("--project-root", type=str, help="项目根目录")
|
||
|
||
subparsers = parser.add_subparsers(dest="command")
|
||
|
||
# 注册别名
|
||
register_parser = subparsers.add_parser("register-alias")
|
||
register_parser.add_argument("--entity", required=True, help="实体ID")
|
||
register_parser.add_argument("--alias", required=True, help="别名")
|
||
register_parser.add_argument("--type", default="角色", help="实体类型（默认：角色）")
|
||
|
||
# 查找别名
|
||
lookup_parser = subparsers.add_parser("lookup")
|
||
lookup_parser.add_argument("--mention", required=True, help="提及文本")
|
||
lookup_parser.add_argument("--type", help="按类型过滤")
|
||
|
||
# 查找所有匹配（一对多）
|
||
lookup_all_parser = subparsers.add_parser("lookup-all")
|
||
lookup_all_parser.add_argument("--mention", required=True, help="提及文本")
|
||
|
||
# 列出别名
|
||
list_parser = subparsers.add_parser("list-aliases")
|
||
list_parser.add_argument("--entity", required=True, help="实体ID")
|
||
list_parser.add_argument("--type", help="实体类型")
|
||
|
||
argv = normalize_global_project_root(sys.argv[1:])
|
||
args = parser.parse_args(argv)
|
||
|
||
# 初始化
|
||
config = None
|
||
if args.project_root:
|
||
# 允许传入“工作区根目录”，统一解析到真正的 book project_root（必须包含 .noma/state.json）
|
||
from project_locator import resolve_project_root
|
||
from .config import DataModulesConfig
|
||
|
||
resolved_root = resolve_project_root(args.project_root)
|
||
config = DataModulesConfig.from_project_root(resolved_root)
|
||
|
||
linker = EntityLinker(config)
|
||
logger = IndexManager(config)
|
||
tool_name = f"entity_linker:{args.command or 'unknown'}"
|
||
|
||
def emit_success(data=None, message: str = "ok"):
|
||
print_success(data, message=message)
|
||
safe_log_tool_call(logger, tool_name=tool_name, success=True)
|
||
|
||
def emit_error(code: str, message: str, suggestion: str | None = None):
|
||
print_error(code, message, suggestion=suggestion)
|
||
safe_log_tool_call(
|
||
logger,
|
||
tool_name=tool_name,
|
||
success=False,
|
||
error_code=code,
|
||
error_message=message,
|
||
)
|
||
|
||
if args.command == "register-alias":
|
||
entity_type = getattr(args, "type", "角色")
|
||
success = linker.register_alias(args.entity, args.alias, entity_type)
|
||
if success:
|
||
emit_success({"entity": args.entity, "alias": args.alias, "type": entity_type}, message="alias_registered")
|
||
else:
|
||
emit_error("ALIAS_EXISTS", "注册失败或已存在")
|
||
|
||
elif args.command == "lookup":
|
||
entity_type = getattr(args, "type", None)
|
||
entity_id = linker.lookup_alias(args.mention, entity_type)
|
||
if entity_id:
|
||
emit_success({"mention": args.mention, "entity": entity_id}, message="lookup")
|
||
else:
|
||
emit_error("NOT_FOUND", f"未找到别名: {args.mention}")
|
||
|
||
elif args.command == "lookup-all":
|
||
matches = linker.lookup_alias_all(args.mention)
|
||
emit_success(matches, message="lookup_all")
|
||
|
||
elif args.command == "list-aliases":
|
||
entity_type = getattr(args, "type", None)
|
||
aliases = linker.get_all_aliases(args.entity, entity_type)
|
||
emit_success(aliases, message="aliases")
|
||
|
||
else:
|
||
emit_error("UNKNOWN_COMMAND", "未指定有效命令", suggestion="请查看 --help")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|