原始文档
    ↓
[1] 文本分块
    ↓
[2] 实体 & 关系抽取（LLM）
    ↓
[3] 知识图谱构建（实体节点 + 关系边）
    ↓
[4] 社区检测（Leiden 算法）
    ↓
[5] 层级社区摘要（LLM）
    ↓
[6] 检索：局部检索（实体子图） + 全局检索（社区摘要）
    ↓
[7] 生成

from openai import OpenAI
import json
import re

client = OpenAI()

def extract_entities_relations(
    text: str,
    entity_types: list[str] = None,
    model: str = "gpt-4o-mini",
) -> dict:
    """
    用 LLM 从文本中抽取实体和关系
    """
    entity_types = entity_types or ["person", "organization", "location", "event", "concept"]

    prompt = f"""你是一个信息抽取专家。请从以下文本中抽取实体和关系。

实体类型：{', '.join(entity_types)}

要求：
1. 抽取所有提到的实体，给出名称、类型和简短描述
2. 抽取实体之间的关系，给出关系类型和描述
3. 合并相同实体（消歧）

返回 JSON 格式：
{{
  "entities": [
    {{"name": "实体名", "type": "类型", "description": "描述"}}
  ],
  "relations": [
    {{"source": "实体A", "target": "实体B", "type": "关系类型", "description": "描述"}}
  ]
}}

文本：
{text}
"""
    resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        response_format={"type": "json_object"},
        temperature=0.0,
    )

    result = json.loads(resp.choices[0].message.content)
    return result


# 批量处理文档块
def build_graph_from_chunks(chunks: list[str]) -> dict:
    """从多个文本块构建知识图谱"""
    all_entities = {}
    all_relations = []

    for i, chunk in enumerate(chunks):
        result = extract_entities_relations(chunk)

        # 合并实体（同名实体合并描述）
        for ent in result["entities"]:
            name = ent["name"].lower().strip()
            if name in all_entities:
                # 合并描述
                existing = all_entities[name]
                existing["description"] += f" | {ent['description']}"
                existing["source_chunks"].append(i)
            else:
                ent["source_chunks"] = [i]
                all_entities[name] = ent

        for rel in result["relations"]:
            rel["source_chunk"] = i
            all_relations.append(rel)

    return {
        "entities": list(all_entities.values()),
        "relations": all_relations,
    }

def entity_resolution(entities: list[dict]) -> list[dict]:
    """
    实体消歧：合并指代同一实体的不同名称
    """
    # 用嵌入相似度 + LLM 判断
    import numpy as np
    from sentence_transformers import SentenceTransformer

    model = SentenceTransformer("BAAI/bge-large-zh-v1.5")

    # 计算实体名+描述的嵌入
    texts = [f"{e['name']} {e['description']}" for e in entities]
    embeddings = model.encode(texts)

    # 找相似实体对
    from sklearn.metrics.pairwise import cosine_similarity
    sim_matrix = cosine_similarity(embeddings)

    merged = set()
    clusters = []

    for i in range(len(entities)):
        if i in merged:
            continue
        cluster = [i]
        for j in range(i + 1, len(entities)):
            if j in merged:
                continue
            if sim_matrix[i][j] > 0.85:  # 高相似度阈值
                cluster.append(j)
                merged.add(j)
        merged.add(i)
        clusters.append(cluster)

    # 合并同簇实体
    resolved = []
    for cluster in clusters:
        primary = entities[cluster[0]]
        if len(cluster) > 1:
            aliases = [entities[c]["name"] for c in cluster[1:]]
            primary["aliases"] = aliases
        resolved.append(primary)

    return resolved

from neo4j import GraphDatabase

