# ============================================================ # 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, )