from dotenv import load_dotenv  # 加载 .env 文件中的环境变量
from pathlib import Path
load_dotenv(Path(__file__).parent.parent / '.env')  # 用脚本文件的绝对路径定位 .env

from langchain_community.embeddings import DashScopeEmbeddings  # 阿里 DashScope 文本向量化模型，用于把文本 chunks 转成向量
from langchain_community.document_loaders import PyPDFLoader  # PDF 文档加载器，把 PDF 解析成带 page_content 和 metadata 的 Document 对象
from langchain_community.vectorstores import Chroma  # Chroma 向量数据库，用于存储向量并做相似度检索
from langchain_text_splitters import RecursiveCharacterTextSplitter  # 递归字符分块器，按分隔符层级把长文本切成小块
from langchain_openai import ChatOpenAI  # DeepSeek API 兼容 OpenAI 接口，用 ChatOpenAI 调用
from langchain_core.prompts import ChatPromptTemplate  # 聊天提示词模板，统一管理 system/user 消息格式
from langchain_core.output_parsers import StrOutputParser  # 输出解析器，把 LLM 返回的 AIMessage 提取为纯字符串
import os  # 标准库，用于读取环境变量（如 DASHSCOPE_API_KEY）


# ========== 第一步：加载文档 ==========
loader = PyPDFLoader('./1.大模型全景认知.pdf')
datas = loader.load()
print(f'原始页数：{len(datas)}')

# ========== 第二步：分块 ==========
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 200,
    chunk_overlap = 40,
    separators=['\n\n','\n','。','?','!',' ','']
)

chunks = text_splitter.split_documents(datas)
print(f'分块后的页数：{len(chunks)}')

# ========== 第四步：向量化 + 存入向量库 ==========
embeddings_model = DashScopeEmbeddings(
    model='text-embedding-v3',
    dashscope_api_key = os.getenv('DASHSCOPE_API_KEY')
)

vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings_model,
    collection_metadata={'hnsw:space': 'cosine'},
    persist_directory='./chroma_db2',
)

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

# ========== 第六步：提问 ==========
query="什么是人工智能，一句话总结"
docs = retriever.invoke(query)
print(docs)
print(f'检索到的文档数：{len(docs)}')

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

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

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

参考资料：
{context}"""),
    ("human", "{query}")
])

llm = ChatOpenAI(
    model=os.getenv('moduel', 'deepseek-chat'),  # .env 中的 moduel 字段
    api_key=os.getenv('OPENAI_API_KEY'),          # .env 中的 OPENAI_API_KEY
    base_url=os.getenv('base_url'),               # https://api.deepseek.com
)

chain = prompt | llm | StrOutputParser()

# 执行整条链，获取回答
answer = chain.invoke({"context": context, "query": query})
print(answer)