Spaces:
Paused
Paused
File size: 3,434 Bytes
4b03eed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | # -*- coding: utf-8 -*-
# pylint: disable=protected-access
"""Unit tests for OllamaEmbeddingModel."""
from dataclasses import asdict
from typing import Any
from unittest import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock
from utils import AnyValue
from agentscope.credential import OllamaCredential
from agentscope.embedding import (
OllamaEmbeddingModel,
EmbeddingResponse,
EmbeddingUsage,
)
A = AnyValue()
def _cred() -> OllamaCredential:
"""Create a test credential."""
return OllamaCredential(host="http://localhost:11434")
def _mock_resp(embeddings: list[list[float]]) -> EmbeddingResponse:
"""Create a mock EmbeddingResponse."""
return EmbeddingResponse(
embeddings=embeddings,
usage=EmbeddingUsage(tokens=len(embeddings), time=0.01),
)
class OllamaListModelsTest(IsolatedAsyncioTestCase):
"""Test list_models for Ollama."""
async def test_list_models_empty(self) -> None:
"""Ollama has no pre-defined YAMLs, returns empty list."""
self.assertEqual(OllamaEmbeddingModel.list_models(), [])
class OllamaEmbeddingCallTest(IsolatedAsyncioTestCase):
"""Test Ollama embedding via mocked _call_api."""
async def test_basic_call(self) -> None:
"""Basic call returns correct embeddings."""
model = OllamaEmbeddingModel(
credential=_cred(),
model="nomic-embed-text",
dimensions=2,
)
model._call_api = AsyncMock(
return_value=_mock_resp([[0.1, 0.2], [0.3, 0.4]]),
)
result = await model(["hello", "world"])
self.assertDictEqual(
asdict(result),
{
"embeddings": [[0.1, 0.2], [0.3, 0.4]],
"id": A,
"created_at": A,
"type": "embedding",
"usage": {"tokens": 2, "time": 0.01, "type": "embedding"},
"source": "api",
},
)
async def test_dimensions_and_host(self) -> None:
"""Dimensions and host are set correctly from constructor."""
model = OllamaEmbeddingModel(
credential=OllamaCredential(host="http://gpu:11434"),
model="test",
dimensions=768,
)
self.assertEqual(model.dimensions, 768)
self.assertEqual(model.host, "http://gpu:11434")
async def test_multi_batch(self) -> None:
"""Batching splits inputs correctly."""
model = OllamaEmbeddingModel(
credential=_cred(),
model="test",
dimensions=1,
)
model.batch_size = 2
call_count = 0
async def _mock(inputs: list[str], **_kw: Any) -> EmbeddingResponse:
nonlocal call_count
call_count += 1
return _mock_resp([[0.1]] * len(inputs))
model._call_api = _mock # type: ignore[assignment]
result = await model(["a", "b", "c"])
self.assertDictEqual(
asdict(result),
{
"embeddings": [[0.1], [0.1], [0.1]],
"id": A,
"created_at": A,
"type": "embedding",
"usage": {"tokens": A, "time": A, "type": "embedding"},
"source": "api",
},
)
self.assertEqual(call_count, 2)
|