Toggle Menu
Prev Up Next

# Source code for langchain_community.retrievers.kay

    from __future__ import annotations

    from typing import Any, List

    from langchain_core.callbacks import CallbackManagerForRetrieverRun
    from langchain_core.documents import Document
    from langchain_core.retrievers import BaseRetriever

    class KayAiRetriever(BaseRetriever):
        """
        Retriever for Kay.ai datasets.

        To work properly, expects you to have KAY_API_KEY env variable set.
        You can get one for free at https://kay.ai/.
        """

        client: Any
        num_contexts: int

        @classmethod
        def create(
            cls,
            dataset_id: str,
            data_types: List[str],
            num_contexts: int = 6,
        ) -> KayAiRetriever:
            """
            Create a KayRetriever given a Kay dataset id and a list of datasources.

            Args:
                dataset_id: A dataset id category in Kay, like "company"
                data_types: A list of datasources present within a dataset. For
                    "company" the corresponding datasources could be
                    ["10-K", "10-Q", "8-K", "PressRelease"].
                num_contexts: The number of documents to retrieve on each query.
                    Defaults to 6.
            """
            try:
                from kay.rag.retrievers import KayRetriever
            except ImportError:
                raise ImportError(
                    "Could not import kay python package. Please install it with "
                    "`pip install kay`.",
                )

            client = KayRetriever(dataset_id, data_types)
            return cls(client=client, num_contexts=num_contexts)

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