ReMe/reme/components/as_embedding/__init__.py

"""AgentScope embedding model wrappers."""

import hashlib
import os
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]
    endpoint_fields = ("base_url", "host")
    endpoint_env: str | None = None
    default_endpoint = ""

    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)

    @property
    def vector_space(self) -> tuple[str, ...]:
        """Return the fields that make vectors from two embedding setups incompatible.

        The wrapper identifies the provider consistently before and after lazy model
        construction.  Model details come from the live provider when one has been
        injected at runtime, otherwise they come from the configured kwargs.
        """
        if self.model is not None:
            return (
                self.backend or self.credential_cls.__name__,
                str(getattr(self.model, "model", self.kwargs.get("model") or "")),
                str(self.dimensions),
                self._model_endpoint(),
            )
        return (
            self.backend or self.credential_cls.__name__,
            str(self.kwargs.get("model") or ""),
            str(self.dimensions),
            self._configured_endpoint(),
        )

    @property
    def vector_space_id(self) -> str:
        """Return a short digest of :attr:`vector_space` for naming persisted vectors.

        Cache consumers use this digest to avoid reusing vectors produced by a
        different embedding setup.
        """
        return hashlib.sha256("\x1f".join(self.vector_space).encode()).hexdigest()[:12]

    @staticmethod
    def _endpoint(credential: Any) -> str:
        """Read the provider endpoint from a credential object or a raw kwargs dict."""
        for field in ("base_url", "host"):
            value = credential.get(field) if isinstance(credential, dict) else getattr(credential, field, None)
            if value:
                return str(value).rstrip("/")
        return ""

    def _configured_endpoint(self) -> str:
        """Resolve the provider endpoint without constructing it eagerly."""
        credential = self.kwargs.get("credential")
        if isinstance(credential, dict):
            fields = getattr(self.credential_cls, "model_fields", {})
            for name in self.endpoint_fields:
                if name in credential:
                    value = credential[name]
                    if value is not None:
                        return str(value).rstrip("/")
                    continue
                field = fields.get(name)
                if field is not None and not field.is_required():
                    value = field.get_default(call_default_factory=True)
                    if value:
                        return str(value).rstrip("/")
        else:
            endpoint = self._endpoint(credential)
            if endpoint:
                return endpoint
        if self.endpoint_env:
            endpoint = os.environ.get(self.endpoint_env)
            if endpoint:
                return endpoint.rstrip("/")
        return self.default_endpoint

    def _model_endpoint(self) -> str:
        """Return the endpoint used by the constructed provider."""
        assert self.model is not None
        endpoint = self._endpoint(getattr(self.model, "credential", self.kwargs.get("credential")))
        return endpoint or self._configured_endpoint()

    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

    def initialize_model(self) -> None:
        """Construct the provider without making a remote request.

        Callers that apply their own request timeout can initialize first so
        one-time SDK imports and client construction do not consume that
        timeout budget. Normal embedding calls remain lazily initialized.
        """
        self._ensure_model()

    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
    endpoint_fields = ("base_url",)
    endpoint_env = "OPENAI_BASE_URL"
    default_endpoint = "https://api.openai.com/v1"

    def _model_endpoint(self) -> str:
        """Read the endpoint the OpenAI client actually resolved."""
        assert self.model is not None
        client = getattr(self.model, "client", None)
        base_url = getattr(client, "base_url", None)
        return str(base_url).rstrip("/") if base_url is not None else super()._model_endpoint()


@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
    endpoint_fields = ("host",)
    endpoint_env = "OLLAMA_HOST"
    default_endpoint = "http://127.0.0.1:11434"


__all__ = [
    "BaseAsEmbedding",
    "OpenAIAsEmbedding",
    "DashScopeAsEmbedding",
    "DashScopeMultiModalAsEmbedding",
    "GeminiAsEmbedding",
    "OllamaAsEmbedding",
]