Underground-Digital/Workflow-Engine

1import os2from time import sleep3 4from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult5from core.model_runtime.model_providers.wenxin.text_embedding.text_embedding import WenxinTextEmbeddingModel6 7 8def test_invoke_embedding_v1():9    sleep(3)10    model = WenxinTextEmbeddingModel()11 12    response = model.invoke(13        model="embedding-v1",14        credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},15        texts=["hello", "你好", "xxxxx"],16        user="abc-123",17    )18 19    assert isinstance(response, TextEmbeddingResult)20    assert len(response.embeddings) == 321    assert isinstance(response.embeddings[0], list)22 23 24def test_invoke_embedding_bge_large_en():25    sleep(3)26    model = WenxinTextEmbeddingModel()27 28    response = model.invoke(29        model="bge-large-en",30        credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},31        texts=["hello", "你好", "xxxxx"],32        user="abc-123",33    )34 35    assert isinstance(response, TextEmbeddingResult)36    assert len(response.embeddings) == 337    assert isinstance(response.embeddings[0], list)38 39 40def test_invoke_embedding_bge_large_zh():41    sleep(3)42    model = WenxinTextEmbeddingModel()43 44    response = model.invoke(45        model="bge-large-zh",46        credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},47        texts=["hello", "你好", "xxxxx"],48        user="abc-123",49    )50 51    assert isinstance(response, TextEmbeddingResult)52    assert len(response.embeddings) == 353    assert isinstance(response.embeddings[0], list)54 55 56def test_invoke_embedding_tao_8k():57    sleep(3)58    model = WenxinTextEmbeddingModel()59 60    response = model.invoke(61        model="tao-8k",62        credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},63        texts=["hello", "你好", "xxxxx"],64        user="abc-123",65    )66 67    assert isinstance(response, TextEmbeddingResult)68    assert len(response.embeddings) == 369    assert isinstance(response.embeddings[0], list)70