dify/api/tests/integration_tests/model_runtime/wenxin/test_embedding.py
Bowen Liang b035c02f78
Some checks are pending
Build and Push API & Web / build (api, DIFY_API_IMAGE_NAME, linux/amd64, build-api-amd64) (push) Waiting to run
Build and Push API & Web / build (api, DIFY_API_IMAGE_NAME, linux/arm64, build-api-arm64) (push) Waiting to run
Build and Push API & Web / build (web, DIFY_WEB_IMAGE_NAME, linux/amd64, build-web-amd64) (push) Waiting to run
Build and Push API & Web / build (web, DIFY_WEB_IMAGE_NAME, linux/arm64, build-web-arm64) (push) Waiting to run
Build and Push API & Web / create-manifest (api, DIFY_API_IMAGE_NAME, merge-api-images) (push) Blocked by required conditions
Build and Push API & Web / create-manifest (web, DIFY_WEB_IMAGE_NAME, merge-web-images) (push) Blocked by required conditions
chore(api/tests): apply ruff reformat #7590 (#7591)
Co-authored-by: -LAN- <laipz8200@outlook.com>
2024-08-23 23:52:25 +08:00

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)