langchain_community.embeddings.octoai_embeddings

from typing import Dict

from langchain_core.pydantic_v1 import Field, SecretStr, root_validator
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env

from langchain_community.embeddings.openai import OpenAIEmbeddings
from langchain_community.utils.openai import is_openai_v1

DEFAULT_API_BASE = "https://text.octoai.run/v1/"
DEFAULT_MODEL = "thenlper/gte-large"


[docs]class OctoAIEmbeddings(OpenAIEmbeddings):
    """OctoAI计算服务嵌入模型。

    有关OctoAI的信息，请访问https://octo.ai/。

    要使用，您应该已安装``openai`` python包，并设置环境变量``OCTOAI_API_TOKEN``为您的API令牌。
    或者，您可以使用octoai_api_token关键字参数。"""

    octoai_api_token: SecretStr = Field(default=None)
    """OctoAI 终端 API 密钥。"""
    endpoint_url: str = Field(default=DEFAULT_API_BASE)
    """API请求的基础URL路径。"""
    model: str = Field(default=DEFAULT_MODEL)
    """要使用的模型名称。"""
    tiktoken_enabled: bool = False
    """将此设置为False，用于嵌入API的非OpenAI实现。"""

    @property
    def _llm_type(self) -> str:
        """嵌入模型的返回类型。"""
        return "octoai-embeddings"

    @property
    def lc_secrets(self) -> Dict[str, str]:
        return {"octoai_api_token": "OCTOAI_API_TOKEN"}

    @root_validator()
    def validate_environment(cls, values: dict) -> dict:
        """验证环境中是否存在API密钥和Python包。"""
        values["endpoint_url"] = get_from_dict_or_env(
            values,
            "endpoint_url",
            "ENDPOINT_URL",
            default=DEFAULT_API_BASE,
        )
        values["octoai_api_token"] = convert_to_secret_str(
            get_from_dict_or_env(values, "octoai_api_token", "OCTOAI_API_TOKEN")
        )
        values["model"] = get_from_dict_or_env(
            values,
            "model",
            "MODEL",
            default=DEFAULT_MODEL,
        )

        try:
            import openai

            if is_openai_v1():
                client_params = {
                    "api_key": values["octoai_api_token"].get_secret_value(),
                    "base_url": values["endpoint_url"],
                }
                if not values.get("client"):
                    values["client"] = openai.OpenAI(**client_params).embeddings
                if not values.get("async_client"):
                    values["async_client"] = openai.AsyncOpenAI(
                        **client_params
                    ).embeddings
            else:
                values["openai_api_base"] = values["endpoint_url"]
                values["openai_api_key"] = values["octoai_api_token"].get_secret_value()
                values["client"] = openai.Embedding
                values["async_client"] = openai.Embedding

        except ImportError:
            raise ImportError(
                "Could not import openai python package. "
                "Please install it with `pip install openai`."
            )

        return values
Source code for langchain_community.embeddings.octoai_embeddings
