FEAT: added streaming to generate referrals pipelines (#153)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: michelle-hadfield-nava

app/src/common/haystack_utils.py

@@ -1,5 +1,4 @@
 import logging
-import uuid
 from typing import Any, Callable, Generator, Sequence
 
 import hayhooks
@@ -68,14 +67,15 @@ def to_chat_messages(
     return messages
 
 
-def create_result_id_hook(pipeline: Pipeline) -> Callable[[dict], Generator]:
-    """Creates a pregenerator hook that generates result_id and yields it as first chunk.
+def create_result_id_hook(pipeline: Pipeline, result_id: str) -> Callable[[dict], Generator]:
+    """Creates a generator hook that yields the result_id as the last chunk.
 
     This hook is specific to pipelines that use the SaveResult component and need to
     return the result_id to the frontend for caching/reference.
 
     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
 
     Raises:
         ValueError: If the pipeline does not have a SaveResult component
@@ -95,14 +95,7 @@ def create_result_id_hook(pipeline: Pipeline) -> Callable[[dict], Generator]:
         )
 
     def hook(pipeline_run_args: dict) -> Generator:
-        result_id = str(uuid.uuid4())
-
-        # Add to pipeline args for SaveResult component
-        if save_result_component_name not in pipeline_run_args:  # checking to prevent KeyError
-            pipeline_run_args[save_result_component_name] = {}
-        pipeline_run_args[save_result_component_name]["result_id"] = result_id
-
-        # Yield as first chunk for frontend
+        # Yield result_id as last chunk for frontend
         yield StreamingChunk(content=f'{{"result_id": "{result_id}"}}\n')
 
     return hook
@@ -127,7 +120,7 @@ def stream_response(
         input_: Any | None = None,
         shorten_output: Callable[[str], str] = lambda resp: resp,
         parent_span_name_suffix: str | None = None,
-        pregenerator_hook: Callable[[dict], Generator] | None = None,
+        generator_hook: Callable[[dict], Generator] | None = None,
     ) -> Generator:
         # Must set using attributes and metadata tracer context before calling tracer.start_as_current_span()
         with using_attributes(user_id=user_id, metadata=metadata):
@@ -141,11 +134,6 @@ def stream_response(
                 try:
                     span.set_input(input_)
 
-                    # Call pregenerator_hook if provided (inside span, before streaming)
-                    if pregenerator_hook:
-                        for chunk in pregenerator_hook(pipeline_run_args):
-                            yield chunk
-
                     # hayhooks.streaming_generator() creates a thread that now inherits OpenTelemetry context
                     # thanks to ThreadingInstrumentor enabled in phoenix_utils.py
                     # streaming_generator() must be run inside the span context to inherit correctly
@@ -165,6 +153,12 @@ def stream_response(
                         # Must yield in the tracer span context to have spans linked as child spans
                         yield chunk
 
+                    # Call generator_hook if provided (inside span, after streaming)
+                    # We do this after streaming all chunks so that the parent-child span relationship is established
+                    if generator_hook:
+                        for chunk in generator_hook(pipeline_run_args):
+                            yield chunk
+
                     logger.info("Successfully streamed %d chunks", chunk_count)
                     response_text = "".join(full_response) if full_response else ""
                     span.set_output(shorten_output(response_text))

app/src/pipelines/generate_action_plan/pipeline_wrapper.py

@@ -1,4 +1,5 @@
 import logging
+import uuid
 from pprint import pformat
 from typing import Generator
 
@@ -148,13 +149,18 @@ def run_chat_completion(self, model: str, messages: list, body: dict) -> Generat
             reasoning_effort=body.get("reasoning_effort", None),
             streaming=True,
         )
+
+        # 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 action plan: %s", pipeline_run_args)
         return self.runner.stream_response(
             pipeline_run_args,
             user_id=user_email,
             metadata={"user_id": user_email},
             input_=[r.name for r in resource_objects],
-            pregenerator_hook=haystack_utils.create_result_id_hook(self.pipeline),
+            generator_hook=haystack_utils.create_result_id_hook(self.pipeline, result_id),
         )
 
 

app/src/pipelines/generate_referrals/pipeline_wrapper.py

@@ -1,15 +1,15 @@
 import json
 import logging
+import uuid
 from enum import Enum
 from pprint import pformat
-from typing import Optional
+from typing import Generator, Optional
 
 import httpx
 from fastapi import HTTPException
 from hayhooks import BasePipelineWrapper
 from haystack import Pipeline
 from haystack.components.builders import ChatPromptBuilder
-from haystack.dataclasses.chat_message import ChatMessage
 from pydantic import BaseModel
 
 from src.app_config import config
@@ -103,17 +103,13 @@ def run_api(
         self, query: str, user_email: str, prompt_version_id: str = "", suffix: str = ""
     ) -> dict:
         # Retrieve the requested prompt (with optional prompt_version_id and/or suffix)
-        try:
-            prompt_template = haystack_utils.get_phoenix_prompt(
-                "generate_referrals", prompt_version_id=prompt_version_id, suffix=suffix
-            )
-        except httpx.HTTPStatusError as he:
-            raise HTTPException(
-                status_code=422,
-                detail=f"The requested prompt version '{prompt_version_id}' with suffix '{suffix}' could not be retrieved due to HTTP status {he.response.status_code}",
-            ) from he
-        pipeline_run_args = self._run_arg_data(
-            query, user_email, prompt_template, region=suffix or "centraltx"
+
+        pipeline_run_args = self.create_pipeline_args(
+            query,
+            user_email,
+            prompt_version_id=prompt_version_id,
+            suffix=suffix,
+            region=suffix or "centraltx",
         )
 
         def extract_output(result: dict) -> list | str:
@@ -135,9 +131,29 @@ def extract_output(result: dict) -> list | str:
         logger.debug("Results: %s", pformat(response, width=160))
         return response
 
-    def _run_arg_data(
-        self, query: str, user_email: str, prompt_template: list[ChatMessage], *, region: str
+    def create_pipeline_args(
+        self,
+        query: str,
+        user_email: str,
+        *,
+        region: str,
+        prompt_version_id: str = "",
+        suffix: str = "",
+        llm_model: str | None = None,
+        reasoning_effort: str | None = None,
+        streaming: bool = False,
     ) -> dict:
+        """Create pipeline run arguments with optional overrides for model, reasoning effort, and streaming."""
+        try:
+            prompt_template = haystack_utils.get_phoenix_prompt(
+                "generate_referrals", prompt_version_id=prompt_version_id, suffix=suffix
+            )
+        except httpx.HTTPStatusError as e:
+            raise HTTPException(
+                status_code=422,
+                detail=f"The requested prompt version '{prompt_version_id}' with suffix '{suffix}' could not be retrieved",
+            ) from e
+
         return {
             "logger": {
                 "messages_list": [{"query": query, "user_email": user_email}],
@@ -148,7 +164,49 @@ def _run_arg_data(
                 "response_json": response_schema,
             },
             "llm": {
-                "model": config.generate_referrals_model_version,
-                "reasoning_effort": config.generate_referrals_reasoning_level,
+                "model": llm_model or config.generate_referrals_model_version,
+                "reasoning_effort": reasoning_effort or config.generate_referrals_reasoning_level,
+                "streaming": streaming,
             },
         }
+
+    # https://docs.haystack.deepset.ai/docs/hayhooks#openai-compatibility
+    # Called for the `{pipeline_name}/chat`, `/chat/completions`, or `/v1/chat/completions` streaming endpoint using Server-Sent Events (SSE)
+    def run_chat_completion(self, model: str, messages: list, body: dict) -> Generator:
+        # Note: 'model' parameter is the pipeline name, not the LLM model
+        assert model == self.name, f"Unexpected model/pipeline name: {model}"
+
+        # Extract custom parameters from the body
+        query = body.get("query", "")
+        user_email = body.get("user_email", "")
+        suffix = body.get("suffix", "")
+
+        if not query:
+            raise ValueError("query parameter is required")
+
+        if not user_email:
+            raise ValueError("user_email parameter is required")
+
+        pipeline_run_args = self.create_pipeline_args(
+            query,
+            user_email,
+            prompt_version_id=body.get("prompt_version_id", ""),
+            suffix=body.get("suffix", ""),
+            region=suffix or "centraltx",
+            llm_model=body.get("llm_model", None),
+            reasoning_effort=body.get("reasoning_effort", None),
+            streaming=True,
+        )
+
+        # 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,
+            user_id=user_email,
+            metadata={"user_id": user_email},
+            input_=[query],
+            generator_hook=haystack_utils.create_result_id_hook(self.pipeline, result_id),
+        )

app/src/pipelines/generate_referrals_rag/pipeline_wrapper.py

@@ -4,7 +4,6 @@
 from haystack.components.builders import ChatPromptBuilder
 from haystack.components.converters import OutputAdapter
 from haystack.components.embedders import SentenceTransformersTextEmbedder
-from haystack.dataclasses.chat_message import ChatMessage
 from haystack_integrations.components.retrievers.chroma import ChromaEmbeddingRetriever
 
 from src.app_config import config
@@ -76,10 +75,33 @@ def _create_pipeline(self) -> Pipeline:
         # pipeline.draw(path="generate_referrals_rag.png")
         return pipeline
 
-    def _run_arg_data(
-        self, query: str, user_email: str, prompt_template: list[ChatMessage], *, region: str
+    def create_pipeline_args(
+        self,
+        query: str,
+        user_email: str,
+        *,
+        prompt_version_id: str = "",
+        suffix: str = "",
+        region: str,
+        llm_model: str | None = None,
+        reasoning_effort: str | None = None,
+        streaming: bool = False,
     ) -> dict:
-        return super()._run_arg_data(query, user_email, prompt_template, region=region) | {
+        """Create pipeline run arguments with optional overrides for model, reasoning effort, and streaming."""
+        # Get base args from parent class
+        base_args = super().create_pipeline_args(
+            query,
+            user_email,
+            prompt_version_id=prompt_version_id,
+            suffix=suffix,
+            region=region,
+            llm_model=llm_model,
+            reasoning_effort=reasoning_effort,
+            streaming=streaming,
+        )
+
+        # Override/add RAG-specific args
+        return base_args | {
             # For querying RAG DB
             "query_embedder": {"text": query},
             "retriever": {
@@ -88,7 +110,9 @@ def _run_arg_data(
             },
             # Override LLM config for RAG pipeline
             "llm": {
-                "model": config.generate_referrals_rag_model_version,
-                "reasoning_effort": config.generate_referrals_rag_reasoning_level,
+                "model": llm_model or config.generate_referrals_rag_model_version,
+                "reasoning_effort": reasoning_effort
+                or config.generate_referrals_rag_reasoning_level,
+                "streaming": streaming,
             },
         }

frontend/src/util/fetchActionPlan.ts

@@ -285,6 +285,7 @@ export async function fetchActionPlanStreaming(
 
     let buffer = "";
     let accumulatedJSON = ""; // Accumulate the full JSON response
+    let lastChunkContent = ""; // Track the last chunk to extract result_id
 
     while (true) {
       const { done, value } = await reader.read();
@@ -340,17 +341,8 @@ export async function fetchActionPlanStreaming(
             // Handle hayhooks response format: choices[0].delta.content
             const content = parsed.choices?.[0]?.delta?.content;
             if (content) {
-              // Check if this is the first message containing result_id
-              if (!resultId && content.includes("result_id")) {
-                try {
-                  const resultIdData = JSON.parse(content);
-                  if (resultIdData.result_id) {
-                    resultId = resultIdData.result_id;
-                  }
-                } catch (e) {
-                  // Not a JSON object, continue processing as regular content
-                }
-              }
+              // Store the last chunk content for result_id extraction
+              lastChunkContent = content;
 
               // Only accumulate content that's not the result_id metadata
               if (!content.includes("result_id")) {
@@ -365,6 +357,23 @@ export async function fetchActionPlanStreaming(
 
             // Check for finish_reason to detect completion
             if (parsed.choices?.[0]?.finish_reason === "stop") {
+              // Extract result_id from the last chunk (the chunk before this stop message)
+              if (!resultId && lastChunkContent.includes("result_id")) {
+                try {
+                  const resultIdData = JSON.parse(lastChunkContent) as {
+                    result_id?: string;
+                  };
+                  if (resultIdData.result_id) {
+                    resultId = resultIdData.result_id;
+                  }
+                } catch (e) {
+                  console.error(
+                    "Failed to extract result_id from last chunk:",
+                    e,
+                  );
+                }
+              }
+
               clearTimeout(timer);
               onComplete();
               break;