fix: don't treat a generator's own "error" field as failure in LLMMetadataExtractor (#12852)

haystack/components/extractors/llm_metadata_extractor.py

@@ -5,10 +5,8 @@
 import copy
 import json
 from asyncio import Semaphore, gather
-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
@@ -297,10 +295,17 @@ class LLMMetadataExtractor:
 
         return all_prompts
 
-    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
+    @staticmethod
+    def _fail(document: Document, error: str, response: ChatMessage | None = None) -> tuple[Document, bool]:
+        """Return a copy of ``document`` with the failure metadata set, flagged as failed."""
+        new_meta = {**document.meta, "metadata_extraction_error": error, "metadata_extraction_response": response}
+        return replace(document, meta=new_meta), False
+
+    def _run_on_thread(
+        self, document: Document, prompt: ChatMessage | None, parent_span: tracing.Span | None = None
+    ) -> tuple[Document, bool]:
         if prompt is None:
-            return {"error": "Document has no content, skipping LLM call."}
+            return self._fail(document, "Document has no content, skipping LLM call.")
 
         try:
             with _trace_chat_generator_run(
@@ -316,13 +321,14 @@ class LLMMetadataExtractor:
                 class_name=self._chat_generator.__class__.__name__,
                 error=e,
             )
-            result = {"error": "LLM failed with exception: " + str(e)}
-        return result
+            return self._fail(document, "LLM failed with exception: " + str(e))
+        return self._process_reply(document, result["replies"][0])
 
-    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
+    async def _run_async(
+        self, document: Document, prompt: ChatMessage | None, parent_span: tracing.Span | None = None
+    ) -> tuple[Document, bool]:
         if prompt is None:
-            return {"error": "Document has no content, skipping LLM call."}
+            return self._fail(document, "Document has no content, skipping LLM call.")
 
         try:
             with _trace_chat_generator_run(
@@ -338,45 +344,33 @@ class LLMMetadataExtractor:
                 class_name=self._chat_generator.__class__.__name__,
                 error=e,
             )
-            result = {"error": "LLM failed with exception: " + str(e)}
-        return result
+            return self._fail(document, "LLM failed with exception: " + str(e))
+        return self._process_reply(document, result["replies"][0])
 
-    def _process_results(
-        self, documents: list[Document], results: Iterable[dict[str, Any]]
-    ) -> tuple[list[Document], list[Document]]:
-        successful_documents = []
-        failed_documents = []
-        for document, result in zip(documents, results, strict=True):
-            new_meta = {**document.meta}
-            if "error" in result:
-                new_meta["metadata_extraction_error"] = result["error"]
-                new_meta["metadata_extraction_response"] = None
-                failed_documents.append(replace(document, meta=new_meta))
-                continue
+    def _process_reply(self, document: Document, reply: ChatMessage) -> tuple[Document, bool]:
+        """
+        Parse one document's LLM reply into metadata.
 
-            reply = result["replies"][0]
-            try:
-                parsed_metadata = _parse_dict_from_json(
-                    reply.text, expected_keys=self.expected_keys, raise_on_failure=True
-                )
-            except (ValueError, json.JSONDecodeError) as e:
-                logger.warning(
-                    "Response from the LLM is not valid JSON or missing expected keys. Received output: {response}",
-                    response=reply.text,
-                )
-                if self.raise_on_failure:
-                    raise
-                new_meta["metadata_extraction_error"] = "Response is not valid JSON or missing keys. Error: " + str(e)
-                new_meta["metadata_extraction_response"] = reply
-                failed_documents.append(replace(document, meta=new_meta))
-                continue
+        Returns (updated_document, True if success else False).
+        """
+        try:
+            parsed_metadata = _parse_dict_from_json(
+                reply.text or "", expected_keys=self.expected_keys, raise_on_failure=True
+            )
+        except (ValueError, json.JSONDecodeError) as e:
+            logger.warning(
+                "Response from the LLM is not valid JSON or missing expected keys. Received output: {response}",
+                response=reply.text,
+            )
+            if self.raise_on_failure:
+                raise
+            return self._fail(document, "Response is not valid JSON or missing keys. Error: " + str(e), response=reply)
 
-            new_meta.update(parsed_metadata)
-            # Remove metadata_extraction_error and metadata_extraction_response if present from previous runs.
-            new_meta.pop("metadata_extraction_error", None)
-            new_meta.pop("metadata_extraction_response", None)
-            successful_documents.append(replace(document, meta=new_meta))
-        return successful_documents, failed_documents
+        new_meta = {**document.meta, **parsed_metadata}
+        # Remove metadata_extraction_error and metadata_extraction_response if present from previous runs.
+        new_meta.pop("metadata_extraction_error", None)
+        new_meta.pop("metadata_extraction_response", None)
+        return replace(document, meta=new_meta), True
 
     @component.output_types(documents=list[Document], failed_documents=list[Document])
     def run(self, documents: list[Document], page_range: list[str | int] | None = None) -> dict[str, Any]:
@@ -421,9 +415,19 @@ class LLMMetadataExtractor:
 
         # Run the LLM on each prompt
         with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            results = executor.map(partial(self._run_on_thread, parent_span=parent_span), all_prompts)
+            results = executor.map(
+                lambda document, prompt: self._run_on_thread(document, prompt, parent_span=parent_span),
+                documents,
+                all_prompts,
+            )
 
-        successful_documents, failed_documents = self._process_results(documents, results)
+        successful_documents = []
+        failed_documents = []
+        for doc, success in results:
+            if success:
+                successful_documents.append(doc)
+            else:
+                failed_documents.append(doc)
 
         return {"documents": successful_documents, "failed_documents": failed_documents}
 
@@ -474,12 +478,20 @@ class LLMMetadataExtractor:
         # 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 def _bounded_run(document: Document, prompt: ChatMessage | None) -> tuple[Document, bool]:
             async with sem:
-                return await self._run_async(prompt, parent_span=parent_span)
+                return await self._run_async(document, prompt, parent_span=parent_span)
 
-        results = await gather(*[_bounded_run(prompt) for prompt in all_prompts])
+        results = await gather(
+            *[_bounded_run(document, prompt) for document, prompt in zip(documents, all_prompts, strict=True)]
+        )
 
-        successful_documents, failed_documents = self._process_results(documents, results)
+        successful_documents = []
+        failed_documents = []
+        for doc, success in results:
+            if success:
+                successful_documents.append(doc)
+            else:
+                failed_documents.append(doc)
 
         return {"documents": successful_documents, "failed_documents": failed_documents}

releasenotes/notes/llm-metadata-extractor-generator-error-key-collision-73b215ffbfef8f5a.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    Fixed ``LLMMetadataExtractor`` incorrectly treating a chat generator output that carries its own
+    ``error`` field as a failed LLM call, which sent every document to ``failed_documents``
+    with ``metadata_extraction_error`` set to ``None``. Such documents are now processed normally.

test/components/extractors/test_llm_metadata_extractor.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import os
+from typing import Any
 from unittest.mock import AsyncMock, Mock
 
 import pytest
@@ -16,6 +17,19 @@ from haystack.dataclasses import ChatMessage
 from haystack.document_stores.in_memory import InMemoryDocumentStore
 
 
+class ErrorKeyChatGenerator:
+    """Wrapper-style generator whose output dict always includes an "error" field."""
+
+    def __init__(self, response: str) -> None:
+        self._response = response
+
+    def run(self, messages: list[ChatMessage], **kwargs: Any) -> dict[str, Any]:
+        return {"replies": [ChatMessage.from_assistant(self._response)], "error": None}
+
+    async def run_async(self, messages: list[ChatMessage], **kwargs: Any) -> dict[str, Any]:
+        return self.run(messages, **kwargs)
+
+
 @pytest.fixture
 def ner_prompt() -> str:
     return """-Goal-
@@ -288,6 +302,23 @@ class TestLLMMetadataExtractor:
         assert result["failed_documents"] == []
         assert result["documents"][0].meta == {"error": "timeout", "severity": "high"}
 
+    def test_run_generator_output_with_error_key_is_not_treated_as_failure(self) -> None:
+        extractor = LLMMetadataExtractor(
+            prompt="prompt {{document.content}}",
+            expected_keys=["topic"],
+            chat_generator=ErrorKeyChatGenerator(response='{"topic": "physics"}'),
+        )
+
+        result = extractor.run(documents=[Document(content="content"), Document(content="")])
+
+        assert len(result["documents"]) == 1
+        assert result["documents"][0].meta == {"topic": "physics"}
+        assert len(result["failed_documents"]) == 1
+        assert result["failed_documents"][0].meta == {
+            "metadata_extraction_error": "Document has no content, skipping LLM call.",
+            "metadata_extraction_response": None,
+        }
+
     def test_run_raises_parse_error_when_raise_on_failure_is_true(self, caplog: pytest.LogCaptureFixture) -> None:
         extractor = LLMMetadataExtractor(
             prompt="prompt {{document.content}}",
@@ -313,6 +344,24 @@ class TestLLMMetadataExtractor:
         assert result["failed_documents"] == []
         assert result["documents"][0].meta == {"error": "timeout", "severity": "high"}
 
+    @pytest.mark.asyncio
+    async def test_run_async_generator_output_with_error_key_is_not_treated_as_failure(self) -> None:
+        extractor = LLMMetadataExtractor(
+            prompt="prompt {{document.content}}",
+            expected_keys=["topic"],
+            chat_generator=ErrorKeyChatGenerator(response='{"topic": "physics"}'),
+        )
+
+        result = await extractor.run_async(documents=[Document(content="content"), Document(content="")])
+
+        assert len(result["documents"]) == 1
+        assert result["documents"][0].meta == {"topic": "physics"}
+        assert len(result["failed_documents"]) == 1
+        assert result["failed_documents"][0].meta == {
+            "metadata_extraction_error": "Document has no content, skipping LLM call.",
+            "metadata_extraction_response": None,
+        }
+
     @pytest.mark.asyncio
     async def test_run_async_raises_parse_error_when_raise_on_failure_is_true(self) -> None:
         extractor = LLMMetadataExtractor(