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


if (
                isinstance(value, (str, int, float, bool))
                and name != self.content_field
            ):
                lc_doc.metadata[name] = value
        return lc_doc
    async def _aget_relevant_documents(
        self,
        query: str,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        raise NotImplementedError


Source code for langchain.retrievers.azure_cognitive_search
"""Retriever wrapper for Azure Cognitive Search."""
from __future__ import annotations
import json
from typing import Any, Dict, List, Optional
import aiohttp
import requests
from pydantic import BaseModel, Extra, root_validator
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
from langchain.utils import get_from_dict_or_env
[docs]class AzureCognitiveSearchRetriever(BaseRetriever, BaseModel):
    """Wrapper around Azure Cognitive Search."""
    service_name: str = ""
    """Name of Azure Cognitive Search service"""
    index_name: str = ""
    """Name of Index inside Azure Cognitive Search service"""
    api_key: str = ""
    """API Key. Both Admin and Query keys work, but for reading data it's
    recommended to use a Query key."""
    api_version: str = "2020-06-30"
    """API version"""
    aiosession: Optional[aiohttp.ClientSession] = None
    """ClientSession, in case we want to reuse connection for better performance."""
    content_key: str = "content"
    """Key in a retrieved result to set as the Document page_content."""
[docs]    class Config:
        extra = Extra.forbid
        arbitrary_types_allowed = True
[docs]    @root_validator(pre=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that service name, index name and api key exists in environment."""
        values["service_name"] = get_from_dict_or_env(
            values, "service_name", "AZURE_COGNITIVE_SEARCH_SERVICE_NAME"
        )


)
        values["index_name"] = get_from_dict_or_env(
            values, "index_name", "AZURE_COGNITIVE_SEARCH_INDEX_NAME"
        )
        values["api_key"] = get_from_dict_or_env(
            values, "api_key", "AZURE_COGNITIVE_SEARCH_API_KEY"
        )
        return values
    def _build_search_url(self, query: str) -> str:
        base_url = f"https://{self.service_name}.search.windows.net/"
        endpoint_path = f"indexes/{self.index_name}/docs?api-version={self.api_version}"
        return base_url + endpoint_path + f"&search={query}"
    @property
    def _headers(self) -> Dict[str, str]:
        return {
            "Content-Type": "application/json",
            "api-key": self.api_key,
        }
    def _search(self, query: str) -> List[dict]:
        search_url = self._build_search_url(query)
        response = requests.get(search_url, headers=self._headers)
        if response.status_code != 200:
            raise Exception(f"Error in search request: {response}")
        return json.loads(response.text)["value"]
    async def _asearch(self, query: str) -> List[dict]:
        search_url = self._build_search_url(query)
        if not self.aiosession:
            async with aiohttp.ClientSession() as session:
                async with session.get(search_url, headers=self._headers) as response:
                    response_json = await response.json()
        else:
            async with self.aiosession.get(
                search_url, headers=self._headers
            ) as response:
                response_json = await response.json()
        return response_json["value"]
    def _get_relevant_documents(


return response_json["value"]
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        search_results = self._search(query)
        return [
            Document(page_content=result.pop(self.content_key), metadata=result)
            for result in search_results
        ]
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        search_results = await self._asearch(query)
        return [
            Document(page_content=result.pop(self.content_key), metadata=result)
            for result in search_results
        ]


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


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


self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        raise NotImplementedError("LlamaIndexGraphRetriever does not support async")


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


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


retriever_docs = [
            await retriever.aget_relevant_documents(
                query, callbacks=run_manager.get_child("retriever_{}".format(i + 1))
            )
            for i, retriever in enumerate(self.retrievers)
        ]
        # Merge the results of the retrievers.
        merged_documents = []
        max_docs = max(len(docs) for docs in retriever_docs)
        for i in range(max_docs):
            for retriever, doc in zip(self.retrievers, retriever_docs):
                if i < len(doc):
                    merged_documents.append(doc[i])
        return merged_documents


Source code for langchain.retrievers.zep
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
if TYPE_CHECKING:
    from zep_python import MemorySearchResult
[docs]class ZepRetriever(BaseRetriever):
    """A Retriever implementation for the Zep long-term memory store. Search your
    user's long-term chat history with Zep.
    Note: You will need to provide the user's `session_id` to use this retriever.
    More on Zep:
    Zep provides long-term conversation storage for LLM apps. The server stores,
    summarizes, embeds, indexes, and enriches conversational AI chat
    histories, and exposes them via simple, low-latency APIs.
    For server installation instructions, see:
    https://docs.getzep.com/deployment/quickstart/
    """
    def __init__(
        self,
        session_id: str,
        url: str,
        api_key: Optional[str] = None,
        top_k: Optional[int] = None,
    ):
        try:
            from zep_python import ZepClient
        except ImportError:
            raise ValueError(
                "Could not import zep-python package. "
                "Please install it with `pip install zep-python`."
            )
        self.zep_client = ZepClient(base_url=url, api_key=api_key)
        self.session_id = session_id
        self.top_k = top_k
    def _search_result_to_doc(
        self, results: List[MemorySearchResult]


self, results: List[MemorySearchResult]
    ) -> List[Document]:
        return [
            Document(
                page_content=r.message.pop("content"),
                metadata={"score": r.dist, **r.message},
            )
            for r in results
            if r.message
        ]
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        metadata: Optional[Dict] = None,
        **kwargs: Any,
    ) -> List[Document]:
        from zep_python import MemorySearchPayload
        payload: MemorySearchPayload = MemorySearchPayload(
            text=query, metadata=metadata
        )
        results: List[MemorySearchResult] = self.zep_client.search_memory(
            self.session_id, payload, limit=self.top_k
        )
        return self._search_result_to_doc(results)
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        metadata: Optional[Dict] = None,
        **kwargs: Any,
    ) -> List[Document]:
        from zep_python import MemorySearchPayload
        payload: MemorySearchPayload = MemorySearchPayload(
            text=query, metadata=metadata
        )
        results: List[MemorySearchResult] = await self.zep_client.asearch_memory(
            self.session_id, payload, limit=self.top_k
        )
        return self._search_result_to_doc(results)


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


default_salience: Optional[float] = None
    """The salience to assign memories not retrieved from the vector store.
    None assigns no salience to documents not fetched from the vector store.
    """
[docs]    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
    def _get_combined_score(
        self,
        document: Document,
        vector_relevance: Optional[float],
        current_time: datetime.datetime,
    ) -> float:
        """Return the combined score for a document."""
        hours_passed = _get_hours_passed(
            current_time,
            document.metadata["last_accessed_at"],
        )
        score = (1.0 - self.decay_rate) ** hours_passed
        for key in self.other_score_keys:
            if key in document.metadata:
                score += document.metadata[key]
        if vector_relevance is not None:
            score += vector_relevance
        return score
[docs]    def get_salient_docs(self, query: str) -> Dict[int, Tuple[Document, float]]:
        """Return documents that are salient to the query."""
        docs_and_scores: List[Tuple[Document, float]]
        docs_and_scores = self.vectorstore.similarity_search_with_relevance_scores(
            query, **self.search_kwargs
        )
        results = {}
        for fetched_doc, relevance in docs_and_scores:
            if "buffer_idx" in fetched_doc.metadata:
                buffer_idx = fetched_doc.metadata["buffer_idx"]
                doc = self.memory_stream[buffer_idx]
                results[buffer_idx] = (doc, relevance)
        return results
    def _get_relevant_documents(
        self,
        query: str,
        *,


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


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


Source code for langchain.retrievers.vespa_retriever
"""Wrapper for retrieving documents from Vespa."""
from __future__ import annotations
import json
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Sequence, Union
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
if TYPE_CHECKING:
    from vespa.application import Vespa
[docs]class VespaRetriever(BaseRetriever):
    """Retriever that uses the Vespa."""
    def __init__(
        self,
        app: Vespa,
        body: Dict,
        content_field: str,
        metadata_fields: Optional[Sequence[str]] = None,
    ):
        """
        Args:
            app: Vespa client.
            body: query body.
            content_field: result field with document contents.
            metadata_fields: result fields to include in document metadata.
        """
        self._application = app
        self._query_body = body
        self._content_field = content_field
        self._metadata_fields = metadata_fields or ()
    def _query(self, body: Dict) -> List[Document]:
        response = self._application.query(body)
        if not str(response.status_code).startswith("2"):
            raise RuntimeError(
                "Could not retrieve data from Vespa. Error code: {}".format(
                    response.status_code
                )
            )
        root = response.json["root"]
        if "errors" in root:
            raise RuntimeError(json.dumps(root["errors"]))
        docs = []
        for child in response.hits:
            page_content = child["fields"].pop(self._content_field, "")


page_content = child["fields"].pop(self._content_field, "")
            if self._metadata_fields == "*":
                metadata = child["fields"]
            else:
                metadata = {mf: child["fields"].get(mf) for mf in self._metadata_fields}
            metadata["id"] = child["id"]
            docs.append(Document(page_content=page_content, metadata=metadata))
        return docs
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        body = self._query_body.copy()
        body["query"] = query
        return self._query(body)
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        raise NotImplementedError
[docs]    def get_relevant_documents_with_filter(
        self, query: str, *, _filter: Optional[str] = None
    ) -> List[Document]:
        body = self._query_body.copy()
        _filter = f" and {_filter}" if _filter else ""
        body["yql"] = body["yql"] + _filter
        body["query"] = query
        return self._query(body)
[docs]    @classmethod
    def from_params(
        cls,
        url: str,
        content_field: str,
        *,
        k: Optional[int] = None,
        metadata_fields: Union[Sequence[str], Literal["*"]] = (),
        sources: Union[Sequence[str], Literal["*"], None] = None,


sources: Union[Sequence[str], Literal["*"], None] = None,
        _filter: Optional[str] = None,
        yql: Optional[str] = None,
        **kwargs: Any,
    ) -> VespaRetriever:
        """Instantiate retriever from params.
        Args:
            url (str): Vespa app URL.
            content_field (str): Field in results to return as Document page_content.
            k (Optional[int]): Number of Documents to return. Defaults to None.
            metadata_fields(Sequence[str] or "*"): Fields in results to include in
                document metadata. Defaults to empty tuple ().
            sources (Sequence[str] or "*" or None): Sources to retrieve
                from. Defaults to None.
            _filter (Optional[str]): Document filter condition expressed in YQL.
                Defaults to None.
            yql (Optional[str]): Full YQL query to be used. Should not be specified
                if _filter or sources are specified. Defaults to None.
            kwargs (Any): Keyword arguments added to query body.
        """
        try:
            from vespa.application import Vespa
        except ImportError:
            raise ImportError(
                "pyvespa is not installed, please install with `pip install pyvespa`"
            )
        app = Vespa(url)
        body = kwargs.copy()
        if yql and (sources or _filter):
            raise ValueError(
                "yql should only be specified if both sources and _filter are not "
                "specified."
            )
        else:
            if metadata_fields == "*":
                _fields = "*"
                body["summary"] = "short"
            else:
                _fields = ", ".join([content_field] + list(metadata_fields or []))


_fields = ", ".join([content_field] + list(metadata_fields or []))
            _sources = ", ".join(sources) if isinstance(sources, Sequence) else "*"
            _filter = f" and {_filter}" if _filter else ""
            yql = f"select {_fields} from sources {_sources} where userQuery(){_filter}"
        body["yql"] = yql
        if k:
            body["hits"] = k
        return cls(app, body, content_field, metadata_fields=metadata_fields)


Source code for langchain.retrievers.remote_retriever
from typing import Any, List, Optional
import aiohttp
import requests
from pydantic import BaseModel
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
[docs]class RemoteLangChainRetriever(BaseRetriever, BaseModel):
    url: str
    headers: Optional[dict] = None
    input_key: str = "message"
    response_key: str = "response"
    page_content_key: str = "page_content"
    metadata_key: str = "metadata"
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        response = requests.post(
            self.url, json={self.input_key: query}, headers=self.headers
        )
        result = response.json()
        return [
            Document(
                page_content=r[self.page_content_key], metadata=r[self.metadata_key]
            )
            for r in result[self.response_key]
        ]
    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.url, headers=self.headers, json={self.input_key: query}
            ) as response:
                result = await response.json()
        return [
            Document(


result = await response.json()
        return [
            Document(
                page_content=r[self.page_content_key], metadata=r[self.metadata_key]
            )
            for r in result[self.response_key]
        ]


Source code for langchain.retrievers.pubmed
"""A retriever that uses PubMed API to retrieve documents."""
from typing import Any, List
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
from langchain.utilities.pupmed import PubMedAPIWrapper
[docs]class PubMedRetriever(BaseRetriever, PubMedAPIWrapper):
    """
    It is effectively a wrapper for PubMedAPIWrapper.
    It wraps load() to get_relevant_documents().
    It uses all PubMedAPIWrapper arguments without any change.
    """
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        return self.load_docs(query=query)
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        raise NotImplementedError


Source code for langchain.retrievers.metal
from typing import Any, List, Optional
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from langchain.schema import BaseRetriever, Document
[docs]class MetalRetriever(BaseRetriever):
    """Retriever that uses the Metal API."""
    def __init__(self, client: Any, params: Optional[dict] = None):
        from metal_sdk.metal import Metal
        if not isinstance(client, Metal):
            raise ValueError(
                "Got unexpected client, should be of type metal_sdk.metal.Metal. "
                f"Instead, got {type(client)}"
            )
        self.client: Metal = client
        self.params = params or {}
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> 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
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        raise NotImplementedError


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


"""Given a user query, use an LLM to write a set of queries.
    Retrieve docs for each query. Rake the unique union of all retrieved docs."""
    def __init__(
        self,
        retriever: BaseRetriever,
        llm_chain: LLMChain,
        verbose: bool = True,
        parser_key: str = "lines",
    ) -> None:
        """Initialize MultiQueryRetriever.
        Args:
            retriever: retriever to query documents from
            llm_chain: llm_chain for query generation
            verbose: show the queries that we generated to the user
            parser_key: attribute name for the parsed output
        Returns:
            MultiQueryRetriever
        """
        self.retriever = retriever
        self.llm_chain = llm_chain
        self.verbose = verbose
        self.parser_key = parser_key
[docs]    @classmethod
    def from_llm(
        cls,
        retriever: BaseRetriever,
        llm: BaseLLM,
        prompt: PromptTemplate = DEFAULT_QUERY_PROMPT,
        parser_key: str = "lines",
    ) -> "MultiQueryRetriever":
        """Initialize from llm using default template.
        Args:
            retriever: retriever to query documents from
            llm: llm for query generation using DEFAULT_QUERY_PROMPT
        Returns:
            MultiQueryRetriever
        """
        output_parser = LineListOutputParser()
        llm_chain = LLMChain(llm=llm, prompt=prompt, output_parser=output_parser)
        return cls(
            retriever=retriever,
            llm_chain=llm_chain,
            parser_key=parser_key,
        )


parser_key=parser_key,
        )
    def _get_relevant_documents(
        self,
        query: str,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        """Get relevated documents given a user query.
        Args:
            question: user query
        Returns:
            Unique union of relevant documents from all generated queries
        """
        queries = self.generate_queries(query, run_manager)
        documents = self.retrieve_documents(queries, run_manager)
        unique_documents = self.unique_union(documents)
        return unique_documents
    async def _aget_relevant_documents(
        self,
        query: str,
        run_manager: AsyncCallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        raise NotImplementedError
[docs]    def generate_queries(
        self, question: str, run_manager: CallbackManagerForRetrieverRun
    ) -> List[str]:
        """Generate queries based upon user input.
        Args:
            question: user query
        Returns:
            List of LLM generated queries that are similar to the user input
        """
        response = self.llm_chain(
            {"question": question}, callbacks=run_manager.get_child()
        )
        lines = getattr(response["text"], self.parser_key, [])
        if self.verbose:
            logger.info(f"Generated queries: {lines}")
        return lines
[docs]    def retrieve_documents(
        self, queries: List[str], run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Run all LLM generated queries.
        Args:
            queries: query list
        Returns:
            List of retrived Documents
        """


queries: query list
        Returns:
            List of retrived Documents
        """
        documents = []
        for query in queries:
            docs = self.retriever.get_relevant_documents(
                query, callbacks=run_manager.get_child()
            )
            documents.extend(docs)
        return documents
[docs]    def unique_union(self, documents: List[Document]) -> List[Document]:
        """Get uniqe Documents.
        Args:
            documents: List of retrived Documents
        Returns:
            List of unique retrived Documents
        """
        # Create a dictionary with page_content as keys to remove duplicates
        # TODO: Add Document ID property (e.g., UUID)
        unique_documents_dict = {
            (doc.page_content, tuple(sorted(doc.metadata.items()))): doc
            for doc in documents
        }
        unique_documents = list(unique_documents_dict.values())
        return unique_documents


Source code for langchain.retrievers.self_query.weaviate
"""Logic for converting internal query language to a valid Weaviate query."""
from typing import Dict, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class WeaviateTranslator(Visitor):
    """Logic for converting 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.chroma
"""Logic for converting internal query language to a valid Chroma query."""
from typing import Dict, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class ChromaTranslator(Visitor):
    """Logic for converting internal query language elements to valid filters."""
    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,
)
[docs]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
[docs]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):
    """Logic for converting 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 = {


Comparator.LIKE,
    ]
    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"),
    }
    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 type(value) is 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.qdrant
"""Logic for converting internal query language to a valid Qdrant query."""
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):
    """Logic for converting internal query language elements to valid filters."""
    def __init__(self, metadata_key: str):
        self.metadata_key = metadata_key
[docs]    def visit_operation(self, operation: Operation) -> rest.Filter:
        from qdrant_client.http import models as rest
        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:
        from qdrant_client.http import models as rest
        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:


) -> 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.pinecone
"""Logic for converting internal query language to a valid Pinecone query."""
from typing import Dict, Tuple, Union
from langchain.chains.query_constructor.ir import (
    Comparator,
    Comparison,
    Operation,
    Operator,
    StructuredQuery,
    Visitor,
)
[docs]class PineconeTranslator(Visitor):
    """Logic for converting internal query language elements to valid filters."""
    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.base
"""Retriever that generates and executes structured queries over its own data source."""
from typing import Any, Dict, List, Optional, Type, cast
from pydantic import BaseModel, Field, root_validator
from langchain import LLMChain
from langchain.base_language import BaseLanguageModel
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
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.retrievers.self_query.chroma import ChromaTranslator
from langchain.retrievers.self_query.myscale import MyScaleTranslator
from langchain.retrievers.self_query.pinecone import PineconeTranslator
from langchain.retrievers.self_query.qdrant import QdrantTranslator
from langchain.retrievers.self_query.weaviate import WeaviateTranslator
from langchain.schema import BaseRetriever, Document
from langchain.vectorstores import (
    Chroma,
    MyScale,
    Pinecone,
    Qdrant,
    VectorStore,
    Weaviate,
)
def _get_builtin_translator(vectorstore: VectorStore) -> Visitor:
    """Get the translator class corresponding to the vector store class."""
    vectorstore_cls = vectorstore.__class__
    BUILTIN_TRANSLATORS: Dict[Type[VectorStore], Type[Visitor]] = {
        Pinecone: PineconeTranslator,
        Chroma: ChromaTranslator,
        Weaviate: WeaviateTranslator,
        Qdrant: QdrantTranslator,
        MyScale: MyScaleTranslator,
    }


MyScale: MyScaleTranslator,
    }
    if vectorstore_cls not in BUILTIN_TRANSLATORS:
        raise ValueError(
            f"Self query retriever with Vector Store type {vectorstore_cls}"
            f" not supported."
        )
    if isinstance(vectorstore, Qdrant):
        return QdrantTranslator(metadata_key=vectorstore.metadata_payload_key)
    elif isinstance(vectorstore, MyScale):
        return MyScaleTranslator(metadata_key=vectorstore.metadata_column)
    return BUILTIN_TRANSLATORS[vectorstore_cls]()
[docs]class SelfQueryRetriever(BaseRetriever, BaseModel):
    """Retriever that wraps around a vector store and uses 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 instead of the revised new query from LLM"""
    use_original_query: bool = False
[docs]    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
[docs]    @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["structured_query_translator"] = _get_builtin_translator(
                values["vectorstore"]
            )
        return values
    def _get_relevant_documents(
        self,
        query: str,
        *,
        run_manager: CallbackManagerForRetrieverRun,
        **kwargs: Any,
    ) -> List[Document]:
        """Get documents relevant for a query.
        Args:
            query: string to find relevant documents for
        Returns:
            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
    async def _aget_relevant_documents(
        self,
        query: str,
        *,
        run_manager: Optional[AsyncCallbackManagerForRetrieverRun],
        **kwargs: Any,
    ) -> List[Document]:
        raise NotImplementedError
