feat: initial commit
This commit is contained in:
@@ -0,0 +1,274 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user