Underground-Digital/Workflow-Engine

1import os
2
3import pytest
4
5from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
6from core.model_runtime.errors.validate import CredentialsValidateFailedError
7from core.model_runtime.model_providers.gitee_ai.text_embedding.text_embedding import GiteeAIEmbeddingModel
8
9
10def test_validate_credentials():
11    model = GiteeAIEmbeddingModel()
12
13    with pytest.raises(CredentialsValidateFailedError):
14        model.validate_credentials(model="bge-large-zh-v1.5", credentials={"api_key": "invalid_key"})
15
16    model.validate_credentials(model="bge-large-zh-v1.5", credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")})
17
18
19def test_invoke_model():
20    model = GiteeAIEmbeddingModel()
21
22    result = model.invoke(
23        model="bge-large-zh-v1.5",
24        credentials={
25            "api_key": os.environ.get("GITEE_AI_API_KEY"),
26        },
27        texts=["hello", "world"],
28        user="user",
29    )
30
31    assert isinstance(result, TextEmbeddingResult)
32    assert len(result.embeddings) == 2
33
34
35def test_get_num_tokens():
36    model = GiteeAIEmbeddingModel()
37
38    num_tokens = model.get_num_tokens(
39        model="bge-large-zh-v1.5",
40        credentials={
41            "api_key": os.environ.get("GITEE_AI_API_KEY"),
42        },
43        texts=["hello", "world"],
44    )
45
46    assert num_tokens == 2
47