feat(langfuse): embedder, retriever and generator as obs. type (#2497)

* add: embedder, retriever and generator as obs. type

* add: mock call obs

* add: improved code readability

* add: else span case

* add: static typing

* fix: hatch run fmt

Author: hansehart

integrations/langfuse/src/haystack_integrations/tracing/langfuse/tracer.py

@@ -10,7 +10,7 @@
 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, Dict, Iterator, List, Literal, Optional, cast
 
 from haystack import default_from_dict, default_to_dict, logging
 from haystack.dataclasses import ChatMessage
@@ -25,27 +25,6 @@
 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,6 +37,9 @@
 _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")
@@ -340,12 +322,26 @@ def create_span(self, context: SpanContext) -> LangfuseSpan:
                 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))
 
@@ -371,14 +367,7 @@ def handle(self, span: LangfuseSpan, component_type: Optional[str]) -> None:
                 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 +385,12 @@ def handle(self, span: LangfuseSpan, component_type: Optional[str]) -> None:
                     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"))
 
 
 class LangfuseTracer(Tracer):

integrations/langfuse/tests/test_tracer.py

@@ -295,6 +295,122 @@ def test_handle_bad_completion_start_time(self, caplog):
             "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
+
 
 class TestCustomSpanHandler:
     def test_handle(self):