src/rag_pipeline.py

# ============================================================
# FILE: src/rag_pipeline.py
# ============================================================
# PURPOSE:
# Orchestrate the full RAG workflow.
#
# FULL FLOW:
#
# documents
# -> clean text
# -> chunks
# -> embeddings
# -> ChromaDB
# -> retrieve relevant chunks
# -> build prompt
# -> call cloud LLM
# -> return grounded answer
#
# AI ENGINEER PRODUCTION TIP:
# Keep orchestration separate from individual components.
# This makes your app easier to test, debug, and deploy.
# ============================================================

import re
import time
from typing import Any, Dict, List

from src.chunker import build_chunks_from_documents
from src.config import AppConfig
from src.document_loader import Document, load_documents
from src.embeddings import EmbeddingModel
from src.llm_client import CloudLLMClient
from src.logging_utils import save_json_output, write_jsonl_event
from src.retriever import Retriever
from src.text_cleaner import clean_text
from src.vector_store import ChromaVectorStore


class RAGPipeline:
    """
    Complete RAG pipeline.
    """

    def __init__(self, config: AppConfig) -> None:
        """
        Initialize all main components.

        Components:
        - embedding model
        - vector store
        - retriever
        - cloud LLM client
        """

        self.config = config

        self.embedding_model = EmbeddingModel(
            model_name=config.embedding_model_name,
            device=config.embedding_device,
        )

        self.vector_store = ChromaVectorStore(
            persist_directory=config.vector_db_folder,
            collection_name=config.collection_name,
            embedding_model_name=config.embedding_model_name,
        )

        self.retriever = Retriever(
            embedding_model=self.embedding_model,
            vector_store=self.vector_store,
        )

        self.llm_client = CloudLLMClient(config=config)

    def load_and_clean_documents(self) -> List[Document]:
        """
        Load documents and clean their text.
        """

        documents = load_documents(
            folder=self.config.data_folder,
            project_root=self.config.project_root,
        )

        cleaned_documents = []

        for document in documents:
            cleaned_text = clean_text(document.text)

            cleaned_documents.append(
                Document(
                    source=document.source,
                    text=cleaned_text,
                    file_type=document.file_type,
                    character_count=len(cleaned_text),
                )
            )

        return cleaned_documents

    def rebuild_vector_database(self) -> Dict[str, Any]:
        """
        Rebuild the vector database from files in data/raw/.

        Development behavior:
        - reset collection
        - re-index all chunks

        Production behavior should eventually:
        - detect changed files
        - upsert only changed chunks
        - preserve document versions
        """

        documents = self.load_and_clean_documents()

        chunks = build_chunks_from_documents(
            documents=documents,
            chunk_size=self.config.chunk_size,
            chunk_overlap=self.config.chunk_overlap,
        )

        self.vector_store.reset_collection()

        if chunks:
            chunk_texts = [chunk.text for chunk in chunks]
            embeddings = self.embedding_model.embed_texts(chunk_texts)

            self.vector_store.add_chunks(
                chunks=chunks,
                embeddings=embeddings,
            )

        result = {
            "documents_loaded": len(documents),
            "chunks_created": len(chunks),
            "vectors_stored": self.vector_store.count(),
            "collection_name": self.config.collection_name,
            "embedding_model": self.config.embedding_model_name,
        }

        write_jsonl_event(
            logs_folder=self.config.logs_folder,
            event={
                "event_type": "vector_database_rebuilt",
                **result,
            },
        )

        return result

    @staticmethod
    def format_context(retrieved_chunks: List[Dict[str, Any]]) -> str:
        """
        Convert retrieved chunks into a readable context block.
        """

        if not retrieved_chunks:
            return "No relevant context was retrieved."

        context_parts = []

        for item in retrieved_chunks:
            context_parts.append(
                f"[Source: {item['source']} | "
                f"Chunk: {item['chunk_index']} | "
                f"Rank: {item['rank']}]\n"
                f"{item['text']}"
            )

        return "\n\n---\n\n".join(context_parts)

    def build_messages(
        self,
        question: str,
        retrieved_chunks: List[Dict[str, Any]],
    ) -> List[Dict[str, str]]:
        """
        Build OpenAI-compatible chat messages.

        The system message defines behavior.
        The user message provides retrieved context and the question.
        """

        context = self.format_context(retrieved_chunks)

        if self.config.require_context_for_answer:
            answer_rule = (
                "Use only the provided context. "
                "If the answer is not in the context, say: "
                "'I do not know from the provided knowledge base.'"
            )
        else:
            answer_rule = (
                "Use the provided context first. "
                "If needed, you may use general knowledge, but clearly say when you do."
            )

        system_message = f"""
            You are KnowFlow AI, a careful Retrieval-Augmented Generation assistant.

            Rules:
            - {answer_rule}
            - Keep the answer simple, clear, and useful.
            - Do not invent facts.
            - Mention the source file when useful.
            - If multiple sources are used, summarize them clearly.
            - If retrieved context is incomplete, be honest.

            Prompt template version:
            {self.config.prompt_template_version}
            """.strip()

        user_message = f"""
            Retrieved context:
            {context}

            User question:
            {question}

            Answer:
            """.strip()

        return [
            {
                "role": "system",
                "content": system_message,
            },
            {
                "role": "user",
                "content": user_message,
            },
        ]

    def ask(self, question: str, top_k: int | None = None) -> Dict[str, Any]:
        """
        Ask a question using the full RAG pipeline.

        Returns structured output:
        - question
        - answer
        - retrieved chunks
        - model info
        - timing
        - raw response
        """

        question = question.strip()

        if not question:
            raise ValueError("Question cannot be empty.")

        if top_k is None:
            top_k = self.config.top_k

        start_time = time.time()

        retrieved_chunks = self.retriever.retrieve(
            question=question,
            top_k=top_k,
        )

        messages = self.build_messages(
            question=question,
            retrieved_chunks=retrieved_chunks,
        )

        llm_result = self.llm_client.chat(messages)

        total_elapsed_seconds = round(time.time() - start_time, 3)

        result = {
            "question": question,
            "answer": llm_result["answer"],
            "retrieved_chunks": retrieved_chunks,
            "messages": messages,
            "raw_response": llm_result["raw_response"],
            "status_code": llm_result["status_code"],
            "llm_elapsed_seconds": llm_result["elapsed_seconds"],
            "total_elapsed_seconds": total_elapsed_seconds,
            "attempts": llm_result["attempts"],
            "provider": self.config.cloud_api_provider,
            "model": self.config.cloud_chat_model,
            "embedding_model": self.config.embedding_model_name,
            "prompt_template_version": self.config.prompt_template_version,
            "top_k": top_k,
        }

        write_jsonl_event(
            logs_folder=self.config.logs_folder,
            event={
                "event_type": "rag_question_answered",
                "question": question,
                "answer_preview": llm_result["answer"][:300],
                "model": self.config.cloud_chat_model,
                "provider": self.config.cloud_api_provider,
                "top_k": top_k,
                "total_elapsed_seconds": total_elapsed_seconds,
                "retrieved_sources": [
                    {
                        "source": item["source"],
                        "chunk_index": item["chunk_index"],
                        "distance": item["distance"],
                    }
                    for item in retrieved_chunks
                ],
            },
        )

        return result

    def save_result(self, result: Dict[str, Any]) -> str:
        """
        Save one RAG result to outputs folder.
        """

        timestamp = time.strftime("%Y%m%d_%H%M%S")
        safe_question = re.sub(r"[^a-zA-Z0-9]+", "_", result["question"][:50]).strip("_")

        file_name = f"rag_result_{timestamp}_{safe_question}.json"

        output_path = save_json_output(
            outputs_folder=self.config.outputs_folder,
            data=result,
            file_name=file_name,
        )

        return str(output_path)

    def debug_retrieval(self, question: str, top_k: int | None = None) -> List[Dict[str, Any]]:
        """
        Retrieve chunks without calling the LLM.

        Use this when debugging RAG quality.

        If retrieved chunks do not contain the answer, fix retrieval first.
        """

        if top_k is None:
            top_k = self.config.top_k

        return self.retriever.retrieve(
            question=question,
            top_k=top_k,
        )