RAG 管道搭建指南：向量数据库 + LLM 集成实战

没有 RAG：
  用户："我们公司的退款政策是什么？"
  LLM："我没有你公司的具体政策信息……"

使用 RAG：
  用户："我们公司的退款政策是什么？"
  系统：[检索 refund-policy.pdf 片段] → 注入到 prompt
  LLM："根据你的政策文档，30天内可以退款……"

from openai import OpenAI
client = OpenAI()

response = client.embeddings.create(
    model="text-embedding-3-small",
    input="How do I reset my password?"
)

vector = response.data[0].embedding
print(f"Dimensions: {len(vector)}")  # 1536
print(f"First 5 values: {vector[:5]}")
# [0.0123, -0.0456, 0.0789, -0.0234, 0.0567]

Layer 3:  A -------- D                    (稀疏，远距离跳跃)
Layer 2:  A --- B -- D --- F              (中等密度)
Layer 1:  A - B - C - D - E - F - G      (密集，短距离跳跃)
Layer 0:  A B C D E F G H I J K L M      (所有节点)

文档加载 → 分块 → Embedding → 存储 → 检索 → 生成

from pathlib import Path

def load_documents(directory: str) -> list[dict]:
    """从目录加载文本文档。"""
    documents = []
    for path in Path(directory).glob("**/*.md"):
        text = path.read_text(encoding="utf-8")
        documents.append({
            "text": text,
            "source": str(path),
            "filename": path.name,
        })
    return documents

docs = load_documents("./knowledge_base")
print(f"Loaded {len(docs)} documents")

def fixed_size_chunk(text: str, chunk_size: int = 500, overlap: int = 100) -> list[str]:
    """将文本切分为固定大小的重叠块。"""
    words = text.split()
    chunks = []
    start = 0
    while start < len(words):
        end = start + chunk_size
        chunk = " ".join(words[start:end])
        chunks.append(chunk)
        start += chunk_size - overlap  # 带重叠步进
    return chunks

def recursive_split(text: str, max_size: int = 500, separators: list[str] = None) -> list[str]:
    """在自然边界处递归分割文本。"""
    if separators is None:
        separators = ["\n\n", "\n", ". ", " "]

    if len(text.split()) <= max_size:
        return [text]

    for sep in separators:
        parts = text.split(sep)
        if len(parts) > 1:
            chunks = []
            current = ""
            for part in parts:
                candidate = current + sep + part if current else part
                if len(candidate.split()) > max_size and current:
                    chunks.append(current.strip())
                    current = part
                else:
                    current = candidate
            if current.strip():
                chunks.append(current.strip())
            # 递归分割仍然过大的块
            result = []
            for chunk in chunks:
                result.extend(recursive_split(chunk, max_size, separators))
            return result

    # 兜底：按词强制分割
    return fixed_size_chunk(text, max_size)

import numpy as np
from openai import OpenAI

client = OpenAI()

def get_embeddings(texts: list[str]) -> list[list[float]]:
    """批量获取文本的 Embedding。"""
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    return [item.embedding for item in response.data]

def semantic_chunk(text: str, threshold: float = 0.8) -> list[str]:
    """基于句子间语义相似度进行分块。"""
    import re
    sentences = re.split(r'(?<=[.!?])\s+', text)
    if len(sentences) <= 1:
        return [text]

    embeddings = get_embeddings(sentences)

    chunks = []
    current_chunk = [sentences[0]]

    for i in range(1, len(sentences)):
        # 当前句和前一句的余弦相似度
        sim = np.dot(embeddings[i], embeddings[i-1]) / (
            np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[i-1])
        )
        if sim < threshold:
            # 检测到话题转换，开始新块
            chunks.append(" ".join(current_chunk))
            current_chunk = [sentences[i]]
        else:
            current_chunk.append(sentences[i])

    if current_chunk:
        chunks.append(" ".join(current_chunk))

    return chunks

import chromadb
from openai import OpenAI

client = OpenAI()
chroma = chromadb.PersistentClient(path="./chroma_db")

# 创建或获取集合
collection = chroma.get_or_create_collection(
    name="knowledge_base",
    metadata={"hnsw:space": "cosine"}  # 使用余弦相似度
)

def embed_and_store(chunks: list[dict]):
    """将分块进行 Embedding 并存入 Chroma。"""
    texts = [c["text"] for c in chunks]

    # 批量 Embedding（OpenAI 每次调用最多支持 2048 个输入）
    batch_size = 2000
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        response = client.embeddings.create(
            model="text-embedding-3-small",
            input=batch
        )
        all_embeddings.extend([item.embedding for item in response.data])

    # 存入 Chroma
    collection.add(
        ids=[f"chunk_{i}" for i in range(len(chunks))],
        embeddings=all_embeddings,
        documents=texts,
        metadatas=[{"source": c["source"]} for c in chunks]
    )
    print(f"Stored {len(chunks)} chunks in Chroma")

