langfuse-haystack 0.6.4__tar.gz → 0.8.0__tar.gz

{langfuse_haystack-0.6.4 → langfuse_haystack-0.8.0}/CHANGELOG.md

@@ -1,5 +1,23 @@
 # Changelog
 
+## [integrations/langfuse-v0.7.0] - 2025-01-21
+
+### 🚀 Features
+
+- LangfuseConnector - add httpx.Client init param (#1308)
+
+### 🐛 Bug Fixes
+
+- End langfuse generation spans properly (#1301)
+
+
+## [integrations/langfuse-v0.6.4] - 2025-01-17
+
+### 🚀 Features
+
+- Add LangfuseConnector secure key management and serialization  (#1287)
+
+
 ## [integrations/langfuse-v0.6.3] - 2025-01-15
 
 ### 🌀 Miscellaneous

{langfuse_haystack-0.6.4 → langfuse_haystack-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfuse-haystack
-Version: 0.6.4
+Version: 0.8.0
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/langfuse#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
 Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/langfuse

{langfuse_haystack-0.6.4 → langfuse_haystack-0.8.0}/pyproject.toml

@@ -67,7 +67,12 @@ python = ["3.8", "3.9", "3.10", "3.11", "3.12"]
 [tool.hatch.envs.lint]
 installer = "uv"
 detached = true
-dependencies = ["pip", "black>=23.1.0", "mypy>=1.0.0", "ruff>=0.0.243"]
+dependencies = [
+    "pip",
+    "black>=23.1.0",
+    "mypy>=1.0.0",
+    "ruff>=0.0.243",
+]
 
 [tool.hatch.envs.lint.scripts]
 typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}"
@@ -167,6 +172,7 @@ module = [
   "haystack_integrations.*",
   "pytest.*",
   "numpy.*",
+  "httpx.*",
 ]
 ignore_missing_imports = true
 

{langfuse_haystack-0.6.4 → langfuse_haystack-0.8.0}/src/haystack_integrations/components/connectors/langfuse/langfuse_connector.py

@@ -1,9 +1,11 @@
 from typing import Any, Dict, Optional
 
+import httpx
 from haystack import component, default_from_dict, default_to_dict, logging, tracing
 from haystack.utils import Secret, deserialize_secrets_inplace
+from haystack.utils.base_serialization import deserialize_class_instance, serialize_class_instance
 
-from haystack_integrations.tracing.langfuse import LangfuseTracer
+from haystack_integrations.tracing.langfuse import LangfuseTracer, SpanHandler
 from langfuse import Langfuse
 
 logger = logging.getLogger(__name__)
@@ -93,6 +95,24 @@ class LangfuseConnector:
         print(response["tracer"]["trace_url"])
     ```
 
+    For advanced use cases, you can also customize how spans are created and processed by
+    providing a custom SpanHandler. This allows you to add custom metrics, set warning levels,
+    or attach additional metadata to your Langfuse traces:
+
+    ```python
+    from haystack_integrations.tracing.langfuse import DefaultSpanHandler, LangfuseSpan
+    from typing import Optional
+
+    class CustomSpanHandler(DefaultSpanHandler):
+
+        def handle(self, span: LangfuseSpan, component_type: Optional[str]) -> None:
+            # Custom span handling logic, customize Langfuse spans however it fits you
+            # see DefaultSpanHandler for how we create and process spans by default
+            pass
+
+    connector = LangfuseConnector(span_handler=CustomSpanHandler())
+    ```
+
     """
 
     def __init__(
@@ -101,6 +121,8 @@ class LangfuseConnector:
         public: bool = False,
         public_key: Optional[Secret] = Secret.from_env_var("LANGFUSE_PUBLIC_KEY"),  # noqa: B008
         secret_key: Optional[Secret] = Secret.from_env_var("LANGFUSE_SECRET_KEY"),  # noqa: B008
+        httpx_client: Optional[httpx.Client] = None,
+        span_handler: Optional[SpanHandler] = None,
     ):
         """
         Initialize the LangfuseConnector component.
@@ -112,18 +134,28 @@ class LangfuseConnector:
             only accessible to the Langfuse account owner. The default is `False`.
         :param public_key: The Langfuse public key. Defaults to reading from LANGFUSE_PUBLIC_KEY environment variable.
         :param secret_key: The Langfuse secret key. Defaults to reading from LANGFUSE_SECRET_KEY environment variable.
+        :param httpx_client: Optional custom httpx.Client instance to use for Langfuse API calls. Note that when
+            deserializing a pipeline from YAML, any custom client is discarded and Langfuse will create its own default
+            client, since HTTPX clients cannot be serialized.
+        :param span_handler: Optional custom handler for processing spans. If None, uses DefaultSpanHandler.
+            The span handler controls how spans are created and processed, allowing customization of span types
+            based on component types and additional processing after spans are yielded. See SpanHandler class for
+            details on implementing custom handlers.
         """
         self.name = name
         self.public = public
         self.secret_key = secret_key
         self.public_key = public_key
+        self.span_handler = span_handler
         self.tracer = LangfuseTracer(
             tracer=Langfuse(
                 secret_key=secret_key.resolve_value() if secret_key else None,
                 public_key=public_key.resolve_value() if public_key else None,
+                httpx_client=httpx_client,
             ),
             name=name,
             public=public,
+            span_handler=span_handler,
         )
         tracing.enable_tracing(self.tracer)
 
@@ -152,12 +184,15 @@ class LangfuseConnector:
 
         :returns: The serialized component as a dictionary.
         """
+        span_handler = serialize_class_instance(self.span_handler) if self.span_handler else None
         return default_to_dict(
             self,
             name=self.name,
             public=self.public,
             secret_key=self.secret_key.to_dict() if self.secret_key else None,
             public_key=self.public_key.to_dict() if self.public_key else None,
+            # Note: httpx_client is not serialized as it's not serializable
+            span_handler=span_handler,
         )
 
     @classmethod
@@ -168,5 +203,9 @@ class LangfuseConnector:
         :param data: The dictionary representation of this component.
         :returns: The deserialized component instance.
         """
-        deserialize_secrets_inplace(data["init_parameters"], keys=["secret_key", "public_key"])
+        init_params = data["init_parameters"]
+        deserialize_secrets_inplace(init_params, keys=["secret_key", "public_key"])
+        init_params["span_handler"] = (
+            deserialize_class_instance(init_params["span_handler"]) if init_params["span_handler"] else None
+        )
         return default_from_dict(cls, data)

langfuse_haystack-0.8.0/src/haystack_integrations/tracing/langfuse/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+from .tracer import DefaultSpanHandler, LangfuseSpan, LangfuseTracer, SpanContext, SpanHandler
+
+__all__ = ["DefaultSpanHandler", "LangfuseSpan", "LangfuseTracer", "SpanContext", "SpanHandler"]

langfuse_haystack-0.8.0/src/haystack_integrations/tracing/langfuse/tracer.py

@@ -0,0 +1,382 @@
+import contextlib
+import os
+from abc import ABC, abstractmethod
+from contextvars import ContextVar
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from haystack import default_from_dict, default_to_dict, logging
+from haystack.dataclasses import ChatMessage
+from haystack.tracing import Span, Tracer
+from haystack.tracing import tracer as proxy_tracer
+from haystack.tracing import utils as tracing_utils
+from typing_extensions import TypeAlias
+
+import langfuse
+from langfuse.client import StatefulGenerationClient, StatefulSpanClient, StatefulTraceClient
+
+# Type alias for Langfuse stateful clients
+LangfuseStatefulClient: TypeAlias = Union[StatefulTraceClient, StatefulSpanClient, StatefulGenerationClient]
+
+
+logger = logging.getLogger(__name__)
+
+HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR = "HAYSTACK_LANGFUSE_ENFORCE_FLUSH"
+_SUPPORTED_GENERATORS = [
+    "AzureOpenAIGenerator",
+    "OpenAIGenerator",
+    "AnthropicGenerator",
+    "HuggingFaceAPIGenerator",
+    "HuggingFaceLocalGenerator",
+    "CohereGenerator",
+]
+_SUPPORTED_CHAT_GENERATORS = [
+    "AzureOpenAIChatGenerator",
+    "OpenAIChatGenerator",
+    "AnthropicChatGenerator",
+    "HuggingFaceAPIChatGenerator",
+    "HuggingFaceLocalChatGenerator",
+    "CohereChatGenerator",
+]
+_ALL_SUPPORTED_GENERATORS = _SUPPORTED_GENERATORS + _SUPPORTED_CHAT_GENERATORS
+
+# These are the keys used by Haystack for traces and span.
+# We keep them here to avoid making typos when using them.
+_PIPELINE_RUN_KEY = "haystack.pipeline.run"
+_COMPONENT_NAME_KEY = "haystack.component.name"
+_COMPONENT_TYPE_KEY = "haystack.component.type"
+_COMPONENT_OUTPUT_KEY = "haystack.component.output"
+
+# Context var used to keep track of tracing related info.
+# This mainly useful for parents spans.
+tracing_context_var: ContextVar[Dict[Any, Any]] = ContextVar("tracing_context")
+
+
+class LangfuseSpan(Span):
+    """
+    Internal class representing a bridge between the Haystack span tracing API and Langfuse.
+    """
+
+    def __init__(self, span: LangfuseStatefulClient) -> None:
+        """
+        Initialize a LangfuseSpan instance.
+
+        :param span: The span instance managed by Langfuse.
+        """
+        self._span = span
+        # locally cache tags
+        self._data: Dict[str, Any] = {}
+
+    def set_tag(self, key: str, value: Any) -> None:
+        """
+        Set a generic tag for this span.
+
+        :param key: The tag key.
+        :param value: The tag value.
+        """
+        coerced_value = tracing_utils.coerce_tag_value(value)
+        self._span.update(metadata={key: coerced_value})
+        self._data[key] = value
+
+    def set_content_tag(self, key: str, value: Any) -> None:
+        """
+        Set a content-specific tag for this span.
+
+        :param key: The content tag key.
+        :param value: The content tag value.
+        """
+        if not proxy_tracer.is_content_tracing_enabled:
+            return
+        if key.endswith(".input"):
+            if "messages" in value:
+                messages = [m.to_openai_dict_format() for m in value["messages"]]
+                self._span.update(input=messages)
+            else:
+                self._span.update(input=value)
+        elif key.endswith(".output"):
+            if "replies" in value:
+                if all(isinstance(r, ChatMessage) for r in value["replies"]):
+                    replies = [m.to_openai_dict_format() for m in value["replies"]]
+                else:
+                    replies = value["replies"]
+                self._span.update(output=replies)
+            else:
+                self._span.update(output=value)
+
+        self._data[key] = value
+
+    def raw_span(self) -> LangfuseStatefulClient:
+        """
+        Return the underlying span instance.
+
+        :return: The Langfuse span instance.
+        """
+        return self._span
+
+    def get_correlation_data_for_logs(self) -> Dict[str, Any]:
+        return {}
+
+
+@dataclass(frozen=True)
+class SpanContext:
+    """
+    Context for creating spans in Langfuse.
+
+    Encapsulates the information needed to create and configure a span in Langfuse tracing.
+    Used by SpanHandler to determine the span type (trace, generation, or default) and its configuration.
+
+    :param name: The name of the span to create. For components, this is typically the component name.
+    :param operation_name: The operation being traced (e.g. "haystack.pipeline.run"). Used to determine
+        if a new trace should be created without warning.
+    :param component_type: The type of component creating the span (e.g. "OpenAIChatGenerator").
+        Can be used to determine the type of span to create.
+    :param tags: Additional metadata to attach to the span. Contains component input/output data
+        and other trace information.
+    :param parent_span: The parent span if this is a child span. If None, a new trace will be created.
+    :param trace_name: The name to use for the trace when creating a parent span. Defaults to "Haystack".
+    :param public: Whether traces should be publicly accessible. Defaults to False.
+    """
+
+    name: str
+    operation_name: str
+    component_type: Optional[str]
+    tags: Dict[str, Any]
+    parent_span: Optional[Span]
+    trace_name: str = "Haystack"
+    public: bool = False
+
+    def __post_init__(self) -> None:
+        """
+        Validate the span context attributes.
+
+        :raises ValueError: If name, operation_name or trace_name are empty
+        :raises TypeError: If tags is not a dictionary
+        """
+        if not self.name:
+            msg = "Span name cannot be empty"
+            raise ValueError(msg)
+        if not self.operation_name:
+            msg = "Operation name cannot be empty"
+            raise ValueError(msg)
+        if not self.trace_name:
+            msg = "Trace name cannot be empty"
+            raise ValueError(msg)
+
+
+class SpanHandler(ABC):
+    """
+    Abstract base class for customizing how Langfuse spans are created and processed.
+
+    This class defines two key extension points:
+    1. create_span: Controls what type of span to create (default or generation)
+    2. handle: Processes the span after component execution (adding metadata, metrics, etc.)
+
+    To implement a custom handler:
+    - Extend this class or DefaultSpanHandler
+    - Override create_span and handle methods. It is more common to override handle.
+    - Pass your handler to LangfuseConnector init method
+    """
+
+    def __init__(self):
+        self.tracer: Optional[langfuse.Langfuse] = None
+
+    def init_tracer(self, tracer: langfuse.Langfuse) -> None:
+        """
+        Initialize with Langfuse tracer. Called internally by LangfuseTracer.
+
+        :param tracer: The Langfuse client instance to use for creating spans
+        """
+        self.tracer = tracer
+
+    @abstractmethod
+    def create_span(self, context: SpanContext) -> LangfuseSpan:
+        """
+        Create a span of appropriate type based on the context.
+
+        This method determines what kind of span to create:
+        - A new trace if there's no parent span
+        - A generation span for LLM components
+        - A default span for other components
+
+        :param context: The context containing all information needed to create the span
+        :returns: A new LangfuseSpan instance configured according to the context
+        """
+        pass
+
+    @abstractmethod
+    def handle(self, span: LangfuseSpan, component_type: Optional[str]) -> None:
+        """
+        Process a span after component execution by attaching metadata and metrics.
+
+        This method is called after the component yields its span, allowing you to:
+        - Extract and attach token usage statistics
+        - Add model information
+        - Record timing data (e.g., time-to-first-token)
+        - Set log levels for quality monitoring
+        - Add custom metrics and observations
+
+        :param span: The span that was yielded by the component
+        :param component_type: The type of component that created the span, used to determine
+            what metadata to extract and how to process it
+        """
+        pass
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "SpanHandler":
+        return default_from_dict(cls, data)
+
+    def to_dict(self) -> Dict[str, Any]:
+        return default_to_dict(self)
+
+
+class DefaultSpanHandler(SpanHandler):
+    """DefaultSpanHandler provides the default Langfuse tracing behavior for Haystack."""
+
+    def create_span(self, context: SpanContext) -> LangfuseSpan:
+        message = "Tracer is not initialized"
+        if self.tracer is None:
+            raise RuntimeError(message)
+        tracing_ctx = tracing_context_var.get({})
+        if not context.parent_span:
+            if context.operation_name != _PIPELINE_RUN_KEY:
+                logger.warning(
+                    "Creating a new trace without a parent span is not recommended for operation '{operation_name}'.",
+                    operation_name=context.operation_name,
+                )
+            # Create a new trace when there's no parent span
+            return LangfuseSpan(
+                self.tracer.trace(
+                    name=context.trace_name,
+                    public=context.public,
+                    id=tracing_ctx.get("trace_id"),
+                    user_id=tracing_ctx.get("user_id"),
+                    session_id=tracing_ctx.get("session_id"),
+                    tags=tracing_ctx.get("tags"),
+                    version=tracing_ctx.get("version"),
+                )
+            )
+        elif context.component_type in _ALL_SUPPORTED_GENERATORS:
+            return LangfuseSpan(context.parent_span.raw_span().generation(name=context.name))
+        else:
+            return LangfuseSpan(context.parent_span.raw_span().span(name=context.name))
+
+    def handle(self, span: LangfuseSpan, component_type: Optional[str]) -> None:
+        if component_type in _SUPPORTED_GENERATORS:
+            meta = span._data.get(_COMPONENT_OUTPUT_KEY, {}).get("meta")
+            if meta:
+                m = meta[0]
+                span._span.update(usage=m.get("usage") or None, model=m.get("model"))
+        elif component_type in _SUPPORTED_CHAT_GENERATORS:
+            replies = span._data.get(_COMPONENT_OUTPUT_KEY, {}).get("replies")
+            if replies:
+                meta = replies[0].meta
+                completion_start_time = meta.get("completion_start_time")
+                if completion_start_time:
+                    try:
+                        completion_start_time = datetime.fromisoformat(completion_start_time)
+                    except ValueError:
+                        logger.error(f"Failed to parse completion_start_time: {completion_start_time}")
+                        completion_start_time = None
+                span._span.update(
+                    usage=meta.get("usage") or None,
+                    model=meta.get("model"),
+                    completion_start_time=completion_start_time,
+                )
+
+
+class LangfuseTracer(Tracer):
+    """
+    Internal class representing a bridge between the Haystack tracer and Langfuse.
+    """
+
+    def __init__(
+        self,
+        tracer: langfuse.Langfuse,
+        name: str = "Haystack",
+        public: bool = False,
+        span_handler: Optional[SpanHandler] = None,
+    ) -> None:
+        """
+        Initialize a LangfuseTracer instance.
+
+        :param tracer: The Langfuse tracer instance.
+        :param name: The name of the pipeline or component. This name will be used to identify the tracing run on the
+            Langfuse dashboard.
+        :param public: Whether the tracing data should be public or private. If set to `True`, the tracing data will
+            be publicly accessible to anyone with the tracing URL. If set to `False`, the tracing data will be private
+            and only accessible to the Langfuse account owner.
+        :param span_handler: Custom handler for processing spans. If None, uses DefaultSpanHandler.
+        """
+        if not proxy_tracer.is_content_tracing_enabled:
+            logger.warning(
+                "Traces will not be logged to Langfuse because Haystack tracing is disabled. "
+                "To enable, set the HAYSTACK_CONTENT_TRACING_ENABLED environment variable to true "
+                "before importing Haystack."
+            )
+        self._tracer = tracer
+        self._context: List[LangfuseSpan] = []
+        self._name = name
+        self._public = public
+        self.enforce_flush = os.getenv(HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR, "true").lower() == "true"
+        self._span_handler = span_handler or DefaultSpanHandler()
+        self._span_handler.init_tracer(tracer)
+
+    @contextlib.contextmanager
+    def trace(
+        self, operation_name: str, tags: Optional[Dict[str, Any]] = None, parent_span: Optional[Span] = None
+    ) -> Iterator[Span]:
+        tags = tags or {}
+        span_name = tags.get(_COMPONENT_NAME_KEY, operation_name)
+        component_type = tags.get(_COMPONENT_TYPE_KEY)
+
+        # Create span using the handler
+        span = self._span_handler.create_span(
+            SpanContext(
+                name=span_name,
+                operation_name=operation_name,
+                component_type=component_type,
+                tags=tags,
+                parent_span=parent_span,
+                trace_name=self._name,
+                public=self._public,
+            )
+        )
+
+        self._context.append(span)
+        span.set_tags(tags)
+
+        yield span
+
+        # Let the span handler process the span
+        self._span_handler.handle(span, component_type)
+
+        raw_span = span.raw_span()
+
+        # In this section, we finalize both regular spans and generation spans created using the LangfuseSpan class.
+        # It's important to end() these spans to ensure they are properly closed and all relevant data is recorded.
+        # Note that we do not call end() on the main trace span itself, as its lifecycle is managed differently.
+        if isinstance(raw_span, (StatefulSpanClient, StatefulGenerationClient)):
+            raw_span.end()
+        self._context.pop()
+
+        if self.enforce_flush:
+            self.flush()
+
+    def flush(self):
+        self._tracer.flush()
+
+    def current_span(self) -> Optional[Span]:
+        """
+        Return the current active span.
+
+        :return: The current span if available, else None.
+        """
+        return self._context[-1] if self._context else None
+
+    def get_trace_url(self) -> str:
+        """
+        Return the URL to the tracing data.
+        :return: The URL to the tracing data.
+        """
+        return self._tracer.get_trace_url()

{langfuse_haystack-0.6.4 → langfuse_haystack-0.8.0}/tests/test_tracing.py

@@ -1,6 +1,7 @@
 import os
 import time
 from urllib.parse import urlparse
+from typing import Optional
 
 import pytest
 import requests
@@ -10,10 +11,13 @@ from haystack.components.generators.chat import OpenAIChatGenerator
 from haystack.dataclasses import ChatMessage
 from haystack.utils import Secret
 from requests.auth import HTTPBasicAuth
+import httpx
 
 from haystack_integrations.components.connectors.langfuse import LangfuseConnector
 from haystack_integrations.components.generators.anthropic import AnthropicChatGenerator
 from haystack_integrations.components.generators.cohere import CohereChatGenerator
+from haystack_integrations.tracing.langfuse import LangfuseSpan, DefaultSpanHandler
+from haystack_integrations.tracing.langfuse.tracer import _COMPONENT_OUTPUT_KEY
 
 # don't remove (or move) this env var setting from here, it's needed to turn tracing on
 os.environ["HAYSTACK_CONTENT_TRACING_ENABLED"] = "true"
@@ -49,6 +53,27 @@ def pipeline_with_secrets(llm_class, expected_trace):
     return pipe
 
 
+@pytest.fixture
+def pipeline_with_custom_client(llm_class, expected_trace):
+    """Pipeline factory using custom httpx client for Langfuse"""
+    pipe = Pipeline()
+    custom_client = httpx.Client(timeout=30.0)  # Custom timeout of 30 seconds
+    pipe.add_component(
+        "tracer",
+        LangfuseConnector(
+            name=f"Chat example - {expected_trace}",
+            public=True,
+            secret_key=Secret.from_env_var("LANGFUSE_SECRET_KEY"),
+            public_key=Secret.from_env_var("LANGFUSE_PUBLIC_KEY"),
+            httpx_client=custom_client,
+        ),
+    )
+    pipe.add_component("prompt_builder", ChatPromptBuilder())
+    pipe.add_component("llm", llm_class())
+    pipe.connect("prompt_builder.prompt", "llm.messages")
+    return pipe
+
+
 @pytest.mark.integration
 @pytest.mark.parametrize(
     "llm_class, env_var, expected_trace",
@@ -58,7 +83,9 @@ def pipeline_with_secrets(llm_class, expected_trace):
         (CohereChatGenerator, "COHERE_API_KEY", "Cohere"),
     ],
 )
-@pytest.mark.parametrize("pipeline_fixture", ["pipeline_with_env_vars", "pipeline_with_secrets"])
+@pytest.mark.parametrize(
+    "pipeline_fixture", ["pipeline_with_env_vars", "pipeline_with_secrets", "pipeline_with_custom_client"]
+)
 def test_tracing_integration(llm_class, env_var, expected_trace, pipeline_fixture, request):
     if not all([os.environ.get("LANGFUSE_SECRET_KEY"), os.environ.get("LANGFUSE_PUBLIC_KEY"), os.environ.get(env_var)]):
         pytest.skip(f"Missing required environment variables: LANGFUSE_SECRET_KEY, LANGFUSE_PUBLIC_KEY, or {env_var}")
@@ -145,3 +172,86 @@ def test_pipeline_serialization(monkeypatch):
 
     # Verify pipeline is the same
     assert new_pipe == pipe
+
+
+class QualityCheckSpanHandler(DefaultSpanHandler):
+    """Extends default handler to add quality checks with warning levels."""
+
+    def handle(self, span: LangfuseSpan, component_type: Optional[str]) -> None:
+        # First do the default handling (model, usage, etc.)
+        super().handle(span, component_type)
+
+        # Then add our custom quality checks
+        if component_type == "OpenAIChatGenerator":
+            output = span._data.get(_COMPONENT_OUTPUT_KEY, {})
+            replies = output.get("replies", [])
+
+            if not replies:
+                span._span.update(level="ERROR", status_message="No response received")
+                return
+
+            reply = replies[0]
+            if "error" in reply.meta:
+                span._span.update(level="ERROR", status_message=f"OpenAI error: {reply.meta['error']}")
+            elif len(reply.text) > 10:
+                span._span.update(level="WARNING", status_message="Response too long (> 10 chars)")
+            else:
+                span._span.update(level="DEFAULT", status_message="Success")
+
+
+@pytest.mark.integration
+def test_custom_span_handler():
+    """Test that custom span handler properly sets Langfuse levels."""
+    if not all(
+        [os.environ.get("LANGFUSE_SECRET_KEY"), os.environ.get("LANGFUSE_PUBLIC_KEY"), os.environ.get("OPENAI_API_KEY")]
+    ):
+        pytest.skip("Missing required environment variables")
+
+    pipe = Pipeline()
+    pipe.add_component(
+        "tracer",
+        LangfuseConnector(
+            name="Quality Check Example",
+            public=True,
+            span_handler=QualityCheckSpanHandler(),
+        ),
+    )
+    pipe.add_component("prompt_builder", ChatPromptBuilder())
+    pipe.add_component("llm", OpenAIChatGenerator())
+    pipe.connect("prompt_builder.prompt", "llm.messages")
+
+    # Test short response
+    messages = [
+        ChatMessage.from_system("Respond with exactly 3 words."),
+        ChatMessage.from_user("What is Berlin?"),
+    ]
+
+    response = pipe.run(
+        data={
+            "prompt_builder": {"template_variables": {}, "template": messages},
+            "tracer": {"invocation_context": {"user_id": "test_user"}},
+        }
+    )
+
+    trace_url = response["tracer"]["trace_url"]
+    uuid = os.path.basename(urlparse(trace_url).path)
+    url = f"https://cloud.langfuse.com/api/public/traces/{uuid}"
+
+    # Poll the Langfuse API
+    attempts = 5
+    delay = 1
+    while attempts >= 0:
+        res = requests.get(
+            url, auth=HTTPBasicAuth(os.environ["LANGFUSE_PUBLIC_KEY"], os.environ["LANGFUSE_SECRET_KEY"])
+        )
+        if attempts > 0 and res.status_code != 200:
+            attempts -= 1
+            time.sleep(delay)
+            delay *= 2
+            continue
+
+        assert res.status_code == 200
+        content = str(res.content)
+        assert "WARNING" in content
+        assert "Response too long" in content
+        break

langfuse_haystack-0.6.4/src/haystack_integrations/tracing/langfuse/__init__.py

@@ -1,6 +0,0 @@
-# SPDX-FileCopyrightText: 2024-present deepset GmbH <info@deepset.ai>
-#
-# SPDX-License-Identifier: Apache-2.0
-from .tracer import LangfuseTracer
-
-__all__ = ["LangfuseTracer"]

langfuse_haystack-0.6.4/src/haystack_integrations/tracing/langfuse/tracer.py

@@ -1,227 +0,0 @@
-import contextlib
-import os
-from contextvars import ContextVar
-from datetime import datetime
-from typing import Any, Dict, Iterator, List, Optional, Union
-
-from haystack import logging
-from haystack.dataclasses import ChatMessage
-from haystack.tracing import Span, Tracer
-from haystack.tracing import tracer as proxy_tracer
-from haystack.tracing import utils as tracing_utils
-
-import langfuse
-
-logger = logging.getLogger(__name__)
-
-HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR = "HAYSTACK_LANGFUSE_ENFORCE_FLUSH"
-_SUPPORTED_GENERATORS = [
-    "AzureOpenAIGenerator",
-    "OpenAIGenerator",
-    "AnthropicGenerator",
-    "HuggingFaceAPIGenerator",
-    "HuggingFaceLocalGenerator",
-    "CohereGenerator",
-]
-_SUPPORTED_CHAT_GENERATORS = [
-    "AzureOpenAIChatGenerator",
-    "OpenAIChatGenerator",
-    "AnthropicChatGenerator",
-    "HuggingFaceAPIChatGenerator",
-    "HuggingFaceLocalChatGenerator",
-    "CohereChatGenerator",
-]
-_ALL_SUPPORTED_GENERATORS = _SUPPORTED_GENERATORS + _SUPPORTED_CHAT_GENERATORS
-
-# These are the keys used by Haystack for traces and span.
-# We keep them here to avoid making typos when using them.
-_PIPELINE_RUN_KEY = "haystack.pipeline.run"
-_COMPONENT_NAME_KEY = "haystack.component.name"
-_COMPONENT_TYPE_KEY = "haystack.component.type"
-_COMPONENT_OUTPUT_KEY = "haystack.component.output"
-
-# Context var used to keep track of tracing related info.
-# This mainly useful for parents spans.
-tracing_context_var: ContextVar[Dict[Any, Any]] = ContextVar("tracing_context")
-
-
-class LangfuseSpan(Span):
-    """
-    Internal class representing a bridge between the Haystack span tracing API and Langfuse.
-    """
-
-    def __init__(self, span: "Union[langfuse.client.StatefulSpanClient, langfuse.client.StatefulTraceClient]") -> None:
-        """
-        Initialize a LangfuseSpan instance.
-
-        :param span: The span instance managed by Langfuse.
-        """
-        self._span = span
-        # locally cache tags
-        self._data: Dict[str, Any] = {}
-
-    def set_tag(self, key: str, value: Any) -> None:
-        """
-        Set a generic tag for this span.
-
-        :param key: The tag key.
-        :param value: The tag value.
-        """
-        coerced_value = tracing_utils.coerce_tag_value(value)
-        self._span.update(metadata={key: coerced_value})
-        self._data[key] = value
-
-    def set_content_tag(self, key: str, value: Any) -> None:
-        """
-        Set a content-specific tag for this span.
-
-        :param key: The content tag key.
-        :param value: The content tag value.
-        """
-        if not proxy_tracer.is_content_tracing_enabled:
-            return
-        if key.endswith(".input"):
-            if "messages" in value:
-                messages = [m.to_openai_dict_format() for m in value["messages"]]
-                self._span.update(input=messages)
-            else:
-                self._span.update(input=value)
-        elif key.endswith(".output"):
-            if "replies" in value:
-                if all(isinstance(r, ChatMessage) for r in value["replies"]):
-                    replies = [m.to_openai_dict_format() for m in value["replies"]]
-                else:
-                    replies = value["replies"]
-                self._span.update(output=replies)
-            else:
-                self._span.update(output=value)
-
-        self._data[key] = value
-
-    def raw_span(self) -> "Union[langfuse.client.StatefulSpanClient, langfuse.client.StatefulTraceClient]":
-        """
-        Return the underlying span instance.
-
-        :return: The Langfuse span instance.
-        """
-        return self._span
-
-    def get_correlation_data_for_logs(self) -> Dict[str, Any]:
-        return {}
-
-
-class LangfuseTracer(Tracer):
-    """
-    Internal class representing a bridge between the Haystack tracer and Langfuse.
-    """
-
-    def __init__(self, tracer: "langfuse.Langfuse", name: str = "Haystack", public: bool = False) -> None:
-        """
-        Initialize a LangfuseTracer instance.
-
-        :param tracer: The Langfuse tracer instance.
-        :param name: The name of the pipeline or component. This name will be used to identify the tracing run on the
-            Langfuse dashboard.
-        :param public: Whether the tracing data should be public or private. If set to `True`, the tracing data will
-        be publicly accessible to anyone with the tracing URL. If set to `False`, the tracing data will be private
-        and only accessible to the Langfuse account owner.
-        """
-        if not proxy_tracer.is_content_tracing_enabled:
-            logger.warning(
-                "Traces will not be logged to Langfuse because Haystack tracing is disabled. "
-                "To enable, set the HAYSTACK_CONTENT_TRACING_ENABLED environment variable to true "
-                "before importing Haystack."
-            )
-        self._tracer = tracer
-        self._context: List[LangfuseSpan] = []
-        self._name = name
-        self._public = public
-        self.enforce_flush = os.getenv(HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR, "true").lower() == "true"
-
-    @contextlib.contextmanager
-    def trace(
-        self, operation_name: str, tags: Optional[Dict[str, Any]] = None, parent_span: Optional[Span] = None
-    ) -> Iterator[Span]:
-        tags = tags or {}
-        span_name = tags.get(_COMPONENT_NAME_KEY, operation_name)
-
-        # Create new span depending whether there's a parent span or not
-        if not parent_span:
-            if operation_name != _PIPELINE_RUN_KEY:
-                logger.warning(
-                    "Creating a new trace without a parent span is not recommended for operation '{operation_name}'.",
-                    operation_name=operation_name,
-                )
-            # Create a new trace if no parent span is provided
-            context = tracing_context_var.get({})
-            span = LangfuseSpan(
-                self._tracer.trace(
-                    name=self._name,
-                    public=self._public,
-                    id=context.get("trace_id"),
-                    user_id=context.get("user_id"),
-                    session_id=context.get("session_id"),
-                    tags=context.get("tags"),
-                    version=context.get("version"),
-                )
-            )
-        elif tags.get(_COMPONENT_TYPE_KEY) in _ALL_SUPPORTED_GENERATORS:
-            span = LangfuseSpan(parent_span.raw_span().generation(name=span_name))
-        else:
-            span = LangfuseSpan(parent_span.raw_span().span(name=span_name))
-
-        self._context.append(span)
-        span.set_tags(tags)
-
-        yield span
-
-        # Update span metadata based on component type
-        if tags.get(_COMPONENT_TYPE_KEY) in _SUPPORTED_GENERATORS:
-            # Haystack returns one meta dict for each message, but the 'usage' value
-            # is always the same, let's just pick the first item
-            meta = span._data.get(_COMPONENT_OUTPUT_KEY, {}).get("meta")
-            if meta:
-                m = meta[0]
-                span._span.update(usage=m.get("usage") or None, model=m.get("model"))
-        elif tags.get(_COMPONENT_TYPE_KEY) in _SUPPORTED_CHAT_GENERATORS:
-            replies = span._data.get(_COMPONENT_OUTPUT_KEY, {}).get("replies")
-            if replies:
-                meta = replies[0].meta
-                completion_start_time = meta.get("completion_start_time")
-                if completion_start_time:
-                    try:
-                        completion_start_time = datetime.fromisoformat(completion_start_time)
-                    except ValueError:
-                        logger.error(f"Failed to parse completion_start_time: {completion_start_time}")
-                        completion_start_time = None
-                span._span.update(
-                    usage=meta.get("usage") or None,
-                    model=meta.get("model"),
-                    completion_start_time=completion_start_time,
-                )
-
-        raw_span = span.raw_span()
-        if isinstance(raw_span, langfuse.client.StatefulSpanClient):
-            raw_span.end()
-        self._context.pop()
-
-        if self.enforce_flush:
-            self.flush()
-
-    def flush(self):
-        self._tracer.flush()
-
-    def current_span(self) -> Optional[Span]:
-        """
-        Return the current active span.
-
-        :return: The current span if available, else None.
-        """
-        return self._context[-1] if self._context else None
-
-    def get_trace_url(self) -> str:
-        """
-        Return the URL to the tracing data.
-        :return: The URL to the tracing data.
-        """
-        return self._tracer.get_trace_url()