RAG_langchain.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
基于 LangChain 的 RAG 实现 – 支持 TXT / PDF / DOCX / PPTX / XLSX / HTML / MD / PY
"""

import sys
from pathlib import Path
from typing import List

try:
    # 优先使用新版本（LangChain 0.3.1+）
    from langchain_ollama import OllamaEmbeddings, OllamaLLM as Ollama
except ImportError:
    # 兼容旧版本
    from langchain_community.embeddings import OllamaEmbeddings
    from langchain_community.llms import Ollama
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_community.document_loaders import (
    TextLoader,
    PyPDFLoader,
    Docx2txtLoader,
    PythonLoader,
)

# 可选的 unstructured 加载器
try:
    from langchain_community.document_loaders import (
        UnstructuredPowerPointLoader,
        UnstructuredExcelLoader,
        UnstructuredHTMLLoader,
        UnstructuredMarkdownLoader,
    )
    UNSTRUCTURED_AVAILABLE = True
except ImportError:
    UNSTRUCTURED_AVAILABLE = False
from langchain.schema import Document
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

# ==== 配置 ====
EMBED_MODEL = "nomic-embed-text"
LLM_MODEL = "qwen3:0.6b"
CHUNK_SIZE = 500  # 字符数（LangChain 用字符而非 token）
CHUNK_OVERLAP = 50
TOP_K = 4

# ==== 文件加载器映射 ====
LOADER_MAP = {
    ".txt": TextLoader,
    ".pdf": PyPDFLoader,
    ".docx": Docx2txtLoader,
    ".py": PythonLoader,
}

# 如果 unstructured 可用，添加更多格式支持
if UNSTRUCTURED_AVAILABLE:
    LOADER_MAP.update({
        ".pptx": UnstructuredPowerPointLoader,
        ".xlsx": UnstructuredExcelLoader,
        ".html": UnstructuredHTMLLoader,
        ".htm": UnstructuredHTMLLoader,
        ".md": UnstructuredMarkdownLoader,
        ".markdown": UnstructuredMarkdownLoader,
    })
else:
    # 使用简单的文本加载器作为后备
    LOADER_MAP.update({
        ".md": TextLoader,
        ".markdown": TextLoader,
        ".html": TextLoader,
        ".htm": TextLoader,
    })


def load_documents(file_paths: List[str]) -> List[Document]:
    """加载多个文件并返回 Document 列表"""
    documents = []
    for file_path in file_paths:
        path = Path(file_path)
        ext = path.suffix.lower()
        
        if ext not in LOADER_MAP:
            print(f" 跳过不支持的文件格式: {file_path}")
            continue
        
        try:
            loader = LOADER_MAP[ext](str(path))
            docs = loader.load()
            print(f"✓ 加载文件: {path.name} ({len(docs)} 个文档)")
            documents.extend(docs)
        except Exception as e:
            print(f"✗ 加载失败 {file_path}: {e}")
    
    return documents


def build_rag_chain(file_paths: List[str]):
    """构建 RAG 链"""
    print("\n 正在加载文件...")
    documents = load_documents(file_paths)
    
    if not documents:
        raise ValueError("未成功加载任何文档！")
    
    print(f"\n 正在切分文档...")
    # LangChain 的文本切分器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", "。", "！", "？", ".", "!", "?", " ", ""],
    )
    splits = text_splitter.split_documents(documents)
    print(f"   已生成 {len(splits)} 个文档片段")
    
    print(f"\n 正在构建向量库 (使用 {EMBED_MODEL})...")
    # 初始化 Ollama 嵌入模型
    embeddings = OllamaEmbeddings(
        model=EMBED_MODEL,
        base_url="http://127.0.0.1:11434"
    )
    
    # 使用 Chroma 向量数据库（内存模式）
    vectorstore = Chroma.from_documents(
        documents=splits,
        embedding=embeddings,
        collection_name="rag_collection"
    )
    print("   向量库构建完成")
    
    print(f"\n 初始化语言模型 ({LLM_MODEL})...")
    # 初始化 Ollama LLM
    llm = Ollama(
        model=LLM_MODEL,
        base_url="http://127.0.0.1:11434",
        temperature=0.2,
        callbacks=[StreamingStdOutCallbackHandler()],
    )
    
    # 自定义 Prompt 模板
    template = """你是一位严谨的中文问答助手，只能依据参考资料作答。

参考资料：
{context}

用户问题：{question}

请根据以上参考资料回答问题。如果参考资料中没有相关信息，请明确说明。
"""
    
    prompt = PromptTemplate(
        template=template,
        input_variables=["context", "question"]
    )
    
    # 构建 RAG 链
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",  # 将所有检索到的文档塞给 LLM
        retriever=vectorstore.as_retriever(search_kwargs={"k": TOP_K}),
        chain_type_kwargs={"prompt": prompt},
        return_source_documents=True,  # 返回源文档以便查看
    )
    
    return qa_chain


def main():
    if len(sys.argv) < 2:
        print("用法: python RAG_langchain.py <文件1> <文件2> ...")
        sys.exit(0)
    
    print("=" * 60)
    print(" LangChain RAG 系统")
    print("=" * 60)
    
    if not UNSTRUCTURED_AVAILABLE:
        print(" 注意: unstructured 未安装，.pptx/.xlsx 将使用基础解析")
        print("   安装命令: pip install 'unstructured[all-docs]'")
        print("=" * 60)
    
    try:
        qa_chain = build_rag_chain(sys.argv[1:])
    except Exception as e:
        print(f"\n 初始化失败: {e}")
        sys.exit(1)
    
    print("\n" + "=" * 60)
    print(" 系统就绪！输入问题开始对话，回车或 Ctrl-C 退出")
    print("=" * 60)
    
    try:
        while True:
            q = input("\n问题> ").strip()
            if not q:
                break
            
            print("\n回答: ", end="", flush=True)
            try:
                result = qa_chain.invoke({"query": q})
                
                # 显示源文档信息
                print("\n\n 参考来源:")
                for i, doc in enumerate(result["source_documents"], 1):
                    source = doc.metadata.get("source", "未知")
                    print(f"   [{i}] {Path(source).name}")
                    
            except Exception as e:
                print(f"\n 查询失败: {e}")
    
    except (KeyboardInterrupt, EOFError):
        print("\n\n 再见！")


if __name__ == "__main__":
    main()