用 Python + Ollama 搭一个本地知识库问答系统，离线也能问文档

# 问答用的生成模型（选一个适合你显存的）
ollama pull qwen2.5:7b

# 文本嵌入模型（用于文档向量化）
ollama pull nomic-embed-text

pip install langchain langchain-community chromadb gradio PyPDF2 python-docx

"""
本地知识库构建脚本
将文档目录中的文件切片、向量化，存入 ChromaDB
"""
import os
import glob
from langchain_community.document_loaders import (
    PyPDFLoader,
    TextLoader,
    UnstructuredMarkdownLoader,
    Docx2txtLoader
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma

# 支持的文档类型和对应的加载器
LOADER_MAP = {
    ".pdf": PyPDFLoader,
    ".txt": TextLoader,
    ".md": UnstructuredMarkdownLoader,
    ".docx": Docx2txtLoader,
}

def load_documents(doc_dir: str):
    """加载指定目录下的所有文档"""
    documents = []
    
    for filepath in glob.glob(os.path.join(doc_dir, "**/*"), recursive=True):
        ext = os.path.splitext(filepath)[1].lower()
        
        if ext not in LOADER_MAP:
            print(f"⏭️  跳过不支持的文件: {filepath}")
            continue
        
        print(f"📄 正在加载: {filepath}")
        loader = LOADER_MAP[ext](filepath)
        documents.extend(loader.load())
    
    return documents


def chunk_documents(documents, chunk_size=500, chunk_overlap=50):
    """将文档切分为小块，便于向量化"""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", "。", ".", "!"]
    )
    chunks = splitter.split_documents(documents)
    print(f"✅ 文档已切分为 {len(chunks)} 个片段")
    return chunks


def build_vector_store(chunks, persist_dir="./chroma_kb"):
    """将文本块向量化并存入 ChromaDB"""
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    
    # 如果已有持久化数据，直接加载
    if os.path.exists(persist_dir):
        print(f"📂 加载已有向量库: {persist_dir}")
        db = Chroma(persist_directory=persist_dir, embedding_function=embeddings)
        # 追加新文档
        db.add_documents(chunks)
    else:
        print(f"🆕 创建新的向量库: {persist_dir}")
        db = Chroma.from_documents(
            documents=chunks,
            embedding=embeddings,
            persist_directory=persist_dir
        )
    
    db.persist()
    print("💾 向量库已保存")
    return db


def main():
    doc_dir = "./documents"  # 把文档放在这里
    
    if not os.path.exists(doc_dir):
        print(f"❌ 找不到文档目录: {doc_dir}")
        print("请在当前目录下创建 documents 文件夹，放入你的文档")
        return
    
    docs = load_documents(doc_dir)
    if not docs:
        print("❌ 没有找到任何可处理的文档")
        return
    
    chunks = chunk_documents(docs)
    build_vector_store(chunks)
    print("\n🎉 知识库构建完成！可以开始提问了。")


if __name__ == "__main__":
    main()

"""
知识库问答脚本
从向量库中检索相关内容，交给大模型生成回答
"""
import os
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.llms import Ollama
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

def load_vectorstore(persist_dir="./chroma_kb"):
    """加载持久化的向量库"""
    embeddings = OllamaEmbeddings(model="nomic-embed-text")
    db = Chroma(persist_directory=persist_dir, embedding_function=embeddings)
    return db


def create_qa_chain(db):
    """创建检索增强问答链"""
    retriever = db.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 4}  # 返回最相关的4个片段
    )
    
    llm = Ollama(model="qwen2.5:7b", temperature=0.2, max_tokens=1024)
    
    # 自定义提示词模板
    prompt_template = """基于以下已知信息，简洁而专业地回答用户的问题。
如果已知信息中不包含答案，请如实回答"抱歉，知识库中没有相关内容"。

已知信息：
{context}

用户问题：{question}

回答："""
    
    prompt = PromptTemplate(
        template=prompt_template,
        input_variables=["context", "question"]
    )
    
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        chain_type_kwargs={"prompt": prompt}
    )
    
    return qa_chain


def ask_question(qa_chain, question):
    """提问并获取回答"""
    result = qa_chain.invoke({"query": question})
    return result["result"]


def main():
    persist_dir = "./chroma_kb"
    
    if not os.path.exists(persist_dir):
        print("❌ 找不到向量库，请先运行 build_knowledge_base.py 构建知识库")
        return
    
    print("📚 正在加载知识库...")
    db = load_vectorstore(persist_dir)
    qa_chain = create_qa_chain(db)
    
    print("\n🤖 知识库问答系统已就绪")
    print("输入问题（输入 quit 退出）:\n")
    
    while True:
        question = input("❓ 你的问题: ").strip()
        if question.lower() in ("quit", "exit", "q"):
            print("👋 再见！")
            break
        if not question:
            continue
        
        print("\n🔍 正在检索和生成回答...")
        answer = ask_question(qa_chain, question)
        print(f"\n💡 {answer}\n")


if __name__ == "__main__":
    main()

"""
Web 界面 — 用 Gradio 快速搭建可视化问答页面
"""
import os
from chat_knowledge_base import load_vectorstore, create_qa_chain

persist_dir = "./chroma_kb"
db = load_vectorstore(persist_dir)
qa_chain = create_qa_chain(db)


def answer(question, history):
    """Gradio 回调函数"""
    if not question.strip():
        return history
    response = ask_question(qa_chain, question)
    history.append((question, response))
    return history


if __name__ == "__main__":
    import gradio as gr
    
    gr.Interface(
        fn=answer,
        inputs=[
            gr.Textbox(placeholder="请输入你的问题...", lines=2),
            gr.Chatbot(label="对话记录", height=400)
        ],
        title="📚 本地知识库问答系统",
        description="基于 Ollama + ChromaDB 的离线 RAG 问答系统。请先用 build_knowledge_base.py 构建知识库。",
        allow_flagging="never"
    ).launch(server_name="0.0.0.0", server_port=7860)

# 1. 把文档放进 documents 目录
mkdir -p documents
cp ~/some-report.pdf documents/
cp ~/notes.md documents/

# 2. 构建知识库
python build_knowledge_base.py

# 3. 启动 Web 界面
python web_ui.py
# 访问 http://localhost:7860