codekingpro/portable-devtools

1"""Anyscale embeddings wrapper."""2 3from __future__ import annotations4 5from typing import Dict, Optional6 7from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env, pre_init8from pydantic import Field, SecretStr9 10from langchain_community.embeddings.openai import OpenAIEmbeddings11from langchain_community.utils.openai import is_openai_v112 13DEFAULT_API_BASE = "https://api.endpoints.anyscale.com/v1"14DEFAULT_MODEL = "thenlper/gte-large"15 16 17class AnyscaleEmbeddings(OpenAIEmbeddings):18    """`Anyscale` Embeddings API."""19 20    anyscale_api_key: Optional[SecretStr] = Field(default=None)21    """AnyScale Endpoints API keys."""22    model: str = Field(default=DEFAULT_MODEL)23    """Model name to use."""24    anyscale_api_base: str = Field(default=DEFAULT_API_BASE)25    """Base URL path for API requests."""26    tiktoken_enabled: bool = False27    """Set this to False for non-OpenAI implementations of the embeddings API"""28    embedding_ctx_length: int = 50029    """The maximum number of tokens to embed at once."""30 31    @property32    def lc_secrets(self) -> Dict[str, str]:33        return {34            "anyscale_api_key": "ANYSCALE_API_KEY",35        }36 37    @pre_init38    def validate_environment(cls, values: dict) -> dict:39        """Validate that api key and python package exists in environment."""40        values["anyscale_api_key"] = convert_to_secret_str(41            get_from_dict_or_env(42                values,43                "anyscale_api_key",44                "ANYSCALE_API_KEY",45            )46        )47        values["anyscale_api_base"] = get_from_dict_or_env(48            values,49            "anyscale_api_base",50            "ANYSCALE_API_BASE",51            default=DEFAULT_API_BASE,52        )53        try:54            import openai55 56        except ImportError:57            raise ImportError(58                "Could not import openai python package. "59                "Please install it with `pip install openai`."60            )61        if is_openai_v1():62            # For backwards compatibility.63            client_params = {64                "api_key": values["anyscale_api_key"].get_secret_value(),65                "base_url": values["anyscale_api_base"],66            }67            values["client"] = openai.OpenAI(**client_params).embeddings68        else:69            values["openai_api_base"] = values["anyscale_api_base"]70            values["openai_api_key"] = values["anyscale_api_key"].get_secret_value()71            values["client"] = openai.Embedding72        return values73 74    @property75    def _llm_type(self) -> str:76        return "anyscale-embedding"77