ZhouYI/03_RAG优化/agent/rag/rag_cli.py

"""RAG 命令行入口。"""

from pathlib import Path

from langchain_core.prompts import ChatPromptTemplate

from agent.config import load_config
from agent.llm import create_llm
from agent.rag.document import load_pdf_document
from agent.rag.embedding import get_embedding_model
from agent.rag.multi_retriever import MultiRetrieverSystem
from agent.rag.splitter import recursive_split

from query_optimization import get_prompt_template
DOCUMENT_PATH = Path("car_info.pdf")


def create_rag_system() -> tuple[MultiRetrieverSystem, object, ChatPromptTemplate]:
    """加载文档并创建检索器和问答链所需组件。"""
    pages = load_pdf_document(str(DOCUMENT_PATH))
    documents = recursive_split(pages)

    config = load_config()
    llm = create_llm(config)
    retriever_system = MultiRetrieverSystem(
        documents=documents,
        embedding_model=get_embedding_model(),
        llm=llm,
    )
    
    return retriever_system, llm


def run_cli() -> None:
    """启动交互式 RAG 命令行。"""
    retriever_system, llm = create_rag_system()
    prompt =get_prompt_template("rewrite")
    mode = "ensemble"
    valid_modes = {"bm25", "vector", "ensemble"}
    print("检索模式：ensemble（可输入 /mode bm25、/mode vector 或 /mode ensemble 切换）")

    while True:
        question = input("请输入问题：").strip()
        if not question:
            continue
        if question.lower() in {"exit", "quit", "q"}:
            break
        if question.startswith("/mode "):
            new_mode = question.removeprefix("/mode ").strip().lower()
            if new_mode in valid_modes:
                mode = new_mode
                print(f"已切换为 {mode} 检索。")
            else:
                print("不支持的检索模式，可选：bm25、vector、ensemble")
            continue

        relevant_docs = retriever_system.search(question, mode)
        context = "\n\n".join(document.page_content for document in relevant_docs)
        answer = (prompt | llm).invoke({"context": context, "question": question})
        print(f"AI 回答：{answer.content}")


if __name__ == "__main__":
    run_cli()