langchain/libs/community/tests/unit_tests/embeddings/test_baichuan.py
maang-h c6b7db6587
community: Add Baichuan Embeddings batch size (#22942)
- **Support batch size** 
Baichuan updates the document, indicating that up to 16 documents can be
imported at a time

- **Standardized model init arg names**
    - baichuan_api_key -> api_key
    - model_name  -> model
2024-06-17 14:11:04 -04:00

19 lines
588 B
Python

from typing import cast
from langchain_core.pydantic_v1 import SecretStr
from langchain_community.embeddings import BaichuanTextEmbeddings
def test_sparkllm_initialization_by_alias() -> None:
# Effective initialization
embeddings = BaichuanTextEmbeddings( # type: ignore[call-arg]
model="embedding_model", # type: ignore[arg-type]
api_key="your-api-key", # type: ignore[arg-type]
)
assert embeddings.model_name == "embedding_model"
assert (
cast(SecretStr, embeddings.baichuan_api_key).get_secret_value()
== "your-api-key"
)