from agent.rag.document import load_pdf_document, clean_pdf_text
from agent.rag import splitter
from agent.rag.embedding import get_embedding_model, get_vectorstore
from agent.llm import create_llm
from agent.config import load_config

from langchain_core.prompts import ChatPromptTemplate

#加载文档
pages = load_pdf_document("car_info.pdf")
#清洗
#documents = [clean_pdf_text(page.page_content) for page in pages]

#分块
#chunks=[splitter.embedding_split(document)for document in documents]
docs=splitter.recursive_split(pages)
#向量化+构建向量数据库

embedding_model = get_embedding_model()
db = get_vectorstore(embedding_model, docs)
#创建检索器
retriever = db.as_retriever(search_kwargs={"k": 3})



context=""
config=load_config()
llm = create_llm(config)
prompt=ChatPromptTemplate.from_messages([
    ("system", "你是⼀个专业的知识库助⼿。请根据以下上下⽂回答问题。"),
    ("user", "根据以下内容回答用户问题，如果无法从中获取答案，请说“抱歉，我无法回答这个问题。”\n\n{context}\n\n用户问题: {question}")
])
while True:
    question = input("请输入: ").strip()
    if not question:
        continue

    if question.lower() in ["exit", "quit", "q"]:
        break
    #查询向量数据库
    relevant_docs = retriever.invoke(question)
    for doc in relevant_docs:
        context += doc.page_content + "\n"

    #生成回答
    chain=prompt|llm
    answer = chain.invoke({"context": context, "question": question})
    print("AI回答:", answer)