Underground-Digital/Workflow-Engine
0
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.openllm.text_embedding.text_embedding import OpenLLMTextEmbeddingModel8 9 10def test_validate_credentials():11 model = OpenLLMTextEmbeddingModel()12 13 with pytest.raises(CredentialsValidateFailedError):14 model.validate_credentials(15 model="NOT IMPORTANT",16 credentials={17 "server_url": "ww" + os.environ.get("OPENLLM_SERVER_URL"),18 },19 )20 21 model.validate_credentials(22 model="NOT IMPORTANT",23 credentials={24 "server_url": os.environ.get("OPENLLM_SERVER_URL"),25 },26 )27 28 29def test_invoke_model():30 model = OpenLLMTextEmbeddingModel()31 32 result = model.invoke(33 model="NOT IMPORTANT",34 credentials={35 "server_url": os.environ.get("OPENLLM_SERVER_URL"),36 },37 texts=["hello", "world"],38 user="abc-123",39 )40 41 assert isinstance(result, TextEmbeddingResult)42 assert len(result.embeddings) == 243 assert result.usage.total_tokens > 044 45 46def test_get_num_tokens():47 model = OpenLLMTextEmbeddingModel()48 49 num_tokens = model.get_num_tokens(50 model="NOT IMPORTANT",51 credentials={52 "server_url": os.environ.get("OPENLLM_SERVER_URL"),53 },54 texts=["hello", "world"],55 )56 57 assert num_tokens == 258 