"""Embeddings sub-client for Yarn SDK."""
from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .client import Client


class Embeddings:
    """Generate vector embeddings via the Yarn API.

    Usage::

        client = yarn.Client(api_key="yarn_...")
        result = client.embeddings.create(
            model="text-embedding-3-small",
            input="Hello world",
        )
        vector = result["data"][0]["embedding"]
    """

    def __init__(self, client: Client) -> None:
        self._client = client

    def create(
        self,
        model: str,
        input: str | list[str],
        encoding_format: str = "float",
    ) -> dict:
        """Create embeddings for the given input.

        Args:
            model: Model ID registered in Yarn (e.g. ``text-embedding-3-small``).
            input: A string or list of strings to embed.
            encoding_format: ``"float"`` (default) or ``"base64"``.

        Returns:
            OpenAI-compatible embedding response with ``data``, ``model``,
            and ``usage`` fields.
        """
        return self._client.post(
            "/v1/embeddings",
            json={
                "model": model,
                "input": input,
                "encoding_format": encoding_format,
            },
        )
