if ids is None:
        # create unique ids using hash of the text
        ids = [hash_text(context) for context in contexts]
    for i in _iterator:
        # find end of batch
        i_end = min(i + batch_size, len(contexts))
        # extract batch
        context_batch = contexts[i:i_end]
        batch_ids = ids[i:i_end]
        metadata_batch = (
            metadatas[i:i_end] if metadatas else [{} for _ in context_batch]
        )
        # add context passages as metadata
        meta = [
            {"context": context, **metadata}
            for context, metadata in zip(context_batch, metadata_batch)
        ]
        # create dense vectors
        dense_embeds = embeddings.embed_documents(context_batch)
        # create sparse vectors
        sparse_embeds = sparse_encoder.encode_documents(context_batch)
        for s in sparse_embeds:
            s["values"] = [float(s1) for s1 in s["values"]]
        vectors = []
        # loop through the data and create dictionaries for upserts
        for doc_id, sparse, dense, metadata in zip(
            batch_ids, sparse_embeds, dense_embeds, meta
        ):
            vectors.append(
                {
                    "id": doc_id,
                    "sparse_values": sparse,
                    "values": dense,
                    "metadata": metadata,
                }
            )
        # upload the documents to the new hybrid index
        index.upsert(vectors, namespace=namespace)
[docs]class PineconeHybridSearchRetriever(BaseRetriever):
    """`Pinecone Hybrid Search` retriever."""
    embeddings: Embeddings
    """Embeddings model to use."""
    """description"""
    sparse_encoder: Any


"""Embeddings model to use."""
    """description"""
    sparse_encoder: Any
    """Sparse encoder to use."""
    index: Any
    """Pinecone index to use."""
    top_k: int = 4
    """Number of documents to return."""
    alpha: float = 0.5
    """Alpha value for hybrid search."""
    namespace: Optional[str] = None
    """Namespace value for index partition."""
    class Config:
        """Configuration for this pydantic object."""
        extra = Extra.forbid
        arbitrary_types_allowed = True
