initial commit

venv/Lib/site-packages/langchain_community/embeddings/premai.py

@@ -0,0 +1,130 @@
+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__)
+
+
+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
+
+    @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
+
+    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
+
+    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]
+
+
+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
+
+
+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)