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.jina.text_embedding.text_embedding import JinaTextEmbeddingModel8 9 10def test_validate_credentials():11    model = JinaTextEmbeddingModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(model="jina-embeddings-v2-base-en", credentials={"api_key": "invalid_key"})15 16    model.validate_credentials(17        model="jina-embeddings-v2-base-en", credentials={"api_key": os.environ.get("JINA_API_KEY")}18    )19 20 21def test_invoke_model():22    model = JinaTextEmbeddingModel()23 24    result = model.invoke(25        model="jina-embeddings-v2-base-en",26        credentials={27            "api_key": os.environ.get("JINA_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 = JinaTextEmbeddingModel()40 41    num_tokens = model.get_num_tokens(42        model="jina-embeddings-v2-base-en",43        credentials={44            "api_key": os.environ.get("JINA_API_KEY"),45        },46        texts=["hello", "world"],47    )48 49    assert num_tokens == 650