|
|
from enum import Enum |
|
|
from typing import Any, Dict, List, Optional |
|
|
|
|
|
from langchain_core.callbacks import ( |
|
|
AsyncCallbackManagerForRetrieverRun, |
|
|
CallbackManagerForRetrieverRun, |
|
|
) |
|
|
from langchain_core.documents import Document |
|
|
from langchain_core.retrievers import BaseRetriever |
|
|
from langchain_core.stores import BaseStore, ByteStore |
|
|
from langchain_core.vectorstores import VectorStore |
|
|
from pydantic import Field, model_validator |
|
|
|
|
|
from langchain.storage._lc_store import create_kv_docstore |
|
|
|
|
|
|
|
|
class SearchType(str, Enum): |
|
|
"""Enumerator of the types of search to perform.""" |
|
|
|
|
|
similarity = "similarity" |
|
|
"""Similarity search.""" |
|
|
similarity_score_threshold = "similarity_score_threshold" |
|
|
"""Similarity search with a score threshold.""" |
|
|
mmr = "mmr" |
|
|
"""Maximal Marginal Relevance reranking of similarity search.""" |
|
|
|
|
|
|
|
|
class MultiVectorRetriever(BaseRetriever): |
|
|
"""Retrieve from a set of multiple embeddings for the same document.""" |
|
|
|
|
|
vectorstore: VectorStore |
|
|
"""The underlying vectorstore to use to store small chunks |
|
|
and their embedding vectors""" |
|
|
byte_store: Optional[ByteStore] = None |
|
|
"""The lower-level backing storage layer for the parent documents""" |
|
|
docstore: BaseStore[str, Document] |
|
|
"""The storage interface for the parent documents""" |
|
|
id_key: str = "doc_id" |
|
|
search_kwargs: dict = Field(default_factory=dict) |
|
|
"""Keyword arguments to pass to the search function.""" |
|
|
search_type: SearchType = SearchType.similarity |
|
|
"""Type of search to perform (similarity / mmr)""" |
|
|
|
|
|
@model_validator(mode="before") |
|
|
@classmethod |
|
|
def shim_docstore(cls, values: Dict) -> Any: |
|
|
byte_store = values.get("byte_store") |
|
|
docstore = values.get("docstore") |
|
|
if byte_store is not None: |
|
|
docstore = create_kv_docstore(byte_store) |
|
|
elif docstore is None: |
|
|
raise Exception("You must pass a `byte_store` parameter.") |
|
|
values["docstore"] = docstore |
|
|
return values |
|
|
|
|
|
def _get_relevant_documents( |
|
|
self, query: str, *, run_manager: CallbackManagerForRetrieverRun |
|
|
) -> List[Document]: |
|
|
"""Get documents relevant to a query. |
|
|
Args: |
|
|
query: String to find relevant documents for |
|
|
run_manager: The callbacks handler to use |
|
|
Returns: |
|
|
List of relevant documents |
|
|
""" |
|
|
if self.search_type == SearchType.mmr: |
|
|
sub_docs = self.vectorstore.max_marginal_relevance_search( |
|
|
query, **self.search_kwargs |
|
|
) |
|
|
elif self.search_type == SearchType.similarity_score_threshold: |
|
|
sub_docs_and_similarities = ( |
|
|
self.vectorstore.similarity_search_with_relevance_scores( |
|
|
query, **self.search_kwargs |
|
|
) |
|
|
) |
|
|
sub_docs = [sub_doc for sub_doc, _ in sub_docs_and_similarities] |
|
|
else: |
|
|
sub_docs = self.vectorstore.similarity_search(query, **self.search_kwargs) |
|
|
|
|
|
|
|
|
ids = [] |
|
|
for d in sub_docs: |
|
|
if self.id_key in d.metadata and d.metadata[self.id_key] not in ids: |
|
|
ids.append(d.metadata[self.id_key]) |
|
|
docs = self.docstore.mget(ids) |
|
|
return [d for d in docs if d is not None] |
|
|
|
|
|
async def _aget_relevant_documents( |
|
|
self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun |
|
|
) -> List[Document]: |
|
|
"""Asynchronously get documents relevant to a query. |
|
|
Args: |
|
|
query: String to find relevant documents for |
|
|
run_manager: The callbacks handler to use |
|
|
Returns: |
|
|
List of relevant documents |
|
|
""" |
|
|
if self.search_type == SearchType.mmr: |
|
|
sub_docs = await self.vectorstore.amax_marginal_relevance_search( |
|
|
query, **self.search_kwargs |
|
|
) |
|
|
elif self.search_type == SearchType.similarity_score_threshold: |
|
|
sub_docs_and_similarities = ( |
|
|
await self.vectorstore.asimilarity_search_with_relevance_scores( |
|
|
query, **self.search_kwargs |
|
|
) |
|
|
) |
|
|
sub_docs = [sub_doc for sub_doc, _ in sub_docs_and_similarities] |
|
|
else: |
|
|
sub_docs = await self.vectorstore.asimilarity_search( |
|
|
query, **self.search_kwargs |
|
|
) |
|
|
|
|
|
|
|
|
ids = [] |
|
|
for d in sub_docs: |
|
|
if self.id_key in d.metadata and d.metadata[self.id_key] not in ids: |
|
|
ids.append(d.metadata[self.id_key]) |
|
|
docs = await self.docstore.amget(ids) |
|
|
return [d for d in docs if d is not None] |
|
|
|
