"Please install it with `pip install zep-python`."
            )
        values["zep_client"] = values.get(
            "zep_client",
            ZepClient(base_url=values["url"], api_key=values.get("api_key")),
        )
        return values
    def _messages_search_result_to_doc(
        self, results: List[MemorySearchResult]
    ) -> List[Document]:
        return [
            Document(
                page_content=r.message.pop("content"),
                metadata={"score": r.dist, **r.message},
            )
            for r in results
            if r.message
        ]
    def _summary_search_result_to_doc(
        self, results: List[MemorySearchResult]
    ) -> List[Document]:
        return [
            Document(
                page_content=r.summary.content,
                metadata={
                    "score": r.dist,
                    "uuid": r.summary.uuid,
                    "created_at": r.summary.created_at,
                    "token_count": r.summary.token_count,
                },
            )
            for r in results
            if r.summary
        ]
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> List[Document]:
        from zep_python.memory import MemorySearchPayload
        if not self.zep_client:
            raise RuntimeError("Zep client not initialized.")
        payload = MemorySearchPayload(
            text=query,
            metadata=metadata,
            search_scope=self.search_scope,
            search_type=self.search_type,
            mmr_lambda=self.mmr_lambda,
        )


mmr_lambda=self.mmr_lambda,
        )
        results: List[MemorySearchResult] = self.zep_client.memory.search_memory(
            self.session_id, payload, limit=self.top_k
        )
        if self.search_scope == SearchScope.summary:
            return self._summary_search_result_to_doc(results)
        return self._messages_search_result_to_doc(results)
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> List[Document]:
        from zep_python.memory import MemorySearchPayload
        if not self.zep_client:
            raise RuntimeError("Zep client not initialized.")
        payload = MemorySearchPayload(
            text=query,
            metadata=metadata,
            search_scope=self.search_scope,
            search_type=self.search_type,
            mmr_lambda=self.mmr_lambda,
        )
        results: List[MemorySearchResult] = await self.zep_client.memory.asearch_memory(
            self.session_id, payload, limit=self.top_k
        )
        if self.search_scope == SearchScope.summary:
            return self._summary_search_result_to_doc(results)
        return self._messages_search_result_to_doc(results)


Source code for langchain.retrievers.kay
from __future__ import annotations
from typing import Any, List
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
[docs]class KayAiRetriever(BaseRetriever):
    """
    Retriever for Kay.ai datasets.
    To work properly, expects you to have KAY_API_KEY env variable set.
    You can get one for free at https://kay.ai/.
    """
    client: Any
    num_contexts: int
[docs]    @classmethod
    def create(
        cls,
        dataset_id: str,
        data_types: List[str],
        num_contexts: int = 6,
    ) -> KayAiRetriever:
        """
        Create a KayRetriever given a Kay dataset id and a list of datasources.
        Args:
            dataset_id: A dataset id category in Kay, like "company"
            data_types: A list of datasources present within a dataset. For
                "company" the corresponding datasources could be
                ["10-K", "10-Q", "8-K", "PressRelease"].
            num_contexts: The number of documents to retrieve on each query.
                Defaults to 6.
        """
        try:
            from kay.rag.retrievers import KayRetriever
        except ImportError:
            raise ImportError(
                "Could not import kay python package. Please install it with "
                "`pip install kay`.",
            )
        client = KayRetriever(dataset_id, data_types)
        return cls(client=client, num_contexts=num_contexts)
    def _get_relevant_documents(


def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        ctxs = self.client.query(query=query, num_context=self.num_contexts)
        docs = []
        for ctx in ctxs:
            page_content = ctx.pop("chunk_embed_text", None)
            if page_content is None:
                continue
            docs.append(Document(page_content=page_content, metadata={**ctx}))
        return docs


Source code for langchain.retrievers.zilliz
import warnings
from typing import Any, Dict, List, Optional
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.pydantic_v1 import root_validator
from langchain.schema import BaseRetriever, Document
from langchain.schema.embeddings import Embeddings
from langchain.vectorstores.zilliz import Zilliz
# TODO: Update to ZillizClient + Hybrid Search when available
[docs]class ZillizRetriever(BaseRetriever):
    """`Zilliz API` retriever."""
    embedding_function: Embeddings
    """The underlying embedding function from which documents will be retrieved."""
    collection_name: str = "LangChainCollection"
    """The name of the collection in Zilliz."""
    connection_args: Optional[Dict[str, Any]] = None
    """The connection arguments for the Zilliz client."""
    consistency_level: str = "Session"
    """The consistency level for the Zilliz client."""
    search_params: Optional[dict] = None
    """The search parameters for the Zilliz client."""
    store: Zilliz
    """The underlying Zilliz store."""
    retriever: BaseRetriever
    """The underlying retriever."""
    @root_validator(pre=True)
    def create_client(cls, values: dict) -> dict:
        values["store"] = Zilliz(
            values["embedding_function"],
            values["collection_name"],
            values["connection_args"],
            values["consistency_level"],
        )
        values["retriever"] = values["store"].as_retriever(
            search_kwargs={"param": values["search_params"]}
        )
        return values
[docs]    def add_texts(


)
        return values
[docs]    def add_texts(
        self, texts: List[str], metadatas: Optional[List[dict]] = None
    ) -> None:
        """Add text to the Zilliz store
        Args:
            texts (List[str]): The text
            metadatas (List[dict]): Metadata dicts, must line up with existing store
        """
        self.store.add_texts(texts, metadatas)
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        return self.retriever.get_relevant_documents(
            query, run_manager=run_manager.get_child(), **kwargs
        )
[docs]def ZillizRetreiver(*args: Any, **kwargs: Any) -> ZillizRetriever:
    """Deprecated ZillizRetreiver.
    Please use ZillizRetriever ('i' before 'e') instead.
    Args:
        *args:
        **kwargs:
    Returns:
        ZillizRetriever
    """
    warnings.warn(
        "ZillizRetreiver will be deprecated in the future. "
        "Please use ZillizRetriever ('i' before 'e') instead.",
        DeprecationWarning,
    )
    return ZillizRetriever(*args, **kwargs)


Source code for langchain.retrievers.contextual_compression
from typing import Any, List
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.retrievers.document_compressors.base import (
    BaseDocumentCompressor,
)
from langchain.schema import BaseRetriever, Document
[docs]class ContextualCompressionRetriever(BaseRetriever):
    """Retriever that wraps a base retriever and compresses the results."""
    base_compressor: BaseDocumentCompressor
    """Compressor for compressing retrieved documents."""
    base_retriever: BaseRetriever
    """Base Retriever to use for getting relevant documents."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        """Get documents relevant for a query.
        Args:
            query: string to find relevant documents for
        Returns:
            Sequence of relevant documents
        """
        docs = self.base_retriever.get_relevant_documents(
            query, callbacks=run_manager.get_child(), **kwargs
        )
        if docs:
            compressed_docs = self.base_compressor.compress_documents(
                docs, query, callbacks=run_manager.get_child()
            )
            return list(compressed_docs)
        else:
            return []
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,


run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        """Get documents relevant for a query.
        Args:
            query: string to find relevant documents for
        Returns:
            List of relevant documents
        """
        docs = await self.base_retriever.aget_relevant_documents(
            query, callbacks=run_manager.get_child(), **kwargs
        )
        if docs:
            compressed_docs = await self.base_compressor.acompress_documents(
                docs, query, callbacks=run_manager.get_child()
            )
            return list(compressed_docs)
        else:
            return []


Source code for langchain.retrievers.knn
"""KNN Retriever.
Largely based on
https://github.com/karpathy/randomfun/blob/master/knn_vs_svm.ipynb"""
from __future__ import annotations
import concurrent.futures
from typing import Any, List, Optional
import numpy as np
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
from langchain.schema.embeddings import Embeddings
[docs]def create_index(contexts: List[str], embeddings: Embeddings) -> np.ndarray:
    """
    Create an index of embeddings for a list of contexts.
    Args:
        contexts: List of contexts to embed.
        embeddings: Embeddings model to use.
    Returns:
        Index of embeddings.
    """
    with concurrent.futures.ThreadPoolExecutor() as executor:
        return np.array(list(executor.map(embeddings.embed_query, contexts)))
[docs]class KNNRetriever(BaseRetriever):
    """`KNN` retriever."""
    embeddings: Embeddings
    """Embeddings model to use."""
    index: Any
    """Index of embeddings."""
    texts: List[str]
    """List of texts to index."""
    k: int = 4
    """Number of results to return."""
    relevancy_threshold: Optional[float] = None
    """Threshold for relevancy."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
[docs]    @classmethod
    def from_texts(
        cls, texts: List[str], embeddings: Embeddings, **kwargs: Any
    ) -> KNNRetriever:
        index = create_index(texts, embeddings)


index = create_index(texts, embeddings)
        return cls(embeddings=embeddings, index=index, texts=texts, **kwargs)
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        query_embeds = np.array(self.embeddings.embed_query(query))
        # calc L2 norm
        index_embeds = self.index / np.sqrt((self.index**2).sum(1, keepdims=True))
        query_embeds = query_embeds / np.sqrt((query_embeds**2).sum())
        similarities = index_embeds.dot(query_embeds)
        sorted_ix = np.argsort(-similarities)
        denominator = np.max(similarities) - np.min(similarities) + 1e-6
        normalized_similarities = (similarities - np.min(similarities)) / denominator
        top_k_results = [
            Document(page_content=self.texts[row])
            for row in sorted_ix[0 : self.k]
            if (
                self.relevancy_threshold is None
                or normalized_similarities[row] >= self.relevancy_threshold
            )
        ]
        return top_k_results


Source code for langchain.retrievers.cohere_rag_retriever
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Dict, List
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.chat_models.base import BaseChatModel
from langchain.pydantic_v1 import Field
from langchain.schema import BaseRetriever, Document, HumanMessage
if TYPE_CHECKING:
    from langchain.schema.messages import BaseMessage
def _get_docs(response: Any) -> List[Document]:
    docs = [
        Document(page_content=doc["snippet"], metadata=doc)
        for doc in response.generation_info["documents"]
    ]
    docs.append(
        Document(
            page_content=response.message.content,
            metadata={
                "type": "model_response",
                "citations": response.generation_info["citations"],
                "search_results": response.generation_info["search_results"],
                "search_queries": response.generation_info["search_queries"],
                "token_count": response.generation_info["token_count"],
            },
        )
    )
    return docs
[docs]class CohereRagRetriever(BaseRetriever):
    """Cohere Chat API with RAG."""
    connectors: List[Dict] = Field(default_factory=lambda: [{"id": "web-search"}])
    """
    When specified, the model's reply will be enriched with information found by
    querying each of the connectors (RAG). These will be returned as langchain
    documents.
    Currently only accepts {"id": "web-search"}.
    """
    llm: BaseChatModel
    """Cohere ChatModel to use."""
    class Config:


"""Cohere ChatModel to use."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
        """Allow arbitrary types."""
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun, **kwargs: Any
    ) -> List[Document]:
        messages: List[List[BaseMessage]] = [[HumanMessage(content=query)]]
        res = self.llm.generate(
            messages,
            connectors=self.connectors,
            callbacks=run_manager.get_child(),
            **kwargs,
        ).generations[0][0]
        return _get_docs(res)
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        messages: List[List[BaseMessage]] = [[HumanMessage(content=query)]]
        res = (
            await self.llm.agenerate(
                messages,
                connectors=self.connectors,
                callbacks=run_manager.get_child(),
                **kwargs,
            )
        ).generations[0][0]
        return _get_docs(res)