class KnowledgeGraph:
    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_entity(self, name: str, entity_type: str, description: str):
        with self.driver.session() as session:
            session.execute_write(
                self._create_entity_tx, name, entity_type, description
            )

    @staticmethod
    def _create_entity_tx(tx, name, entity_type, description):
        query = """
        MERGE (e:Entity {name: $name})
        SET e.type = $type, e.description = $description
        """
        tx.run(query, name=name, type=entity_type, description=description)

    def create_relation(self, source: str, target: str, rel_type: str, description: str):
        with self.driver.session() as session:
            session.execute_write(
                self._create_relation_tx, source, target, rel_type, description
            )

    @staticmethod
    def _create_relation_tx(tx, source, target, rel_type, description):
        query = """
        MATCH (s:Entity {name: $source})
        MATCH (t:Entity {name: $target})
        MERGE (s)-[r:RELATION {type: $type}]->(t)
        SET r.description = $description
        """
        tx.run(query, source=source, target=target, type=rel_type, description=description)

    def build_from_graph_data(self, graph_data: dict):
        """批量导入图谱数据"""
        with self.driver.session() as session:
            # 创建实体
            for ent in graph_data["entities"]:
                session.execute_write(
                    self._create_entity_tx,
                    ent["name"], ent["type"], ent["description"]
                )
            # 创建关系
            for rel in graph_data["relations"]:
                session.execute_write(
                    self._create_relation_tx,
                    rel["source"], rel["target"],
                    rel["type"], rel.get("description", "")
                )

    def get_entity_subgraph(self, entity_name: str, hops: int = 2) -> dict:
        """获取实体的 N 跳子图"""
        with self.driver.session() as session:
            result = session.run("""
                MATCH path = (e:Entity {name: $name})-[*1..$hops]-(related)
                RETURN path
            """, name=entity_name, hops=hops)

            nodes, edges = set(), []
            for record in result:
                path = record["path"]
                for node in path.nodes:
                    nodes.add(node["name"])
                for rel in path.relationships:
                    edges.append({
                        "source": rel.start_node["name"],
                        "target": rel.end_node["name"],
                        "type": rel["type"],
                    })
            return {"nodes": list(nodes), "edges": edges}

import networkx as nx

class InMemoryGraph:
    """基于 NetworkX 的内存知识图谱"""
    def __init__(self):
        self.graph = nx.MultiDiGraph()

    def add_entity(self, name: str, **attrs):
        self.graph.add_node(name, **attrs)

    def add_relation(self, source: str, target: str, rel_type: str, **attrs):
        self.graph.add_edge(source, target, type=rel_type, **attrs)

    def get_subgraph(self, entity: str, hops: int = 2) -> nx.Graph:
        """获取 N 跳子图"""
        if entity not in self.graph:
            return nx.Graph()
        nodes = nx.single_source_shortest_path_length(self.graph, entity, cutoff=hops)
        subgraph = self.graph.subgraph(nodes.keys()).copy()
        return subgraph

    def get_all_entities(self) -> list[str]:
        return list(self.graph.nodes())

    def get_stats(self) -> dict:
        return {
            "nodes": self.graph.number_of_nodes(),
            "edges": self.graph.number_of_edges(),
        }

import community as community_louvain
import networkx as nx
from typing import List, Dict, Tuple

def detect_communities(
    graph: nx.Graph,
    resolution: float = 1.0,
    min_community_size: int = 3,
) -> Dict[str, int]:
    """
    使用 Louvain/Leiden 算法检测社区
    返回: {entity_name: community_id}
    """
    # 转为无向图
    undirected = graph.to_undirected()

    # Louvain 社区检测
    partition = community_louvain.best_partition(
        undirected,
        resolution=resolution,
        random_state=42,
    )

    # 过滤小社区
    community_sizes = {}
    for node, comm in partition.items():
        community_sizes[comm] = community_sizes.get(comm, 0) + 1

    valid_communities = {
        comm for comm, size in community_sizes.items()
        if size >= min_community_size
    }

    return {
        node: comm for node, comm in partition.items()
        if comm in valid_communities
    }


