codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4from typing import Any, Callable, Dict, List, Optional, Union5 6from langchain_core.embeddings import Embeddings7from langchain_core.language_models.llms import create_base_retry_decorator8from langchain_core.utils import get_from_dict_or_env, pre_init9from pydantic import BaseModel, SecretStr10 11logger = logging.getLogger(__name__)12 13 14class PremAIEmbeddings(BaseModel, Embeddings):15    """Prem's Embedding APIs"""16 17    project_id: int18    """The project ID in which the experiments or deployments are carried out. 19    You can find all your projects here: https://app.premai.io/projects/"""20 21    premai_api_key: Optional[SecretStr] = None22    """Prem AI API Key. Get it here: https://app.premai.io/api_keys/"""23 24    model: str25    """The Embedding model to choose from"""26 27    show_progress_bar: bool = False28    """Whether to show a tqdm progress bar. Must have `tqdm` installed."""29 30    max_retries: int = 131    """Max number of retries for tenacity"""32 33    client: Any34 35    @pre_init36    def validate_environments(cls, values: Dict) -> Dict:37        """Validate that the package is installed and that the API token is valid"""38        try:39            from premai import Prem40        except ImportError as error:41            raise ImportError(42                "Could not import Prem Python package."43                "Please install it with: `pip install premai`"44            ) from error45 46        try:47            premai_api_key = get_from_dict_or_env(48                values, "premai_api_key", "PREMAI_API_KEY"49            )50            values["client"] = Prem(api_key=premai_api_key)51        except Exception as error:52            raise ValueError("Your API Key is incorrect. Please try again.") from error53        return values54 55    def embed_query(self, text: str) -> List[float]:56        """Embed query text"""57        embeddings = embed_with_retry(58            self, model=self.model, project_id=self.project_id, input=text59        )60        return embeddings.data[0].embedding61 62    def embed_documents(self, texts: List[str]) -> List[List[float]]:63        embeddings = embed_with_retry(64            self, model=self.model, project_id=self.project_id, input=texts65        ).data66 67        return [embedding.embedding for embedding in embeddings]68 69 70def create_prem_retry_decorator(71    embedder: PremAIEmbeddings,72    *,73    max_retries: int = 1,74) -> Callable[[Any], Any]:75    """Create a retry decorator for PremAIEmbeddings.76 77    Args:78        embedder (PremAIEmbeddings): The PremAIEmbeddings instance79        max_retries (int): The maximum number of retries80 81    Returns:82        Callable[[Any], Any]: The retry decorator83    """84    import premai.models85 86    errors = [87        premai.models.api_response_validation_error.APIResponseValidationError,88        premai.models.conflict_error.ConflictError,89        premai.models.model_not_found_error.ModelNotFoundError,90        premai.models.permission_denied_error.PermissionDeniedError,91        premai.models.provider_api_connection_error.ProviderAPIConnectionError,92        premai.models.provider_api_status_error.ProviderAPIStatusError,93        premai.models.provider_api_timeout_error.ProviderAPITimeoutError,94        premai.models.provider_internal_server_error.ProviderInternalServerError,95        premai.models.provider_not_found_error.ProviderNotFoundError,96        premai.models.rate_limit_error.RateLimitError,97        premai.models.unprocessable_entity_error.UnprocessableEntityError,98        premai.models.validation_error.ValidationError,99    ]100 101    decorator = create_base_retry_decorator(102        error_types=errors, max_retries=max_retries, run_manager=None103    )104    return decorator105 106 107def embed_with_retry(108    embedder: PremAIEmbeddings,109    model: str,110    project_id: int,111    input: Union[str, List[str]],112) -> Any:113    """Using tenacity for retry in embedding calls"""114    retry_decorator = create_prem_retry_decorator(115        embedder, max_retries=embedder.max_retries116    )117 118    @retry_decorator119    def _embed_with_retry(120        embedder: PremAIEmbeddings,121        project_id: int,122        model: str,123        input: Union[str, List[str]],124    ) -> Any:125        embedding_response = embedder.client.embeddings.create(126            project_id=project_id, model=model, input=input127        )128        return embedding_response129 130    return _embed_with_retry(embedder, project_id=project_id, model=model, input=input)131