codekingpro/portable-devtools

1"""Wrapper around Bookend AI embedding models."""2 3import json4from typing import Any, List5 6import requests7from langchain_core.embeddings import Embeddings8from pydantic import BaseModel, ConfigDict, Field9 10API_URL = "https://api.bookend.ai/"11DEFAULT_TASK = "embeddings"12PATH = "/models/predict"13 14 15class BookendEmbeddings(BaseModel, Embeddings):16    """Bookend AI sentence_transformers embedding models.17 18    Example:19        .. code-block:: python20 21            from langchain_community.embeddings import BookendEmbeddings22 23            bookend = BookendEmbeddings(24                domain={domain}25                api_token={api_token}26                model_id={model_id}27            )28            bookend.embed_documents([29                "Please put on these earmuffs because I can't you hear.",30                "Baby wipes are made of chocolate stardust.",31            ])32            bookend.embed_query(33                "She only paints with bold colors; she does not like pastels."34            )35    """36 37    domain: str38    """Request for a domain at https://bookend.ai/ to use this embeddings module."""39    api_token: str40    """Request for an API token at https://bookend.ai/ to use this embeddings module."""41    model_id: str42    """Embeddings model ID to use."""43    auth_header: dict = Field(default_factory=dict)44 45    model_config = ConfigDict(protected_namespaces=())46 47    def __init__(self, **kwargs: Any):48        super().__init__(**kwargs)49        self.auth_header = {"Authorization": "Basic {}".format(self.api_token)}50 51    def embed_documents(self, texts: List[str]) -> List[List[float]]:52        """Embed documents using a Bookend deployed embeddings model.53 54        Args:55            texts: The list of texts to embed.56 57        Returns:58            List of embeddings, one for each text.59        """60        result = []61        headers = self.auth_header62        headers["Content-Type"] = "application/json; charset=utf-8"63        params = {64            "model_id": self.model_id,65            "task": DEFAULT_TASK,66        }67 68        for text in texts:69            data = json.dumps(70                {71                    "text": text,72                    "question": None,73                    "context": None,74                    "instruction": None,75                }76            )77            r = requests.request(78                "POST",79                API_URL + self.domain + PATH,80                headers=headers,81                params=params,82                data=data,83            )84            result.append(r.json()[0]["data"])85 86        return result87 88    def embed_query(self, text: str) -> List[float]:89        """Embed a query using a Bookend deployed embeddings model.90 91        Args:92            text: The text to embed.93 94        Returns:95            Embeddings for the text.96        """97        return self.embed_documents([text])[0]98