从平铺走到分层：AI分层RAG笔记

# 平铺式 RAG 的典型实现
def flat_rag_search(query, embedding_model, vector_db):
    # 1. 直接 query 搜最相似的 chunk
    results = vector_db.search(query, top_k=5)

    # 2. 拼接成 prompt
    context = "\n".join([r.text for r in results])

    # 3. 扔给 LLM
    response = llm.generate(f"根据以下内容回答问题：\n{context}\n\n问题：{query}")

    return response

def split_documents(docs, chunk_size=512, overlap=50):
    chunks = []
    for doc in docs:
        for i in range(0, len(doc), chunk_size - overlap):
            chunks.append(doc[i:i + chunk_size])
    return chunks

from dataclasses import dataclass
from typing import List, Optional

@dataclass
class DocumentNode:
    id: str
    text: str
    level: int  # 1=文档, 2=章节, 3=小节, 4=段落
    embedding: Optional[List[float]] = None
    parent: Optional['DocumentNode'] = None
    children: List['DocumentNode'] = None

def parse_markdown_to_hierarchy(markdown_text):
    """解析 Markdown 文档为层次结构"""
    nodes = []
    stack = []  # 记录当前层级的节点

    lines = markdown_text.split('\n')
    current_text = []

    for line in lines:
        if line.startswith('#'):
            # 保存当前段落
            if current_text:
                paragraph = '\n'.join(current_text).strip()
                if paragraph:
                    add_paragraph_node(stack, paragraph)
                current_text = []

            # 处理标题
            level = len(line) - len(line.lstrip('#'))
            title = line.lstrip('#').strip()

            # 创建新节点
            node = DocumentNode(
                id=f"node_{len(nodes)}",
                text=title,
                level=level
            )

            # 处理父子关系
            while stack and stack[-1].level >= level:
                stack.pop()

            if stack:
                node.parent = stack[-1]
                if not stack[-1].children:
                    stack[-1].children = []
                stack[-1].children.append(node)

            nodes.append(node)
            stack.append(node)
        else:
            current_text.append(line)

    # 保存最后一段
    if current_text:
        paragraph = '\n'.join(current_text).strip()
        if paragraph:
            add_paragraph_node(stack, paragraph)

    return nodes[0] if nodes else None  # 返回根节点

def add_paragraph_node(stack, text):
    """添加段落节点"""
    if not stack:
        return

    paragraph = DocumentNode(
        id=f"node_{len(stack) * 1000 + len(stack[-1].children) if stack[-1].children else 0}",
        text=text,
        level=4,  # 段落级
        parent=stack[-1]
    )

    if not stack[-1].children:
        stack[-1].children = []
    stack[-1].children.append(paragraph)

import numpy as np
from sentence_transformers import SentenceTransformer

embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

def embed_hierarchy(root_node):
    """为层次化文档生成 embedding"""
    if not root_node:
        return

    # BFS 遍历
    queue = [root_node]

    while queue:
        node = queue.pop(0)
        node.embedding = embedding_model.encode(node.text)

        if node.children:
            queue.extend(node.children)

def store_hierarchy(root_node, vector_db):
    """存储到向量数据库，并记录层次关系"""
    if not root_node:
        return

    queue = [root_node]

    while queue:
        node = queue.pop(0)

        # 存储当前节点
        vector_db.add(
            id=node.id,
            embedding=node.embedding,
            metadata={
                'text': node.text,
                'level': node.level,
                'parent_id': node.parent.id if node.parent else None,
                'children_ids': [c.id for c in node.children] if node.children else []
            }
        )

        if node.children:
            queue.extend(node.children)

def hierarchical_search(query, vector_db, top_k=5):
    """分层检索"""
    query_embedding = embedding_model.encode(query)

    # 1. 判断查询意图（简单启发式）
    intent = classify_query_intent(query)

    # 2. 根据意图选择搜索层级
    if intent == "overview":
        search_levels = [2]  # 章节级
    elif intent == "detail":
        search_levels = [4]  # 段落级
    else:  # "mixed" 或不确定
        search_levels = [2, 3, 4]  # 多层级搜索

    # 3. 在目标层级搜索
    results = vector_db.search(
        query_embedding,
        top_k=top_k * 2,  # 多取一些，后续去重
        filter={'level': {'$in': search_levels}}
    )

    # 4. 去重和扩展（如果需要父节点上下文）
    results = deduplicate_and_expand(results, vector_db)

    # 5. 按相关性排序
    results = sorted(results, key=lambda x: x['score'], reverse=True)

    return results[:top_k]

