mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-09 06:53:59 +00:00
community[minor]: Prem AI langchain integration (#19113)
### Prem SDK integration in LangChain This PR adds the integration with [PremAI's](https://www.premai.io/) prem-sdk with langchain. User can now access to deployed models (llms/embeddings) and use it with langchain's ecosystem. This PR adds the following: ### This PR adds the following: - [x] Add chat support - [X] Adding embedding support - [X] writing integration tests - [X] writing tests for chat - [X] writing tests for embedding - [X] writing unit tests - [X] writing tests for chat - [X] writing tests for embedding - [X] Adding documentation - [X] writing documentation for chat - [X] writing documentation for embedding - [X] run `make test` - [X] run `make lint`, `make lint_diff` - [X] Final checks (spell check, lint, format and overall testing) --------- Co-authored-by: Anindyadeep Sannigrahi <anindyadeepsannigrahi@Anindyadeeps-MacBook-Pro.local> Co-authored-by: Bagatur <baskaryan@gmail.com> Co-authored-by: Erick Friis <erick@langchain.dev> Co-authored-by: Bagatur <22008038+baskaryan@users.noreply.github.com>
This commit is contained in:
121
libs/community/langchain_community/embeddings/premai.py
Normal file
121
libs/community/langchain_community/embeddings/premai.py
Normal file
@@ -0,0 +1,121 @@
|
||||
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.pydantic_v1 import BaseModel, SecretStr, root_validator
|
||||
from langchain_core.utils import get_from_dict_or_env
|
||||
|
||||
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
|
||||
|
||||
@root_validator()
|
||||
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]:
|
||||
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)
|
Reference in New Issue
Block a user