Harrison/bookend ai (#14258)

Co-authored-by: stvhu-bookend <142813359+stvhu-bookend@users.noreply.github.com>

libs/langchain/langchain/embeddings/__init__.py

@@ -22,6 +22,7 @@ from langchain.embeddings.awa import AwaEmbeddings
 from langchain.embeddings.azure_openai import AzureOpenAIEmbeddings
 from langchain.embeddings.baidu_qianfan_endpoint import QianfanEmbeddingsEndpoint
 from langchain.embeddings.bedrock import BedrockEmbeddings
+from langchain.embeddings.bookend import BookendEmbeddings
 from langchain.embeddings.cache import CacheBackedEmbeddings
 from langchain.embeddings.clarifai import ClarifaiEmbeddings
 from langchain.embeddings.cohere import CohereEmbeddings
@@ -127,6 +128,7 @@ __all__ = [
     "QianfanEmbeddingsEndpoint",
     "JohnSnowLabsEmbeddings",
     "VoyageEmbeddings",
+    "BookendEmbeddings",
 ]
 
 

libs/langchain/langchain/embeddings/bookend.py

@@ -0,0 +1,91 @@
+"""Wrapper around Bookend AI embedding models."""
+
+import json
+from typing import Any, List
+
+import requests
+
+from langchain.pydantic_v1 import BaseModel, Field
+from langchain.schema.embeddings import Embeddings
+
+API_URL = "https://api.bookend.ai/"
+DEFAULT_TASK = "embeddings"
+PATH = "/models/predict"
+
+
+class BookendEmbeddings(BaseModel, Embeddings):
+    """Bookend AI sentence_transformers embedding models.
+
+    Example:
+        .. code-block:: python
+
+            from langchain.embeddings import BookendEmbeddings
+
+            bookend = BookendEmbeddings(
+                domain={domain}
+                api_token={api_token}
+                model_id={model_id}
+            )
+            bookend.embed_documents([
+                "Please put on these earmuffs because I can't you hear.",
+                "Baby wipes are made of chocolate stardust.",
+            ])
+            bookend.embed_query(
+                "She only paints with bold colors; she does not like pastels."
+            )
+    """
+
+    domain: str
+    """Request for a domain at https://bookend.ai/ to use this embeddings module."""
+    api_token: str
+    """Request for an API token at https://bookend.ai/ to use this embeddings module."""
+    model_id: str
+    """Embeddings model ID to use."""
+    auth_header: dict = Field(default_factory=dict)
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        self.auth_header = {"Authorization": "Basic {}".format(self.api_token)}
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Embed documents using a Bookend deployed embeddings model.
+
+        Args:
+            texts: The list of texts to embed.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        result = []
+        headers = self.auth_header
+        headers["Content-Type"] = "application/json; charset=utf-8"
+        params = {
+            "model_id": self.model_id,
+            "task": DEFAULT_TASK,
+        }
+
+        for text in texts:
+            data = json.dumps(
+                {"text": text, "question": None, "context": None, "instruction": None}
+            )
+            r = requests.request(
+                "POST",
+                API_URL + self.domain + PATH,
+                headers=headers,
+                params=params,
+                data=data,
+            )
+            result.append(r.json()[0]["data"])
+
+        return result
+
+    def embed_query(self, text: str) -> List[float]:
+        """Embed a query using a Bookend deployed embeddings model.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embeddings for the text.
+        """
+        return self.embed_documents([text])[0]

libs/langchain/tests/integration_tests/embeddings/test_bookend.py

@@ -0,0 +1,27 @@
+"""Test Bookend AI embeddings."""
+from langchain.embeddings.bookend import BookendEmbeddings
+
+
+def test_bookend_embedding_documents() -> None:
+    """Test Bookend AI embeddings for documents."""
+    documents = ["foo bar", "bar foo"]
+    embedding = BookendEmbeddings(
+        domain="<bookend_domain>",
+        api_token="<bookend_api_token>",
+        model_id="<bookend_embeddings_model_id>",
+    )
+    output = embedding.embed_documents(documents)
+    assert len(output) == 2
+    assert len(output[0]) == 768
+
+
+def test_bookend_embedding_query() -> None:
+    """Test Bookend AI embeddings for query."""
+    document = "foo bar"
+    embedding = BookendEmbeddings(
+        domain="<bookend_domain>",
+        api_token="<bookend_api_token>",
+        model_id="<bookend_embeddings_model_id>",
+    )
+    output = embedding.embed_query(document)
+    assert len(output) == 768

libs/langchain/tests/unit_tests/embeddings/test_imports.py

@@ -53,6 +53,7 @@ EXPECTED_ALL = [
     "QianfanEmbeddingsEndpoint",
     "JohnSnowLabsEmbeddings",
     "VoyageEmbeddings",
+    "BookendEmbeddings",
 ]