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.ollama.text_embedding.text_embedding import OllamaEmbeddingModel8 9 10def test_validate_credentials():11    model = OllamaEmbeddingModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(15            model="mistral:text",16            credentials={17                "base_url": "http://localhost:21434",18                "mode": "chat",19                "context_size": 4096,20            },21        )22 23    model.validate_credentials(24        model="mistral:text",25        credentials={26            "base_url": os.environ.get("OLLAMA_BASE_URL"),27            "mode": "chat",28            "context_size": 4096,29        },30    )31 32 33def test_invoke_model():34    model = OllamaEmbeddingModel()35 36    result = model.invoke(37        model="mistral:text",38        credentials={39            "base_url": os.environ.get("OLLAMA_BASE_URL"),40            "mode": "chat",41            "context_size": 4096,42        },43        texts=["hello", "world"],44        user="abc-123",45    )46 47    assert isinstance(result, TextEmbeddingResult)48    assert len(result.embeddings) == 249    assert result.usage.total_tokens == 250 51 52def test_get_num_tokens():53    model = OllamaEmbeddingModel()54 55    num_tokens = model.get_num_tokens(56        model="mistral:text",57        credentials={58            "base_url": os.environ.get("OLLAMA_BASE_URL"),59            "mode": "chat",60            "context_size": 4096,61        },62        texts=["hello", "world"],63    )64 65    assert num_tokens == 266