[docs]    @classmethod
    def from_llm(
        cls,
        llm: BaseLanguageModel,
        vectorstore: VectorStore,
        document_contents: str,
        metadata_field_info: List[AttributeInfo],


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"
            ] = 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.document_compressors.cohere_rerank
from __future__ import annotations
from typing import TYPE_CHECKING, Dict, Optional, Sequence
from pydantic import Extra, root_validator
from langchain.callbacks.manager import Callbacks
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):
    client: Client
    top_n: int = 3
    model: str = "rerank-english-v2.0"
[docs]    class Config:
        """Configuration for this pydantic object."""
        extra = Extra.forbid
        arbitrary_types_allowed = True
[docs]    @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. "
                "Please install it with `pip install cohere`."
            )
        return values
[docs]    def compress_documents(
        self,
        documents: Sequence[Document],
        query: str,


self,
        documents: Sequence[Document],
        query: str,
        callbacks: Optional[Callbacks] = None,
    ) -> Sequence[Document]:
        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 import BasePromptTemplate, LLMChain, PromptTemplate
from langchain.base_language import BaseLanguageModel
from langchain.callbacks.manager import Callbacks
from langchain.output_parsers.boolean import BooleanOutputParser
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
from langchain.retrievers.document_compressors.chain_filter_prompt import (
    prompt_template,
)
from langchain.schema import Document
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."""
        filtered_docs = []
        for doc in documents:


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":
        _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.embeddings_filter
"""Document compressor that uses embeddings to drop documents unrelated to the query."""
from typing import Callable, Dict, Optional, Sequence
import numpy as np
from pydantic import root_validator
from langchain.callbacks.manager import Callbacks
from langchain.document_transformers import (
    _get_embeddings_from_stateful_docs,
    get_stateful_documents,
)
from langchain.embeddings.base import Embeddings
from langchain.math_utils import cosine_similarity
from langchain.retrievers.document_compressors.base import (
    BaseDocumentCompressor,
)
from langchain.schema import Document
[docs]class EmbeddingsFilter(BaseDocumentCompressor):
    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."""