def classify_query_intent(query):
    """简单查询意图分类"""
    # 这是一个简化版本，实际可以用 LLM 分类
    overview_keywords = ['架构', '模块', '整体', '概览', '结构', '组成']
    detail_keywords = ['配置', '参数', '如何', '怎么', '实现', '代码']

    if any(kw in query for kw in overview_keywords):
        return "overview"
    elif any(kw in query for kw in detail_keywords):
        return "detail"
    else:
        return "mixed"

def deduplicate_and_expand(results, vector_db):
    """去重并扩展必要的父节点上下文"""
    seen = set()
    final_results = []

    for result in results:
        parent_id = result['metadata'].get('parent_id')

        # 去重：如果父节点已经在结果中，跳过子节点
        if parent_id in seen:
            continue

        final_results.append(result)
        seen.add(result['id'])

        # 如果是段落级且父节点是章节，可能需要父节点提供上下文
        if result['metadata']['level'] == 4 and parent_id:
            parent = vector_db.get(parent_id)
            if parent and parent['metadata']['level'] == 2:
                # 在合适位置插入父节点
                insert_idx = final_results.index(result)
                final_results.insert(insert_idx, parent)
                seen.add(parent_id)

    return final_results

def hierarchical_rag(query, vector_db, llm):
    """完整的分层 RAG 流程"""
    # 1. 分层检索
    results = hierarchical_search(query, vector_db, top_k=5)

    # 2. 构建上下文
    context_parts = []
    for r in results:
        level_name = {2: '章节', 3: '小节', 4: '段落'}.get(r['metadata']['level'], '内容')
        context_parts.append(f"【{level_name}】{r['metadata']['text']}")

    context = '\n\n'.join(context_parts)

    # 3. 生成回答
    prompt = f"""请根据以下文档内容回答问题。注意：
1. 直接基于提供的内容回答
2. 如果信息不足，明确说明
3. 保持回答简洁

文档内容：
{context}

问题：{query}

回答："""

    response = llm.generate(prompt)

    return response, results

# 标题一
内容一

## 标题二

    ### 小标题（前面有空格，解析失败）
内容二

from markdown_it import MarkdownIt

md = MarkdownIt()
tokens = md.parse(markdown_text)

# 根据 heading tokens 构建层次结构
headings = [t for t in tokens if t.type == 'heading_open']

def embed_with_level_context(node):
    """为不同层级的节点添加上下文前缀"""
    level_prefixes = {
        1: "文档标题：",
        2: "章节内容：",
        3: "小节内容：",
        4: "具体细节："
    }

    prefixed_text = level_prefixes.get(node.level, "") + node.text
    node.embedding = embedding_model.encode(prefixed_text)

def classify_with_llm(query):
    """用 LLM 分类查询意图"""
    prompt = f"""判断以下查询需要什么粒度的信息，只返回类别名称：

查询：{query}

类别说明：
- overview：需要整体架构、模块组成等宏观信息
- detail：需要具体配置、代码实现、参数说明等微观信息
- mixed：需要混合信息或不确定

类别："""

    response = llm.generate(prompt, temperature=0)
    return response.strip().lower()

def expand_context(node, vector_db, max_depth=2):
    """扩展父节点上下文，限制深度"""
    expanded = [node]
    current = node.parent
    depth = 0

    while current and depth < max_depth:
        # 判断是否真的需要父节点（比如子节点信息不完整）
        if needs_parent_context(node, current):
            expanded.insert(0, current)
        current = current.parent
        depth += 1

    return expanded

def needs_parent_context(child, parent):
    """判断是否需要父节点提供上下文"""
    # 简单启发式：如果子节点太短（< 100 字），可能需要父节点
    return len(child.text) < 100

from functools import lru_cache
import concurrent.futures

@lru_cache(maxsize=1000)
def get_node_embedding(node_id):
    """缓存 embedding"""
    return vector_db.get_embedding(node_id)

def parallel_search_by_level(query_embedding, levels, vector_db):
    """并行查询不同层级"""
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = []
        for level in levels:
            future = executor.submit(
                vector_db.search,
                query_embedding,
                top_k=5,
                filter={'level': level}
            )
            futures.append(future)

        results = []
        for future in concurrent.futures.as_completed(futures):
            results.extend(future.result())

    return results