本地文档问答机器人：Ollama + LangChain，离线也能用 AI 读文件

# 下载并运行 qwen2.5 模型（中文能力强）
ollama pull qwen2.5:7b

# 下载嵌入模型
ollama pull nomic-embed-text

# 验证模型是否就绪
ollama list

ollama serve &

pip install langchain langchain-community langchain-ollama chromadb pymupdf python-dotenv

import os
import sys
from pathlib import Path
from langchain_community.document_loaders import PyMuPDFLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_ollama import OllamaLLM
from langchain.chains import RetrievalQA

# 配置
MODEL_NAME = "qwen2.5:7b"
EMBEDDING_MODEL = "nomic-embed-text"
VECTOR_DB_PATH = "./chroma_db"
DOCS_DIR = "./documents"


def load_documents(doc_dir):
    """加载目录下的所有文档"""
    documents = []
    doc_path = Path(doc_dir)

    if not doc_path.exists():
        print(f"⚠️  文档目录不存在：{doc_dir}")
        print("创建目录并放入 PDF/TXT 文件后重试")
        doc_path.mkdir(parents=True, exist_ok=True)
        return documents

    for file_path in doc_path.iterdir():
        if file_path.suffix == ".pdf":
            print(f"📄 加载 PDF：{file_path.name}")
            loader = PyMuPDFLoader(str(file_path))
            documents.extend(loader.load())
        elif file_path.suffix in (".txt", ".md"):
            print(f"📝 加载文本：{file_path.name}")
            loader = TextLoader(str(file_path))
            documents.append(loader.load()[0])

    print(f"✅ 共加载 {len(documents)} 个文档块")
    return documents


def build_vector_store(docs):
    """切片文档并构建向量索引"""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        separators=["\n\n", "\n", "。", "！", "？", " ", ""]
    )
    chunks = splitter.split_documents(docs)
    print(f"🔪 切分为 {len(chunks)} 个片段")

    embeddings = OllamaEmbeddings(
        model=EMBEDDING_MODEL,
        base_url="http://localhost:11434"
    )

    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=VECTOR_DB_PATH
    )
    vectorstore.persist()
    print(f"💾 向量索引已保存至 {VECTOR_DB_PATH}")
    return vectorstore


def ask_question(vectorstore, question):
    """向文档提问"""
    llm = OllamaLLM(
        model=MODEL_NAME,
        base_url="http://localhost:11434",
        temperature=0.2,
        num_ctx=4096
    )

    retriever = vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 3}
    )

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=retriever,
        return_source_documents=True
    )

    result = qa_chain.invoke({"query": question})
    print(f"\n❓ 问题：{question}")
    print(f"🤖 回答：{result['result']}")
    print(f"\n📎 引用来源：")
    for doc in result["source_documents"][:2]:
        print(f"   - {doc.metadata.get('source', '未知')}")


def main():
    print("=" * 50)
    print("📚 本地文档问答机器人")
    print("=" * 50)

    # 加载文档
    docs = load_documents(DOCS_DIR)

    if not docs:
        print("没有可处理的文档，请先将文件放入 documents/ 目录")
        sys.exit(1)

    # 构建向量库
    vectorstore = build_vector_store(docs)

    # 交互式问答
    print("\n开始提问（输入 'quit' 退出）：")
    while True:
        try:
            question = input("\n❓ ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n再见！")
            break

        if not question:
            continue
        if question.lower() in ("quit", "exit", "退出"):
            print("再见！")
            break

        ask_question(vectorstore, question)


if __name__ == "__main__":
    main()

mkdir -p documents
cat > documents/sample.md << 'EOF'
# 产品使用手册 v2.0

## 快速开始
1. 注册账号后进入控制台
2. 在「项目设置」中创建新项目
3. 获取 API Key 并开始调用接口

## 计费说明
基础版每月免费额度 1000 次调用。超出部分按 0.01 元/次计费。

## 常见问题
Q: API 响应超时怎么办？
A: 检查网络连接，确认请求体不超过 1MB。如仍超时请提交工单。
EOF

python docs_qa.py

# 在 load_documents 函数中添加过滤逻辑
EXCLUDE_PATTERNS = ["草稿", "废弃", "临时"]
for file_path in doc_path.iterdir():
    if any(pattern in file_path.name for pattern in EXCLUDE_PATTERNS):
        print(f"⏭️  跳过：{file_path.name}")
        continue
    # ... 正常加载