[docs]    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
[docs]    @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.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 import LLMChain, PromptTemplate
from langchain.base_language import BaseLanguageModel
from langchain.callbacks.manager import Callbacks
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
[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):
    llm_chain: LLMChain
    """LLM wrapper to use for compressing documents."""
    get_input: Callable[[str, Document], dict] = default_get_input


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,
        get_input: Optional[Callable[[str, Document], str]] = 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.base
"""Interface for retrieved document compressors."""
from abc import ABC, abstractmethod
from inspect import signature
from typing import List, Optional, Sequence, Union
from pydantic import BaseModel
from langchain.callbacks.manager import Callbacks
from langchain.schema import BaseDocumentTransformer, Document
[docs]class BaseDocumentCompressor(BaseModel, ABC):
    """Base abstraction interface for document compression."""
[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."""
[docs]    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 = (


if isinstance(_transformer, BaseDocumentCompressor):
                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.utilities.wolfram_alpha
"""Util that calls WolframAlpha."""
from typing import Any, Dict, Optional
from pydantic import BaseModel, Extra, root_validator
from langchain.utils import get_from_dict_or_env
[docs]class WolframAlphaAPIWrapper(BaseModel):
    """Wrapper for Wolfram Alpha.
    Docs for using:
    1. Go to wolfram alpha and sign up for a developer account
    2. Create an app and get your APP ID
    3. Save your APP ID into WOLFRAM_ALPHA_APPID env variable
    4. pip install wolframalpha
    """
    wolfram_client: Any  #: :meta private:
    wolfram_alpha_appid: Optional[str] = None
[docs]    class Config:
        """Configuration for this pydantic object."""
        extra = Extra.forbid
[docs]    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that api key and python package exists in environment."""
        wolfram_alpha_appid = get_from_dict_or_env(
            values, "wolfram_alpha_appid", "WOLFRAM_ALPHA_APPID"
        )
        values["wolfram_alpha_appid"] = wolfram_alpha_appid
        try:
            import wolframalpha
        except ImportError:
            raise ImportError(
                "wolframalpha is not installed. "
                "Please install it with `pip install wolframalpha`"
            )
        client = wolframalpha.Client(wolfram_alpha_appid)
        values["wolfram_client"] = client
        return values
[docs]    def run(self, query: str) -> str:
        """Run query through WolframAlpha and parse result."""


"""Run query through WolframAlpha and parse result."""
        res = self.wolfram_client.query(query)
        try:
            assumption = next(res.pods).text
            answer = next(res.results).text
        except StopIteration:
            return "Wolfram Alpha wasn't able to answer it"
        if answer is None or answer == "":
            # We don't want to return the assumption alone if answer is empty
            return "No good Wolfram Alpha Result was found"
        else:
            return f"Assumption: {assumption} \nAnswer: {answer}"


Source code for langchain.utilities.bing_search
"""Util that calls Bing Search.
In order to set this up, follow instructions at:
https://levelup.gitconnected.com/api-tutorial-how-to-use-bing-web-search-api-in-python-4165d5592a7e
"""
from typing import Dict, List
import requests
from pydantic import BaseModel, Extra, root_validator
from langchain.utils import get_from_dict_or_env
[docs]class BingSearchAPIWrapper(BaseModel):
    """Wrapper for Bing Search API.
    In order to set this up, follow instructions at:
    https://levelup.gitconnected.com/api-tutorial-how-to-use-bing-web-search-api-in-python-4165d5592a7e
    """
    bing_subscription_key: str
    bing_search_url: str
    k: int = 10
[docs]    class Config:
        """Configuration for this pydantic object."""
        extra = Extra.forbid
    def _bing_search_results(self, search_term: str, count: int) -> List[dict]:
        headers = {"Ocp-Apim-Subscription-Key": self.bing_subscription_key}
        params = {
            "q": search_term,
            "count": count,
            "textDecorations": True,
            "textFormat": "HTML",
        }
        response = requests.get(
            self.bing_search_url, headers=headers, params=params  # type: ignore
        )
        response.raise_for_status()
        search_results = response.json()
        return search_results["webPages"]["value"]
[docs]    @root_validator(pre=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that api key and endpoint exists in environment."""


"""Validate that api key and endpoint exists in environment."""
        bing_subscription_key = get_from_dict_or_env(
            values, "bing_subscription_key", "BING_SUBSCRIPTION_KEY"
        )
        values["bing_subscription_key"] = bing_subscription_key
        bing_search_url = get_from_dict_or_env(
            values,
            "bing_search_url",
            "BING_SEARCH_URL",
            # default="https://api.bing.microsoft.com/v7.0/search",
        )
        values["bing_search_url"] = bing_search_url
        return values
[docs]    def run(self, query: str) -> str:
        """Run query through BingSearch and parse result."""
        snippets = []
        results = self._bing_search_results(query, count=self.k)
        if len(results) == 0:
            return "No good Bing Search Result was found"
        for result in results:
            snippets.append(result["snippet"])
        return " ".join(snippets)
[docs]    def results(self, query: str, num_results: int) -> List[Dict]:
        """Run query through BingSearch and return metadata.
        Args:
            query: The query to search for.
            num_results: The number of results to return.
        Returns:
            A list of dictionaries with the following keys:
                snippet - The description of the result.
                title - The title of the result.
                link - The link to the result.
        """
        metadata_results = []
        results = self._bing_search_results(query, count=num_results)
        if len(results) == 0:
            return [{"Result": "No good Bing Search Result was found"}]
        for result in results:
            metadata_result = {
                "snippet": result["snippet"],


metadata_result = {
                "snippet": result["snippet"],
                "title": result["name"],
                "link": result["url"],
            }
            metadata_results.append(metadata_result)
        return metadata_results


Source code for langchain.utilities.zapier
"""Util that can interact with Zapier NLA.
Full docs here: https://nla.zapier.com/start/
Note: this wrapper currently only implemented the `api_key` auth method for testing
and server-side production use cases (using the developer's connected accounts on
Zapier.com)
For use-cases where LangChain + Zapier NLA is powering a user-facing application, and
LangChain needs access to the end-user's connected accounts on Zapier.com, you'll need
to use oauth. Review the full docs above and reach out to nla@zapier.com for
developer support.
"""
import json
from typing import Any, Dict, List, Optional
import aiohttp
import requests
from pydantic import BaseModel, Extra, root_validator
from requests import Request, Session
from langchain.utils import get_from_dict_or_env
[docs]class ZapierNLAWrapper(BaseModel):
    """Wrapper for Zapier NLA.
    Full docs here: https://nla.zapier.com/start/
    This wrapper supports both API Key and OAuth Credential auth methods. API Key
    is the fastest way to get started using this wrapper.
    Call this wrapper with either `zapier_nla_api_key` or
    `zapier_nla_oauth_access_token` arguments, or set the `ZAPIER_NLA_API_KEY`
    environment variable. If both arguments are set, the Access Token will take
    precedence.
    For use-cases where LangChain + Zapier NLA is powering a user-facing application,
    and LangChain needs access to the end-user's connected accounts on Zapier.com,
    you'll need to use OAuth. Review the full docs above to learn how to create
    your own provider and generate credentials.
    """


your own provider and generate credentials.
    """
    zapier_nla_api_key: str
    zapier_nla_oauth_access_token: str
    zapier_nla_api_base: str = "https://nla.zapier.com/api/v1/"
[docs]    class Config:
        """Configuration for this pydantic object."""
        extra = Extra.forbid
    def _format_headers(self) -> Dict[str, str]:
        """Format headers for requests."""
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
        }
        if self.zapier_nla_oauth_access_token:
            headers.update(
                {"Authorization": f"Bearer {self.zapier_nla_oauth_access_token}"}
            )
        else:
            headers.update({"X-API-Key": self.zapier_nla_api_key})
        return headers
    def _get_session(self) -> Session:
        session = requests.Session()
        session.headers.update(self._format_headers())
        return session
    async def _arequest(self, method: str, url: str, **kwargs: Any) -> Dict[str, Any]:
        """Make an async request."""
        async with aiohttp.ClientSession(headers=self._format_headers()) as session:
            async with session.request(method, url, **kwargs) as response:
                response.raise_for_status()
                return await response.json()
    def _create_action_payload(  # type: ignore[no-untyped-def]
        self, instructions: str, params: Optional[Dict] = None, preview_only=False
    ) -> Dict:
        """Create a payload for an action."""
        data = params if params else {}
        data.update(
            {
                "instructions": instructions,
            }
        )


{
                "instructions": instructions,
            }
        )
        if preview_only:
            data.update({"preview_only": True})
        return data
    def _create_action_url(self, action_id: str) -> str:
        """Create a url for an action."""
        return self.zapier_nla_api_base + f"exposed/{action_id}/execute/"
    def _create_action_request(  # type: ignore[no-untyped-def]
        self,
        action_id: str,
        instructions: str,
        params: Optional[Dict] = None,
        preview_only=False,
    ) -> Request:
        data = self._create_action_payload(instructions, params, preview_only)
        return Request(
            "POST",
            self._create_action_url(action_id),
            json=data,
        )
