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.oci.text_embedding.text_embedding import OCITextEmbeddingModel8 9 10def test_validate_credentials():11    model = OCITextEmbeddingModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(15            model="cohere.embed-multilingual-v3.0",16            credentials={"oci_config_content": "invalid_key", "oci_key_content": "invalid_key"},17        )18 19    model.validate_credentials(20        model="cohere.embed-multilingual-v3.0",21        credentials={22            "oci_config_content": os.environ.get("OCI_CONFIG_CONTENT"),23            "oci_key_content": os.environ.get("OCI_KEY_CONTENT"),24        },25    )26 27 28def test_invoke_model():29    model = OCITextEmbeddingModel()30 31    result = model.invoke(32        model="cohere.embed-multilingual-v3.0",33        credentials={34            "oci_config_content": os.environ.get("OCI_CONFIG_CONTENT"),35            "oci_key_content": os.environ.get("OCI_KEY_CONTENT"),36        },37        texts=["hello", "world", " ".join(["long_text"] * 100), " ".join(["another_long_text"] * 100)],38        user="abc-123",39    )40 41    assert isinstance(result, TextEmbeddingResult)42    assert len(result.embeddings) == 443    # assert result.usage.total_tokens == 81144 45 46def test_get_num_tokens():47    model = OCITextEmbeddingModel()48 49    num_tokens = model.get_num_tokens(50        model="cohere.embed-multilingual-v3.0",51        credentials={52            "oci_config_content": os.environ.get("OCI_CONFIG_CONTENT"),53            "oci_key_content": os.environ.get("OCI_KEY_CONTENT"),54        },55        texts=["hello", "world"],56    )57 58    assert num_tokens == 259