feat: add optional BM25 retrieval confidence metadata

Author: SeCuReDmE-main-dev

haystack/components/retrievers/in_memory/bm25_retriever.py

@@ -44,6 +44,7 @@ def __init__(
         filters: dict[str, Any] | None = None,
         top_k: int = 10,
         scale_score: bool = False,
+        include_confidence: bool = False,
         filter_policy: FilterPolicy = FilterPolicy.REPLACE,
     ) -> None:
         """
@@ -58,6 +59,10 @@ def __init__(
         :param scale_score:
             When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
             When `False`, uses raw similarity scores.
+        :param include_confidence:
+            When `True`, adds optional retrieval confidence metadata to returned documents when `scale_score` is also
+            `True`. The metadata is exposed via `Document.meta["retrieval_confidence"]` and
+            `Document.meta["retrieval_confidence_source"]`.
         :param filter_policy: The filter policy to apply during retrieval.
         Filter policy determines how filters are applied when retrieving documents. You can choose:
         - `REPLACE` (default): Overrides the initialization filters with the filters specified at runtime.
@@ -78,6 +83,7 @@ def __init__(
         self.filters = filters
         self.top_k = top_k
         self.scale_score = scale_score
+        self.include_confidence = include_confidence
         self.filter_policy = filter_policy
 
     def _get_telemetry_data(self) -> dict[str, Any]:
@@ -99,6 +105,7 @@ def to_dict(self) -> dict[str, Any]:
             filters=self.filters,
             top_k=self.top_k,
             scale_score=self.scale_score,
+            include_confidence=self.include_confidence,
             filter_policy=self.filter_policy.value,
         )
 
@@ -124,6 +131,7 @@ def run(
         filters: dict[str, Any] | None = None,
         top_k: int | None = None,
         scale_score: bool | None = None,
+        include_confidence: bool | None = None,
     ) -> dict[str, list[Document]]:
         """
         Run the InMemoryBM25Retriever on the given input data.
