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.cohere.text_embedding.text_embedding import CohereTextEmbeddingModel8 9 10def test_validate_credentials():11    model = CohereTextEmbeddingModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(model="embed-multilingual-v3.0", credentials={"api_key": "invalid_key"})15 16    model.validate_credentials(17        model="embed-multilingual-v3.0", credentials={"api_key": os.environ.get("COHERE_API_KEY")}18    )19 20 21def test_invoke_model():22    model = CohereTextEmbeddingModel()23 24    result = model.invoke(25        model="embed-multilingual-v3.0",26        credentials={"api_key": os.environ.get("COHERE_API_KEY")},27        texts=["hello", "world", " ".join(["long_text"] * 100), " ".join(["another_long_text"] * 100)],28        user="abc-123",29    )30 31    assert isinstance(result, TextEmbeddingResult)32    assert len(result.embeddings) == 433    assert result.usage.total_tokens == 81134 35 36def test_get_num_tokens():37    model = CohereTextEmbeddingModel()38 39    num_tokens = model.get_num_tokens(40        model="embed-multilingual-v3.0",41        credentials={"api_key": os.environ.get("COHERE_API_KEY")},42        texts=["hello", "world"],43    )44 45    assert num_tokens == 346