[docs]    def add_texts(
        self,
        texts: List[str],
        ids: Optional[List[str]] = None,
        metadatas: Optional[List[dict]] = None,
        namespace: Optional[str] = None,
    ) -> None:
        create_index(
            texts,
            self.index,
            self.embeddings,
            self.sparse_encoder,
            ids=ids,
            metadatas=metadatas,
            namespace=namespace,
        )
    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that api key and python package exists in environment."""
        try:
            from pinecone_text.hybrid import hybrid_convex_scale  # noqa:F401
            from pinecone_text.sparse.base_sparse_encoder import (
                BaseSparseEncoder,  # noqa:F401
            )
        except ImportError:
            raise ValueError(
                "Could not import pinecone_text python package. "
                "Please install it with `pip install pinecone_text`."
            )
        return values
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun


self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        from pinecone_text.hybrid import hybrid_convex_scale
        sparse_vec = self.sparse_encoder.encode_queries(query)
        # convert the question into a dense vector
        dense_vec = self.embeddings.embed_query(query)
        # scale alpha with hybrid_scale
        dense_vec, sparse_vec = hybrid_convex_scale(dense_vec, sparse_vec, self.alpha)
        sparse_vec["values"] = [float(s1) for s1 in sparse_vec["values"]]
        # query pinecone with the query parameters
        result = self.index.query(
            vector=dense_vec,
            sparse_vector=sparse_vec,
            top_k=self.top_k,
            include_metadata=True,
            namespace=self.namespace,
        )
        final_result = []
        for res in result["matches"]:
            context = res["metadata"].pop("context")
            final_result.append(
                Document(page_content=context, metadata=res["metadata"])
            )
        # return search results as json
        return final_result


Source code for langchain.retrievers.chaindesk
from typing import Any, List, Optional
import aiohttp
import requests
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
[docs]class ChaindeskRetriever(BaseRetriever):
    """`Chaindesk API` retriever."""
    datastore_url: str
    top_k: Optional[int]
    api_key: Optional[str]
    def __init__(
        self,
        datastore_url: str,
        top_k: Optional[int] = None,
        api_key: Optional[str] = None,
    ):
        self.datastore_url = datastore_url
        self.api_key = api_key
        self.top_k = top_k
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> 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"]
        ]


)
            for r in data["results"]
        ]
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        async with aiohttp.ClientSession() as session:
            async with session.request(
                "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 {}
                    ),
                },
            ) 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.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.
    """
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        return self.load(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.metal
from typing import Any, List, Optional
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.pydantic_v1 import root_validator
from langchain.schema import BaseRetriever, Document
[docs]class MetalRetriever(BaseRetriever):
    """`Metal API` retriever."""
    client: Any
    """The Metal client to use."""
    params: Optional[dict] = None
    """The parameters to pass to the Metal client."""
    @root_validator(pre=True)
    def validate_client(cls, values: dict) -> dict:
        """Validate that the client is of the correct type."""
        from metal_sdk.metal import Metal
        if "client" in values:
            client = values["client"]
            if not isinstance(client, Metal):
                raise ValueError(
                    "Got unexpected client, should be of type metal_sdk.metal.Metal. "
                    f"Instead, got {type(client)}"
                )
        values["params"] = values.get("params", {})
        return values
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        results = self.client.search({"text": query}, **self.params)
        final_results = []
        for r in results["data"]:
            metadata = {k: v for k, v in r.items() if k != "text"}
            final_results.append(Document(page_content=r["text"], metadata=metadata))
        return final_results


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("\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 key words 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.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.self_query.deeplake
"""Logic for converting internal query language to a valid Chroma query."""
from typing import Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
COMPARATOR_TO_TQL = {
    Comparator.EQ: "==",
    Comparator.GT: ">",
    Comparator.GTE: ">=",
    Comparator.LT: "<",
    Comparator.LTE: "<=",
}
OPERATOR_TO_TQL = {
    Operator.AND: "and",
    Operator.OR: "or",
    Operator.NOT: "NOT",
}
[docs]def can_cast_to_float(string: str) -> bool:
    """Check if a string can be cast to a float."""
    try:
        float(string)
        return True
    except ValueError:
        return False
[docs]class DeepLakeTranslator(Visitor):
    """Translate `DeepLake` internal query language elements to valid filters."""
    allowed_operators = [Operator.AND, Operator.OR, Operator.NOT]
    """Subset of allowed logical operators."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
    ]
    """Subset of allowed logical comparators."""
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        if isinstance(func, Operator):
            value = OPERATOR_TO_TQL[func.value]  # type: ignore
        elif isinstance(func, Comparator):
            value = COMPARATOR_TO_TQL[func.value]  # type: ignore


value = COMPARATOR_TO_TQL[func.value]  # type: ignore
        return f"{value}"
[docs]    def visit_operation(self, operation: Operation) -> str:
        args = [arg.accept(self) for arg in operation.arguments]
        operator = self._format_func(operation.operator)
        return "(" + (" " + operator + " ").join(args) + ")"
[docs]    def visit_comparison(self, comparison: Comparison) -> str:
        comparator = self._format_func(comparison.comparator)
        values = comparison.value
        if isinstance(values, list):
            tql = []
            for value in values:
                comparison.value = value
                tql.append(self.visit_comparison(comparison))
            return "(" + (" or ").join(tql) + ")"
        if not can_cast_to_float(comparison.value):
            values = f"'{values}'"
        return f"metadata['{comparison.attribute}'] {comparator} {values}"
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            tqL = f"SELECT * WHERE {structured_query.filter.accept(self)}"
            kwargs = {"tql": tqL}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.timescalevector
from __future__ import annotations
from typing import TYPE_CHECKING, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
if TYPE_CHECKING:
    from timescale_vector import client
[docs]class TimescaleVectorTranslator(Visitor):
    """Translate the internal query language elements to valid filters."""
    allowed_operators = [Operator.AND, Operator.OR, Operator.NOT]
    """Subset of allowed logical operators."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
    ]
    COMPARATOR_MAP = {
        Comparator.EQ: "==",
        Comparator.GT: ">",
        Comparator.GTE: ">=",
        Comparator.LT: "<",
        Comparator.LTE: "<=",
    }
    OPERATOR_MAP = {Operator.AND: "AND", Operator.OR: "OR", Operator.NOT: "NOT"}
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        if isinstance(func, Operator):
            value = self.OPERATOR_MAP[func.value]  # type: ignore
        elif isinstance(func, Comparator):
            value = self.COMPARATOR_MAP[func.value]  # type: ignore
        return f"{value}"
[docs]    def visit_operation(self, operation: Operation) -> client.Predicates:
        try:
            from timescale_vector import client
        except ImportError as e:
            raise ImportError(
                "Cannot import timescale-vector. Please install with `pip install "


"Cannot import timescale-vector. Please install with `pip install "
                "timescale-vector`."
            ) from e
        args = [arg.accept(self) for arg in operation.arguments]
        return client.Predicates(*args, operator=self._format_func(operation.operator))
[docs]    def visit_comparison(self, comparison: Comparison) -> client.Predicates:
        try:
            from timescale_vector import client
        except ImportError as e:
            raise ImportError(
                "Cannot import timescale-vector. Please install with `pip install "
                "timescale-vector`."
            ) from e
        return client.Predicates(
            (
                comparison.attribute,
                self._format_func(comparison.comparator),
                comparison.value,
            )
        )
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"predicates": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.dashvector
"""Logic for converting internal query language to a valid DashVector query."""
from typing import Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class DashvectorTranslator(Visitor):
    """Logic for converting internal query language elements to valid filters."""
    allowed_operators = [Operator.AND, Operator.OR]
    allowed_comparators = [
        Comparator.EQ,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
        Comparator.LIKE,
    ]
    map_dict = {
        Operator.AND: " AND ",
        Operator.OR: " OR ",
        Comparator.EQ: " = ",
        Comparator.GT: " > ",
        Comparator.GTE: " >= ",
        Comparator.LT: " < ",
        Comparator.LTE: " <= ",
        Comparator.LIKE: " LIKE ",
    }
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        return self.map_dict[func]
[docs]    def visit_operation(self, operation: Operation) -> str:
        args = [arg.accept(self) for arg in operation.arguments]
        return self._format_func(operation.operator).join(args)
[docs]    def visit_comparison(self, comparison: Comparison) -> str:
        value = comparison.value
        if isinstance(value, str):
            if comparison.comparator == Comparator.LIKE:
                value = f"'%{value}%'"
            else:
                value = f"'{value}'"
        return (


else:
                value = f"'{value}'"
        return (
            f"{comparison.attribute}{self._format_func(comparison.comparator)}{value}"
        )
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"filter": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.qdrant
from __future__ import annotations
from typing import TYPE_CHECKING, Tuple
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
if TYPE_CHECKING:
    from qdrant_client.http import models as rest
[docs]class QdrantTranslator(Visitor):
    """Translate `Qdrant` internal query language elements to valid filters."""
    allowed_comparators = (
        Comparator.EQ,
        Comparator.LT,
        Comparator.LTE,
        Comparator.GT,
        Comparator.GTE,
    )
    """Subset of allowed logical comparators."""
[docs]    def __init__(self, metadata_key: str):
        self.metadata_key = metadata_key
[docs]    def visit_operation(self, operation: Operation) -> rest.Filter:
        try:
            from qdrant_client.http import models as rest
        except ImportError as e:
            raise ImportError(
                "Cannot import qdrant_client. Please install with `pip install "
                "qdrant-client`."
            ) from e
        args = [arg.accept(self) for arg in operation.arguments]
        operator = {
            Operator.AND: "must",
            Operator.OR: "should",
            Operator.NOT: "must_not",
        }[operation.operator]
        return rest.Filter(**{operator: args})
[docs]    def visit_comparison(self, comparison: Comparison) -> rest.FieldCondition:
        try:
            from qdrant_client.http import models as rest
        except ImportError as e:
            raise ImportError(
                "Cannot import qdrant_client. Please install with `pip install "


"Cannot import qdrant_client. Please install with `pip install "
                "qdrant-client`."
            ) from e
        self._validate_func(comparison.comparator)
        attribute = self.metadata_key + "." + comparison.attribute
        if comparison.comparator == Comparator.EQ:
            return rest.FieldCondition(
                key=attribute, match=rest.MatchValue(value=comparison.value)
            )
        kwargs = {comparison.comparator.value: comparison.value}
        return rest.FieldCondition(key=attribute, range=rest.Range(**kwargs))
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        try:
            from qdrant_client.http import models as rest
        except ImportError as e:
            raise ImportError(
                "Cannot import qdrant_client. Please install with `pip install "
                "qdrant-client`."
            ) from e
        if structured_query.filter is None:
            kwargs = {}
        else:
            filter = structured_query.filter.accept(self)
            if isinstance(filter, rest.FieldCondition):
                filter = rest.Filter(must=[filter])
            kwargs = {"filter": filter}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.milvus
"""Logic for converting internal query language to a valid Milvus query."""
from typing import Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
COMPARATOR_TO_BER = {
    Comparator.EQ: "==",
    Comparator.GT: ">",
    Comparator.GTE: ">=",
    Comparator.LT: "<",
    Comparator.LTE: "<=",
}
UNARY_OPERATORS = [Operator.NOT]
[docs]def process_value(value: Union[int, float, str]) -> str:
    # required for comparators involving strings
    if isinstance(value, str):
        # If the value is already a string, add double quotes
        return f'"{value}"'
    else:
        # If the valueis not a string, convert it to a string without double quotes
        return str(value)
[docs]class MilvusTranslator(Visitor):
    """Translate Milvus internal query language elements to valid filters."""
    """Subset of allowed logical operators."""
    allowed_operators = [Operator.AND, Operator.NOT, Operator.OR]
    """Subset of allowed logical comparators."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
    ]
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        value = func.value
        if isinstance(func, Comparator):
            value = COMPARATOR_TO_BER[func]
        return f"{value}"


value = COMPARATOR_TO_BER[func]
        return f"{value}"
[docs]    def visit_operation(self, operation: Operation) -> str:
        if operation.operator in UNARY_OPERATORS and len(operation.arguments) == 1:
            operator = self._format_func(operation.operator)
            return operator + "(" + operation.arguments[0].accept(self) + ")"
        elif operation.operator in UNARY_OPERATORS:
            raise ValueError(
                f'"{operation.operator.value}" can have only one argument in Milvus'
            )
        else:
            args = [arg.accept(self) for arg in operation.arguments]
            operator = self._format_func(operation.operator)
            return "(" + (" " + operator + " ").join(args) + ")"
[docs]    def visit_comparison(self, comparison: Comparison) -> str:
        comparator = self._format_func(comparison.comparator)
        processed_value = process_value(comparison.value)
        attribute = comparison.attribute
        return "( " + attribute + " " + comparator + " " + processed_value + " )"
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"expr": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.supabase
from typing import Any, Dict, Tuple
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class SupabaseVectorTranslator(Visitor):
    """Translate Langchain filters to Supabase PostgREST filters."""
    allowed_operators = [Operator.AND, Operator.OR]
    """Subset of allowed logical operators."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.NE,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
        Comparator.LIKE,
    ]
    """Subset of allowed logical comparators."""
    metadata_column = "metadata"
    def _map_comparator(self, comparator: Comparator) -> str:
        """
        Maps Langchain comparator to PostgREST comparator:
        https://postgrest.org/en/stable/references/api/tables_views.html#operators
        """
        postgrest_comparator = {
            Comparator.EQ: "eq",
            Comparator.NE: "neq",
            Comparator.GT: "gt",
            Comparator.GTE: "gte",
            Comparator.LT: "lt",
            Comparator.LTE: "lte",
            Comparator.LIKE: "like",
        }.get(comparator)
        if postgrest_comparator is None:
            raise Exception(
                f"Comparator '{comparator}' is not currently "
                "supported in Supabase Vector"
            )
        return postgrest_comparator
    def _get_json_operator(self, value: Any) -> str:
        if isinstance(value, str):
            return "->>"


if isinstance(value, str):
            return "->>"
        else:
            return "->"
[docs]    def visit_operation(self, operation: Operation) -> str:
        args = [arg.accept(self) for arg in operation.arguments]
        return f"{operation.operator.value}({','.join(args)})"
[docs]    def visit_comparison(self, comparison: Comparison) -> str:
        if isinstance(comparison.value, list):
            return self.visit_operation(
                Operation(
                    operator=Operator.AND,
                    arguments=(
                        Comparison(
                            comparator=comparison.comparator,
                            attribute=comparison.attribute,
                            value=value,
                        )
                        for value in comparison.value
                    ),
                )
            )
        return ".".join(
            [
                f"{self.metadata_column}{self._get_json_operator(comparison.value)}{comparison.attribute}",
                f"{self._map_comparator(comparison.comparator)}",
                f"{comparison.value}",
            ]
        )
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, Dict[str, str]]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"postgrest_filter": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.base
"""Retriever that generates and executes structured queries over its own data source."""
from typing import Any, Dict, List, Optional, Type, cast
from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.chains import LLMChain
from langchain.chains.query_constructor.base import load_query_constructor_chain
from langchain.chains.query_constructor.ir import StructuredQuery, Visitor
from langchain.chains.query_constructor.schema import AttributeInfo
from langchain.pydantic_v1 import BaseModel, Field, root_validator
from langchain.retrievers.self_query.chroma import ChromaTranslator
from langchain.retrievers.self_query.dashvector import DashvectorTranslator
from langchain.retrievers.self_query.deeplake import DeepLakeTranslator
from langchain.retrievers.self_query.elasticsearch import ElasticsearchTranslator
from langchain.retrievers.self_query.milvus import MilvusTranslator
from langchain.retrievers.self_query.myscale import MyScaleTranslator
from langchain.retrievers.self_query.opensearch import OpenSearchTranslator
from langchain.retrievers.self_query.pinecone import PineconeTranslator
from langchain.retrievers.self_query.qdrant import QdrantTranslator
from langchain.retrievers.self_query.redis import RedisTranslator
from langchain.retrievers.self_query.supabase import SupabaseVectorTranslator
from langchain.retrievers.self_query.timescalevector import TimescaleVectorTranslator
from langchain.retrievers.self_query.vectara import VectaraTranslator
from langchain.retrievers.self_query.weaviate import WeaviateTranslator
from langchain.schema import BaseRetriever, Document
from langchain.schema.language_model import BaseLanguageModel
from langchain.vectorstores import (
    Chroma,
    DashVector,
    DeepLake,
    ElasticsearchStore,


Chroma,
    DashVector,
    DeepLake,
    ElasticsearchStore,
    Milvus,
    MyScale,
    OpenSearchVectorSearch,
    Pinecone,
    Qdrant,
    Redis,
    SupabaseVectorStore,
    TimescaleVector,
    Vectara,
    VectorStore,
    Weaviate,
)
def _get_builtin_translator(vectorstore: VectorStore) -> Visitor:
    """Get the translator class corresponding to the vector store class."""
    BUILTIN_TRANSLATORS: Dict[Type[VectorStore], Type[Visitor]] = {
        Pinecone: PineconeTranslator,
        Chroma: ChromaTranslator,
        DashVector: DashvectorTranslator,
        Weaviate: WeaviateTranslator,
        Vectara: VectaraTranslator,
        Qdrant: QdrantTranslator,
        MyScale: MyScaleTranslator,
        DeepLake: DeepLakeTranslator,
        ElasticsearchStore: ElasticsearchTranslator,
        Milvus: MilvusTranslator,
        SupabaseVectorStore: SupabaseVectorTranslator,
        TimescaleVector: TimescaleVectorTranslator,
        OpenSearchVectorSearch: OpenSearchTranslator,
    }
    if isinstance(vectorstore, Qdrant):
        return QdrantTranslator(metadata_key=vectorstore.metadata_payload_key)
    elif isinstance(vectorstore, MyScale):
        return MyScaleTranslator(metadata_key=vectorstore.metadata_column)
    elif isinstance(vectorstore, Redis):
        return RedisTranslator.from_vectorstore(vectorstore)
    elif vectorstore.__class__ in BUILTIN_TRANSLATORS:
        return BUILTIN_TRANSLATORS[vectorstore.__class__]()
    else:
        raise ValueError(
            f"Self query retriever with Vector Store type {vectorstore.__class__}"


f"Self query retriever with Vector Store type {vectorstore.__class__}"
            f" not supported."
        )
