feat: add internal tracing to components that use chat generators (#12075)

Co-authored-by: David S. Batista <dsbatista@gmail.com>

Author: sjrl

haystack/components/evaluators/llm_evaluator.py

@@ -13,6 +13,7 @@
 from haystack.components.builders import PromptBuilder
 from haystack.components.generators.chat.openai import OpenAIChatGenerator
 from haystack.components.generators.chat.types import ChatGenerator
+from haystack.components.generators.utils import _trace_chat_generator_run
 from haystack.core.serialization import component_to_dict
 from haystack.dataclasses.chat_message import ChatMessage
 from haystack.utils import deserialize_chatgenerator_inplace, deserialize_type, serialize_type
@@ -232,7 +233,9 @@ def run(self, **inputs: Any) -> dict[str, Any]:
             prompt = self.builder.run(**input_names_to_values)
             messages = [ChatMessage.from_user(prompt["prompt"])]
             try:
-                result = self._chat_generator.run(messages=messages)
+                with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span:
+                    result = self._chat_generator.run(messages=messages)
+                    span.set_content_tag("haystack.component.output", result)
             except Exception as e:
                 if self.raise_on_failure:
                     raise ValueError(f"Error while generating response for prompt: {prompt}. Error: {e}") from e
@@ -301,15 +304,17 @@ async def run_async(self, **inputs: Any) -> dict[str, Any]:
             prompt = self.builder.run(**input_names_to_values)
             messages = [ChatMessage.from_user(prompt["prompt"])]
             try:
-                if generator_has_async:
-                    result = await self._chat_generator.run_async(messages=messages)  # type: ignore[attr-defined]
-                else:
-                    logger.debug(
-                        "{generator_type} does not implement 'run_async'."
-                        " Running the synchronous 'run' method in a thread to avoid blocking the event loop.",
-                        generator_type=type(self._chat_generator).__name__,
-                    )
-                    result = await asyncio.to_thread(self._chat_generator.run, messages=messages)
+                with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span:
+                    if generator_has_async:
+                        result = await self._chat_generator.run_async(messages=messages)  # type: ignore[attr-defined]
+                    else:
+                        logger.debug(
+                            "{generator_type} does not implement 'run_async'."
+                            " Running the synchronous 'run' method in a thread to avoid blocking the event loop.",
+                            generator_type=type(self._chat_generator).__name__,
+                        )
+                        result = await asyncio.to_thread(self._chat_generator.run, messages=messages)
+                    span.set_content_tag("haystack.component.output", result)
             except Exception as e:
                 if self.raise_on_failure:
                     raise ValueError(f"Error while generating response for prompt: {prompt}. Error: {e}") from e

haystack/components/extractors/image/llm_document_content_extractor.py

@@ -6,14 +6,16 @@
 import json
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import replace
+from functools import partial
 from typing import Any, Literal
 
 from jinja2 import meta
 from jinja2.sandbox import SandboxedEnvironment
 
-from haystack import Document, component, default_from_dict, default_to_dict, logging
+from haystack import Document, component, default_from_dict, default_to_dict, logging, tracing
 from haystack.components.converters.image.document_to_image import DocumentToImageContent
 from haystack.components.generators.chat.types import ChatGenerator
+from haystack.components.generators.utils import _trace_chat_generator_run
 from haystack.core.serialization import component_to_dict
 from haystack.dataclasses import ImageContent, TextContent
 from haystack.dataclasses.chat_message import ChatMessage
@@ -267,11 +269,14 @@ def _process_response(response_text: str) -> tuple[str | None, dict[str, Any], s
         meta_updates = {k: v for k, v in parsed.items() if k != DOCUMENT_CONTENT_KEY}
         return content, meta_updates, None
 
-    def _run_on_thread(self, image_content: ImageContent | None) -> dict[str, Any]:
+    def _run_on_thread(
+        self, image_content: ImageContent | None, parent_span: tracing.Span | None = None
+    ) -> dict[str, Any]:
         """
         Execute the LLM inference in a separate thread for each document.
 
         :param image_content: The image content for one document, or None if conversion failed.
+        :param parent_span: Span to nest the generator span under, captured on the calling thread.
         :returns:
             The LLM response if successful, or a dictionary with an "error" key on failure.
         """
@@ -282,7 +287,11 @@ def _run_on_thread(self, image_content: ImageContent | None) -> dict[str, Any]:
         message = ChatMessage.from_user(content_parts=[TextContent(text=self.prompt), image_content])
 
         try:
-            result = self._chat_generator.run(messages=[message])
+            with _trace_chat_generator_run(
+                self._chat_generator, {"messages": [message]}, parent_span=parent_span
+            ) as span:
+                result = self._chat_generator.run(messages=[message])
+                span.set_content_tag("haystack.component.output", result)
         except Exception as e:
             if self.raise_on_failure:
                 raise e
@@ -295,11 +304,14 @@ def _run_on_thread(self, image_content: ImageContent | None) -> dict[str, Any]:
 
         return result
 
-    async def _run_async(self, image_content: ImageContent | None) -> dict[str, Any]:
+    async def _run_async(
+        self, image_content: ImageContent | None, parent_span: tracing.Span | None = None
+    ) -> dict[str, Any]:
         """
         Execute the LLM inference asynchronously for each document.
 
         :param image_content: The image content for one document, or None if conversion failed.
+        :param parent_span: Span to nest the generator span under, captured on the calling task.
         :returns:
             The LLM response if successful, or a dictionary with an "error" key on failure.
         """
@@ -310,7 +322,11 @@ async def _run_async(self, image_content: ImageContent | None) -> dict[str, Any]
         message = ChatMessage.from_user(content_parts=[TextContent(text=self.prompt), image_content])
 
         try:
-            result = await _execute_component_async(self._chat_generator, messages=[message])
+            with _trace_chat_generator_run(
+                self._chat_generator, {"messages": [message]}, parent_span=parent_span
+            ) as span:
+                result = await _execute_component_async(self._chat_generator, messages=[message])
+                span.set_content_tag("haystack.component.output", result)
         except Exception as e:
             if self.raise_on_failure:
                 raise e
@@ -366,8 +382,11 @@ def run(self, documents: list[Document]) -> dict[str, list[Document]]:
 
         image_contents = self._document_to_image_content.run(documents=documents)["image_contents"]
 
+        # Capture the current span here so worker threads nest their generator spans under the component span.
+        parent_span = tracing.tracer.current_span()
+
         with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            results = executor.map(self._run_on_thread, image_contents)
+            results = executor.map(partial(self._run_on_thread, parent_span=parent_span), image_contents)
 
         successful_documents = []
         failed_documents = []
@@ -401,12 +420,15 @@ async def run_async(self, documents: list[Document]) -> dict[str, list[Document]
 
         image_contents = self._document_to_image_content.run(documents=documents)["image_contents"]
 
+        # Capture the current span here so concurrent tasks nest their generator spans under the component span.
+        parent_span = tracing.tracer.current_span()
+
         # Run the LLM on each image content, bounding concurrency per task so max_workers is enforced.
         sem = asyncio.Semaphore(max(1, self.max_workers))
 
         async def _bounded_run(image_content: ImageContent | None) -> dict[str, Any]:
             async with sem:
-                return await self._run_async(image_content)
+                return await self._run_async(image_content, parent_span=parent_span)
 
         results = await asyncio.gather(*[_bounded_run(image_content) for image_content in image_contents])
 

haystack/components/extractors/llm_metadata_extractor.py

@@ -8,14 +8,16 @@
 from collections.abc import Iterable
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import replace
+from functools import partial
 from typing import Any
 
 from jinja2 import meta
 from jinja2.sandbox import SandboxedEnvironment
 
-from haystack import Document, component, default_from_dict, default_to_dict, logging
+from haystack import Document, component, default_from_dict, default_to_dict, logging, tracing
 from haystack.components.builders import PromptBuilder
 from haystack.components.generators.chat.types import ChatGenerator
+from haystack.components.generators.utils import _trace_chat_generator_run
 from haystack.components.preprocessors import DocumentSplitter
 from haystack.core.serialization import component_to_dict
 from haystack.dataclasses import ChatMessage
@@ -309,13 +311,17 @@ def _prepare_prompts(
 
         return all_prompts
 
-    def _run_on_thread(self, prompt: ChatMessage | None) -> dict[str, Any]:
+    def _run_on_thread(self, prompt: ChatMessage | None, parent_span: tracing.Span | None = None) -> dict[str, Any]:
         # If prompt is None, return an error dictionary
         if prompt is None:
             return {"error": "Document has no content, skipping LLM call."}
 
         try:
-            result = self._chat_generator.run(messages=[prompt])
+            with _trace_chat_generator_run(
+                self._chat_generator, {"messages": [prompt]}, parent_span=parent_span
+            ) as span:
+                result = self._chat_generator.run(messages=[prompt])
+                span.set_content_tag("haystack.component.output", result)
         except Exception as e:
             if self.raise_on_failure:
                 raise e
@@ -327,13 +333,17 @@ def _run_on_thread(self, prompt: ChatMessage | None) -> dict[str, Any]:
             result = {"error": "LLM failed with exception: " + str(e)}
         return result
 
-    async def _run_async(self, prompt: ChatMessage | None) -> dict[str, Any]:
+    async def _run_async(self, prompt: ChatMessage | None, parent_span: tracing.Span | None = None) -> dict[str, Any]:
         # If prompt is None, return an error dictionary
         if prompt is None:
             return {"error": "Document has no content, skipping LLM call."}
 
         try:
-            result = await _execute_component_async(self._chat_generator, messages=[prompt])
+            with _trace_chat_generator_run(
+                self._chat_generator, {"messages": [prompt]}, parent_span=parent_span
+            ) as span:
+                result = await _execute_component_async(self._chat_generator, messages=[prompt])
+                span.set_content_tag("haystack.component.output", result)
         except Exception as e:
             if self.raise_on_failure:
                 raise e
@@ -411,9 +421,12 @@ def run(self, documents: list[Document], page_range: list[str | int] | None = No
         # Create ChatMessage prompts for each document
         all_prompts = self._prepare_prompts(documents=documents, expanded_range=expanded_range)
 
+        # Capture the current span here so worker threads nest their generator spans under the component span.
+        parent_span = tracing.tracer.current_span()
+
         # Run the LLM on each prompt
         with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            results = executor.map(self._run_on_thread, all_prompts)
+            results = executor.map(partial(self._run_on_thread, parent_span=parent_span), all_prompts)
 
         successful_documents, failed_documents = self._process_results(documents, results)
 
@@ -460,12 +473,15 @@ async def run_async(self, documents: list[Document], page_range: list[str | int]
         # Create ChatMessage prompts for each document
         all_prompts = self._prepare_prompts(documents=documents, expanded_range=expanded_range)
 
+        # Capture the current span here so concurrent tasks nest their generator spans under the component span.
+        parent_span = tracing.tracer.current_span()
+
         # Run the LLM on each prompt, bounding concurrency per task so max_workers is enforced.
         sem = Semaphore(max(1, self.max_workers))
 
         async def _bounded_run(prompt: ChatMessage | None) -> dict[str, Any]:
             async with sem:
-                return await self._run_async(prompt)
+                return await self._run_async(prompt, parent_span=parent_span)
 
         results = await gather(*[_bounded_run(prompt) for prompt in all_prompts])
 

haystack/components/generators/utils.py

@@ -2,15 +2,53 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import contextlib
 import json
+from collections.abc import Iterator
 from typing import Any
 
-from haystack import logging
+from haystack import logging, tracing
 from haystack.dataclasses import ChatMessage, ReasoningContent, StreamingChunk, ToolCall
 
 logger = logging.getLogger(__name__)
 
 
+@contextlib.contextmanager
+def _trace_chat_generator_run(
+    chat_generator: Any, generator_inputs: dict[str, Any], parent_span: tracing.Span | None = None
+) -> Iterator[tracing.Span]:
+    """
+    Open a tracing span around a ChatGenerator call made internally by another component.
+
+    Components that embed a ChatGenerator but do not return its `ChatMessage` replies (for example rankers,
+    extractors or evaluators) would otherwise discard the LLM token usage carried in `reply.meta["usage"]`.
+    Wrapping the internal call in this span re-exposes that usage to tracers via the `haystack.component.output`
+    content tag, mirroring how the `Pipeline` traces its top-level components.
+
+    The caller is responsible for setting the output tag inside the context, so it is skipped when the call fails:
+
+    ```python
+    with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span:
+        result = self._chat_generator.run(messages=messages)
+        span.set_content_tag("haystack.component.output", result)
+    ```
+
+    :param chat_generator: The ChatGenerator being invoked.
+    :param generator_inputs: The inputs passed to the generator, recorded as the span input content tag.
+    :param parent_span: Explicit parent span. Defaults to the current span. Pass it explicitly when the generator
+        runs in a worker thread, where the ambient span context does not propagate.
+    """
+    # Fall back to the active span so same-thread callers get correct nesting without passing a parent explicitly.
+    parent_span = parent_span or tracing.tracer.current_span()
+    with tracing.tracer.trace(
+        "haystack.chat_generator.run",
+        tags={"haystack.component.name": "chat_generator", "haystack.component.type": type(chat_generator).__name__},
+        parent_span=parent_span,
+    ) as span:
+        span.set_content_tag("haystack.component.input", generator_inputs)
+        yield span
+
+
 def print_streaming_chunk(chunk: StreamingChunk) -> None:
     """
     Callback function to handle and display streaming output chunks.

haystack/components/query/query_expander.py

@@ -8,6 +8,7 @@
 from haystack.components.builders.prompt_builder import PromptBuilder
 from haystack.components.generators.chat.openai import OpenAIChatGenerator
 from haystack.components.generators.chat.types import ChatGenerator
+from haystack.components.generators.utils import _trace_chat_generator_run
 from haystack.core.component import component
 from haystack.core.serialization import component_to_dict
 from haystack.dataclasses.chat_message import ChatMessage
@@ -209,7 +210,10 @@ def run(self, query: str, n_expansions: int | None = None) -> dict[str, list[str
 
         try:
             prompt_result = self._prompt_builder.run(query=query.strip(), n_expansions=expansion_count)
-            generator_result = self.chat_generator.run(messages=[ChatMessage.from_user(prompt_result["prompt"])])
+            messages = [ChatMessage.from_user(prompt_result["prompt"])]
+            with _trace_chat_generator_run(self.chat_generator, {"messages": messages}) as span:
+                generator_result = self.chat_generator.run(messages=messages)
+                span.set_content_tag("haystack.component.output", generator_result)
 
             if not generator_result.get("replies") or len(generator_result["replies"]) == 0:
                 logger.warning("ChatGenerator returned no replies for query: {query}", query=query)
@@ -276,9 +280,10 @@ async def run_async(self, query: str, n_expansions: int | None = None) -> dict[s
 
         try:
             prompt_result = self._prompt_builder.run(query=query.strip(), n_expansions=expansion_count)
-            generator_result = await _execute_component_async(
-                self.chat_generator, messages=[ChatMessage.from_user(prompt_result["prompt"])]
-            )
+            messages = [ChatMessage.from_user(prompt_result["prompt"])]
+            with _trace_chat_generator_run(self.chat_generator, {"messages": messages}) as span:
+                generator_result = await _execute_component_async(self.chat_generator, messages=messages)
+                span.set_content_tag("haystack.component.output", generator_result)
 
             if not generator_result.get("replies") or len(generator_result["replies"]) == 0:
                 logger.warning("ChatGenerator returned no replies for query: {query}", query=query)

haystack/components/rankers/llm_ranker.py

@@ -8,6 +8,7 @@
 from haystack.components.builders import PromptBuilder
 from haystack.components.generators.chat.openai import OpenAIChatGenerator
 from haystack.components.generators.chat.types import ChatGenerator
+from haystack.components.generators.utils import _trace_chat_generator_run
 from haystack.core.serialization import component_to_dict
 from haystack.dataclasses import ChatMessage
 from haystack.utils import deserialize_chatgenerator_inplace
@@ -259,9 +260,12 @@ def run(self, query: str, documents: list[Document], top_k: int | None = None) -
         self.warm_up()
 
         prompt = self._prompt_builder.run(query=query.strip(), documents=deduplicated_documents)
+        messages = [ChatMessage.from_user(prompt["prompt"])]
 
         try:
-            result = self._chat_generator.run(messages=[ChatMessage.from_user(prompt["prompt"])])
+            with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span:
+                result = self._chat_generator.run(messages=messages)
+                span.set_content_tag("haystack.component.output", result)
         except Exception as exc:
             if self.raise_on_failure:
                 raise
@@ -323,11 +327,12 @@ async def run_async(
         await self.warm_up_async()
 
         prompt = self._prompt_builder.run(query=query.strip(), documents=deduplicated_documents)
+        messages = [ChatMessage.from_user(prompt["prompt"])]
 
         try:
-            result = await _execute_component_async(
-                self._chat_generator, messages=[ChatMessage.from_user(prompt["prompt"])]
-            )
+            with _trace_chat_generator_run(self._chat_generator, {"messages": messages}) as span:
+                result = await _execute_component_async(self._chat_generator, messages=messages)
+                span.set_content_tag("haystack.component.output", result)
         except Exception as exc:
             if self.raise_on_failure:
                 raise