{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "711d55a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "根据上下文，员工被辞退时“不支付任何经济补偿金”。\n"
     ]
    }
   ],
   "source": [
    "from langchain_community.document_loaders import Docx2txtLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_community.embeddings import DashScopeEmbeddings\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_community.chat_models import ChatTongyi\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "qwen_api_key = os.getenv(\"QWEN_API_KEY\")\n",
    "\n",
    "# ========== 第一步：加载文档 ==========\n",
    "loader = Docx2txtLoader(\"./人事管理流程.docx\")\n",
    "pages = loader.load()\n",
    "\n",
    "# ========== 第二步：清洗数据（可选，根据文档质量决定）==========\n",
    "# clean_pages = [clean_pdf_text(page) for page in pages]\n",
    "\n",
    "# ========== 第三步：分块 ==========\n",
    "splitter = RecursiveCharacterTextSplitter(\n",
    "    # 分隔符优先级：段落 → 换行 → 句号 → 空格 → 硬切\n",
    "    separators=[\"\\n\\n\", \"\\n\", \"。\", \"！\", \"？\", \" \", \"\"],\n",
    "    # 每个块最大 50 字符\n",
    "    chunk_size=50,\n",
    "    # 相邻块重叠 10 字符（chunk_size 的 20%）\n",
    "    chunk_overlap=10,\n",
    "    # 长度计算函数\n",
    "    length_function=len\n",
    ")\n",
    "docs = splitter.split_documents(pages)\n",
    "\n",
    "# ========== 第四步：向量化 + 存入向量库 ==========\n",
    "embedding_model = DashScopeEmbeddings(\n",
    "    model=\"text-embedding-v3\",\n",
    "    dashscope_api_key=qwen_api_key\n",
    ")\n",
    "vectorstore = Chroma.from_documents(\n",
    "    documents=docs,\n",
    "    embedding=embedding_model,\n",
    "    persist_directory=\"./knowledge_db\"\n",
    ")\n",
    "\n",
    "# ========== 第五步：创建检索器 ==========\n",
    "retriever = vectorstore.as_retriever(search_kwargs={\"k\": 3})\n",
    "\n",
    "# ========== 第六步：提问 ==========\n",
    "query = \"员工被辞退的补偿有什么？\"\n",
    "relevant_docs = retriever.invoke(query)\n",
    "\n",
    "# ========== 第七步：生成回答 ==========\n",
    "context = \"\\n\\n---\\n\\n\".join([d.page_content for d in relevant_docs])\n",
    "\n",
    "prompt = ChatPromptTemplate.from_template(\"\"\"\n",
    "你是一个专业的知识库助手。请根据以下上下文回答问题。\n",
    "\n",
    "**规则：**\n",
    "- 只基于提供的上下文回答，不要编造\n",
    "- 如果上下文中没有相关信息，直接说「根据现有资料，我找不到这个问题的答案」\n",
    "- 回答要简洁直接，引用原文时用引号\n",
    "\n",
    "**上下文：**\n",
    "{context}\n",
    "\n",
    "**问题：**\n",
    "{question}\n",
    "\"\"\")\n",
    "\n",
    "llm = ChatTongyi(model=\"qwen-plus\", dashscope_api_key=qwen_api_key)\n",
    "chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "answer = chain.invoke({\"context\": context, \"question\": query})\n",
    "print(answer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "0701RAG (3.11.x)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}