query.py

import os
import re
import sys
from vector_store import create_or_load_db
from retriever import get_retriever
from qa_chain import build_qa_chain

def sanitize_filename(text, max_len=50):
    """把问题转换成合法的文件名"""
    # 只保留中文、字母、数字、空格、连字符
    text = re.sub(r'[^\w\s\-]', '', text)
    # 压缩空格
    text = re.sub(r'\s+', '_', text.strip())
    # 限制长度
    return text[:max_len]

if __name__ == "__main__":
    vectordb = create_or_load_db()
    retriever = get_retriever(vectordb)
    qa = build_qa_chain(retriever)

    question = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else input("请输入问题：")
    result = qa.invoke({"query": question})

    answer = result["result"]
    source_docs = result["source_documents"]

    # ====== 输出到终端 ======
    print("\n" + "="*50)
    print("【回答】")
    print(answer)
    print("\n" + "-"*30)
    print("📍 参考定位")
    for i, doc in enumerate(source_docs, 1):
        chapter = doc.metadata.get("章节", "")
        section = doc.metadata.get("节", "")
        sub_section = doc.metadata.get("小节", "")
        source_file = doc.metadata.get("source", "").split("/")[-1]
        location = f"{source_file} "
        if chapter: location += f"第{chapter}章"
        if section: location += f" > {section}"
        if sub_section: location += f" > {sub_section}"
        snippet = doc.page_content[:30].replace("\n", " ")
        print(f"[{i}] {location}\n    ↳ {snippet}...")

    # ====== 保存为 Markdown 文件 ======
    safe_name = sanitize_filename(question)
    filepath = f"./{safe_name}.md"
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(f"# {question}\n\n")
        f.write(answer)
        f.write("\n\n---\n## 参考定位\n")
        for i, doc in enumerate(source_docs, 1):
            chapter = doc.metadata.get("章节", "")
            section = doc.metadata.get("节", "")
            sub_section = doc.metadata.get("小节", "")
            source_file = doc.metadata.get("source", "").split("/")[-1]
            location = f"{source_file} "
            if chapter: location += f"第{chapter}章"
            if section: location += f" > {section}"
            if sub_section: location += f" > {sub_section}"
            snippet = doc.page_content[:30].replace("\n", " ")
            f.write(f"- [{i}] {location}\n  ↳ {snippet}...\n")
    print(f"\n📁 回答已保存至：{filepath}")