from __future__ import annotations
import logging
from typing import Any, Callable, Dict, List, Optional, Union
from langchain_core.embeddings import Embeddings
from langchain_core.language_models.llms import create_base_retry_decorator
from langchain_core.utils import get_from_dict_or_env, pre_init
from pydantic import BaseModel, SecretStr
logger = logging.getLogger(__name__)
[docs]
class PremAIEmbeddings(BaseModel, Embeddings):
"""Prem's Embedding APIs"""
project_id: int
"""The project ID in which the experiments or deployments are carried out.
You can find all your projects here: https://app.premai.io/projects/"""
premai_api_key: Optional[SecretStr] = None
"""Prem AI API Key. Get it here: https://app.premai.io/api_keys/"""
model: str
"""The Embedding model to choose from"""
show_progress_bar: bool = False
"""Whether to show a tqdm progress bar. Must have `tqdm` installed."""
max_retries: int = 1
"""Max number of retries for tenacity"""
client: Any
[docs]
@pre_init
def validate_environments(cls, values: Dict) -> Dict:
"""Validate that the package is installed and that the API token is valid"""
try:
from premai import Prem
except ImportError as error:
raise ImportError(
"Could not import Prem Python package."
"Please install it with: `pip install premai`"
) from error
try:
premai_api_key = get_from_dict_or_env(
values, "premai_api_key", "PREMAI_API_KEY"
)
values["client"] = Prem(api_key=premai_api_key)
except Exception as error:
raise ValueError("Your API Key is incorrect. Please try again.") from error
return values
[docs]
def embed_query(self, text: str) -> List[float]:
"""Embed query text"""
embeddings = embed_with_retry(
self, model=self.model, project_id=self.project_id, input=text
)
return embeddings.data[0].embedding
[docs]
def embed_documents(self, texts: List[str]) -> List[List[float]]:
embeddings = embed_with_retry(
self, model=self.model, project_id=self.project_id, input=texts
).data
return [embedding.embedding for embedding in embeddings]
[docs]
def create_prem_retry_decorator(
embedder: PremAIEmbeddings,
*,
max_retries: int = 1,
) -> Callable[[Any], Any]:
"""Create a retry decorator for PremAIEmbeddings.
Args:
embedder (PremAIEmbeddings): The PremAIEmbeddings instance
max_retries (int): The maximum number of retries
Returns:
Callable[[Any], Any]: The retry decorator
"""
import premai.models
errors = [
premai.models.api_response_validation_error.APIResponseValidationError,
premai.models.conflict_error.ConflictError,
premai.models.model_not_found_error.ModelNotFoundError,
premai.models.permission_denied_error.PermissionDeniedError,
premai.models.provider_api_connection_error.ProviderAPIConnectionError,
premai.models.provider_api_status_error.ProviderAPIStatusError,
premai.models.provider_api_timeout_error.ProviderAPITimeoutError,
premai.models.provider_internal_server_error.ProviderInternalServerError,
premai.models.provider_not_found_error.ProviderNotFoundError,
premai.models.rate_limit_error.RateLimitError,
premai.models.unprocessable_entity_error.UnprocessableEntityError,
premai.models.validation_error.ValidationError,
]
decorator = create_base_retry_decorator(
error_types=errors, max_retries=max_retries, run_manager=None
)
return decorator
[docs]
def embed_with_retry(
embedder: PremAIEmbeddings,
model: str,
project_id: int,
input: Union[str, List[str]],
) -> Any:
"""Using tenacity for retry in embedding calls"""
retry_decorator = create_prem_retry_decorator(
embedder, max_retries=embedder.max_retries
)
@retry_decorator
def _embed_with_retry(
embedder: PremAIEmbeddings,
project_id: int,
model: str,
input: Union[str, List[str]],
) -> Any:
embedding_response = embedder.client.embeddings.create(
project_id=project_id, model=model, input=input
)
return embedding_response
return _embed_with_retry(embedder, project_id=project_id, model=model, input=input)
