main_GraphBuild.py

'''
此部分代码负责构建知识图谱
1、读取数据
2、切分文本
3、LLM抽三元组
4、解析三元组
5、写入Kùzu
'''

import os
import re
from typing import Dict, List

# LangChain相关导入
from langchain.prompts import ChatPromptTemplate
from langchain.schema import HumanMessage, SystemMessage

# 自定义模块导入
from Prompts.prompt_graph import (
    graph_extraction, 
    graph_extraction_parameter_default,
    TUPLE_DELIMITER,
    RECORD_DELIMITER,
    COMPLETION_DELIMITER,
    system_prompt
)
from Models.LLM_Models import build_model, stream_invoke, remove_think

# 导入 kuzu_db_path 和 KuzuGraph
from DataBase.config import kuzu_db_path
from DataBase.GraphData.create_graph_knowledge import TripleDataParser, KuzuGraph


# 初始化LLM模型
llm_model = build_model(mode="ollama")

# 使用 KuzuGraph
graph = KuzuGraph(kuzu_db_path)

# 初始化解析器
dataparse = TripleDataParser(
    TUPLE_DELIMITER,
    RECORD_DELIMITER,
    COMPLETION_DELIMITER
)


def build_prompt(text: str, entity_types: str = None) -> str:
    """
    构建完整的提示词
    
    :param text: 输入文本
    :param entity_types: 实体类型列表，如"组织、人物、日期、文件"
    :return: 格式化后的提示词
    """
    params = graph_extraction_parameter_default
    params['input_text'] = text
    
    if entity_types is not None:
        params['entity_types'] = entity_types
        
    return graph_extraction.format(**params)


def read_file(file_path: str) -> str:
    """
    读取TXT文件内容
    
    :param file_path: 文件路径
    :return: 文件内容字符串
    :raises FileNotFoundError: 当文件不存在时抛出异常
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"❌ 文件未找到：{file_path}")
        
    with open(file_path, "r", encoding="utf-8") as f:
        return f.read().strip()


def split_text_by_newline(text: str) -> list:
    """
    将文本按换行符分割，每个非空行作为一个列表元素
    
    :param text: 原始文本（字符串）
    :return: 分割后的非空行列表
    """
    return [line.strip() for line in text.splitlines() if line.strip()]


def extract_and_write_graph_data(
    graph_class,
    entity_types: str,
    text_lst: list,
    system_prompt: str,
    batch: int
) -> None:
    """
    提取文本中的图谱数据并写入数据库
    
    :param graph_class: 图谱数据库操作类实例
    :param entity_types: 实体类型列表
    :param text_lst: 文本列表
    :param system_prompt: 系统提示词
    :param batch: 批量处理大小
    """
    # 构建聊天提示模板
    chat_template = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", "{user_input}"),
    ])

    for i, text in enumerate(text_lst):
        # 构建提示词并获取模型响应
        prompt = build_prompt(text, entity_types=entity_types)
        formatted_messages = chat_template.format_messages(user_input=prompt)
        
        # LLM抽三元组
        response = llm_model.invoke(formatted_messages).content
        print("\n===== LLM RESPONSE =====")
        print(response)
        print("========================\n")
        
        # 去除无关思维链输出
        response = remove_think(response)
        if isinstance(response, tuple):
            response = response[0]
        print(type(response))
        print(response)

        #进一步进行文本清洗，根据实际情况调整
        response = (
            response
            .replace("```", "") # 去除 Markdown 代码块标记
            .replace("\r", "") ## 去除 Windows 风格回车符
        )

        response = response.strip()# 去除首尾空白字符
        
        # 解析三元组
        extract_graph_text = {"graph_text": response, "textid": str(i)}
        entities, relationships = dataparse.batch_parse_data(
            extract_graph_text, 
            batch=batch
        )
        # 写入Kùzu数据库
        graph_class.write_entities_and_relationships(entities, relationships)

    # 处理缓冲区中剩余的数据
    if dataparse._buffer:
        entities, relationships = dataparse.flush()
        graph_class.write_entities_and_relationships(entities, relationships)


def write_graph_demo() -> None:
    """图谱数据写入示例函数"""
    # 配置参数
    text_file_path =  './example_data.txt'   
    entity_types = "组织、人物、文件"
    batch_size = 1

    # 调用函数读取数据
    text_file = read_file(text_file_path)
    # 调用函数按换行符切分文本
    text_lst = split_text_by_newline(text_file)
    
    # 调用函数提取文本中的图谱数据并写入数据库
    extract_and_write_graph_data(
        graph, 
        entity_types, 
        text_lst, 
        system_prompt, 
        batch_size
    )
    
    # 关闭数据库连接（KuzuGraph 的 close 是空操作，但保留兼容性）
    graph.close()


if __name__ == "__main__":
    write_graph_demo()