codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4from typing import (5    Any,6    Callable,7    Dict,8    List,9    Optional,10)11 12from langchain_core.embeddings import Embeddings13from langchain_core.utils import get_from_dict_or_env14from pydantic import BaseModel, ConfigDict, model_validator15from requests.exceptions import HTTPError16from tenacity import (17    before_sleep_log,18    retry,19    retry_if_exception_type,20    stop_after_attempt,21    wait_exponential,22)23 24logger = logging.getLogger(__name__)25 26BATCH_SIZE = {27    "text-embedding-v1": 25,28    "text-embedding-v2": 25,29    "text-embedding-v3": 10,30    "text-embedding-v4": 10,31}32 33 34def _create_retry_decorator(embeddings: DashScopeEmbeddings) -> Callable[[Any], Any]:35    multiplier = 136    min_seconds = 137    max_seconds = 438    # Wait 2^x * 1 second between each retry starting with39    # 1 seconds, then up to 4 seconds, then 4 seconds afterwards40    return retry(41        reraise=True,42        stop=stop_after_attempt(embeddings.max_retries),43        wait=wait_exponential(multiplier, min=min_seconds, max=max_seconds),44        retry=(retry_if_exception_type(HTTPError)),45        before_sleep=before_sleep_log(logger, logging.WARNING),46    )47 48 49def embed_with_retry(embeddings: DashScopeEmbeddings, **kwargs: Any) -> Any:50    """Use tenacity to retry the embedding call."""51    retry_decorator = _create_retry_decorator(embeddings)52 53    @retry_decorator54    def _embed_with_retry(**kwargs: Any) -> Any:55        result = []56        i = 057        input_data = kwargs["input"]58        input_len = len(input_data) if isinstance(input_data, list) else 159        batch_size = BATCH_SIZE.get(kwargs["model"], 25)60        while i < input_len:61            kwargs["input"] = (62                input_data[i : i + batch_size]63                if isinstance(input_data, list)64                else input_data65            )66            resp = embeddings.client.call(**kwargs)67            if resp.status_code == 200:68                result += resp.output["embeddings"]69            elif resp.status_code in [400, 401]:70                raise ValueError(71                    f"status_code: {resp.status_code} \n "72                    f"code: {resp.code} \n message: {resp.message}"73                )74            else:75                raise HTTPError(76                    f"HTTP error occurred: status_code: {resp.status_code} \n "77                    f"code: {resp.code} \n message: {resp.message}",78                    response=resp,79                )80            i += batch_size81        return result82 83    return _embed_with_retry(**kwargs)84 85 86class DashScopeEmbeddings(BaseModel, Embeddings):87    """DashScope embedding models.88 89    To use, you should have the ``dashscope`` python package installed, and the90    environment variable ``DASHSCOPE_API_KEY`` set with your API key or pass it91    as a named parameter to the constructor.92 93    Example:94        .. code-block:: python95 96            from langchain_community.embeddings import DashScopeEmbeddings97            embeddings = DashScopeEmbeddings(dashscope_api_key="my-api-key")98 99    Example:100        .. code-block:: python101 102            import os103            os.environ["DASHSCOPE_API_KEY"] = "your DashScope API KEY"104 105            from langchain_community.embeddings.dashscope import DashScopeEmbeddings106            embeddings = DashScopeEmbeddings(107                model="text-embedding-v1",108            )109            text = "This is a test query."110            query_result = embeddings.embed_query(text)111 112    """113 114    client: Any = None  #: :meta private:115    """The DashScope client."""116    model: str = "text-embedding-v1"117    dashscope_api_key: Optional[str] = None118    max_retries: int = 5119    """Maximum number of retries to make when generating."""120 121    model_config = ConfigDict(122        extra="forbid",123    )124 125    @model_validator(mode="before")126    @classmethod127    def validate_environment(cls, values: Dict) -> Any:128        import dashscope129 130        """Validate that api key and python package exists in environment."""131        values["dashscope_api_key"] = get_from_dict_or_env(132            values, "dashscope_api_key", "DASHSCOPE_API_KEY"133        )134        dashscope.api_key = values["dashscope_api_key"]135        try:136            import dashscope137 138            values["client"] = dashscope.TextEmbedding139        except ImportError:140            raise ImportError(141                "Could not import dashscope python package. "142                "Please install it with `pip install dashscope`."143            )144        return values145 146    def embed_documents(self, texts: List[str]) -> List[List[float]]:147        """Call out to DashScope's embedding endpoint for embedding search docs.148 149        Args:150            texts: The list of texts to embed.151 152        Returns:153            List of embeddings, one for each text.154        """155        embeddings = embed_with_retry(156            self, input=texts, text_type="document", model=self.model157        )158        embedding_list = [item["embedding"] for item in embeddings]159        return embedding_list160 161    def embed_query(self, text: str) -> List[float]:162        """Call out to DashScope's embedding endpoint for embedding query text.163 164        Args:165            text: The text to embed.166 167        Returns:168            Embedding for the text.169        """170        embedding = embed_with_retry(171            self, input=text, text_type="query", model=self.model172        )[0]["embedding"]173        return embedding174