Source code for langchain_community.embeddings.cohere

from typing import Any, Dict, List, Optional

from langchain_core._api.deprecation import deprecated
from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, Extra, root_validator
from langchain_core.utils import get_from_dict_or_env

from langchain_community.llms.cohere import _create_retry_decorator


@deprecated(
    since="0.0.30",
    removal="0.3.0",
    alternative_import="langchain_cohere.CohereEmbeddings",
)
class CohereEmbeddings(BaseModel, Embeddings):
    """嵌入模型。

    要使用，您应该安装``cohere`` python包，并设置环境变量``COHERE_API_KEY``为您的API密钥，或将其作为命名参数传递给构造函数。

    示例:
        .. code-block:: python

            from langchain_community.embeddings import CohereEmbeddings
            cohere = CohereEmbeddings(
                model="embed-english-light-v3.0",
                cohere_api_key="my-api-key"
            )
    """

    client: Any  #: :meta private:
    """Cohere客户端。"""
    async_client: Any  #: :meta private:
    """异步客户端。"""
    model: str = "embed-english-v2.0"
    """要使用的模型名称。"""

    truncate: Optional[str] = None
    """从开头或结尾截断过长的嵌入（"NONE"|"START"|"END"）"""

    cohere_api_key: Optional[str] = None

    max_retries: int = 3
    """生成时最大的重试次数。"""
    request_timeout: Optional[float] = None
    """Cohere API请求的超时时间（秒）。"""
    user_agent: str = "langchain"
    """用于发出请求的应用程序的标识符。"""

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证环境中是否存在API密钥和Python包。"""
        cohere_api_key = get_from_dict_or_env(
            values, "cohere_api_key", "COHERE_API_KEY"
        )
        request_timeout = values.get("request_timeout")

        try:
            import cohere

            client_name = values["user_agent"]
            values["client"] = cohere.Client(
                cohere_api_key,
                timeout=request_timeout,
                client_name=client_name,
            )
            values["async_client"] = cohere.AsyncClient(
                cohere_api_key,
                timeout=request_timeout,
                client_name=client_name,
            )
        except ImportError:
            raise ImportError(
                "Could not import cohere python package. "
                "Please install it with `pip install cohere`."
            )
        return values

    def embed_with_retry(self, **kwargs: Any) -> Any:
        """使用tenacity来重试嵌入调用。"""
        retry_decorator = _create_retry_decorator(self.max_retries)

        @retry_decorator
        def _embed_with_retry(**kwargs: Any) -> Any:
            return self.client.embed(**kwargs)

        return _embed_with_retry(**kwargs)

    def aembed_with_retry(self, **kwargs: Any) -> Any:
        """使用tenacity来重试嵌入调用。"""
        retry_decorator = _create_retry_decorator(self.max_retries)

        @retry_decorator
        async def _embed_with_retry(**kwargs: Any) -> Any:
            return await self.async_client.embed(**kwargs)

        return _embed_with_retry(**kwargs)

    def embed(
        self, texts: List[str], *, input_type: Optional[str] = None
    ) -> List[List[float]]:
        embeddings = self.embed_with_retry(
            model=self.model,
            texts=texts,
            input_type=input_type,
            truncate=self.truncate,
        ).embeddings
        return [list(map(float, e)) for e in embeddings]

    async def aembed(
        self, texts: List[str], *, input_type: Optional[str] = None
    ) -> List[List[float]]:
        embeddings = (
            await self.aembed_with_retry(
                model=self.model,
                texts=texts,
                input_type=input_type,
                truncate=self.truncate,
            )
        ).embeddings
        return [list(map(float, e)) for e in embeddings]

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """嵌入文档文本列表。

参数：
    texts：要嵌入的文本列表。

返回：
    嵌入列表，每个文本对应一个嵌入。
"""
        return self.embed(texts, input_type="search_document")

    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
        """异步调用到Cohere的嵌入端点。

参数：
    texts: 要嵌入的文本列表。

返回：
    每个文本的嵌入列表。
"""
        return await self.aembed(texts, input_type="search_document")

    def embed_query(self, text: str) -> List[float]:
        """调用Cohere的嵌入端点。

参数：
    text：要嵌入的文本。

返回：
    文本的嵌入。
"""
        return self.embed([text], input_type="search_query")[0]

    async def aembed_query(self, text: str) -> List[float]:
        """异步调用到Cohere的嵌入端点。

参数：
    text: 要嵌入的文本。

返回：
    文本的嵌入。
"""
        return (await self.aembed([text], input_type="search_query"))[0]