Source code for langchain.retrievers.kendra
import re
from abc import ABC, abstractmethod
from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Union
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.docstore.document import Document
from langchain.pydantic_v1 import BaseModel, Extra, root_validator, validator
from langchain.schema import BaseRetriever
[docs]def clean_excerpt(excerpt: str) -> str:
    """Clean an excerpt from Kendra.
    Args:
        excerpt: The excerpt to clean.
    Returns:
        The cleaned excerpt.
    """
    if not excerpt:
        return excerpt
    res = re.sub(r"\s+", " ", excerpt).replace("...", "")
    return res
[docs]def combined_text(item: "ResultItem") -> str:
    """Combine a ResultItem title and excerpt into a single string.
    Args:
        item: the ResultItem of a Kendra search.
    Returns:
        A combined text of the title and excerpt of the given item.
    """
    text = ""
    title = item.get_title()
    if title:
        text += f"Document Title: {title}\n"
    excerpt = clean_excerpt(item.get_excerpt())
    if excerpt:
        text += f"Document Excerpt: \n{excerpt}\n"
    return text
DocumentAttributeValueType = Union[str, int, List[str], None]
"""Possible types of a DocumentAttributeValue.
Dates are also represented as str.
"""
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class Highlight(BaseModel, extra=Extra.allow):  # type: ignore[call-arg]


"""Information that highlights the keywords in the excerpt."""
    BeginOffset: int
    """The zero-based location in the excerpt where the highlight starts."""
    EndOffset: int
    """The zero-based location in the excerpt where the highlight ends."""
    TopAnswer: Optional[bool]
    """Indicates whether the result is the best one."""
    Type: Optional[str]
    """The highlight type: STANDARD or THESAURUS_SYNONYM."""
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class TextWithHighLights(BaseModel, extra=Extra.allow):  # type: ignore[call-arg]
    """Text with highlights."""
    Text: str
    """The text."""
    Highlights: Optional[Any]
    """The highlights."""
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class AdditionalResultAttributeValue(  # type: ignore[call-arg]
    BaseModel, extra=Extra.allow
):
    """Value of an additional result attribute."""
    TextWithHighlightsValue: TextWithHighLights
    """The text with highlights value."""
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class AdditionalResultAttribute(BaseModel, extra=Extra.allow):  # type: ignore[call-arg]
    """Additional result attribute."""
    Key: str
    """The key of the attribute."""
    ValueType: Literal["TEXT_WITH_HIGHLIGHTS_VALUE"]
    """The type of the value."""
    Value: AdditionalResultAttributeValue
    """The value of the attribute."""
[docs]    def get_value_text(self) -> str:
        return self.Value.TextWithHighlightsValue.Text


return self.Value.TextWithHighlightsValue.Text
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class DocumentAttributeValue(BaseModel, extra=Extra.allow):  # type: ignore[call-arg]
    """Value of a document attribute."""
    DateValue: Optional[str]
    """The date expressed as an ISO 8601 string."""
    LongValue: Optional[int]
    """The long value."""
    StringListValue: Optional[List[str]]
    """The string list value."""
    StringValue: Optional[str]
    """The string value."""
    @property
    def value(self) -> DocumentAttributeValueType:
        """The only defined document attribute value or None.
        According to Amazon Kendra, you can only provide one
        value for a document attribute.
        """
        if self.DateValue:
            return self.DateValue
        if self.LongValue:
            return self.LongValue
        if self.StringListValue:
            return self.StringListValue
        if self.StringValue:
            return self.StringValue
        return None
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class DocumentAttribute(BaseModel, extra=Extra.allow):  # type: ignore[call-arg]
    """Document attribute."""
    Key: str
    """The key of the attribute."""
    Value: DocumentAttributeValue
    """The value of the attribute."""
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class ResultItem(BaseModel, ABC, extra=Extra.allow):  # type: ignore[call-arg]
    """Base class of a result item."""
    Id: Optional[str]
    """The ID of the relevant result item."""


Id: Optional[str]
    """The ID of the relevant result item."""
    DocumentId: Optional[str]
    """The document ID."""
    DocumentURI: Optional[str]
    """The document URI."""
    DocumentAttributes: Optional[List[DocumentAttribute]] = []
    """The document attributes."""
[docs]    @abstractmethod
    def get_title(self) -> str:
        """Document title."""
[docs]    @abstractmethod
    def get_excerpt(self) -> str:
        """Document excerpt or passage original content as retrieved by Kendra."""
[docs]    def get_additional_metadata(self) -> dict:
        """Document additional metadata dict.
        This returns any extra metadata except these:
            * result_id
            * document_id
            * source
            * title
            * excerpt
            * document_attributes
        """
        return {}
[docs]    def get_document_attributes_dict(self) -> Dict[str, DocumentAttributeValueType]:
        """Document attributes dict."""
        return {attr.Key: attr.Value.value for attr in (self.DocumentAttributes or [])}
[docs]    def to_doc(
        self, page_content_formatter: Callable[["ResultItem"], str] = combined_text
    ) -> Document:
        """Converts this item to a Document."""
        page_content = page_content_formatter(self)
        metadata = self.get_additional_metadata()
        metadata.update(
            {
                "result_id": self.Id,
                "document_id": self.DocumentId,
                "source": self.DocumentURI,
                "title": self.get_title(),
                "excerpt": self.get_excerpt(),
                "document_attributes": self.get_document_attributes_dict(),
            }
        )
        return Document(page_content=page_content, metadata=metadata)
[docs]class QueryResultItem(ResultItem):


[docs]class QueryResultItem(ResultItem):
    """Query API result item."""
    DocumentTitle: TextWithHighLights
    """The document title."""
    FeedbackToken: Optional[str]
    """Identifies a particular result from a particular query."""
    Format: Optional[str]
    """
    If the Type is ANSWER, then format is either:
        * TABLE: a table excerpt is returned in TableExcerpt;
        * TEXT: a text excerpt is returned in DocumentExcerpt.
    """
    Type: Optional[str]
    """Type of result: DOCUMENT or QUESTION_ANSWER or ANSWER"""
    AdditionalAttributes: Optional[List[AdditionalResultAttribute]] = []
    """One or more additional attributes associated with the result."""
    DocumentExcerpt: Optional[TextWithHighLights]
    """Excerpt of the document text."""
[docs]    def get_title(self) -> str:
        return self.DocumentTitle.Text
[docs]    def get_attribute_value(self) -> str:
        if not self.AdditionalAttributes:
            return ""
        if not self.AdditionalAttributes[0]:
            return ""
        else:
            return self.AdditionalAttributes[0].get_value_text()
[docs]    def get_excerpt(self) -> str:
        if (
            self.AdditionalAttributes
            and self.AdditionalAttributes[0].Key == "AnswerText"
        ):
            excerpt = self.get_attribute_value()
        elif self.DocumentExcerpt:
            excerpt = self.DocumentExcerpt.Text
        else:
            excerpt = ""
        return excerpt
[docs]    def get_additional_metadata(self) -> dict:
        additional_metadata = {"type": self.Type}
        return additional_metadata
[docs]class RetrieveResultItem(ResultItem):
    """Retrieve API result item."""


[docs]class RetrieveResultItem(ResultItem):
    """Retrieve API result item."""
    DocumentTitle: Optional[str]
    """The document title."""
    Content: Optional[str]
    """The content of the item."""
[docs]    def get_title(self) -> str:
        return self.DocumentTitle or ""