def build_hierarchy(
    graph: nx.Graph,
    partition: Dict[str, int],
    max_levels: int = 3,
) -> List[Dict]:
    """
    构建层级社区结构
    Level 0: 所有实体
    Level 1: 一级社区
    Level 2: 社区中的子社区
    """
    hierarchy = []

    # Level 0: 全局
    hierarchy.append({
        "level": 0,
        "community_id": 0,
        "entities": list(graph.nodes()),
        "size": graph.number_of_nodes(),
    })

    # Level 1+: 社区划分
    current_level = 1
    current_partition = partition

    while current_level <= max_levels:
        communities = {}
        for node, comm_id in current_partition.items():
            if comm_id not in communities:
                communities[comm_id] = []
            communities[comm_id].append(node)

        if len(communities) <= 1:
            break

        for comm_id, entities in communities.items():
            hierarchy.append({
                "level": current_level,
                "community_id": comm_id,
                "entities": entities,
                "size": len(entities),
            })

            # 在社区内继续细分
            if len(entities) > 10:
                subgraph = graph.subgraph(entities)
                sub_partition = community_louvain.best_partition(
                    subgraph.to_undirected(),
                    resolution=1.5,  # 更高分辨率 = 更小社区
                )
                # 更新 partition 供下一层使用
                for node in entities:
                    current_partition[node] = f"{comm_id}_{sub_partition[node]}"

        current_level += 1

    return hierarchy

def generate_community_summaries(
    graph: nx.Graph,
    hierarchy: List[Dict],
    llm_client,
) -> List[Dict]:
    """
    为每个社区生成摘要
    """
    for community in hierarchy:
        if community["level"] == 0:
            # 全局摘要
            community["summary"] = generate_global_summary(graph, llm_client)
            continue

        entities = community["entities"]
        # 收集社区内的实体描述和关系
        entity_descs = []
        for ent in entities:
            node_data = graph.nodes[ent]
            entity_descs.append(f"- {ent} ({node_data.get('type', 'unknown')}): {node_data.get('description', '')}")

        # 收集社区内部关系
        internal_edges = []
        for u, v, data in graph.subgraph(entities).edges(data=True):
            internal_edges.append(f"- {u} --[{data['type']}]--> {v}: {data.get('description', '')}")

        # LLM 生成摘要
        prompt = f"""请为以下知识图谱社区生成结构化摘要。

社区实体：
{chr(10).join(entity_descs[:50])}

社区内部关系：
{chr(10).join(internal_edges[:50])}

请输出：
1. 社区主题（一句话）
2. 关键实体列表（5-10个）
3. 主要关系模式（2-3条）
4. 社区整体描述（3-5句话）
"""
        resp = llm_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
        )
        community["summary"] = resp.choices[0].message.content

    return hierarchy

def local_search(
    query: str,
    graph: nx.Graph,
    embed_model,
    vector_index,           # 实体描述的向量索引
    community_summaries: list,
    top_k_entities: int = 5,
    hops: int = 2,
) -> dict:
    """
    局部检索：找到相关实体 → 扩展子图 → 返回结构化上下文
    """
    # Step 1: 向量检索找到相关实体
    entity_scores = vector_index.search(query, top_k=top_k_entities)
    seed_entities = [s["entity"] for s in entity_scores]

    # Step 2: 扩展子图
    subgraph_nodes = set()
    subgraph_edges = []
    for entity in seed_entities:
        if entity in graph:
            sub = nx.single_source_shortest_path_length(graph, entity, cutoff=hops)
            for node in sub:
                subgraph_nodes.add(node)
            for u, v, data in graph.subgraph(sub.keys()).edges(data=True):
                subgraph_edges.append({"source": u, "target": v, **data})

    # Step 3: 构建上下文
    entity_context = []
    for node in subgraph_nodes:
        node_data = graph.nodes[node]
        entity_context.append(f"- {node} ({node_data.get('type', '')}): {node_data.get('description', '')}")

    relation_context = [f"- {e['source']} --[{e['type']}]--> {e['target']}" for e in subgraph_edges]

    return {
        "entities": list(subgraph_nodes),
        "entity_descriptions": entity_context,
        "relations": relation_context,
        "subgraph_size": len(subgraph_nodes),
    }

