RAG小实战 - 技术栈

安装依赖

pip install sentence-transformers openai
2.单文件 rag_mini.py
python 复制代码
import os

from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI

# 1. 加载文档并切分成片段
loader = DirectoryLoader("docs", glob="**/*.txt")  # 你的文档目录
raw_docs = loader.load()
splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
texts = splitter.split_documents(raw_docs)

# 2. 嵌入并构建向量索引（FAISS）
embedding_model_path = "BAAI/bge-small-zh-v1.5"
embeddings = HuggingFaceEmbeddings(
    model_name=embedding_model_path,
    model_kwargs={'device': 'cpu'},  # 如果有 GPU 可改为 'cuda'
    encode_kwargs={'normalize_embeddings': True}
)
vectorstore = FAISS.from_documents(texts, embeddings)

# 3. 构建 Retriever
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 2})

# 4. LLM
llm = ChatOpenAI(
    temperature=0,
    model="glm-4.5",
    openai_api_key=os.getenv("ZAI_API_KEY"),
    openai_api_base="https://open.bigmodel.cn/api/paas/v4/"
)

# 5. Prompt + Combine Docs Chain
prompt = ChatPromptTemplate.from_template("""
你是一个问答助手，请基于以下上下文简要地使用纯文本回答用户问题：
<context>
{context}
</context>

问题: {input}
""")

combine_docs_chain = create_stuff_documents_chain(llm, prompt)

# 6. 创建 Retrieval Chain
chain = create_retrieval_chain(
    retriever=retriever,
    combine_docs_chain=combine_docs_chain
)

# 7. 调用
out = chain.invoke({"input": "阿玛兰妲·乌尔苏拉怎么样了？"})
print("context:", out["context"])
print("answer:", out["answer"])
文本
运行效果