codekingpro/portable-devtools
114k
1"""A vector store stores embedded data and performs vector search.2 3One of the most common ways to store and search over unstructured data is to4embed it and store the resulting embedding vectors, and then query the store5and retrieve the data that are 'most similar' to the embedded query.6"""7 8from __future__ import annotations9 10import logging11import math12import warnings13from abc import ABC, abstractmethod14from itertools import cycle15from typing import (16 TYPE_CHECKING,17 Any,18 ClassVar,19 TypeVar,20)21 22from pydantic import ConfigDict, Field, model_validator23from typing_extensions import Self, override24 25from langchain_core.documents import Document26from langchain_core.embeddings import Embeddings27from langchain_core.retrievers import BaseRetriever, LangSmithRetrieverParams28from langchain_core.runnables.config import run_in_executor29 30if TYPE_CHECKING:31 from collections.abc import Callable, Collection, Iterable, Iterator, Sequence32 33 from langchain_core.callbacks.manager import (34 AsyncCallbackManagerForRetrieverRun,35 CallbackManagerForRetrieverRun,36 )37 38logger = logging.getLogger(__name__)39 40VST = TypeVar("VST", bound="VectorStore")41 42 43class VectorStore(ABC):44 """Interface for vector store."""45 46 def add_texts(47 self,48 texts: Iterable[str],49 metadatas: list[dict] | None = None,50 *,51 ids: list[str] | None = None,52 **kwargs: Any,53 ) -> list[str]:54 """Run more texts through the embeddings and add to the `VectorStore`.55 56 Args:57 texts: Iterable of strings to add to the `VectorStore`.58 metadatas: Optional list of metadatas associated with the texts.59 ids: Optional list of IDs associated with the texts.60 **kwargs: `VectorStore` specific parameters.61 62 One of the kwargs should be `ids` which is a list of ids63 associated with the texts.64 65 Returns:66 List of IDs from adding the texts into the `VectorStore`.67 68 Raises:69 ValueError: If the number of metadatas does not match the number of texts.70 ValueError: If the number of IDs does not match the number of texts.71 """72 if type(self).add_documents != VectorStore.add_documents:73 # This condition is triggered if the subclass has provided74 # an implementation of the upsert method.75 # The existing add_texts76 texts_: Sequence[str] = (77 texts if isinstance(texts, (list, tuple)) else list(texts)78 )79 if metadatas and len(metadatas) != len(texts_):80 msg = (81 "The number of metadatas must match the number of texts."82 f"Got {len(metadatas)} metadatas and {len(texts_)} texts."83 )84 raise ValueError(msg)85 metadatas_ = iter(metadatas) if metadatas else cycle([{}])86 ids_: Iterator[str | None] = iter(ids) if ids else cycle([None])87 docs = [88 Document(id=id_, page_content=text, metadata=metadata_)89 for text, metadata_, id_ in zip(texts, metadatas_, ids_, strict=False)90 ]91 if ids is not None:92 # For backward compatibility93 kwargs["ids"] = ids94 95 return self.add_documents(docs, **kwargs)96 msg = f"`add_texts` has not been implemented for {self.__class__.__name__} "97 raise NotImplementedError(msg)98 99 @property100 def embeddings(self) -> Embeddings | None:101 """Access the query embedding object if available."""102 logger.debug(103 "The embeddings property has not been implemented for %s",104 self.__class__.__name__,105 )106 return None107 108 def delete(self, ids: list[str] | None = None, **kwargs: Any) -> bool | None:109 """Delete by vector ID or other criteria.110 111 Args:112 ids: List of IDs to delete. If `None`, delete all.113 **kwargs: Other keyword arguments that subclasses might use.114 115 Returns:116 `True` if deletion is successful, `False` otherwise, `None` if not117 implemented.118 """119 msg = "delete method must be implemented by subclass."120 raise NotImplementedError(msg)121 122 def get_by_ids(self, ids: Sequence[str], /) -> list[Document]:123 """Get documents by their IDs.124 125 The returned documents are expected to have the ID field set to the ID of the126 document in the vector store.127 128 Fewer documents may be returned than requested if some IDs are not found or129 if there are duplicated IDs.130 131 Users should not assume that the order of the returned documents matches132 the order of the input IDs. Instead, users should rely on the ID field of the133 returned documents.134 135 This method should **NOT** raise exceptions if no documents are found for136 some IDs.137 138 Args:139 ids: List of IDs to retrieve.140 141 Returns:142 List of `Document` objects.143 """144 msg = f"{self.__class__.__name__} does not yet support get_by_ids."145 raise NotImplementedError(msg)146 147 # Implementations should override this method to provide an async native version.148 async def aget_by_ids(self, ids: Sequence[str], /) -> list[Document]:149 """Async get documents by their IDs.150 151 The returned documents are expected to have the ID field set to the ID of the152 document in the vector store.153 154 Fewer documents may be returned than requested if some IDs are not found or155 if there are duplicated IDs.156 157 Users should not assume that the order of the returned documents matches158 the order of the input IDs. Instead, users should rely on the ID field of the159 returned documents.160 161 This method should **NOT** raise exceptions if no documents are found for162 some IDs.163 164 Args:165 ids: List of IDs to retrieve.166 167 Returns:168 List of `Document` objects.169 """170 return await run_in_executor(None, self.get_by_ids, ids)171 172 async def adelete(self, ids: list[str] | None = None, **kwargs: Any) -> bool | None:173 """Async delete by vector ID or other criteria.174 175 Args:176 ids: List of IDs to delete. If `None`, delete all.177 **kwargs: Other keyword arguments that subclasses might use.178 179 Returns:180 `True` if deletion is successful, `False` otherwise, `None` if not181 implemented.182 """183 return await run_in_executor(None, self.delete, ids, **kwargs)184 185 async def aadd_texts(186 self,187 texts: Iterable[str],188 metadatas: list[dict] | None = None,189 *,190 ids: list[str] | None = None,191 **kwargs: Any,192 ) -> list[str]:193 """Async run more texts through the embeddings and add to the `VectorStore`.194 195 Args:196 texts: Iterable of strings to add to the `VectorStore`.197 metadatas: Optional list of metadatas associated with the texts.198 ids: Optional list199 **kwargs: `VectorStore` specific parameters.200 201 Returns:202 List of IDs from adding the texts into the `VectorStore`.203 204 Raises:205 ValueError: If the number of metadatas does not match the number of texts.206 ValueError: If the number of IDs does not match the number of texts.207 """208 if ids is not None:209 # For backward compatibility210 kwargs["ids"] = ids211 if type(self).aadd_documents != VectorStore.aadd_documents:212 # This condition is triggered if the subclass has provided213 # an implementation of the upsert method.214 # The existing add_texts215 texts_: Sequence[str] = (216 texts if isinstance(texts, (list, tuple)) else list(texts)217 )218 if metadatas and len(metadatas) != len(texts_):219 msg = (220 "The number of metadatas must match the number of texts."221 f"Got {len(metadatas)} metadatas and {len(texts_)} texts."222 )223 raise ValueError(msg)224 metadatas_ = iter(metadatas) if metadatas else cycle([{}])225 ids_: Iterator[str | None] = iter(ids) if ids else cycle([None])226 227 docs = [228 Document(id=id_, page_content=text, metadata=metadata_)229 for text, metadata_, id_ in zip(texts, metadatas_, ids_, strict=False)230 ]231 return await self.aadd_documents(docs, **kwargs)232 return await run_in_executor(None, self.add_texts, texts, metadatas, **kwargs)233 234 def add_documents(self, documents: list[Document], **kwargs: Any) -> list[str]:235 """Add or update documents in the `VectorStore`.236 237 Args:238 documents: Documents to add to the `VectorStore`.239 **kwargs: Additional keyword arguments.240 241 If kwargs contains IDs and documents contain ids, the IDs in the kwargs242 will receive precedence.243 244 Returns:245 List of IDs of the added texts.246 """247 if type(self).add_texts != VectorStore.add_texts:248 if "ids" not in kwargs:249 ids = [doc.id for doc in documents]250 251 # If there's at least one valid ID, we'll assume that IDs252 # should be used.253 if any(ids):254 kwargs["ids"] = ids255 256 texts = [doc.page_content for doc in documents]257 metadatas = [doc.metadata for doc in documents]258 return self.add_texts(texts, metadatas, **kwargs)259 msg = (260 f"`add_documents` and `add_texts` has not been implemented "261 f"for {self.__class__.__name__} "262 )263 raise NotImplementedError(msg)264 265 async def aadd_documents(266 self, documents: list[Document], **kwargs: Any267 ) -> list[str]:268 """Async run more documents through the embeddings and add to the `VectorStore`.269 270 Args:271 documents: Documents to add to the `VectorStore`.272 **kwargs: Additional keyword arguments.273 274 Returns:275 List of IDs of the added texts.276 """277 # If the async method has been overridden, we'll use that.278 if type(self).aadd_texts != VectorStore.aadd_texts:279 if "ids" not in kwargs:280 ids = [doc.id for doc in documents]281 282 # If there's at least one valid ID, we'll assume that IDs283 # should be used.284 if any(ids):285 kwargs["ids"] = ids286 287 texts = [doc.page_content for doc in documents]288 metadatas = [doc.metadata for doc in documents]289 return await self.aadd_texts(texts, metadatas, **kwargs)290 291 return await run_in_executor(None, self.add_documents, documents, **kwargs)292 293 def search(self, query: str, search_type: str, **kwargs: Any) -> list[Document]:294 """Return docs most similar to query using a specified search type.295 296 Args:297 query: Input text.298 search_type: Type of search to perform.299 300 Can be `'similarity'`, `'mmr'`, or `'similarity_score_threshold'`.301 **kwargs: Arguments to pass to the search method.302 303 Returns:304 List of `Document` objects most similar to the query.305 306 Raises:307 ValueError: If `search_type` is not one of `'similarity'`,308 `'mmr'`, or `'similarity_score_threshold'`.309 """310 if search_type == "similarity":311 return self.similarity_search(query, **kwargs)312 if search_type == "similarity_score_threshold":313 docs_and_similarities = self.similarity_search_with_relevance_scores(314 query, **kwargs315 )316 return [doc for doc, _ in docs_and_similarities]317 if search_type == "mmr":318 return self.max_marginal_relevance_search(query, **kwargs)319 msg = (320 f"search_type of {search_type} not allowed. Expected "321 "search_type to be 'similarity', 'similarity_score_threshold'"322 " or 'mmr'."323 )324 raise ValueError(msg)325 326 async def asearch(327 self, query: str, search_type: str, **kwargs: Any328 ) -> list[Document]:329 """Async return docs most similar to query using a specified search type.330 331 Args:332 query: Input text.333 search_type: Type of search to perform.334 335 Can be `'similarity'`, `'mmr'`, or `'similarity_score_threshold'`.336 **kwargs: Arguments to pass to the search method.337 338 Returns:339 List of `Document` objects most similar to the query.340 341 Raises:342 ValueError: If `search_type` is not one of `'similarity'`,343 `'mmr'`, or `'similarity_score_threshold'`.344 """345 if search_type == "similarity":346 return await self.asimilarity_search(query, **kwargs)347 if search_type == "similarity_score_threshold":348 docs_and_similarities = await self.asimilarity_search_with_relevance_scores(349 query, **kwargs350 )351 return [doc for doc, _ in docs_and_similarities]352 if search_type == "mmr":353 return await self.amax_marginal_relevance_search(query, **kwargs)354 msg = (355 f"search_type of {search_type} not allowed. Expected "356 "search_type to be 'similarity', 'similarity_score_threshold' or 'mmr'."357 )358 raise ValueError(msg)359 360 @abstractmethod361 def similarity_search(362 self, query: str, k: int = 4, **kwargs: Any363 ) -> list[Document]:364 """Return docs most similar to query.365 366 Args:367 query: Input text.368 k: Number of `Document` objects to return.369 **kwargs: Arguments to pass to the search method.370 371 Returns:372 List of `Document` objects most similar to the query.373 """374 375 @staticmethod376 def _euclidean_relevance_score_fn(distance: float) -> float:377 """Return a similarity score on a scale [0, 1]."""378 # The 'correct' relevance function379 # may differ depending on a few things, including:380 # - the distance / similarity metric used by the VectorStore381 # - the scale of your embeddings (OpenAI's are unit normed. Many382 # others are not!)383 # - embedding dimensionality384 # - etc.385 # This function converts the Euclidean norm of normalized embeddings386 # (0 is most similar, sqrt(2) most dissimilar)387 # to a similarity function (0 to 1)388 return 1.0 - distance / math.sqrt(2)389 390 @staticmethod391 def _cosine_relevance_score_fn(distance: float) -> float:392 """Normalize the distance to a score on a scale [0, 1]."""393 return 1.0 - distance394 395 @staticmethod396 def _max_inner_product_relevance_score_fn(distance: float) -> float:397 """Normalize the distance to a score on a scale [0, 1]."""398 if distance > 0:399 return 1.0 - distance400 401 return -1.0 * distance402 403 def _select_relevance_score_fn(self) -> Callable[[float], float]:404 """The 'correct' relevance function.405 406 May differ depending on a few things, including:407 408 - The distance / similarity metric used by the VectorStore409 - The scale of your embeddings (OpenAI's are unit normed. Many others are not!)410 - Embedding dimensionality411 - etc.412 413 Vectorstores should define their own selection-based method of relevance.414 """415 raise NotImplementedError416 417 def similarity_search_with_score(418 self, *args: Any, **kwargs: Any419 ) -> list[tuple[Document, float]]:420 """Run similarity search with distance.421 422 Args:423 *args: Arguments to pass to the search method.424 **kwargs: Arguments to pass to the search method.425 426 Returns:427 List of tuples of `(doc, similarity_score)`.428 """429 raise NotImplementedError430 431 async def asimilarity_search_with_score(432 self, *args: Any, **kwargs: Any433 ) -> list[tuple[Document, float]]:434 """Async run similarity search with distance.435 436 Args:437 *args: Arguments to pass to the search method.438 **kwargs: Arguments to pass to the search method.439 440 Returns:441 List of tuples of `(doc, similarity_score)`.442 """443 # This is a temporary workaround to make the similarity search444 # asynchronous. The proper solution is to make the similarity search445 # asynchronous in the vector store implementations.446 return await run_in_executor(447 None, self.similarity_search_with_score, *args, **kwargs448 )449 450 def _similarity_search_with_relevance_scores(451 self,452 query: str,453 k: int = 4,454 **kwargs: Any,455 ) -> list[tuple[Document, float]]:456 """Default similarity search with relevance scores.457 458 Modify if necessary in subclass.459 Return docs and relevance scores in the range `[0, 1]`.460 461 `0` is dissimilar, `1` is most similar.462 463 Args:464 query: Input text.465 k: Number of `Document` objects to return.466 **kwargs: Kwargs to be passed to similarity search.467 468 Should include `score_threshold`, an optional floating point value469 between `0` to `1` to filter the resulting set of retrieved docs.470 471 Returns:472 List of tuples of `(doc, similarity_score)`473 """474 relevance_score_fn = self._select_relevance_score_fn()475 docs_and_scores = self.similarity_search_with_score(query, k, **kwargs)476 return [(doc, relevance_score_fn(score)) for doc, score in docs_and_scores]477 478 async def _asimilarity_search_with_relevance_scores(479 self,480 query: str,481 k: int = 4,482 **kwargs: Any,483 ) -> list[tuple[Document, float]]:484 """Default similarity search with relevance scores.485 486 Modify if necessary in subclass.487 Return docs and relevance scores in the range `[0, 1]`.488 489 `0` is dissimilar, `1` is most similar.490 491 Args:492 query: Input text.493 k: Number of `Document` objects to return.494 **kwargs: Kwargs to be passed to similarity search.495 496 Should include `score_threshold`, an optional floating point value497 between `0` to `1` to filter the resulting set of retrieved docs.498 499 Returns:500 List of tuples of `(doc, similarity_score)`501 """502 relevance_score_fn = self._select_relevance_score_fn()503 docs_and_scores = await self.asimilarity_search_with_score(query, k, **kwargs)504 return [(doc, relevance_score_fn(score)) for doc, score in docs_and_scores]505 506 def similarity_search_with_relevance_scores(507 self,508 query: str,509 k: int = 4,510 **kwargs: Any,511 ) -> list[tuple[Document, float]]:512 """Return docs and relevance scores in the range `[0, 1]`.513 514 `0` is dissimilar, `1` is most similar.515 516 Args:517 query: Input text.518 k: Number of `Document` objects to return.519 **kwargs: Kwargs to be passed to similarity search.520 521 Should include `score_threshold`, an optional floating point value522 between `0` to `1` to filter the resulting set of retrieved docs.523 524 Returns:525 List of tuples of `(doc, similarity_score)`.526 """527 score_threshold = kwargs.pop("score_threshold", None)528 529 docs_and_similarities = self._similarity_search_with_relevance_scores(530 query, k=k, **kwargs531 )532 if any(533 similarity < 0.0 or similarity > 1.0534 for _, similarity in docs_and_similarities535 ):536 warnings.warn(537 "Relevance scores must be between"538 f" 0 and 1, got {docs_and_similarities}",539 stacklevel=2,540 )541 542 if score_threshold is not None:543 docs_and_similarities = [544 (doc, similarity)545 for doc, similarity in docs_and_similarities546 if similarity >= score_threshold547 ]548 if len(docs_and_similarities) == 0:549 logger.warning(550 "No relevant docs were retrieved using the "551 "relevance score threshold %s",552 score_threshold,553 )554 return docs_and_similarities555 556 async def asimilarity_search_with_relevance_scores(557 self,558 query: str,559 k: int = 4,560 **kwargs: Any,561 ) -> list[tuple[Document, float]]:562 """Async return docs and relevance scores in the range `[0, 1]`.563 564 `0` is dissimilar, `1` is most similar.565 566 Args:567 query: Input text.568 k: Number of `Document` objects to return.569 **kwargs: Kwargs to be passed to similarity search.570 571 Should include `score_threshold`, an optional floating point value572 between `0` to `1` to filter the resulting set of retrieved docs.573 574 Returns:575 List of tuples of `(doc, similarity_score)`576 """577 score_threshold = kwargs.pop("score_threshold", None)578 579 docs_and_similarities = await self._asimilarity_search_with_relevance_scores(580 query, k=k, **kwargs581 )582 if any(583 similarity < 0.0 or similarity > 1.0584 for _, similarity in docs_and_similarities585 ):586 warnings.warn(587 "Relevance scores must be between"588 f" 0 and 1, got {docs_and_similarities}",589 stacklevel=2,590 )591 592 if score_threshold is not None:593 docs_and_similarities = [594 (doc, similarity)595 for doc, similarity in docs_and_similarities596 if similarity >= score_threshold597 ]598 if len(docs_and_similarities) == 0:599 logger.warning(600 "No relevant docs were retrieved using the "601 "relevance score threshold %s",602 score_threshold,603 )604 return docs_and_similarities605 606 async def asimilarity_search(607 self, query: str, k: int = 4, **kwargs: Any608 ) -> list[Document]:609 """Async return docs most similar to query.610 611 Args:612 query: Input text.613 k: Number of `Document` objects to return.614 **kwargs: Arguments to pass to the search method.615 616 Returns:617 List of `Document` objects most similar to the query.618 """619 # This is a temporary workaround to make the similarity search620 # asynchronous. The proper solution is to make the similarity search621 # asynchronous in the vector store implementations.622 return await run_in_executor(None, self.similarity_search, query, k=k, **kwargs)623 624 def similarity_search_by_vector(625 self, embedding: list[float], k: int = 4, **kwargs: Any626 ) -> list[Document]:627 """Return docs most similar to embedding vector.628 629 Args:630 embedding: Embedding to look up documents similar to.631 k: Number of `Document` objects to return.632 **kwargs: Arguments to pass to the search method.633 634 Returns:635 List of `Document` objects most similar to the query vector.636 """637 raise NotImplementedError638 639 async def asimilarity_search_by_vector(640 self, embedding: list[float], k: int = 4, **kwargs: Any641 ) -> list[Document]:642 """Async return docs most similar to embedding vector.643 644 Args:645 embedding: Embedding to look up documents similar to.646 k: Number of `Document` objects to return.647 **kwargs: Arguments to pass to the search method.648 649 Returns:650 List of `Document` objects most similar to the query vector.651 """652 # This is a temporary workaround to make the similarity search653 # asynchronous. The proper solution is to make the similarity search654 # asynchronous in the vector store implementations.655 return await run_in_executor(656 None, self.similarity_search_by_vector, embedding, k=k, **kwargs657 )658 659 def max_marginal_relevance_search(660 self,661 query: str,662 k: int = 4,663 fetch_k: int = 20,664 lambda_mult: float = 0.5,665 **kwargs: Any,666 ) -> list[Document]:667 """Return docs selected using the maximal marginal relevance.668 669 Maximal marginal relevance optimizes for similarity to query AND diversity670 among selected documents.671 672 Args:673 query: Text to look up documents similar to.674 k: Number of `Document` objects to return.675 fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.676 lambda_mult: Number between `0` and `1` that determines the degree of677 diversity among the results with `0` corresponding to maximum diversity678 and `1` to minimum diversity.679 **kwargs: Arguments to pass to the search method.680 681 Returns:682 List of `Document` objects selected by maximal marginal relevance.683 """684 raise NotImplementedError685 686 async def amax_marginal_relevance_search(687 self,688 query: str,689 k: int = 4,690 fetch_k: int = 20,691 lambda_mult: float = 0.5,692 **kwargs: Any,693 ) -> list[Document]:694 """Async return docs selected using the maximal marginal relevance.695 696 Maximal marginal relevance optimizes for similarity to query AND diversity697 among selected documents.698 699 Args:700 query: Text to look up documents similar to.701 k: Number of `Document` objects to return.702 fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.703 lambda_mult: Number between `0` and `1` that determines the degree of704 diversity among the results with `0` corresponding to maximum diversity705 and `1` to minimum diversity.706 **kwargs: Arguments to pass to the search method.707 708 Returns:709 List of `Document` objects selected by maximal marginal relevance.710 """711 # This is a temporary workaround to make the similarity search712 # asynchronous. The proper solution is to make the similarity search713 # asynchronous in the vector store implementations.714 return await run_in_executor(715 None,716 self.max_marginal_relevance_search,717 query,718 k=k,719 fetch_k=fetch_k,720 lambda_mult=lambda_mult,721 **kwargs,722 )723 724 def max_marginal_relevance_search_by_vector(725 self,726 embedding: list[float],727 k: int = 4,728 fetch_k: int = 20,729 lambda_mult: float = 0.5,730 **kwargs: Any,731 ) -> list[Document]:732 """Return docs selected using the maximal marginal relevance.733 734 Maximal marginal relevance optimizes for similarity to query AND diversity735 among selected documents.736 737 Args:738 embedding: Embedding to look up documents similar to.739 k: Number of `Document` objects to return.740 fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.741 lambda_mult: Number between `0` and `1` that determines the degree of742 diversity among the results with `0` corresponding to maximum diversity743 and `1` to minimum diversity.744 **kwargs: Arguments to pass to the search method.745 746 Returns:747 List of `Document` objects selected by maximal marginal relevance.748 """749 raise NotImplementedError750 751 async def amax_marginal_relevance_search_by_vector(752 self,753 embedding: list[float],754 k: int = 4,755 fetch_k: int = 20,756 lambda_mult: float = 0.5,757 **kwargs: Any,758 ) -> list[Document]:759 """Async return docs selected using the maximal marginal relevance.760 761 Maximal marginal relevance optimizes for similarity to query AND diversity762 among selected documents.763 764 Args:765 embedding: Embedding to look up documents similar to.766 k: Number of `Document` objects to return.767 fetch_k: Number of `Document` objects to fetch to pass to MMR algorithm.768 lambda_mult: Number between `0` and `1` that determines the degree of769 diversity among the results with `0` corresponding to maximum diversity770 and `1` to minimum diversity.771 **kwargs: Arguments to pass to the search method.772 773 Returns:774 List of `Document` objects selected by maximal marginal relevance.775 """776 return await run_in_executor(777 None,778 self.max_marginal_relevance_search_by_vector,779 embedding,780 k=k,781 fetch_k=fetch_k,782 lambda_mult=lambda_mult,783 **kwargs,784 )785 786 @classmethod787 def from_documents(788 cls,789 documents: list[Document],790 embedding: Embeddings,791 **kwargs: Any,792 ) -> Self:793 """Return `VectorStore` initialized from documents and embeddings.794 795 Args:796 documents: List of `Document` objects to add to the `VectorStore`.797 embedding: Embedding function to use.798 **kwargs: Additional keyword arguments.799 800 Returns:801 `VectorStore` initialized from documents and embeddings.802 """803 texts = [d.page_content for d in documents]804 metadatas = [d.metadata for d in documents]805 806 if "ids" not in kwargs:807 ids = [doc.id for doc in documents]808 809 # If there's at least one valid ID, we'll assume that IDs810 # should be used.811 if any(ids):812 kwargs["ids"] = ids813 814 return cls.from_texts(texts, embedding, metadatas=metadatas, **kwargs)815 816 @classmethod817 async def afrom_documents(818 cls,819 documents: list[Document],820 embedding: Embeddings,821 **kwargs: Any,822 ) -> Self:823 """Async return `VectorStore` initialized from documents and embeddings.824 825 Args:826 documents: List of `Document` objects to add to the `VectorStore`.827 embedding: Embedding function to use.828 **kwargs: Additional keyword arguments.829 830 Returns:831 `VectorStore` initialized from documents and embeddings.832 """833 texts = [d.page_content for d in documents]834 metadatas = [d.metadata for d in documents]835 836 if "ids" not in kwargs:837 ids = [doc.id for doc in documents]838 839 # If there's at least one valid ID, we'll assume that IDs840 # should be used.841 if any(ids):842 kwargs["ids"] = ids843 844 return await cls.afrom_texts(texts, embedding, metadatas=metadatas, **kwargs)845 846 @classmethod847 @abstractmethod848 def from_texts(849 cls: type[VST],850 texts: list[str],851 embedding: Embeddings,852 metadatas: list[dict] | None = None,853 *,854 ids: list[str] | None = None,855 **kwargs: Any,856 ) -> VST:857 """Return `VectorStore` initialized from texts and embeddings.858 859 Args:860 texts: Texts to add to the `VectorStore`.861 embedding: Embedding function to use.862 metadatas: Optional list of metadatas associated with the texts.863 ids: Optional list of IDs associated with the texts.864 **kwargs: Additional keyword arguments.865 866 Returns:867 `VectorStore` initialized from texts and embeddings.868 """869 870 @classmethod871 async def afrom_texts(872 cls,873 texts: list[str],874 embedding: Embeddings,875 metadatas: list[dict] | None = None,876 *,877 ids: list[str] | None = None,878 **kwargs: Any,879 ) -> Self:880 """Async return `VectorStore` initialized from texts and embeddings.881 882 Args:883 texts: Texts to add to the `VectorStore`.884 embedding: Embedding function to use.885 metadatas: Optional list of metadatas associated with the texts.886 ids: Optional list of IDs associated with the texts.887 **kwargs: Additional keyword arguments.888 889 Returns:890 `VectorStore` initialized from texts and embeddings.891 """892 if ids is not None:893 kwargs["ids"] = ids894 return await run_in_executor(895 None, cls.from_texts, texts, embedding, metadatas, **kwargs896 )897 898 def _get_retriever_tags(self) -> list[str]:899 """Get tags for retriever."""900 tags = [self.__class__.__name__]901 if self.embeddings:902 tags.append(self.embeddings.__class__.__name__)903 return tags904 905 def as_retriever(self, **kwargs: Any) -> VectorStoreRetriever:906 """Return `VectorStoreRetriever` initialized from this `VectorStore`.907 908 Args:909 **kwargs: Keyword arguments to pass to the search function.910 911 Can include:912 913 * `search_type`: Defines the type of search that the Retriever should914 perform. Can be `'similarity'` (default), `'mmr'`, or915 `'similarity_score_threshold'`.916 * `search_kwargs`: Keyword arguments to pass to the search function.917 918 Can include things like:919 920 * `k`: Amount of documents to return (Default: `4`)921 * `score_threshold`: Minimum relevance threshold922 for `similarity_score_threshold`923 * `fetch_k`: Amount of documents to pass to MMR algorithm924 (Default: `20`)925 * `lambda_mult`: Diversity of results returned by MMR;926 `1` for minimum diversity and 0 for maximum. (Default: `0.5`)927 * `filter`: Filter by document metadata928 929 Returns:930 Retriever class for `VectorStore`.931 932 Examples:933 ```python934 # Retrieve more documents with higher diversity935 # Useful if your dataset has many similar documents936 docsearch.as_retriever(937 search_type="mmr", search_kwargs={"k": 6, "lambda_mult": 0.25}938 )939 940 # Fetch more documents for the MMR algorithm to consider941 # But only return the top 5942 docsearch.as_retriever(search_type="mmr", search_kwargs={"k": 5, "fetch_k": 50})943 944 # Only retrieve documents that have a relevance score945 # Above a certain threshold946 docsearch.as_retriever(947 search_type="similarity_score_threshold",948 search_kwargs={"score_threshold": 0.8},949 )950 951 # Only get the single most similar document from the dataset952 docsearch.as_retriever(search_kwargs={"k": 1})953 954 # Use a filter to only retrieve documents from a specific paper955 docsearch.as_retriever(956 search_kwargs={"filter": {"paper_title": "GPT-4 Technical Report"}}957 )958 ```959 """960 tags = kwargs.pop("tags", None) or [*self._get_retriever_tags()]961 return VectorStoreRetriever(vectorstore=self, tags=tags, **kwargs)962 963 964class VectorStoreRetriever(BaseRetriever):965 """Base Retriever class for VectorStore."""966 967 vectorstore: VectorStore968 """VectorStore to use for retrieval."""969 970 search_type: str = "similarity"971 """Type of search to perform."""972 973 search_kwargs: dict = Field(default_factory=dict)974 """Keyword arguments to pass to the search function."""975 976 allowed_search_types: ClassVar[Collection[str]] = (977 "similarity",978 "similarity_score_threshold",979 "mmr",980 )981 982 model_config = ConfigDict(983 arbitrary_types_allowed=True,984 )985 986 @model_validator(mode="before")987 @classmethod988 def validate_search_type(cls, values: dict) -> Any:989 """Validate search type.990 991 Args:992 values: Values to validate.993 994 Returns:995 Validated values.996 997 Raises:998 ValueError: If `search_type` is not one of the allowed search types.999 ValueError: If `score_threshold` is not specified with a float value(`0~1`)1000 """1001 search_type = values.get("search_type", "similarity")1002 if search_type not in cls.allowed_search_types:1003 msg = (1004 f"search_type of {search_type} not allowed. Valid values are: "1005 f"{cls.allowed_search_types}"1006 )1007 raise ValueError(msg)1008 if search_type == "similarity_score_threshold":1009 score_threshold = values.get("search_kwargs", {}).get("score_threshold")1010 if (score_threshold is None) or (not isinstance(score_threshold, float)):1011 msg = (1012 "`score_threshold` is not specified with a float value(0~1) "1013 "in `search_kwargs`."1014 )1015 raise ValueError(msg)1016 return values1017 1018 def _get_ls_params(self, **kwargs: Any) -> LangSmithRetrieverParams:1019 """Get standard params for tracing."""1020 kwargs_ = self.search_kwargs | kwargs1021 1022 ls_params = super()._get_ls_params(**kwargs_)1023 1024 ls_params["ls_vector_store_provider"] = self.vectorstore.__class__.__name__1025 1026 if self.vectorstore.embeddings:1027 ls_params["ls_embedding_provider"] = (1028 self.vectorstore.embeddings.__class__.__name__1029 )1030 elif hasattr(self.vectorstore, "embedding") and isinstance(1031 self.vectorstore.embedding, Embeddings1032 ):1033 ls_params["ls_embedding_provider"] = (1034 self.vectorstore.embedding.__class__.__name__1035 )1036 1037 return ls_params1038 1039 @override1040 def _get_relevant_documents(1041 self, query: str, *, run_manager: CallbackManagerForRetrieverRun, **kwargs: Any1042 ) -> list[Document]:1043 kwargs_ = self.search_kwargs | kwargs1044 if self.search_type == "similarity":1045 docs = self.vectorstore.similarity_search(query, **kwargs_)1046 elif self.search_type == "similarity_score_threshold":1047 docs_and_similarities = (1048 self.vectorstore.similarity_search_with_relevance_scores(1049 query, **kwargs_1050 )1051 )1052 docs = [doc for doc, _ in docs_and_similarities]1053 elif self.search_type == "mmr":1054 docs = self.vectorstore.max_marginal_relevance_search(query, **kwargs_)1055 else:1056 msg = f"search_type of {self.search_type} not allowed."1057 raise ValueError(msg)1058 return docs1059 1060 @override1061 async def _aget_relevant_documents(1062 self,1063 query: str,1064 *,1065 run_manager: AsyncCallbackManagerForRetrieverRun,1066 **kwargs: Any,1067 ) -> list[Document]:1068 kwargs_ = self.search_kwargs | kwargs1069 if self.search_type == "similarity":1070 docs = await self.vectorstore.asimilarity_search(query, **kwargs_)1071 elif self.search_type == "similarity_score_threshold":1072 docs_and_similarities = (1073 await self.vectorstore.asimilarity_search_with_relevance_scores(1074 query, **kwargs_1075 )1076 )1077 docs = [doc for doc, _ in docs_and_similarities]1078 elif self.search_type == "mmr":1079 docs = await self.vectorstore.amax_marginal_relevance_search(1080 query, **kwargs_1081 )1082 else:1083 msg = f"search_type of {self.search_type} not allowed."1084 raise ValueError(msg)1085 return docs1086 1087 def add_documents(self, documents: list[Document], **kwargs: Any) -> list[str]:1088 """Add documents to the `VectorStore`.1089 1090 Args:1091 documents: Documents to add to the `VectorStore`.1092 **kwargs: Other keyword arguments that subclasses might use.1093 1094 Returns:1095 List of IDs of the added texts.1096 """1097 return self.vectorstore.add_documents(documents, **kwargs)1098 1099 async def aadd_documents(1100 self, documents: list[Document], **kwargs: Any1101 ) -> list[str]:1102 """Async add documents to the `VectorStore`.1103 1104 Args:1105 documents: Documents to add to the `VectorStore`.1106 **kwargs: Other keyword arguments that subclasses might use.1107 1108 Returns:1109 List of IDs of the added texts.1110 """1111 return await self.vectorstore.aadd_documents(documents, **kwargs)1112 