langfuse-haystack 3.2.1__tar.gz → 3.3.1__tar.gz

{langfuse_haystack-3.2.1 → langfuse_haystack-3.3.1}/CHANGELOG.md

@@ -1,5 +1,22 @@
 # Changelog
 
+## [integrations/langfuse-v3.3.0] - 2025-11-21
+
+### 🚀 Features
+
+- *(langfuse)* Embedder, retriever and generator as obs. type (#2497)
+
+### 🌀 Miscellaneous
+
+- Enhancement: Adopt PEP 585 type hinting (part 4) (#2527)
+- *(langfuse)* Log levels (#2522)
+
+## [integrations/langfuse-v3.2.1] - 2025-11-07
+
+### 🌀 Miscellaneous
+
+- Chore: Upgrade langfuse dep, observation types require version>=3.3.1 (#2493)
+
 ## [integrations/langfuse-v3.2.0] - 2025-11-07
 
 ### 🐛 Bug Fixes

{langfuse_haystack-3.2.1 → langfuse_haystack-3.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfuse-haystack
-Version: 3.2.1
+Version: 3.3.1
 Summary: Langfuse integration for Haystack
 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

{langfuse_haystack-3.2.1 → langfuse_haystack-3.3.1}/pyproject.toml

@@ -46,7 +46,7 @@ installer = "uv"
 dependencies = ["haystack-pydoc-tools", "ruff"]
 
 [tool.hatch.envs.default.scripts]
-docs = ["pydoc-markdown pydoc/config.yml"]
+docs = ["pydoc-markdown pydoc/config_docusaurus.yml"]
 fmt = "ruff check --fix {args} && ruff format {args}"
 fmt-check = "ruff check {args} && ruff format --check {args}"
 
@@ -82,7 +82,7 @@ allow-direct-references = true
 
 
 [tool.ruff]
-target-version = "py38"
+target-version = "py39"
 line-length = 120
 
 [tool.ruff.lint]

{langfuse_haystack-3.2.1 → langfuse_haystack-3.3.1}/src/haystack_integrations/components/connectors/langfuse/langfuse_connector.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import httpx
 from haystack import component, default_from_dict, default_to_dict, logging, tracing
@@ -124,7 +124,7 @@ class LangfuseConnector:
         span_handler: Optional[SpanHandler] = None,
         *,
         host: Optional[str] = None,
-        langfuse_client_kwargs: Optional[Dict[str, Any]] = None,
+        langfuse_client_kwargs: Optional[dict[str, Any]] = None,
     ) -> None:
         """
         Initialize the LangfuseConnector component.
@@ -172,7 +172,7 @@ class LangfuseConnector:
         tracing.enable_tracing(self.tracer)
 
     @component.output_types(name=str, trace_url=str, trace_id=str)
-    def run(self, invocation_context: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
+    def run(self, invocation_context: Optional[dict[str, Any]] = None) -> dict[str, str]:
         """
         Runs the LangfuseConnector component.
 
@@ -191,7 +191,7 @@ class LangfuseConnector:
         )
         return {"name": self.name, "trace_url": self.tracer.get_trace_url(), "trace_id": self.tracer.get_trace_id()}
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """
         Serialize this component to a dictionary.
 
@@ -218,7 +218,7 @@ class LangfuseConnector:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "LangfuseConnector":
+    def from_dict(cls, data: dict[str, Any]) -> "LangfuseConnector":
         """
         Deserialize this component from a dictionary.
 

{langfuse_haystack-3.2.1 → langfuse_haystack-3.3.1}/src/haystack_integrations/tracing/langfuse/tracer.py

@@ -4,13 +4,15 @@
 
 import contextlib
 import os
+import sys
 from abc import ABC, abstractmethod
 from collections import Counter
+from collections.abc import Iterator
 from contextlib import AbstractContextManager
 from contextvars import ContextVar
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Dict, Iterator, List, Literal, Optional
+from typing import Any, Literal, Optional, cast
 
 from haystack import default_from_dict, default_to_dict, logging
 from haystack.dataclasses import ChatMessage
@@ -25,27 +27,6 @@ from langfuse.types import TraceMetadata
 logger = logging.getLogger(__name__)
 
 HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR = "HAYSTACK_LANGFUSE_ENFORCE_FLUSH"
-_SUPPORTED_GENERATORS = [
-    "AzureOpenAIGenerator",
-    "OpenAIGenerator",
-    "AnthropicGenerator",
-    "HuggingFaceAPIGenerator",
-    "HuggingFaceLocalGenerator",
-    "CohereGenerator",
-    "OllamaGenerator",
-]
-_SUPPORTED_CHAT_GENERATORS = [
-    "AmazonBedrockChatGenerator",
-    "AzureOpenAIChatGenerator",
-    "OpenAIChatGenerator",
-    "AnthropicChatGenerator",
-    "HuggingFaceAPIChatGenerator",
-    "HuggingFaceLocalChatGenerator",
-    "CohereChatGenerator",
-    "OllamaChatGenerator",
-    "GoogleGenAIChatGenerator",
-]
-_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.
@@ -58,13 +39,16 @@ _COMPONENT_TYPE_KEY = "haystack.component.type"
 _COMPONENT_OUTPUT_KEY = "haystack.component.output"
 _COMPONENT_INPUT_KEY = "haystack.component.input"
 
+# Type alias for observation span types
+ObservationSpanType = Literal["tool", "agent", "retriever", "embedding", "generation"]
+
 # External session metadata for trace correlation (Haystack system)
 # Stores trace_id, user_id, session_id, tags, version for root trace creation
-tracing_context_var: ContextVar[Dict[Any, Any]] = ContextVar("tracing_context")
+tracing_context_var: ContextVar[dict[Any, Any]] = ContextVar("tracing_context")
 
 # Internal span execution hierarchy for our tracer
 # Manages parent-child relationships and prevents cross-request span interleaving
-span_stack_var: ContextVar[Optional[List["LangfuseSpan"]]] = ContextVar("span_stack", default=None)
+span_stack_var: ContextVar[Optional[list["LangfuseSpan"]]] = ContextVar("span_stack", default=None)
 
 
 class LangfuseSpan(Span):
@@ -81,7 +65,7 @@ class LangfuseSpan(Span):
         `langfuse.get_client().start_as_current_observation`.
         """
         self._span = context_manager.__enter__()
-        self._data: Dict[str, Any] = {}
+        self._data: dict[str, Any] = {}
         self._context_manager = context_manager
 
     def set_tag(self, key: str, value: Any) -> None:
@@ -132,7 +116,7 @@ class LangfuseSpan(Span):
         """
         return self._span
 
-    def get_data(self) -> Dict[str, Any]:
+    def get_data(self) -> dict[str, Any]:
         """
         Return the data associated with the span.
 
@@ -140,7 +124,7 @@ class LangfuseSpan(Span):
         """
         return self._data
 
-    def get_correlation_data_for_logs(self) -> Dict[str, Any]:
+    def get_correlation_data_for_logs(self) -> dict[str, Any]:
         return {}
 
 
@@ -167,7 +151,7 @@ class SpanContext:
     name: str
     operation_name: str
     component_type: Optional[str]
-    tags: Dict[str, Any]
+    tags: dict[str, Any]
     parent_span: Optional[Span]
     trace_name: str = "Haystack"
     public: bool = False
@@ -249,50 +233,39 @@ class SpanHandler(ABC):
         pass
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "SpanHandler":
+    def from_dict(cls, data: dict[str, Any]) -> "SpanHandler":
         return default_from_dict(cls, data)
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         return default_to_dict(self)
 
 
-def _sanitize_usage_data(usage: Dict[str, Any]) -> Dict[str, Any]:
+def _sanitize_usage_data(usage: dict[str, Any]) -> dict[str, Any]:
     """
-    Sanitize usage data for Langfuse by flattening to a single-level dictionary.
+    Sanitize usage data for Langfuse by converting provider-specific keys to Langfuse standard keys.
 
-    Langfuse's usage_details must be a flat dictionary with only numeric values. This function:
-    - Flattens nested dictionaries using dot notation (e.g., cache_creation.input_tokens)
-    - Keeps int and float values
-    - Skips None, boolean, string, and other non-numeric types
+    Langfuse expects usage_details with standard keys: input_tokens, output_tokens, and total_tokens.
+    This function converts provider-specific keys to Langfuse's expected format:
+    - prompt_tokens -> input_tokens
+    - completion_tokens -> output_tokens
+    - total_tokens -> total_tokens (preserved as-is)
 
     :param usage: Raw usage dictionary from the provider.
-    :returns: Flat dictionary with only numeric values (int or float).
+    :returns: Dictionary with Langfuse standard keys (input_tokens, output_tokens, total_tokens).
     """
     if not isinstance(usage, dict):
         return {}
 
-    sanitized: Dict[str, Any] = {}
-
-    def _flatten(data: Dict[str, Any], prefix: str = "") -> None:
-        """Recursively flatten nested dictionaries."""
-        for key, value in data.items():
-            full_key = f"{prefix}.{key}" if prefix else key
-
-            if value is None:
-                # Skip None values (e.g., Anthropic's server_tool_use)
-                continue
-            elif isinstance(value, bool):
-                # Skip boolean values
-                continue
-            elif isinstance(value, (int, float)):
-                # Keep numeric values
-                sanitized[full_key] = value
-            elif isinstance(value, dict):
-                # Recursively flatten nested dicts
-                _flatten(value, full_key)
-            # Skip strings and other non-numeric types (e.g., Anthropic's service_tier)
-
-    _flatten(usage)
+    # Start with Langfuse standard keys from usage if present
+    sanitized: dict[str, Any] = {
+        k: v for k, v in usage.items() if k in ("input_tokens", "output_tokens", "total_tokens")
+    }
+    # Convert provider format to Langfuse standard keys if not already present
+    if "input_tokens" not in sanitized and "prompt_tokens" in usage:
+        sanitized["input_tokens"] = usage["prompt_tokens"]
+    if "output_tokens" not in sanitized and "completion_tokens" in usage:
+        sanitized["output_tokens"] = usage["completion_tokens"]
+
     return sanitized
 
 
@@ -316,7 +289,9 @@ class DefaultSpanHandler(SpanHandler):
             )
             # Create a new trace when there's no parent span
             span_context_manager = self.tracer.start_as_current_observation(
-                name=context.trace_name, version=tracing_ctx.get("version"), as_type=root_span_type
+                name=context.trace_name,
+                version=tracing_ctx.get("version"),
+                as_type=root_span_type,
             )
 
             # Create LangfuseSpan which will handle entering the context manager
@@ -340,12 +315,26 @@ class DefaultSpanHandler(SpanHandler):
                 span._span.update_trace(**trace_attrs)
 
             return span
-        elif context.component_type == "ToolInvoker":
-            return LangfuseSpan(self.tracer.start_as_current_observation(name=context.name, as_type="tool"))
+
+        span_type = None
+
+        if context.component_type == "ToolInvoker":
+            span_type = "tool"
         elif context.operation_name == "haystack.agent.run":
-            return LangfuseSpan(self.tracer.start_as_current_observation(name=context.name, as_type="agent"))
-        elif context.component_type in _ALL_SUPPORTED_GENERATORS:
-            return LangfuseSpan(self.tracer.start_as_current_observation(name=context.name, as_type="generation"))
+            span_type = "agent"
+        elif context.component_type and context.component_type.endswith("Retriever"):
+            span_type = "retriever"
+        elif context.component_type and context.component_type.endswith("Embedder"):
+            span_type = "embedding"
+        elif context.component_type and context.component_type.endswith("Generator"):
+            span_type = "generation"
+
+        if span_type:
+            return LangfuseSpan(
+                self.tracer.start_as_current_observation(
+                    name=context.name, as_type=cast(ObservationSpanType, span_type)
+                )
+            )
         else:
             return LangfuseSpan(self.tracer.start_as_current_span(name=context.name))
 
@@ -358,7 +347,7 @@ class DefaultSpanHandler(SpanHandler):
             span.raw_span().update_trace(input=coerced_input, output=coerced_output)
         # special case for ToolInvoker (to update the span name to be: `original_component_name - [tool_names]`)
         if component_type == "ToolInvoker":
-            tool_names: List[str] = []
+            tool_names: list[str] = []
             messages = span.get_data().get(_COMPONENT_INPUT_KEY, {}).get("messages", [])
             for message in messages:
                 if isinstance(message, ChatMessage) and message.tool_calls:
@@ -371,14 +360,7 @@ class DefaultSpanHandler(SpanHandler):
                 formatted_names = [f"{name} (x{count})" if count > 1 else name for name, count in tool_counts.items()]
                 span.raw_span().update(name=f"{tool_invoker_name} - {sorted(formatted_names)}")
 
-        if component_type in _SUPPORTED_GENERATORS:
-            meta = span.get_data().get(_COMPONENT_OUTPUT_KEY, {}).get("meta")
-            if meta:
-                usage = meta[0].get("usage")
-                sanitized_usage = _sanitize_usage_data(usage) if usage else None
-                span.raw_span().update(usage_details=sanitized_usage, model=meta[0].get("model"))
-
-        if component_type in _SUPPORTED_CHAT_GENERATORS:
+        if component_type and component_type.endswith("ChatGenerator"):
             replies = span.get_data().get(_COMPONENT_OUTPUT_KEY, {}).get("replies")
             if replies:
                 meta = replies[0].meta
@@ -396,6 +378,36 @@ class DefaultSpanHandler(SpanHandler):
                     model=meta.get("model"),
                     completion_start_time=completion_start_time,
                 )
+        elif component_type and component_type.endswith("Generator"):
+            meta = span.get_data().get(_COMPONENT_OUTPUT_KEY, {}).get("meta")
+            if meta:
+                usage = meta[0].get("usage")
+                sanitized_usage = _sanitize_usage_data(usage) if usage else None
+                span.raw_span().update(usage_details=sanitized_usage, model=meta[0].get("model"))
+        elif component_type and component_type.endswith("Embedder"):
+            # Extract usage data from embedder output
+            output = span.get_data().get(_COMPONENT_OUTPUT_KEY, {})
+            meta = output.get("meta")
+
+            if meta and isinstance(meta, dict):
+                # Build update parameters with available data
+                update_params: dict[str, Any] = {}
+
+                # Try both common formats: 'usage' (OpenAI) or 'billed_units' (Cohere)
+                usage = meta.get("usage") or meta.get("billed_units")
+                if usage:
+                    sanitized_usage = _sanitize_usage_data(usage)
+                    if sanitized_usage:
+                        update_params["usage_details"] = sanitized_usage
+
+                # Some embedders may provide model information
+                model = meta.get("model")
+                if model and isinstance(model, str):
+                    update_params["model"] = model
+
+                # Single update call if we have data to update
+                if update_params:
+                    span.raw_span().update(**update_params)
 
 
 class LangfuseTracer(Tracer):
@@ -429,7 +441,7 @@ class LangfuseTracer(Tracer):
             )
         self._tracer = tracer
         # Keep _context as deprecated shim to avoid AttributeError if anyone uses it
