mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-12 03:50:59 +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
118 lines
3.3 KiB
Python
118 lines
3.3 KiB
Python
"""AgentScope embedding model wrappers."""
|
|
|
|
from typing import Any
|
|
|
|
from agentscope.credential import (
|
|
CredentialBase,
|
|
DashScopeCredential,
|
|
GeminiCredential,
|
|
OllamaCredential,
|
|
OpenAICredential,
|
|
)
|
|
from agentscope.embedding import (
|
|
EmbeddingModelBase,
|
|
)
|
|
|
|
from ..base_component import BaseComponent
|
|
from ..component_registry import R
|
|
from ...enumeration import ComponentEnum
|
|
|
|
|
|
class BaseAsEmbedding(BaseComponent):
|
|
"""Base wrapper for AgentScope embedding models."""
|
|
|
|
component_type = ComponentEnum.AS_EMBEDDING
|
|
credential_cls: type[CredentialBase]
|
|
|
|
def __init__(self, **kwargs) -> None:
|
|
super().__init__(**kwargs)
|
|
self.model: EmbeddingModelBase[Any] | None = None
|
|
|
|
@property
|
|
def dimensions(self) -> int:
|
|
"""Return configured dimensions without forcing provider construction."""
|
|
if self.model is not None:
|
|
return self.model.dimensions
|
|
dimensions = self.kwargs.get("dimensions")
|
|
if dimensions is None:
|
|
raise RuntimeError("Embedding dimensions are required before provider initialization.")
|
|
return int(dimensions)
|
|
|
|
async def __call__(self, inputs: list[Any], **kwargs) -> list[list[float]]:
|
|
self._ensure_model()
|
|
assert self.model is not None
|
|
response = await self.model(inputs, **kwargs) # pylint: disable=not-callable
|
|
return response.embeddings
|
|
|
|
async def _start(self) -> None:
|
|
"""Defer provider construction until the first remote embedding call."""
|
|
return None
|
|
|
|
def _ensure_model(self) -> None:
|
|
"""Construct the provider on demand while keeping dimensions locally available."""
|
|
if self.model is not None:
|
|
return
|
|
|
|
kwargs = dict(self.kwargs)
|
|
credential = self.credential_cls(**kwargs.pop("credential", {}))
|
|
|
|
model_cls = self.credential_cls.get_embedding_model_class()
|
|
if model_cls is None:
|
|
raise ValueError(f"{self.credential_cls.__name__} does not support embeddings.")
|
|
|
|
dimensions = self.dimensions
|
|
kwargs.pop("dimensions", None)
|
|
params_dict = kwargs.pop("parameters", None)
|
|
parameters = model_cls.Parameters(**params_dict) if params_dict else None
|
|
|
|
self.model = model_cls(
|
|
credential=credential,
|
|
dimensions=dimensions,
|
|
parameters=parameters,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
@R.register("openai")
|
|
class OpenAIAsEmbedding(BaseAsEmbedding):
|
|
"""OpenAI embedding model wrapper."""
|
|
|
|
credential_cls = OpenAICredential
|
|
|
|
|
|
@R.register("dashscope")
|
|
class DashScopeAsEmbedding(BaseAsEmbedding):
|
|
"""DashScope embedding model wrapper."""
|
|
|
|
credential_cls = DashScopeCredential
|
|
|
|
|
|
@R.register("dashscope_multimodal")
|
|
class DashScopeMultiModalAsEmbedding(BaseAsEmbedding):
|
|
"""DashScope multimodal embedding model wrapper."""
|
|
|
|
credential_cls = DashScopeCredential
|
|
|
|
|
|
@R.register("gemini")
|
|
class GeminiAsEmbedding(BaseAsEmbedding):
|
|
"""Gemini embedding model wrapper."""
|
|
|
|
credential_cls = GeminiCredential
|
|
|
|
|
|
@R.register("ollama")
|
|
class OllamaAsEmbedding(BaseAsEmbedding):
|
|
"""Ollama embedding model wrapper."""
|
|
|
|
credential_cls = OllamaCredential
|
|
|
|
|
|
__all__ = [
|
|
"BaseAsEmbedding",
|
|
"OpenAIAsEmbedding",
|
|
"DashScopeAsEmbedding",
|
|
"DashScopeMultiModalAsEmbedding",
|
|
"GeminiAsEmbedding",
|
|
"OllamaAsEmbedding",
|
|
]
|