knowledge-graph-kit 0.1.0__py3-none-any.whl

knowledge_graph_kit/extractor.py

@@ -0,0 +1,419 @@
+"""
+extractor.py — Schema 引导的实体与关系抽取
+
+流程:
+  1. 解析 schema.txt 得到本体定义
+  2. 对每个 Chunk，构造含 schema 约束的 prompt
+  3. 调用 LLM（OpenAI function calling）抽取实体和关系
+  4. 输出结构化结果（Pydantic 模型）
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+from openai import OpenAI
+from pydantic import BaseModel, Field
+
+# Windows 终端 UTF-8
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding="utf-8")  # type: ignore[attr-defined]
+    except Exception:
+        pass
+
+# ── 惰性客户端 ────────────────────────────────────────────────
+_client_instance: OpenAI | None = None
+_MODEL: str = "gpt-4o-mini"
+
+
+def configure(api_key: str, base_url: str | None = None, model: str = "gpt-4o-mini") -> None:
+    """显式配置 OpenAI 客户端（在程序化使用时调用）。"""
+    global _client_instance, _MODEL
+    _client_instance = OpenAI(api_key=api_key, base_url=base_url, timeout=90)
+    _MODEL = model
+
+
+def _get_client() -> OpenAI:
+    """惰性初始化 OpenAI 客户端，从环境变量读取配置。"""
+    global _client_instance, _MODEL
+    if _client_instance is None:
+        api_key = os.getenv("OPENAI_API_KEY", "")
+        if not api_key:
+            raise RuntimeError(
+                "OPENAI_API_KEY 未设置。请在 .env 中配置或调用 configure()。"
+            )
+        base_url = os.getenv("OPENAI_BASE_URL")
+        model = os.getenv("LLM_MODEL_NAME", "gpt-4o-mini")
+        _client_instance = OpenAI(api_key=api_key, base_url=base_url, timeout=90)
+        _MODEL = model
+    return _client_instance
+
+
+# ═══════════════════════════════════════════════════════════
+# 1. 结构化输出 Schema（Pydantic）
+# ═══════════════════════════════════════════════════════════
+
+class Entity(BaseModel):
+    """抽取出的实体"""
+    name: str = Field(description="实体名称，如 '解构赋值'、'let关键字'")
+    type: str = Field(description="实体类型: Topic / Concept / Skill / CodeExample / TextSegment / Question")
+    properties: dict[str, Any] = Field(default_factory=dict, description="实体属性键值对，如 {'level': '基础', 'difficulty': '容易'}")
+
+
+class Relation(BaseModel):
+    """抽取出的关系"""
+    source: str = Field(description="源实体名称（必须与某个 Entity.name 一致）")
+    target: str = Field(description="目标实体名称")
+    type: str = Field(description="关系类型: hasPrerequisite / isA / commonlyConfusedWith / teaches / hasExample / assessedBy")
+    properties: dict[str, Any] = Field(default_factory=dict, description="关系属性，如 {'confidence': 0.9}")
+
+
+class ExtractionResult(BaseModel):
+    """单个 Chunk 的抽取结果"""
+    entities: list[Entity] = Field(description="从文本中抽取的实体列表")
+    relations: list[Relation] = Field(description="从文本中抽取的关系列表")
+
+
+# ═══════════════════════════════════════════════════════════
+# 2. Schema 解析
+# ═══════════════════════════════════════════════════════════
+
+def parse_schema(schema_path: str | Path) -> dict[str, Any]:
+    """解析 schema.txt，返回结构化本体定义"""
+    schema_path = Path(schema_path)
+    text = schema_path.read_text(encoding="utf-8")
+
+    schema_info: dict[str, Any] = {
+        "node_types": {},
+        "relation_types": {},
+    }
+
+    current_section: Optional[str] = None
+    current_rel_type: Optional[str] = None
+
+    for raw_line in text.splitlines():
+        line = raw_line.strip()
+        if not line or line.startswith("#"):
+            continue
+
+        # ── 切换区块 ──
+        if "节点类型" in line and ":" in line:
+            current_section = "node"
+            continue
+        if "关系类型" in line and ":" in line:
+            current_section = "rel"
+            continue
+
+        if current_section == "node":
+            # 匹配 "- Topic:             # 知识主题"
+            if line.startswith("- "):
+                content = line[2:]
+                # 提取类型名（冒号或空格或括号前的内容）
+                type_name = content.split(":")[0].split("#")[0].split("(")[0].strip()
+                # 提取描述（# 后面的内容，去掉 "Topic的子类" 等后缀信息）
+                desc = ""
+                if "#" in content:
+                    desc = content.split("#", 1)[1].strip()
+                schema_info["node_types"][type_name] = {
+                    "description": desc,
+                    "properties": [],
+                }
+                continue
+
+            # 匹配属性行: "      属性: [name, level(认知层次), difficulty, estimatedTime]"
+            if "属性:" in line:
+                props_str = line.split("属性:", 1)[1].strip()
+                if props_str.startswith("[") and "]" in props_str:
+                    props_str = props_str[1:props_str.index("]")]
+                    props = [p.strip() for p in props_str.split(",")]
+                    # 找出当前最近的节点类型来关联属性
+                    if schema_info["node_types"]:
+                        last_type = list(schema_info["node_types"].keys())[-1]
+                        schema_info["node_types"][last_type]["properties"] = props
+                continue
+
+        elif current_section == "rel":
+            # 关系定义行的两种格式:
+            #   带 "- " 前缀:  "- hasPrerequisite:     # 知识A是B的前置"
+            #   不带 "- " 前缀: "hasPrerequisite:     # 知识A是B的前置"
+            rel_match = None
+            for prefix in ("- ", ""):
+                if line.startswith(prefix):
+                    probe = line[len(prefix):] if prefix else line
+                    if ":" in probe:
+                        candidate_name = probe.split(":")[0].strip()
+                        # 避免误匹配 "domain:" / "range:" 行
+                        if candidate_name in ("domain", "range"):
+                            continue
+                        rel_match = (candidate_name, probe)
+                        break
+            if rel_match:
+                rel_name, probe = rel_match
+                desc = ""
+                if "#" in probe:
+                    desc = probe.split("#", 1)[1].strip()
+                schema_info["relation_types"][rel_name] = {
+                    "description": desc,
+                    "domain": "",
+                    "range": "",
+                }
+                current_rel_type = rel_name
+                continue
+
+            # 匹配 domain/range 行: "      domain: Topic, range: Topic"
+            if current_rel_type:
+                if "domain:" in line:
+                    d_val = line.split("domain:", 1)[1].split(",")[0].strip()
+                    schema_info["relation_types"][current_rel_type]["domain"] = d_val
+                if "range:" in line:
+                    r_val = line.split("range:", 1)[1].split(",")[0].strip()
+                    schema_info["relation_types"][current_rel_type]["range"] = r_val
+
+    return schema_info
+
+
+# ═══════════════════════════════════════════════════════════
+# 3. Prompt 构建
+# ═══════════════════════════════════════════════════════════
+
+def build_system_prompt(schema: dict[str, Any]) -> str:
+    """根据 schema 构建系统提示词"""
+    # 节点类型描述
+    node_lines = []
+    for name, info in schema["node_types"].items():
+        desc = f" ({info['description']})" if info["description"] else ""
+        props = f" 属性: {', '.join(info['properties'])}" if info["properties"] else ""
+        node_lines.append(f"  - {name}{desc}")
+        if props:
+            node_lines.append(f"    {props}")
+
+    # 关系类型描述
+    rel_lines = []
+    for name, info in schema["relation_types"].items():
+        desc = f" ({info['description']})" if info["description"] else ""
+        dom = f" domain: {info['domain']}" if info["domain"] else ""
+        ran = f" range: {info['range']}" if info["range"] else ""
+        rel_lines.append(f"  - {name}{desc}{dom}{ran}")
+
+    system_prompt = f"""你是一个知识抽取专家。请从教材文本中抽取符合以下本体定义的实体和关系。
+
+## 节点类型（实体）
+{chr(10).join(node_lines)}
+
+## 关系类型
+{chr(10).join(rel_lines)}
+
+## 抽取规则
+1. 只抽取文本中**明确提到**或**可以明确推导**的实体和关系
+2. 实体名称要**精确**，使用教材中的术语（如 "let关键字" 而非 "let"）
+3. ⚠️ **实体类型必须严格从上面定义的节点类型中选择**，不允许发明任何新的类型（如 "Tool"、"Variable" 等）——如果拿不准，用 Concept
+4. ⚠️ **关系类型必须严格从上面定义的关系类型中选择**，不允许发明新的关系名
+5. ⚠️ **关系方向必须遵守 domain→range 定义**，例如 hasExample 的 domain 是 Topic，range 是 CodeExample，即 (知识点)→[hasExample]→(代码示例)，**不能反过来**
+6. ⚠️ **hasPrerequisite** 的方向是 (前置知识)→[hasPrerequisite]→(后续知识)，即 A 是 B 的前置知识 → (A)→[hasPrerequisite]→(B)
+7. CodeExample 的 codeSnippet 属性保存代码内容，language 属性保存编程语言
+8. 对于 Concept（原子概念），properties 中可包含 level、difficulty 等
+9. 注意层级：示例代码用 CodeExample，知识要点用 Concept 或 Topic
+10. 不要重复抽取同一个实体（同名+同类型视为同一实体）
+11. 习题中的题目用 Question 类型
+"""
+    return system_prompt
+
+
+def build_user_prompt(chunk_id: str, title: str, text: str, examples: list[str]) -> str:
+    """构建用户提示词（单块内容）"""
+    example_info = ""
+    if examples:
+        example_info = f"\n本块包含示例: {', '.join(examples)}"
+
+    user_prompt = f"""请从以下教材段落中抽取实体和关系。
+
+## Chunk 信息
+- ID: {chunk_id}
+- 标题: {title}{example_info}
+
+## 文本内容
+{text}
+
+请严格按照本体定义的类型进行抽取。
+"""
+    return user_prompt
+
+
+# ═══════════════════════════════════════════════════════════
+# 4. LLM 抽取
+# ═══════════════════════════════════════════════════════════
+
+def extract_from_chunk(
+    chunk_id: str,
+    title: str,
+    text: str,
+    examples: list[str],
+    system_prompt: str,
+    client: OpenAI | None = None,
+) -> ExtractionResult:
+    """对单个 Chunk 调用 LLM 抽取实体和关系（JSON mode）"""
+    llm_client = client or _get_client()
+    user_prompt = build_user_prompt(chunk_id, title, text, examples)
+
+    # JSON mode 要求 system prompt 中出现 "JSON" 关键词
+    json_system = system_prompt + "\n\n请以 JSON 格式输出，严格按照以下结构：\n{\n  \"entities\": [{\"name\": \"...\", \"type\": \"Topic|Concept|Skill|CodeExample|TextSegment|Question\", \"properties\": {}}],\n  \"relations\": [{\"source\": \"实体名称\", \"target\": \"实体名称\", \"type\": \"hasPrerequisite|isA|commonlyConfusedWith|teaches|hasExample|assessedBy\", \"properties\": {}}]\n}\n\n⚠️ type 字段只允许以上列表中的值，不能自己发明。"
+
+    for attempt in range(3):
+        try:
+            response = llm_client.chat.completions.create(
+                model=_MODEL,
+                messages=[
+                    {"role": "system", "content": json_system},
+                    {"role": "user", "content": user_prompt},
+                ],
+                response_format={"type": "json_object"},
+                temperature=0.1,
+                max_tokens=4096,
+            )
+            raw = response.choices[0].message.content
+            if not raw:
+                raise ValueError("LLM 返回空内容")
+            data = json.loads(raw)
+            # 用 Pydantic 校验
+            result = ExtractionResult.model_validate(data)
+            return result
+        except Exception as e:
+            if attempt < 2:
+                wait = (attempt + 1) * 3
+                print(f"  ⏳ 重试 {attempt + 1}/3 ({wait}s)... 错误: {e}")
+                import time
+                time.sleep(wait)
+            else:
+                print(f"  ❌ [{chunk_id}] 抽取失败 (3次重试): {e}")
+                return ExtractionResult(entities=[], relations=[])
+
+
+# ═══════════════════════════════════════════════════════════
+# 5. 批量抽取 + 合并
+# ═══════════════════════════════════════════════════════════
+
+def merge_results(results: list[ExtractionResult]) -> ExtractionResult:
+    """合并多个 Chunk 的抽取结果（去重）"""
+    seen_entities: set[tuple[str, str]] = set()  # (name, type)
+    seen_relations: set[tuple[str, str, str]] = set()  # (source, target, type)
+
+    all_entities: list[Entity] = []
+    all_relations: list[Relation] = []
+
+    for r in results:
+        for e in r.entities:
+            key = (e.name.strip(), e.type.strip())
+            if key not in seen_entities:
+                seen_entities.add(key)
+                all_entities.append(e)
+        for rel in r.relations:
+            key = (rel.source.strip(), rel.target.strip(), rel.type.strip())
+            if key not in seen_relations:
+                seen_relations.add(key)
+                all_relations.append(rel)
+
+    return ExtractionResult(entities=all_entities, relations=all_relations)
+
+
+# ═══════════════════════════════════════════════════════════
+# 6. CLI 入口
+# ═══════════════════════════════════════════════════════════
+
+def main(argv: list[str] | None = None) -> None:
+    """CLI 入口：解析 schema → 加载 chunk → LLM 抽取 → 输出结果。
+
+    用法: kg-extractor <txt_path>
+          kg-extractor （使用 KG_EXTRACTOR_INPUT / KG_SCHEMA_PATH 环境变量）
+    """
+    if argv is None:
+        argv = sys.argv[1:]
+
+    from dotenv import load_dotenv
+    load_dotenv()
+
+    # 1. Schema 路径：包内数据或环境变量覆盖
+    schema_path_str = os.environ.get("KG_SCHEMA_PATH")
+    if schema_path_str:
+        schema_path = Path(schema_path_str)
+    else:
+        from importlib.resources import files
+        schema_path = files("knowledge_graph_kit").joinpath("schema.txt")
+
+    print(f"📋 加载本体: {schema_path}")
+    schema = parse_schema(schema_path)
+    node_types = list(schema["node_types"].keys())
+    rel_types = list(schema["relation_types"].keys())
+    print(f"   节点类型: {', '.join(node_types)}")
+    print(f"   关系类型: {', '.join(rel_types)}")
+
+    # 2. 构建 system prompt
+    system_prompt = build_system_prompt(schema)
+
+    # 3. 输入文本路径：CLI 参数或环境变量
+    if len(argv) >= 1:
+        txt_path = argv[0]
+    elif os.environ.get("KG_EXTRACTOR_INPUT"):
+        txt_path = os.environ["KG_EXTRACTOR_INPUT"]
+    else:
+        print("用法: kg-extractor <txt_path>", file=sys.stderr)
+        print("      或设置 KG_EXTRACTOR_INPUT 环境变量", file=sys.stderr)
+        sys.exit(1)
+
+    from .chunker import chunk_file  # type: ignore[import-untyped]
+
+    print(f"\n📄 加载文本: {txt_path}")
+    chunks = chunk_file(txt_path)
+    print(f"   共 {len(chunks)} 个 Chunk")
+
+    # 只处理有实际内容的 chunk（跳过纯标题）
+    non_empty = [c for c in chunks if len(c.text.strip()) > 20]
+    print(f"   有效 Chunk: {len(non_empty)} 个")
+
+    # 4. 逐个抽取（测试阶段限制数量）
+    max_chunks = int(os.environ.get("EXTRACT_CHUNKS", "3"))
+    test_chunks = non_empty[:max_chunks]
+
+    all_results: list[ExtractionResult] = []
+    for i, c in enumerate(test_chunks):
+        text_len = len(c.text)
+        print(f"\n🔍 [{i+1}/{len(test_chunks)}] 抽取: [{c.chunk_id}] {c.title} ({text_len} 字符)")
+        chunk_text = c.text[:6000] if text_len > 6000 else c.text
+        result = extract_from_chunk(
+            chunk_id=c.chunk_id,
+            title=c.title,
+            text=chunk_text,
+            examples=[f"{a}-{b}" for a, b in c.examples],
+            system_prompt=system_prompt,
+        )
+        print(f"   → 实体: {len(result.entities)} 个, 关系: {len(result.relations)} 条")
+        for e in result.entities:
+            props_str = f" {dict(list(e.properties.items())[:3])}" if e.properties else ""
+            print(f"    实体: [{e.type}] {e.name}{props_str}")
+        for rel in result.relations:
+            print(f"    关系: ({rel.source})-[{rel.type}]->({rel.target})")
+        all_results.append(result)
+
+    # 5. 合并去重
+    merged = merge_results(all_results)
+    print(f"\n{'='*50}")
+    print(f"📊 合并结果:")
+    print(f"   实体: {len(merged.entities)} 个")
+    print(f"   关系: {len(merged.relations)} 条")
+
+    # 6. 保存结果
+    out_path_str = os.environ.get("KG_EXTRACTOR_OUTPUT", "extraction_result.json")
+    out_path = Path(out_path_str)
+    out_path.write_text(
+        merged.model_dump_json(indent=2, ensure_ascii=False),
+        encoding="utf-8",
+    )
+    print(f"✅ 已保存到: {out_path}")
+
+
+if __name__ == "__main__":
+    main()

knowledge_graph_kit/neo4j_writer.py

@@ -0,0 +1,273 @@
+"""
+neo4j_writer.py — 将抽取结果写入 Neo4j
+
+流程:
+  1. 读取清洗后的 extraction_result_clean.json
+  2. 连接 Neo4j，创建约束
+  3. 创建节点（实体类型 → Neo4j Label）
+  4. 创建关系
+  5. 添加 __Entity__ 标签（LlamaIndex 兼容）
+  6. 打印统计信息验证
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+from neo4j import GraphDatabase
+
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding="utf-8")  # type: ignore[attr-defined]
+    except Exception:
+        pass
+
+# ── Neo4j 配置 ──────────────────────────────────────────────
+NEO4J_URI = os.getenv("NEO4J_URI", "bolt://localhost:7687")
+NEO4J_USER = os.getenv("NEO4J_USERNAME", "neo4j")
+NEO4J_PASS = os.getenv("NEO4J_PASSWORD", "12345678")
+
+# ── 类型 → Label 映射 ──────────────────────────────────────
+# 实体类型名 → Neo4j 主标签 + 是否添加 Topic 父标签
+LABEL_MAP: dict[str, tuple[str, bool]] = {
+    "Topic": ("Topic", False),
+    "Concept": ("Concept", True),   # 同时打 :Topic 标签（继承关系）
+    "Skill": ("Skill", True),       # 同时打 :Topic 标签
+    "CodeExample": ("CodeExample", False),
+    "TextSegment": ("TextSegment", False),
+    "Question": ("Question", False),
+    "Video": ("Video", False),
+    "Student": ("Student", False),
+    "StudyRecord": ("StudyRecord", False),
+}
+
+
+# ═══════════════════════════════════════════════════════════
+# 1. Neo4j 操作
+# ═══════════════════════════════════════════════════════════
+
+class Neo4jWriter:
+    """封装 Neo4j 写入操作"""
+
+    def __init__(self, uri: str, user: str, password: str):
+        self.driver = GraphDatabase.driver(uri, auth=(user, password))
+
+    def close(self):
+        self.driver.close()
+
+    # ── 约束 ──────────────────────────────────────────────
+
+    def create_constraints(self):
+        with self.driver.session() as sess:
+            # 所有实体节点共用的唯一约束
+            sess.run("CREATE CONSTRAINT IF NOT EXISTS FOR (e:Entity) REQUIRE e.name IS UNIQUE")
+            # 节点类型上的索引（可选）
+            sess.run("CREATE INDEX IF NOT EXISTS FOR (e:Entity) ON (e.entity_type)")
+            print("  ✅ 约束已创建")
+
+    # ── 清空 ──────────────────────────────────────────────
+
+    def clear_all(self):
+        with self.driver.session() as sess:
+            sess.run("MATCH (n) DETACH DELETE n")
+        print("  🧹 已清空所有数据")
+
+    # ── 写入节点 ──────────────────────────────────────────
+
+    def write_entities(self, entities: list[dict]) -> int:
+        count = 0
+        with self.driver.session() as sess:
+            for e in entities:
+                name = e["name"]
+                etype = e["type"]
+                props = {k: v for k, v in e.get("properties", {}).items() if v is not None}
+
+                main_label, add_topic_label = LABEL_MAP.get(etype, (etype, False))
+
+                # 构建标签字符串
+                labels = [main_label]
+                if add_topic_label:
+                    labels.append("Topic")
+
+                # 动态构建 MERGE 语句（标签不能作为参数传入）
+                label_str = ":" + ":".join(labels)
+                cypher = f"""
+                    MERGE (n:Entity {{name: $name}})
+                    ON CREATE SET
+                        n:{main_label},
+                        n.entity_type = $etype,
+                        n.id = $name
+                    SET n += $props
+                """
+                # 如果是 Topic 子类，额外增加标签
+                if add_topic_label:
+                    cypher = f"""
+                        MERGE (n:Entity {{name: $name}})
+                        ON CREATE SET
+                            n:{main_label},
+                            n:Topic,
+                            n.entity_type = $etype,
+                            n.id = $name
+                        SET n += $props
+                    """
+
+                sess.run(cypher, name=name, etype=etype, props=props)
+                count += 1
+
+        return count
+
+    # ── 写入关系 ──────────────────────────────────────────
+
+    def write_relations(self, relations: list[dict]) -> int:
+        count = 0
+        with self.driver.session() as sess:
+            for rel in relations:
+                src = rel["source"]
+                tgt = rel["target"]
+                rtype = rel["type"]
+                props = {k: v for k, v in rel.get("properties", {}).items() if v is not None}
+
+                cypher = f"""
+                    MATCH (a:Entity {{name: $src}})
+                    MATCH (b:Entity {{name: $tgt}})
+                    MERGE (a)-[r:{rtype}]->(b)
+                    SET r += $props
+                """
+                try:
+                    sess.run(cypher, src=src, tgt=tgt, props=props)
+                    count += 1
+                except Exception as ex:
+                    print(f"  ⚠️  关系写入失败: ({src})-[{rtype}]->({tgt}) → {ex}")
+
+        return count
+
+    # ── 兼容 LlamaIndex ──────────────────────────────────
+
+    def ensure_llamaindex_compat(self):
+        """添加 __Entity__ 标签以兼容 LlamaIndex PropertyGraphIndex"""
+        with self.driver.session() as sess:
+            sess.run("MATCH (e:Entity) SET e:`__Entity__`")
+            sess.run("MATCH (e:Entity) SET e.id = e.name")
+        print("  🔄 LlamaIndex 兼容标签已添加")
+
+    # ── 统计 ──────────────────────────────────────────────
+
+    def print_stats(self):
+        with self.driver.session() as sess:
+            # 节点数
+            total = sess.run("MATCH (n:Entity) RETURN count(n) AS c").single()["c"]
+            # 标签分布
+            labels = sess.run(
+                "MATCH (n:Entity) RETURN n.entity_type AS type, count(*) AS cnt ORDER BY cnt DESC"
+            ).data()
+            # 关系分布
+            rels = sess.run(
+                "CALL db.relationshipTypes() YIELD relationshipType AS type "
+                "RETURN type ORDER BY type"
+            ).values()
+            rel_counts = []
+            for rt in rels:
+                cnt = sess.run(
+                    f"MATCH ()-[r:{rt[0]}]->() RETURN count(r) AS c"
+                ).single()["c"]
+                rel_counts.append((rt[0], cnt))
+
+            print(f"\n{'='*50}")
+            print(f"📊 Neo4j 统计:")
+            print(f"   节点总数: {total}")
+            print(f"   类型分布:")
+            for row in labels:
+                print(f"     - {row['type']:12s}: {row['cnt']}")
+            print(f"   关系:")
+            for name, cnt in rel_counts:
+                print(f"     - {name}: {cnt}")
+
+
+# ═══════════════════════════════════════════════════════════
+# 2. 主流程
+# ═══════════════════════════════════════════════════════════
+
+def main(argv: list[str] | None = None) -> None:
+    """CLI 入口：将抽取结果写入 Neo4j。
+
+    用法: kg-neo4j-writer <input.json>
+          kg-neo4j-writer （使用 KG_NEO4J_INPUT 环境变量）
+          NEO4J_URI / NEO4J_USERNAME / NEO4J_PASSWORD 环境变量用于连接配置
+    """
+    if argv is None:
+        argv = sys.argv[1:]
+
+    from dotenv import load_dotenv
+    load_dotenv()
+
+    # 1. 加载数据
+    if len(argv) >= 1:
+        src_path = Path(argv[0])
+    elif os.environ.get("KG_NEO4J_INPUT"):
+        src_path = Path(os.environ["KG_NEO4J_INPUT"])
+    else:
+        for candidate in ["extraction_result_clean.json", "extraction_result.json"]:
+            p = Path(candidate)
+            if p.exists():
+                src_path = p
+                break
+        else:
+            print("用法: kg-neo4j-writer <input.json>", file=sys.stderr)
+            print("      或设置 KG_NEO4J_INPUT 环境变量", file=sys.stderr)
+            sys.exit(1)
+
+    data = json.loads(src_path.read_text(encoding="utf-8"))
+    entities: list[dict] = data["entities"]
+    relations: list[dict] = data["relations"]
+
+    print(f"📥 加载数据: {src_path.name}")
+    print(f"   实体: {len(entities)} 个")
+    print(f"   关系: {len(relations)} 条")
+
+    # 2. 连接 Neo4j
+    print(f"\n🔗 连接 Neo4j: {NEO4J_URI}")
+    writer = Neo4jWriter(NEO4J_URI, NEO4J_USER, NEO4J_PASS)
+    try:
+        # 测试连接
+        writer.driver.verify_connectivity()
+        print("   ✅ 连接成功")
+    except Exception as e:
+        print(f"   ❌ 连接失败: {e}")
+        writer.close()
+        return
+
+    # 3. 创建约束
+    print("\n📌 创建约束...")
+    writer.create_constraints()
+
+    # 4. 清空旧数据（可选，通过环境变量控制）
+    if os.environ.get("NEO4J_CLEAR", "").lower() in ("1", "true", "yes"):
+        writer.clear_all()
+
+    # 5. 写入节点
+    print("\n📝 写入节点...")
+    n_count = writer.write_entities(entities)
+    print(f"   ✅ 写入 {n_count} 个节点")
+
+    # 6. 写入关系
+    print("\n🔗 写入关系...")
+    r_count = writer.write_relations(relations)
+    print(f"   ✅ 写入 {r_count} 条关系")
+
+    # 7. LlamaIndex 兼容
+    print("\n🔄 添加 LlamaIndex 兼容标签...")
+    writer.ensure_llamaindex_compat()
+
+    # 8. 统计
+    writer.print_stats()
+
+    writer.close()
+    print("\n🎉 写入完成！")
+
+
+if __name__ == "__main__":
+    main()