-        self._context: List[LangfuseSpan] = []
+        self._context: list[LangfuseSpan] = []
         self._name = name
         self._public = public
         self.enforce_flush = os.getenv(HAYSTACK_LANGFUSE_ENFORCE_FLUSH_ENV_VAR, "true").lower() == "true"
@@ -438,7 +450,7 @@ class LangfuseTracer(Tracer):
 
     @contextlib.contextmanager
     def trace(
-        self, operation_name: str, tags: Optional[Dict[str, Any]] = None, parent_span: Optional[Span] = None
+        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)
@@ -470,16 +482,44 @@ class LangfuseTracer(Tracer):
 
         try:
             yield span
-        finally:
-            # Always clean up context, even if nested operations fail
+        except Exception:
+            # Exception occurred - capture exception info and pass to __exit__
+            # This allows Langfuse/OpenTelemetry to properly mark the span with ERROR level
+            exc_info = sys.exc_info()
             try:
                 # Process span data (may fail with nested pipeline exceptions)
                 self._span_handler.handle(span, component_type)
 
-                # End span (may fail if span data is corrupted)
+                # End span with exception info (may fail if span data is corrupted)
+                raw_span = span.raw_span()
+                if span._context_manager is not None:
+                    # Pass actual exception info to mark span as failed with ERROR level
+                    span._context_manager.__exit__(*exc_info)
+                elif hasattr(raw_span, "end"):
+                    # Only call end() if it's not a context manager
+                    raw_span.end()
+            except Exception as cleanup_error:
+                # Log cleanup errors but don't let them corrupt context
+                logger.warning(
+                    "Error during span cleanup for {operation_name}: {cleanup_error}",
+                    operation_name=operation_name,
+                    cleanup_error=cleanup_error,
+                )
+
+            # Re-raise the original exception
+            raise
+        else:
+            # No exception - clean exit with success status
+            # This preserves any manually-set log levels (WARNING, DEBUG)
+            try:
+                # Process span data
+                self._span_handler.handle(span, component_type)
+
+                # End span successfully
                 raw_span = span.raw_span()
                 # In v3, we need to properly exit context managers
                 if span._context_manager is not None:
