AI知识图谱RAG实践笔记

def extract_entities_relations(chunk_text):
    prompt = f"""
从以下文本中提取实体和关系，按照给定的类型约束输出 JSON。

文本：
{chunk_text}

实体类型：服务、组件、配置、错误、API、产品、版本
关系类型：调用、依赖、导致、修复、包含、兼容

输出格式：
{{
  "entities": [{{"text": "...", "type": "..."}}],
  "relations": [{{"source": "...", "target": "...", "relation": "..."}}]
}}
"""
    response = llm.generate(prompt, temperature=0.1)
    return parse_json(response)

import networkx as nx

def build_graph(extractions):
    G = nx.DiGraph()
    for extraction in extractions:
        for entity in extraction["entities"]:
            G.add_node(entity["text"], type=entity["type"])
        for relation in extraction["relations"]:
            source = relation["source"]
            target = relation["target"]
            rel_type = relation["relation"]
            if G.has_node(source) and G.has_node(target):
                if G.has_edge(source, target):
                    G[source][target]["relations"].append(rel_type)
                else:
                    G.add_edge(source, target, relations=[rel_type])
    return G

import json

def save_graph(G, filepath):
    # 转换为可序列化格式
    graph_data = nx.node_link_data(G)
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(graph_data, f, ensure_ascii=False, indent=2)

def extract_question_entities(question):
    prompt = f"""
从以下问题中提取涉及的技术实体，只返回实体列表。

问题：{question}

实体类型：服务、组件、配置、错误、API、产品、版本

输出格式：["实体1", "实体2", ...]
"""
    response = llm.generate(prompt, temperature=0)
    entities = parse_list(response)
    return entities

def retrieve_context(G, entities, max_hops=2):
    context_nodes = set(entities)
    for entity in entities:
        if G.has_node(entity):
            # 1-2 跳邻居
            for node, distance in nx.single_source_shortest_path_length(
                G, entity, cutoff=max_hops
            ).items():
                context_nodes.add(node)

    # 提取子图
    subgraph = G.subgraph(context_nodes).copy()
    return subgraph

def graph_to_text(G, start_entities):
    visited = set()
    result = []

    def bfs(start):
        queue = [(start, 0)]
        while queue:
            node, depth = queue.pop(0)
            if node in visited:
                continue
            visited.add(node)
            node_type = G.nodes[node].get("type", "未知")
            result.append(f"【{node_type}】{node}")

            for neighbor in G.neighbors(node):
                if neighbor not in visited:
                    edge_data = G[node][neighbor]
                    for rel in edge_data.get("relations", []):
                        result.append(f"→ {rel} → 【{G.nodes[neighbor].get('type', '未知')}】{neighbor}")
                    queue.append((neighbor, depth + 1))

    for entity in start_entities:
        bfs(entity)

    return "\n".join(result)

def generate_answer(question, graph_context):
    prompt = f"""
基于以下知识图谱上下文回答问题，只回答问题中的内容，不要编造。

问题：{question}

知识图谱上下文：
{graph_context}

答案：
"""
    response = llm.generate(prompt, temperature=0.3)
    return response

def prune_graph(G, min_degree=2):
    # 删除低度节点
    low_degree_nodes = [node for node, degree in G.degree() if degree <= min_degree]
    G.remove_nodes_from(low_degree_nodes)
    return G

G.add_node(entity, type=entity_type, sources=[doc_id])
G.add_edge(source, target, relations=[rel_type], sources=[doc_id])

import matplotlib.pyplot as plt
import matplotlib.font_manager as fm

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 数据
methods = ['向量检索', '图谱RAG']
accuracy = [0.68, 0.82]

fig, ax = plt.subplots(figsize=(8, 5))
bars = ax.bar(methods, accuracy, color=['#6c757d', '#28a745'])

# 添加数值标签
for bar, acc in zip(bars, accuracy):
    height = bar.get_height()
    ax.text(bar.get_x() + bar.get_width()/2., height,
            f'{acc:.0%}',
            ha='center', va='bottom', fontsize=12)

ax.set_ylabel('准确度', fontsize=12)
ax.set_ylim(0, 1)
ax.set_title('图谱RAG vs 向量检索准确度对比', fontsize=14, pad=20)
ax.grid(axis='y', alpha=0.3)

plt.tight_layout()
plt.savefig('/tmp/accuracy_comparison.png', dpi=300, bbox_inches='tight')
plt.close()