[docs]    def get_excerpt(self) -> str:
        return self.Content or ""
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class QueryResult(BaseModel, extra=Extra.allow):  # type: ignore[call-arg]
    """`Amazon Kendra Query API` search result.
    It is composed of:
        * Relevant suggested answers: either a text excerpt or table excerpt.
        * Matching FAQs or questions-answer from your FAQ file.
        * Documents including an excerpt of each document with its title.
    """
    ResultItems: List[QueryResultItem]
    """The result items."""
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class RetrieveResult(BaseModel, extra=Extra.allow):  # type: ignore[call-arg]
    """`Amazon Kendra Retrieve API` search result.
    It is composed of:
        * relevant passages or text excerpts given an input query.
    """
    QueryId: str
    """The ID of the query."""
    ResultItems: List[RetrieveResultItem]
    """The result items."""
[docs]class AmazonKendraRetriever(BaseRetriever):
    """`Amazon Kendra Index` retriever.
    Args:
        index_id: Kendra index id
        region_name: The aws region e.g., `us-west-2`.
            Fallsback to AWS_DEFAULT_REGION env variable


Fallsback to AWS_DEFAULT_REGION env variable
            or region specified in ~/.aws/config.
        credentials_profile_name: The name of the profile in the ~/.aws/credentials
            or ~/.aws/config files, which has either access keys or role information
            specified. If not specified, the default credential profile or, if on an
            EC2 instance, credentials from IMDS will be used.
        top_k: No of results to return
        attribute_filter: Additional filtering of results based on metadata
            See: https://docs.aws.amazon.com/kendra/latest/APIReference
        page_content_formatter: generates the Document page_content
            allowing access to all result item attributes. By default, it uses
            the item's title and excerpt.
        client: boto3 client for Kendra
        user_context: Provides information about the user context
            See: https://docs.aws.amazon.com/kendra/latest/APIReference
    Example:
        .. code-block:: python
            retriever = AmazonKendraRetriever(
                index_id="c0806df7-e76b-4bce-9b5c-d5582f6b1a03"
            )
    """
    index_id: str
    region_name: Optional[str] = None
    credentials_profile_name: Optional[str] = None
    top_k: int = 3
    attribute_filter: Optional[Dict] = None
    page_content_formatter: Callable[[ResultItem], str] = combined_text
    client: Any
    user_context: Optional[Dict] = None
    @validator("top_k")
    def validate_top_k(cls, value: int) -> int:
        if value < 0:
            raise ValueError(f"top_k ({value}) cannot be negative.")
        return value
    @root_validator(pre=True)


return value
    @root_validator(pre=True)
    def create_client(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        if values.get("client") is not None:
            return values
        try:
            import boto3
            if values.get("credentials_profile_name"):
                session = boto3.Session(profile_name=values["credentials_profile_name"])
            else:
                # use default credentials
                session = boto3.Session()
            client_params = {}
            if values.get("region_name"):
                client_params["region_name"] = values["region_name"]
            values["client"] = session.client("kendra", **client_params)
            return values
        except ImportError:
            raise ModuleNotFoundError(
                "Could not import boto3 python package. "
                "Please install it with `pip install boto3`."
            )
        except Exception as e:
            raise ValueError(
                "Could not load credentials to authenticate with AWS client. "
                "Please check that credentials in the specified "
                "profile name are valid."
            ) from e
    def _kendra_query(self, query: str) -> Sequence[ResultItem]:
        kendra_kwargs = {
            "IndexId": self.index_id,
            "QueryText": query.strip(),
            "PageSize": self.top_k,
        }
        if self.attribute_filter is not None:
            kendra_kwargs["AttributeFilter"] = self.attribute_filter
        if self.user_context is not None:
            kendra_kwargs["UserContext"] = self.user_context
        response = self.client.retrieve(**kendra_kwargs)
        r_result = RetrieveResult.parse_obj(response)
        if r_result.ResultItems:
            return r_result.ResultItems
        # Retrieve API returned 0 results, fall back to Query API


# Retrieve API returned 0 results, fall back to Query API
        response = self.client.query(**kendra_kwargs)
        q_result = QueryResult.parse_obj(response)
        return q_result.ResultItems
    def _get_top_k_docs(self, result_items: Sequence[ResultItem]) -> List[Document]:
        top_docs = [
            item.to_doc(self.page_content_formatter)
            for item in result_items[: self.top_k]
        ]
        return top_docs
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """Run search on Kendra index and get top k documents
        Example:
        .. code-block:: python
            docs = retriever.get_relevant_documents('This is my query')
        """
        result_items = self._kendra_query(query)
        top_k_docs = self._get_top_k_docs(result_items)
        return top_k_docs


Source code for langchain.retrievers.bm25
from __future__ import annotations
from typing import Any, Callable, Dict, Iterable, List, Optional
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
[docs]def default_preprocessing_func(text: str) -> List[str]:
    return text.split()
[docs]class BM25Retriever(BaseRetriever):
    """`BM25` retriever without Elasticsearch."""
    vectorizer: Any
    """ BM25 vectorizer."""
    docs: List[Document]
    """ List of documents."""
    k: int = 4
    """ Number of documents to return."""
    preprocess_func: Callable[[str], List[str]] = default_preprocessing_func
    """ Preprocessing function to use on the text before BM25 vectorization."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
[docs]    @classmethod
    def from_texts(
        cls,
        texts: Iterable[str],
        metadatas: Optional[Iterable[dict]] = None,
        bm25_params: Optional[Dict[str, Any]] = None,
        preprocess_func: Callable[[str], List[str]] = default_preprocessing_func,
        **kwargs: Any,
    ) -> BM25Retriever:
        """
        Create a BM25Retriever from a list of texts.
        Args:
            texts: A list of texts to vectorize.
            metadatas: A list of metadata dicts to associate with each text.
            bm25_params: Parameters to pass to the BM25 vectorizer.
            preprocess_func: A function to preprocess each text before vectorization.
            **kwargs: Any other arguments to pass to the retriever.


**kwargs: Any other arguments to pass to the retriever.
        Returns:
            A BM25Retriever instance.
        """
        try:
            from rank_bm25 import BM25Okapi
        except ImportError:
            raise ImportError(
                "Could not import rank_bm25, please install with `pip install "
                "rank_bm25`."
            )
        texts_processed = [preprocess_func(t) for t in texts]
        bm25_params = bm25_params or {}
        vectorizer = BM25Okapi(texts_processed, **bm25_params)
        metadatas = metadatas or ({} for _ in texts)
        docs = [Document(page_content=t, metadata=m) for t, m in zip(texts, metadatas)]
        return cls(
            vectorizer=vectorizer, docs=docs, preprocess_func=preprocess_func, **kwargs
        )
[docs]    @classmethod
    def from_documents(
        cls,
        documents: Iterable[Document],
        *,
        bm25_params: Optional[Dict[str, Any]] = None,
        preprocess_func: Callable[[str], List[str]] = default_preprocessing_func,
        **kwargs: Any,
    ) -> BM25Retriever:
        """
        Create a BM25Retriever from a list of Documents.
        Args:
            documents: A list of Documents to vectorize.
            bm25_params: Parameters to pass to the BM25 vectorizer.
            preprocess_func: A function to preprocess each text before vectorization.
            **kwargs: Any other arguments to pass to the retriever.
        Returns:
            A BM25Retriever instance.
        """


Returns:
            A BM25Retriever instance.
        """
        texts, metadatas = zip(*((d.page_content, d.metadata) for d in documents))
        return cls.from_texts(
            texts=texts,
            bm25_params=bm25_params,
            metadatas=metadatas,
            preprocess_func=preprocess_func,
            **kwargs,
        )
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        processed_query = self.preprocess_func(query)
        return_docs = self.vectorizer.get_top_n(processed_query, self.docs, n=self.k)
        return return_docs


Source code for langchain.retrievers.tfidf
from __future__ import annotations
import pickle
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
[docs]class TFIDFRetriever(BaseRetriever):
    """`TF-IDF` retriever.
    Largely based on
    https://github.com/asvskartheek/Text-Retrieval/blob/master/TF-IDF%20Search%20Engine%20(SKLEARN).ipynb
    """
    vectorizer: Any
    """TF-IDF vectorizer."""
    docs: List[Document]
    """Documents."""
    tfidf_array: Any
    """TF-IDF array."""
    k: int = 4
    """Number of documents to return."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
[docs]    @classmethod
    def from_texts(
        cls,
        texts: Iterable[str],
        metadatas: Optional[Iterable[dict]] = None,
        tfidf_params: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> TFIDFRetriever:
        try:
            from sklearn.feature_extraction.text import TfidfVectorizer
        except ImportError:
            raise ImportError(
                "Could not import scikit-learn, please install with `pip install "
                "scikit-learn`."
            )
        tfidf_params = tfidf_params or {}
        vectorizer = TfidfVectorizer(**tfidf_params)
        tfidf_array = vectorizer.fit_transform(texts)


tfidf_array = vectorizer.fit_transform(texts)
        metadatas = metadatas or ({} for _ in texts)
        docs = [Document(page_content=t, metadata=m) for t, m in zip(texts, metadatas)]
        return cls(vectorizer=vectorizer, docs=docs, tfidf_array=tfidf_array, **kwargs)
[docs]    @classmethod
    def from_documents(
        cls,
        documents: Iterable[Document],
        *,
        tfidf_params: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> TFIDFRetriever:
        texts, metadatas = zip(*((d.page_content, d.metadata) for d in documents))
        return cls.from_texts(
            texts=texts, tfidf_params=tfidf_params, metadatas=metadatas, **kwargs
        )
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        from sklearn.metrics.pairwise import cosine_similarity
        query_vec = self.vectorizer.transform(
            [query]
        )  # Ip -- (n_docs,x), Op -- (n_docs,n_Feats)
        results = cosine_similarity(self.tfidf_array, query_vec).reshape(
            (-1,)
        )  # Op -- (n_docs,1) -- Cosine Sim with each doc
        return_docs = [self.docs[i] for i in results.argsort()[-self.k :][::-1]]
        return return_docs
[docs]    def save_local(
        self,
        folder_path: str,
        file_name: str = "tfidf_vectorizer",
    ) -> None:
        try:
            import joblib


) -> None:
        try:
            import joblib
        except ImportError:
            raise ImportError(
                "Could not import joblib, please install with `pip install joblib`."
            )
        path = Path(folder_path)
        path.mkdir(exist_ok=True, parents=True)
        # Save vectorizer with joblib dump.
        joblib.dump(self.vectorizer, path / f"{file_name}.joblib")
        # Save docs and tfidf array as pickle.
        with open(path / f"{file_name}.pkl", "wb") as f:
            pickle.dump((self.docs, self.tfidf_array), f)
[docs]    @classmethod
    def load_local(
        cls,
        folder_path: str,
        file_name: str = "tfidf_vectorizer",
    ) -> TFIDFRetriever:
        try:
            import joblib
        except ImportError:
            raise ImportError(
                "Could not import joblib, please install with `pip install joblib`."
            )
        path = Path(folder_path)
        # Load vectorizer with joblib load.
        vectorizer = joblib.load(path / f"{file_name}.joblib")
        # Load docs and tfidf array as pickle.
        with open(path / f"{file_name}.pkl", "rb") as f:
            docs, tfidf_array = pickle.load(f)
        return cls(vectorizer=vectorizer, docs=docs, tfidf_array=tfidf_array)


Source code for langchain.retrievers.multi_query
import asyncio
import logging
from typing import List, Sequence
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.chains.llm import LLMChain
from langchain.llms.base import BaseLLM
from langchain.output_parsers.pydantic import PydanticOutputParser
from langchain.prompts.prompt import PromptTemplate
from langchain.pydantic_v1 import BaseModel, Field
from langchain.schema import BaseRetriever, Document
logger = logging.getLogger(__name__)
[docs]class LineList(BaseModel):
    """List of lines."""
    lines: List[str] = Field(description="Lines of text")
    """List of lines."""
[docs]class LineListOutputParser(PydanticOutputParser):
    """Output parser for a list of lines."""
    def __init__(self) -> None:
        super().__init__(pydantic_object=LineList)
[docs]    def parse(self, text: str) -> LineList:
        lines = text.strip().split("\n")
        return LineList(lines=lines)
# Default prompt
DEFAULT_QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template="""You are an AI language model assistant. Your task is
    to generate 3 different versions of the given user
    question to retrieve relevant documents from a vector  database.
    By generating multiple perspectives on the user question,
    your goal is to help the user overcome some of the limitations
    of distance-based similarity search. Provide these alternative
    questions separated by newlines. Original question: {question}""",
)
def _unique_documents(documents: Sequence[Document]) -> List[Document]:


)
def _unique_documents(documents: Sequence[Document]) -> List[Document]:
    return [doc for i, doc in enumerate(documents) if doc not in documents[:i]]
[docs]class MultiQueryRetriever(BaseRetriever):
    """Given a query, use an LLM to write a set of queries.
    Retrieve docs for each query. Return the unique union of all retrieved docs.
    """
    retriever: BaseRetriever
    llm_chain: LLMChain
    verbose: bool = True
    parser_key: str = "lines"
    include_original: bool = False
    """Whether to include the original query in the list of generated queries."""
[docs]    @classmethod
    def from_llm(
        cls,
        retriever: BaseRetriever,
        llm: BaseLLM,
        prompt: PromptTemplate = DEFAULT_QUERY_PROMPT,
        parser_key: str = "lines",
        include_original: bool = False,
    ) -> "MultiQueryRetriever":
        """Initialize from llm using default template.
        Args:
            retriever: retriever to query documents from
            llm: llm for query generation using DEFAULT_QUERY_PROMPT
            include_original: Whether to include the original query in the list of
                generated queries.
        Returns:
            MultiQueryRetriever
        """
        output_parser = LineListOutputParser()
        llm_chain = LLMChain(llm=llm, prompt=prompt, output_parser=output_parser)
        return cls(
            retriever=retriever,
            llm_chain=llm_chain,
            parser_key=parser_key,
            include_original=include_original,
        )
    async def _aget_relevant_documents(
        self,


)
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """Get relevant documents given a user query.
        Args:
            question: user query
        Returns:
            Unique union of relevant documents from all generated queries
        """
        queries = await self.agenerate_queries(query, run_manager)
        if self.include_original:
            queries.append(query)
        documents = await self.aretrieve_documents(queries, run_manager)
        return self.unique_union(documents)
[docs]    async def agenerate_queries(
        self, question: str, run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[str]:
        """Generate queries based upon user input.
        Args:
            question: user query
        Returns:
            List of LLM generated queries that are similar to the user input
        """
        response = await self.llm_chain.acall(
            inputs={"question": question}, callbacks=run_manager.get_child()
        )
        lines = getattr(response["text"], self.parser_key, [])
        if self.verbose:
            logger.info(f"Generated queries: {lines}")
        return lines
[docs]    async def aretrieve_documents(
        self, queries: List[str], run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Run all LLM generated queries.
        Args:
            queries: query list
        Returns:
            List of retrieved Documents
        """
        document_lists = await asyncio.gather(
            *(
                self.retriever.aget_relevant_documents(
                    query, callbacks=run_manager.get_child()
                )
                for query in queries
            )


)
                for query in queries
            )
        )
        return [doc for docs in document_lists for doc in docs]
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """Get relevant documents given a user query.
        Args:
            question: user query
        Returns:
            Unique union of relevant documents from all generated queries
        """
        queries = self.generate_queries(query, run_manager)
        if self.include_original:
            queries.append(query)
        documents = self.retrieve_documents(queries, run_manager)
        return self.unique_union(documents)