+                    # No exception - pass None to indicate success
                     span._context_manager.__exit__(None, None, None)
                 elif hasattr(raw_span, "end"):
                     # Only call end() if it's not a context manager
@@ -491,9 +531,9 @@ class LangfuseTracer(Tracer):
                     operation_name=operation_name,
                     cleanup_error=cleanup_error,
                 )
-            finally:
-                # Restore previous span stack using saved token - ensures proper cleanup
-                span_stack_var.reset(token)
+        finally:
+            # Restore previous span stack using saved token
+            span_stack_var.reset(token)
 
             if self.enforce_flush:
                 self.flush()

{langfuse_haystack-3.2.1 → langfuse_haystack-3.3.1}/tests/test_tracer.py

@@ -204,20 +204,16 @@ class TestSanitizeUsageData:
             "completion_tokens": 449,
         }
         result = _sanitize_usage_data(usage)
-        assert result == {
-            "cache_creation.ephemeral_1h_input_tokens": 0,
-            "cache_creation.ephemeral_5m_input_tokens": 0,
-            "cache_creation_input_tokens": 0,
-            "cache_read_input_tokens": 0,
-            "prompt_tokens": 25,
-            "completion_tokens": 449,
-        }
+        assert result["input_tokens"] == 25
+        assert result["output_tokens"] == 449
 
     def test_openai_usage_preserved(self):
         """Test OpenAI/Cohere flat dict with only numeric values works unchanged"""
         usage = {"prompt_tokens": 29, "completion_tokens": 267, "total_tokens": 296}
         result = _sanitize_usage_data(usage)
