Ollama + LangChain 搭建本地 RAG 知识库：上传 PDF 让 AI 基于你的私人资料回答

ollama pull qwen2.5:7b
ollama pull nomic-embed-text

mkdir ~/rag-demo && cd ~/rag-demo
python3 -m venv venv && source venv/bin/activate
pip install langchain langchain-community langchain-ollama chromadb pypdf

#!/usr/bin/env python3
import os, sys, shutil, subprocess
from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_community.vectorstores import Chroma
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

MODEL = "qwen2.5:7b"
EMBED_MODEL = "nomic-embed-text"
DB_PATH = "./chroma_db"
DOCS_DIR = "./docs"

def ensure_model(name):
    r = subprocess.run(["ollama", "list"], capture_output=True, text=True)
    if name not in r.stdout:
        print(f"下载模型 {name}...")
        subprocess.run(["ollama", "pull", name], check=True)

def build_db():
    print(f"扫描文档目录: {DOCS_DIR}")
    loader = DirectoryLoader(DOCS_DIR, glob="**/*")
    docs = loader.load()
    print(f"加载了 {len(docs)} 个文档")

    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500, chunk_overlap=50,
        separators=["nn", "n", "。", "！", "？", ".", "!", "?", " ", ""]
    )
    chunks = splitter.split_documents(docs)
    print(f"切分成 {len(chunks)} 个文本块")

    if os.path.exists(DB_PATH):
        shutil.rmtree(DB_PATH)

    Chroma.from_documents(
        documents=chunks,
        embedding=OllamaEmbeddings(model=EMBED_MODEL),
        persist_directory=DB_PATH,
    )
    print(f"向量库已保存")

def load_db():
    return Chroma(
        persist_directory=DB_PATH,
        embedding_function=OllamaEmbeddings(model=EMBED_MODEL),
    )

def main():
    ensure_model(MODEL)
    ensure_model(EMBED_MODEL)

    if not os.path.exists(DB_PATH) or "--rebuild" in sys.argv:
        build_db()

    vectorstore = load_db()
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    llm = ChatOllama(model=MODEL, temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "基于以下参考资料回答问题。如果资料中没有相关信息，直接说"文档中没有相关内容"。nn参考资料：n{context}"),
        ("human", "{input}"),
    ])

    chain = create_retrieval_chain(
        retriever,
        create_stuff_documents_chain(llm, prompt)
    )

    print("输入问题开始查询（quit 退出 / rebuild 重建索引）")
    while True:
        q = input("n你的问题：").strip()
        if q.lower() in ('quit', 'exit', '退出'):
            break
        if q.lower() == 'rebuild':
            build_db()
            continue

        print("⏳ 回答中...")
        result = chain.invoke({"input": q})
        print(f"n💬 {result['answer']}")

        print("n📎 引用来源：")
        for i, doc in enumerate(result["context"], 1):
            print(f"  [{i}] {doc.metadata.get('source', 'unknown')}")
            print(f"      {doc.page_content[:200]}...")

if __name__ == "__main__":
    main()

mkdir docs
cp /path/to/your/document.pdf docs/

cd ~/rag-demo
source venv/bin/activate
python3 rag.py

pip install gradio

import gradio as gr

def answer(q):
    result = chain.invoke({"input": q})
    return result["answer"]

gr.Interface(fn=answer, inputs=gr.Textbox(label="问题"), outputs=gr.Textbox(label="回答")).launch()