{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Hybrid RAG Pipeline with Breakpoints\n", "\n", "This notebook demonstrates how to setup breakpoints in a Haystack pipeline. In this case, we will set up break points in a hybrid retrieval-augmented generation (RAG) pipeline. The pipeline combines BM25 and embedding-based retrieval methods, then uses a transformer-based reranker and an LLM to generate answers." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Install packages" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "%%bash\n", "\n", "pip install haystack-ai>=2.16.0\n", "pip install \"transformers[torch,sentencepiece]\"\n", "pip install sentence-transformers-haystack" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup OpenAI API keys" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "import os\n", "from getpass import getpass\n", "\n", "# Enable writing pipeline snapshots to disk when a breakpoint is hit\n", "os.environ[\"HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED\"] = \"true\"\n", "\n", "if \"OPENAI_API_KEY\" not in os.environ:\n", " os.environ[\"OPENAI_API_KEY\"] = getpass(\"Enter OpenAI API key:\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Import Required Libraries\n", "\n", "First, let's import all the necessary components from Haystack." ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "from haystack import Document, Pipeline\n", "from haystack.components.builders import AnswerBuilder, ChatPromptBuilder\n", "from haystack_integrations.components.embedders.sentence_transformers import SentenceTransformersDocumentEmbedder, SentenceTransformersTextEmbedder\n", "from haystack.components.generators.chat import OpenAIChatGenerator\n", "from haystack.components.joiners import DocumentJoiner\n", "from haystack_integrations.components.rankers.sentence_transformers import SentenceTransformersSimilarityRanker\n", "from haystack.components.retrievers.in_memory import InMemoryBM25Retriever, InMemoryEmbeddingRetriever\n", "from haystack.components.writers import DocumentWriter\n", "from haystack.dataclasses import ChatMessage\n", "from haystack.document_stores.in_memory import InMemoryDocumentStore\n", "from haystack.document_stores.types import DuplicatePolicy" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Document Store Initializations\n", "\n", "Let's create a simple document store with some sample documents and their embeddings." ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "def indexing():\n", " \"\"\"\n", " Indexing documents in a DocumentStore.\n", " \"\"\"\n", "\n", " print(\"Indexing documents...\")\n", "\n", " # Create sample documents\n", " documents = [\n", " Document(content=\"My name is Jean and I live in Paris. The weather today is 25°C.\"),\n", " Document(content=\"My name is Mark and I live in Berlin. The weather today is 15°C.\"),\n", " Document(content=\"My name is Giorgio and I live in Rome. The weather today is 30°C.\"),\n", " ]\n", "\n", " # Initialize document store and components\n", " document_store = InMemoryDocumentStore()\n", " doc_writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.SKIP)\n", " doc_embedder = SentenceTransformersDocumentEmbedder(model=\"intfloat/e5-base-v2\", progress_bar=False)\n", "\n", " # Build and run the ingestion pipeline\n", " ingestion_pipe = Pipeline()\n", " ingestion_pipe.add_component(instance=doc_embedder, name=\"doc_embedder\")\n", " ingestion_pipe.add_component(instance=doc_writer, name=\"doc_writer\")\n", "\n", " ingestion_pipe.connect(\"doc_embedder.documents\", \"doc_writer.documents\")\n", " ingestion_pipe.run({\"doc_embedder\": {\"documents\": documents}})\n", "\n", " return document_store" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## A Hybrid Retrieval Pipeline\n", "\n", "Now let's build a hybrid RAG pipeline." ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "def hybrid_retrieval(doc_store):\n", " \"\"\"\n", " A simple pipeline for hybrid retrieval using BM25 and embeddings.\n", " \"\"\"\n", "\n", " # Initialize query embedder\n", " query_embedder = SentenceTransformersTextEmbedder(model=\"intfloat/e5-base-v2\", progress_bar=False)\n", "\n", " # Define the prompt template for the LLM\n", " template = [\n", " ChatMessage.from_system(\n", " \"You are a helpful AI assistant. Answer the following question based on the given context information only. If the context is empty or just a '\\n' answer with None, example: 'None'.\"\n", " ),\n", " ChatMessage.from_user(\n", " \"\"\"\n", " Context:\n", " {% for document in documents %}\n", " {{ document.content }}\n", " {% endfor %}\n", " \n", " Question: {{question}}\n", " \"\"\"\n", " )\n", " ]\n", "\n", " \n", " # Build the RAG pipeline\n", " rag_pipeline = Pipeline()\n", " \n", " # Add components to the pipeline\n", " rag_pipeline.add_component(instance=InMemoryBM25Retriever(document_store=doc_store), name=\"bm25_retriever\")\n", " rag_pipeline.add_component(instance=query_embedder, name=\"query_embedder\")\n", " rag_pipeline.add_component(instance=InMemoryEmbeddingRetriever(document_store=doc_store), name=\"embedding_retriever\")\n", " rag_pipeline.add_component(instance=DocumentJoiner(sort_by_score=False), name=\"doc_joiner\")\n", " rag_pipeline.add_component(instance=SentenceTransformersSimilarityRanker(model=\"intfloat/simlm-msmarco-reranker\", top_k=5), name=\"ranker\") \n", " rag_pipeline.add_component(instance=ChatPromptBuilder(template=template, required_variables=[\"question\", \"documents\"]), name=\"prompt_builder\", ) \n", " rag_pipeline.add_component(instance=OpenAIChatGenerator(), name=\"llm\")\n", " rag_pipeline.add_component(instance=AnswerBuilder(), name=\"answer_builder\")\n", "\n", " # Connect the components\n", " rag_pipeline.connect(\"query_embedder\", \"embedding_retriever.query_embedding\")\n", " rag_pipeline.connect(\"embedding_retriever\", \"doc_joiner.documents\")\n", " rag_pipeline.connect(\"bm25_retriever\", \"doc_joiner.documents\")\n", " rag_pipeline.connect(\"doc_joiner\", \"ranker.documents\")\n", " rag_pipeline.connect(\"ranker\", \"prompt_builder.documents\")\n", " rag_pipeline.connect(\"prompt_builder\", \"llm\")\n", " rag_pipeline.connect(\"llm.replies\", \"answer_builder.replies\") \n", " rag_pipeline.connect(\"doc_joiner\", \"answer_builder.documents\")\n", "\n", " return rag_pipeline" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Running the pipeline with breakpoints\n", "\n", "Now we demonstrate how to set breakpoints in a Haystack pipeline to inspect and debug the pipeline execution at specific points. Breakpoints allow you to pause execution, save the current state of pipeline, and later resume from where you left off.\n", "\n", "We'll run the pipeline with a breakpoint set at the `query_embedder` component. This will save the pipeline state before executing the `query_embedder` and raise `PipelineBreakpointException` to stop execution.\n" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "from haystack.dataclasses.breakpoints import Breakpoint\n", "\n", "break_point = Breakpoint(component_name=\"query_embedder\", visit_count=0, snapshot_file_path=\"snapshots/\")" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Indexing documents...\n" ] }, { "name": "stderr", "output_type": "stream", "text": "\rLoading weights: 0%| | 0/199 [00:00, _content=[TextContent(text='Mark lives in Berlin.')], _name=None, _meta={'model': 'gpt-5-mini-2025-08-07', 'index': 0, 'finish_reason': 'stop', 'usage': {'completion_tokens': 78, 'prompt_tokens': 122, 'total_tokens': 200, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 64, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}})]}\n" ] } ], "source": [ " # Load the pipeline_snapshot and continue execution\n", "from haystack.core.pipeline.breakpoint import load_pipeline_snapshot\n", "\n", "import glob, os\n", "latest = max(glob.glob(\"snapshots/query_embedder_*.json\"), key=os.path.getctime)\n", "snapshot = load_pipeline_snapshot(latest)\n", "result = pipeline.run(data={}, pipeline_snapshot=snapshot)\n", " \n", "# Print the results\n", "print(result['answer_builder']['answers'][0].data)\n", "print(result['answer_builder']['answers'][0].meta)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Advanced Use Cases for Pipeline Breakpoints\n", "\n", "Here are some advanced scenarios where pipeline breakpoints can be particularly valuable:\n", "1. Set a breakpoint at the LLM to try results of different prompts and iterate in real time.\n", "\n", "2. Place a breakpoint after the document retriever to examine and modify retrieved documents.\n", "\n", "3. Set a breakpoint before a component to inject gold-standard inputs and isolate whether issues stem from input quality or downstream logic." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To demonstrate the use case stated in point 1, we reuse the same query pipeline with a new question. First, we run the pipeline with the prompt that we originally passed to the prompt_builder. Then, we define a breakpoint at the prompt_builder to try an alternative prompt. This allows us to compare the results generated by different prompts without running the whole pipeline again." ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Indexing documents...\n" ] }, { "name": "stderr", "output_type": "stream", "text": "\rLoading weights: 0%| | 0/201 [00:00