-        assert result == {"prompt_tokens": 29, "completion_tokens": 267, "total_tokens": 296}
+        assert result["input_tokens"] == 29
+        assert result["output_tokens"] == 267
+        assert result["total_tokens"] == 296
 
     def test_empty_and_invalid_input(self):
         """Test edge cases return empty dict"""
@@ -295,6 +291,207 @@ class TestDefaultSpanHandler:
             "completion_start_time": None,
         }
 
+    def test_create_span_custom_chat_generator(self):
+        """Test that custom chat generators create 'generation' span type."""
+        mock_client = Mock()
+        mock_client.start_as_current_span = Mock(return_value=MockContextManager())
+        mock_client.start_as_current_observation = Mock(return_value=MockContextManager())
+
+        handler = DefaultSpanHandler()
+        handler.init_tracer(mock_client)
+
+        context = SpanContext(
+            name="MistralChatGenerator",
+            operation_name="haystack.component.run",
+            component_type="MistralChatGenerator",
+            tags={},
+            parent_span=LangfuseSpan(mock_client.start_as_current_span()),
+        )
+
+        span = handler.create_span(context)
+        assert isinstance(span, LangfuseSpan)
+        mock_client.start_as_current_observation.assert_called_once_with(
+            name="MistralChatGenerator", as_type="generation"
+        )
+
+    def test_create_span_custom_generator(self):
+        """Test that custom generators create 'generation' span type."""
+        mock_client = Mock()
+        mock_client.start_as_current_span = Mock(return_value=MockContextManager())
+        mock_client.start_as_current_observation = Mock(return_value=MockContextManager())
+
+        handler = DefaultSpanHandler()
+        handler.init_tracer(mock_client)
+
+        context = SpanContext(
+            name="CustomAPIGenerator",
+            operation_name="haystack.component.run",
+            component_type="CustomAPIGenerator",
+            tags={},
+            parent_span=LangfuseSpan(mock_client.start_as_current_span()),
+        )
+
+        span = handler.create_span(context)
+        assert isinstance(span, LangfuseSpan)
+        mock_client.start_as_current_observation.assert_called_once_with(
+            name="CustomAPIGenerator", as_type="generation"
+        )
+
+    def test_create_span_retriever(self):
+        """Test that retrievers create 'retriever' span type."""
+        mock_client = Mock()
+        mock_client.start_as_current_span = Mock(return_value=MockContextManager())
+        mock_client.start_as_current_observation = Mock(return_value=MockContextManager())
+
+        handler = DefaultSpanHandler()
+        handler.init_tracer(mock_client)
+
+        context = SpanContext(
+            name="InMemoryBM25Retriever",
+            operation_name="haystack.component.run",
+            component_type="InMemoryBM25Retriever",
+            tags={},
+            parent_span=LangfuseSpan(mock_client.start_as_current_span()),
+        )
+
+        span = handler.create_span(context)
+        assert isinstance(span, LangfuseSpan)
+        mock_client.start_as_current_observation.assert_called_once_with(
+            name="InMemoryBM25Retriever", as_type="retriever"
+        )
+
+    def test_create_span_embedder(self):
+        """Test that embedders create 'embedding' span type."""
+        mock_client = Mock()
+        mock_client.start_as_current_span = Mock(return_value=MockContextManager())
+        mock_client.start_as_current_observation = Mock(return_value=MockContextManager())
+
+        handler = DefaultSpanHandler()
+        handler.init_tracer(mock_client)
+
+        context = SpanContext(
+            name="SentenceTransformersDocumentEmbedder",
+            operation_name="haystack.component.run",
+            component_type="SentenceTransformersDocumentEmbedder",
+            tags={},
+            parent_span=LangfuseSpan(mock_client.start_as_current_span()),
+        )
+
+        span = handler.create_span(context)
+        assert isinstance(span, LangfuseSpan)
+        mock_client.start_as_current_observation.assert_called_once_with(
+            name="SentenceTransformersDocumentEmbedder", as_type="embedding"
+        )
+
+    def test_create_span_non_component(self):
+        """Test that non-matching components create regular spans."""
+        mock_client = Mock()
+        mock_client.start_as_current_span = Mock(return_value=MockContextManager())
+        mock_client.start_as_current_observation = Mock(return_value=MockContextManager())
+
+        handler = DefaultSpanHandler()
+        handler.init_tracer(mock_client)
+
+        context = SpanContext(
+            name="DocumentJoiner",
+            operation_name="haystack.component.run",
+            component_type="DocumentJoiner",
+            tags={},
+            parent_span=LangfuseSpan(mock_client.start_as_current_span()),
+        )
+
+        span = handler.create_span(context)
+        assert isinstance(span, LangfuseSpan)
+        # Non-matching components should use start_as_current_span, not start_as_current_observation
+        mock_client.start_as_current_observation.assert_not_called()
+        # Verify start_as_current_span was called for the actual span creation (not just parent)
+        assert mock_client.start_as_current_span.call_count == 2  # Once for parent, once for the span
+
+    def test_handle_embedder_with_openai_format(self):
+        """Test that embedder usage is extracted in OpenAI format."""
+        mock_span = Mock()
+        mock_span.raw_span.return_value = mock_span
+        mock_span.get_data.return_value = {
+            "haystack.component.type": "OpenAITextEmbedder",
+            "haystack.component.output": {
+                "embedding": [0.1, 0.2, 0.3],
+                "meta": {"model": "custom-model", "usage": {"prompt_tokens": 15, "total_tokens": 15}},
+            },
+        }
+
+        handler = DefaultSpanHandler()
+        handler.handle(mock_span, component_type="OpenAITextEmbedder")
+
+        assert mock_span.update.call_count == 1
+        update_args = mock_span.update.call_args_list[0][1]
+        assert update_args["model"] == "custom-model"
+        assert update_args["usage_details"] == {"input_tokens": 15, "total_tokens": 15}
+
+    def test_handle_embedder_with_cohere_format(self):
+        """Test that embedder usage is extracted in Cohere billed_units format."""
+        mock_span = Mock()
+        mock_span.raw_span.return_value = mock_span
+        mock_span.get_data.return_value = {
+            "haystack.component.type": "CohereTextEmbedder",
+            "haystack.component.output": {
+                "embedding": [0.1, 0.2, 0.3],
+                "meta": {"api_version": {"version": "1"}, "billed_units": {"input_tokens": 4}},
+            },
+        }
+
+        handler = DefaultSpanHandler()
+        handler.handle(mock_span, component_type="CohereTextEmbedder")
+
+        assert mock_span.update.call_count == 1
+        assert mock_span.update.call_args_list[0][1] == {"usage_details": {"input_tokens": 4}}
+
+    def test_handle_embedder_without_usage(self):
+        """Test that embedders without usage data are handled gracefully."""
+        mock_span = Mock()
+        mock_span.raw_span.return_value = mock_span
+        mock_span.get_data.return_value = {
+            "haystack.component.type": "SentenceTransformersTextEmbedder",
+            "haystack.component.output": {
+                "embedding": [0.1, 0.2, 0.3],
+                "meta": {},  # No usage data
+            },
+        }
+
+        handler = DefaultSpanHandler()
+        handler.handle(mock_span, component_type="SentenceTransformersTextEmbedder")
+
+        # Should not call update when no usage data is available
+        assert mock_span.update.call_count == 0
+
+    def test_handle_embedder_with_nested_usage(self):
+        """Test that embedders with nested usage data are sanitized correctly."""
+        mock_span = Mock()
+        mock_span.raw_span.return_value = mock_span
+        mock_span.get_data.return_value = {
+            "haystack.component.type": "CustomEmbedder",
+            "haystack.component.output": {
+                "embedding": [0.1, 0.2, 0.3],
+                "meta": {
+                    "model": "custom-model",
+                    "usage": {
+                        "cache_creation": {"input_tokens": 10},
+                        "cache_read": {"input_tokens": 5},
+                        "total_tokens": 15,
+                    },
+                },
+            },
+        }
+
+        handler = DefaultSpanHandler()
+        handler.handle(mock_span, component_type="CustomEmbedder")
+
+        assert mock_span.update.call_count == 1
+        # Only adds total_tokens as Langfuse standard key (no prompt_tokens/completion_tokens to convert)
+        assert mock_span.update.call_args_list[0][1] == {
+            "usage_details": {"total_tokens": 15},
+            "model": "custom-model",
+        }
+
 
 class TestCustomSpanHandler:
     def test_handle(self):