@@ -137,6 +145,9 @@ def run(
         :param scale_score:
             When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
             When `False`, uses raw similarity scores.
+        :param include_confidence:
+            When `True`, adds optional retrieval confidence metadata to returned documents when `scale_score` is also
+            `True`. When `False`, no retrieval confidence metadata is added.
         :returns:
             The retrieved documents.
 
@@ -151,8 +162,12 @@ def run(
             top_k = self.top_k
         if scale_score is None:
             scale_score = self.scale_score
+        if include_confidence is None:
+            include_confidence = self.include_confidence
 
         docs = self.document_store.bm25_retrieval(query=query, filters=filters, top_k=top_k, scale_score=scale_score)
+        if include_confidence and scale_score:
+            self._add_confidence_metadata(docs)
         return {"documents": docs}
 
     @component.output_types(documents=list[Document])
@@ -162,6 +177,7 @@ async def run_async(
         filters: dict[str, Any] | None = None,
         top_k: int | None = None,
         scale_score: bool | None = None,
+        include_confidence: bool | None = None,
     ) -> dict[str, list[Document]]:
         """
         Run the InMemoryBM25Retriever on the given input data.
@@ -175,6 +191,9 @@ async def run_async(
         :param scale_score:
             When `True`, scales the score of retrieved documents to a range of 0 to 1, where 1 means extremely relevant.
             When `False`, uses raw similarity scores.
+        :param include_confidence:
+            When `True`, adds optional retrieval confidence metadata to returned documents when `scale_score` is also
+            `True`. When `False`, no retrieval confidence metadata is added.
         :returns:
             The retrieved documents.
 
@@ -189,8 +208,20 @@ async def run_async(
             top_k = self.top_k
         if scale_score is None:
             scale_score = self.scale_score
+        if include_confidence is None:
+            include_confidence = self.include_confidence
 
         docs = await self.document_store.bm25_retrieval_async(
             query=query, filters=filters, top_k=top_k, scale_score=scale_score
         )
+        if include_confidence and scale_score:
+            self._add_confidence_metadata(docs)
         return {"documents": docs}
+
+    @staticmethod
+    def _add_confidence_metadata(documents: list[Document]) -> None:
+        for document in documents:
+            if document.score is None:
+                continue
+            document.meta["retrieval_confidence"] = document.score
+            document.meta["retrieval_confidence_source"] = "bm25_scaled_score"

releasenotes/notes/add-bm25-retrieval-confidence-metadata-phase2.yaml

@@ -0,0 +1,8 @@
+---
+enhancements:
+  - |
+    Adds an optional `include_confidence` parameter to `InMemoryBM25Retriever`.
+    When enabled together with `scale_score=True`, returned documents now expose
+    BM25 retrieval confidence metadata via `Document.meta["retrieval_confidence"]`
+    and `Document.meta["retrieval_confidence_source"]` without changing
+    `Document.score` semantics.

test/components/retrievers/test_in_memory_bm25_retriever.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import asyncio
 from typing import Any
 
 import pytest
@@ -31,12 +32,16 @@ def test_init_default(self, in_memory_doc_store):
         assert retriever.filters is None
         assert retriever.top_k == 10
         assert retriever.scale_score is False
+        assert retriever.include_confidence is False
 
     def test_init_with_parameters(self, in_memory_doc_store):
-        retriever = InMemoryBM25Retriever(in_memory_doc_store, filters={"name": "test.txt"}, top_k=5, scale_score=True)
+        retriever = InMemoryBM25Retriever(
+            in_memory_doc_store, filters={"name": "test.txt"}, top_k=5, scale_score=True, include_confidence=True
+        )
         assert retriever.filters == {"name": "test.txt"}
         assert retriever.top_k == 5
         assert retriever.scale_score
+        assert retriever.include_confidence is True
 
     def test_init_with_invalid_top_k_parameter(self, in_memory_doc_store):
         with pytest.raises(ValueError):
@@ -56,6 +61,7 @@ def test_to_dict(self):
                 "filters": None,
                 "top_k": 10,
                 "scale_score": False,
+                "include_confidence": False,
                 "filter_policy": "replace",
             },
         }
@@ -78,6 +84,7 @@ def test_to_dict_with_custom_init_parameters(self):
                 "filters": {"name": "test.txt"},
                 "top_k": 5,
                 "scale_score": True,
+                "include_confidence": False,
                 "filter_policy": "replace",
             },
         }
@@ -99,6 +106,7 @@ def test_from_dict(self):
         assert component.filters == {"name": "test.txt"}
         assert component.top_k == 5
         assert component.scale_score is False
+        assert component.include_confidence is False
         assert component.filter_policy == FilterPolicy.REPLACE
 
     def test_from_dict_without_docstore(self):
@@ -140,6 +148,38 @@ def test_invalid_run_wrong_store_type(self):
         with pytest.raises(TypeError, match="document_store must be an instance of InMemoryDocumentStore"):
             InMemoryBM25Retriever(SomeOtherDocumentStore())
 
+    def test_run_with_include_confidence_adds_metadata_when_scaled(self, in_memory_doc_store, mock_docs):
+        in_memory_doc_store.write_documents(mock_docs)
+
+        retriever = InMemoryBM25Retriever(in_memory_doc_store, top_k=3, scale_score=True, include_confidence=True)
+        result = retriever.run(query="PHP")
+
+        first_document = result["documents"][0]
+        assert first_document.score is not None
+        assert first_document.meta["retrieval_confidence"] == first_document.score
+        assert first_document.meta["retrieval_confidence_source"] == "bm25_scaled_score"
+
+    def test_run_with_include_confidence_does_not_add_metadata_when_not_scaled(self, in_memory_doc_store, mock_docs):
+        in_memory_doc_store.write_documents(mock_docs)
+
+        retriever = InMemoryBM25Retriever(in_memory_doc_store, top_k=3, include_confidence=True)
+        result = retriever.run(query="PHP")
+
+        first_document = result["documents"][0]
+        assert "retrieval_confidence" not in first_document.meta
+        assert "retrieval_confidence_source" not in first_document.meta
+
+    def test_run_async_with_include_confidence_matches_sync_behavior(self, in_memory_doc_store, mock_docs):
+        in_memory_doc_store.write_documents(mock_docs)
+
+        retriever = InMemoryBM25Retriever(in_memory_doc_store, top_k=3, scale_score=True, include_confidence=True)
+        result = asyncio.run(retriever.run_async(query="PHP"))
+
+        first_document = result["documents"][0]
+        assert first_document.score is not None
+        assert first_document.meta["retrieval_confidence"] == first_document.score
+        assert first_document.meta["retrieval_confidence_source"] == "bm25_scaled_score"
+
     @pytest.mark.integration
     @pytest.mark.parametrize(
         "query, query_result",