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.replicate.text_embedding.text_embedding import ReplicateEmbeddingModel8 9 10def test_validate_credentials_one():11    model = ReplicateEmbeddingModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(15            model="replicate/all-mpnet-base-v2",16            credentials={17                "replicate_api_token": "invalid_key",18                "model_version": "b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305",19            },20        )21 22    model.validate_credentials(23        model="replicate/all-mpnet-base-v2",24        credentials={25            "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),26            "model_version": "b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305",27        },28    )29 30 31def test_validate_credentials_two():32    model = ReplicateEmbeddingModel()33 34    with pytest.raises(CredentialsValidateFailedError):35        model.validate_credentials(36            model="nateraw/bge-large-en-v1.5",37            credentials={38                "replicate_api_token": "invalid_key",39                "model_version": "9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",40            },41        )42 43    model.validate_credentials(44        model="nateraw/bge-large-en-v1.5",45        credentials={46            "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),47            "model_version": "9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",48        },49    )50 51 52def test_invoke_model_one():53    model = ReplicateEmbeddingModel()54 55    result = model.invoke(56        model="nateraw/bge-large-en-v1.5",57        credentials={58            "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),59            "model_version": "9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",60        },61        texts=["hello", "world"],62        user="abc-123",63    )64 65    assert isinstance(result, TextEmbeddingResult)66    assert len(result.embeddings) == 267    assert result.usage.total_tokens == 268 69 70def test_invoke_model_two():71    model = ReplicateEmbeddingModel()72 73    result = model.invoke(74        model="andreasjansson/clip-features",75        credentials={76            "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),77            "model_version": "75b33f253f7714a281ad3e9b28f63e3232d583716ef6718f2e46641077ea040a",78        },79        texts=["hello", "world"],80        user="abc-123",81    )82 83    assert isinstance(result, TextEmbeddingResult)84    assert len(result.embeddings) == 285    assert result.usage.total_tokens == 286 87 88def test_invoke_model_three():89    model = ReplicateEmbeddingModel()90 91    result = model.invoke(92        model="replicate/all-mpnet-base-v2",93        credentials={94            "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),95            "model_version": "b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305",96        },97        texts=["hello", "world"],98        user="abc-123",99    )100 101    assert isinstance(result, TextEmbeddingResult)102    assert len(result.embeddings) == 2103    assert result.usage.total_tokens == 2104 105 106def test_invoke_model_four():107    model = ReplicateEmbeddingModel()108 109    result = model.invoke(110        model="nateraw/jina-embeddings-v2-base-en",111        credentials={112            "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),113            "model_version": "f8367a1c072ba2bc28af549d1faeacfe9b88b3f0e475add7a75091dac507f79e",114        },115        texts=["hello", "world"],116        user="abc-123",117    )118 119    assert isinstance(result, TextEmbeddingResult)120    assert len(result.embeddings) == 2121    assert result.usage.total_tokens == 2122 123 124def test_get_num_tokens():125    model = ReplicateEmbeddingModel()126 127    num_tokens = model.get_num_tokens(128        model="nateraw/jina-embeddings-v2-base-en",129        credentials={130            "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),131            "model_version": "f8367a1c072ba2bc28af549d1faeacfe9b88b3f0e475add7a75091dac507f79e",132        },133        texts=["hello", "world"],134    )135 136    assert num_tokens == 2137