@@ -338,13 +535,22 @@ class TestLangfuseTracer:
         mock_raw_span.metadata = {"tag1": "value1", "tag2": "value2"}
 
         with patch("haystack_integrations.tracing.langfuse.tracer.LangfuseSpan") as mock_langfuse_span:
+            mock_context_manager = MockContextManager()
+            mock_context_manager._span = mock_raw_span
+
             mock_span_instance = mock_langfuse_span.return_value
             mock_span_instance.raw_span.return_value = mock_raw_span
+            # Return a proper dict to prevent MagicMock from being truthy in handle() checks.
+            # When get_data() returns a MagicMock, `span.get_data().get(key) is not None` is True
+            # because MagicMock().get() returns another MagicMock (truthy). This triggers
+            # tracing_utils.coerce_tag_value() with MagicMock objects, which can hang on
+            # Linux Python 3.9/3.13 due to platform-specific MagicMock iteration behavior.
+            mock_span_instance.get_data.return_value = {}
+            mock_span_instance._context_manager = mock_context_manager
 
-            mock_context_manager = MockContextManager()
-            mock_context_manager._span = mock_raw_span
             mock_tracer = MagicMock()
             mock_tracer.start_as_current_span.return_value = mock_context_manager
+            mock_tracer.start_as_current_observation.return_value = mock_context_manager
 
             tracer = LangfuseTracer(tracer=mock_tracer, name="Haystack", public=False)
 