[docs]class SelfQueryRetriever(BaseRetriever, BaseModel):
    """Retriever that uses a vector store and an LLM to generate
    the vector store queries."""
    vectorstore: VectorStore
    """The underlying vector store from which documents will be retrieved."""
    llm_chain: LLMChain
    """The LLMChain for generating the vector store queries."""
    search_type: str = "similarity"
    """The search type to perform on the vector store."""
    search_kwargs: dict = Field(default_factory=dict)
    """Keyword arguments to pass in to the vector store search."""
    structured_query_translator: Visitor
    """Translator for turning internal query language into vectorstore search params."""
    verbose: bool = False
    use_original_query: bool = False
    """Use original query instead of the revised new query from LLM"""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
    @root_validator(pre=True)
    def validate_translator(cls, values: Dict) -> Dict:
        """Validate translator."""
        if "structured_query_translator" not in values:
            values["structured_query_translator"] = _get_builtin_translator(
                values["vectorstore"]
            )
        return values
    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
        """


Returns:
            List of relevant documents
        """
        inputs = self.llm_chain.prep_inputs({"query": query})
        structured_query = cast(
            StructuredQuery,
            self.llm_chain.predict_and_parse(
                callbacks=run_manager.get_child(), **inputs
            ),
        )
        if self.verbose:
            print(structured_query)
        new_query, new_kwargs = self.structured_query_translator.visit_structured_query(
            structured_query
        )
        if structured_query.limit is not None:
            new_kwargs["k"] = structured_query.limit
        if self.use_original_query:
            new_query = query
        search_kwargs = {**self.search_kwargs, **new_kwargs}
        docs = self.vectorstore.search(new_query, self.search_type, **search_kwargs)
        return docs
