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.zhipuai.text_embedding.text_embedding import ZhipuAITextEmbeddingModel8 9 10def test_validate_credentials():11    model = ZhipuAITextEmbeddingModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(model="text_embedding", credentials={"api_key": "invalid_key"})15 16    model.validate_credentials(model="text_embedding", credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")})17 18 19def test_invoke_model():20    model = ZhipuAITextEmbeddingModel()21 22    result = model.invoke(23        model="text_embedding",24        credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")},25        texts=["hello", "world"],26        user="abc-123",27    )28 29    assert isinstance(result, TextEmbeddingResult)30    assert len(result.embeddings) == 231    assert result.usage.total_tokens > 032 33 34def test_get_num_tokens():35    model = ZhipuAITextEmbeddingModel()36 37    num_tokens = model.get_num_tokens(38        model="text_embedding", credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")}, texts=["hello", "world"]39    )40 41    assert num_tokens == 242