rag_pipeline.py

"""
Full RAG pipeline orchestration.

Two entry points:
  - build_knowledge_base(): run once to index documents
  - answer(question, profile, language): called per user query
"""

from config import DATA_RAW_DIR, DISCLAIMER
from data_loader import load_all_documents
from preprocessing import preprocess_documents
from chunking import chunk_documents
from embeddings import embed_texts
from vector_store import build_index, save_index
from retriever import retrieve, reset_index_cache
from prompt_builder import build_prompt, format_sources
from generator import generate_answer


def build_knowledge_base() -> None:
    """
    Full pipeline to build and persist the vector index.
    Run this once after adding/changing documents in data/raw/.
    """
    print("=== Building knowledge base ===")
    docs = load_all_documents(DATA_RAW_DIR)
    docs = preprocess_documents(docs)
    chunks = chunk_documents(docs)
    texts = [c["text"] for c in chunks]
    embeddings = embed_texts(texts)
    index = build_index(embeddings)
    save_index(index, chunks)
    reset_index_cache()
    print("=== Knowledge base ready ===")


def answer(question: str, profile: str, language: str = "English") -> dict:
    """
    Full RAG query pipeline.
    Returns a dict with keys: answer, sources, disclaimer, error.
    """
    if not question.strip():
        return {"answer": "", "sources": "", "disclaimer": "", "error": "Please enter a question."}

    try:
        chunks = retrieve(question)
        if not chunks:
            return {
                "answer": "No relevant information was found in the knowledge base for this question.",
                "sources": "",
                "disclaimer": DISCLAIMER,
                "error": "",
            }

        prompt = build_prompt(question, chunks, profile, language)
        generated = generate_answer(prompt)
        sources = format_sources(chunks)

        return {
            "answer": generated.strip(),
            "sources": sources,
            "disclaimer": DISCLAIMER,
            "error": "",
        }

    except FileNotFoundError as e:
        return {
            "answer": "",
            "sources": "",
            "disclaimer": "",
            "error": str(e),
        }
    except Exception as e:
        return {
            "answer": "",
            "sources": "",
            "disclaimer": "",
            "error": f"Unexpected error: {e}",
        }