[docs]    @classmethod
    def from_llm(
        cls,
        llm: BaseLanguageModel,
        vectorstore: VectorStore,
        document_contents: str,
        metadata_field_info: List[AttributeInfo],
        structured_query_translator: Optional[Visitor] = None,
        chain_kwargs: Optional[Dict] = None,
        enable_limit: bool = False,
        use_original_query: bool = False,
        **kwargs: Any,
    ) -> "SelfQueryRetriever":
        if structured_query_translator is None:
            structured_query_translator = _get_builtin_translator(vectorstore)
        chain_kwargs = chain_kwargs or {}
        if "allowed_comparators" not in chain_kwargs:
            chain_kwargs[
                "allowed_comparators"
            ] = structured_query_translator.allowed_comparators
        if "allowed_operators" not in chain_kwargs:
            chain_kwargs[
                "allowed_operators"


chain_kwargs[
                "allowed_operators"
            ] = structured_query_translator.allowed_operators
        llm_chain = load_query_constructor_chain(
            llm,
            document_contents,
            metadata_field_info,
            enable_limit=enable_limit,
            **chain_kwargs,
        )
        return cls(
            llm_chain=llm_chain,
            vectorstore=vectorstore,
            use_original_query=use_original_query,
            structured_query_translator=structured_query_translator,
            **kwargs,
        )


Source code for langchain.retrievers.self_query.weaviate
from typing import Dict, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class WeaviateTranslator(Visitor):
    """Translate `Weaviate` internal query language elements to valid filters."""
    allowed_operators = [Operator.AND, Operator.OR]
    """Subset of allowed logical operators."""
    allowed_comparators = [Comparator.EQ]
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        # https://weaviate.io/developers/weaviate/api/graphql/filters
        map_dict = {Operator.AND: "And", Operator.OR: "Or", Comparator.EQ: "Equal"}
        return map_dict[func]
[docs]    def visit_operation(self, operation: Operation) -> Dict:
        args = [arg.accept(self) for arg in operation.arguments]
        return {"operator": self._format_func(operation.operator), "operands": args}
[docs]    def visit_comparison(self, comparison: Comparison) -> Dict:
        return {
            "path": [comparison.attribute],
            "operator": self._format_func(comparison.comparator),
            "valueText": comparison.value,
        }
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"where_filter": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.redis
from __future__ import annotations
from typing import Any, Tuple
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
from langchain.vectorstores.redis import Redis
from langchain.vectorstores.redis.filters import (
    RedisFilterExpression,
    RedisFilterField,
    RedisFilterOperator,
    RedisNum,
    RedisTag,
    RedisText,
)
from langchain.vectorstores.redis.schema import RedisModel
_COMPARATOR_TO_BUILTIN_METHOD = {
    Comparator.EQ: "__eq__",
    Comparator.NE: "__ne__",
    Comparator.LT: "__lt__",
    Comparator.GT: "__gt__",
    Comparator.LTE: "__le__",
    Comparator.GTE: "__ge__",
    Comparator.CONTAIN: "__eq__",
    Comparator.LIKE: "__mod__",
}
[docs]class RedisTranslator(Visitor):
    """Translate"""
    allowed_comparators = (
        Comparator.EQ,
        Comparator.NE,
        Comparator.LT,
        Comparator.LTE,
        Comparator.GT,
        Comparator.GTE,
        Comparator.CONTAIN,
        Comparator.LIKE,
    )
    """Subset of allowed logical comparators."""
    allowed_operators = (Operator.AND, Operator.OR)
    """Subset of allowed logical operators."""
[docs]    def __init__(self, schema: RedisModel) -> None:
        self._schema = schema
    def _attribute_to_filter_field(self, attribute: str) -> RedisFilterField:
        if attribute in [tf.name for tf in self._schema.text]:
            return RedisText(attribute)


return RedisText(attribute)
        elif attribute in [tf.name for tf in self._schema.tag or []]:
            return RedisTag(attribute)
        elif attribute in [tf.name for tf in self._schema.numeric or []]:
            return RedisNum(attribute)
        else:
            raise ValueError(
                f"Invalid attribute {attribute} not in vector store schema. Schema is:"
                f"\n{self._schema.as_dict()}"
            )
[docs]    def visit_comparison(self, comparison: Comparison) -> RedisFilterExpression:
        filter_field = self._attribute_to_filter_field(comparison.attribute)
        comparison_method = _COMPARATOR_TO_BUILTIN_METHOD[comparison.comparator]
        return getattr(filter_field, comparison_method)(comparison.value)
[docs]    def visit_operation(self, operation: Operation) -> Any:
        left = operation.arguments[0].accept(self)
        if len(operation.arguments) > 2:
            right = self.visit_operation(
                Operation(
                    operator=operation.operator, arguments=operation.arguments[1:]
                )
            )
        else:
            right = operation.arguments[1].accept(self)
        redis_operator = (
            RedisFilterOperator.OR
            if operation.operator == Operator.OR
            else RedisFilterOperator.AND
        )
        return RedisFilterExpression(operator=redis_operator, left=left, right=right)
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"filter": structured_query.filter.accept(self)}
        return structured_query.query, kwargs
[docs]    @classmethod
    def from_vectorstore(cls, vectorstore: Redis) -> RedisTranslator:


def from_vectorstore(cls, vectorstore: Redis) -> RedisTranslator:
        return cls(vectorstore._schema)


Source code for langchain.retrievers.self_query.vectara
from typing import Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]def process_value(value: Union[int, float, str]) -> str:
    if isinstance(value, str):
        return f"'{value}'"
    else:
        return str(value)
[docs]class VectaraTranslator(Visitor):
    """Translate `Vectara` internal query language elements to valid filters."""
    allowed_operators = [Operator.AND, Operator.OR]
    """Subset of allowed logical operators."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.NE,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
    ]
    """Subset of allowed logical comparators."""
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        map_dict = {
            Operator.AND: " and ",
            Operator.OR: " or ",
            Comparator.EQ: "=",
            Comparator.NE: "!=",
            Comparator.GT: ">",
            Comparator.GTE: ">=",
            Comparator.LT: "<",
            Comparator.LTE: "<=",
        }
        self._validate_func(func)
        return map_dict[func]
[docs]    def visit_operation(self, operation: Operation) -> str:
        args = [arg.accept(self) for arg in operation.arguments]
        operator = self._format_func(operation.operator)
        return "( " + operator.join(args) + " )"
[docs]    def visit_comparison(self, comparison: Comparison) -> str:


[docs]    def visit_comparison(self, comparison: Comparison) -> str:
        comparator = self._format_func(comparison.comparator)
        processed_value = process_value(comparison.value)
        attribute = comparison.attribute
        return (
            "( " + "doc." + attribute + " " + comparator + " " + processed_value + " )"
        )
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"filter": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.elasticsearch
from typing import Dict, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class ElasticsearchTranslator(Visitor):
    """Translate `Elasticsearch` internal query language elements to valid filters."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
        Comparator.CONTAIN,
        Comparator.LIKE,
    ]
    """Subset of allowed logical comparators."""
    allowed_operators = [Operator.AND, Operator.OR, Operator.NOT]
    """Subset of allowed logical operators."""
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        map_dict = {
            Operator.OR: "should",
            Operator.NOT: "must_not",
            Operator.AND: "must",
            Comparator.EQ: "term",
            Comparator.GT: "gt",
            Comparator.GTE: "gte",
            Comparator.LT: "lt",
            Comparator.LTE: "lte",
            Comparator.CONTAIN: "match",
            Comparator.LIKE: "fuzzy",
        }
        return map_dict[func]
