feat: capture RAG documents and pass as context to action plan (#198)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: michelle-hadfield-nava

app/src/common/components.py

@@ -11,9 +11,10 @@
 
 import json
 import logging
+import threading
 from json import JSONDecodeError
 from pprint import pformat
-from typing import Any, Callable, List, Optional, TypeVar
+from typing import Any, Callable, ClassVar, List, Optional, TypeVar
 from uuid import UUID
 
 from fastapi import UploadFile
@@ -675,6 +676,35 @@ def parse_json_if_possible(self, content: Any) -> Any:
         return content
 
 
+@component
+class DocumentCapture:
+    """Passes Haystack Documents through unchanged while capturing their content as strings.
+
+    Acts as a side-effect component: documents flow to the next component as normal,
+    but their text content is stored in a class-level dict keyed by result_id so that
+    the content can be retrieved later (e.g. in a streaming generator hook).
+
+    Thread-safe: multiple concurrent pipeline runs are isolated by result_id.
+    """
+
+    _storage: ClassVar[dict[str, list[str]]] = {}
+    _lock: ClassVar[threading.Lock] = threading.Lock()
+
+    @component.output_types(documents=list[Document])
+    def run(self, documents: list[Document], result_id: str = "") -> dict:
+        if result_id:
+            contents = [d.content for d in documents if d.content]
+            with self.__class__._lock:
+                self.__class__._storage[result_id] = contents
+        return {"documents": documents}
+
+    @classmethod
+    def pop(cls, result_id: str) -> list[str]:
+        """Retrieve and remove stored document content for the given result_id."""
+        with cls._lock:
+            return cls._storage.pop(result_id, [])
+
+
 @component
 class DocumentMetadataAdder:
     def __init__(self, metadata: dict[str, Any]) -> None:

app/src/common/haystack_utils.py

@@ -1,3 +1,4 @@
+import json
 import logging
 from typing import Any, Callable, Generator, Sequence
 
@@ -67,7 +68,11 @@ def to_chat_messages(
     return messages
 
 
-def create_result_id_hook(pipeline: Pipeline, result_id: str) -> Callable[[dict], Generator]:
+def create_result_id_hook(
+    pipeline: Pipeline,
+    result_id: str,
+    pop_documents_fn: Callable[[], dict] | None = None,
+) -> Callable[[dict], Generator]:
     """Creates a generator hook that yields the result_id as the last chunk in a streaming response,
     i.e., TracedPipelineRunner.stream_response() calls a generator_hook() after pipeline.run() completes.
 
@@ -77,6 +82,8 @@ def create_result_id_hook(pipeline: Pipeline, result_id: str) -> Callable[[dict]
     Args:
         pipeline: The Haystack pipeline to check for SaveResult component
         result_id: The result_id that will be used by SaveResult and yielded to frontend
+        pop_documents_fn: Optional callable that retrieves and clears captured documents,
+            returning them as a dict to merge into the yielded JSON
 
     Raises:
         ValueError: If the pipeline does not have a SaveResult component
@@ -96,8 +103,10 @@ def create_result_id_hook(pipeline: Pipeline, result_id: str) -> Callable[[dict]
         )
 
     def hook(pipeline_run_args: dict) -> Generator:
-        # Yield result_id as last chunk for frontend
-        yield StreamingChunk(content=f'{{"result_id": "{result_id}"}}\n')
+        data: dict = {"result_id": result_id}
+        if pop_documents_fn:
+            data.update(pop_documents_fn())
+        yield StreamingChunk(content=json.dumps(data) + "\n")
 
     return hook
 
@@ -132,6 +141,7 @@ def stream_response(
         shorten_output: Callable[[str], str] = lambda resp: resp,
         parent_span_name_suffix: str | None = None,
         generator_hook: Callable[[dict], Generator] | None = None,
+        cleanup_fn: Callable[[], Any] | None = None,
     ) -> Generator:
         """
         Run the pipeline with tracing and return a streaming response using hayhooks.streaming_generator().
@@ -140,6 +150,9 @@ def stream_response(
         The parent_span_name_suffix is appended to the parent span name for easier region identification in Phoenix.
         The generator_hook can be used to yield additional chunks after the main pipeline.run() completes,
         such as yielding the result_id for reference by the frontend.
+        The cleanup_fn is called in a finally block, guaranteeing execution even when the pipeline raises
+        before the generator_hook runs (e.g., to release class-level component state like DocumentCapture entries).
+        Its return value is ignored.
         """
         # Must set using attributes and metadata tracer context before calling tracer.start_as_current_span()
         with using_attributes(user_id=user_id, metadata=metadata):
@@ -191,6 +204,9 @@ def stream_response(
                     span.set_status(Status(StatusCode.ERROR, str(e)))
                     span.record_exception(e)
                     raise HTTPException(status_code=500, detail=str(e)) from e
+                finally:
+                    if cleanup_fn:
+                        cleanup_fn()
 
     def return_response(
         self,

app/src/pipelines/generate_action_plan/pipeline_wrapper.py

@@ -45,7 +45,7 @@ def setup(self) -> None:
 
         pipeline.add_component(
             instance=ChatPromptBuilder(
-                variables=["resources", "action_plan_json", "user_query"],
+                variables=["resources", "action_plan_json", "user_query", "context_documents"],
             ),
             name="prompt_builder",
         )
@@ -68,16 +68,20 @@ def run_api(
         resources: list[Resource] | list[dict],
         user_email: str,
         user_query: str,
+        context_documents: list[str] | None = None,
     ) -> dict:
         """
         Generate an action plan based on the given resources.
         The user query provides more context to the generation process.
+        The optional context_documents are RAG source documents from the referrals pipeline,
+        passed as additional context to the LLM prompt.
         """
         resource_objects = get_resources(resources)
         pipeline_run_args = self.create_pipeline_args(
             user_email,
             resource_objects,
             user_query,
+            context_documents=context_documents,
         )
         response = self.runner.return_response(
             pipeline_run_args,
@@ -100,6 +104,7 @@ def create_pipeline_args(
         resource_objects: list[Resource],
         user_query: str,
         *,
+        context_documents: list[str] | None = None,
         llm_model: str | None = None,
         reasoning_effort: str | None = None,
         streaming: bool = False,
@@ -116,6 +121,7 @@ def create_pipeline_args(
                 "resources": format_resources(resource_objects),
                 "action_plan_json": action_plan_as_json,
                 "user_query": user_query,
+                "context_documents": format_context_documents(context_documents),
             },
             "llm": {
                 "model": llm_model or config.generate_action_plan_model_version,
@@ -135,6 +141,7 @@ def run_chat_completion(self, model: str, messages: list, body: dict) -> Generat
         resources = body.get("resources", [])
         user_email = body.get("user_email", "")
         user_query = body.get("user_query", "")
+        context_documents = body.get("context_documents", None)
 
         if not resources:
             raise ValueError("resources parameter is required")
@@ -146,6 +153,7 @@ def run_chat_completion(self, model: str, messages: list, body: dict) -> Generat
             user_email,
             resource_objects,
             user_query,
+            context_documents=context_documents,
             llm_model=body.get("llm_model", None),
             reasoning_effort=body.get("reasoning_effort", None),
             streaming=True,
@@ -174,6 +182,14 @@ def get_resources(resources: list[Resource] | list[dict]) -> list[Resource]:
     return [Resource(**res) for res in resources]  # type: ignore[arg-type]
 
 
+def format_context_documents(documents: list[str] | None) -> str:
+    """Format a list of retrieved RAG documents into a readable string for the LLM prompt."""
+    if not documents:
+        return ""
+    doc_lines = "\n\n".join(f"- {doc}" for doc in documents)
+    return f"Source Documents:\n{doc_lines}"
+
+
 def format_resources(resources: list[Resource]) -> str:
     """Format a list of Resource objects into a readable string."""
     formatted_resources = []

app/src/pipelines/generate_referrals_rag/pipeline_wrapper.py

@@ -105,8 +105,11 @@ def _create_pipeline(self) -> Pipeline:
                 output_type=list,
             ),
         )
+        pipeline.add_component("document_capture", components.DocumentCapture())
+
         pipeline.connect("query_embedder.embedding", "retriever.query_embedding")
-        pipeline.connect("retriever.documents", "output_adapter")
+        pipeline.connect("retriever.documents", "document_capture.documents")
+        pipeline.connect("document_capture.documents", "output_adapter")
 
         pipeline.add_component(
             "prompt_builder",
@@ -173,7 +176,7 @@ def run_api(
             pipeline_run_args,
             user_id=user_email,
             metadata={"user_id": user_email},
-            include_outputs_from={"llm", "save_result"},
+            include_outputs_from={"llm", "save_result", "retriever"},
             input_=query,
             extract_output=extract_output,
             parent_span_name_suffix=suffix,
@@ -192,6 +195,7 @@ def create_pipeline_args(
         llm_model: str | None = None,
         reasoning_effort: str | None = None,
         streaming: bool = False,
+        result_id: str | None = None,
     ) -> dict:
         """Create pipeline run arguments with optional overrides for model, reasoning effort, and streaming."""
         assert suffix, "suffix is required"
@@ -207,7 +211,7 @@ def create_pipeline_args(
                 detail=f"The requested prompt version '{prompt_version_id}' with suffix '{suffix}' could not be retrieved",
             ) from e
 
-        return {
+        args = {
             "logger": {
                 "messages_list": [{"query": query, "user_email": user_email}],
             },
@@ -230,6 +234,10 @@ def create_pipeline_args(
                 "filters": {"field": "region", "operator": "==", "value": region},
             },
         }
+        if result_id:
+            args["save_result"] = {"result_id": result_id}
+            args["document_capture"] = {"result_id": result_id}
+        return args
 
     # https://docs.haystack.deepset.ai/docs/hayhooks#openai-compatibility
     # This function is called for the `{pipeline_name}/chat`, `/chat/completions`, or `/v1/chat/completions` streaming endpoint using Server-Sent Events (SSE)
@@ -249,6 +257,8 @@ def run_chat_completion(self, model: str, messages: list, body: dict) -> Generat
         if not user_email:
             raise ValueError("user_email parameter is required")
 
+        # Generate result_id upfront to pass to both SaveResult and DocumentCapture via pipeline args
+        result_id = str(uuid.uuid4())
         pipeline_run_args = self.create_pipeline_args(
             query,
             user_email,
@@ -258,12 +268,9 @@ def run_chat_completion(self, model: str, messages: list, body: dict) -> Generat
             llm_model=body.get("llm_model", None),
             reasoning_effort=body.get("reasoning_effort", None),
             streaming=True,
+            result_id=result_id,
         )
 
-        # Generate result_id upfront to pass to both SaveResult and the hook
-        result_id = str(uuid.uuid4())
-        pipeline_run_args["save_result"] = {"result_id": result_id}
-
         logger.info("Streaming referrals: %s", pipeline_run_args)
         return self.runner.stream_response(
             pipeline_run_args,
@@ -272,5 +279,13 @@ def run_chat_completion(self, model: str, messages: list, body: dict) -> Generat
             input_=query,
             shorten_output=shorten_output,
             parent_span_name_suffix=suffix,
-            generator_hook=haystack_utils.create_result_id_hook(self.pipeline, result_id),
+            generator_hook=haystack_utils.create_result_id_hook(
+                self.pipeline,
+                result_id,
+                pop_documents_fn=lambda: {"documents": components.DocumentCapture.pop(result_id)},
+            ),
+            # No-op if the hook already ran (pop returns [] for missing keys).
+            # Ensures the DocumentCapture entry is always removed even when the pipeline
+            # raises before the hook executes.
+            cleanup_fn=lambda: components.DocumentCapture.pop(result_id),
         )

app/tests/src/common/test_components.py

@@ -7,6 +7,7 @@
 
 from src.adapters import db
 from src.common.components import (
+    DocumentCapture,
     EmailResponses,
     LlmOutputValidator,
     LoadResultOptional,
@@ -607,6 +608,109 @@ def test_RemoveResourcesForEmail_preserves_other_fields():
     assert resources_dict["resources"][0]["name"] == "Resource A"
 
 
+def test_DocumentCapture_stores_and_pops_content():
+    """DocumentCapture stores document content and pop() retrieves and removes it."""
+    from haystack import Document
+
+    component = DocumentCapture()
+    docs = [Document(content="First doc"), Document(content="Second doc")]
+
+    output = component.run(documents=docs, result_id="test-id-1")
+
+    # Documents pass through unchanged
+    assert output["documents"] == docs
+
+    # Content was stored
+    retrieved = DocumentCapture.pop("test-id-1")
+    assert retrieved == ["First doc", "Second doc"]
+
+    # Second pop returns empty (entry removed)
+    assert DocumentCapture.pop("test-id-1") == []
+
+
+def test_DocumentCapture_no_op_without_result_id():
+    """DocumentCapture does not store anything when result_id is empty."""
+    from haystack import Document
+
+    component = DocumentCapture()
+    docs = [Document(content="Some content")]
+
+    output = component.run(documents=docs, result_id="")
+
+    # Documents still pass through
+    assert output["documents"] == docs
+
+    # Nothing stored (pop of any key returns [])
+    assert DocumentCapture.pop("") == []
+
+
+def test_DocumentCapture_pop_nonexistent_key():
+    """pop() returns an empty list for an unknown result_id."""
+    assert DocumentCapture.pop("nonexistent-key") == []
+
+
+def test_DocumentCapture_skips_documents_without_content():
+    """DocumentCapture only stores documents that have non-empty content."""
+    from haystack import Document
+
+    component = DocumentCapture()
+    docs = [Document(content="Has content"), Document(content=None), Document(content="")]
+
+    component.run(documents=docs, result_id="test-id-2")
+
+    retrieved = DocumentCapture.pop("test-id-2")
+    assert retrieved == ["Has content"]
+
+
+def test_DocumentCapture_concurrent_isolation():
+    """Multiple concurrent pipeline runs store and retrieve documents without cross-contamination.
+
+    Verifies the thread-safety contract stated in components.py:1-10: pipeline components
+    are shared across threads and must be safe for concurrent use.
+    """
+    import concurrent.futures
+    import time
+
+    from haystack import Document
+
+    component = DocumentCapture()
+    n_runs = 20
+
+    def run_and_pop(run_id: str) -> list[str]:
+        docs = [Document(content=f"{run_id}_doc_{i}") for i in range(5)]
+        component.run(documents=docs, result_id=run_id)
+        # Brief sleep to increase thread interleaving
+        time.sleep(0.01)
+        return DocumentCapture.pop(run_id)
+
+    errors: list[Exception] = []
+    results: dict[str, list[str]] = {}
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=n_runs) as executor:
+        future_to_run_id = {
+            executor.submit(run_and_pop, f"concurrent-run-{i}"): f"concurrent-run-{i}"
+            for i in range(n_runs)
+        }
+        for future in concurrent.futures.as_completed(future_to_run_id):
+            run_id = future_to_run_id[future]
+            try:
+                results[run_id] = future.result()
+            except Exception as e:
+                errors.append(e)
+
+    assert not errors, f"Concurrent runs raised errors: {errors}"
+    assert len(results) == n_runs
+
+    for i in range(n_runs):
+        run_id = f"concurrent-run-{i}"
+        expected = [f"{run_id}_doc_{j}" for j in range(5)]
+        assert results[run_id] == expected, f"{run_id} got wrong docs: {results[run_id]}"
+
+    # All entries should have been popped — storage is clean
+    for i in range(n_runs):
+        assert DocumentCapture.pop(f"concurrent-run-{i}") == []
+
+
 def test_RemoveResourcesForEmail_mixed_valid_invalid_exclusions(caplog):
     """Test RemoveResourcesForEmail with mix of valid and invalid exclusion names."""
     import logging