# 导入核心依赖（修复所有模块路径）
import os
# 环境变量管理（建议使用 python-dotenv）
from dotenv import load_dotenv
load_dotenv()  # 加载 .env 文件中的环境变量

# --------------------------
# 1. 文本分割器：从 langchain_text_splitters 导入（v0.1.0+ 新版本路径）
# --------------------------
from langchain_text_splitters import CharacterTextSplitter

# --------------------------
# 2. 文档加载器：仍在 langchain_community 中
# --------------------------
from langchain_community.document_loaders import TextLoader

# --------------------------
# 3. 嵌入模型：根据版本选择路径
# --------------------------
# 方式1：OpenAI 嵌入（需安装 langchain-openai）
from langchain_openai import OpenAIEmbeddings
# 方式2：本地 HuggingFace 嵌入（需安装 sentence-transformers）
# from langchain_community.embeddings import HuggingFaceEmbeddings

# --------------------------
# 4. 向量存储：仍在 langchain_community 中
# --------------------------
from langchain_community.vectorstores import Chroma

# --------------------------
# 5. 大模型：根据版本选择路径
# --------------------------
# 方式1：OpenAI Chat 模型（需安装 langchain-openai）
from langchain_openai import ChatOpenAI
# 方式2：本地 Ollama 模型（需安装 langchain-community 和 ollama）
# from langchain_community.llms import Ollama

# --------------------------
# 6. RetrievalQA：仍在 langchain_community 中
# --------------------------
from langchain_community.chains import RetrievalQA

# 配置环境变量（也可通过 .env 文件设置）
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your-openai-api-key")  # 替换为你的 API Key

def load_and_split_documents(file_path="docs/knowledge.txt"):
    """加载文档并分割为小片段"""
    # 确保文档目录存在
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    # 若文档不存在，创建临时测试文档
    if not os.path.exists(file_path):
        with open(file_path, "w", encoding="utf-8") as f:
            f.write("LangChain 是一个用于构建 LLM 应用的框架，支持检索增强生成（RAG）等功能。\n")
            f.write("RetrievalQA 是 LangChain 中用于检索问答的核心链，在 v0.1.0 后迁移到了 langchain_community 模块。\n")
            f.write("CharacterTextSplitter 在 v0.1.0 后迁移到了 langchain_text_splitters 模块。\n")
    # 1. 加载文档
    loader = TextLoader(file_path, encoding="utf-8")
    documents = loader.load()
    
    # 2. 分割文档（修复后的文本分割器）
    text_splitter = CharacterTextSplitter(
        chunk_size=1000,  # 每个片段的字符数
        chunk_overlap=200,  # 片段间重叠字符数
        separator="\n"
    )
    split_docs = text_splitter.split_documents(documents)
    return split_docs

def create_vector_store(docs):
    """创建向量存储（使用 Chroma 本地向量库）"""
    # 初始化嵌入模型
    # 方式1：OpenAI 嵌入
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    # 方式2：本地 HuggingFace 嵌入
    # embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
    
    # 创建向量库并存储文档
    vector_store = Chroma.from_documents(
        documents=docs,
        embedding=embeddings,
        persist_directory="./chroma_db"  # 向量库持久化路径
    )
    # Chroma v0.4+ 不再需要手动 persist()
    return vector_store

def create_retrieval_qa_chain(vector_store):
    """创建 RetrievalQA 检索问答链"""
    # 1. 初始化大模型
    # 方式1：OpenAI GPT-3.5/4
    llm = ChatOpenAI(
        model_name="gpt-3.5-turbo",
        temperature=0.1  # 温度越低，回答越稳定
    )
    # 方式2：本地 Ollama 模型（如 Llama3）
    # llm = Ollama(model="llama3:8b", temperature=0.1)
    
    # 2. 创建检索器（从向量库中检索相关文档）
    retriever = vector_store.as_retriever(
        search_kwargs={"k": 3}  # 检索前 3 个最相关的文档
    )
    
    # 3. 创建 RetrievalQA 链
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",  # 链类型：stuff/map_reduce/refine 等
        retriever=retriever,
        return_source_documents=True  # 是否返回检索到的源文档
    )
    return qa_chain

if __name__ == "__main__":
    # 1. 加载并分割文档
    docs = load_and_split_documents()
    # 2. 创建向量存储
    vector_store = create_vector_store(docs)
    # 3. 创建 RetrievalQA 链
    qa_chain = create_retrieval_qa_chain(vector_store)
    
    # 4. 测试问答
    query = "CharacterTextSplitter 在 LangChain 新版本中被迁移到了哪里？"
    result = qa_chain.invoke({"query": query})  # v0.1.0+ 推荐使用 invoke() 替代 __call__
    
    # 输出结果
    print("===== 问题 =====")
    print(query)
    print("\n===== 回答 =====")
    print(result["result"])
    print("\n===== 参考文档 =====")
    for i, doc in enumerate(result["source_documents"]):
        print(f"\n文档 {i+1}：")
        print(f"内容：{doc.page_content}")
        print(f"来源：{doc.metadata['source']}")