[docs]    def visit_operation(self, operation: Operation) -> Dict:
        args = [arg.accept(self) for arg in operation.arguments]
        return {"bool": {self._format_func(operation.operator): args}}
[docs]    def visit_comparison(self, comparison: Comparison) -> Dict:
        # ElasticsearchStore filters require to target
        # the metadata object field


# ElasticsearchStore filters require to target
        # the metadata object field
        field = f"metadata.{comparison.attribute}"
        is_range_comparator = comparison.comparator in [
            Comparator.GT,
            Comparator.GTE,
            Comparator.LT,
            Comparator.LTE,
        ]
        if is_range_comparator:
            return {
                "range": {
                    field: {self._format_func(comparison.comparator): comparison.value}
                }
            }
        if comparison.comparator == Comparator.LIKE:
            return {
                self._format_func(comparison.comparator): {
                    field: {"value": comparison.value, "fuzziness": "AUTO"}
                }
            }
        if comparison.comparator == Comparator.CONTAIN:
            return {self._format_func(comparison.comparator): {field: comparison.value}}
        # we assume that if the value is a string,
        # we want to use the keyword field
        field = f"{field}.keyword" if isinstance(comparison.value, str) else field
        return {self._format_func(comparison.comparator): {field: comparison.value}}
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"filter": [structured_query.filter.accept(self)]}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.pinecone
from typing import Dict, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class PineconeTranslator(Visitor):
    """Translate `Pinecone` internal query language elements to valid filters."""
    allowed_comparators = (
        Comparator.EQ,
        Comparator.NE,
        Comparator.LT,
        Comparator.LTE,
        Comparator.GT,
        Comparator.GTE,
    )
    """Subset of allowed logical comparators."""
    allowed_operators = (Operator.AND, Operator.OR)
    """Subset of allowed logical operators."""
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        return f"${func.value}"
[docs]    def visit_operation(self, operation: Operation) -> Dict:
        args = [arg.accept(self) for arg in operation.arguments]
        return {self._format_func(operation.operator): args}
[docs]    def visit_comparison(self, comparison: Comparison) -> Dict:
        return {
            comparison.attribute: {
                self._format_func(comparison.comparator): comparison.value
            }
        }
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"filter": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.myscale
import datetime
import re
from typing import Any, Callable, Dict, Tuple
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
def _DEFAULT_COMPOSER(op_name: str) -> Callable:
    """
    Default composer for logical operators.
    Args:
        op_name: Name of the operator.
    Returns:
        Callable that takes a list of arguments and returns a string.
    """
    def f(*args: Any) -> str:
        args_: map[str] = map(str, args)
        return f" {op_name} ".join(args_)
    return f
def _FUNCTION_COMPOSER(op_name: str) -> Callable:
    """
    Composer for functions.
    Args:
        op_name: Name of the function.
    Returns:
        Callable that takes a list of arguments and returns a string.
    """
    def f(*args: Any) -> str:
        args_: map[str] = map(str, args)
        return f"{op_name}({','.join(args_)})"
    return f
[docs]class MyScaleTranslator(Visitor):
    """Translate `MyScale` internal query language elements to valid filters."""
    allowed_operators = [Operator.AND, Operator.OR, Operator.NOT]
    """Subset of allowed logical operators."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
        Comparator.CONTAIN,
        Comparator.LIKE,
    ]
    map_dict = {
        Operator.AND: _DEFAULT_COMPOSER("AND"),


map_dict = {
        Operator.AND: _DEFAULT_COMPOSER("AND"),
        Operator.OR: _DEFAULT_COMPOSER("OR"),
        Operator.NOT: _DEFAULT_COMPOSER("NOT"),
        Comparator.EQ: _DEFAULT_COMPOSER("="),
        Comparator.GT: _DEFAULT_COMPOSER(">"),
        Comparator.GTE: _DEFAULT_COMPOSER(">="),
        Comparator.LT: _DEFAULT_COMPOSER("<"),
        Comparator.LTE: _DEFAULT_COMPOSER("<="),
        Comparator.CONTAIN: _FUNCTION_COMPOSER("has"),
        Comparator.LIKE: _DEFAULT_COMPOSER("ILIKE"),
    }
[docs]    def __init__(self, metadata_key: str = "metadata") -> None:
        super().__init__()
        self.metadata_key = metadata_key
[docs]    def visit_operation(self, operation: Operation) -> Dict:
        args = [arg.accept(self) for arg in operation.arguments]
        func = operation.operator
        self._validate_func(func)
        return self.map_dict[func](*args)
[docs]    def visit_comparison(self, comparison: Comparison) -> Dict:
        regex = "\((.*?)\)"
        matched = re.search("\(\w+\)", comparison.attribute)
        # If arbitrary function is applied to an attribute
        if matched:
            attr = re.sub(
                regex,
                f"({self.metadata_key}.{matched.group(0)[1:-1]})",
                comparison.attribute,
            )
        else:
            attr = f"{self.metadata_key}.{comparison.attribute}"
        value = comparison.value
        comp = comparison.comparator
        value = f"'{value}'" if isinstance(value, str) else value
        # convert timestamp for datetime objects
        if type(value) is datetime.date:


# convert timestamp for datetime objects
        if type(value) is datetime.date:
            attr = f"parseDateTime32BestEffort({attr})"
            value = f"parseDateTime32BestEffort('{value.strftime('%Y-%m-%d')}')"
        # string pattern match
        if comp is Comparator.LIKE:
            value = f"'%{value[1:-1]}%'"
        return self.map_dict[comp](attr, value)
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        print(structured_query)
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"where_str": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.opensearch
from typing import Dict, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class OpenSearchTranslator(Visitor):
    """Translate `OpenSearch` internal query domain-specific
    language elements to valid filters."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.LT,
        Comparator.LTE,
        Comparator.GT,
        Comparator.GTE,
        Comparator.CONTAIN,
        Comparator.LIKE,
    ]
    """Subset of allowed logical comparators."""
    allowed_operators = [Operator.AND, Operator.OR, Operator.NOT]
    """Subset of allowed logical operators."""
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        comp_operator_map = {
            Comparator.EQ: "term",
            Comparator.LT: "lt",
            Comparator.LTE: "lte",
            Comparator.GT: "gt",
            Comparator.GTE: "gte",
            Comparator.CONTAIN: "match",
            Comparator.LIKE: "fuzzy",
            Operator.AND: "must",
            Operator.OR: "should",
            Operator.NOT: "must_not",
        }
        return comp_operator_map[func]
