Spaces:
Paused
Paused
File size: 6,525 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | # -*- coding: utf-8 -*-
# pylint: disable=protected-access
"""Unit tests for GeminiEmbeddingModel."""
from dataclasses import asdict
from typing import Any
from unittest import IsolatedAsyncioTestCase
from unittest.mock import AsyncMock
from utils import AnyValue
from agentscope.embedding import (
GeminiEmbeddingModel,
EmbeddingResponse,
EmbeddingUsage,
)
from agentscope.message import DataBlock, Base64Source
A = AnyValue()
def _mock_resp(embeddings: list[list[float]]) -> EmbeddingResponse:
"""Create a mock EmbeddingResponse."""
return EmbeddingResponse(
embeddings=embeddings,
usage=EmbeddingUsage(tokens=len(embeddings), time=0.01),
)
def _img() -> DataBlock:
"""Create a test image DataBlock."""
return DataBlock(
source=Base64Source(data="aWltYWdl", media_type="image/png"),
)
class GeminiListModelsTest(IsolatedAsyncioTestCase):
"""Test list_models for Gemini."""
async def test_list_models(self) -> None:
"""Should list 2 models."""
cards = GeminiEmbeddingModel.list_models()
names = sorted(c.name for c in cards)
self.assertEqual(names, ["gemini-embedding-001", "gemini-embedding-2"])
async def test_text_model_card(self) -> None:
"""gemini-embedding-001 is text-only with 2048 context."""
cards = GeminiEmbeddingModel.list_models()
card = next(c for c in cards if c.name == "gemini-embedding-001")
self.assertDictEqual(
card.model_dump(),
{
"type": "embedding_model",
"name": "gemini-embedding-001",
"label": "Gemini Embedding 001",
"status": "active",
"input_types": ["text/plain"],
"output_types": ["application/x-embedding"],
"dimensions": 3072,
"supported_dimensions": [3072, 1536, 768, 512, 256, 128],
"context_size": 2048,
"parameter_schema": {
"type": "object",
"properties": {},
"required": [],
},
"parameter_overrides": {},
},
)
async def test_multimodal_model_card(self) -> None:
"""gemini-embedding-2 is multimodal with 8192 context."""
cards = GeminiEmbeddingModel.list_models()
card = next(c for c in cards if c.name == "gemini-embedding-2")
self.assertIn("image/png", card.input_types)
self.assertIn("application/pdf", card.input_types)
self.assertEqual(card.context_size, 8192)
self.assertEqual(card.supported_dimensions, [3072, 1536, 768])
class GeminiTextCallTest(IsolatedAsyncioTestCase):
"""Test Gemini text embedding via mocked _call_text."""
def _make_text_model(self) -> GeminiEmbeddingModel:
"""Create a text-mode model bypassing __init__ (no genai)."""
model = GeminiEmbeddingModel.__new__(GeminiEmbeddingModel)
model.model = "gemini-embedding-001"
model.dimensions = 2
model.context_size = 2048
model.batch_size = 100
model.max_retries = 3
model.retry_delay = 1.0
model._is_multimodal = False
model.embedding_cache = None
return model
async def test_text_call(self) -> None:
"""Text mode delegates to _call_text."""
model = self._make_text_model()
model._call_text = 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_text_rejects_datablock(self) -> None:
"""Text mode rejects DataBlock inputs."""
model = self._make_text_model()
with self.assertRaises(ValueError):
await GeminiEmbeddingModel._call_text(model, [_img()])
class GeminiMultimodalCallTest(IsolatedAsyncioTestCase):
"""Test Gemini multimodal embedding via mocked _call_multimodal."""
def _make_multimodal_model(self) -> GeminiEmbeddingModel:
"""Create a multimodal-mode model bypassing __init__."""
model = GeminiEmbeddingModel.__new__(GeminiEmbeddingModel)
model.model = "gemini-embedding-2"
model.dimensions = 1
model.context_size = 8192
model.batch_size = 100
model.max_retries = 3
model.retry_delay = 1.0
model._is_multimodal = True
model.embedding_cache = None
from agentscope.embedding._gemini._model import _MultimodalLimits
model._limits = _MultimodalLimits(
max_elements=20,
max_images=6,
max_videos=1,
max_audios=1,
max_pdfs=1,
)
return model
async def test_multimodal_delegates(self) -> None:
"""Multimodal mode delegates to _call_multimodal."""
model = self._make_multimodal_model()
model._call_multimodal = AsyncMock(
return_value=_mock_resp([[0.1], [0.2]]),
)
result = await model(["hello", "world"])
self.assertDictEqual(
asdict(result),
{
"embeddings": [[0.1], [0.2]],
"id": A,
"created_at": A,
"type": "embedding",
"usage": {"tokens": 2, "time": 0.01, "type": "embedding"},
"source": "api",
},
)
async def test_multimodal_batching_by_image_limit(self) -> None:
"""8 images with max_images=6 produces 2 batches (6+2)."""
model = self._make_multimodal_model()
call_count = 0
async def _mock(inputs: list, **_kw: Any) -> EmbeddingResponse:
nonlocal call_count
call_count += 1
return _mock_resp([[0.1]] * len(inputs))
model._call_multimodal = _mock # type: ignore[assignment]
result = await model([_img() for _ in range(8)])
self.assertEqual(result["embeddings"], [[0.1]] * 8)
self.assertEqual(call_count, 2)
|