使用Python和FAISS搭建一个本地RAG知识库

pip install sentence-transformers faiss-cpu

documents = [
    "Docker 容器是一种轻量级的虚拟化技术，它将应用及其依赖打包在一个隔离环境中运行。",
    "Kubernetes 是一个容器编排平台，用于自动化部署、扩展和管理容器化应用。",
    "Python 的 requests 库是一个简单易用的 HTTP 客户端库，支持 GET、POST 等常用方法。",
    "Git 是一个分布式版本控制系统，支持分支管理、合并和提交历史记录。",
    "Docker Compose 允许通过 YAML 文件定义和运行多容器 Docker 应用。",
]

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

# 加载预训练 embedding 模型（首次运行会自动下载，约 420MB）
model = SentenceTransformer('all-MiniLM-L6-v2')

# 对文档进行向量化
embeddings = model.encode(documents, show_progress_bar=True)
print(f"向量维度: {embeddings.shape}")  # 输出: (5, 768)

dimension = embeddings.shape[1]  # 768

# 使用 L2 距离（欧几里得距离）构建索引
index = faiss.IndexFlatL2(dimension)

# FAISS 要求 float32 格式
embeddings_32 = embeddings.astype('float32')
index.add(embeddings_32)  # 向索引中写入向量

print(f"已存储 {index.ntotal} 条向量")

def search(query: str, top_k: int = 3) -> list:
    """根据问题检索最相关的文档"""
    query_embedding = model.encode([query])
    query_32 = query_embedding.astype('float32')

    # FAISS 检索：返回距离最近的前 top_k 个
    distances, indices = index.search(query_32, top_k)

    results = []
    for i, (dist, idx) in enumerate(zip(distances[0], indices[0])):
        results.append({
            "rank": i + 1,
            "distance": round(float(dist), 4),
            "content": documents[idx]
        })
    return results

# 测试检索
query = "如何管理容器化应用？"
results = search(query, top_k=2)

for r in results:
    print(f"[{r['rank']}] 相似度分数: {r['distance']}")
    print(f"    内容: {r['content']}\n")

import requests

def generate_answer(query: str, retrieved_docs: list) -> str:
    """基于检索结果生成回答"""
    context = "\n".join([d['content'] for d in retrieved_docs])

    prompt = f"""你是一个专业的技术助手。请基于以下参考资料回答问题。
如果资料中没有答案，请明确说"参考资料中未提及"。

参考资料：
{context}

问题：{query}
"""

    # 调用 Ollama（需先安装并启动 ollama serve）
    response = requests.post(
        "http://localhost:11434/api/generate",
        json={"model": "qwen2.5:7b", "prompt": prompt, "stream": False}
    )
    return response.json().get("response", "")

# 完整 RAG 流程
query = "Docker 和 Kubernetes 有什么关系？"
docs = search(query, top_k=3)
answer = generate_answer(query, docs)
print(answer)