Kay retriever (#10657)

- **Description**: Adding retrievers for [kay.ai](https://kay.ai) and
SEC filings powered by Kay and Cybersyn. Kay provides context as a
service: it's an API built for RAG.
- **Issue**: N/A
- **Dependencies**: Just added a dep to the
[kay](https://pypi.org/project/kay/) package
- **Tag maintainer**: @baskaryan @hwchase17 Discussed in slack
- **Twtter handle:** [@vishalrohra_](https://twitter.com/vishalrohra_)

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>

libs/langchain/langchain/retrievers/__init__.py

@@ -30,6 +30,7 @@ from langchain.retrievers.ensemble import EnsembleRetriever
 from langchain.retrievers.google_cloud_enterprise_search import (
     GoogleCloudEnterpriseSearchRetriever,
 )
+from langchain.retrievers.kay import KayAiRetriever
 from langchain.retrievers.kendra import AmazonKendraRetriever
 from langchain.retrievers.knn import KNNRetriever
 from langchain.retrievers.llama_index import (
@@ -68,6 +69,7 @@ __all__ = [
     "ChaindeskRetriever",
     "ElasticSearchBM25Retriever",
     "GoogleCloudEnterpriseSearchRetriever",
+    "KayAiRetriever",
     "KNNRetriever",
     "LlamaIndexGraphRetriever",
     "LlamaIndexRetriever",

libs/langchain/langchain/retrievers/kay.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import Any, List
+
+from langchain.callbacks.manager import CallbackManagerForRetrieverRun
+from langchain.schema import BaseRetriever, Document
+
+
+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

libs/langchain/tests/integration_tests/retrievers/test_kay.py

@@ -0,0 +1,24 @@
+"""Integration test for Kay.ai API Wrapper."""
+import pytest
+
+from langchain.retrievers import KayAiRetriever
+from langchain.schema import Document
+
+
+@pytest.mark.requires("kay")
+def test_kay_retriever() -> None:
+    retriever = KayAiRetriever.create(
+        dataset_id="company",
+        data_types=["10-K", "10-Q", "8-K", "PressRelease"],
+        num_contexts=3,
+    )
+    docs = retriever.get_relevant_documents(
+        "What were the biggest strategy changes and partnerships made by Roku "
+        "in 2023?",
+    )
+    assert len(docs) == 3
+    for doc in docs:
+        assert isinstance(doc, Document)
+        assert doc.page_content
+        assert doc.metadata
+        assert len(list(doc.metadata.items())) > 0