diff --git a/haystack/components/evaluators/llm_evaluator.py b/haystack/components/evaluators/llm_evaluator.py index 505571fb72e..6dd6a5d3595 100644 --- a/haystack/components/evaluators/llm_evaluator.py +++ b/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 diff --git a/haystack/components/extractors/image/llm_document_content_extractor.py b/haystack/components/extractors/image/llm_document_content_extractor.py index d1e09433c53..b96a2e2d0f3 100644 --- a/haystack/components/extractors/image/llm_document_content_extractor.py +++ b/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]) diff --git a/haystack/components/extractors/llm_metadata_extractor.py b/haystack/components/extractors/llm_metadata_extractor.py index abb04e7d8a6..8d8ee17f56b 100644 --- a/haystack/components/extractors/llm_metadata_extractor.py +++ b/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]) diff --git a/haystack/components/generators/utils.py b/haystack/components/generators/utils.py index 25bf764c15e..9ab8b84eb0e 100644 --- a/haystack/components/generators/utils.py +++ b/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. diff --git a/haystack/components/query/query_expander.py b/haystack/components/query/query_expander.py index 5d952a6a25d..5279db12abe 100644 --- a/haystack/components/query/query_expander.py +++ b/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) diff --git a/haystack/components/rankers/llm_ranker.py b/haystack/components/rankers/llm_ranker.py index 8f1c85a8fb1..6ad029ce173 100644 --- a/haystack/components/rankers/llm_ranker.py +++ b/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 diff --git a/haystack/components/routers/llm_messages_router.py b/haystack/components/routers/llm_messages_router.py index 3ad79673937..0ae31b15210 100644 --- a/haystack/components/routers/llm_messages_router.py +++ b/haystack/components/routers/llm_messages_router.py @@ -7,6 +7,7 @@ from haystack import component, default_from_dict, default_to_dict 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, ChatRole from haystack.utils import deserialize_chatgenerator_inplace @@ -142,7 +143,10 @@ def run(self, messages: list[ChatMessage]) -> dict[str, str | list[ChatMessage]] messages_for_inference.append(ChatMessage.from_system(self._system_prompt)) messages_for_inference.extend(messages) - chat_generator_text = self._chat_generator.run(messages=messages_for_inference)["replies"][0].text + with _trace_chat_generator_run(self._chat_generator, {"messages": messages_for_inference}) as span: + generator_result = self._chat_generator.run(messages=messages_for_inference) + span.set_content_tag("haystack.component.output", generator_result) + chat_generator_text = generator_result["replies"][0].text output = {"chat_generator_text": chat_generator_text} @@ -188,7 +192,9 @@ async def run_async(self, messages: list[ChatMessage]) -> dict[str, str | list[C messages_for_inference.append(ChatMessage.from_system(self._system_prompt)) messages_for_inference.extend(messages) - generator_result = await _execute_component_async(self._chat_generator, messages=messages_for_inference) + with _trace_chat_generator_run(self._chat_generator, {"messages": messages_for_inference}) as span: + generator_result = await _execute_component_async(self._chat_generator, messages=messages_for_inference) + span.set_content_tag("haystack.component.output", generator_result) chat_generator_text = generator_result["replies"][0].text output = {"chat_generator_text": chat_generator_text} diff --git a/releasenotes/notes/trace-internal-chat-generator-token-usage-50bbe296e0bd505c.yaml b/releasenotes/notes/trace-internal-chat-generator-token-usage-50bbe296e0bd505c.yaml new file mode 100644 index 00000000000..313d27f3c38 --- /dev/null +++ b/releasenotes/notes/trace-internal-chat-generator-token-usage-50bbe296e0bd505c.yaml @@ -0,0 +1,10 @@ +--- +enhancements: + - | + ``LLMEvaluator``, ``LLMRanker``, ``QueryExpander``, ``LLMMetadataExtractor``, + ``LLMDocumentContentExtractor`` and ``LLMMessagesRouter`` now wrap their internal ``ChatGenerator`` calls + in a ``haystack.chat_generator.run`` tracing span. These components do not return ``ChatMessage`` objects, + so the LLM token usage carried in ``reply.meta["usage"]`` was previously lost to tracers. The new span + exposes the generator's replies via the ``haystack.component.output`` tag, so token usage is now visible + in traces (requires content tracing to be enabled). When the generator runs across threads, the span is + nested under the component's span. diff --git a/test/components/evaluators/test_llm_evaluator.py b/test/components/evaluators/test_llm_evaluator.py index ee2c512f857..27471ebf260 100644 --- a/test/components/evaluators/test_llm_evaluator.py +++ b/test/components/evaluators/test_llm_evaluator.py @@ -8,6 +8,7 @@ from haystack import Pipeline from haystack.components.evaluators import LLMEvaluator +from haystack.components.generators.chat import MockChatGenerator from haystack.components.generators.chat.openai import OpenAIChatGenerator from haystack.dataclasses.chat_message import ChatMessage @@ -636,3 +637,40 @@ async def test_lifecycle_is_safe_when_chat_generator_lacks_methods(self): await evaluator.warm_up_async() evaluator.close() await evaluator.close_async() + + +class TestLLMEvaluatorTracing: + def test_run_traces_chat_generator_token_usage(self, spying_tracer): + evaluator = LLMEvaluator( + instructions="test-instruction", + inputs=[("predicted_answers", list[str])], + outputs=["score"], + examples=[{"inputs": {"predicted_answers": "Answer"}, "outputs": {"score": 1}}], + chat_generator=MockChatGenerator('{"score": 1}'), + ) + + evaluator.run(predicted_answers=["Football is the most popular sport."]) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 1 + output = gen_spans[0].tags["haystack.component.output"] + assert output["replies"][0].meta["usage"]["total_tokens"] > 0 + + +class TestLLMEvaluatorTracingAsync: + @pytest.mark.asyncio + async def test_run_async_traces_chat_generator_token_usage(self, spying_tracer): + evaluator = LLMEvaluator( + instructions="test-instruction", + inputs=[("predicted_answers", list[str])], + outputs=["score"], + examples=[{"inputs": {"predicted_answers": "Answer"}, "outputs": {"score": 1}}], + chat_generator=MockChatGenerator('{"score": 1}'), + ) + + await evaluator.run_async(predicted_answers=["Football is the most popular sport."]) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 1 + output = gen_spans[0].tags["haystack.component.output"] + assert output["replies"][0].meta["usage"]["total_tokens"] > 0 diff --git a/test/components/extractors/image/test_llm_document_content_extractor.py b/test/components/extractors/image/test_llm_document_content_extractor.py index 8edaaad7858..aae03fa8739 100644 --- a/test/components/extractors/image/test_llm_document_content_extractor.py +++ b/test/components/extractors/image/test_llm_document_content_extractor.py @@ -676,3 +676,66 @@ async def test_lifecycle_is_safe_when_chat_generator_lacks_methods(self): await extractor.warm_up_async() extractor.close() await extractor.close_async() + + +class TestLLMDocumentContentExtractorTracing: + @patch.object(DocumentToImageContent, "run") + def test_run_traces_token_usage_and_nests_per_document_spans(self, mock_doc_to_image_run, spying_tracer): + # Two documents are processed on worker threads. Each generator span must expose the reply's token usage and + # nest under the span active when run() was called, not under another document's generator span. + mock_doc_to_image_run.return_value = { + "image_contents": [ + ImageContent.from_file_path("./test/test_files/images/apple.jpg"), + ImageContent.from_file_path("./test/test_files/images/apple.jpg"), + ] + } + extractor = LLMDocumentContentExtractor( + chat_generator=MockChatGenerator('{"document_content": "Extracted content"}') + ) + + documents = [ + Document(content="", meta={"file_path": "/path/to/a.jpg"}), + Document(content="", meta={"file_path": "/path/to/b.jpg"}), + ] + with spying_tracer.trace("parent") as parent_span: + extractor.run(documents=documents) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 2 + assert all(span.parent_span is parent_span for span in gen_spans) + assert all( + span.tags["haystack.component.output"]["replies"][0].meta["usage"]["total_tokens"] > 0 for span in gen_spans + ) + + +class TestLLMDocumentContentExtractorTracingAsync: + @pytest.mark.asyncio + @patch.object(DocumentToImageContent, "run") + async def test_run_async_traces_token_usage_and_nests_per_document_spans( + self, mock_doc_to_image_run, spying_tracer + ): + # Two documents are processed concurrently. Each generator span must expose the reply's token usage and nest + # under the span active when run_async() was called, not under another document's generator span. + mock_doc_to_image_run.return_value = { + "image_contents": [ + ImageContent.from_file_path("./test/test_files/images/apple.jpg"), + ImageContent.from_file_path("./test/test_files/images/apple.jpg"), + ] + } + extractor = LLMDocumentContentExtractor( + chat_generator=MockChatGenerator('{"document_content": "Extracted content"}') + ) + + documents = [ + Document(content="", meta={"file_path": "/path/to/a.jpg"}), + Document(content="", meta={"file_path": "/path/to/b.jpg"}), + ] + with spying_tracer.trace("parent") as parent_span: + await extractor.run_async(documents=documents) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 2 + assert all(span.parent_span is parent_span for span in gen_spans) + assert all( + span.tags["haystack.component.output"]["replies"][0].meta["usage"]["total_tokens"] > 0 for span in gen_spans + ) diff --git a/test/components/extractors/test_llm_metadata_extractor.py b/test/components/extractors/test_llm_metadata_extractor.py index 457f03b5376..431e81698cb 100644 --- a/test/components/extractors/test_llm_metadata_extractor.py +++ b/test/components/extractors/test_llm_metadata_extractor.py @@ -10,7 +10,7 @@ from haystack import Document, Pipeline from haystack.components.extractors import LLMMetadataExtractor -from haystack.components.generators.chat import OpenAIChatGenerator +from haystack.components.generators.chat import MockChatGenerator, OpenAIChatGenerator from haystack.components.writers import DocumentWriter from haystack.dataclasses import ChatMessage from haystack.document_stores.in_memory import InMemoryDocumentStore @@ -569,3 +569,44 @@ async def test_lifecycle_is_safe_when_inner_lacks_methods(self): await extractor.warm_up_async() extractor.close() await extractor.close_async() + + +class TestLLMMetadataExtractorTracing: + def test_run_traces_token_usage_and_nests_per_document_spans(self, spying_tracer): + # Two documents are processed on worker threads. Each generator span must expose the reply's token usage and + # nest under the span active when run() was called, not under another document's generator span. + extractor = LLMMetadataExtractor( + prompt="Extract entities from: {{ document.content }}", chat_generator=MockChatGenerator('{"entities": []}') + ) + + documents = [Document(content="deepset is in Berlin."), Document(content="Paris is in France.")] + with spying_tracer.trace("parent") as parent_span: + extractor.run(documents=documents) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 2 + assert all(span.parent_span is parent_span for span in gen_spans) + assert all( + span.tags["haystack.component.output"]["replies"][0].meta["usage"]["total_tokens"] > 0 for span in gen_spans + ) + + +class TestLLMMetadataExtractorTracingAsync: + @pytest.mark.asyncio + async def test_run_async_traces_token_usage_and_nests_per_document_spans(self, spying_tracer): + # Two documents are processed concurrently. Each generator span must expose the reply's token usage and nest + # under the span active when run_async() was called, not under another document's generator span. + extractor = LLMMetadataExtractor( + prompt="Extract entities from: {{ document.content }}", chat_generator=MockChatGenerator('{"entities": []}') + ) + + documents = [Document(content="deepset is in Berlin."), Document(content="Paris is in France.")] + with spying_tracer.trace("parent") as parent_span: + await extractor.run_async(documents=documents) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 2 + assert all(span.parent_span is parent_span for span in gen_spans) + assert all( + span.tags["haystack.component.output"]["replies"][0].meta["usage"]["total_tokens"] > 0 for span in gen_spans + ) diff --git a/test/components/query/test_query_expander.py b/test/components/query/test_query_expander.py index dea0f7bfacd..c932f6af257 100644 --- a/test/components/query/test_query_expander.py +++ b/test/components/query/test_query_expander.py @@ -533,3 +533,28 @@ def test_lifecycle_is_safe_when_chat_generator_lacks_methods(self): expander = QueryExpander(chat_generator=chat_generator) expander.warm_up() expander.close() + + +class TestQueryExpanderTracing: + def test_run_traces_chat_generator_token_usage(self, spying_tracer): + expander = QueryExpander(chat_generator=MockChatGenerator('{"queries": ["alt1", "alt2"]}')) + + expander.run(query="green energy sources") + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 1 + output = gen_spans[0].tags["haystack.component.output"] + assert output["replies"][0].meta["usage"]["total_tokens"] > 0 + + +class TestQueryExpanderTracingAsync: + @pytest.mark.asyncio + async def test_run_async_traces_chat_generator_token_usage(self, spying_tracer): + expander = QueryExpander(chat_generator=MockChatGenerator('{"queries": ["alt1", "alt2"]}')) + + await expander.run_async(query="green energy sources") + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 1 + output = gen_spans[0].tags["haystack.component.output"] + assert output["replies"][0].meta["usage"]["total_tokens"] > 0 diff --git a/test/components/rankers/test_llm_ranker.py b/test/components/rankers/test_llm_ranker.py index 683e1211474..f769a16272d 100644 --- a/test/components/rankers/test_llm_ranker.py +++ b/test/components/rankers/test_llm_ranker.py @@ -466,3 +466,28 @@ def test_lifecycle_is_safe_when_chat_generator_lacks_methods(self): ranker = LLMRanker(chat_generator=chat_generator) ranker.warm_up() ranker.close() + + +class TestLLMRankerTracing: + def test_run_traces_chat_generator_token_usage(self, spying_tracer): + ranker = LLMRanker(chat_generator=MockChatGenerator('{"documents": [{"index": 0}]}')) + + ranker.run(query="capital of Germany", documents=[Document(content="Berlin")]) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 1 + output = gen_spans[0].tags["haystack.component.output"] + assert output["replies"][0].meta["usage"]["total_tokens"] > 0 + + +class TestLLMRankerTracingAsync: + @pytest.mark.asyncio + async def test_run_async_traces_chat_generator_token_usage(self, spying_tracer): + ranker = LLMRanker(chat_generator=MockChatGenerator('{"documents": [{"index": 0}]}')) + + await ranker.run_async(query="capital of Germany", documents=[Document(content="Berlin")]) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 1 + output = gen_spans[0].tags["haystack.component.output"] + assert output["replies"][0].meta["usage"]["total_tokens"] > 0 diff --git a/test/components/routers/test_llm_messages_router.py b/test/components/routers/test_llm_messages_router.py index e9669a9714a..93c28c7c1f9 100644 --- a/test/components/routers/test_llm_messages_router.py +++ b/test/components/routers/test_llm_messages_router.py @@ -4,36 +4,20 @@ import os import re -from typing import Any from unittest.mock import AsyncMock, Mock import pytest +from haystack.components.generators.chat import MockChatGenerator from haystack.components.generators.chat.openai import OpenAIChatGenerator from haystack.components.routers.llm_messages_router import LLMMessagesRouter -from haystack.core.serialization import default_from_dict, default_to_dict from haystack.dataclasses import ChatMessage -class MockChatGenerator: - def __init__(self, return_text: str = "safe"): - self.return_text = return_text - - def run(self, messages: list[ChatMessage]) -> dict[str, Any]: - return {"replies": [ChatMessage.from_assistant(self.return_text)]} - - def to_dict(self) -> dict[str, Any]: - return default_to_dict(self, return_text=self.return_text) - - @classmethod - def from_dict(cls, data: dict[str, Any]) -> "MockChatGenerator": - return default_from_dict(cls, data) - - class TestLLMMessagesRouter: def test_init(self): system_prompt = "Classify the messages as safe or unsafe." - chat_generator = MockChatGenerator() + chat_generator = MockChatGenerator("safe") router = LLMMessagesRouter( chat_generator=chat_generator, @@ -49,7 +33,7 @@ def test_init(self): assert router._compiled_patterns == [re.compile(pattern) for pattern in ["safe", "unsafe"]] def test_init_errors(self): - chat_generator = MockChatGenerator() + chat_generator = MockChatGenerator("safe") with pytest.raises(ValueError): LLMMessagesRouter(chat_generator=chat_generator, output_names=[], output_patterns=["pattern1", "pattern2"]) @@ -64,7 +48,9 @@ def test_init_errors(self): def test_run_input_errors(self): router = LLMMessagesRouter( - chat_generator=MockChatGenerator(), output_names=["safe", "unsafe"], output_patterns=["safe", "unsafe"] + chat_generator=MockChatGenerator("safe"), + output_names=["safe", "unsafe"], + output_patterns=["safe", "unsafe"], ) with pytest.raises(ValueError): @@ -74,8 +60,10 @@ def test_run_input_errors(self): router.run([ChatMessage.from_system("You are a helpful assistant.")]) def test_run_no_warm_up_with_unwarmable_chat_generator(self): + chat_generator = Mock(spec=["run"]) + chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("safe")]} router = LLMMessagesRouter( - chat_generator=MockChatGenerator(), output_names=["safe", "unsafe"], output_patterns=["safe", "unsafe"] + chat_generator=chat_generator, output_names=["safe", "unsafe"], output_patterns=["safe", "unsafe"] ) router.run([ChatMessage.from_user("Hello")]) @@ -94,7 +82,7 @@ def mock_run(messages): def test_run(self): router = LLMMessagesRouter( - chat_generator=MockChatGenerator(return_text="safe"), + chat_generator=MockChatGenerator("safe"), output_names=["safe", "unsafe"], output_patterns=["safe", "unsafe"], ) @@ -127,7 +115,7 @@ def test_run_with_system_prompt(self): def test_run_unmatched_output(self): router = LLMMessagesRouter( - chat_generator=MockChatGenerator(return_text="irrelevant"), + chat_generator=MockChatGenerator("irrelevant"), output_names=["safe", "unsafe"], output_patterns=["safe", "unsafe"], ) @@ -141,7 +129,7 @@ def test_run_unmatched_output(self): assert "unsafe" not in result def test_to_dict(self): - chat_generator = MockChatGenerator(return_text="safe") + chat_generator = MockChatGenerator("safe") router = LLMMessagesRouter( chat_generator=chat_generator, output_names=["safe", "unsafe"], output_patterns=["safe", "unsafe"] @@ -156,7 +144,7 @@ def test_to_dict(self): assert result["init_parameters"]["system_prompt"] is None def test_from_dict(self): - chat_generator = MockChatGenerator(return_text="safe") + chat_generator = MockChatGenerator("safe") data = { "type": "haystack.components.routers.llm_messages_router.LLMMessagesRouter", @@ -235,8 +223,9 @@ async def test_run_async_unmatched_output(self): @pytest.mark.asyncio async def test_run_async_fallback_to_sync_run(self): - # MockChatGenerator defines only a synchronous `run`, so the utility falls back to it. - chat_generator = MockChatGenerator(return_text="safe") + # A chat generator that defines only a synchronous `run`, so the utility falls back to it. + chat_generator = Mock(spec=["run"]) + chat_generator.run.return_value = {"replies": [ChatMessage.from_assistant("safe")]} assert not hasattr(chat_generator, "run_async") router = LLMMessagesRouter( chat_generator=chat_generator, output_names=["safe", "unsafe"], output_patterns=["safe", "unsafe"] @@ -345,3 +334,32 @@ def test_lifecycle_is_safe_when_chat_generator_lacks_methods(self): router = self._make_router(chat_generator) router.warm_up() router.close() + + +class TestLLMMessagesRouterTracing: + def test_run_traces_chat_generator_token_usage(self, spying_tracer): + router = LLMMessagesRouter( + chat_generator=MockChatGenerator("safe"), output_names=["safe"], output_patterns=["safe"] + ) + + router.run(messages=[ChatMessage.from_user("How to bake bread?")]) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 1 + output = gen_spans[0].tags["haystack.component.output"] + assert output["replies"][0].meta["usage"]["total_tokens"] > 0 + + +class TestLLMMessagesRouterTracingAsync: + @pytest.mark.asyncio + async def test_run_async_traces_chat_generator_token_usage(self, spying_tracer): + router = LLMMessagesRouter( + chat_generator=MockChatGenerator("safe"), output_names=["safe"], output_patterns=["safe"] + ) + + await router.run_async(messages=[ChatMessage.from_user("How to bake bread?")]) + + gen_spans = [s for s in spying_tracer.spans if s.operation_name == "haystack.chat_generator.run"] + assert len(gen_spans) == 1 + output = gen_spans[0].tags["haystack.component.output"] + assert output["replies"][0].meta["usage"]["total_tokens"] > 0