codekingpro/portable-devtools

1from typing import Any, Dict, List2 3import requests4from langchain_core._api.deprecation import deprecated5from langchain_core.embeddings import Embeddings6from pydantic import BaseModel, ConfigDict7 8DEFAULT_MODEL_NAME = "@cf/baai/bge-base-en-v1.5"9 10 11@deprecated(12    since="0.3.23",13    removal="1.0",14    alternative_import="langchain_cloudflare.CloudflareWorkersAIEmbeddings",15)16class CloudflareWorkersAIEmbeddings(BaseModel, Embeddings):17    """Cloudflare Workers AI embedding model.18 19    To use, you need to provide an API token and20    account ID to access Cloudflare Workers AI.21 22    Example:23        .. code-block:: python24 25            from langchain_community.embeddings import CloudflareWorkersAIEmbeddings26 27            account_id = "my_account_id"28            api_token = "my_secret_api_token"29            model_name =  "@cf/baai/bge-small-en-v1.5"30 31            cf = CloudflareWorkersAIEmbeddings(32                account_id=account_id,33                api_token=api_token,34                model_name=model_name35            )36    """37 38    api_base_url: str = "https://api.cloudflare.com/client/v4/accounts"39    account_id: str40    api_token: str41    model_name: str = DEFAULT_MODEL_NAME42    batch_size: int = 5043    strip_new_lines: bool = True44    headers: Dict[str, str] = {"Authorization": "Bearer "}45 46    def __init__(self, **kwargs: Any):47        """Initialize the Cloudflare Workers AI client."""48        super().__init__(**kwargs)49 50        self.headers = {"Authorization": f"Bearer {self.api_token}"}51 52    model_config = ConfigDict(extra="forbid", protected_namespaces=())53 54    def embed_documents(self, texts: List[str]) -> List[List[float]]:55        """Compute doc embeddings using Cloudflare Workers AI.56 57        Args:58            texts: The list of texts to embed.59 60        Returns:61            List of embeddings, one for each text.62        """63        if self.strip_new_lines:64            texts = [text.replace("\n", " ") for text in texts]65 66        batches = [67            texts[i : i + self.batch_size]68            for i in range(0, len(texts), self.batch_size)69        ]70        embeddings = []71 72        for batch in batches:73            response = requests.post(74                f"{self.api_base_url}/{self.account_id}/ai/run/{self.model_name}",75                headers=self.headers,76                json={"text": batch},77            )78            embeddings.extend(response.json()["result"]["data"])79 80        return embeddings81 82    def embed_query(self, text: str) -> List[float]:83        """Compute query embeddings using Cloudflare Workers AI.84 85        Args:86            text: The text to embed.87 88        Returns:89            Embeddings for the text.90        """91        text = text.replace("\n", " ") if self.strip_new_lines else text92        response = requests.post(93            f"{self.api_base_url}/{self.account_id}/ai/run/{self.model_name}",94            headers=self.headers,95            json={"text": [text]},96        )97        return response.json()["result"]["data"][0]98