src/agentscope/embedding/_embedding_base.py

# -*- coding: utf-8 -*-
"""The embedding model base class."""
from __future__ import annotations

import asyncio
import inspect
from abc import abstractmethod
from pathlib import Path
from typing import Any, Generic, TypeVar, Type, Union

from pydantic import BaseModel, ConfigDict

from ._embedding_model_card import EmbeddingModelCard
from ._embedding_response import EmbeddingResponse
from ._embedding_usage import EmbeddingUsage
from .._logging import logger
from ..credential import CredentialBase
from ..message import DataBlock, TextBlock

#: Type variable for embedding input elements.
#:
#: Bound to the union of all element shapes the framework supports;
#: each concrete subclass narrows it to its accepted input via
#: ``class Foo(EmbeddingModelBase[str | TextBlock]): ...`` so callers
#: get accurate IDE completion and type-checking on ``__call__``.
#:
#: ``__call__`` accepts ``TextBlock`` and unpacks it to ``.text``
#: before invoking :meth:`_call_api`; therefore ``_call_api``'s input
#: type is intentionally decoupled from :data:`InputT` (it is typed
#: ``list[Any]`` on the base and narrowed by each subclass).
InputT = TypeVar("InputT", bound=Union[str, TextBlock, DataBlock])


class EmbeddingModelBase(Generic[InputT]):
    """Base class for embedding models.

    Generic over :data:`InputT` so that text-only subclasses
    (``EmbeddingModelBase[str]``) and multimodal subclasses
    (``EmbeddingModelBase[str | TextBlock | DataBlock]``) expose the
    correct ``inputs`` type to the IDE.

    Follows the same pattern as :class:`~agentscope.model.ChatModelBase`:

    - ``__call__`` splits inputs into batches of size
      :attr:`batch_size`, calls :meth:`_call_api` for each batch
      **concurrently** via :func:`asyncio.gather`, and merges the
      results.  Each batch call is wrapped with retry logic.
    - Subclasses only implement :meth:`_call_api` for a **single
      batch** — no batching or retry code needed.
    - Each subclass may override :meth:`_get_retryable_exceptions` to
      declare provider-specific retriable errors.
    """

    class Parameters(BaseModel):
        """Provider-specific tunables for embedding models.

        Intentionally empty in the base — ``dimensions`` is a contract
        property of an embedding model (its output vector size), not a
        tunable knob, so it lives directly on the instance via the
        required :paramref:`__init__.dimensions` argument.  Subclasses
        extend this class to expose **non-dimensional** knobs (e.g.
        Gemini's ``task_type``, Dashscope's ``text_type``).

        ``extra="allow"`` is set so old persisted configs (where
        ``dimensions`` lived in ``parameters``) keep deserialising
        without raising; :meth:`EmbeddingModelBase.__init__` extracts
        it back out for backward compatibility.
        """

        model_config = ConfigDict(extra="allow")

    credential: CredentialBase
    """The API credential."""

    model: str
    """The embedding model name."""

    dimensions: int
    """The output embedding vector dimensions.

    Set directly from the :paramref:`__init__.dimensions` argument —
    a required, first-class field rather than something derived from
    :attr:`parameters`.
    """

    context_size: int
    """Maximum input length (in tokens) per single input item."""

    batch_size: int
    """Maximum number of input items per API call."""

    max_retries: int
    """The maximum number of retries for the underlying API."""

    retry_delay: float
    """Seconds to sleep between retry attempts."""

    supports_multimodal: bool = False
    """Whether this model instance accepts :class:`DataBlock` inputs in
    addition to text.  Text-only models keep the default ``False``;
    multimodal subclasses must set it to ``True`` (per instance when
    routing depends on the model name)."""

    def __init__(
        self,
        credential: CredentialBase,
        model: str,
        dimensions: int | None,
        parameters: BaseModel | None,
        context_size: int,
        batch_size: int,
        max_retries: int,
        retry_delay: float,
    ) -> None:
        """Initialize the embedding model base class.

        Args:
            credential (`CredentialBase`):
                The API credential used for authentication.
            model (`str`):
                The name of the embedding model.
            dimensions (`int | None`):
                The output embedding vector dimensions for this
                instance.  Required and first-class — see the class
                docstring for the rationale of keeping ``dimensions``
                outside :class:`Parameters`.  For backward compatibility
                with older configs that stored ``dimensions`` inside
                :class:`Parameters`, ``None`` is accepted at the
                signature level and is back-filled from
                ``parameters.dimensions`` when present.
            parameters (`BaseModel | None`):
                Provider-specific non-dimensional parameters.  When
                ``None``, the default ``Parameters()`` is used.
            context_size (`int`):
                Maximum input length (in tokens) per single input item.
            batch_size (`int`):
                Maximum number of input items per API call.  When
                ``__call__`` receives more items, it splits them into
                batches and calls :meth:`_call_api` concurrently.
            max_retries (`int`):
                The maximum number of retries for each batch API call.
                Only exceptions listed in
                :meth:`_get_retryable_exceptions` count against this
                budget.
            retry_delay (`float`):
                Seconds to sleep between retry attempts.
        """
        resolved_parameters = parameters or self.Parameters()
        # Backward-compat: older session/KB configs persisted
        # ``dimensions`` inside ``parameters``.  Promote it to the
        # constructor argument when the caller did not pass one
        # explicitly, then strip it from the parameters object so it
        # never reaches provider-specific request payloads.
        param_dump = resolved_parameters.model_dump()
        legacy_dimensions = param_dump.pop("dimensions", None)
        if dimensions is None:
            if legacy_dimensions is None:
                raise ValueError(
                    "dimensions is required: pass it explicitly to "
                    "EmbeddingModelBase.__init__ or include it in the "
                    "legacy `parameters` mapping.",
                )
            dimensions = int(legacy_dimensions)
            resolved_parameters = type(resolved_parameters)(**param_dump)
        elif legacy_dimensions is not None:
            # Both routes set it — explicit constructor wins, strip the
            # legacy mirror so it can't drift.
            resolved_parameters = type(resolved_parameters)(**param_dump)

        if dimensions <= 0:
            raise ValueError(
                f"dimensions must be a positive integer, got {dimensions}.",
            )

        self.credential = credential
        self.model = model
        self.dimensions = dimensions
        self.parameters = resolved_parameters
        self.context_size = context_size
        self.batch_size = batch_size
        self.max_retries = max_retries
        self.retry_delay = retry_delay

    @classmethod
    def _get_retryable_exceptions(cls) -> tuple[Type[Exception], ...]:
        """Return exception types that should trigger a retry.

        Defaults to an empty tuple (no retries).  Subclasses can
        override to declare provider-specific retryable exceptions.
        """
        return ()

    # ------------------------------------------------------------------
    # Public API — batching + concurrent retry
    # ------------------------------------------------------------------

    async def __call__(
        self,
        inputs: list[InputT],
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Embed a list of inputs with automatic batching and retry.

        The inputs are split into chunks of :attr:`batch_size`.  All
        chunks are dispatched **concurrently** via
        :func:`asyncio.gather`.  Each chunk is individually retried up
        to ``max_retries`` times on retryable errors.  Results are
        merged into a single :class:`EmbeddingResponse` preserving the
        original input order.

        Args:
            inputs (`list[InputT]`):
                The input data to embed.  For text-only models this is
                ``list[str]``; for multimodal models it is
                ``list[str | TextBlock | DataBlock]``.  Any
                :class:`TextBlock` items are transparently unpacked to
                their ``.text`` field on entry, so subclasses' batching
                and ``_call_api`` only have to handle ``str`` (and
                ``DataBlock`` for multimodal variants).
            **kwargs:
                Additional keyword arguments forwarded to
                :meth:`_call_api`.

        Returns:
            `EmbeddingResponse`:
                A merged response containing embeddings for all inputs.
        """
        if not inputs:
            return EmbeddingResponse(
                embeddings=[],
                usage=EmbeddingUsage(tokens=0, time=0),
            )

        normalized: list[Any] = [
            item.text if isinstance(item, TextBlock) else item
            for item in inputs
        ]

        # Split into batches.
        batches = [
            normalized[i : i + self.batch_size]
            for i in range(0, len(normalized), self.batch_size)
        ]

        if len(batches) > 1:
            logger.info(
                "Embedding %d inputs in %d batches (batch_size=%d) "
                "for model %s.",
                len(normalized),
                len(batches),
                self.batch_size,
                self.model,
            )

        # Dispatch all batches concurrently, each with retry.
        results: list[EmbeddingResponse] = await asyncio.gather(
            *(self._call_with_retry(batch, **kwargs) for batch in batches),
        )

        return self._merge_responses(results)

    # ------------------------------------------------------------------
    # Internal — merge multiple batch responses
    # ------------------------------------------------------------------

    @staticmethod
    def _merge_responses(
        responses: list[EmbeddingResponse],
    ) -> EmbeddingResponse:
        """Merge multiple batch :class:`EmbeddingResponse` objects into
        one, preserving input order.

        Args:
            responses (`list[EmbeddingResponse]`):
                Batch responses to merge.

        Returns:
            `EmbeddingResponse`: The merged response.
        """
        if len(responses) == 1:
            return responses[0]

        all_embeddings: list = []
        total_tokens = 0
        total_time = 0.0

        for resp in responses:
            all_embeddings.extend(resp.embeddings)
            if resp.usage:
                total_time += resp.usage.time
                if resp.usage.tokens:
                    total_tokens += resp.usage.tokens

        return EmbeddingResponse(
            embeddings=all_embeddings,
            usage=EmbeddingUsage(
                tokens=total_tokens,
                time=total_time,
            ),
            source="api",
        )

    # ------------------------------------------------------------------
    # Internal — retry wrapper for a single batch
    # ------------------------------------------------------------------

    async def _call_with_retry(
        self,
        inputs: list[Any],
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Call :meth:`_call_api` with retry logic for a single batch.

        Args:
            inputs (`list[Any]`):
                A single batch of inputs (size ≤ ``batch_size``), already
                normalised by :meth:`__call__` (any :class:`TextBlock`
                items unpacked to their ``.text``). Typed as
                ``list[Any]`` because the concrete element shape depends
                on the subclass — see :meth:`_call_api`.
            **kwargs:
                Forwarded to :meth:`_call_api`.
        """
        retryable = tuple(self._get_retryable_exceptions())
        last_error: Exception | None = None

        for attempt in range(self.max_retries + 1):
            try:
                return await self._call_api(inputs, **kwargs)
            except Exception as e:
                if not isinstance(e, retryable):
                    raise
                last_error = e
                if attempt < self.max_retries:
                    logger.warning(
                        "Batch attempt %d failed for embedding model "
                        "%s: %s. Retrying in %.1fs...",
                        attempt + 1,
                        self.model,
                        str(e),
                        self.retry_delay,
                    )
                    await asyncio.sleep(self.retry_delay)
                else:
                    logger.warning(
                        "All %d attempt(s) failed for a batch of "
                        "embedding model %s.",
                        self.max_retries + 1,
                        self.model,
                    )

        if last_error is not None:
            raise last_error
        raise RuntimeError(
            f"Failed to call embedding model {self.model} after "
            f"{self.max_retries + 1} retries.",
        )

    # ------------------------------------------------------------------
    # Abstract — subclasses implement this for a single batch
    # ------------------------------------------------------------------

    @abstractmethod
    async def _call_api(
        self,
        inputs: list[Any],
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Call the underlying embedding API for a **single batch**.

        Subclasses must implement this method.  The batch splitting,
        concurrency, and retry logic are handled by :meth:`__call__`
        — this method only needs to handle one API call.

        .. note::
            The parameter is typed ``list[Any]`` rather than
            ``list[InputT]`` because :meth:`__call__` unpacks
            :class:`TextBlock` items to their ``.text`` field *before*
            dispatching to this method.  The element shape this method
            actually receives is therefore subclass-specific
            (:data:`InputT` minus :class:`TextBlock`).  Subclasses
            should override with their concrete narrower type, e.g.
            ``list[str]`` for text-only models or
            ``list[str | DataBlock]`` for multimodal ones.

        Args:
            inputs (`list[Any]`):
                A batch of inputs (guaranteed ``len(inputs) <=
                self.batch_size``).
            **kwargs:
                Additional keyword arguments.

        Returns:
            `EmbeddingResponse`:
                The embedding response for this batch.
        """

    # ------------------------------------------------------------------
    # Model card discovery
    # ------------------------------------------------------------------

    @classmethod
    def list_models(
        cls,
        custom_yaml_dir: str | None = None,
    ) -> list[EmbeddingModelCard]:
        """List candidate embedding models from YAML files.

        Each concrete subclass should live in its own provider
        subdirectory (e.g. ``embedding/_openai/_model.py``) with a
        sibling ``_models/`` directory containing YAML files — identical
        to the layout used by :class:`~agentscope.model.ChatModelBase`.

        Args:
            custom_yaml_dir (`str | None`):
                Override the YAML directory.

        Returns:
            `list[EmbeddingModelCard]`:
                A list of embedding model cards.
        """
        if custom_yaml_dir is None:
            subclass_file = Path(inspect.getfile(cls))
            yaml_dir = subclass_file.parent / "_models"
        else:
            yaml_dir = Path(custom_yaml_dir)

        if not yaml_dir.is_dir():
            return []

        yaml_files = list(yaml_dir.glob("*.yaml"))

        model_cards = []
        for yaml_file in yaml_files:
            try:
                card = EmbeddingModelCard.from_yaml(
                    yaml_path=str(yaml_file),
                    parameter_class=cls.Parameters,
                )
                model_cards.append(card)
            except Exception as e:
                logger.warning(
                    "Failed to load embedding model card %s: %s",
                    yaml_file,
                    str(e),
                )
                continue

        return model_cards