langchain_community.embeddings.llm_rails — 🦜🔗 LangChain 0.2.0

""" 这个文件是用于LLMRails嵌入的。"""
from typing import Dict, List, Optional

import requests
from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import BaseModel, Extra, SecretStr, root_validator
from langchain_core.utils import convert_to_secret_str, get_from_dict_or_env


[docs]class LLMRailsEmbeddings(BaseModel, Embeddings):
    """LLMRails嵌入模型。

    要使用，应该设置环境变量``LLM_RAILS_API_KEY``为您的API密钥，或将其作为构造函数的命名参数传递。

    模型可以是["embedding-english-v1", "embedding-multi-v1"]

    示例:
        .. code-block:: python

            from langchain_community.embeddings import LLMRailsEmbeddings
            cohere = LLMRailsEmbeddings(
                model="embedding-english-v1", api_key="my-api-key"
            )"""

    model: str = "embedding-english-v1"
    """要使用的模型名称。"""

    api_key: Optional[SecretStr] = None
    """LLMRails API密钥。"""

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

        extra = Extra.forbid

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

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

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

返回：
    每个文本的嵌入列表。
"""
        response = requests.post(
            "https://api.llmrails.com/v1/embeddings",
            headers={"X-API-KEY": self.api_key.get_secret_value()},  # type: ignore[union-attr]
            json={"input": texts, "model": self.model},
            timeout=60,
        )
        return [item["embedding"] for item in response.json()["data"]]

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

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

返回：
    文本的嵌入。
"""
        return self.embed_documents([text])[0]
Source code for langchain_community.embeddings.llm_rails
