Source code for langchain_nvidia_ai_endpoints.embeddings
"""Embeddings Components Derived from NVEModel/Embeddings"""
from typing import Any, List, Literal, Optional
from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
import langchain_nvidia_ai_endpoints._common as nvai_common
[docs]class NVIDIAEmbeddings(BaseModel, Embeddings):
    """NVIDIA's AI Foundation Retriever Question-Answering Asymmetric Model."""
    client: nvai_common.NVEModel = Field(nvai_common.NVEModel)
    model: str = Field(
        ..., description="The embedding model to use. Example: nvolveqa_40k"
    )
    max_length: int = Field(2048, ge=1, le=2048)
    max_batch_size: int = Field(default=50)
    model_type: Optional[Literal["passage", "query"]] = Field(
        "passage", description="The type of text to be embedded."
    )
    @root_validator(pre=True)
    def _validate_client(cls, values: Any) -> Any:
        if "client" not in values:
            values["client"] = nvai_common.NVEModel()
        return values
    @property
    def available_models(self) -> dict:
        """Map the available models that can be invoked."""
        return self.client.available_models
    def _embed(
        self, texts: List[str], model_type: Literal["passage", "query"]
    ) -> List[List[float]]:
        """Embed a single text entry to either passage or query type"""
        response = self.client.get_req(
            model_name=self.model,
            payload={
                "input": texts,