RAG 实战:从零搭建企业私有知识库

pip install langchain langchain-community langchain-openai \
chromadb pypdf tiktoken

# 1. 加载文档
loader = TextLoader("docs/product_intro.txt")
documents = loader.load()

# 2. 切分文档
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500, chunk_overlap=50
)
chunks = text_splitter.split_documents(documents)

# 3. 向量化并存到向量数据库
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_documents(chunks, embeddings)

# 4. 创建 LLM
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# 5. 创建 RAG 链
qa = RetrievalQA.from_chain_type(
llm=llm,
retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
return_source_documents=True,
)

# 6. 提问
result = qa.invoke({"query": "我们的产品有什么特点?"})
print("答案:", result["result"])
print("\n来源:")
for i, doc in enumerate(result["source_documents"]):
print(f" [{i+1}] {doc.page_content[:100]}…")

vectorstore = Chroma(persist_directory="./chroma", embedding_function=OpenAIEmbeddings())

def file_hash(path: Path) -> str:
"""计算文件 hash,作为是否更新过的依据"""
return hashlib.md5(path.read_bytes()).hexdigest()

existing_meta = {m["hash"]: m["id"] for m in vectorstore.get()["metadatas"] if m and "hash" in m}

for path in Path("docs").rglob("*"):
if path.suffix not in [".pdf", ".txt", ".md", ".docx"]:
continue
h = file_hash(path)
if h in existing_meta:
print(f"跳过(未改动): {path}")
continue
# 读取、切除、入库
docs = load_document(path)
vectorstore.add_documents(docs, metadatas=[{"source": str(path), "hash": h}])
print(f"已入库: {path}")

loader = PyPDFLoader("docs/manual.pdf")
pages = loader.load_and_split()

loader = NotionDBLoader(
integration_token="secret_xxx",
database_id="你的 Notion DB ID",
)
docs = loader.load()

loader = FeishuLoader(
app_id="your_app_id",
app_secret="your_app_secret",
dir_token="文件夹 token",
)
docs = loader.load()

def load_directory(directory: str):
"""批量加载一个目录下所有支持的文档"""
documents = []
for file_path in Path(directory).rglob("*"):
if file_path.suffix == ".pdf":
loader = PyPDFLoader(str(file_path))
elif file_path.suffix in [".txt", ".md"]:
loader = TextLoader(str(file_path))
elif file_path.suffix == ".docx":
from langchain_community.document_loaders import UnstructuredWordDocumentLoader
loader = UnstructuredWordDocumentLoader(str(file_path))
else:
continue

documents.extend(loader.load())
return documents

documents = load_directory("docs/")
print(f"加载了 {len(documents)} 个文档块")

# 按 Markdown 标题切分
splitter = MarkdownTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = splitter.split_text(markdown_text)

splitter = SemanticChunker(
embeddings,
breakpoint_threshold_type="percentile",
)
chunks = splitter.split_documents(documents)

# 关键词检索
bm25_retriever = BM25Retriever.from_documents(chunks)

# 向量检索
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

# 混合(70% 向量, 30% 关键词)
ensemble_retriever = EnsembleRetriever(
retrievers=[vector_retriever, bm25_retriever],
weights=[0.7, 0.3],
)
qa = RetrievalQA.from_chain_type(
llm=llm,
retriever=ensemble_retriever,
)

# 召回 10 个文档
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 10})

# 重排到前 3
compressor = CohereRerank(top_n=3)
compression_retriever = ContextualCompressionRetriever(
base_compressor=compressor,
base_retriever=base_retriever,
)
qa = RetrievalQA.from_chain_type(
llm=llm,
retriever=compression_retriever,
)

rewrite_prompt = PromptTemplate.from_template(
"将以下用户问题改写为更适合检索关键字词汇库的形式。仅输出改写后的查询,不要其他内容。\n\n原始问题: {question}\n改写后:"
)

def rewrite_query(question: str, llm) -> str:
return llm.invoke(rewrite_prompt.format(question=question)).content

# 使用
original_q = "那个能看代码的东西怎么用?"
rewritten_q = rewrite_query(original_q, llm)
print(f"原始: {original_q}")
print(f"改写后: {rewritten_q}")
# 可能输出: "开发者使用 AI 代码代理的方法是什么"

def evaluate(test_cases):
correct = 0
for case in test_cases:
result = qa.invoke({"query": case["q"]})
sources = [doc.metadata.get("source") for doc in result["source_documents"]]
if case["expected_source"] in sources:
correct += 1
print(f"准确率: {correct}/{len(test_cases)} = {correct/len(test_cases)*100:.0f}%")

evaluate(test_cases)

load_dotenv()

def load_documents(directory: str):
docs = []
for path in Path(directory).rglob("*"):
if path.suffix == ".pdf":
docs.extend(PyPDFLoader(str(path)).load())
elif path.suffix in [".txt", ".md"]:
docs.extend(TextLoader(str(path)).load())
elif path.suffix == ".docx":
docs.extend(UnstructuredWordDocumentLoader(str(path)).load())
return docs

def main():
documents = load_documents("docs")
print(f"加载 {len(documents)} 个文档")

splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = splitter.split_documents(documents)
print(f"切分成 {len(chunks)} 个 chunk")

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_documents(
chunks, embeddings, persist_directory="./chroma"
)
vectorstore.persist()
print(f"向量库已保存到 ./chroma")

if __name__ == "__main__":
main()

load_dotenv()

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma(persist_directory="./chroma", embedding_function=embeddings)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# 自定义 prompt 模板
prompt_template = """基于以下参考资料回答问题。
参考资料:{context}
问题:{question}
答案要简洁、有依据。如果参考资料里没有答案,说"我不知道"。"""

PROMPT = PromptTemplate(
template=prompt_template,
input_variables=["context", "question"],
)

qa = RetrievalQA.from_chain_type(
llm=llm,
retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
chain_type_kwargs={"prompt": PROMPT},
return_source_documents=True,
)

def ask(question: str):
result = qa.invoke({"query": question})
print(f"\n问: {question}")
print(f"答: {result['result']}")
print(f"来源: {[doc.metadata.get('source', '?') for doc in result['source_documents']]}")

if __name__ == "__main__":
questions = [
"年假可以休几天?",
"加班费是怎么计算的?",
"工资什么时候发?",
"我刚入职,需要办理什么手续?",
]
for q in questions:
ask(q)