| | 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] |
| |
|
