feat: initial commit

2026-06-23 20:29:02 +08:00
commit ea8f6066c4
217 changed files with 60754 additions and 0 deletions
@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+EntityLinker extra tests + CLI
+"""
+
+import sys
+
+import pytest
+
+from data_modules.entity_linker import EntityLinker, main as linker_main
+from data_modules.index_manager import IndexManager, EntityMeta
+
+
+@pytest.fixture
+def temp_project(tmp_path):
+    from data_modules.config import DataModulesConfig
+    cfg = DataModulesConfig.from_project_root(tmp_path)
+    cfg.ensure_dirs()
+    return cfg
+
+
+def test_process_extraction_and_register_new_entities(temp_project):
+    linker = EntityLinker(temp_project)
+    idx = IndexManager(temp_project)
+    idx.upsert_entity(
+        EntityMeta(
+            id="xiaoyan",
+            type="角色",
+            canonical_name="萧炎",
+            current={},
+            first_appearance=1,
+            last_appearance=1,
+        )
+    )
+
+    results, warnings = linker.process_extraction_result(
+        [
+            {
+                "mention": "萧炎",
+                "candidates": ["xiaoyan"],
+                "suggested": "xiaoyan",
+                "confidence": 0.7,
+            },
+            {
+                "mention": "宗主",
+                "candidates": ["zongzhu"],
+                "suggested": "zongzhu",
+                "confidence": 0.4,
+            },
+        ]
+    )
+
+    assert len(results) == 2
+    assert len(warnings) == 2
+
+    registered = linker.register_new_entities(
+        [
+            {
+                "suggested_id": "hongyi",
+                "name": "红衣女子",
+                "type": "角色",
+                "mentions": ["红衣", "女子"],
+            }
+        ]
+    )
+    assert registered == ["hongyi"]
+    aliases = idx.get_entity_aliases("hongyi")
+    assert "红衣女子" in aliases
+
+
+def test_entity_linker_cli(temp_project, monkeypatch, capsys):
+    idx = IndexManager(temp_project)
+    idx.upsert_entity(
+        EntityMeta(
+            id="xiaoyan",
+            type="角色",
+            canonical_name="萧炎",
+            current={},
+            first_appearance=1,
+            last_appearance=1,
+        )
+    )
+
+    def run_cli(args):
+        monkeypatch.setattr(sys, "argv", ["entity_linker"] + args)
+        linker_main()
+
+    root = str(temp_project.project_root)
+
+    run_cli(["--project-root", root, "register-alias", "--entity", "xiaoyan", "--alias", "炎帝"])
+    run_cli(["--project-root", root, "lookup", "--mention", "炎帝"])
+    run_cli(["--project-root", root, "lookup", "--mention", "不存在"])
+    run_cli(["--project-root", root, "lookup-all", "--mention", "炎帝"])
+    run_cli(["--project-root", root, "list-aliases", "--entity", "xiaoyan"])
+
+    capsys.readouterr()