langfuse-haystack 3.2.0__tar.gz → 3.3.0__tar.gz

{langfuse_haystack-3.2.0 → langfuse_haystack-3.3.0}/CHANGELOG.md

@@ -1,5 +1,26 @@
 # Changelog
 
+## [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
+
+- Flatten usage_details dict (#2491)
+
+### ⚙️ CI
+
+- Change pytest command (#2475)
+
+### 🌀 Miscellaneous
+
+- Usage_details instead usage (#2481)
+- Feat: Langfuse - add support for tool and agent observation types (#2490)
+
 ## [integrations/langfuse-v3.1.0] - 2025-10-24
 
 ### 🐛 Bug Fixes

{langfuse_haystack-3.2.0 → langfuse_haystack-3.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langfuse-haystack
-Version: 3.2.0
+Version: 3.3.0
 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
@@ -19,7 +19,7 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.9
 Requires-Dist: haystack-ai>=2.17.1
-Requires-Dist: langfuse<4.0.0,>=3.3.0
+Requires-Dist: langfuse<4.0.0,>=3.3.1
 Description-Content-Type: text/markdown
 
 # langfuse-haystack

{langfuse_haystack-3.2.0 → langfuse_haystack-3.3.0}/pyproject.toml

@@ -22,7 +22,7 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.17.1", "langfuse>=3.3.0, <4.0.0"]
+dependencies = ["haystack-ai>=2.17.1", "langfuse>=3.3.1, <4.0.0"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/langfuse#readme"
@@ -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.0 → langfuse_haystack-3.3.0}/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.0 → langfuse_haystack-3.3.0}/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,14 +233,14 @@ 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.
 
@@ -271,9 +255,9 @@ def _sanitize_usage_data(usage: Dict[str, Any]) -> Dict[str, Any]:
     if not isinstance(usage, dict):
         return {}
 
-    sanitized: Dict[str, Any] = {}
+    sanitized: dict[str, Any] = {}
 
-    def _flatten(data: Dict[str, Any], prefix: str = "") -> None:
+    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
@@ -316,7 +300,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 +326,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 +358,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 +371,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 +389,12 @@ 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"))
 
 
 class LangfuseTracer(Tracer):
@@ -429,7 +428,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 +437,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 +469,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 +518,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.0 → langfuse_haystack-3.3.0}/tests/test_tracer.py

@@ -295,6 +295,122 @@ 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
+
 
 class TestCustomSpanHandler:
     def test_handle(self):
@@ -574,3 +690,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.0 → langfuse_haystack-3.3.0}/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
@@ -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"}}}