[docs]    @root_validator(pre=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate that api key exists in environment."""
        zapier_nla_api_key_default = None
        # If there is a oauth_access_key passed in the values
        # we don't need a nla_api_key it can be blank
        if "zapier_nla_oauth_access_token" in values:
            zapier_nla_api_key_default = ""
        else:
            values["zapier_nla_oauth_access_token"] = ""
        # we require at least one API Key
        zapier_nla_api_key = get_from_dict_or_env(
            values,
            "zapier_nla_api_key",
            "ZAPIER_NLA_API_KEY",
            zapier_nla_api_key_default,
        )
        values["zapier_nla_api_key"] = zapier_nla_api_key
        return values


return values
[docs]    async def alist(self) -> List[Dict]:
        """Returns a list of all exposed (enabled) actions associated with
        current user (associated with the set api_key). Change your exposed
        actions here: https://nla.zapier.com/demo/start/
        The return list can be empty if no actions exposed. Else will contain
        a list of action objects:
        [{
            "id": str,
            "description": str,
            "params": Dict[str, str]
        }]
        `params` will always contain an `instructions` key, the only required
        param. All others optional and if provided will override any AI guesses
        (see "understanding the AI guessing flow" here:
        https://nla.zapier.com/api/v1/docs)
        """
        response = await self._arequest("GET", self.zapier_nla_api_base + "exposed/")
        return response["results"]
[docs]    def list(self) -> List[Dict]:
        """Returns a list of all exposed (enabled) actions associated with
        current user (associated with the set api_key). Change your exposed
        actions here: https://nla.zapier.com/demo/start/
        The return list can be empty if no actions exposed. Else will contain
        a list of action objects:
        [{
            "id": str,
            "description": str,
            "params": Dict[str, str]
        }]
        `params` will always contain an `instructions` key, the only required
        param. All others optional and if provided will override any AI guesses
        (see "understanding the AI guessing flow" here:
        https://nla.zapier.com/docs/using-the-api#ai-guessing)
        """


"""
        session = self._get_session()
        try:
            response = session.get(self.zapier_nla_api_base + "exposed/")
            response.raise_for_status()
        except requests.HTTPError as http_err:
            if response.status_code == 401:
                if self.zapier_nla_oauth_access_token:
                    raise requests.HTTPError(
                        f"An unauthorized response occurred. Check that your "
                        f"access token is correct and doesn't need to be "
                        f"refreshed. Err: {http_err}"
                    )
                raise requests.HTTPError(
                    f"An unauthorized response occurred. Check that your api "
                    f"key is correct. Err: {http_err}"
                )
            raise http_err
        return response.json()["results"]
[docs]    def run(
        self, action_id: str, instructions: str, params: Optional[Dict] = None
    ) -> Dict:
        """Executes an action that is identified by action_id, must be exposed
        (enabled) by the current user (associated with the set api_key). Change
        your exposed actions here: https://nla.zapier.com/demo/start/
        The return JSON is guaranteed to be less than ~500 words (350
        tokens) making it safe to inject into the prompt of another LLM
        call.
        """
        session = self._get_session()
        request = self._create_action_request(action_id, instructions, params)
        response = session.send(session.prepare_request(request))
        response.raise_for_status()
        return response.json()["result"]
[docs]    async def arun(
        self, action_id: str, instructions: str, params: Optional[Dict] = None
    ) -> Dict:


) -> Dict:
        """Executes an action that is identified by action_id, must be exposed
        (enabled) by the current user (associated with the set api_key). Change
        your exposed actions here: https://nla.zapier.com/demo/start/
        The return JSON is guaranteed to be less than ~500 words (350
        tokens) making it safe to inject into the prompt of another LLM
        call.
        """
        response = await self._arequest(
            "POST",
            self._create_action_url(action_id),
            json=self._create_action_payload(instructions, params),
        )
        return response["result"]
[docs]    def preview(
        self, action_id: str, instructions: str, params: Optional[Dict] = None
    ) -> Dict:
        """Same as run, but instead of actually executing the action, will
        instead return a preview of params that have been guessed by the AI in
        case you need to explicitly review before executing."""
        session = self._get_session()
        params = params if params else {}
        params.update({"preview_only": True})
        request = self._create_action_request(action_id, instructions, params, True)
        response = session.send(session.prepare_request(request))
        response.raise_for_status()
        return response.json()["input_params"]
[docs]    async def apreview(
        self, action_id: str, instructions: str, params: Optional[Dict] = None
    ) -> Dict:
        """Same as run, but instead of actually executing the action, will
        instead return a preview of params that have been guessed by the AI in
        case you need to explicitly review before executing."""
        response = await self._arequest(
            "POST",


response = await self._arequest(
            "POST",
            self._create_action_url(action_id),
            json=self._create_action_payload(instructions, params, preview_only=True),
        )
        return response["result"]
[docs]    def run_as_str(self, *args, **kwargs) -> str:  # type: ignore[no-untyped-def]
        """Same as run, but returns a stringified version of the JSON for
        insertting back into an LLM."""
        data = self.run(*args, **kwargs)
        return json.dumps(data)
[docs]    async def arun_as_str(self, *args, **kwargs) -> str:  # type: ignore[no-untyped-def]
        """Same as run, but returns a stringified version of the JSON for
        insertting back into an LLM."""
        data = await self.arun(*args, **kwargs)
        return json.dumps(data)
[docs]    def preview_as_str(self, *args, **kwargs) -> str:  # type: ignore[no-untyped-def]
        """Same as preview, but returns a stringified version of the JSON for
        insertting back into an LLM."""
        data = self.preview(*args, **kwargs)
        return json.dumps(data)
[docs]    async def apreview_as_str(  # type: ignore[no-untyped-def]
        self, *args, **kwargs
    ) -> str:
        """Same as preview, but returns a stringified version of the JSON for
        insertting back into an LLM."""
        data = await self.apreview(*args, **kwargs)
        return json.dumps(data)
[docs]    def list_as_str(self) -> str:  # type: ignore[no-untyped-def]


"""Same as list, but returns a stringified version of the JSON for
        insertting back into an LLM."""
        actions = self.list()
        return json.dumps(actions)
[docs]    async def alist_as_str(self) -> str:  # type: ignore[no-untyped-def]
        """Same as list, but returns a stringified version of the JSON for
        insertting back into an LLM."""
        actions = await self.alist()
        return json.dumps(actions)


Source code for langchain.utilities.apify
from typing import Any, Callable, Dict, Optional
from pydantic import BaseModel, root_validator
from langchain.document_loaders import ApifyDatasetLoader
from langchain.document_loaders.base import Document
from langchain.utils import get_from_dict_or_env
[docs]class ApifyWrapper(BaseModel):
    """Wrapper around Apify.
    To use, you should have the ``apify-client`` python package installed,
    and the environment variable ``APIFY_API_TOKEN`` set with your API key, or pass
    `apify_api_token` as a named parameter to the constructor.
    """
    apify_client: Any
    apify_client_async: Any
[docs]    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """Validate environment.
        Validate that an Apify API token is set and the apify-client
        Python package exists in the current environment.
        """
        apify_api_token = get_from_dict_or_env(
            values, "apify_api_token", "APIFY_API_TOKEN"
        )
        try:
            from apify_client import ApifyClient, ApifyClientAsync
            values["apify_client"] = ApifyClient(apify_api_token)
            values["apify_client_async"] = ApifyClientAsync(apify_api_token)
        except ImportError:
            raise ValueError(
                "Could not import apify-client Python package. "
                "Please install it with `pip install apify-client`."
            )
        return values
[docs]    def call_actor(
        self,
        actor_id: str,
        run_input: Dict,
        dataset_mapping_function: Callable[[Dict], Document],
        *,
        build: Optional[str] = None,


*,
        build: Optional[str] = None,
        memory_mbytes: Optional[int] = None,
        timeout_secs: Optional[int] = None,
    ) -> ApifyDatasetLoader:
        """Run an Actor on the Apify platform and wait for results to be ready.
        Args:
            actor_id (str): The ID or name of the Actor on the Apify platform.
            run_input (Dict): The input object of the Actor that you're trying to run.
            dataset_mapping_function (Callable): A function that takes a single
                dictionary (an Apify dataset item) and converts it to an
                instance of the Document class.
            build (str, optional): Optionally specifies the actor build to run.
                It can be either a build tag or build number.
            memory_mbytes (int, optional): Optional memory limit for the run,
                in megabytes.
            timeout_secs (int, optional): Optional timeout for the run, in seconds.
        Returns:
            ApifyDatasetLoader: A loader that will fetch the records from the
                Actor run's default dataset.
        """
        actor_call = self.apify_client.actor(actor_id).call(
            run_input=run_input,
            build=build,
            memory_mbytes=memory_mbytes,
            timeout_secs=timeout_secs,
        )
        return ApifyDatasetLoader(
            dataset_id=actor_call["defaultDatasetId"],
            dataset_mapping_function=dataset_mapping_function,
        )
[docs]    async def acall_actor(
        self,
        actor_id: str,
        run_input: Dict,
        dataset_mapping_function: Callable[[Dict], Document],
        *,
        build: Optional[str] = None,
        memory_mbytes: Optional[int] = None,


memory_mbytes: Optional[int] = None,
        timeout_secs: Optional[int] = None,
    ) -> ApifyDatasetLoader:
        """Run an Actor on the Apify platform and wait for results to be ready.
        Args:
            actor_id (str): The ID or name of the Actor on the Apify platform.
            run_input (Dict): The input object of the Actor that you're trying to run.
            dataset_mapping_function (Callable): A function that takes a single
                dictionary (an Apify dataset item) and converts it to
                an instance of the Document class.
            build (str, optional): Optionally specifies the actor build to run.
                It can be either a build tag or build number.
            memory_mbytes (int, optional): Optional memory limit for the run,
                in megabytes.
            timeout_secs (int, optional): Optional timeout for the run, in seconds.
        Returns:
            ApifyDatasetLoader: A loader that will fetch the records from the
                Actor run's default dataset.
        """
        actor_call = await self.apify_client_async.actor(actor_id).call(
            run_input=run_input,
            build=build,
            memory_mbytes=memory_mbytes,
            timeout_secs=timeout_secs,
        )
        return ApifyDatasetLoader(
            dataset_id=actor_call["defaultDatasetId"],
            dataset_mapping_function=dataset_mapping_function,
        )
[docs]    def call_actor_task(
        self,
        task_id: str,
        task_input: Dict,
        dataset_mapping_function: Callable[[Dict], Document],
        *,
        build: Optional[str] = None,
        memory_mbytes: Optional[int] = None,
        timeout_secs: Optional[int] = None,


timeout_secs: Optional[int] = None,
    ) -> ApifyDatasetLoader:
        """Run a saved Actor task on Apify and wait for results to be ready.
        Args:
            task_id (str): The ID or name of the task on the Apify platform.
            task_input (Dict): The input object of the task that you're trying to run.
                Overrides the task's saved input.
            dataset_mapping_function (Callable): A function that takes a single
                dictionary (an Apify dataset item) and converts it to an
                instance of the Document class.
            build (str, optional): Optionally specifies the actor build to run.
                It can be either a build tag or build number.
            memory_mbytes (int, optional): Optional memory limit for the run,
                in megabytes.
            timeout_secs (int, optional): Optional timeout for the run, in seconds.
        Returns:
            ApifyDatasetLoader: A loader that will fetch the records from the
                task run's default dataset.
        """
        task_call = self.apify_client.task(task_id).call(
            task_input=task_input,
            build=build,
            memory_mbytes=memory_mbytes,
            timeout_secs=timeout_secs,
        )
        return ApifyDatasetLoader(
            dataset_id=task_call["defaultDatasetId"],
            dataset_mapping_function=dataset_mapping_function,
        )
[docs]    async def acall_actor_task(
        self,
        task_id: str,
        task_input: Dict,
        dataset_mapping_function: Callable[[Dict], Document],
        *,
        build: Optional[str] = None,
        memory_mbytes: Optional[int] = None,
        timeout_secs: Optional[int] = None,
    ) -> ApifyDatasetLoader:


timeout_secs: Optional[int] = None,
    ) -> ApifyDatasetLoader:
        """Run a saved Actor task on Apify and wait for results to be ready.
        Args:
            task_id (str): The ID or name of the task on the Apify platform.
            task_input (Dict): The input object of the task that you're trying to run.
                Overrides the task's saved input.
            dataset_mapping_function (Callable): A function that takes a single
                dictionary (an Apify dataset item) and converts it to an
                instance of the Document class.
            build (str, optional): Optionally specifies the actor build to run.
                It can be either a build tag or build number.
            memory_mbytes (int, optional): Optional memory limit for the run,
                in megabytes.
            timeout_secs (int, optional): Optional timeout for the run, in seconds.
        Returns:
            ApifyDatasetLoader: A loader that will fetch the records from the
                task run's default dataset.
        """
        task_call = await self.apify_client_async.task(task_id).call(
            task_input=task_input,
            build=build,
            memory_mbytes=memory_mbytes,
            timeout_secs=timeout_secs,
        )
        return ApifyDatasetLoader(
            dataset_id=task_call["defaultDatasetId"],
            dataset_mapping_function=dataset_mapping_function,
        )


Source code for langchain.utilities.powerbi
"""Wrapper around a Power BI endpoint."""
from __future__ import annotations
import asyncio
import logging
import os
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
import aiohttp
import requests
from aiohttp import ServerTimeoutError
from pydantic import BaseModel, Field, root_validator, validator
from requests.exceptions import Timeout
_LOGGER = logging.getLogger(__name__)
BASE_URL = os.getenv("POWERBI_BASE_URL", "https://api.powerbi.com/v1.0/myorg")
if TYPE_CHECKING:
    from azure.core.credentials import TokenCredential
[docs]class PowerBIDataset(BaseModel):
    """Create PowerBI engine from dataset ID and credential or token.
    Use either the credential or a supplied token to authenticate.
    If both are supplied the credential is used to generate a token.
    The impersonated_user_name is the UPN of a user to be impersonated.
    If the model is not RLS enabled, this will be ignored.
    """
    dataset_id: str
    table_names: List[str]
    group_id: Optional[str] = None
    credential: Optional[TokenCredential] = None
    token: Optional[str] = None
    impersonated_user_name: Optional[str] = None
    sample_rows_in_table_info: int = Field(default=1, gt=0, le=10)
    schemas: Dict[str, str] = Field(default_factory=dict)
    aiosession: Optional[aiohttp.ClientSession] = None
[docs]    class Config:
        """Configuration for this pydantic object."""
        arbitrary_types_allowed = True
[docs]    @validator("table_names", allow_reuse=True)
    def fix_table_names(cls, table_names: List[str]) -> List[str]: