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.siliconflow.text_embedding.text_embedding import (8    SiliconflowTextEmbeddingModel,9)10 11 12def test_validate_credentials():13    model = SiliconflowTextEmbeddingModel()14 15    with pytest.raises(CredentialsValidateFailedError):16        model.validate_credentials(17            model="BAAI/bge-large-zh-v1.5",18            credentials={"api_key": "invalid_key"},19        )20 21    model.validate_credentials(22        model="BAAI/bge-large-zh-v1.5",23        credentials={24            "api_key": os.environ.get("API_KEY"),25        },26    )27 28 29def test_invoke_model():30    model = SiliconflowTextEmbeddingModel()31 32    result = model.invoke(33        model="BAAI/bge-large-zh-v1.5",34        credentials={35            "api_key": os.environ.get("API_KEY"),36        },37        texts=[38            "hello",39            "world",40        ],41        user="abc-123",42    )43 44    assert isinstance(result, TextEmbeddingResult)45    assert len(result.embeddings) == 246    assert result.usage.total_tokens == 647 48 49def test_get_num_tokens():50    model = SiliconflowTextEmbeddingModel()51 52    num_tokens = model.get_num_tokens(53        model="BAAI/bge-large-zh-v1.5",54        credentials={55            "api_key": os.environ.get("API_KEY"),56        },57        texts=["hello", "world"],58    )59 60    assert num_tokens == 261