@@ -574,3 +780,54 @@ class TestLangfuseTracer:
         assert task2_spans[1][2] == task2_inner  # current_span during inner
         assert task2_spans[2][2] == task2_outer  # current_span after inner
         assert task2_spans[3][2] is None  # current_span after outer
+
+    def test_trace_exception_handling(self):
+        """
+        Test that exceptions are properly captured and passed to span __exit__.
+
+        This verifies the new exception handling behavior where:
+        - Exception case: __exit__() receives (exc_type, exc_val, exc_tb)
+        - Success case: __exit__() receives (None, None, None)
+        """
+        # Create a mock context manager that tracks how __exit__ was called
+        mock_exit_calls = []
+
+        class TrackingContextManager:
+            def __init__(self):
+                self._span = MockSpan()
+
+            def __enter__(self):
+                return self._span
+
+            def __exit__(self, exc_type, exc_val, exc_tb):
+                # Track what was passed to __exit__
+                mock_exit_calls.append((exc_type, exc_val, exc_tb))
+                return False  # Don't suppress exceptions
+
+        mock_client = MockLangfuseClient()
+        mock_client._mock_context_manager = TrackingContextManager()
+
+        tracer = LangfuseTracer(tracer=mock_client, name="Test", public=False)
+
+        # Test 1: Exception case - __exit__ should receive exception info
+        mock_exit_calls.clear()
+        error_msg = "test error"
+        with pytest.raises(ValueError, match="test error"):
+            with tracer.trace("test_operation"):
+                raise ValueError(error_msg)
+
+        assert len(mock_exit_calls) == 1
+        assert mock_exit_calls[0][0] is ValueError  # exc_type
+        assert str(mock_exit_calls[0][1]) == error_msg  # exc_val
+        assert mock_exit_calls[0][2] is not None  # exc_tb (traceback)
+
+        # Test 2: Success case - __exit__ should receive (None, None, None)
+        mock_exit_calls.clear()
+        with tracer.trace("test_operation"):
+            pass  # No exception
+
+        assert len(mock_exit_calls) == 1
+        assert mock_exit_calls[0] == (None, None, None)
+
+        # Test 3: Verify span stack is cleaned up after exception
+        assert tracer.current_span() is None

