mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-11 03:40:03 +00:00
2a85c36fa9
- Changed dimensions property to avoid forcing provider construction - Added _ensure_model method to construct provider on demand - Modified __call__ to ensure model exists before use - Updated _start to defer provider initialization - Removed eager health check during startup - Added compact embedding serialization with base64 encoding - Implemented batch processing for vector search with heap-based ranking - Added document_ids property to keyword index interface - Updated chunk persistence to handle legacy JSON embeddings - Optimized memory usage by avoiding materialization of metadata in document_ids
67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
"""Tests for lazy AgentScope embedding provider construction."""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
from reme.components.as_embedding import BaseAsEmbedding
|
|
|
|
|
|
class FakeModel:
|
|
"""Minimal async embedding model used to observe construction."""
|
|
|
|
constructions = 0
|
|
|
|
class Parameters:
|
|
"""Accept arbitrary provider parameters."""
|
|
|
|
def __init__(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
|
|
def __init__(self, credential, dimensions, parameters=None, **kwargs):
|
|
type(self).constructions += 1
|
|
self.credential = credential
|
|
self.dimensions = dimensions
|
|
self.parameters = parameters
|
|
self.kwargs = kwargs
|
|
|
|
async def __call__(self, inputs, **_kwargs):
|
|
return SimpleNamespace(embeddings=[[float(index)] * self.dimensions for index, _ in enumerate(inputs)])
|
|
|
|
|
|
class FakeCredential:
|
|
"""Credential boundary that resolves to ``FakeModel``."""
|
|
|
|
def __init__(self, **kwargs):
|
|
self.kwargs = kwargs
|
|
|
|
@staticmethod
|
|
def get_embedding_model_class():
|
|
"""Return the fake provider model class."""
|
|
return FakeModel
|
|
|
|
|
|
class LazyAsEmbedding(BaseAsEmbedding):
|
|
"""Concrete wrapper backed by the local fakes."""
|
|
|
|
credential_cls = FakeCredential
|
|
|
|
|
|
def test_provider_is_constructed_once_on_first_call():
|
|
"""Start and dimension inspection stay local; the first request builds the model."""
|
|
|
|
async def go():
|
|
FakeModel.constructions = 0
|
|
embedding = LazyAsEmbedding(dimensions=3, credential={"token": "test"}, parameters={"mode": "test"})
|
|
|
|
await embedding.start()
|
|
assert embedding.model is None
|
|
assert embedding.dimensions == 3
|
|
assert FakeModel.constructions == 0
|
|
|
|
assert await embedding(["first"]) == [[0.0, 0.0, 0.0]]
|
|
assert await embedding(["second"]) == [[0.0, 0.0, 0.0]]
|
|
assert FakeModel.constructions == 1
|
|
|
|
await embedding.close()
|
|
|
|
asyncio.run(go())
|