FEAT: stream action plan response from backend (#124)

Co-authored-by: Yoom Lam <yoom@navapbc.com>

Author: michelle-hadfield-nava

app/src/common/components.py

@@ -20,6 +20,7 @@
 from haystack.core.component.types import Variadic
 from haystack.dataclasses.byte_stream import ByteStream
 from haystack.dataclasses.chat_message import ChatMessage
+from haystack.dataclasses.streaming_chunk import StreamingChunk
 from openai import OpenAI
 from pydantic import BaseModel, ValidationError
 
@@ -187,13 +188,22 @@ def run(self, result_id: str) -> dict:
 class OpenAIWebSearchGenerator:
     """Searches the web using OpenAI's web search capabilities and generates a response."""
 
+    def __init__(self) -> None:
+        """
+        Initialize the OpenAI web search generator.
+        """
+
+        # Declare this attribute so it can be set when streaming_generator() is called
+        self.streaming_callback: Callable | None = None
+
     @component.output_types(replies=List[ChatMessage])
     def run(
         self,
         messages: list[ChatMessage],
         domain: str | None = None,
         model: str = config.default_openai_model_version,
         reasoning_effort: str = config.default_openai_reasoning_level,
+        streaming: bool = False,
     ) -> dict:
         """
         Run the OpenAI web search generator.
@@ -207,10 +217,11 @@ def run(
         """
 
         logger.info(
-            "Calling OpenAI API with web_search, model=%s, domain=%s, reasoning_effort=%s",
+            "Calling OpenAI API with web_search, model=%s, domain=%s, reasoning_effort=%s, streaming=%s",
             model,
             domain,
             reasoning_effort,
+            streaming,
         )
 
         assert len(messages) == 1
@@ -228,11 +239,72 @@ def run(
             api_params["tools"][0]["filters"] = {"allowed_domains": [domain]}
 
         client = OpenAI()
-        response = client.responses.create(**api_params)
 
-        logger.debug("Response: %s", pformat(response.output_text, width=160))
+        # Use streaming if callback is provided
+        if streaming:
+            logger.info(
+                "Starting OpenAI streaming request (model=%s, reasoning_effort=%s)",
+                model,
+                reasoning_effort,
+            )
+            api_params["stream"] = True
 
-        return {"replies": [ChatMessage.from_assistant(response.output_text)]}
+            try:
+                response = client.responses.create(**api_params)
+            except Exception as e:
+                logger.error("Failed to create OpenAI stream: %s", e, exc_info=True)
+                raise
+
+            # Collect full response while streaming
+            full_text = ""
+            chunk_count = 0
+
+            try:
+                for openai_chunk in response:
+                    chunk_count += 1
+                    chunk_text = ""
+
+                    # Extract text from OpenAI Responses API events
+                    if hasattr(openai_chunk, "type"):
+                        # Check delta attribute (for text delta events)
+                        if not chunk_text and hasattr(openai_chunk, "delta"):
+                            delta = openai_chunk.delta
+                            if isinstance(delta, str):
+                                chunk_text = delta
+                            elif isinstance(delta, list):
+                                chunk_text = "".join(str(item) for item in delta)
+                            elif hasattr(delta, "content"):
+                                chunk_text = delta.content or ""
+                            elif hasattr(delta, "text"):
+                                chunk_text = delta.text or ""
+
+                    # Fallback for non-Responses API format
+                    if not chunk_text and hasattr(openai_chunk, "output_text"):
+                        chunk_text = openai_chunk.output_text or ""
+
+                    if chunk_text:
+                        full_text += chunk_text
+                        # Convert to Haystack StreamingChunk and call the callback
+                        streaming_chunk = StreamingChunk(content=chunk_text)
+                        assert (
+                            self.streaming_callback is not None
+                        ), "Expected streaming_callback to be set by Hayhooks"
+                        self.streaming_callback(streaming_chunk)
+
+            except Exception as e:
+                logger.error("Error during streaming: %s", e, exc_info=True)
+                raise
+
+            logger.info("Streaming complete: %d chunks, %d characters", chunk_count, len(full_text))
+            if not full_text:
+                logger.warning("No text collected during streaming")
+
+            return {"replies": [ChatMessage.from_assistant(full_text)]}
+        else:
+            # Non-streaming response
+            response = client.responses.create(**api_params)
+            logger.debug("Response: %s", pformat(response.output_text, width=160))
+            return {"replies": [ChatMessage.from_assistant(response.output_text)]}
 
 
 EMAIL_INTRO = """\

app/src/common/haystack_utils.py

@@ -4,7 +4,7 @@
 import hayhooks
 from haystack import Pipeline
 from haystack.dataclasses.chat_message import ChatMessage
-from openinference.instrumentation import using_metadata
+from openinference.instrumentation import using_attributes
 from opentelemetry.trace import Span
 from opentelemetry.trace.status import Status, StatusCode
 from phoenix.client.__generated__ import v1
@@ -74,12 +74,13 @@ def stream_response(
         self,
         pipeline_run_args: dict,
         *,
+        user_id: str,
         metadata: dict[str, Any],
         input_: Any | None = None,
         shorten_output: Callable[[str], str] = lambda resp: resp,
     ) -> Generator:
-        with using_metadata(metadata):
-            # Must set using_metadata context before calling tracer.start_as_current_span()
+        # Must set using attributes and metadata tracer context before calling tracer.start_as_current_span()
+        with using_attributes(user_id=user_id, metadata=metadata):
             with phoenix_utils.tracer().start_as_current_span(  # pylint: disable=not-context-manager,unexpected-keyword-arg
                 self.parent_span_name, openinference_span_kind="chain"
             ) as span:
@@ -120,13 +121,14 @@ def return_response(
         self,
         pipeline_run_args: dict,
         *,
+        user_id: str,
         metadata: dict[str, Any],
         input_: Any | None = None,
         include_outputs_from: set[str] | None = None,
         extract_output: Callable[[Any], Any] = lambda resp: resp,
     ) -> dict:
         # Must set using_metadata context before calling tracer.start_as_current_span()
-        with using_metadata(metadata):
+        with using_attributes(user_id=user_id, metadata=metadata):
             with phoenix_utils.tracer().start_as_current_span(  # pylint: disable=not-context-manager,unexpected-keyword-arg
                 self.parent_span_name, openinference_span_kind="chain"
             ) as span:

app/src/pipelines/generate_action_plan/pipeline_wrapper.py

@@ -1,15 +1,14 @@
 import logging
 from pprint import pformat
+from typing import Generator
 
 from hayhooks import BasePipelineWrapper
 from haystack import Pipeline
 from haystack.components.builders import ChatPromptBuilder
-from openinference.instrumentation import _tracers, using_attributes, using_metadata
-from opentelemetry.trace.status import Status, StatusCode
 from pydantic import BaseModel
 
 from src.app_config import config
-from src.common import haystack_utils, phoenix_utils
+from src.common import haystack_utils
 from src.common.components import (
     LlmOutputValidator,
     OpenAIWebSearchGenerator,
@@ -19,7 +18,6 @@
 from src.pipelines.generate_referrals.pipeline_wrapper import Resource
 
 logger = logging.getLogger(__name__)
-tracer = phoenix_utils.tracer_provider.get_tracer(__name__)
 
 
 class ActionPlan(BaseModel):
@@ -67,51 +65,94 @@ def setup(self) -> None:
         pipeline.connect("llm", "logger")
 
         self.pipeline = pipeline
+        self.runner = haystack_utils.TracedPipelineRunner(self.name, self.pipeline)
 
     # Called for the `generate-action-plan/run` endpoint
     def run_api(
         self, resources: list[Resource] | list[dict], user_email: str, user_query: str
     ) -> dict:
         resource_objects = get_resources(resources)
-
-        with using_attributes(user_id=user_email), using_metadata({"user_id": user_email}):
-            # Must set using_metadata context before calling tracer.start_as_current_span()
-            assert isinstance(tracer, _tracers.OITracer), f"Got unexpected {type(tracer)}"
-            with tracer.start_as_current_span(  # pylint: disable=not-context-manager,unexpected-keyword-arg
-                self.name, openinference_span_kind="chain"
-            ) as span:
-                result = self._run(resource_objects, user_email, user_query)
-                span.set_input([r.name for r in resource_objects])
-                span.set_output(result["response"])
-                span.set_status(Status(StatusCode.OK))
-                return result
-
-    def _run(self, resource_objects: list[Resource], user_email: str, user_query: str) -> dict:
-        response = self.pipeline.run(
-            {
-                "logger": {
-                    "messages_list": [
-                        {"resource_count": len(resource_objects), "user_email": user_email}
-                    ],
-                },
-                "prompt_builder": {
-                    "resources": format_resources(resource_objects),
-                    "action_plan_json": action_plan_as_json,
-                    "user_query": user_query,
-                },
-                "llm": {
-                    "model": config.generate_action_plan_model_version,
-                    "reasoning_effort": config.generate_action_plan_reasoning_level,
-                },
-            },
+        pipeline_run_args = self.create_pipeline_args(
+            user_email,
+            resource_objects,
+            user_query,
+        )
+        response = self.runner.return_response(
+            pipeline_run_args,
+            user_id=user_email,
+            metadata={"user_id": user_email},
             include_outputs_from={"llm", "save_result"},
+            input_=[r.name for r in resource_objects],
+            extract_output=lambda response: response["llm"]["replies"][0]._content[0].text,
         )
         logger.debug("Results: %s", pformat(response, width=160))
+
         return {
             "response": response["llm"]["replies"][0]._content[0].text,
             "save_result": response["save_result"],
         }
 
+    def create_pipeline_args(
+        self,
+        user_email: str,
+        resource_objects: list[Resource],
+        user_query: str,
+        *,
+        llm_model: str | None = None,
+        reasoning_effort: str | None = None,
+        streaming: bool = False,
+    ) -> dict:
+        return {
+            "logger": {
+                "messages_list": [
+                    {"resource_count": len(resource_objects), "user_email": user_email}
+                ],
+            },
+            "prompt_builder": {
+                "resources": format_resources(resource_objects),
+                "action_plan_json": action_plan_as_json,
+                "user_query": user_query,
+            },
+            "llm": {
+                "model": llm_model or config.generate_action_plan_model_version,
+                "reasoning_effort": reasoning_effort or config.generate_action_plan_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
+        resources = body.get("resources", [])
+        user_email = body.get("user_email", "")
+        user_query = body.get("user_query", "")
+
+        if not resources:
+            raise ValueError("resources parameter is required")
+        if not user_email:
+            raise ValueError("user_email parameter is required")
+
+        resource_objects = get_resources(resources)
+        pipeline_run_args = self.create_pipeline_args(
+            user_email,
+            resource_objects,
+            user_query,
+            llm_model=body.get("llm_model", None),
+            reasoning_effort=body.get("reasoning_effort", None),
+            streaming=True,
+        )
+        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],
+        )
+
 
 def get_resources(resources: list[Resource] | list[dict]) -> list[Resource]:
     """Ensure we have a list of Resource objects."""

app/src/pipelines/sample_pipeline/pipeline_wrapper.py

@@ -56,6 +56,7 @@ def run_api(self, question: str) -> dict:
         }
         response = self.runner.return_response(
             pipeline_run_args,
+            user_id=user_id,
             metadata={"user_id": user_id},
             include_outputs_from={"echo_component", "echo_component2"},
             input_=question,
@@ -87,7 +88,7 @@ def run_chat_completion(self, model: str, messages: list, body: dict) -> Generat
 
         user_id = "someone@example.com"
         return self.runner.stream_response(
-            pipeline_run_args, metadata={"user_id": user_id}, input_=question
+            pipeline_run_args, user_id=user_id, metadata={"user_id": user_id}, input_=question
         )
 
     def create_pipeline_args(self, location: str, messages: list[ChatMessage]) -> dict: