feat: Add optional suffix to parent_span_name (#148)

Author: yoomlam

app/src/common/haystack_utils.py

@@ -2,7 +2,9 @@
 from typing import Any, Callable, Generator, Sequence
 
 import hayhooks
+from fastapi import HTTPException
 from haystack import Pipeline
+from haystack.core.errors import PipelineRuntimeError
 from haystack.dataclasses.chat_message import ChatMessage
 from openinference.instrumentation import using_attributes
 from opentelemetry.trace import Span
@@ -81,11 +83,15 @@ def stream_response(
         metadata: dict[str, Any],
         input_: Any | None = None,
         shorten_output: Callable[[str], str] = lambda resp: resp,
+        parent_span_name_suffix: str | 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):
             with phoenix_utils.tracer().start_as_current_span(  # pylint: disable=not-context-manager,unexpected-keyword-arg
-                self.parent_span_name, openinference_span_kind="chain"
+                f"{self.parent_span_name}--{parent_span_name_suffix}"
+                if parent_span_name_suffix
+                else self.parent_span_name,
+                openinference_span_kind="chain",
             ) as span:
                 assert isinstance(span, Span), f"Got unexpected {type(span)}"
                 try:
@@ -114,11 +120,16 @@ def stream_response(
                     response_text = "".join(full_response) if full_response else ""
                     span.set_output(shorten_output(response_text))
                     span.set_status(Status(StatusCode.OK))
+                except PipelineRuntimeError as e:
+                    logger.error("PipelineRuntimeError: %s", e, exc_info=True)
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise HTTPException(status_code=500, detail=str(e)) from e
                 except Exception as e:
                     logger.error("Error during streaming: %s", e, exc_info=True)
                     span.set_status(Status(StatusCode.ERROR, str(e)))
                     span.record_exception(e)
-                    raise
+                    raise HTTPException(status_code=500, detail=str(e)) from e
 
     def return_response(
         self,
@@ -128,12 +139,16 @@ def return_response(
         metadata: dict[str, Any],
         input_: Any | None = None,
         include_outputs_from: set[str] | None = None,
-        extract_output: Callable[[Any], Any] = lambda resp: resp,
+        extract_output: Callable[[dict], Any] = lambda resp: resp,
+        parent_span_name_suffix: str | None = None,
     ) -> dict:
         # Must set using_metadata 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"
+                f"{self.parent_span_name}--{parent_span_name_suffix}"
+                if parent_span_name_suffix
+                else self.parent_span_name,
+                openinference_span_kind="chain",
             ) as span:
                 try:
                     result = self.pipeline.run(
@@ -145,8 +160,13 @@ def return_response(
                     span.set_output(extract_output(result))
                     span.set_status(Status(StatusCode.OK))
                     return result
+                except PipelineRuntimeError as e:
+                    logger.error("PipelineRuntimeError: %s", e, exc_info=True)
+                    span.set_status(Status(StatusCode.ERROR, str(e)))
+                    span.record_exception(e)
+                    raise HTTPException(status_code=500, detail=str(e)) from e
                 except Exception as e:
                     logger.error("Error during pipeline run: %s", e, exc_info=True)
                     span.set_status(Status(StatusCode.ERROR, str(e)))
                     span.record_exception(e)
-                    raise
+                    raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}") from e

app/src/pipelines/generate_referrals/pipeline_wrapper.py

@@ -9,10 +9,7 @@
 from hayhooks import BasePipelineWrapper
 from haystack import Pipeline
 from haystack.components.builders import ChatPromptBuilder
-from haystack.core.errors import PipelineRuntimeError
 from haystack.dataclasses.chat_message import ChatMessage
-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
@@ -63,6 +60,10 @@ class PipelineWrapper(BasePipelineWrapper):
     name = "generate_referrals"
 
     def setup(self) -> None:
+        self.pipeline = self._create_pipeline()
+        self.runner = haystack_utils.TracedPipelineRunner(self.name, self.pipeline)
+
+    def _create_pipeline(self) -> Pipeline:
         # Do not rely on max_runs_per_component strictly, i.e., a component may run max_runs_per_component+1 times.
         # The component_visits counter for max_runs_per_component is reset with each call to pipeline.run()
         pipeline = Pipeline(max_runs_per_component=3)
@@ -96,31 +97,11 @@ def setup(self) -> None:
 
         pipeline.add_component("logger", components.ReadableLogger())
         pipeline.connect("output_validator.valid_replies", "logger")
-
-        self.pipeline = pipeline
+        return pipeline
 
     # Called for the `generate-referrals/run` endpoint
     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, suffix)
-                span.set_input(query)
-                try:
-                    resp_obj = json.loads(result["llm"]["replies"][-1].text)
-                    span.set_output([r["name"] for r in resp_obj["resources"]])
-                except (KeyError, IndexError):
-                    span.set_output(result["llm"]["replies"][-1].text)
-                span.set_status(Status(StatusCode.OK))
-                return result
-
-    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:
@@ -132,20 +113,26 @@ def _run(
                 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
-
-        try:
-            response = self.pipeline.run(
-                self._run_arg_data(query, user_email, prompt_template),
-                include_outputs_from={"llm", "save_result"},
-            )
-            logger.debug("Results: %s", pformat(response, width=160))
-            return response
-        except PipelineRuntimeError as re:
-            logger.error("PipelineRuntimeError: %s", re, exc_info=True)
-            raise HTTPException(status_code=500, detail=str(re)) from re
-        except Exception as e:
-            logger.error("Error %s: %s", type(e), e, exc_info=True)
-            raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}") from e
+        pipeline_run_args = self._run_arg_data(query, user_email, prompt_template)
+
+        def extract_output(result: dict) -> list | str:
+            try:
+                resp_obj = json.loads(result["llm"]["replies"][-1].text)
+                return [r["name"] for r in resp_obj["resources"]]
+            except (KeyError, IndexError):
+                return result["llm"]["replies"][-1].text
+
+        response = self.runner.return_response(
+            pipeline_run_args,
+            user_id=user_email,
+            metadata={"user_id": user_email},
+            include_outputs_from={"llm", "save_result"},
+            input_=query,
+            extract_output=extract_output,
+            parent_span_name_suffix=suffix,
+        )
+        logger.debug("Results: %s", pformat(response, width=160))
+        return response
 
     def _run_arg_data(
         self, query: str, user_email: str, prompt_template: list[ChatMessage]

app/src/pipelines/generate_referrals_rag/pipeline_wrapper.py

@@ -21,7 +21,7 @@
 class PipelineWrapper(GenerateReferralsPipelineWrapper):
     name = "generate_referrals_rag"
 
-    def setup(self) -> None:
+    def _create_pipeline(self) -> Pipeline:
         pipeline = Pipeline(max_runs_per_component=3)
 
         # Replace LoadSupports() with retrieval from vector DB
@@ -75,7 +75,7 @@ def setup(self) -> None:
         pipeline.connect("output_validator.valid_replies", "logger")
 
         # pipeline.draw(path="generate_referrals_rag.png")
-        self.pipeline = pipeline
+        return pipeline
 
     def _run_arg_data(
         self, query: str, user_email: str, prompt_template: list[ChatMessage]