Source code for langchain_community.embeddings.zhipuai

from typing import Any, Dict, List, Optional

from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
from langchain_core.utils import get_from_dict_or_env


[docs]class ZhipuAIEmbeddings(BaseModel, Embeddings): """ZhipuAI embedding model integration. Setup: To use, you should have the ``zhipuai`` python package installed, and the environment variable ``ZHIPU_API_KEY`` set with your API KEY. More instructions about ZhipuAi Embeddings, you can get it from https://open.bigmodel.cn/dev/api#vector .. code-block:: bash pip install -U zhipuai export ZHIPU_API_KEY="your-api-key" Key init args — completion params: model: Optional[str] Name of ZhipuAI model to use. api_key: str Automatically inferred from env var `ZHIPU_API_KEY` if not provided. See full list of supported init args and their descriptions in the params section. Instantiate: .. code-block:: python from langchain_community.embeddings import ZhipuAIEmbeddings embed = ZhipuAIEmbeddings( model="embedding-2", # api_key="...", ) Embed single text: .. code-block:: python input_text = "The meaning of life is 42" embed.embed_query(input_text) .. code-block:: python [-0.003832892, 0.049372625, -0.035413884, -0.019301128, 0.0068899863, 0.01248398, -0.022153955, 0.006623926, 0.00778216, 0.009558191, ...] Embed multiple text: .. code-block:: python input_texts = ["This is a test query1.", "This is a test query2."] embed.embed_documents(input_texts) .. code-block:: python [ [0.0083934665, 0.037985895, -0.06684559, -0.039616987, 0.015481004, -0.023952313, ...], [-0.02713102, -0.005470169, 0.032321047, 0.042484466, 0.023290444, 0.02170547, ...] ] """ # noqa: E501 client: Any = Field(default=None, exclude=True) #: :meta private: model: str = Field(default="embedding-2") """Model name""" api_key: str """Automatically inferred from env var `ZHIPU_API_KEY` if not provided.""" dimensions: Optional[int] = None """The number of dimensions the resulting output embeddings should have. Only supported in `embedding-3` and later models. """ @root_validator(pre=True) def validate_environment(cls, values: Dict) -> Dict: """Validate that auth token exists in environment.""" values["api_key"] = get_from_dict_or_env(values, "api_key", "ZHIPUAI_API_KEY") try: from zhipuai import ZhipuAI values["client"] = ZhipuAI(api_key=values["api_key"]) except ImportError: raise ImportError( "Could not import zhipuai python package." "Please install it with `pip install zhipuai`." ) return values
[docs] def embed_query(self, text: str) -> List[float]: """ Embeds a text using the AutoVOT algorithm. Args: text: A text to embed. Returns: Input document's embedded list. """ resp = self.embed_documents([text]) return resp[0]
[docs] def embed_documents(self, texts: List[str]) -> List[List[float]]: """ Embeds a list of text documents using the AutoVOT algorithm. Args: texts: A list of text documents to embed. Returns: A list of embeddings for each document in the input list. Each embedding is represented as a list of float values. """ if self.dimensions is not None: resp = self.client.embeddings.create( model=self.model, input=texts, dimensions=self.dimensions, ) else: resp = self.client.embeddings.create(model=self.model, input=texts) embeddings = [r.embedding for r in resp.data] return embeddings