def global_search(
    query: str,
    community_summaries: list[dict],
    embed_model,
    llm_client,
    top_k_communities: int = 5,
) -> dict:
    """
    全局检索：找到相关社区摘要 → 聚合答案
    适合回答「整个文档集的主要主题」等全局性问题
    """
    # Step 1: 检索相关社区
    summaries = [c["summary"] for c in community_summaries if c["level"] >= 1]
    summary_embeddings = embed_model.encode(summaries)
    query_emb = embed_model.encode([query])

    from sklearn.metrics.pairwise import cosine_similarity
    scores = cosine_similarity(query_emb, summary_embeddings)[0]
    top_indices = scores.argsort()[-top_k_communities:][::-1]

    # Step 2: 每个社区生成局部答案
    partial_answers = []
    for idx in top_indices:
        comm = community_summaries[idx + 1]  # 跳过 level 0
        prompt = f"""基于以下社区摘要回答问题。如果信息不足以回答，请说明。

社区摘要：
{comm['summary']}

问题：{query}

答案："""
        resp = llm_client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
        )
        partial_answers.append({
            "community_id": comm["community_id"],
            "answer": resp.choices[0].message.content,
            "score": scores[idx],
        })

    # Step 3: 聚合答案
    combined = "\n\n".join([
        f"[社区 {a['community_id']}]: {a['answer']}"
        for a in partial_answers
    ])
    final_prompt = f"""以下是多个社区对同一问题的回答，请综合生成最终答案。

问题：{query}

各社区回答：
{combined}

最终答案："""
    final_resp = llm_client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": final_prompt}],
    )

    return {
        "answer": final_resp.choices[0].message.content,
        "partial_answers": partial_answers,
        "communities_used": len(partial_answers),
    }

def hybrid_graph_rag(
    query: str,
    vector_index,       # 传统向量索引
    graph: nx.Graph,    # 知识图谱
    community_summaries: list,
    embed_model,
    llm_client,
    top_k_docs: int = 5,
    top_k_entities: int = 5,
    use_global: bool = False,
) -> str:
    """
    混合检索：同时使用向量检索和图谱检索
    """
    # 1. 传统 RAG 检索
    doc_results = vector_index.search(query, top_k=top_k_docs)
    doc_context = "\n\n".join([r["content"] for r in doc_results])

    # 2. GraphRAG 检索
    if use_global:
        graph_result = global_search(query, community_summaries, embed_model, llm_client)
        graph_context = graph_result["answer"]
    else:
        graph_result = local_search(query, graph, embed_model, vector_index)
        graph_context = "\n".join(graph_result["entity_descriptions"] + graph_result["relations"])

    # 3. 融合上下文
    prompt = f"""基于以下信息回答问题。

文档检索结果：
{doc_context}

知识图谱信息：
{graph_context}

问题：{query}

回答："""
    resp = llm_client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
    )
    return resp.choices[0].message.content

# 安装
pip install graphrag

# 初始化项目
graphrag init --root ./my-graphrag

# 配置 settings.yaml

# settings.yaml
encoding_model: cl100k_base
llm:
  model: gpt-4o-mini
  api_key: ${OPENAI_API_KEY}
chunks:
  size: 1200
  overlap: 100
entity_extraction:
  max_gleanings: 1
community_reports:
  max_length: 2000

# 构建索引（从文档到知识图谱）
graphrag index --root ./my-graphrag

# 查询
graphrag query --root ./my-graphrag \
  --method local \
  --query "谁是张三的合作伙伴？"

graphrag query --root ./my-graphrag \
  --method global \
  --query "这些文档的主要主题是什么？"

# 1. 使用小模型做实体抽取
extract_model = "gpt-4o-mini"  # 而非 gpt-4o

# 2. 增大分块大小，减少 LLM 调用
chunk_size = 1200  # 而非 512

# 3. 并行处理
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor(max_workers=10) as pool:
    results = list(pool.map(extract_entities_relations, chunks))

# 4. 缓存抽取结果，增量更新

# 1. 预计算常见查询的子图
# 2. 社区摘要预嵌入
# 3. 使用更快的 LLM 做社区摘要（gpt-4o-mini）
# 4. Global 查询用 map-reduce 并行