[draft]: propagate score through retriever

libs/community/langchain_community/vectorstores/redis/base.py

@@ -1444,7 +1444,11 @@ class Config:
         arbitrary_types_allowed = True
 
     def _get_relevant_documents(
-        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,
+        include_score: bool = False,
     ) -> List[Document]:
         if self.search_type == "similarity":
             docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
@@ -1472,7 +1476,11 @@ def _get_relevant_documents(
         return docs
 
     async def _aget_relevant_documents(
-        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
+        self,
+        query: str,
+        *,
+        run_manager: AsyncCallbackManagerForRetrieverRun,
+        include_score: bool = False,
     ) -> List[Document]:
         if self.search_type == "similarity":
             docs = await self.vectorstore.asimilarity_search(

libs/community/tests/unit_tests/vectorstores/test_databricks_vector_search.py

@@ -1,10 +1,11 @@
 import itertools
 import random
 import uuid
-from typing import Dict, List, Optional, Set
+from typing import Dict, List, Optional, Set, cast
 from unittest.mock import MagicMock, patch
 
 import pytest
+from langchain_core.documents import DocumentSearchHit
 
 from langchain_community.vectorstores import DatabricksVectorSearch
 from tests.integration_tests.vectorstores.fake_embeddings import (
@@ -598,6 +599,13 @@ def test_similarity_score_threshold(index_details: dict, threshold: float) -> No
         assert len(search_result) == len(fake_texts)
     else:
         assert len(search_result) == 0
+    result_with_scores = cast(
+        List[DocumentSearchHit], retriever.invoke(query, include_score=True)
+    )
+    for idx, result in enumerate(result_with_scores):
+        assert result.score >= threshold
+        assert result.page_content == search_result[idx].page_content
+        assert result.metadata == search_result[idx].metadata
 
 
 @pytest.mark.requires("databricks", "databricks.vector_search")

libs/core/langchain_core/documents/__init__.py

@@ -2,8 +2,13 @@
 and their transformations.
 
 """
-from langchain_core.documents.base import Document
+from langchain_core.documents.base import Document, DocumentSearchHit
 from langchain_core.documents.compressor import BaseDocumentCompressor
 from langchain_core.documents.transformers import BaseDocumentTransformer
 
-__all__ = ["Document", "BaseDocumentTransformer", "BaseDocumentCompressor"]
+__all__ = [
+    "Document",
+    "DocumentSearchHit",
+    "BaseDocumentTransformer",
+    "BaseDocumentCompressor",
+]

libs/core/langchain_core/documents/base.py

@@ -30,3 +30,21 @@ def is_lc_serializable(cls) -> bool:
     def get_lc_namespace(cls) -> List[str]:
         """Get the namespace of the langchain object."""
         return ["langchain", "schema", "document"]
+
+
+class DocumentSearchHit(Document):
+    """Class for storing a document and fields associated with retrieval."""
+
+    score: float
+    """Score associated with the document's relevance to a query."""
+    type: Literal["DocumentSearchHit"] = "DocumentSearchHit"  # type: ignore[assignment] # noqa: E501
+
+    @classmethod
+    def is_lc_serializable(cls) -> bool:
+        """Return whether this class is serializable."""
+        return True
+
+    @classmethod
+    def get_lc_namespace(cls) -> List[str]:
+        """Get the namespace of the langchain object."""
+        return ["langchain", "schema", "document_search_hit"]

libs/core/langchain_core/load/mapping.py

@@ -157,6 +157,12 @@
         "base",
         "Document",
     ),
+    ("langchain", "schema", "document_search_hit", "DocumentSearchHit"): (
+        "langchain_core",
+        "documents",
+        "base",
+        "DocumentSearchHit",
+    ),
     ("langchain", "output_parsers", "fix", "OutputFixingParser"): (
         "langchain",
         "output_parsers",
@@ -666,6 +672,12 @@
         "base",
         "Document",
     ),
+    ("langchain_core", "documents", "base", "DocumentSearchHit"): (
+        "langchain_core",
+        "documents",
+        "base",
+        "DocumentSearchHit",
+    ),
     ("langchain_core", "prompts", "chat", "AIMessagePromptTemplate"): (
         "langchain_core",
         "prompts",

libs/core/langchain_core/vectorstores.py

@@ -39,6 +39,7 @@
     TypeVar,
 )
 
+from langchain_core.documents import DocumentSearchHit
 from langchain_core.embeddings import Embeddings
 from langchain_core.pydantic_v1 import Field, root_validator
 from langchain_core.retrievers import BaseRetriever
@@ -690,8 +691,17 @@ def validate_search_type(cls, values: Dict) -> Dict:
         return values
 
     def _get_relevant_documents(
-        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,
+        include_score: bool = False,
     ) -> List[Document]:
+        if include_score and self.search_type != "similarity_score_threshold":
+            raise ValueError(
+                "include_score is only supported "
+                "for search_type=similarity_score_threshold"
+            )
         if self.search_type == "similarity":
             docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
         elif self.search_type == "similarity_score_threshold":
@@ -700,6 +710,15 @@ def _get_relevant_documents(
                     query, **self.search_kwargs
                 )
             )
+            if include_score:
+                return [
+                    DocumentSearchHit(
+                        page_content=doc.page_content,
+                        metadata=doc.metadata,
+                        score=score,
+                    )
+                    for doc, score in docs_and_similarities
+                ]
             docs = [doc for doc, _ in docs_and_similarities]
         elif self.search_type == "mmr":
             docs = self.vectorstore.max_marginal_relevance_search(
@@ -710,8 +729,17 @@ def _get_relevant_documents(
         return docs
 
     async def _aget_relevant_documents(
-        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
+        self,
+        query: str,
+        *,
+        run_manager: AsyncCallbackManagerForRetrieverRun,
+        include_score: bool = False,
     ) -> List[Document]:
+        if include_score and self.search_type != "similarity_score_threshold":
+            raise ValueError(
+                "include_score is only supported "
+                "for search_type=similarity_score_threshold"
+            )
         if self.search_type == "similarity":
             docs = await self.vectorstore.asimilarity_search(
                 query, **self.search_kwargs
@@ -722,6 +750,15 @@ async def _aget_relevant_documents(
                     query, **self.search_kwargs
                 )
             )
+            if include_score:
+                return [
+                    DocumentSearchHit(
+                        page_content=doc.page_content,
+                        metadata=doc.metadata,
+                        score=score,
+                    )
+                    for doc, score in docs_and_similarities
+                ]
             docs = [doc for doc, _ in docs_and_similarities]
         elif self.search_type == "mmr":
             docs = await self.vectorstore.amax_marginal_relevance_search(

libs/core/tests/unit_tests/documents/test_imports.py

@@ -1,6 +1,11 @@
 from langchain_core.documents import __all__
 
-EXPECTED_ALL = ["Document", "BaseDocumentTransformer", "BaseDocumentCompressor"]
+EXPECTED_ALL = [
+    "Document",
+    "DocumentSearchHit",
+    "BaseDocumentTransformer",
+    "BaseDocumentCompressor",
+]
 
 
 def test_all_imports() -> None:

libs/langchain/langchain/retrievers/self_query/base.py

@@ -33,7 +33,7 @@
     AsyncCallbackManagerForRetrieverRun,
     CallbackManagerForRetrieverRun,
 )
-from langchain_core.documents import Document
+from langchain_core.documents import Document, DocumentSearchHit
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.pydantic_v1 import Field, root_validator
 from langchain_core.retrievers import BaseRetriever
@@ -192,19 +192,47 @@ def _prepare_query(
         return new_query, search_kwargs
 
     def _get_docs_with_query(
-        self, query: str, search_kwargs: Dict[str, Any]
+        self, query: str, search_kwargs: Dict[str, Any], include_score: bool = False
     ) -> List[Document]:
-        docs = self.vectorstore.search(query, self.search_type, **search_kwargs)
+        if include_score:
+            docs_and_scores = self.vectorstore.similarity_search_with_score(
+                query, **search_kwargs
+            )
+            return [
+                DocumentSearchHit(
+                    page_content=doc.page_content, metadata=doc.metadata, score=score
+                )
+                for doc, score in docs_and_scores
+            ]
+        else:
+            docs = self.vectorstore.search(query, self.search_type, **search_kwargs)
         return docs
 
     async def _aget_docs_with_query(
-        self, query: str, search_kwargs: Dict[str, Any]
+        self, query: str, search_kwargs: Dict[str, Any], include_score: bool = False
     ) -> List[Document]:
-        docs = await self.vectorstore.asearch(query, self.search_type, **search_kwargs)
+        if include_score:
+            docs_and_scores = await self.vectorstore.asimilarity_search_with_score(
+                query, **search_kwargs
+            )
+            return [
+                DocumentSearchHit(
+                    page_content=doc.page_content, metadata=doc.metadata, score=score
+                )
+                for doc, score in docs_and_scores
+            ]
+        else:
+            docs = await self.vectorstore.asearch(
+                query, self.search_type, **search_kwargs
+            )
         return docs
 
     def _get_relevant_documents(
-        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,
+        include_score: bool = False,
     ) -> List[Document]:
         """Get documents relevant for a query.
 
@@ -220,11 +248,17 @@ def _get_relevant_documents(
         if self.verbose:
             logger.info(f"Generated Query: {structured_query}")
         new_query, search_kwargs = self._prepare_query(query, structured_query)
-        docs = self._get_docs_with_query(new_query, search_kwargs)
+        docs = self._get_docs_with_query(
+            new_query, search_kwargs, include_score=include_score
+        )
         return docs
 
     async def _aget_relevant_documents(
-        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
+        self,
+        query: str,
+        *,
+        run_manager: AsyncCallbackManagerForRetrieverRun,
+        include_score: bool = False,
     ) -> List[Document]:
         """Get documents relevant for a query.
 
@@ -240,7 +274,9 @@ async def _aget_relevant_documents(
         if self.verbose:
             logger.info(f"Generated Query: {structured_query}")
         new_query, search_kwargs = self._prepare_query(query, structured_query)
-        docs = await self._aget_docs_with_query(new_query, search_kwargs)
+        docs = await self._aget_docs_with_query(
+            new_query, search_kwargs, include_score=include_score
+        )
         return docs
 
     @classmethod

libs/langchain/pyproject.toml

@@ -12,7 +12,7 @@ langchain-server = "langchain.server:main"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
-langchain-core = "^0.1.48"
+langchain-core = "^0.1.52"
 langchain-text-splitters = ">=0.0.1,<0.1"
 langchain-community = ">=0.0.37,<0.1"
 langsmith = "^0.1.17"