{langfuse_haystack-3.2.1 → langfuse_haystack-3.3.1}/tests/test_tracing.py

@@ -5,7 +5,7 @@
 import json
 import os
 import time
-from typing import Any, Dict, List
+from typing import Any
 from urllib.parse import urlparse
 
 import pytest
@@ -38,7 +38,7 @@ os.environ.setdefault("LANGFUSE_HOST", "https://cloud.langfuse.com")
 def poll_langfuse(url: str):
     """Utility function to poll Langfuse API until the trace is ready"""
     # Initial wait for trace creation
-    time.sleep(10)
+    time.sleep(30)
 
     auth = HTTPBasicAuth(os.environ["LANGFUSE_PUBLIC_KEY"], os.environ["LANGFUSE_SECRET_KEY"])
 
@@ -137,8 +137,8 @@ def test_tracing_with_sub_pipelines():
             self.sub_pipeline = Pipeline()
             self.sub_pipeline.add_component("llm", OpenAIChatGenerator())
 
-        @component.output_types(replies=List[ChatMessage])
-        def run(self, messages: List[ChatMessage]) -> Dict[str, Any]:
+        @component.output_types(replies=list[ChatMessage])
+        def run(self, messages: list[ChatMessage]) -> dict[str, Any]:
             return {"replies": self.sub_pipeline.run(data={"llm": {"messages": messages}})["llm"]["replies"]}
 
     @component
