mirror of
https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git
synced 2025-06-30 03:35:13 +08:00
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
import os
|
|
from time import sleep
|
|
|
|
from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
|
|
from core.model_runtime.model_providers.wenxin.text_embedding.text_embedding import WenxinTextEmbeddingModel
|
|
|
|
|
|
def test_invoke_embedding_v1():
|
|
sleep(3)
|
|
model = WenxinTextEmbeddingModel()
|
|
|
|
response = model.invoke(
|
|
model="embedding-v1",
|
|
credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
|
|
texts=["hello", "你好", "xxxxx"],
|
|
user="abc-123",
|
|
)
|
|
|
|
assert isinstance(response, TextEmbeddingResult)
|
|
assert len(response.embeddings) == 3
|
|
assert isinstance(response.embeddings[0], list)
|
|
|
|
|
|
def test_invoke_embedding_bge_large_en():
|
|
sleep(3)
|
|
model = WenxinTextEmbeddingModel()
|
|
|
|
response = model.invoke(
|
|
model="bge-large-en",
|
|
credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
|
|
texts=["hello", "你好", "xxxxx"],
|
|
user="abc-123",
|
|
)
|
|
|
|
assert isinstance(response, TextEmbeddingResult)
|
|
assert len(response.embeddings) == 3
|
|
assert isinstance(response.embeddings[0], list)
|
|
|
|
|
|
def test_invoke_embedding_bge_large_zh():
|
|
sleep(3)
|
|
model = WenxinTextEmbeddingModel()
|
|
|
|
response = model.invoke(
|
|
model="bge-large-zh",
|
|
credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
|
|
texts=["hello", "你好", "xxxxx"],
|
|
user="abc-123",
|
|
)
|
|
|
|
assert isinstance(response, TextEmbeddingResult)
|
|
assert len(response.embeddings) == 3
|
|
assert isinstance(response.embeddings[0], list)
|
|
|
|
|
|
def test_invoke_embedding_tao_8k():
|
|
sleep(3)
|
|
model = WenxinTextEmbeddingModel()
|
|
|
|
response = model.invoke(
|
|
model="tao-8k",
|
|
credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
|
|
texts=["hello", "你好", "xxxxx"],
|
|
user="abc-123",
|
|
)
|
|
|
|
assert isinstance(response, TextEmbeddingResult)
|
|
assert len(response.embeddings) == 3
|
|
assert isinstance(response.embeddings[0], list)
|