Team Ai
Apppublic

Underground-Digital/Workflow-Engine

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes
test_embedding.py58 linesDownload Raw Back to openllm
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