[docs]    def generate_queries(
        self, question: str, run_manager: CallbackManagerForRetrieverRun
    ) -> List[str]:
        """Generate queries based upon user input.
        Args:
            question: user query
        Returns:
            List of LLM generated queries that are similar to the user input
        """
        response = self.llm_chain(
            {"question": question}, callbacks=run_manager.get_child()
        )
        lines = getattr(response["text"], self.parser_key, [])
        if self.verbose:
            logger.info(f"Generated queries: {lines}")
        return lines
[docs]    def retrieve_documents(
        self, queries: List[str], run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Run all LLM generated queries.
        Args:
            queries: query list
        Returns:
            List of retrieved Documents
        """
        documents = []
        for query in queries:
            docs = self.retriever.get_relevant_documents(


for query in queries:
            docs = self.retriever.get_relevant_documents(
                query, callbacks=run_manager.get_child()
            )
            documents.extend(docs)
        return documents
[docs]    def unique_union(self, documents: List[Document]) -> List[Document]:
        """Get unique Documents.
        Args:
            documents: List of retrieved Documents
        Returns:
            List of unique retrieved Documents
        """
        return _unique_documents(documents)


Source code for langchain.retrievers.google_cloud_documentai_warehouse
"""Retriever wrapper for Google Cloud Document AI Warehouse."""
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.docstore.document import Document
from langchain.pydantic_v1 import root_validator
from langchain.schema import BaseRetriever
from langchain.utilities.vertexai import get_client_info
from langchain.utils import get_from_dict_or_env
if TYPE_CHECKING:
    from google.cloud.contentwarehouse_v1 import (
        DocumentServiceClient,
        RequestMetadata,
        SearchDocumentsRequest,
    )
    from google.cloud.contentwarehouse_v1.services.document_service.pagers import (
        SearchDocumentsPager,
    )
[docs]class GoogleDocumentAIWarehouseRetriever(BaseRetriever):
    """A retriever based on Document AI Warehouse.
    Documents should be created and documents should be uploaded
        in a separate flow, and this retriever uses only Document AI
        schema_id provided to search for revelant documents.
    More info: https://cloud.google.com/document-ai-warehouse.
    """
    location: str = "us"
    """Google Cloud location where Document AI Warehouse is placed."""
    project_number: str
    """Google Cloud project number, should contain digits only."""
    schema_id: Optional[str] = None
    """Document AI Warehouse schema to query against.
    If nothing is provided, all documents in the project will be searched."""
    qa_size_limit: int = 5
    """The limit on the number of documents returned."""
    client: "DocumentServiceClient" = None  #: :meta private:
    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:


@root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validates the environment."""
        try:  # noqa: F401
            from google.cloud.contentwarehouse_v1 import DocumentServiceClient
        except ImportError as exc:
            raise ImportError(
                "google.cloud.contentwarehouse is not installed."
                "Please install it with pip install google-cloud-contentwarehouse"
            ) from exc
        values["project_number"] = get_from_dict_or_env(
            values, "project_number", "PROJECT_NUMBER"
        )
        values["client"] = DocumentServiceClient(
            client_info=get_client_info(module="document-ai-warehouse")
        )
        return values
    def _prepare_request_metadata(self, user_ldap: str) -> "RequestMetadata":
        from google.cloud.contentwarehouse_v1 import RequestMetadata, UserInfo
        user_info = UserInfo(id=f"user:{user_ldap}")
        return RequestMetadata(user_info=user_info)
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun, **kwargs: Any
    ) -> List[Document]:
        request = self._prepare_search_request(query, **kwargs)
        response = self.client.search_documents(request=request)
        return self._parse_search_response(response=response)
    def _prepare_search_request(
        self, query: str, **kwargs: Any
    ) -> "SearchDocumentsRequest":
        from google.cloud.contentwarehouse_v1 import (
            DocumentQuery,
            SearchDocumentsRequest,
        )
        try:
            user_ldap = kwargs["user_ldap"]
        except KeyError:
            raise ValueError("Argument user_ldap should be provided!")
        request_metadata = self._prepare_request_metadata(user_ldap=user_ldap)
        schemas = []