# 准备带元数据的分块
chunks = []
for doc in docs:
    doc_chunks = recursive_split(doc["text"], max_size=400)
    for chunk_text in doc_chunks:
        chunks.append({"text": chunk_text, "source": doc["source"]})

embed_and_store(chunks)

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from openai import OpenAI

client = OpenAI()
qdrant = QdrantClient(url="http://localhost:6333")

# 创建集合
qdrant.create_collection(
    collection_name="knowledge_base",
    vectors_config=VectorParams(
        size=1536,  # text-embedding-3-small 的维度
        distance=Distance.COSINE,
    ),
)

def embed_and_store_qdrant(chunks: list[dict]):
    """将分块进行 Embedding 并存入 Qdrant。"""
    texts = [c["text"] for c in chunks]

    # 获取 Embedding
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    embeddings = [item.embedding for item in response.data]

    # 创建数据点
    points = [
        PointStruct(
            id=i,
            vector=embeddings[i],
            payload={
                "text": texts[i],
                "source": chunks[i]["source"],
            }
        )
        for i in range(len(chunks))
    ]

    # 批量写入
    batch_size = 100
    for i in range(0, len(points), batch_size):
        qdrant.upsert(
            collection_name="knowledge_base",
            points=points[i:i + batch_size]
        )

    print(f"Stored {len(chunks)} chunks in Qdrant")

def retrieve(query: str, top_k: int = 5) -> list[dict]:
    """从 Chroma 中检索与查询相关的分块。"""
    # 对查询进行 Embedding
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=query
    )
    query_embedding = response.data[0].embedding

    # 搜索 Chroma
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        include=["documents", "metadatas", "distances"]
    )

    # 格式化结果
    retrieved = []
    for i in range(len(results["documents"][0])):
        retrieved.append({
            "text": results["documents"][0][i],
            "source": results["metadatas"][0][i]["source"],
            "score": 1 - results["distances"][0][i],  # 将距离转换为相似度
        })

    return retrieved

def generate_answer(query: str, context_chunks: list[dict]) -> str:
    """使用检索到的上下文生成回答。"""
    # 构建上下文字符串
    context = "\n\n---\n\n".join([
        f"[Source: {c['source']}]\n{c['text']}"
        for c in context_chunks
    ])

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": (
                    "You are a helpful assistant. Answer the user's question "
                    "based on the provided context. If the context does not "
                    "contain relevant information, say so. Always cite your "
                    "sources."
                )
            },
            {
                "role": "user",
                "content": f"Context:\n{context}\n\nQuestion: {query}"
            }
        ],
        temperature=0.1,  # 低温度以获得更准确的回答
    )

    return response.choices[0].message.content

def rag_query(question: str) -> str:
    """完整的 RAG 管道：检索上下文并生成回答。"""
    # 第一步：检索相关分块
    chunks = retrieve(question, top_k=5)

    # 第二步：过滤低相关性结果
    relevant_chunks = [c for c in chunks if c["score"] > 0.7]

    if not relevant_chunks:
        return "I could not find relevant information to answer your question."

    # 第三步：生成回答
    answer = generate_answer(question, relevant_chunks)

    # 第四步：附加来源
    sources = set(c["source"] for c in relevant_chunks)
    answer += f"\n\nSources: {', '.join(sources)}"

    return answer

# 使用示例
result = rag_query("What is our company's vacation policy?")
print(result)

def rerank(query: str, chunks: list[dict], top_k: int = 3) -> list[dict]:
    """使用 LLM 作为评判者对分块重排序。"""
    scored = []
    for chunk in chunks:
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": (
                    f"Rate how relevant this text is to the question "
                    f"on a scale of 0-10.\n\n"
                    f"Question: {query}\n\n"
                    f"Text: {chunk['text']}\n\n"
                    f"Score (just the number):"
                )
            }],
            max_tokens=3,
            temperature=0,
        )
        try:
            score = float(response.choices[0].message.content.strip())
            scored.append({**chunk, "rerank_score": score})
        except ValueError:
            scored.append({**chunk, "rerank_score": 0})

    scored.sort(key=lambda x: x["rerank_score"], reverse=True)
    return scored[:top_k]

def expand_query(query: str) -> list[str]:
    """生成查询变体以提高召回率。"""
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{
            "role": "user",
            "content": (
                f"Generate 3 different ways to ask this question. "
                f"Return only the questions, one per line.\n\n"
                f"Question: {query}"
            )
        }],
        temperature=0.7,
    )
    variations = response.choices[0].message.content.strip().split("\n")
    return [query] + [v.strip() for v in variations if v.strip()]