from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.chat_models import ChatTongyi
from langchain_core.output_parsers import StrOutputParser

# ========== 第一步：加载文档 ==========
loader = PyMuPDFLoader("./2222.pdf")
pages = loader.load()

# ========== 第二步：清洗数据（可选，根据文档质量决定）==========
# clean_pages = [clean_pdf_text(page) for page in pages]

# ========== 第三步：分块 ==========
splitter = RecursiveCharacterTextSplitter(
    # 分隔符优先级：段落 → 换行 → 句号 → 空格 → 硬切
    #separators=["\n\n", "\n", "。", "！", "？", " ", ""],
    separators = ["\n\n", "\n", "。", "！", "？", " "],
    # 每个块最大 50 字符
    chunk_size=100,
    # 相邻块重叠 10 字符（chunk_size 的 20%）
    chunk_overlap=10,
    # 长度计算函数
    length_function=len
)
docs = splitter.split_documents(pages)

# ========== 第四步：向量化 + 存入向量库 ==========

vectorstore = Chroma.from_documents(
    documents=docs,
    embedding=embedding_model,
    persist_directory="./test1_db"
)

# ========== 第五步：创建检索器 ==========
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# ========== 第六步：提问 ==========
query = "团队管理"
relevant_docs = retriever.invoke(query)

# ========== 第七步：生成回答 ==========
context = "\n\n---\n\n".join([d.page_content for d in relevant_docs])

prompt = ChatPromptTemplate.from_template("""
你是一个专业的知识库助手。请根据以下上下文回答问题。

**规则：**
- 只基于提供的上下文回答，不要编造
- 如果上下文中没有相关信息，直接说「根据现有资料，我找不到这个问题的答案」
- 回答要简洁直接，引用原文时用引号

**上下文：**
{context}

**问题：**
{question}
""")

llm = ChatOpenAI(
    model_name="kimi-k2.6",
    api_key=ds_open_key,                    # 在 platform.deepseek.com 获取
    base_url=ds_open_url     # DeepSeek API 地址
)
chain = prompt | llm | StrOutputParser()

answer = chain.invoke({"context": context, "question": query})
print(answer)

根据上下文，团队管理的相关内容如下：

**提高团队绩效的措施：**
- "改进方法提高效率"
- "采用快速跟进"
- "制定奖惩规则，提高团队绩效"

**团队管理中存在的问题：**
- "没有明确各团队成员的职责"
- "没有制定团队基本规则"
- "对于团队冲突没有及时、合理解决"