src/agentscope/embedding/_gemini/_model.py

# -*- coding: utf-8 -*-
"""The Google Gemini embedding model.

Handles both text-only and multimodal models under a single class.
``gemini-embedding-001`` accepts ``list[str | TextBlock]``.
``gemini-embedding-2`` additionally accepts
:class:`~agentscope.message.DataBlock` (images, video, audio, PDF).
The model name determines the API call style.

Text payloads may be passed either as bare ``str`` or as
:class:`~agentscope.message.TextBlock` — the latter is unpacked to its
``.text`` field on entry so the rest of the pipeline only deals with
``str`` and ``DataBlock``.
"""
from __future__ import annotations

import asyncio
from dataclasses import dataclass
from datetime import datetime
from typing import Any

from .._cache_base import EmbeddingCacheBase
from .._embedding_response import EmbeddingResponse
from .._embedding_usage import EmbeddingUsage
from .._embedding_base import EmbeddingModelBase
from ..._logging import logger
from ...credential import CredentialBase
from ...message import DataBlock, TextBlock

#: Model name prefixes that use the multimodal API path.
_MULTIMODAL_PREFIXES = ("gemini-embedding-2",)


@dataclass
class _MultimodalLimits:
    """Per-request constraints for Gemini multimodal embedding."""

    max_elements: int = 20
    """Maximum total content elements per API call."""

    max_images: int = 6
    """Maximum image elements per API call."""

    max_videos: int = 1
    """Maximum video elements per API call."""

    max_audios: int = 1
    """Maximum audio elements per API call."""

    max_pdfs: int = 1
    """Maximum PDF documents per API call."""


_MODEL_LIMITS: dict[str, _MultimodalLimits] = {
    "gemini-embedding-2": _MultimodalLimits(
        max_elements=20,
        max_images=6,
        max_videos=1,
        max_audios=1,
        max_pdfs=1,
    ),
}

_DEFAULT_LIMITS = _MultimodalLimits()