schemas = []
        if self.schema_id:
            schemas.append(
                self.client.document_schema_path(
                    project=self.project_number,
                    location=self.location,
                    document_schema=self.schema_id,
                )
            )
        return SearchDocumentsRequest(
            parent=self.client.common_location_path(self.project_number, self.location),
            request_metadata=request_metadata,
            document_query=DocumentQuery(
                query=query, is_nl_query=True, document_schema_names=schemas
            ),
            qa_size_limit=self.qa_size_limit,
        )
    def _parse_search_response(
        self, response: "SearchDocumentsPager"
    ) -> List[Document]:
        documents = []
        for doc in response.matching_documents:
            metadata = {
                "title": doc.document.title,
                "source": doc.document.raw_document_path,
            }
            documents.append(
                Document(page_content=doc.search_text_snippet, metadata=metadata)
            )
        return documents


Source code for langchain.retrievers.arcee
from typing import Any, Dict, List, Optional
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.docstore.document import Document
from langchain.pydantic_v1 import Extra, root_validator
from langchain.schema import BaseRetriever
from langchain.utilities.arcee import ArceeWrapper, DALMFilter
from langchain.utils import get_from_dict_or_env
[docs]class ArceeRetriever(BaseRetriever):
    """Document retriever for Arcee's Domain Adapted Language Models (DALMs).
    To use, set the ``ARCEE_API_KEY`` environment variable with your Arcee API key,
    or pass ``arcee_api_key`` as a named parameter.
    Example:
        .. code-block:: python
            from langchain.retrievers import ArceeRetriever
            retriever = ArceeRetriever(
                model="DALM-PubMed",
                arcee_api_key="ARCEE-API-KEY"
            )
            documents = retriever.get_relevant_documents("AI-driven music therapy")
    """
    _client: Optional[ArceeWrapper] = None  #: :meta private:
    """Arcee client."""
    arcee_api_key: str = ""
    """Arcee API Key"""
    model: str
    """Arcee DALM name"""
    arcee_api_url: str = "https://api.arcee.ai"
    """Arcee API URL"""
    arcee_api_version: str = "v2"
    """Arcee API Version"""
    arcee_app_url: str = "https://app.arcee.ai"
    """Arcee App URL"""
    model_kwargs: Optional[Dict[str, Any]] = None
    """Keyword arguments to pass to the model."""


"""Keyword arguments to pass to the model."""
    class Config:
        """Configuration for this pydantic object."""
        extra = Extra.forbid
        underscore_attrs_are_private = True
    def __init__(self, **data: Any) -> None:
        """Initializes private fields."""
        super().__init__(**data)
        self._client = ArceeWrapper(
            arcee_api_key=self.arcee_api_key,
            arcee_api_url=self.arcee_api_url,
            arcee_api_version=self.arcee_api_version,
            model_kwargs=self.model_kwargs,
            model_name=self.model,
        )
        self._client.validate_model_training_status()
    @root_validator()
    def validate_environments(cls, values: Dict) -> Dict:
        """Validate Arcee environment variables."""
        # validate env vars
        values["arcee_api_key"] = get_from_dict_or_env(
            values,
            "arcee_api_key",
            "ARCEE_API_KEY",
        )
        values["arcee_api_url"] = get_from_dict_or_env(
            values,
            "arcee_api_url",
            "ARCEE_API_URL",
        )
        values["arcee_app_url"] = get_from_dict_or_env(
            values,
            "arcee_app_url",
            "ARCEE_APP_URL",
        )
        values["arcee_api_version"] = get_from_dict_or_env(
            values,
            "arcee_api_version",
            "ARCEE_API_VERSION",
        )
        # validate model kwargs
        if values["model_kwargs"]:
            kw = values["model_kwargs"]
            # validate size
            if kw.get("size") is not None:
                if not kw.get("size") >= 0:


if not kw.get("size") >= 0:
                    raise ValueError("`size` must not be negative.")
            # validate filters
            if kw.get("filters") is not None:
                if not isinstance(kw.get("filters"), List):
                    raise ValueError("`filters` must be a list.")
                for f in kw.get("filters"):
                    DALMFilter(**f)
        return values
    def _get_relevant_documents(
        self, query: str, run_manager: CallbackManagerForRetrieverRun, **kwargs: Any
    ) -> List[Document]:
        """Retrieve {size} contexts with your retriever for a given query
        Args:
            query: Query to submit to the model
            size: The max number of context results to retrieve.
            Defaults to 3. (Can be less if filters are provided).
            filters: Filters to apply to the context dataset.
        """
        try:
            if not self._client:
                raise ValueError("Client is not initialized.")
            return self._client.retrieve(query=query, **kwargs)
        except Exception as e:
            raise ValueError(f"Error while retrieving documents: {e}") from e


