feat: Enable streaming in sample_pipeline (#138)

Author: yoomlam

app/poetry.lock

@@ -25,6 +25,7 @@ haystack-ai = "^2.16.1"
 hayhooks = "1.0.1"
 openai = "^1.99.8"
 openinference-instrumentation-haystack = "^0.1.24"
+opentelemetry-instrumentation-threading = "^0.59b0"
 arize-phoenix-client = "^1.15.3"
 arize-phoenix-otel = "^0.13.0"
 amazon-bedrock-haystack = "^3.10.0"

app/pyproject.toml

@@ -23,7 +23,7 @@ class AppConfig(PydanticBaseEnvConfig):
     phoenix_collector_endpoint: str = "https://phoenix:6006"
     batch_otel: bool = True
 
-    redact_pii: bool = True
+    redact_pii: bool = False
 
     aws_ses_from_email: str = "no-reply@test.com"
 

app/src/app_config.py

@@ -366,7 +366,7 @@ def parse_json_if_possible(self, content: Any) -> Any:
             try:
                 return json.loads(content.text)
             except JSONDecodeError:
-                logger.warning("Failed to parse content as JSON: %s", content.text, exc_info=True)
+                logger.warning("Failed to parse content as JSON: %s", content.text)
                 return content.text
 
         return content

app/src/common/components.py

@@ -1,6 +1,12 @@
-from typing import Sequence
+import logging
+from typing import Any, Callable, Generator, Sequence
 
+import hayhooks
+from haystack import Pipeline
 from haystack.dataclasses.chat_message import ChatMessage
+from openinference.instrumentation import using_metadata
+from opentelemetry.trace import Span
+from opentelemetry.trace.status import Status, StatusCode
 from phoenix.client.__generated__ import v1
 
 from src.common import phoenix_utils
@@ -12,7 +18,7 @@ def get_phoenix_prompt(prompt_name: str, prompt_version_id: str = "") -> list[Ch
 
 
 def to_chat_messages(
-    msg_list: Sequence[dict | v1.PromptMessage | ChatMessage],
+    msg_list: Sequence[str | dict | v1.PromptMessage | ChatMessage],
 ) -> list[ChatMessage]:
     """Convert a list of dicts or Phoenix PromptMessage to a list of Haystack ChatMessage."""
     messages = []
@@ -26,11 +32,17 @@ def to_chat_messages(
         role = msg["role"]
         content = msg["content"]
 
-        assert isinstance(content, list), f"Expected list content, got {type(content)}: {content}"
-        assert len(content) == 1, f"Expected single content, got {len(content)} items: {content}"
-        assert content[0]["type"] == "text", f"Expected text content, got {content[0]['type']}"
-        assert "text" in content[0], f"Expected 'text' in content[0], got {content[0]}"
-        text = content[0]["text"]
+        if isinstance(content, str):
+            text = content
+        elif isinstance(content, list):
+            assert (
+                len(content) == 1
+            ), f"Expected single content, got {len(content)} items: {content}"
+            assert content[0]["type"] == "text", f"Expected text content, got {content[0]['type']}"
+            assert "text" in content[0], f"Expected 'text' in content[0], got {content[0]}"
+            text = content[0]["text"]
+        else:
+            raise ValueError(f"Unexpected content type: {type(content)} for message {msg}")
 
         if role == "system":
             assert isinstance(text, str), f"Expected string, got {type(text)}"
@@ -46,3 +58,90 @@ def to_chat_messages(
         messages.append(chat_msg)
 
     return messages
+
+
+logger = logging.getLogger(__name__)
+
+
+class TracedPipelineRunner:
+    """Helper class to run Haystack pipelines with OpenInference tracing."""
+
+    def __init__(self, parent_span_name: str, pipeline: Pipeline) -> None:
+        self.parent_span_name = parent_span_name
+        self.pipeline = pipeline
+
+    def stream_response(
+        self,
+        pipeline_run_args: dict,
+        *,
+        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()
+            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:
+                assert isinstance(span, Span), f"Got unexpected {type(span)}"
+                try:
+                    span.set_input(input_)
+
+                    # 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
+                    generator = hayhooks.streaming_generator(
+                        pipeline=self.pipeline,
+                        pipeline_run_args=pipeline_run_args,
+                    )
+
+                    # Stream chunks with each 'yield' call
+                    chunk_count = 0
+                    full_response = []
+                    for chunk in generator:
+                        chunk_count += 1
+                        if hasattr(chunk, "content"):
+                            full_response.append(chunk.content)
+                        # Must yield chunks one by one for SSE
+                        # Must yield in the tracer span context to have spans linked as child spans
+                        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))
+                    span.set_status(Status(StatusCode.OK))
+                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
+
+    def return_response(
+        self,
+        pipeline_run_args: dict,
+        *,
+        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 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:
+                try:
+                    result = self.pipeline.run(
+                        pipeline_run_args,
+                        include_outputs_from=include_outputs_from,
+                    )
+                    # Shorter than span.set_attribute(SpanAttributes.INPUT_VALUE, ...)
+                    span.set_input(input_)
+                    span.set_output(extract_output(result))
+                    span.set_status(Status(StatusCode.OK))
+                    return result
+                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

app/src/common/haystack_utils.py

@@ -5,6 +5,8 @@
 import httpx
 import opentelemetry.exporter.otlp.proto.http.trace_exporter as otel_trace_exporter
 import phoenix.otel
+from openinference.instrumentation import OITracer
+from opentelemetry.instrumentation.threading import ThreadingInstrumentor
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.trace import NoOpTracerProvider, TracerProvider
 
@@ -19,6 +21,51 @@
 logger = logging.getLogger(__name__)
 
 
+class ContextDetachErrorFilter(logging.Filter):
+    """Filter to suppress harmless OpenTelemetry context detach errors in threaded environments.
+
+    When using OpenInference instrumentation with threading (e.g., hayhooks.streaming_generator),
+    context tokens may be created in one thread and detached in another during cleanup,
+    causing harmless "Failed to detach context" or "was created in a different Context" errors.
+
+    These errors don't affect functionality - the spans are created correctly and traces work as expected.
+    This filter suppresses these specific error messages to reduce log noise.
+
+    Related issue: https://github.com/Arize-ai/openinference/issues/306
+    This is the recommended pattern until OpenInference fixes the upstream issue.
+    Similar approaches are used by other projects facing this issue:
+    - Google ADK (https://github.com/google/adk-python/issues/1670)
+    - Agno framework (https://github.com/agno-agi/agno/issues/5208)
+    - Langfuse (https://github.com/langfuse/langfuse/issues/8316)
+    """
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        message = record.getMessage()
+
+        # Suppress "Failed to detach context" errors
+        if "Failed to detach context" in message:
+            return False
+
+        # Suppress "was created in a different Context" errors
+        if "was created in a different Context" in message:
+            return False
+
+        # Also check exception info if present
+        if record.exc_info and record.exc_info[0] is ValueError:
+            exc_str = str(record.exc_info[1])
+            if "was created in a different Context" in exc_str:
+                return False
+
+        return True
+
+
+def _suppress_context_detach_errors() -> None:
+    """Add filter to suppress harmless OpenTelemetry context detach errors."""
+    otel_context_logger = logging.getLogger("opentelemetry.context")
+    otel_context_logger.addFilter(ContextDetachErrorFilter())
+    logger.info("Added filter to suppress harmless OpenTelemetry context detach errors")
+
+
 def _create_client(
     url: str = config.phoenix_collector_endpoint, api_key: str | None = None
 ) -> Client:
@@ -66,11 +113,25 @@ def configure_phoenix(only_if_alive: bool = True) -> None:
         auto_instrument=True,
     )
 
+    # Enable threading instrumentation to propagate OpenTelemetry context across threads
+    # This fixes the issue where hayhooks.streaming_generator() creates threads without
+    # inheriting the parent span context, causing orphaned spans in Phoenix traces.
+    # https://github.com/langfuse/langfuse/issues/8316#issuecomment-3154235201
+    ThreadingInstrumentor().instrument()
+    logger.info(
+        "Threading instrumentation enabled for OpenTelemetry context propagation during streaming"
+    )
+
+    # Suppress harmless context detach errors that occur when OpenInference instrumentation
+    # tries to clean up context tokens in different threads during streaming.
+    _suppress_context_detach_errors()
+
     if config.redact_pii:
         phoenix_api_key = os.environ.get("PHOENIX_API_KEY")
         span_exporter = otel_trace_exporter.OTLPSpanExporter(
             endpoint=trace_endpoint, headers={"Authorization": f"Bearer {phoenix_api_key}"}
         )
+
         # Create the PII redacting processor with the OTLP exporter
         pii_processor = PresidioRedactionSpanProcessor(span_exporter)
         # Add the pii processor to the otel instance
@@ -79,6 +140,18 @@ def configure_phoenix(only_if_alive: bool = True) -> None:
         tracer_provider.add_span_processor(pii_processor)
 
 
+_tracer: OITracer | None = None
+
+
+def tracer() -> OITracer:
+    global _tracer
+    if _tracer is None:
+        new_tracer = tracer_provider.get_tracer(__name__)
+        assert isinstance(new_tracer, OITracer), f"Got unexpected {type(new_tracer)}"
+        _tracer = new_tracer
+    return _tracer
+
+
 def get_prompt_template(prompt_name: str, prompt_version_id: str = "") -> PromptVersion:
     """Retrieve a prompt template from Phoenix by name.
     https://arize.com/docs/phoenix/sdk-api-reference/python/overview#prompt-management