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.openai.text_embedding.text_embedding import OpenAITextEmbeddingModel8from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock9 10 11@pytest.mark.parametrize("setup_openai_mock", [["text_embedding"]], indirect=True)12def test_validate_credentials(setup_openai_mock):13    model = OpenAITextEmbeddingModel()14 15    with pytest.raises(CredentialsValidateFailedError):16        model.validate_credentials(model="text-embedding-ada-002", credentials={"openai_api_key": "invalid_key"})17 18    model.validate_credentials(19        model="text-embedding-ada-002", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}20    )21 22 23@pytest.mark.parametrize("setup_openai_mock", [["text_embedding"]], indirect=True)24def test_invoke_model(setup_openai_mock):25    model = OpenAITextEmbeddingModel()26 27    result = model.invoke(28        model="text-embedding-ada-002",29        credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY"), "openai_api_base": "https://api.openai.com"},30        texts=["hello", "world", " ".join(["long_text"] * 100), " ".join(["another_long_text"] * 100)],31        user="abc-123",32    )33 34    assert isinstance(result, TextEmbeddingResult)35    assert len(result.embeddings) == 436    assert result.usage.total_tokens == 237 38 39def test_get_num_tokens():40    model = OpenAITextEmbeddingModel()41 42    num_tokens = model.get_num_tokens(43        model="text-embedding-ada-002",44        credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY"), "openai_api_base": "https://api.openai.com"},45        texts=["hello", "world"],46    )47 48    assert num_tokens == 249