Source code for langchain.retrievers.databerry
from typing import List, Optional
import aiohttp
import requests
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
[docs]class DataberryRetriever(BaseRetriever):
    """`Databerry API` retriever."""
    datastore_url: str
    top_k: Optional[int]
    api_key: Optional[str]
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        response = requests.post(
            self.datastore_url,
            json={
                "query": query,
                **({"topK": self.top_k} if self.top_k is not None else {}),
            },
            headers={
                "Content-Type": "application/json",
                **(
                    {"Authorization": f"Bearer {self.api_key}"}
                    if self.api_key is not None
                    else {}
                ),
            },
        )
        data = response.json()
        return [
            Document(
                page_content=r["text"],
                metadata={"source": r["source"], "score": r["score"]},
            )
            for r in data["results"]
        ]
    async def _aget_relevant_documents(
        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[Document]:
        async with aiohttp.ClientSession() as session:
            async with session.request(
                "POST",
                self.datastore_url,
                json={
                    "query": query,


self.datastore_url,
                json={
                    "query": query,
                    **({"topK": self.top_k} if self.top_k is not None else {}),
                },
                headers={
                    "Content-Type": "application/json",
                    **(
                        {"Authorization": f"Bearer {self.api_key}"}
                        if self.api_key is not None
                        else {}
                    ),
                },
            ) as response:
                data = await response.json()
        return [
            Document(
                page_content=r["text"],
                metadata={"source": r["source"], "score": r["score"]},
            )
            for r in data["results"]
        ]


Source code for langchain.retrievers.multi_vector
from typing import List
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.pydantic_v1 import Field
from langchain.schema import BaseRetriever, BaseStore, Document
from langchain.schema.vectorstore import VectorStore
[docs]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"""
    docstore: BaseStore[str, Document]
    """The storage layer for the parent documents"""
    id_key: str = "doc_id"
    search_kwargs: dict = Field(default_factory=dict)
    """Keyword arguments to pass to the search function."""
    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
        """
        sub_docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
        # We do this to maintain the order of the ids that are returned
        ids = []
        for d in sub_docs:
            if 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]


Source code for langchain.retrievers.llama_index
from typing import Any, Dict, List, cast
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.pydantic_v1 import Field
from langchain.schema import BaseRetriever, Document
[docs]class LlamaIndexRetriever(BaseRetriever):
    """`LlamaIndex` retriever.
    It is used for the question-answering with sources over
    an LlamaIndex data structure."""
    index: Any
    """LlamaIndex index to query."""
    query_kwargs: Dict = Field(default_factory=dict)
    """Keyword arguments to pass to the query method."""
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Get documents relevant for a query."""
        try:
            from llama_index.indices.base import BaseGPTIndex
            from llama_index.response.schema import Response
        except ImportError:
            raise ImportError(
                "You need to install `pip install llama-index` to use this retriever."
            )
        index = cast(BaseGPTIndex, self.index)
        response = index.query(query, response_mode="no_text", **self.query_kwargs)
        response = cast(Response, response)
        # parse source nodes
        docs = []
        for source_node in response.source_nodes:
            metadata = source_node.extra_info or {}
            docs.append(
                Document(page_content=source_node.source_text, metadata=metadata)
            )
        return docs
[docs]class LlamaIndexGraphRetriever(BaseRetriever):
    """`LlamaIndex` graph data structure retriever.
    It is used for question-answering with sources over an LlamaIndex


It is used for question-answering with sources over an LlamaIndex
    graph data structure."""
    graph: Any
    """LlamaIndex graph to query."""
    query_configs: List[Dict] = Field(default_factory=list)
    """List of query configs to pass to the query method."""
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Get documents relevant for a query."""
        try:
            from llama_index.composability.graph import (
                QUERY_CONFIG_TYPE,
                ComposableGraph,
            )
            from llama_index.response.schema import Response
        except ImportError:
            raise ImportError(
                "You need to install `pip install llama-index` to use this retriever."
            )
        graph = cast(ComposableGraph, self.graph)
        # for now, inject response_mode="no_text" into query configs
        for query_config in self.query_configs:
            query_config["response_mode"] = "no_text"
        query_configs = cast(List[QUERY_CONFIG_TYPE], self.query_configs)
        response = graph.query(query, query_configs=query_configs)
        response = cast(Response, response)
        # parse source nodes
        docs = []
        for source_node in response.source_nodes:
            metadata = source_node.extra_info or {}
            docs.append(
                Document(page_content=source_node.source_text, metadata=metadata)
            )
        return docs


Source code for langchain.retrievers.arxiv
from typing import List
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
from langchain.utilities.arxiv import ArxivAPIWrapper
[docs]class ArxivRetriever(BaseRetriever, ArxivAPIWrapper):
    """`Arxiv` retriever.
    It wraps load() to get_relevant_documents().
    It uses all ArxivAPIWrapper arguments without any change.
    """
    get_full_documents: bool = False
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        if self.get_full_documents:
            return self.load(query=query)
        else:
            return self.get_summaries_as_docs(query)


Source code for langchain.retrievers.docarray
from enum import Enum
from typing import Any, Dict, List, Optional, Union
import numpy as np
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
from langchain.schema.embeddings import Embeddings
from langchain.vectorstores.utils import maximal_marginal_relevance
[docs]class SearchType(str, Enum):
    """Enumerator of the types of search to perform."""
    similarity = "similarity"
    mmr = "mmr"
[docs]class DocArrayRetriever(BaseRetriever):
    """`DocArray Document Indices` retriever.
    Currently, it supports 5 backends:
    InMemoryExactNNIndex, HnswDocumentIndex, QdrantDocumentIndex,
    ElasticDocIndex, and WeaviateDocumentIndex.
    Args:
        index: One of the above-mentioned index instances
        embeddings: Embedding model to represent text as vectors
        search_field: Field to consider for searching in the documents.
            Should be an embedding/vector/tensor.
        content_field: Field that represents the main content in your document schema.
            Will be used as a `page_content`. Everything else will go into `metadata`.
        search_type: Type of search to perform (similarity / mmr)
        filters: Filters applied for document retrieval.
        top_k: Number of documents to return
    """
    index: Any
    embeddings: Embeddings
    search_field: str
    content_field: str
    search_type: SearchType = SearchType.similarity
    top_k: int = 1
    filters: Optional[Any] = None
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True


"""Configuration for this pydantic object."""
        arbitrary_types_allowed = True
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """Get documents relevant for a query.
        Args:
            query: string to find relevant documents for
        Returns:
            List of relevant documents
        """
        query_emb = np.array(self.embeddings.embed_query(query))
        if self.search_type == SearchType.similarity:
            results = self._similarity_search(query_emb)
        elif self.search_type == SearchType.mmr:
            results = self._mmr_search(query_emb)
        else:
            raise ValueError(
                f"Search type {self.search_type} does not exist. "
                f"Choose either 'similarity' or 'mmr'."
            )
        return results
    def _search(
        self, query_emb: np.ndarray, top_k: int
    ) -> List[Union[Dict[str, Any], Any]]:
        """
        Perform a search using the query embedding and return top_k documents.
        Args:
            query_emb: Query represented as an embedding
            top_k: Number of documents to return
        Returns:
            A list of top_k documents matching the query
        """
        from docarray.index import ElasticDocIndex, WeaviateDocumentIndex
        filter_args = {}
        search_field = self.search_field
        if isinstance(self.index, WeaviateDocumentIndex):
            filter_args["where_filter"] = self.filters
            search_field = ""
        elif isinstance(self.index, ElasticDocIndex):
            filter_args["query"] = self.filters
        else:
            filter_args["filter_query"] = self.filters


else:
            filter_args["filter_query"] = self.filters
        if self.filters:
            query = (
                self.index.build_query()  # get empty query object
                .find(
                    query=query_emb, search_field=search_field
                )  # add vector similarity search
                .filter(**filter_args)  # add filter search
                .build(limit=top_k)  # build the query
            )
            # execute the combined query and return the results
            docs = self.index.execute_query(query)
            if hasattr(docs, "documents"):
                docs = docs.documents
            docs = docs[:top_k]
        else:
            docs = self.index.find(
                query=query_emb, search_field=search_field, limit=top_k
            ).documents
        return docs
    def _similarity_search(self, query_emb: np.ndarray) -> List[Document]:
        """
        Perform a similarity search.
        Args:
            query_emb: Query represented as an embedding
        Returns:
            A list of documents most similar to the query
        """
        docs = self._search(query_emb=query_emb, top_k=self.top_k)
        results = [self._docarray_to_langchain_doc(doc) for doc in docs]
        return results
    def _mmr_search(self, query_emb: np.ndarray) -> List[Document]:
        """
        Perform a maximal marginal relevance (mmr) search.
        Args:
            query_emb: Query represented as an embedding
        Returns:
            A list of diverse documents related to the query
        """
        docs = self._search(query_emb=query_emb, top_k=20)
        mmr_selected = maximal_marginal_relevance(
            query_emb,
            [
                doc[self.search_field]
                if isinstance(doc, dict)


[
                doc[self.search_field]
                if isinstance(doc, dict)
                else getattr(doc, self.search_field)
                for doc in docs
            ],
            k=self.top_k,
        )
        results = [self._docarray_to_langchain_doc(docs[idx]) for idx in mmr_selected]
        return results
    def _docarray_to_langchain_doc(self, doc: Union[Dict[str, Any], Any]) -> Document:
        """
        Convert a DocArray document (which also might be a dict)
        to a langchain document format.
        DocArray document can contain arbitrary fields, so the mapping is done
        in the following way:
        page_content <-> content_field
        metadata <-> all other fields excluding
            tensors and embeddings (so float, int, string)
        Args:
            doc: DocArray document
        Returns:
            Document in langchain format
        Raises:
            ValueError: If the document doesn't contain the content field
        """
        fields = doc.keys() if isinstance(doc, dict) else doc.__fields__
        if self.content_field not in fields:
            raise ValueError(
                f"Document does not contain the content field - {self.content_field}."
            )
        lc_doc = Document(
            page_content=doc[self.content_field]
            if isinstance(doc, dict)
            else getattr(doc, self.content_field)
        )
        for name in fields:
            value = doc[name] if isinstance(doc, dict) else getattr(doc, name)
            if (
                isinstance(value, (str, int, float, bool))
                and name != self.content_field
            ):
                lc_doc.metadata[name] = value
        return lc_doc


Source code for langchain.retrievers.milvus
"""Milvus Retriever"""
import warnings
from typing import Any, Dict, List, Optional
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.pydantic_v1 import root_validator
from langchain.schema import BaseRetriever, Document
from langchain.schema.embeddings import Embeddings
from langchain.vectorstores.milvus import Milvus
# TODO: Update to MilvusClient + Hybrid Search when available
[docs]class MilvusRetriever(BaseRetriever):
    """`Milvus API` retriever."""
    embedding_function: Embeddings
    collection_name: str = "LangChainCollection"
    connection_args: Optional[Dict[str, Any]] = None
    consistency_level: str = "Session"
    search_params: Optional[dict] = None
    store: Milvus
    retriever: BaseRetriever
    @root_validator(pre=True)
    def create_retriever(cls, values: Dict) -> Dict:
        """Create the Milvus store and retriever."""
        values["store"] = Milvus(
            values["embedding_function"],
            values["collection_name"],
            values["connection_args"],
            values["consistency_level"],
        )
        values["retriever"] = values["store"].as_retriever(
            search_kwargs={"param": values["search_params"]}
        )
        return values
[docs]    def add_texts(
        self, texts: List[str], metadatas: Optional[List[dict]] = None
    ) -> None:
        """Add text to the Milvus store
        Args:
            texts (List[str]): The text


Args:
            texts (List[str]): The text
            metadatas (List[dict]): Metadata dicts, must line up with existing store
        """
        self.store.add_texts(texts, metadatas)
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        return self.retriever.get_relevant_documents(
            query, run_manager=run_manager.get_child(), **kwargs
        )
[docs]def MilvusRetreiver(*args: Any, **kwargs: Any) -> MilvusRetriever:
    """Deprecated MilvusRetreiver. Please use MilvusRetriever ('i' before 'e') instead.
    Args:
        *args:
        **kwargs:
    Returns:
        MilvusRetriever
    """
    warnings.warn(
        "MilvusRetreiver will be deprecated in the future. "
        "Please use MilvusRetriever ('i' before 'e') instead.",
        DeprecationWarning,
    )
    return MilvusRetriever(*args, **kwargs)


Source code for langchain.retrievers.wikipedia
from typing import List
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
from langchain.utilities.wikipedia import WikipediaAPIWrapper
[docs]class WikipediaRetriever(BaseRetriever, WikipediaAPIWrapper):
    """`Wikipedia API` retriever.
    It wraps load() to get_relevant_documents().
    It uses all WikipediaAPIWrapper arguments without any change.
    """
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        return self.load(query=query)


Source code for langchain.retrievers.parent_document_retriever
import uuid
from typing import List, Optional
from langchain.retrievers import MultiVectorRetriever
from langchain.schema.document import Document
from langchain.text_splitter import TextSplitter
[docs]class ParentDocumentRetriever(MultiVectorRetriever):
    """Retrieve small chunks then retrieve their parent documents.
    When splitting documents for retrieval, there are often conflicting desires:
    1. You may want to have small documents, so that their embeddings can most
        accurately reflect their meaning. If too long, then the embeddings can
        lose meaning.
    2. You want to have long enough documents that the context of each chunk is
        retained.
    The ParentDocumentRetriever strikes that balance by splitting and storing
    small chunks of data. During retrieval, it first fetches the small chunks
    but then looks up the parent ids for those chunks and returns those larger
    documents.
    Note that "parent document" refers to the document that a small chunk
    originated from. This can either be the whole raw document OR a larger
    chunk.
    Examples:
        .. code-block:: python
            # Imports
            from langchain.vectorstores import Chroma
            from langchain.embeddings import OpenAIEmbeddings
            from langchain.text_splitter import RecursiveCharacterTextSplitter
            from langchain.storage import InMemoryStore
            # This text splitter is used to create the parent documents
            parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000)
            # This text splitter is used to create the child documents
            # It should create documents smaller than the parent
            child_splitter = RecursiveCharacterTextSplitter(chunk_size=400)
            # The vectorstore to use to index the child chunks


# The vectorstore to use to index the child chunks
            vectorstore = Chroma(embedding_function=OpenAIEmbeddings())
            # The storage layer for the parent documents
            store = InMemoryStore()
            # Initialize the retriever
            retriever = ParentDocumentRetriever(
                vectorstore=vectorstore,
                docstore=store,
                child_splitter=child_splitter,
                parent_splitter=parent_splitter,
            )
    """
    child_splitter: TextSplitter
    """The text splitter to use to create child documents."""
    """The key to use to track the parent id. This will be stored in the
    metadata of child documents."""
    parent_splitter: Optional[TextSplitter] = None
    """The text splitter to use to create parent documents.
    If none, then the parent documents will be the raw documents passed in."""
[docs]    def add_documents(
        self,
        documents: List[Document],
        ids: Optional[List[str]] = None,
        add_to_docstore: bool = True,
    ) -> None:
        """Adds documents to the docstore and vectorstores.
        Args:
            documents: List of documents to add
            ids: Optional list of ids for documents. If provided should be the same
                length as the list of documents. Can provided if parent documents
                are already in the document store and you don't want to re-add
                to the docstore. If not provided, random UUIDs will be used as
                ids.
            add_to_docstore: Boolean of whether to add documents to docstore.
                This can be false if and only if `ids` are provided. You may want
                to set this to False if the documents are already in the docstore


to set this to False if the documents are already in the docstore
                and you don't want to re-add them.
        """
        if self.parent_splitter is not None:
            documents = self.parent_splitter.split_documents(documents)
        if ids is None:
            doc_ids = [str(uuid.uuid4()) for _ in documents]
            if not add_to_docstore:
                raise ValueError(
                    "If ids are not passed in, `add_to_docstore` MUST be True"
                )
        else:
            if len(documents) != len(ids):
                raise ValueError(
                    "Got uneven list of documents and ids. "
                    "If `ids` is provided, should be same length as `documents`."
                )
            doc_ids = ids
        docs = []
        full_docs = []
        for i, doc in enumerate(documents):
            _id = doc_ids[i]
            sub_docs = self.child_splitter.split_documents([doc])
            for _doc in sub_docs:
                _doc.metadata[self.id_key] = _id
            docs.extend(sub_docs)
            full_docs.append((_id, doc))
        self.vectorstore.add_documents(docs)
        if add_to_docstore:
            self.docstore.mset(full_docs)


Source code for langchain.retrievers.time_weighted_retriever
import datetime
from copy import deepcopy
from typing import Any, Dict, List, Optional, Tuple
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.pydantic_v1 import Field
from langchain.schema import BaseRetriever, Document
from langchain.schema.vectorstore import VectorStore
def _get_hours_passed(time: datetime.datetime, ref_time: datetime.datetime) -> float:
    """Get the hours passed between two datetimes."""
    return (time - ref_time).total_seconds() / 3600
[docs]class TimeWeightedVectorStoreRetriever(BaseRetriever):
    """Retriever that combines embedding similarity with
    recency in retrieving values."""
    vectorstore: VectorStore
    """The vectorstore to store documents and determine salience."""
    search_kwargs: dict = Field(default_factory=lambda: dict(k=100))
    """Keyword arguments to pass to the vectorstore similarity search."""
    # TODO: abstract as a queue
    memory_stream: List[Document] = Field(default_factory=list)
    """The memory_stream of documents to search through."""
    decay_rate: float = Field(default=0.01)
    """The exponential decay factor used as (1.0-decay_rate)**(hrs_passed)."""
    k: int = 4
    """The maximum number of documents to retrieve in a given call."""
    other_score_keys: List[str] = []
    """Other keys in the metadata to factor into the score, e.g. 'importance'."""
    default_salience: Optional[float] = None
    """The salience to assign memories not retrieved from the vector store.
    None assigns no salience to documents not fetched from the vector store.
    """
    class Config:


"""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
    def _document_get_date(self, field: str, document: Document) -> datetime.datetime:
        """Return the value of the date field of a document."""
        if field in document.metadata:
            if isinstance(document.metadata[field], float):
                return datetime.datetime.fromtimestamp(document.metadata[field])
            return document.metadata[field]
        return datetime.datetime.now()
    def _get_combined_score(
        self,
        document: Document,
        vector_relevance: Optional[float],
        current_time: datetime.datetime,
    ) -> float:
        """Return the combined score for a document."""
        hours_passed = _get_hours_passed(
            current_time,
            self._document_get_date("last_accessed_at", document),
        )
        score = (1.0 - self.decay_rate) ** hours_passed
        for key in self.other_score_keys:
            if key in document.metadata:
                score += document.metadata[key]
        if vector_relevance is not None:
            score += vector_relevance
        return score
[docs]    def get_salient_docs(self, query: str) -> Dict[int, Tuple[Document, float]]:
        """Return documents that are salient to the query."""
        docs_and_scores: List[Tuple[Document, float]]
        docs_and_scores = self.vectorstore.similarity_search_with_relevance_scores(
            query, **self.search_kwargs
        )
        results = {}
        for fetched_doc, relevance in docs_and_scores:
            if "buffer_idx" in fetched_doc.metadata:
                buffer_idx = fetched_doc.metadata["buffer_idx"]
                doc = self.memory_stream[buffer_idx]
                results[buffer_idx] = (doc, relevance)


results[buffer_idx] = (doc, relevance)
        return results
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Return documents that are relevant to the query."""
        current_time = datetime.datetime.now()
        docs_and_scores = {
            doc.metadata["buffer_idx"]: (doc, self.default_salience)
            for doc in self.memory_stream[-self.k :]
        }
        # If a doc is considered salient, update the salience score
        docs_and_scores.update(self.get_salient_docs(query))
        rescored_docs = [
            (doc, self._get_combined_score(doc, relevance, current_time))
            for doc, relevance in docs_and_scores.values()
        ]
        rescored_docs.sort(key=lambda x: x[1], reverse=True)
        result = []
        # Ensure frequently accessed memories aren't forgotten
        for doc, _ in rescored_docs[: self.k]:
            # TODO: Update vector store doc once `update` method is exposed.
            buffered_doc = self.memory_stream[doc.metadata["buffer_idx"]]
            buffered_doc.metadata["last_accessed_at"] = current_time
            result.append(buffered_doc)
        return result
[docs]    def add_documents(self, documents: List[Document], **kwargs: Any) -> List[str]:
        """Add documents to vectorstore."""
        current_time = kwargs.get("current_time")
        if current_time is None:
            current_time = datetime.datetime.now()
        # Avoid mutating input documents
        dup_docs = [deepcopy(d) for d in documents]
        for i, doc in enumerate(dup_docs):
            if "last_accessed_at" not in doc.metadata:


if "last_accessed_at" not in doc.metadata:
                doc.metadata["last_accessed_at"] = current_time
            if "created_at" not in doc.metadata:
                doc.metadata["created_at"] = current_time
            doc.metadata["buffer_idx"] = len(self.memory_stream) + i
        self.memory_stream.extend(dup_docs)
        return self.vectorstore.add_documents(dup_docs, **kwargs)
[docs]    async def aadd_documents(
        self, documents: List[Document], **kwargs: Any
    ) -> List[str]:
        """Add documents to vectorstore."""
        current_time = kwargs.get("current_time")
        if current_time is None:
            current_time = datetime.datetime.now()
        # Avoid mutating input documents
        dup_docs = [deepcopy(d) for d in documents]
        for i, doc in enumerate(dup_docs):
            if "last_accessed_at" not in doc.metadata:
                doc.metadata["last_accessed_at"] = current_time
            if "created_at" not in doc.metadata:
                doc.metadata["created_at"] = current_time
            doc.metadata["buffer_idx"] = len(self.memory_stream) + i
        self.memory_stream.extend(dup_docs)
        return await self.vectorstore.aadd_documents(dup_docs, **kwargs)


Source code for langchain.retrievers.chatgpt_plugin_retriever
from __future__ import annotations
from typing import List, Optional
import aiohttp
import requests
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
[docs]class ChatGPTPluginRetriever(BaseRetriever):
    """`ChatGPT plugin` retriever."""
    url: str
    """URL of the ChatGPT plugin."""
    bearer_token: str
    """Bearer token for the ChatGPT plugin."""
    top_k: int = 3
    """Number of documents to return."""
    filter: Optional[dict] = None
    """Filter to apply to the results."""
    aiosession: Optional[aiohttp.ClientSession] = None
    """Aiohttp session to use for requests."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
        """Allow arbitrary types."""
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        url, json, headers = self._create_request(query)
        response = requests.post(url, json=json, headers=headers)
        results = response.json()["results"][0]["results"]
        docs = []
        for d in results:
            content = d.pop("text")
            metadata = d.pop("metadata", d)
            if metadata.get("source_id"):
                metadata["source"] = metadata.pop("source_id")
            docs.append(Document(page_content=content, metadata=metadata))
        return docs
    async def _aget_relevant_documents(


return docs
    async def _aget_relevant_documents(
        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[Document]:
        url, json, headers = self._create_request(query)
        if not self.aiosession:
            async with aiohttp.ClientSession() as session:
                async with session.post(url, headers=headers, json=json) as response:
                    res = await response.json()
        else:
            async with self.aiosession.post(
                url, headers=headers, json=json
            ) as response:
                res = await response.json()
        results = res["results"][0]["results"]
        docs = []
        for d in results:
            content = d.pop("text")
            metadata = d.pop("metadata", d)
            if metadata.get("source_id"):
                metadata["source"] = metadata.pop("source_id")
            docs.append(Document(page_content=content, metadata=metadata))
        return docs
    def _create_request(self, query: str) -> tuple[str, dict, dict]:
        url = f"{self.url}/query"
        json = {
            "queries": [
                {
                    "query": query,
                    "filter": self.filter,
                    "top_k": self.top_k,
                }
            ]
        }
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.bearer_token}",
        }
        return url, json, headers


Source code for langchain.retrievers.merger_retriever
import asyncio
from typing import List
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
[docs]class MergerRetriever(BaseRetriever):
    """Retriever that merges the results of multiple retrievers."""
    retrievers: List[BaseRetriever]
    """A list of retrievers to merge."""
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """
        Get the relevant documents for a given query.
        Args:
            query: The query to search for.
        Returns:
            A list of relevant documents.
        """
        # Merge the results of the retrievers.
        merged_documents = self.merge_documents(query, run_manager)
        return merged_documents
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """
        Asynchronously get the relevant documents for a given query.
        Args:
            query: The query to search for.
        Returns:
            A list of relevant documents.
        """
        # Merge the results of the retrievers.
        merged_documents = await self.amerge_documents(query, run_manager)
        return merged_documents
[docs]    def merge_documents(
        self, query: str, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """
        Merge the results of the retrievers.


"""
        Merge the results of the retrievers.
        Args:
            query: The query to search for.
        Returns:
            A list of merged documents.
        """
        # Get the results of all retrievers.
        retriever_docs = [
            retriever.get_relevant_documents(
                query, callbacks=run_manager.get_child("retriever_{}".format(i + 1))
            )
            for i, retriever in enumerate(self.retrievers)
        ]
        # Merge the results of the retrievers.
        merged_documents = []
        max_docs = max(len(docs) for docs in retriever_docs)
        for i in range(max_docs):
            for retriever, doc in zip(self.retrievers, retriever_docs):
                if i < len(doc):
                    merged_documents.append(doc[i])
        return merged_documents
[docs]    async def amerge_documents(
        self, query: str, run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[Document]:
        """
        Asynchronously merge the results of the retrievers.
        Args:
            query: The query to search for.
        Returns:
            A list of merged documents.
        """
        # Get the results of all retrievers.
        retriever_docs = await asyncio.gather(
            *(
                retriever.aget_relevant_documents(
                    query, callbacks=run_manager.get_child("retriever_{}".format(i + 1))
                )
                for i, retriever in enumerate(self.retrievers)
            )
        )
        # Merge the results of the retrievers.
        merged_documents = []
        max_docs = max(len(docs) for docs in retriever_docs)
        for i in range(max_docs):


for i in range(max_docs):
            for retriever, doc in zip(self.retrievers, retriever_docs):
                if i < len(doc):
                    merged_documents.append(doc[i])
        return merged_documents


Source code for langchain.retrievers.pubmed
from typing import List
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
from langchain.utilities.pubmed import PubMedAPIWrapper
[docs]class PubMedRetriever(BaseRetriever, PubMedAPIWrapper):
    """`PubMed API` retriever.
    It wraps load() to get_relevant_documents().
    It uses all PubMedAPIWrapper arguments without any change.
    """
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        return self.load_docs(query=query)


Source code for langchain.retrievers.ensemble
"""
Ensemble retriever that ensemble the results of
multiple retrievers by using weighted  Reciprocal Rank Fusion
"""
from typing import Any, Dict, List
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.pydantic_v1 import root_validator
from langchain.schema import BaseRetriever, Document
[docs]class EnsembleRetriever(BaseRetriever):
    """Retriever that ensembles the multiple retrievers.
    It uses a rank fusion.
    Args:
        retrievers: A list of retrievers to ensemble.
        weights: A list of weights corresponding to the retrievers. Defaults to equal
            weighting for all retrievers.
        c: A constant added to the rank, controlling the balance between the importance
            of high-ranked items and the consideration given to lower-ranked items.
            Default is 60.
    """
    retrievers: List[BaseRetriever]
    weights: List[float]
    c: int = 60
    @root_validator(pre=True)
    def set_weights(cls, values: Dict[str, Any]) -> Dict[str, Any]:
        if not values.get("weights"):
            n_retrievers = len(values["retrievers"])
            values["weights"] = [1 / n_retrievers] * n_retrievers
        return values
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """
        Get the relevant documents for a given query.
        Args:
            query: The query to search for.
        Returns:


Args:
            query: The query to search for.
        Returns:
            A list of reranked documents.
        """
        # Get fused result of the retrievers.
        fused_documents = self.rank_fusion(query, run_manager)
        return fused_documents
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
    ) -> List[Document]:
        """
        Asynchronously get the relevant documents for a given query.
        Args:
            query: The query to search for.
        Returns:
            A list of reranked documents.
        """
        # Get fused result of the retrievers.
        fused_documents = await self.arank_fusion(query, run_manager)
        return fused_documents
[docs]    def rank_fusion(
        self, query: str, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """
        Retrieve the results of the retrievers and use rank_fusion_func to get
        the final result.
        Args:
            query: The query to search for.
        Returns:
            A list of reranked documents.
        """
        # Get the results of all retrievers.
        retriever_docs = [
            retriever.get_relevant_documents(
                query, callbacks=run_manager.get_child(tag=f"retriever_{i+1}")
            )
            for i, retriever in enumerate(self.retrievers)
        ]
        # apply rank fusion
        fused_documents = self.weighted_reciprocal_rank(retriever_docs)
        return fused_documents
[docs]    async def arank_fusion(
        self, query: str, run_manager: AsyncCallbackManagerForRetrieverRun


self, query: str, run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[Document]:
        """
        Asynchronously retrieve the results of the retrievers
        and use rank_fusion_func to get the final result.
        Args:
            query: The query to search for.
        Returns:
            A list of reranked documents.
        """
        # Get the results of all retrievers.
        retriever_docs = [
            await retriever.aget_relevant_documents(
                query, callbacks=run_manager.get_child(tag=f"retriever_{i+1}")
            )
            for i, retriever in enumerate(self.retrievers)
        ]
        # apply rank fusion
        fused_documents = self.weighted_reciprocal_rank(retriever_docs)
        return fused_documents
[docs]    def weighted_reciprocal_rank(
        self, doc_lists: List[List[Document]]
    ) -> List[Document]:
        """
        Perform weighted Reciprocal Rank Fusion on multiple rank lists.
        You can find more details about RRF here:
        https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf
        Args:
            doc_lists: A list of rank lists, where each rank list contains unique items.
        Returns:
            list: The final aggregated list of items sorted by their weighted RRF
                    scores in descending order.
        """
        if len(doc_lists) != len(self.weights):
            raise ValueError(
                "Number of rank lists must be equal to the number of weights."
            )
        # Create a union of all unique documents in the input doc_lists
        all_documents = set()
        for doc_list in doc_lists:
            for doc in doc_list:


for doc_list in doc_lists:
            for doc in doc_list:
                all_documents.add(doc.page_content)
        # Initialize the RRF score dictionary for each document
        rrf_score_dic = {doc: 0.0 for doc in all_documents}
        # Calculate RRF scores for each document
        for doc_list, weight in zip(doc_lists, self.weights):
            for rank, doc in enumerate(doc_list, start=1):
                rrf_score = weight * (1 / (rank + self.c))
                rrf_score_dic[doc.page_content] += rrf_score
        # Sort documents by their RRF scores in descending order
        sorted_documents = sorted(
            rrf_score_dic.keys(), key=lambda x: rrf_score_dic[x], reverse=True
        )
        # Map the sorted page_content back to the original document objects
        page_content_to_doc_map = {
            doc.page_content: doc for doc_list in doc_lists for doc in doc_list
        }
        sorted_docs = [
            page_content_to_doc_map[page_content] for page_content in sorted_documents
        ]
        return sorted_docs


Source code for langchain.retrievers.weaviate_hybrid_search
from __future__ import annotations
from typing import Any, Dict, List, Optional, cast
from uuid import uuid4
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.docstore.document import Document
from langchain.pydantic_v1 import root_validator
from langchain.schema import BaseRetriever
[docs]class WeaviateHybridSearchRetriever(BaseRetriever):
    """`Weaviate hybrid search` retriever.
    See the documentation:
      https://weaviate.io/blog/hybrid-search-explained
    """
    client: Any
    """keyword arguments to pass to the Weaviate client."""
    index_name: str
    """The name of the index to use."""
    text_key: str
    """The name of the text key to use."""
    alpha: float = 0.5
    """The weight of the text key in the hybrid search."""
    k: int = 4
    """The number of results to return."""
    attributes: List[str]
    """The attributes to return in the results."""
    create_schema_if_missing: bool = True
    """Whether to create the schema if it doesn't exist."""
    @root_validator(pre=True)
    def validate_client(
        cls,
        values: Dict[str, Any],
    ) -> Dict[str, Any]:
        try:
            import weaviate
        except ImportError:
            raise ImportError(
                "Could not import weaviate python package. "
                "Please install it with `pip install weaviate-client`."
            )
        if not isinstance(values["client"], weaviate.Client):
            client = values["client"]
            raise ValueError(