Underground-Digital/Workflow-Engine

1import os2 3import pytest4 5from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult6from core.model_runtime.errors.validate import CredentialsValidateFailedError7from core.model_runtime.model_providers.baichuan.text_embedding.text_embedding import BaichuanTextEmbeddingModel8 9 10def test_validate_credentials():11    model = BaichuanTextEmbeddingModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(model="baichuan-text-embedding", credentials={"api_key": "invalid_key"})15 16    model.validate_credentials(17        model="baichuan-text-embedding", credentials={"api_key": os.environ.get("BAICHUAN_API_KEY")}18    )19 20 21def test_invoke_model():22    model = BaichuanTextEmbeddingModel()23 24    result = model.invoke(25        model="baichuan-text-embedding",26        credentials={27            "api_key": os.environ.get("BAICHUAN_API_KEY"),28        },29        texts=["hello", "world"],30        user="abc-123",31    )32 33    assert isinstance(result, TextEmbeddingResult)34    assert len(result.embeddings) == 235    assert result.usage.total_tokens == 636 37 38def test_get_num_tokens():39    model = BaichuanTextEmbeddingModel()40 41    num_tokens = model.get_num_tokens(42        model="baichuan-text-embedding",43        credentials={44            "api_key": os.environ.get("BAICHUAN_API_KEY"),45        },46        texts=["hello", "world"],47    )48 49    assert num_tokens == 250 51 52def test_max_chunks():53    model = BaichuanTextEmbeddingModel()54 55    result = model.invoke(56        model="baichuan-text-embedding",57        credentials={58            "api_key": os.environ.get("BAICHUAN_API_KEY"),59        },60        texts=[61            "hello",62            "world",63            "hello",64            "world",65            "hello",66            "world",67            "hello",68            "world",69            "hello",70            "world",71            "hello",72            "world",73            "hello",74            "world",75            "hello",76            "world",77            "hello",78            "world",79            "hello",80            "world",81            "hello",82            "world",83        ],84    )85 86    assert isinstance(result, TextEmbeddingResult)87    assert len(result.embeddings) == 2288