class GeminiEmbeddingModel(EmbeddingModelBase[str | TextBlock | DataBlock]):
    """Unified Google Gemini embedding model.

    Routes to the text-only or multimodal Gemini API based on the
    model name.

    - **Text mode** (``gemini-embedding-001``): uses the base class's
      simple batch splitting + concurrent retry.  The Gemini API
      accepts a list of strings and returns individual embeddings.
    - **Multimodal mode** (``gemini-embedding-2``): overrides
      ``__call__`` with content-aware batching (respecting per-model
      limits on images, videos, audios, PDFs).  Each input is wrapped
      in a ``Content`` object so the API returns separate embeddings.

    Key API differences from other providers:

    - Dimensions are controlled via ``output_dimensionality`` in the
      ``config`` parameter (not a top-level ``dimensions`` field).
    - ``gemini-embedding-001`` supports ``task_type`` in config;
      ``gemini-embedding-2`` uses prompt prefixes instead.
    """

    #: Text-mode batch size.  Gemini docs don't specify an explicit
    #: limit; we use a conservative default.
    _TEXT_BATCH_SIZE: int = 100

    def __init__(
        self,
        credential: CredentialBase,
        model: str,
        dimensions: int | None,
        parameters: "GeminiEmbeddingModel.Parameters | None" = None,
        embedding_cache: EmbeddingCacheBase | None = None,
        context_size: int = 8192,
        max_retries: int = 3,
        retry_delay: float = 1.0,
    ) -> None:
        """Initialize the Gemini embedding model.

        Args:
            credential (`CredentialBase`):
                A :class:`~agentscope.credential.GeminiCredential`
                instance providing the API key.
            model (`str`):
                The embedding model name (e.g.
                ``"gemini-embedding-001"`` or
                ``"gemini-embedding-2"``).
            dimensions (`int | None`):
                The output embedding vector dimensions.  Required at
                the contract level — see :class:`EmbeddingModelBase`
                for the rationale.  ``None`` is accepted only for
                backward compatibility with legacy configs that
                persisted ``dimensions`` inside ``parameters``.
            parameters (`GeminiEmbeddingModel.Parameters | None`, \
            defaults to ``None``):
                Provider-specific non-dimensional parameters.  Currently
                empty for Gemini.
            embedding_cache (`EmbeddingCacheBase | None`, defaults to \
            ``None``):
                Optional embedding cache.
            context_size (`int`, defaults to ``8192``):
                Maximum input tokens.  2048 for ``gemini-embedding-001``,
                8192 for ``gemini-embedding-2``.
            max_retries (`int`, defaults to ``3``):
                Number of retries on transient failures.
            retry_delay (`float`, defaults to ``1.0``):
                Seconds between retry attempts.
        """
        from google import genai

        self._is_multimodal: bool = model.startswith(_MULTIMODAL_PREFIXES)

        super().__init__(
            credential=credential,
            model=model,
            dimensions=dimensions,
            parameters=parameters,
            context_size=context_size,
            batch_size=self._TEXT_BATCH_SIZE,
            max_retries=max_retries,
            retry_delay=retry_delay,
        )
        self.supports_multimodal = self._is_multimodal

        self.client: genai.Client = genai.Client(
            api_key=credential.api_key.get_secret_value(),
        )
        self.embedding_cache: EmbeddingCacheBase | None = embedding_cache

        if self._is_multimodal:
            self._limits = _MODEL_LIMITS.get(model, _DEFAULT_LIMITS)

    # ------------------------------------------------------------------
    # __call__ — override for multimodal content-aware batching
    # ------------------------------------------------------------------

    async def __call__(
        self,
        inputs: list[str | TextBlock | DataBlock],
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Embed inputs with batching and retry.

        For text models, delegates to the base class.  For multimodal
        models, performs content-aware batching that respects per-model
        limits on images, videos, audios, and PDFs.

        Args:
            inputs (`list[str | TextBlock | DataBlock]`):
                The input data to embed.  ``TextBlock`` items are
                unpacked to their ``.text`` field on entry, so the
                remainder of the pipeline only sees ``str`` and
                ``DataBlock``.
            **kwargs:
                Forwarded to the Gemini API config.

        Returns:
            `EmbeddingResponse`: Merged response for all inputs.
        """
        normalized: list[str | DataBlock] = [
            item.text if isinstance(item, TextBlock) else item
            for item in inputs
        ]

        if not self._is_multimodal:
            return await super().__call__(normalized, **kwargs)

        batches = self._split_multimodal_batches(normalized)

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

        results: list[EmbeddingResponse] = await asyncio.gather(
            *(self._call_with_retry(batch, **kwargs) for batch in batches),
        )

        return self._merge_responses(results)

    def _split_multimodal_batches(
        self,
        inputs: list[str | DataBlock],
    ) -> list[list[str | DataBlock]]:
        """Split inputs into batches respecting Gemini multimodal limits.

        Greedy: keep adding items until any constraint would be
        violated, then start a new batch.

        Args:
            inputs (`list[str | DataBlock]`):
                All inputs to split.

        Returns:
            `list[list[str | DataBlock]]`: List of batches.
        """
        limits = self._limits
        batches: list[list[str | DataBlock]] = []
        current: list[str | DataBlock] = []
        n_elem = 0
        n_img = 0
        n_vid = 0
        n_aud = 0
        n_pdf = 0

        for item in inputs:
            is_img = is_vid = is_aud = is_pdf = False
            if isinstance(item, DataBlock):
                mt = item.source.media_type
                is_img = mt.startswith("image/")
                is_vid = mt.startswith("video/")
                is_aud = mt.startswith("audio/")
                is_pdf = mt == "application/pdf"

            would_exceed = (
                n_elem + 1 > limits.max_elements
                or (is_img and n_img + 1 > limits.max_images)
                or (is_vid and n_vid + 1 > limits.max_videos)
                or (is_aud and n_aud + 1 > limits.max_audios)
                or (is_pdf and n_pdf + 1 > limits.max_pdfs)
            )

            if would_exceed and current:
                batches.append(current)
                current = []
                n_elem = n_img = n_vid = n_aud = n_pdf = 0

            current.append(item)
            n_elem += 1
            n_img += is_img
            n_vid += is_vid
            n_aud += is_aud
            n_pdf += is_pdf

        if current:
            batches.append(current)

        return batches

    # ------------------------------------------------------------------
    # _call_api — single batch
    # ------------------------------------------------------------------

    async def _call_api(
        self,
        inputs: list[str | DataBlock],
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Route to text or multimodal Gemini API for a single batch.

        Args:
            inputs (`list[str | DataBlock]`):
                A single batch of inputs.
            **kwargs:
                Extra keyword arguments merged into the Gemini
                ``EmbedContentConfig``.

        Returns:
            `EmbeddingResponse`: Embedding vectors and usage info.
        """
        if self._is_multimodal:
            return await self._call_multimodal(inputs, **kwargs)
        return await self._call_text(inputs, **kwargs)

    # ------------------------------------------------------------------
    # Text API (gemini-embedding-001)
    # ------------------------------------------------------------------

    async def _call_text(
        self,
        inputs: list[str | DataBlock],
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Call the Gemini text embedding API for a single batch.

        Passes the list of strings directly to ``embed_content``,
        which returns one embedding per string.

        Args:
            inputs (`list[str | DataBlock]`):
                Must all be ``str``; raises ``ValueError`` otherwise.
            **kwargs:
                Merged into ``EmbedContentConfig`` (e.g.
                ``task_type``).

        Returns:
            `EmbeddingResponse`: Embedding vectors and usage info.
        """
        from google.genai import types

        texts: list[str] = []
        for item in inputs:
            if not isinstance(item, str):
                raise ValueError(
                    f"Text embedding model {self.model!r} only accepts "
                    f"str inputs, got {type(item).__name__}.",
                )
            texts.append(item)

        config = types.EmbedContentConfig(
            output_dimensionality=self.dimensions,
            **kwargs,
        )

        cache_key = {
            "model": self.model,
            "contents": texts,
            "output_dimensionality": self.dimensions,
            **kwargs,
        }

        if self.embedding_cache:
            cached = await self.embedding_cache.retrieve(
                identifier=cache_key,
            )
            if cached:
                return EmbeddingResponse(
                    embeddings=cached,
                    usage=EmbeddingUsage(tokens=0, time=0),
                    source="cache",
                )

        start_time = datetime.now()
        response = self.client.models.embed_content(
            model=self.model,
            contents=texts,
            config=config,
        )
        time = (datetime.now() - start_time).total_seconds()

        embeddings = [item.values for item in response.embeddings]

        if self.embedding_cache:
            await self.embedding_cache.store(
                identifier=cache_key,
                embeddings=embeddings,
            )

        return EmbeddingResponse(
            embeddings=embeddings,
            usage=EmbeddingUsage(time=time),
        )

    # ------------------------------------------------------------------
    # Multimodal API (gemini-embedding-2)
    # ------------------------------------------------------------------

    async def _call_multimodal(
        self,
        inputs: list[str | DataBlock],
        **kwargs: Any,
    ) -> EmbeddingResponse:
        """Call the Gemini multimodal embedding API for a single batch.

        Each input is wrapped in a separate ``Content`` object so the
        API returns one embedding per input (not one aggregated
        embedding).

        Args:
            inputs (`list[str | DataBlock]`):
                ``str`` for text, ``DataBlock`` for images / video /
                audio / PDF.
            **kwargs:
                Merged into ``EmbedContentConfig``.

        Returns:
            `EmbeddingResponse`: Embedding vectors and usage info.
        """
        from google.genai import types

        contents: list[types.Content] = []
        for item in inputs:
            if isinstance(item, str):
                contents.append(
                    types.Content(
                        parts=[types.Part.from_text(text=item)],
                    ),
                )
            elif isinstance(item, DataBlock):
                contents.append(
                    types.Content(
                        parts=[self._data_block_to_part(item)],
                    ),
                )
            else:
                raise ValueError(
                    f"Invalid input: {item!r}. Expected str or DataBlock.",
                )

        config = types.EmbedContentConfig(
            output_dimensionality=self.dimensions,
            **kwargs,
        )

        start_time = datetime.now()
        response = self.client.models.embed_content(
            model=self.model,
            contents=contents,
            config=config,
        )
        time = (datetime.now() - start_time).total_seconds()

        embeddings = [item.values for item in response.embeddings]

        return EmbeddingResponse(
            embeddings=embeddings,
            usage=EmbeddingUsage(time=time),
        )

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _data_block_to_part(block: DataBlock) -> Any:
        """Convert a :class:`~agentscope.message.DataBlock` to a Gemini
        ``Part`` object.

        Args:
            block (`DataBlock`):
                A data block with ``Base64Source`` or ``URLSource``.

        Returns:
            A ``google.genai.types.Part`` instance.

        Raises:
            `ValueError`: If the source type is unsupported.
        """
        from google.genai import types
        from ...message import Base64Source, URLSource

        source = block.source

        if isinstance(source, Base64Source):
            import base64

            return types.Part.from_bytes(
                data=base64.b64decode(source.data),
                mime_type=source.media_type,
            )

        if isinstance(source, URLSource):
            # Gemini SDK doesn't have a direct from_url for
            # embed_content; download or use File API.
            # For now, raise — callers should use Base64Source.
            raise ValueError(
                "Gemini embedding API requires inline data "
                "(Base64Source). URLSource is not directly supported "
                f"for embedding. Got URL: {source.url}",
            )

        raise ValueError(
            f"Unsupported source type {type(source).__name__} "
            f"in DataBlock.",
        )