FEAT: add prompt suffix handling, restore ActionPlan streaming (#145)

michelle-hadfield-nava · Copilot · yoomlam · web-flow · commit e0c3541c6b9d · 2026-01-14T14:28:29.000-07:00
Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
Co-authored-by: Yoom Lam &lt;yoom@navapbc.com&gt;
diff --git a/app/src/app_config.py b/app/src/app_config.py
@@ -39,8 +39,9 @@ def db_session(self) -> db.Session:
     # so they are not unique across different Phoenix instances.
     PROMPT_VERSIONS: dict = {
         "extract_supports": "UHJvbXB0VmVyc2lvbjo0Ng==",
-        "generate_referrals": "UHJvbXB0VmVyc2lvbjo3NA==",
-        "generate_action_plan": "UHJvbXB0VmVyc2lvbjo1Mg==",
+        "generate_referrals": "UHJvbXB0VmVyc2lvbjo3NA==",  # if no suffix, the default Austin area prompt will be used
+        "generate_referrals_keystone": "UHJvbXB0VmVyc2lvbjo3Mg==",
+        "generate_action_plan": "UHJvbXB0VmVyc2lvbjo3NQ==",
         "crawl_gcta": "UHJvbXB0VmVyc2lvbjozNg==",
         "crawl_indeed": "UHJvbXB0VmVyc2lvbjozNA==",
     }
diff --git a/app/src/common/components.py b/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 = """\
diff --git a/app/src/common/haystack_utils.py b/app/src/common/haystack_utils.py
@@ -4,16 +4,19 @@
 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
 
 from src.common import phoenix_utils
 
 
-def get_phoenix_prompt(prompt_name: str, prompt_version_id: str = "") -> list[ChatMessage]:
-    prompt_ver = phoenix_utils.get_prompt_template(prompt_name, prompt_version_id)
+def get_phoenix_prompt(
+    prompt_name: str, prompt_version_id: str = "", suffix: str = ""
+) -> list[ChatMessage]:
+    full_prompt_name = f"{prompt_name}_{suffix}" if suffix else prompt_name
+    prompt_ver = phoenix_utils.get_prompt_template(full_prompt_name, prompt_version_id)
     return to_chat_messages(prompt_ver._template["messages"])
 
 
@@ -74,12 +77,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 +124,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:
diff --git a/app/src/pipelines/generate_action_plan/pipeline_wrapper.py b/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):
@@ -48,11 +46,9 @@ def setup(self) -> None:
         pipeline = Pipeline()
         pipeline.add_component("llm", create_websearch())
 
-        prompt_template = haystack_utils.get_phoenix_prompt("generate_action_plan")
         pipeline.add_component(
             instance=ChatPromptBuilder(
-                template=prompt_template,
-                required_variables=["resources", "action_plan_json", "user_query"],
+                variables=["resources", "action_plan_json", "user_query"],
             ),
             name="prompt_builder",
         )
@@ -67,51 +63,99 @@ 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
+        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:
+        prompt_template = haystack_utils.get_phoenix_prompt("generate_action_plan")
+        return {
+            "logger": {
+                "messages_list": [
+                    {"resource_count": len(resource_objects), "user_email": user_email}
+                ],
+            },
+            "prompt_builder": {
+                "template": prompt_template,
+                "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."""
diff --git a/app/src/pipelines/generate_referrals/pipeline_wrapper.py b/app/src/pipelines/generate_referrals/pipeline_wrapper.py
@@ -100,14 +100,16 @@ def setup(self) -> None:
         self.pipeline = pipeline
 
     # Called for the `generate-referrals/run` endpoint
-    def run_api(self, query: str, user_email: str, prompt_version_id: str = "") -> dict:
+    def run_api(
+        self, query: str, user_email: str, prompt_version_id: str = "", suffix: str = ""
+    ) -> dict:
         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(query, user_email, prompt_version_id)
+                result = self._run(query, user_email, prompt_version_id, suffix)
                 span.set_input(query)
                 try:
                     resp_obj = json.loads(result["llm"]["replies"][-1].text)
@@ -117,16 +119,18 @@ def run_api(self, query: str, user_email: str, prompt_version_id: str = "") -> d
                 span.set_status(Status(StatusCode.OK))
                 return result
 
-    def _run(self, query: str, user_email: str, prompt_version_id: str = "") -> dict:
-        # Retrieve the requested prompt_version_id and error if requested prompt version is not found
+    def _run(
+        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
+                "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}' could not be retrieved due to HTTP status {he.response.status_code}",
+                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
 
         try:
diff --git a/app/src/pipelines/sample_pipeline/pipeline_wrapper.py b/app/src/pipelines/sample_pipeline/pipeline_wrapper.py
