S SmartDocs
系列: Ricky python 113 行 · 更新於 2026-04-07

main.py

Ricky/RAG_project/main.py

"""
Entry point: startup checks, index load/build, and interactive CLI loop.
All retrieval, confidence, and answer logic lives in dedicated modules.
"""
import logging
import os

from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(
    level=logging.WARNING,
    format="%(levelname)s %(name)s: %(message)s",
)

from llm.providers import check_ollama_embedding_ready, choose_llm, warmup_llm
from ingestion.pipeline import load_or_build_database, create_keyword_retriever
from service.chat import answer_one_turn
from io.voice import to_console_text, init_voice_input, get_voice_question
from config.runtime import is_warmup_enabled, is_voice_input_enabled


def run_chat_loop(vectorstore, keyword_retriever, llm, llm_provider: str) -> None:
    """Interactive question-answer terminal loop."""
    voice_enabled = is_voice_input_enabled()
    recognizer, microphone = None, None
    if voice_enabled:
        recognizer, microphone = init_voice_input()
        if not recognizer:
            print("[WARN] Voice input unavailable, falling back to keyboard input.")
            print("[TIP] Install: pip install SpeechRecognition PyAudio")
            voice_enabled = False

    print("\n=== RAG Q&A System ===")
    print(f"[INFO] Answer model: {llm_provider}")
    if voice_enabled:
        print("Type 'quit' to exit; press Enter to use voice input\n")
    else:
        print("Type 'quit' to exit\n")

    while True:
        try:
            raw = input("Your question: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nGoodbye!")
            break

        if voice_enabled and raw == "":
            question = get_voice_question(recognizer, microphone)
            if not question:
                continue
        else:
            question = raw

        if question.lower() == "quit":
            print("Goodbye!")
            break
        if not question:
            print("[WARN] Please enter a question, or type 'quit' to exit.")
            continue

        try:
            result = answer_one_turn(
                vectorstore, keyword_retriever, llm, llm_provider, question
            )

            # Print answer (the service layer already streamed the LLM output
            # for the generative path; here we handle other branches)
            if result.fallback_mode == "no_docs":
                print(f"\nAnswer: {to_console_text(result.response)}")
                print("[INFO] No document chunks retrieved")
            elif result.fallback_mode == "image":
                print(f"\nAnswer: {to_console_text(result.response)}")
            elif result.fallback_mode == "image_low_confidence":
                print(f"\nAnswer: {to_console_text(result.response)}")
            elif result.fallback_mode == "snippet":
                print(f"\nAnswer: {to_console_text(result.response)}")
            elif result.fallback_mode == "extractive":
                print(f"\nAnswer: {to_console_text(result.response)}")
            elif result.fallback_mode == "llm":
                if not result.response.strip():
                    print("I cannot find relevant information in the documents.", end="")
            # fallback_mode="" means an error was raised (handled below)

            print(f"\n[INFO] Response time: {result.elapsed:.1f}s\n")

        except Exception as e:
            print(f"\n[ERROR] Failed to generate answer: {e}")
            print("[TIP] Please verify the answer model is available and try again.\n")


def main() -> None:
    if not check_ollama_embedding_ready():
        return

    llm, llm_provider = choose_llm()
    if llm is None:
        return

    if is_warmup_enabled():
        warmup_llm()

    vectorstore, chunks_for_keyword = load_or_build_database()
    if vectorstore is None:
        return

    keyword_retriever = create_keyword_retriever(chunks_for_keyword)
    run_chat_loop(vectorstore, keyword_retriever, llm, llm_provider)


if __name__ == "__main__":
    main()

相關文章