[docs]    def visit_operation(self, operation: Operation) -> Dict:
        args = [arg.accept(self) for arg in operation.arguments]
        return {"bool": {self._format_func(operation.operator): args}}
[docs]    def visit_comparison(self, comparison: Comparison) -> Dict:
        field = f"metadata.{comparison.attribute}"


field = f"metadata.{comparison.attribute}"
        if comparison.comparator in [
            Comparator.LT,
            Comparator.LTE,
            Comparator.GT,
            Comparator.GTE,
        ]:
            return {
                "range": {
                    field: {self._format_func(comparison.comparator): comparison.value}
                }
            }
        if comparison.comparator == Comparator.LIKE:
            return {
                self._format_func(comparison.comparator): {
                    field: {"value": comparison.value}
                }
            }
        field = f"{field}.keyword" if isinstance(comparison.value, str) else field
        return {self._format_func(comparison.comparator): {field: comparison.value}}
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"filter": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.self_query.chroma
from typing import Dict, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class ChromaTranslator(Visitor):
    """Translate `Chroma` internal query language elements to valid filters."""
    allowed_operators = [Operator.AND, Operator.OR]
    """Subset of allowed logical operators."""
    allowed_comparators = [
        Comparator.EQ,
        Comparator.NE,
        Comparator.GT,
        Comparator.GTE,
        Comparator.LT,
        Comparator.LTE,
    ]
    """Subset of allowed logical comparators."""
    def _format_func(self, func: Union[Operator, Comparator]) -> str:
        self._validate_func(func)
        return f"${func.value}"
[docs]    def visit_operation(self, operation: Operation) -> Dict:
        args = [arg.accept(self) for arg in operation.arguments]
        return {self._format_func(operation.operator): args}
[docs]    def visit_comparison(self, comparison: Comparison) -> Dict:
        return {
            comparison.attribute: {
                self._format_func(comparison.comparator): comparison.value
            }
        }
[docs]    def visit_structured_query(
        self, structured_query: StructuredQuery
    ) -> Tuple[str, dict]:
        if structured_query.filter is None:
            kwargs = {}
        else:
            kwargs = {"filter": structured_query.filter.accept(self)}
        return structured_query.query, kwargs


Source code for langchain.retrievers.document_compressors.cohere_rerank
from __future__ import annotations
from typing import TYPE_CHECKING, Dict, Optional, Sequence
from langchain.callbacks.manager import Callbacks
from langchain.pydantic_v1 import Extra, root_validator
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
from langchain.schema import Document
from langchain.utils import get_from_dict_or_env
if TYPE_CHECKING:
    from cohere import Client
else:
    # We do to avoid pydantic annotation issues when actually instantiating
    # while keeping this import optional
    try:
        from cohere import Client
    except ImportError:
        pass
