mirror of
https://github.com/hwchase17/langchain.git
synced 2025-05-01 13:26:15 +00:00
43deed2a95
- In the in ` embedding-3 ` and later models of Zhipu AI, it is supported to specify the dimensions parameter of Embedding. Ref: https://bigmodel.cn/dev/api#text_embedding-3 . - Add test case for `embedding-3` model by assigning dimensions.
32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
"""Test ZhipuAI Text Embedding."""
|
|
|
|
from langchain_community.embeddings.zhipuai import ZhipuAIEmbeddings
|
|
|
|
|
|
def test_zhipuai_embedding_documents() -> None:
|
|
"""Test ZhipuAI Text Embedding for documents."""
|
|
documents = ["This is a test query1.", "This is a test query2."]
|
|
embedding = ZhipuAIEmbeddings() # type: ignore[call-arg]
|
|
res = embedding.embed_documents(documents)
|
|
assert len(res) == 2 # type: ignore[arg-type]
|
|
assert len(res[0]) == 1024 # type: ignore[index]
|
|
|
|
|
|
def test_zhipuai_embedding_query() -> None:
|
|
"""Test ZhipuAI Text Embedding for query."""
|
|
document = "This is a test query."
|
|
embedding = ZhipuAIEmbeddings() # type: ignore[call-arg]
|
|
res = embedding.embed_query(document)
|
|
assert len(res) == 1024 # type: ignore[arg-type]
|
|
|
|
|
|
def test_zhipuai_embedding_dimensions() -> None:
|
|
"""Test ZhipuAI Text Embedding for query by assigning dimensions"""
|
|
document = "This is a test query."
|
|
embedding = ZhipuAIEmbeddings(
|
|
model="embedding-3",
|
|
dimensions=2048,
|
|
) # type: ignore[call-arg]
|
|
res = embedding.embed_query(document)
|
|
assert len(res) == 2048 # type: ignore[arg-type]
|