Source code for langchain_community.embeddings.baichuan

from typing import Any, Dict, List, Optional

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

BAICHUAN_API_URL: str = "http://api.baichuan-ai.com/v1/embeddings"

# BaichuanTextEmbeddings is an embedding model provided by Baichuan Inc. (https://www.baichuan-ai.com/home).
# As of today (Jan 25th, 2024) BaichuanTextEmbeddings ranks #1 in C-MTEB
# (Chinese Multi-Task Embedding Benchmark) leaderboard.
# Leaderboard (Under Overall -> Chinese section): https://huggingface.co/spaces/mteb/leaderboard

# Official Website: https://platform.baichuan-ai.com/docs/text-Embedding
# An API-key is required to use this embedding model. You can get one by registering
# at https://platform.baichuan-ai.com/docs/text-Embedding.
# BaichuanTextEmbeddings support 512 token window and preduces vectors with
# 1024 dimensions.


# NOTE!! BaichuanTextEmbeddings only supports Chinese text embedding.
# Multi-language support is coming soon.
class BaichuanTextEmbeddings(BaseModel, Embeddings):
    """Baichuan文本嵌入模型。"""

    session: Any  #: :meta private:
    model_name: str = "Baichuan-Text-Embedding"
    baichuan_api_key: Optional[SecretStr] = None

    @root_validator(allow_reuse=True)
    def validate_environment(cls, values: Dict) -> Dict:
        """确保身份验证令牌存在于环境中。"""
        try:
            baichuan_api_key = convert_to_secret_str(
                get_from_dict_or_env(values, "baichuan_api_key", "BAICHUAN_API_KEY")
            )
        except ValueError as original_exc:
            try:
                baichuan_api_key = convert_to_secret_str(
                    get_from_dict_or_env(
                        values, "baichuan_auth_token", "BAICHUAN_AUTH_TOKEN"
                    )
                )
            except ValueError:
                raise original_exc
        session = requests.Session()
        session.headers.update(
            {
                "Authorization": f"Bearer {baichuan_api_key.get_secret_value()}",
                "Accept-Encoding": "identity",
                "Content-type": "application/json",
            }
        )
        values["session"] = session
        return values

    def _embed(self, texts: List[str]) -> Optional[List[List[float]]]:
        """内部方法，调用Baichuan嵌入API并返回嵌入结果。

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

返回：
    代表嵌入结果的浮点数列表的列表，如果发生错误则返回None。
"""
        try:
            response = self.session.post(
                BAICHUAN_API_URL, json={"input": texts, "model": self.model_name}
            )
            # Check if the response status code indicates success
            if response.status_code == 200:
                resp = response.json()
                embeddings = resp.get("data", [])
                # Sort resulting embeddings by index
                sorted_embeddings = sorted(embeddings, key=lambda e: e.get("index", 0))
                # Return just the embeddings
                return [result.get("embedding", []) for result in sorted_embeddings]
            else:
                # Log error or handle unsuccessful response appropriately
                print(  # noqa: T201
                    f"Error: Received status code {response.status_code} from "
                    "embedding API"
                )
                return None
        except Exception as e:
            # Log the exception or handle it as needed
            print(f"Exception occurred while trying to get embeddings: {str(e)}")  # noqa: T201
            return None

    def embed_documents(self, texts: List[str]) -> Optional[List[List[float]]]:  # type: ignore[override]
        """公共方法，用于获取文档列表的嵌入。

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

返回：
    返回嵌入列表，每个文本对应一个嵌入，如果发生错误则返回None。
"""
        return self._embed(texts)

    def embed_query(self, text: str) -> Optional[List[float]]:  # type: ignore[override]
        """公共方法，用于获取单个查询文本的嵌入。

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

返回：
    文本的嵌入，如果发生错误则返回None。
"""
        result = self._embed([text])
        return result[0] if result is not None else None