feat: Add name for Phoenix traces (#99)

yoomlam · web-flow · commit 2cddebc048d2 · 2025-11-05T21:33:13.000-06:00
diff --git a/app/src/common/phoenix_utils.py b/app/src/common/phoenix_utils.py
@@ -6,6 +6,7 @@
 import opentelemetry.exporter.otlp.proto.http.trace_exporter as otel_trace_exporter
 import phoenix.otel
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.trace import NoOpTracerProvider, TracerProvider
 
 # https://docs.arize.com/phoenix/tracing/integrations-tracing/haystack
 # Arize's Phoenix observability platform
@@ -38,6 +39,9 @@ def service_alive() -> bool:
     return False
 
 
+tracer_provider: TracerProvider = NoOpTracerProvider()
+
+
 def configure_phoenix(only_if_alive: bool = True) -> None:
     "Set only_if_alive=True to fail fast if Phoenix is not reachable."
     if only_if_alive and not service_alive():
@@ -54,6 +58,7 @@ def configure_phoenix(only_if_alive: bool = True) -> None:
     logger.info("Using phoenix.otel.register with batch_otel=%s", config.batch_otel)
     # This uses PHOENIX_COLLECTOR_ENDPOINT and PHOENIX_PROJECT_NAME env variables
     # and PHOENIX_API_KEY to handle authentication to Phoenix.
+    global tracer_provider
     tracer_provider = phoenix.otel.register(
         endpoint=trace_endpoint,
         batch=config.batch_otel,
diff --git a/app/src/pipelines/email_result/pipeline_wrapper.py b/app/src/pipelines/email_result/pipeline_wrapper.py
@@ -5,11 +5,13 @@
 from hayhooks import BasePipelineWrapper
 from haystack import Pipeline
 from haystack.core.errors import PipelineRuntimeError
-from openinference.instrumentation import using_metadata
+from openinference.instrumentation import _tracers, using_metadata
+from opentelemetry.trace.status import Status, StatusCode
 
-from src.common import components
+from src.common import components, phoenix_utils
 
 logger = logging.getLogger(__name__)
+tracer = phoenix_utils.tracer_provider.get_tracer(__name__)
 
 
 class PipelineWrapper(BasePipelineWrapper):
@@ -28,34 +30,46 @@ def setup(self) -> None:
 
     def run_api(self, result_id: str, email: str) -> dict:
         with using_metadata({"email": email}):
-            try:
-                response = self.pipeline.run(
-                    {
-                        "logger": {
-                            "messages_list": [{"result_id": result_id, "email": email}],
-                        },
-                        "load_result": {
-                            "result_id": result_id,
-                        },
-                        "email_result": {
-                            "email": 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(result_id, email)
+                span.set_input(result_id)
+                span.set_output(result["email_result"]["status"])
+                span.set_status(Status(StatusCode.OK))
+                return result
+
+    def _run(self, result_id: str, email: str) -> dict:
+        try:
+            response = self.pipeline.run(
+                {
+                    "logger": {
+                        "messages_list": [{"result_id": result_id, "email": email}],
                     },
-                    include_outputs_from={"email_result"},
-                )
-                logger.debug("Results: %s", pformat(response, width=160))
-                return response
-            except PipelineRuntimeError as re:
-                error_msg = str(re)
-                if re.component_type == components.LoadResult:
-                    if "Invalid JSON format in result" in error_msg:
-                        status_code = 500  # Internal error
-                    else:
-                        status_code = 400  # User error
-                else:
+                    "load_result": {
+                        "result_id": result_id,
+                    },
+                    "email_result": {
+                        "email": email,
+                    },
+                },
+                include_outputs_from={"email_result"},
+            )
+            logger.debug("Results: %s", pformat(response, width=160))
+            return response
+        except PipelineRuntimeError as re:
+            error_msg = str(re)
+            if re.component_type == components.LoadResult:
+                if "Invalid JSON format in result" in error_msg:
                     status_code = 500  # Internal error
+                else:
+                    status_code = 400  # User error
+            else:
+                status_code = 500  # Internal error
 
-                raise HTTPException(
-                    status_code=status_code,
-                    detail=f"Error occurred: {error_msg}",
-                ) from re
+            raise HTTPException(
+                status_code=status_code,
+                detail=f"Error occurred: {error_msg}",
+            ) from re
diff --git a/app/src/pipelines/generate_action_plan/pipeline_wrapper.py b/app/src/pipelines/generate_action_plan/pipeline_wrapper.py
@@ -4,14 +4,16 @@
 from hayhooks import BasePipelineWrapper
 from haystack import Pipeline
 from haystack.components.builders import ChatPromptBuilder
-from openinference.instrumentation import using_attributes, using_metadata
+from openinference.instrumentation import _tracers, using_attributes, using_metadata
+from opentelemetry.trace.status import Status, StatusCode
 from pydantic import BaseModel
 
-from src.common import haystack_utils
+from src.common import haystack_utils, phoenix_utils
 from src.common.components import OpenAIWebSearchGenerator, ReadableLogger
 from src.pipelines.generate_referrals.pipeline_wrapper import Resource
 
 logger = logging.getLogger(__name__)
+tracer = phoenix_utils.tracer_provider.get_tracer(__name__)
 
 
 class ActionPlan(BaseModel):
@@ -59,23 +61,35 @@ def run_api(self, resources: list[Resource] | list[dict], user_email: str) -> di
         resource_objects = get_resources(resources)
 
         with using_attributes(user_id=user_email), using_metadata({"user_id": user_email}):
-            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,
-                    },
-                    "llm": {"model": "gpt-5-mini", "reasoning_effort": "low"},
+            # 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)
+                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) -> dict:
+        response = self.pipeline.run(
+            {
+                "logger": {
+                    "messages_list": [
+                        {"resource_count": len(resource_objects), "user_email": user_email}
+                    ],
                 },
-                include_outputs_from={"llm"},
-            )
-            logger.debug("Results: %s", pformat(response, width=160))
-            return {"response": response["llm"]["replies"][0]._content[0].text}
+                "prompt_builder": {
+                    "resources": format_resources(resource_objects),
+                    "action_plan_json": action_plan_as_json,
+                },
+                "llm": {"model": "gpt-5-mini", "reasoning_effort": "low"},
+            },
+            include_outputs_from={"llm"},
+        )
+        logger.debug("Results: %s", pformat(response, width=160))
+        return {"response": response["llm"]["replies"][0]._content[0].text}
 
 
 def get_resources(resources: list[Resource] | list[dict]) -> list[Resource]:
diff --git a/app/src/pipelines/generate_referrals/pipeline_wrapper.py b/app/src/pipelines/generate_referrals/pipeline_wrapper.py
@@ -1,3 +1,4 @@
+import json
 import logging
 from enum import Enum
 from pprint import pformat
@@ -9,12 +10,14 @@
 from haystack import Pipeline
 from haystack.components.builders import ChatPromptBuilder
 from haystack.core.errors import PipelineRuntimeError
-from openinference.instrumentation import using_attributes, using_metadata
+from openinference.instrumentation import _tracers, using_attributes, using_metadata
+from opentelemetry.trace.status import Status, StatusCode
 from pydantic import BaseModel
 
-from src.common import components, haystack_utils
+from src.common import components, haystack_utils, phoenix_utils
 
 logger = logging.getLogger(__name__)
+tracer = phoenix_utils.tracer_provider.get_tracer(__name__)
 
 
 class ReferralType(str, Enum):
@@ -97,37 +100,53 @@ def setup(self) -> None:
     # Called for the `generate-referrals/run` endpoint
     def run_api(self, query: str, user_email: str, prompt_version_id: str = "") -> dict:
         with using_attributes(user_id=user_email), using_metadata({"user_id": user_email}):
-            # Retrieve the requested prompt_version_id and error if requested prompt version is not found
-            try:
-                prompt_template = haystack_utils.get_phoenix_prompt(
-                    "generate_referrals", prompt_version_id
-                )
-            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}",
-                ) from he
-
-            try:
-                response = self.pipeline.run(
-                    {
-                        "logger": {
-                            "messages_list": [{"query": query, "user_email": user_email}],
-                        },
-                        "prompt_builder": {
-                            "template": prompt_template,
-                            "query": query,
-                            "response_json": response_schema,
-                        },
-                        "llm": {"model": "gpt-5-mini", "reasoning_effort": "low"},
+            # 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)
+                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 = "") -> dict:
+        # Retrieve the requested prompt_version_id and error if requested prompt version is not found
+        try:
+            prompt_template = haystack_utils.get_phoenix_prompt(
+                "generate_referrals", prompt_version_id
+            )
+        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}",
+            ) from he
+
+        try:
+            response = self.pipeline.run(
+                {
+                    "logger": {
+                        "messages_list": [{"query": query, "user_email": user_email}],
                     },
-                    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
+                    "prompt_builder": {
+                        "template": prompt_template,
+                        "query": query,
+                        "response_json": response_schema,
+                    },
+                    "llm": {"model": "gpt-5-mini", "reasoning_effort": "low"},
+                },
+                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
diff --git a/app/src/pipelines/generate_referrals_from_doc/pipeline_wrapper.py b/app/src/pipelines/generate_referrals_from_doc/pipeline_wrapper.py
@@ -1,3 +1,4 @@
+import json
 import logging
 from pprint import pformat
 from typing import List, Optional
@@ -7,16 +8,18 @@
 from haystack import Pipeline
 from haystack.components.builders import ChatPromptBuilder
 from haystack.components.converters import OutputAdapter, PyPDFToDocument
-from openinference.instrumentation import using_metadata
+from openinference.instrumentation import _tracers, using_metadata
+from opentelemetry.trace.status import Status, StatusCode
 
-from src.common import components, haystack_utils
+from src.common import components, haystack_utils, phoenix_utils
 from src.pipelines.generate_referrals.pipeline_wrapper import response_schema
 
 logger = logging.getLogger(__name__)
+tracer = phoenix_utils.tracer_provider.get_tracer(__name__)
 
 
 class PipelineWrapper(BasePipelineWrapper):
-    name = "generate_referrals_from_document"
+    name = "generate_referrals_from_doc"
 
     def setup(self) -> None:
         pipeline = Pipeline()
@@ -67,18 +70,34 @@ def run_api(self, user_email: str, files: Optional[List[UploadFile]] = None) ->
             raise HTTPException(status_code=400, detail="No files provided for processing.")
 
         with using_metadata({"user_id": user_email}):
-            response = self.pipeline.run(
-                {
-                    "logger": {
-                        "messages_list": [{"filenames": [file.filename for file in files]}],
-                    },
-                    "files_to_bytestreams": {"files": files},
-                    "prompt_builder": {
-                        "response_json": response_schema,
-                    },
-                    "llm": {"model": "gpt-5-mini", "reasoning_effort": "low"},
+            # 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(files)
+                span.set_input([file.filename for file in files])
+                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, files: List[UploadFile]) -> dict:
+        response = self.pipeline.run(
+            {
+                "logger": {
+                    "messages_list": [{"filenames": [file.filename for file in files]}],
+                },
+                "files_to_bytestreams": {"files": files},
+                "prompt_builder": {
+                    "response_json": response_schema,
                 },
-                include_outputs_from={"llm"},
-            )
-            logger.info("Pipeline result: %s", pformat(response, width=160))
-            return response
+                "llm": {"model": "gpt-5-mini", "reasoning_effort": "low"},
+            },
+            include_outputs_from={"llm"},
+        )
+        logger.info("Pipeline result: %s", pformat(response, width=160))
+        return response
diff --git a/app/src/pipelines/hello_bedrock/pipeline_wrapper.py b/app/src/pipelines/hello_bedrock/pipeline_wrapper.py
@@ -10,7 +10,6 @@
 
 logger = logging.getLogger(__name__)
 
-
 system_prompt = (
     "Your role is to say hello to the name provided by the user, if no name is found politely inform the user."
     "Assure them any PII is handled securely in AWS Bedrock. You should only greet the user, do not respond "
diff --git a/app/src/pipelines/sample_pipeline/pipeline_wrapper.py b/app/src/pipelines/sample_pipeline/pipeline_wrapper.py
