ai-lib-python 0.5.0__py3-none-any.whl

ai_lib_python/embeddings/client.py

@@ -0,0 +1,339 @@
+"""
+Embedding client for generating embeddings.
+
+Provides a unified interface for embedding generation.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from ai_lib_python.embeddings.types import (
+    Embedding,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    EmbeddingUsage,
+)
+from ai_lib_python.protocol import ProtocolLoader
+from ai_lib_python.transport import HttpTransport
+
+if TYPE_CHECKING:
+    from ai_lib_python.protocol.manifest import ProtocolManifest
+
+
+class EmbeddingClient:
+    """Client for generating embeddings.
+
+    Supports batch embedding generation with automatic chunking
+    and rate limit handling.
+
+    Example:
+        >>> client = await EmbeddingClient.create("openai/text-embedding-3-small")
+        >>> response = await client.embed("Hello, world!")
+        >>> print(response.first.vector[:5])
+
+        >>> # Batch embedding
+        >>> texts = ["Hello", "World", "Test"]
+        >>> response = await client.embed_batch(texts)
+        >>> for emb in response.embeddings:
+        ...     print(f"Text {emb.index}: {len(emb.vector)} dimensions")
+    """
+
+    def __init__(
+        self,
+        manifest: ProtocolManifest,
+        transport: HttpTransport,
+        model_id: str,
+    ) -> None:
+        """Initialize embedding client.
+
+        Args:
+            manifest: Protocol manifest
+            transport: HTTP transport
+            model_id: Model identifier
+        """
+        self._manifest = manifest
+        self._transport = transport
+        self._model_id = model_id
+
+    @classmethod
+    async def create(
+        cls,
+        model: str,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        dimensions: int | None = None,
+    ) -> EmbeddingClient:
+        """Create an embedding client.
+
+        Args:
+            model: Model identifier (e.g., "openai/text-embedding-3-small")
+            api_key: API key (optional, uses environment)
+            base_url: Base URL override
+            dimensions: Output dimensions override
+
+        Returns:
+            EmbeddingClient instance
+        """
+        return await (
+            EmbeddingClientBuilder()
+            .model(model)
+            .api_key(api_key)
+            .base_url(base_url)
+            .dimensions(dimensions)
+            .build()
+        )
+
+    @classmethod
+    def builder(cls) -> EmbeddingClientBuilder:
+        """Get a builder for creating embedding clients.
+
+        Returns:
+            EmbeddingClientBuilder instance
+        """
+        return EmbeddingClientBuilder()
+
+    async def embed(
+        self,
+        text: str,
+        dimensions: int | None = None,
+    ) -> EmbeddingResponse:
+        """Generate embedding for a single text.
+
+        Args:
+            text: Text to embed
+            dimensions: Output dimensions (if supported)
+
+        Returns:
+            EmbeddingResponse with single embedding
+        """
+        request = EmbeddingRequest(
+            input=text,
+            model=self._model_id,
+            dimensions=dimensions,
+        )
+        return await self._execute(request)
+
+    async def embed_batch(
+        self,
+        texts: list[str],
+        dimensions: int | None = None,
+        batch_size: int = 100,
+    ) -> EmbeddingResponse:
+        """Generate embeddings for multiple texts.
+
+        Args:
+            texts: List of texts to embed
+            dimensions: Output dimensions (if supported)
+            batch_size: Maximum texts per API call
+
+        Returns:
+            EmbeddingResponse with all embeddings
+        """
+        if len(texts) <= batch_size:
+            request = EmbeddingRequest(
+                input=texts,
+                model=self._model_id,
+                dimensions=dimensions,
+            )
+            return await self._execute(request)
+
+        # Process in batches
+        all_embeddings: list[Embedding] = []
+        total_usage = EmbeddingUsage()
+
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i : i + batch_size]
+            request = EmbeddingRequest(
+                input=batch,
+                model=self._model_id,
+                dimensions=dimensions,
+            )
+            response = await self._execute(request)
+
+            # Adjust indices for combined result
+            for emb in response.embeddings:
+                all_embeddings.append(
+                    Embedding(
+                        index=i + emb.index,
+                        vector=emb.vector,
+                        object_type=emb.object_type,
+                    )
+                )
+
+            total_usage.prompt_tokens += response.usage.prompt_tokens
+            total_usage.total_tokens += response.usage.total_tokens
+
+        return EmbeddingResponse(
+            embeddings=all_embeddings,
+            model=self._model_id,
+            usage=total_usage,
+        )
+
+    async def _execute(self, request: EmbeddingRequest) -> EmbeddingResponse:
+        """Execute an embedding request.
+
+        Args:
+            request: Embedding request
+
+        Returns:
+            EmbeddingResponse
+        """
+        endpoint = self._get_embedding_endpoint()
+        payload = request.to_dict()
+
+        response = await self._transport.post(endpoint, json=payload)
+        data = response.json()
+
+        return EmbeddingResponse.from_openai_format(data)
+
+    def _get_embedding_endpoint(self) -> str:
+        """Get the embedding API endpoint.
+
+        Returns:
+            Endpoint path
+        """
+        # Try to get from manifest, default to OpenAI-style
+        if hasattr(self._manifest, "embedding_endpoint"):
+            return self._manifest.embedding_endpoint
+        return "/v1/embeddings"
+
+    @property
+    def model(self) -> str:
+        """Get the model identifier."""
+        return self._model_id
+
+    @property
+    def provider(self) -> str:
+        """Get the provider identifier."""
+        return self._manifest.id
+
+    async def close(self) -> None:
+        """Close the client."""
+        await self._transport.close()
+
+    async def __aenter__(self) -> EmbeddingClient:
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(
+        self, exc_type: Any, exc_val: Any, exc_tb: Any
+    ) -> None:
+        """Async context manager exit."""
+        await self.close()
+
+
+class EmbeddingClientBuilder:
+    """Builder for creating EmbeddingClient instances.
+
+    Example:
+        >>> client = await (
+        ...     EmbeddingClient.builder()
+        ...     .model("openai/text-embedding-3-small")
+        ...     .dimensions(512)
+        ...     .build()
+        ... )
+    """
+
+    def __init__(self) -> None:
+        """Initialize builder."""
+        self._model: str | None = None
+        self._api_key: str | None = None
+        self._base_url: str | None = None
+        self._dimensions: int | None = None
+        self._timeout: float | None = None
+
+    def model(self, model: str) -> EmbeddingClientBuilder:
+        """Set the model.
+
+        Args:
+            model: Model identifier (e.g., "openai/text-embedding-3-small")
+
+        Returns:
+            Self for chaining
+        """
+        self._model = model
+        return self
+
+    def api_key(self, api_key: str | None) -> EmbeddingClientBuilder:
+        """Set the API key.
+
+        Args:
+            api_key: API key
+
+        Returns:
+            Self for chaining
+        """
+        self._api_key = api_key
+        return self
+
+    def base_url(self, base_url: str | None) -> EmbeddingClientBuilder:
+        """Set the base URL.
+
+        Args:
+            base_url: Base URL override
+
+        Returns:
+            Self for chaining
+        """
+        self._base_url = base_url
+        return self
+
+    def dimensions(self, dimensions: int | None) -> EmbeddingClientBuilder:
+        """Set the output dimensions.
+
+        Args:
+            dimensions: Output dimensions
+
+        Returns:
+            Self for chaining
+        """
+        self._dimensions = dimensions
+        return self
+
+    def timeout(self, timeout: float) -> EmbeddingClientBuilder:
+        """Set the request timeout.
+
+        Args:
+            timeout: Timeout in seconds
+
+        Returns:
+            Self for chaining
+        """
+        self._timeout = timeout
+        return self
+
+    async def build(self) -> EmbeddingClient:
+        """Build the embedding client.
+
+        Returns:
+            EmbeddingClient instance
+
+        Raises:
+            ValueError: If model is not set
+        """
+        if not self._model:
+            raise ValueError("Model must be set before building")
+
+        # Parse model identifier
+        parts = self._model.split("/")
+        provider_id = parts[0] if len(parts) >= 2 else "openai"
+        model_id = parts[-1]
+
+        # Load protocol manifest
+        loader = ProtocolLoader()
+        manifest = await loader.load_provider(provider_id)
+
+        # Create transport
+        transport = HttpTransport.from_manifest(
+            manifest,
+            api_key=self._api_key,
+            base_url_override=self._base_url,
+            timeout=self._timeout,
+        )
+
+        return EmbeddingClient(
+            manifest=manifest,
+            transport=transport,
+            model_id=model_id,
+        )

ai_lib_python/embeddings/types.py

@@ -0,0 +1,234 @@
+"""
+Embedding types and data models.
+
+Defines the core types for embedding operations.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class EmbeddingModel(str, Enum):
+    """Standard embedding models."""
+
+    # OpenAI
+    TEXT_EMBEDDING_3_SMALL = "text-embedding-3-small"
+    TEXT_EMBEDDING_3_LARGE = "text-embedding-3-large"
+    TEXT_EMBEDDING_ADA_002 = "text-embedding-ada-002"
+
+    # Cohere
+    EMBED_ENGLISH_V3 = "embed-english-v3.0"
+    EMBED_MULTILINGUAL_V3 = "embed-multilingual-v3.0"
+    EMBED_ENGLISH_LIGHT_V3 = "embed-english-light-v3.0"
+
+    # Voyage
+    VOYAGE_LARGE_2 = "voyage-large-2"
+    VOYAGE_CODE_2 = "voyage-code-2"
+
+    # Google
+    TEXT_EMBEDDING_004 = "text-embedding-004"
+
+    @property
+    def dimensions(self) -> int:
+        """Get default dimensions for the model."""
+        dimensions_map = {
+            "text-embedding-3-small": 1536,
+            "text-embedding-3-large": 3072,
+            "text-embedding-ada-002": 1536,
+            "embed-english-v3.0": 1024,
+            "embed-multilingual-v3.0": 1024,
+            "embed-english-light-v3.0": 384,
+            "voyage-large-2": 1536,
+            "voyage-code-2": 1536,
+            "text-embedding-004": 768,
+        }
+        return dimensions_map.get(self.value, 1536)
+
+    @property
+    def max_tokens(self) -> int:
+        """Get maximum input tokens for the model."""
+        max_tokens_map = {
+            "text-embedding-3-small": 8191,
+            "text-embedding-3-large": 8191,
+            "text-embedding-ada-002": 8191,
+            "embed-english-v3.0": 512,
+            "embed-multilingual-v3.0": 512,
+            "embed-english-light-v3.0": 512,
+            "voyage-large-2": 16000,
+            "voyage-code-2": 16000,
+            "text-embedding-004": 2048,
+        }
+        return max_tokens_map.get(self.value, 8191)
+
+
+@dataclass
+class EmbeddingUsage:
+    """Token usage information for embedding request.
+
+    Attributes:
+        prompt_tokens: Number of tokens in the input
+        total_tokens: Total tokens used
+    """
+
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+
+    def to_dict(self) -> dict[str, int]:
+        """Convert to dictionary."""
+        return {
+            "prompt_tokens": self.prompt_tokens,
+            "total_tokens": self.total_tokens,
+        }
+
+
+@dataclass
+class Embedding:
+    """A single embedding result.
+
+    Attributes:
+        index: Index in the batch
+        vector: The embedding vector
+        object_type: Object type (always "embedding")
+    """
+
+    index: int
+    vector: list[float]
+    object_type: str = "embedding"
+
+    @property
+    def dimensions(self) -> int:
+        """Get the dimensionality of the embedding."""
+        return len(self.vector)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "index": self.index,
+            "embedding": self.vector,
+            "object": self.object_type,
+        }
+
+    @classmethod
+    def from_openai_format(cls, data: dict[str, Any]) -> Embedding:
+        """Create from OpenAI API response format.
+
+        Args:
+            data: OpenAI embedding object
+
+        Returns:
+            Embedding instance
+        """
+        return cls(
+            index=data.get("index", 0),
+            vector=data.get("embedding", []),
+            object_type=data.get("object", "embedding"),
+        )
+
+
+@dataclass
+class EmbeddingRequest:
+    """Request for generating embeddings.
+
+    Attributes:
+        input: Text or list of texts to embed
+        model: Model to use
+        dimensions: Output dimensions (if supported)
+        encoding_format: Output format ("float" or "base64")
+        user: User identifier for tracking
+    """
+
+    input: str | list[str]
+    model: str = "text-embedding-3-small"
+    dimensions: int | None = None
+    encoding_format: str = "float"
+    user: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to API request format."""
+        data: dict[str, Any] = {
+            "input": self.input,
+            "model": self.model,
+        }
+        if self.dimensions is not None:
+            data["dimensions"] = self.dimensions
+        if self.encoding_format != "float":
+            data["encoding_format"] = self.encoding_format
+        if self.user:
+            data["user"] = self.user
+        return data
+
+    @property
+    def is_batch(self) -> bool:
+        """Check if this is a batch request."""
+        return isinstance(self.input, list)
+
+    @property
+    def batch_size(self) -> int:
+        """Get the number of inputs."""
+        if isinstance(self.input, list):
+            return len(self.input)
+        return 1
+
+
+@dataclass
+class EmbeddingResponse:
+    """Response from embedding generation.
+
+    Attributes:
+        embeddings: List of embedding results
+        model: Model used
+        usage: Token usage information
+        object_type: Object type (always "list")
+    """
+
+    embeddings: list[Embedding] = field(default_factory=list)
+    model: str = ""
+    usage: EmbeddingUsage = field(default_factory=EmbeddingUsage)
+    object_type: str = "list"
+
+    @property
+    def first(self) -> Embedding | None:
+        """Get the first embedding."""
+        return self.embeddings[0] if self.embeddings else None
+
+    @property
+    def vectors(self) -> list[list[float]]:
+        """Get all vectors."""
+        return [e.vector for e in self.embeddings]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "object": self.object_type,
+            "data": [e.to_dict() for e in self.embeddings],
+            "model": self.model,
+            "usage": self.usage.to_dict(),
+        }
+
+    @classmethod
+    def from_openai_format(cls, data: dict[str, Any]) -> EmbeddingResponse:
+        """Create from OpenAI API response format.
+
+        Args:
+            data: OpenAI embeddings response
+
+        Returns:
+            EmbeddingResponse instance
+        """
+        embeddings = [
+            Embedding.from_openai_format(e) for e in data.get("data", [])
+        ]
+        usage_data = data.get("usage", {})
+        usage = EmbeddingUsage(
+            prompt_tokens=usage_data.get("prompt_tokens", 0),
+            total_tokens=usage_data.get("total_tokens", 0),
+        )
+        return cls(
+            embeddings=embeddings,
+            model=data.get("model", ""),
+            usage=usage,
+            object_type=data.get("object", "list"),
+        )