Source code for langchain_community.embeddings.edenai

from typing import Any, Dict, List, Optional

from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import (
    BaseModel,
    Extra,
    Field,
    SecretStr,
    root_validator,
)
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env

from langchain_community.utilities.requests import Requests


class EdenAiEmbeddings(BaseModel, Embeddings):
    """EdenAI嵌入。
    环境变量``EDENAI_API_KEY``设置为您的API密钥，或作为命名参数传递。"""

    edenai_api_key: Optional[SecretStr] = Field(None, description="EdenAI API Token")

    provider: str = "openai"
    """使用的嵌入提供程序（例如：openai、google等）。"""

    model: Optional[str] = None
    """上述提供商的模型名称（例如：'gpt-3.5-turbo-instruct' 适用于 openai）
可用模型显示在 https://docs.edenai.co/ 的 'available providers' 下。"""

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

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证环境中是否存在API密钥。"""
        values["edenai_api_key"] = convert_to_secret_str(
            get_from_dict_or_env(values, "edenai_api_key", "EDENAI_API_KEY")
        )
        return values

    @staticmethod
    def get_user_agent() -> str:
        from langchain_community import __version__

        return f"langchain/{__version__}"

    def _generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """使用EdenAi API计算嵌入。"""
        url = "https://api.edenai.run/v2/text/embeddings"

        headers = {
            "accept": "application/json",
            "content-type": "application/json",
            "authorization": f"Bearer {self.edenai_api_key.get_secret_value()}",  # type: ignore[union-attr]
            "User-Agent": self.get_user_agent(),
        }

        payload: Dict[str, Any] = {"texts": texts, "providers": self.provider}

        if self.model is not None:
            payload["settings"] = {self.provider: self.model}

        request = Requests(headers=headers)
        response = request.post(url=url, data=payload)
        if response.status_code >= 500:
            raise Exception(f"EdenAI Server: Error {response.status_code}")
        elif response.status_code >= 400:
            raise ValueError(f"EdenAI received an invalid payload: {response.text}")
        elif response.status_code != 200:
            raise Exception(
                f"EdenAI returned an unexpected response with status "
                f"{response.status_code}: {response.text}"
            )

        temp = response.json()

        provider_response = temp[self.provider]
        if provider_response.get("status") == "fail":
            err_msg = provider_response.get("error", {}).get("message")
            raise Exception(err_msg)

        embeddings = []
        for embed_item in temp[self.provider]["items"]:
            embedding = embed_item["embedding"]

            embeddings.append(embedding)

        return embeddings

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

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

返回：
    每个文本的嵌入列表。
"""

        return self._generate_embeddings(texts)

    def embed_query(self, text: str) -> List[float]:
        """嵌入一个查询使用EdenAI。

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

返回：
    文本的嵌入。
"""
        return self._generate_embeddings([text])[0]