[docs]class CohereRerank(BaseDocumentCompressor):
    """Document compressor that uses `Cohere Rerank API`."""
    client: Client
    """Cohere client to use for compressing documents."""
    top_n: int = 3
    """Number of documents to return."""
    model: str = "rerank-english-v2.0"
    """Model to use for reranking."""
    class Config:
        """Configuration for this pydantic object."""
        extra = Extra.forbid
        arbitrary_types_allowed = True
    @root_validator(pre=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that api key and python package exists in environment."""
        cohere_api_key = get_from_dict_or_env(
            values, "cohere_api_key", "COHERE_API_KEY"
        )
        try:
            import cohere
            values["client"] = cohere.Client(cohere_api_key)
        except ImportError:
            raise ImportError(
                "Could not import cohere python package. "


raise ImportError(
                "Could not import cohere python package. "
                "Please install it with `pip install cohere`."
            )
        return values
[docs]    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """
        Compress documents using Cohere's rerank API.
        Args:
            documents: A sequence of documents to compress.
            query: The query to use for compressing the documents.
            callbacks: Callbacks to run during the compression process.
        Returns:
            A sequence of compressed documents.
        """
        if len(documents) == 0:  # to avoid empty api call
            return []
        doc_list = list(documents)
        _docs = [d.page_content for d in doc_list]
        results = self.client.rerank(
            model=self.model, query=query, documents=_docs, top_n=self.top_n
        )
        final_results = []
        for r in results:
            doc = doc_list[r.index]
            doc.metadata["relevance_score"] = r.relevance_score
            final_results.append(doc)
        return final_results
[docs]    async def acompress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        raise NotImplementedError()


Source code for langchain.retrievers.document_compressors.chain_filter
"""Filter that uses an LLM to drop documents that aren't relevant to the query."""
from typing import Any, Callable, Dict, Optional, Sequence
from langchain.callbacks.manager import Callbacks
from langchain.chains import LLMChain
from langchain.output_parsers.boolean import BooleanOutputParser
from langchain.prompts import PromptTemplate
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
from langchain.retrievers.document_compressors.chain_filter_prompt import (
    prompt_template,
)
from langchain.schema import BasePromptTemplate, Document
from langchain.schema.language_model import BaseLanguageModel
def _get_default_chain_prompt() -> PromptTemplate:
    return PromptTemplate(
        template=prompt_template,
        input_variables=["question", "context"],
        output_parser=BooleanOutputParser(),
    )
[docs]def default_get_input(query: str, doc: Document) -> Dict[str, Any]:
    """Return the compression chain input."""
    return {"question": query, "context": doc.page_content}
[docs]class LLMChainFilter(BaseDocumentCompressor):
    """Filter that drops documents that aren't relevant to the query."""
    llm_chain: LLMChain
    """LLM wrapper to use for filtering documents.
    The chain prompt is expected to have a BooleanOutputParser."""
    get_input: Callable[[str, Document], dict] = default_get_input
    """Callable for constructing the chain input from the query and a Document."""
[docs]    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Filter down documents based on their relevance to the query."""


"""Filter down documents based on their relevance to the query."""
        filtered_docs = []
        for doc in documents:
            _input = self.get_input(query, doc)
            include_doc = self.llm_chain.predict_and_parse(
                **_input, callbacks=callbacks
            )
            if include_doc:
                filtered_docs.append(doc)
        return filtered_docs
[docs]    async def acompress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Filter down documents."""
        raise NotImplementedError()
[docs]    @classmethod
    def from_llm(
        cls,
        llm: BaseLanguageModel,
        prompt: Optional[BasePromptTemplate] = None,
        **kwargs: Any
    ) -> "LLMChainFilter":
        """Create a LLMChainFilter from a language model.
        Args:
            llm: The language model to use for filtering.
            prompt: The prompt to use for the filter.
            **kwargs: Additional arguments to pass to the constructor.
        Returns:
            A LLMChainFilter that uses the given language model.
        """
        _prompt = prompt if prompt is not None else _get_default_chain_prompt()
        llm_chain = LLMChain(llm=llm, prompt=_prompt)
        return cls(llm_chain=llm_chain, **kwargs)


Source code for langchain.retrievers.document_compressors.base
from abc import ABC, abstractmethod
from inspect import signature
from typing import List, Optional, Sequence, Union
from langchain.callbacks.manager import Callbacks
from langchain.pydantic_v1 import BaseModel
from langchain.schema import BaseDocumentTransformer, Document
[docs]class BaseDocumentCompressor(BaseModel, ABC):
    """Base class for document compressors."""
[docs]    @abstractmethod
    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Compress retrieved documents given the query context."""
[docs]    @abstractmethod
    async def acompress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Compress retrieved documents given the query context."""
[docs]class DocumentCompressorPipeline(BaseDocumentCompressor):
    """Document compressor that uses a pipeline of Transformers."""
    transformers: List[Union[BaseDocumentTransformer, BaseDocumentCompressor]]
    """List of document filters that are chained together and run in sequence."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
[docs]    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Transform a list of documents."""
        for _transformer in self.transformers:
            if isinstance(_transformer, BaseDocumentCompressor):
                accepts_callbacks = (
                    signature(_transformer.compress_documents).parameters.get(


accepts_callbacks = (
                    signature(_transformer.compress_documents).parameters.get(
                        "callbacks"
                    )
                    is not None
                )
                if accepts_callbacks:
                    documents = _transformer.compress_documents(
                        documents, query, callbacks=callbacks
                    )
                else:
                    documents = _transformer.compress_documents(documents, query)
            elif isinstance(_transformer, BaseDocumentTransformer):
                documents = _transformer.transform_documents(documents)
            else:
                raise ValueError(f"Got unexpected transformer type: {_transformer}")
        return documents
[docs]    async def acompress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Compress retrieved documents given the query context."""
        for _transformer in self.transformers:
            if isinstance(_transformer, BaseDocumentCompressor):
                accepts_callbacks = (
                    signature(_transformer.acompress_documents).parameters.get(
                        "callbacks"
                    )
                    is not None
                )
                if accepts_callbacks:
                    documents = await _transformer.acompress_documents(
                        documents, query, callbacks=callbacks
                    )
                else:
                    documents = await _transformer.acompress_documents(documents, query)
            elif isinstance(_transformer, BaseDocumentTransformer):
                documents = await _transformer.atransform_documents(documents)
            else:
                raise ValueError(f"Got unexpected transformer type: {_transformer}")
        return documents


Source code for langchain.retrievers.document_compressors.chain_extract
"""DocumentFilter that uses an LLM chain to extract the relevant parts of documents."""
from __future__ import annotations
import asyncio
from typing import Any, Callable, Dict, Optional, Sequence
from langchain.callbacks.manager import Callbacks
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
from langchain.retrievers.document_compressors.chain_extract_prompt import (
    prompt_template,
)
from langchain.schema import BaseOutputParser, Document
from langchain.schema.language_model import BaseLanguageModel
[docs]def default_get_input(query: str, doc: Document) -> Dict[str, Any]:
    """Return the compression chain input."""
    return {"question": query, "context": doc.page_content}
[docs]class NoOutputParser(BaseOutputParser[str]):
    """Parse outputs that could return a null string of some sort."""
    no_output_str: str = "NO_OUTPUT"
[docs]    def parse(self, text: str) -> str:
        cleaned_text = text.strip()
        if cleaned_text == self.no_output_str:
            return ""
        return cleaned_text
def _get_default_chain_prompt() -> PromptTemplate:
    output_parser = NoOutputParser()
    template = prompt_template.format(no_output_str=output_parser.no_output_str)
    return PromptTemplate(
        template=template,
        input_variables=["question", "context"],
        output_parser=output_parser,
    )
[docs]class LLMChainExtractor(BaseDocumentCompressor):
    """Document compressor that uses an LLM chain to extract
    the relevant parts of documents."""
    llm_chain: LLMChain
    """LLM wrapper to use for compressing documents."""


"""LLM wrapper to use for compressing documents."""
    get_input: Callable[[str, Document], dict] = default_get_input
    """Callable for constructing the chain input from the query and a Document."""
[docs]    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Compress page content of raw documents."""
        compressed_docs = []
        for doc in documents:
            _input = self.get_input(query, doc)
            output = self.llm_chain.predict_and_parse(**_input, callbacks=callbacks)
            if len(output) == 0:
                continue
            compressed_docs.append(Document(page_content=output, metadata=doc.metadata))
        return compressed_docs
[docs]    async def acompress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Compress page content of raw documents asynchronously."""
        outputs = await asyncio.gather(
            *[
                self.llm_chain.apredict_and_parse(
                    **self.get_input(query, doc), callbacks=callbacks
                )
                for doc in documents
            ]
        )
        compressed_docs = []
        for i, doc in enumerate(documents):
            if len(outputs[i]) == 0:
                continue
            compressed_docs.append(
                Document(page_content=outputs[i], metadata=doc.metadata)
            )
        return compressed_docs
[docs]    @classmethod
    def from_llm(
        cls,
        llm: BaseLanguageModel,
        prompt: Optional[PromptTemplate] = None,


prompt: Optional[PromptTemplate] = None,
        get_input: Optional[Callable[[str, Document], str]] = None,
        llm_chain_kwargs: Optional[dict] = None,
    ) -> LLMChainExtractor:
        """Initialize from LLM."""
        _prompt = prompt if prompt is not None else _get_default_chain_prompt()
        _get_input = get_input if get_input is not None else default_get_input
        llm_chain = LLMChain(llm=llm, prompt=_prompt, **(llm_chain_kwargs or {}))
        return cls(llm_chain=llm_chain, get_input=_get_input)


Source code for langchain.retrievers.document_compressors.embeddings_filter
from typing import Callable, Dict, Optional, Sequence
import numpy as np
from langchain.callbacks.manager import Callbacks
from langchain.document_transformers.embeddings_redundant_filter import (
    _get_embeddings_from_stateful_docs,
    get_stateful_documents,
)
from langchain.pydantic_v1 import root_validator
from langchain.retrievers.document_compressors.base import (
    BaseDocumentCompressor,
)
from langchain.schema import Document
from langchain.schema.embeddings import Embeddings
from langchain.utils.math import cosine_similarity
[docs]class EmbeddingsFilter(BaseDocumentCompressor):
    """Document compressor that uses embeddings to drop documents
    unrelated to the query."""
    embeddings: Embeddings
    """Embeddings to use for embedding document contents and queries."""
    similarity_fn: Callable = cosine_similarity
    """Similarity function for comparing documents. Function expected to take as input
    two matrices (List[List[float]]) and return a matrix of scores where higher values
    indicate greater similarity."""
    k: Optional[int] = 20
    """The number of relevant documents to return. Can be set to None, in which case
    `similarity_threshold` must be specified. Defaults to 20."""
    similarity_threshold: Optional[float]
    """Threshold for determining when two documents are similar enough
    to be considered redundant. Defaults to None, must be specified if `k` is set
    to None."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
    @root_validator()
    def validate_params(cls, values: Dict) -> Dict:
        """Validate similarity parameters."""
        if values["k"] is None and values["similarity_threshold"] is None:


if values["k"] is None and values["similarity_threshold"] is None:
            raise ValueError("Must specify one of `k` or `similarity_threshold`.")
        return values
[docs]    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Filter documents based on similarity of their embeddings to the query."""
        stateful_documents = get_stateful_documents(documents)
        embedded_documents = _get_embeddings_from_stateful_docs(
            self.embeddings, stateful_documents
        )
        embedded_query = self.embeddings.embed_query(query)
        similarity = self.similarity_fn([embedded_query], embedded_documents)[0]
        included_idxs = np.arange(len(embedded_documents))
        if self.k is not None:
            included_idxs = np.argsort(similarity)[::-1][: self.k]
        if self.similarity_threshold is not None:
            similar_enough = np.where(
                similarity[included_idxs] > self.similarity_threshold
            )
            included_idxs = included_idxs[similar_enough]
        return [stateful_documents[i] for i in included_idxs]
[docs]    async def acompress_documents(
        self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        """Filter down documents."""
        raise NotImplementedError()


Source code for langchain.document_transformers.doctran_text_qa
from typing import Any, Optional, Sequence
from langchain.schema import BaseDocumentTransformer, Document
from langchain.utils import get_from_env
[docs]class DoctranQATransformer(BaseDocumentTransformer):
    """Extract QA from text documents using doctran.
    Arguments:
        openai_api_key: OpenAI API key. Can also be specified via environment variable
            ``OPENAI_API_KEY``.
    Example:
        .. code-block:: python
            from langchain.document_transformers import DoctranQATransformer
            # Pass in openai_api_key or set env var OPENAI_API_KEY
            qa_transformer = DoctranQATransformer()
            transformed_document = await qa_transformer.atransform_documents(documents)
    """
[docs]    def __init__(
        self,
        openai_api_key: Optional[str] = None,
        openai_api_model: Optional[str] = None,
    ) -> None:
        self.openai_api_key = openai_api_key or get_from_env(
            "openai_api_key", "OPENAI_API_KEY"
        )
        self.openai_api_model = openai_api_model or get_from_env(
            "openai_api_model", "OPENAI_API_MODEL"
        )
[docs]    def transform_documents(
        self, documents: Sequence[Document], **kwargs: Any
    ) -> Sequence[Document]:
        raise NotImplementedError
[docs]    async def atransform_documents(
        self, documents: Sequence[Document], **kwargs: Any
    ) -> Sequence[Document]:
        """Extracts QA from text documents using doctran."""
        try:
            from doctran import Doctran
            doctran = Doctran(


from doctran import Doctran
            doctran = Doctran(
                openai_api_key=self.openai_api_key, openai_model=self.openai_api_model
            )
        except ImportError:
            raise ImportError(
                "Install doctran to use this parser. (pip install doctran)"
            )
        for d in documents:
            doctran_doc = (
                await doctran.parse(content=d.page_content).interrogate().execute()
            )
            questions_and_answers = doctran_doc.extracted_properties.get(
                "questions_and_answers"
            )
            d.metadata["questions_and_answers"] = questions_and_answers
        return documents


Source code for langchain.document_transformers.embeddings_redundant_filter
"""Transform documents"""
from typing import Any, Callable, List, Sequence
import numpy as np
from langchain.pydantic_v1 import BaseModel, Field
from langchain.schema import BaseDocumentTransformer, Document
from langchain.schema.embeddings import Embeddings
from langchain.utils.math import cosine_similarity
class _DocumentWithState(Document):
    """Wrapper for a document that includes arbitrary state."""
    state: dict = Field(default_factory=dict)
    """State associated with the document."""
    def to_document(self) -> Document:
        """Convert the DocumentWithState to a Document."""
        return Document(page_content=self.page_content, metadata=self.metadata)
    @classmethod
    def from_document(cls, doc: Document) -> "_DocumentWithState":
        """Create a DocumentWithState from a Document."""
        if isinstance(doc, cls):
            return doc
        return cls(page_content=doc.page_content, metadata=doc.metadata)
[docs]def get_stateful_documents(
    documents: Sequence[Document],
) -> Sequence[_DocumentWithState]:
    """Convert a list of documents to a list of documents with state.
    Args:
        documents: The documents to convert.
    Returns:
        A list of documents with state.
    """
    return [_DocumentWithState.from_document(doc) for doc in documents]
def _filter_similar_embeddings(
    embedded_documents: List[List[float]], similarity_fn: Callable, threshold: float
) -> List[int]:
    """Filter redundant documents based on the similarity of their embeddings."""
    similarity = np.tril(similarity_fn(embedded_documents, embedded_documents), k=-1)
    redundant = np.where(similarity > threshold)
    redundant_stacked = np.column_stack(redundant)


redundant_stacked = np.column_stack(redundant)
    redundant_sorted = np.argsort(similarity[redundant])[::-1]
    included_idxs = set(range(len(embedded_documents)))
    for first_idx, second_idx in redundant_stacked[redundant_sorted]:
        if first_idx in included_idxs and second_idx in included_idxs:
            # Default to dropping the second document of any highly similar pair.
            included_idxs.remove(second_idx)
    return list(sorted(included_idxs))
def _get_embeddings_from_stateful_docs(
    embeddings: Embeddings, documents: Sequence[_DocumentWithState]
) -> List[List[float]]:
    if len(documents) and "embedded_doc" in documents[0].state:
        embedded_documents = [doc.state["embedded_doc"] for doc in documents]
    else:
        embedded_documents = embeddings.embed_documents(
            [d.page_content for d in documents]
        )
        for doc, embedding in zip(documents, embedded_documents):
            doc.state["embedded_doc"] = embedding
    return embedded_documents
def _filter_cluster_embeddings(
    embedded_documents: List[List[float]],
    num_clusters: int,
    num_closest: int,
    random_state: int,
    remove_duplicates: bool,
) -> List[int]:
    """Filter documents based on proximity of their embeddings to clusters."""
    try:
        from sklearn.cluster import KMeans
    except ImportError:
        raise ImportError(
            "sklearn package not found, please install it with "
            "`pip install scikit-learn`"
        )
    kmeans = KMeans(n_clusters=num_clusters, random_state=random_state).fit(
        embedded_documents
    )
    closest_indices = []
    # Loop through the number of clusters you have


)
    closest_indices = []
    # Loop through the number of clusters you have
    for i in range(num_clusters):
        # Get the list of distances from that particular cluster center
        distances = np.linalg.norm(
            embedded_documents - kmeans.cluster_centers_[i], axis=1
        )
        # Find the indices of the two unique closest ones
        # (using argsort to find the smallest 2 distances)
        if remove_duplicates:
            # Only add not duplicated vectors.
            closest_indices_sorted = [
                x
                for x in np.argsort(distances)[:num_closest]
                if x not in closest_indices
            ]
        else:
            # Skip duplicates and add the next closest vector.
            closest_indices_sorted = [
                x for x in np.argsort(distances) if x not in closest_indices
            ][:num_closest]
        # Append that position closest indices list
        closest_indices.extend(closest_indices_sorted)
    return closest_indices
[docs]class EmbeddingsRedundantFilter(BaseDocumentTransformer, BaseModel):
    """Filter that drops redundant documents by comparing their embeddings."""
    embeddings: Embeddings
    """Embeddings to use for embedding document contents."""
    similarity_fn: Callable = cosine_similarity
    """Similarity function for comparing documents. Function expected to take as input
    two matrices (List[List[float]]) and return a matrix of scores where higher values
    indicate greater similarity."""
    similarity_threshold: float = 0.95
    """Threshold for determining when two documents are similar enough
    to be considered redundant."""
    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
[docs]    def transform_documents(
        self, documents: Sequence[Document], **kwargs: Any
    ) -> Sequence[Document]: