# DISK 项目组件关系文档

本文档详细描述了 DISK（Domain Incremental conStruction of Knowledge graph）项目的核心组件、方法及其依赖关系。

## 1. 核心控制器 (Main Controller)

### `disk.py` -> `DISK` 类
作为项目的主入口，负责协调整个知识图谱构建流水线（蒸馏 -> 抽取 -> 合并 -> 管理）。

*   **方法**:
    *   `build_knowledge_graph(pdf_path, ...)`: 核心流水线方法。执行流程：文本蒸馏 -> 语言检测 -> 并行/串行抽取实体与关系 -> 分批两两归并合并 -> 增量写入 KG 管理器。
    *   `_extract_parallel(texts, max_workers)`: 使用线程池并行调用 Extractor 处理文本块。
    *   `_extract_serial(texts)`: 顺序处理文本块。
    *   `_batch_merge(results, ...)`: 采用两两归并算法对提取出的实体和关系进行并行合并。
    *   `visualize_knowledge_graph(...)`: 将构建好的 KG 写入 Neo4j 数据库。
*   **依赖组件**:
    *   `distiller.PDFDistiller`: 用于文档蒸馏。
    *   `extractor.Extractor`: 用于信息抽取。
    *   `merger.Merger`: 用于语义合并。
    *   `manager.KGManager`: 用于管理内存中的知识图谱。
    *   `models.neo4j_connector.Neo4jConnector`: 用于数据库持久化。
    *   `utils.lang_detect`: 用于自动识别文档语言。

---

## 2. 蒸馏组件 (Distillation)

### `distiller/pdf_distiller.py` -> `PDFDistiller` 类
负责从 PDF 文档中提取原始信息并进行清洗。

*   **方法**:
    *   `extract_text_blocks(file_path)`: 提取页面文本块，合并短块，并根据规则过滤无效内容。
    *   `is_valid_block(text)`: 验证文本块有效性（检查字符、词数、参考文献格式、句子完整性）。
    *   `extract_images_and_ocr(file_path)`: 提取文档图片并调用 OCR 模型识别文字。
    *   `extract_tables(file_path)`: 使用 `pdfplumber` 提取表格数据。
*   **依赖组件**:
    *   `fitz` (PyMuPDF): 文本和图片提取。
    *   `pdfplumber`: 表格提取。
    *   `rapidocr_onnxruntime`: OCR 识别。
    *   `jieba`: 中文分词辅助验证。

---

## 3. 抽取组件 (Extraction)

### `extractor/extractor.py` -> `Extractor` 类
利用 LLM 将非结构化文本转换为结构化的实体和关系，并生成向量嵌入。

*   **方法**:
    *   `extract_relations_and_entities(text, pdf_path)`: 调用 LLM 一次性提取关系和实体，并生成对应的 Embedding。
    *   `embed_entity(entity)`: 为单个实体生成加权 Embedding（基于名称、标签和描述），支持缓存。
    *   `embed_relations(relations)`: 为提取出的关系生成 Embedding。
*   **依赖组件**:
    *   `utils.parser.Parser`: 封装 LangChain 逻辑，负责 JSON 解析。
    *   `utils.prompts`: 获取对应语言的 Prompt 模板。
    *   `utils.schemas`: 定义 Pydantic 数据结构。
    *   `models.Entity` / `models.Relation`: 实例化数据模型。

---

## 4. 合并组件 (Merging)

### `merger/merger.py` -> `Merger` 类
基于语义相似度对提取结果进行合并，解决实体冲突和重复。

*   **方法**:
    *   `merge(entities1, entities2, relations1, relations2)`: 计算两组实体间的余弦相似度矩阵，合并超过阈值的实体。
    *   `update_and_merge_relations(...)`: 实体合并后，同步更新关系中的实体引用。
*   **依赖组件**:
    *   `numpy`: 高效计算余弦相似度矩阵。
    *   `models.Entity` / `models.Relation`: 处理数据实例。

---

## 5. 管理组件 (Management)

### `manager/kg_manager.py` -> `KGManager` 类
维护内存中的 `KnowledgeGraph` 实例。

*   **方法**:
    *   `add_entities(entities)`: 增量添加实体，自动去重。
    *   `add_relations(relations)`: 增量添加关系，自动去重。
*   **依赖组件**:
    *   `models.KnowledgeGraph`: 核心数据容器。

---

## 6. 模型与连接器 (Models & Connectors)

### `models/knowledge_graph.py`
定义核心数据结构：`Entity`、`Relation`、`KnowledgeGraph`。

### `models/neo4j_connector.py` -> `Neo4jConnector` 类
负责将内存中的 KG 持久化到图数据库。

*   **方法**:
    *   `create_entities(entities)`: 在 Neo4j 中创建或更新节点（使用 `MERGE`）。
    *   `create_relations(relations)`: 在 Neo4j 中创建或更新边，支持 `graph_id` 隔离。
*   **依赖组件**:
    *   `neo4j` (官方驱动): 执行 Cypher 查询。

---

## 7. 工具类 (Utilities)

*   **`utils/parser.py`**: 核心解析器，基于 LangChain 的 `JsonOutputParser`。
*   **`utils/prompts.py`**: 存储中英文 Prompt 模板。
*   **`utils/schemas.py`**: 存储 Pydantic 结构定义。
*   **`utils/lang_detect.py`**: 结合 FastText 或正则表达式检测文档语言。
*   **`utils/checkpoint_helper.py`**: 处理构建过程中的断点续传。

---

## 组件依赖关系图 (简图)

```mermaid
graph TD
    DISK --> PDFDistiller
    DISK --> Extractor
    DISK --> Merger
    DISK --> KGManager
    DISK --> Neo4jConnector

    Extractor --> Parser
    Extractor --> Models
    Extractor --> Utils(Prompts/Schemas)

    Merger --> Models
    KGManager --> Models
    Neo4jConnector --> Models

    Parser --> LLM/Embeddings
```
