dify-plugin-sdks/python/tests/interfaces/model/test_text_embedding.py

from collections.abc import Mapping
from decimal import Decimal

from dify_plugin.entities.model import EmbeddingInputType, ModelType
from dify_plugin.entities.model.text_embedding import EmbeddingUsage, TextEmbeddingResult
from dify_plugin.errors.model import InvokeError
from dify_plugin.interfaces.model.text_embedding_model import TextEmbeddingModel
from tests.interfaces.model.utils import prepare_model_factory


class MockTextEmbeddingModel(TextEmbeddingModel):
    def _invoke(
        self,
        model: str,
        credentials: dict,
        texts: list[str],
        user: str | None = None,
        input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
    ) -> TextEmbeddingResult:
        return TextEmbeddingResult(
            model=model,
            usage=EmbeddingUsage(
                tokens=0,
                total_tokens=0,
                unit_price=Decimal(0),
                price_unit=Decimal(0),
                total_price=Decimal(0),
                currency="USD",
                latency=0,
            ),
            embeddings=[[0.0] * 1536 for _ in texts],
        )

    @property
    def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:
        return {}

    def validate_credentials(self, model: str, credentials: Mapping) -> None:
        pass

    def get_num_tokens(self, model: str, credentials: dict, texts: list[str]) -> list[int]:
        return [0] * len(texts)


# test both constructor and invoke
def test_text_embedding():
    model_factory = prepare_model_factory({ModelType.TEXT_EMBEDDING: MockTextEmbeddingModel})
    instance = model_factory.get_instance(ModelType.TEXT_EMBEDDING)
    assert isinstance(instance, MockTextEmbeddingModel)
    instance.invoke(model="test", credentials={}, texts=["test"])