@@ -149,8 +149,8 @@ def test_tracing_with_sub_pipelines():
             self.sub_pipeline.add_component("sub_llm", SubGenerator())
             self.sub_pipeline.connect("prompt_builder.prompt", "sub_llm.messages")
 
-        @component.output_types(replies=List[ChatMessage])
-        def run(self, messages: List[ChatMessage]) -> Dict[str, Any]:
+        @component.output_types(replies=list[ChatMessage])
+        def run(self, messages: list[ChatMessage]) -> dict[str, Any]:
             return {
                 "replies": self.sub_pipeline.run(
                     data={"prompt_builder": {"template": messages, "template_variables": {"location": "Berlin"}}}
@@ -195,8 +195,8 @@ def test_tracing_with_sub_pipelines():
     # There should be two observations for the haystack.pipeline.run span: one for each sub pipeline
     # Main pipeline is stored under the name "Sub-pipeline example"
     assert len(haystack_pipeline_run_observations) == 2
-    assert "prompt_builder" in str(haystack_pipeline_run_observations[0])
-    assert "llm" in str(haystack_pipeline_run_observations[1])
+    # Verify both observations are pipeline runs (less brittle than checking for component names)
+    assert all(obs["name"] == "haystack.pipeline.run" for obs in haystack_pipeline_run_observations)
 
 
 @pytest.mark.skipif(

langfuse_haystack-3.2.1/pydoc/config.yml

@@ -1,30 +0,0 @@
-loaders:
-  - type: haystack_pydoc_tools.loaders.CustomPythonLoader
-    search_path: [../src]
-    modules: [
-      "haystack_integrations.components.connectors.langfuse.langfuse_connector",
-      "haystack_integrations.tracing.langfuse.tracer",
-    ]
-    ignore_when_discovered: ["__init__"]
-processors:
-  - type: filter
-    expression:
-    documented_only: true
-    do_not_filter_modules: false
-    skip_empty_modules: true
-  - type: smart
-  - type: crossref
-renderer:
-  type: haystack_pydoc_tools.renderers.ReadmeIntegrationRenderer
-  excerpt: Langfuse integration for Haystack
-  category_slug: integrations-api
-  title: langfuse
-  slug: integrations-langfuse
-  order: 136
-  markdown:
-    descriptive_class_title: false
-    classdef_code_block: false
-    descriptive_module_title: true
-    add_method_class_prefix: true
-    add_member_class_prefix: false
-    filename: _readme_langfuse.md