# conda activate rag
# pip install py2neo pyvis
from py2neo import Graph
from pyvis.network import Network
from Config import Config

# 连接Neo4j
graph = Graph(Config.NEO4J_URI, auth=Config.NEO4J_AUTH)

# 修改后的查询语句，使用INCLUDES关系类型
query = """
MATCH p=()-[r:INCLUDES]->() 
RETURN p LIMIT 250
"""
data = graph.run(query).data()

# 创建网络图
net = Network(height="750px", width="100%", notebook=True, cdn_resources='in_line')

# 添加节点和边
for item in data:
    # 提取起始节点
    start_node = item['p'].start_node
    # 提取结束节点
    end_node = item['p'].end_node
    
    net.add_node(start_node.identity,
                label=start_node['name'],
                title=start_node.get('description', ''),
                group=list(start_node.labels)[0])
    net.add_node(end_node.identity,
                label=end_node['name'],
                title=end_node.get('description', ''),
                group=list(end_node.labels)[0])
    net.add_edge(start_node.identity,
                end_node.identity,
                title="INCLUDES")

# 生成HTML文件
with open("knowledge_graph.html", "w", encoding='utf-8') as f:
    f.write(net.generate_html())

# 打开HTML文件
import webbrowser
webbrowser.open("knowledge_graph.html")