hammad-python 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

hammad/{ai/embeddings/client/litellm_embeddings_client.py → genai/embedding_models/embedding_model.py}

@@ -1,6 +1,8 @@
-"""hammad.ai.embeddings.client.litellm_embeddings_client"""
+"""hammad.genai.embedding_models.embedding_model"""
 
-from typing import Any, List, Literal, Optional
+import asyncio
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, TYPE_CHECKING
 import sys
 
 if sys.version_info >= (3, 12):
@@ -8,110 +10,32 @@ if sys.version_info >= (3, 12):
 else:
     from typing_extensions import TypedDict
 
-
-from .base_embeddings_client import BaseEmbeddingsClient
-from ..types import (
+if TYPE_CHECKING:
+    try:
+        from litellm import EmbeddingResponse as _LitellmEmbeddingResponse
+    except ImportError:
+        _LitellmEmbeddingResponse = Any
+
+from ..language_models.language_model import _AIProvider
+from .embedding_model_request import EmbeddingModelRequest
+from .embedding_model_name import EmbeddingModelName
+from .embedding_model_response import (
     Embedding,
     EmbeddingUsage,
-    EmbeddingResponse,
+    EmbeddingModelResponse,
 )
-from ....formatting.text.converters import convert_to_text
-from ..._utils import get_litellm
+from ...formatting.text import convert_to_text
+
 
 __all__ = (
-    "LiteLlmEmbeddingsClient",
-    "LiteLlmEmbeddingModel",
-    "LiteLlmEmbeddingModelSettings",
+    "EmbeddingModel",
+    "EmbeddingModelError",
 )
 
 
-LiteLlmEmbeddingModel = Literal[
-    # OpenAI Embedding Models
-    "text-embedding-3-small",
-    "text-embedding-3-large",
-    "text-embedding-ada-002",
-    # OpenAI Compatible Embedding Models
-    "openai/text-embedding-3-small",
-    "openai/text-embedding-3-large",
-    "openai/text-embedding-ada-002",
-    # Bedrock Embedding Models
-    "amazon.titan-embed-text-v1",
-    "cohere.embed-english-v3",
-    "cohere.embed-multilingual-v3",
-    # Cohere Embedding Models
-    "embed-english-v3.0",
-    "embed-english-light-v3.0",
-    "embed-multilingual-v3.0",
-    "embed-multilingual-light-v3.0",
-    "embed-english-v2.0",
-    "embed-english-light-v2.0",
-    "embed-multilingual-v2.0",
-    # NVIDIA NIM Embedding Models
-    "nvidia_nim/NV-Embed-QA",
-    "nvidia_nim/nvidia/nv-embed-v1",
-    "nvidia_nim/nvidia/nv-embedqa-mistral-7b-v2",
-    "nvidia_nim/nvidia/nv-embedqa-e5-v5",
-    "nvidia_nim/nvidia/embed-qa-4",
-    "nvidia_nim/nvidia/llama-3.2-nv-embedqa-1b-v1",
-    "nvidia_nim/nvidia/llama-3.2-nv-embedqa-1b-v2",
-    "nvidia_nim/snowflake/arctic-embed-l",
-    "nvidia_nim/baai/bge-m3",
-    # HuggingFace Embedding Models
-    "huggingface/microsoft/codebert-base",
-    "huggingface/BAAI/bge-large-zh",
-    # Mistral AI Embedding Models
-    "mistral/mistral-embed",
-    # Gemini AI Embedding Models
-    "gemini/text-embedding-004",
-    # Vertex AI Embedding Models
-    "vertex_ai/textembedding-gecko",
-    "vertex_ai/textembedding-gecko-multilingual",
-    "vertex_ai/textembedding-gecko-multilingual@001",
-    "vertex_ai/textembedding-gecko@001",
-    "vertex_ai/textembedding-gecko@003",
-    "vertex_ai/text-embedding-preview-0409",
-    "vertex_ai/text-multilingual-embedding-preview-0409",
-    # Voyage AI Embedding Models
-    "voyage/voyage-01",
-    "voyage/voyage-lite-01",
-    "voyage/voyage-lite-01-instruct",
-    # Nebius AI Studio Embedding Models
-    "nebius/BAAI/bge-en-icl",
-    "nebius/BAAI/bge-multilingual-gemma2",
-    "nebius/intfloat/e5-mistral-7b-instruct",
-    # Ollama Embedding Models
-    "ollama/granite-embedding:30m",
-    "ollama/granite-embedding:278m",
-    "ollama/snowflake-arctic-embed2",
-    "ollama/bge-large",
-    "ollama/paraphrase-multilingual",
-    "ollama/bge-m3",
-    "ollama/snowflake-arctic-embed",
-    "ollama/mxbai-embed-large",
-    "ollama/all-minilm",
-    "ollama/nomic-embed-text",
-]
-"""Common embedding models supported by `litellm`."""
-
-
-class LiteLlmEmbeddingModelSettings(TypedDict):
-    """Valid settings for the `litellm` embedding models."""
-
-    model: LiteLlmEmbeddingModel | str
-    dimensions: Optional[int]
-    encoding_format: Optional[str]
-    timeout: Optional[int]
-    api_base: Optional[str]
-    api_version: Optional[str]
-    api_key: Optional[str]
-    api_type: Optional[str]
-    caching: bool
-    user: Optional[str]
-
-
-class LiteLlmEmbeddingError(Exception):
+class EmbeddingModelError(Exception):
     """Exception raised when an error occurs while generating embeddings
-    using `litellm`."""
+    using an embedding model."""
 
     def __init__(self, message: str, response: Any):
         self.message = message
@@ -119,8 +43,8 @@ class LiteLlmEmbeddingError(Exception):
         super().__init__(self.message)
 
 
-def _parse_litellm_response_to_embedding_response(response: Any) -> EmbeddingResponse:
-    """Parse the response from `litellm` to an `EmbeddingResponse` object."""
+def _parse_litellm_response_to_embedding_model_response(response: "_LitellmEmbeddingResponse") -> EmbeddingModelResponse:
+    """Parse the response from `litellm` to an `EmbeddingModelResponse` object."""
     try:
         embedding_data: List[Embedding] = []
 
@@ -132,27 +56,29 @@ def _parse_litellm_response_to_embedding_response(response: Any) -> EmbeddingRes
             prompt_tokens=response.usage.prompt_tokens,
             total_tokens=response.usage.total_tokens,
         )
-        return EmbeddingResponse(
+        return EmbeddingModelResponse(
             data=embedding_data,
             model=response.model,
             object="list",
             usage=usage,
         )
     except Exception as e:
-        raise LiteLlmEmbeddingError(
+        raise EmbeddingModelError(
             f"Failed to parse litellm response to embedding response: {e}",
             response,
         )
 
 
-class LiteLlmEmbeddingsClient(BaseEmbeddingsClient):
+@dataclass
+class EmbeddingModel:
     """Embeddings provider client that utilizes the `litellm` module
     when generating embeddings."""
-
-    @staticmethod
-    async def async_embed(
+    
+    model: EmbeddingModelName = "openai/text-embedding-3-small"
+    
+    async def async_run(
+        self,
         input: List[Any] | Any,
-        model: LiteLlmEmbeddingModel | str,
         dimensions: Optional[int] = None,
         encoding_format: Optional[str] = None,
         timeout=600,
@@ -163,13 +89,12 @@ class LiteLlmEmbeddingsClient(BaseEmbeddingsClient):
         caching: bool = False,
         user: Optional[str] = None,
         format: bool = False,
-    ) -> Embedding:
+    ) -> EmbeddingModelResponse:
         """Asynchronously generate embeddings for the given input using
         a valid `litellm` model.
 
         Args:
             input (List[Any] | Any) : The input text / content to generate embeddings for.
-            model (LiteLlmEmbeddingModel | str) : The model to use for generating embeddings.
             dimensions (Optional[int]) : The number of dimensions for the embedding.
             encoding_format (Optional[str]) : The format to return the embeddings in. (e.g. "float", "base64")
             timeout (int) : The timeout for the request.
@@ -182,7 +107,7 @@ class LiteLlmEmbeddingsClient(BaseEmbeddingsClient):
             format (bool) : Whether to format each non-string input as a markdown string.
 
         Returns:
-            Embedding : The embedding generated for the given input.
+            EmbeddingModelResponse : The embedding response generated for the given input.
         """
         if not isinstance(input, list):
             input = [input]
@@ -192,16 +117,16 @@ class LiteLlmEmbeddingsClient(BaseEmbeddingsClient):
                 try:
                     i = convert_to_text(i)
                 except Exception as e:
-                    raise LiteLlmEmbeddingError(
+                    raise EmbeddingModelError(
                         f"Failed to format input to text: {e}",
                         i,
                     )
 
-        async_embedding_fn = get_litellm().aembedding
+        async_embedding_fn = _AIProvider.get_litellm().aembedding
 
         try:
             response = await async_embedding_fn(
-                model=model,
+                model=self.model,
                 input=input,
                 dimensions=dimensions,
                 encoding_format=encoding_format,
@@ -214,14 +139,13 @@ class LiteLlmEmbeddingsClient(BaseEmbeddingsClient):
                 user=user,
             )
         except Exception as e:
-            raise e
+            raise EmbeddingModelError(f"Error in embedding model request: {e}", response=None) from e
 
-        return _parse_litellm_response_to_embedding_response(response)
+        return _parse_litellm_response_to_embedding_model_response(response)
 
-    @staticmethod
-    def embed(
+    def run(
+        self,
         input: List[Any] | Any,
-        model: LiteLlmEmbeddingModel | str,
         dimensions: Optional[int] = None,
         encoding_format: Optional[str] = None,
         timeout=600,
@@ -232,13 +156,12 @@ class LiteLlmEmbeddingsClient(BaseEmbeddingsClient):
         caching: bool = False,
         user: Optional[str] = None,
         format: bool = False,
-    ) -> Embedding:
+    ) -> EmbeddingModelResponse:
         """Generate embeddings for the given input using
         a valid `litellm` model.
 
         Args:
             input (List[Any] | Any) : The input text / content to generate embeddings for.
-            model (LiteLlmEmbeddingModel | str) : The model to use for generating embeddings.
             dimensions (Optional[int]) : The number of dimensions for the embedding.
             encoding_format (Optional[str]) : The format to return the embeddings in. (e.g. "float", "base64")
             timeout (int) : The timeout for the request.
@@ -251,26 +174,10 @@ class LiteLlmEmbeddingsClient(BaseEmbeddingsClient):
             format (bool) : Whether to format each non-string input as a markdown string.
 
         Returns:
-            Embedding : The embedding generated for the given input.
+            EmbeddingModelResponse : The embedding response generated for the given input.
         """
-        if not isinstance(input, list):
-            input = [input]
-
-        if format:
-            for i in input:
-                try:
-                    i = convert_to_text(i)
-                except Exception as e:
-                    raise LiteLlmEmbeddingError(
-                        f"Failed to format input to text: {e}",
-                        i,
-                    )
-
-        sync_embedding_fn = get_litellm().embedding
-
-        try:
-            response = sync_embedding_fn(
-                model=model,
+        return asyncio.run(
+            self.async_run(
                 input=input,
                 dimensions=dimensions,
                 encoding_format=encoding_format,
@@ -281,8 +188,6 @@ class LiteLlmEmbeddingsClient(BaseEmbeddingsClient):
                 api_type=api_type,
                 caching=caching,
                 user=user,
+                format=format,
             )
-        except Exception as e:
-            raise e
-
-        return _parse_litellm_response_to_embedding_response(response)
+        )

hammad/genai/embedding_models/embedding_model_name.py

@@ -0,0 +1,77 @@
+"""hammad.genai.embedding_models.embedding_model_name"""
+
+from typing import Literal
+
+
+__all__ = (
+    "EmbeddingModelName",
+)
+
+
+EmbeddingModelName = Literal[
+    # OpenAI Embedding Models
+    "text-embedding-3-small",
+    "text-embedding-3-large",
+    "text-embedding-ada-002",
+    # OpenAI Compatible Embedding Models
+    "openai/text-embedding-3-small",
+    "openai/text-embedding-3-large",
+    "openai/text-embedding-ada-002",
+    # Bedrock Embedding Models
+    "amazon.titan-embed-text-v1",
+    "cohere.embed-english-v3",
+    "cohere.embed-multilingual-v3",
+    # Cohere Embedding Models
+    "embed-english-v3.0",
+    "embed-english-light-v3.0",
+    "embed-multilingual-v3.0",
+    "embed-multilingual-light-v3.0",
+    "embed-english-v2.0",
+    "embed-english-light-v2.0",
+    "embed-multilingual-v2.0",
+    # NVIDIA NIM Embedding Models
+    "nvidia_nim/NV-Embed-QA",
+    "nvidia_nim/nvidia/nv-embed-v1",
+    "nvidia_nim/nvidia/nv-embedqa-mistral-7b-v2",
+    "nvidia_nim/nvidia/nv-embedqa-e5-v5",
+    "nvidia_nim/nvidia/embed-qa-4",
+    "nvidia_nim/nvidia/llama-3.2-nv-embedqa-1b-v1",
+    "nvidia_nim/nvidia/llama-3.2-nv-embedqa-1b-v2",
+    "nvidia_nim/snowflake/arctic-embed-l",
+    "nvidia_nim/baai/bge-m3",
+    # HuggingFace Embedding Models
+    "huggingface/microsoft/codebert-base",
+    "huggingface/BAAI/bge-large-zh",
+    # Mistral AI Embedding Models
+    "mistral/mistral-embed",
+    # Gemini AI Embedding Models
+    "gemini/text-embedding-004",
+    # Vertex AI Embedding Models
+    "vertex_ai/textembedding-gecko",
+    "vertex_ai/textembedding-gecko-multilingual",
+    "vertex_ai/textembedding-gecko-multilingual@001",
+    "vertex_ai/textembedding-gecko@001",
+    "vertex_ai/textembedding-gecko@003",
+    "vertex_ai/text-embedding-preview-0409",
+    "vertex_ai/text-multilingual-embedding-preview-0409",
+    # Voyage AI Embedding Models
+    "voyage/voyage-01",
+    "voyage/voyage-lite-01",
+    "voyage/voyage-lite-01-instruct",
+    # Nebius AI Studio Embedding Models
+    "nebius/BAAI/bge-en-icl",
+    "nebius/BAAI/bge-multilingual-gemma2",
+    "nebius/intfloat/e5-mistral-7b-instruct",
+    # Ollama Embedding Models
+    "ollama/granite-embedding:30m",
+    "ollama/granite-embedding:278m",
+    "ollama/snowflake-arctic-embed2",
+    "ollama/bge-large",
+    "ollama/paraphrase-multilingual",
+    "ollama/bge-m3",
+    "ollama/snowflake-arctic-embed",
+    "ollama/mxbai-embed-large",
+    "ollama/all-minilm",
+    "ollama/nomic-embed-text",
+]
+"""Common embedding models supported by `litellm`."""

hammad/genai/embedding_models/embedding_model_request.py

@@ -0,0 +1,65 @@
+"""hammad.genai.embedding_models.embedding_model_request"""
+
+import sys
+if sys.version_info >= (3, 12):
+    from typing import TypedDict, Required, NotRequired
+else:
+    from typing_extensions import TypedDict, Required, NotRequired
+
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    Literal,
+)
+
+from .embedding_model_name import EmbeddingModelName
+
+__all__ = [
+    "EmbeddingModelRequest",
+]
+
+
+class EmbeddingModelRequest(TypedDict, total=False):
+    """A request to an embedding model."""
+
+    input: List[Any] | Any
+    """The input items to embed."""
+
+    model: EmbeddingModelName | str
+    """The embedding model to use."""
+
+    format: bool = False
+    """Whether to format each non-string input as a markdown string."""
+
+    # LiteLLM Settings
+    dimensions: Optional[int] = None
+    """The dimensions of the embedding."""
+
+    encoding_format: Optional[str] = None
+    """The encoding format of the embedding."""
+
+    timeout: Optional[int] = None
+    """The timeout for the embedding request."""
+
+    api_base: Optional[str] = None
+    """The API base for the embedding request."""
+
+    api_version: Optional[str] = None
+    """The API version for the embedding request."""
+
+    api_key: Optional[str] = None
+    """The API key for the embedding request."""
+
+    api_type: Optional[str] = None
+    """The API type for the embedding request."""
+
+    caching: bool = False
+    """Whether to cache the embedding request."""
+
+    user: Optional[str] = None
+    """The user for the embedding request."""

hammad/{ai/embeddings/types.py → genai/embedding_models/embedding_model_response.py}

@@ -1,4 +1,4 @@
-"""hammad.ai.embeddings.types"""
+"""hammad.genai.embedding_models.embedding_model_response"""
 
 from typing import List, Literal
 
@@ -41,7 +41,7 @@ class EmbeddingUsage(BaseModel):
     """The total number of tokens used by the request."""
 
 
-class EmbeddingResponse(BaseModel):
+class EmbeddingModelResponse(BaseModel):
     data: List[Embedding]
     """The list of embeddings generated by the model."""
 
@@ -61,7 +61,7 @@ class EmbeddingResponse(BaseModel):
 
     def __str__(self) -> str:
         return (
-            "Embedding Response:\n"
+            "EmbeddingModelResponse:\n"
             f">>> Model: {self.model}\n"
             f">>> Dimensions: {self.dimensions}\n"
             f">>> Usage: {self.usage}\n"

hammad/genai/embedding_models/run.py

@@ -0,0 +1,161 @@
+"""hammad.genai.embedding_models.run
+
+Standalone functions for running embedding models with full parameter typing.
+"""
+
+from typing import Any, List, Optional, overload, Union
+
+from .embedding_model_name import EmbeddingModelName
+from .embedding_model import EmbeddingModel
+from .embedding_model_response import EmbeddingModelResponse
+
+__all__ = [
+    "run_embedding_model",
+    "async_run_embedding_model",
+]
+
+
+# Overloads for run_embedding_model
+@overload
+def run_embedding_model(
+    input: List[Any] | Any,
+    *,
+    # Provider settings
+    model: EmbeddingModelName = "openai/text-embedding-3-small",
+    api_base: Optional[str] = None,
+    api_key: Optional[str] = None,
+    api_version: Optional[str] = None,
+    api_type: Optional[str] = None,
+    # Extended settings
+    dimensions: Optional[int] = None,
+    encoding_format: Optional[str] = None,
+    timeout: int = 600,
+    caching: bool = False,
+    user: Optional[str] = None,
+    format: bool = False,
+) -> EmbeddingModelResponse: ...
+
+
+def run_embedding_model(
+    input: List[Any] | Any,
+    *,
+    # Provider settings
+    model: EmbeddingModelName = "openai/text-embedding-3-small",
+    api_base: Optional[str] = None,
+    api_key: Optional[str] = None,
+    api_version: Optional[str] = None,
+    api_type: Optional[str] = None,
+    # Extended settings
+    dimensions: Optional[int] = None,
+    encoding_format: Optional[str] = None,
+    timeout: int = 600,
+    caching: bool = False,
+    user: Optional[str] = None,
+    format: bool = False,
+) -> EmbeddingModelResponse:
+    """Run an embedding model with the given input.
+    
+    Args:
+        input: The input text/content to generate embeddings for
+        model: The embedding model to use
+        api_base: The base URL for the API
+        api_key: The API key to use for the request
+        api_version: The version of the API
+        api_type: The API type to use for the request
+        dimensions: The number of dimensions for the embedding
+        encoding_format: The format to return the embeddings in
+        timeout: The timeout for the request
+        caching: Whether to cache the request
+        user: The user to use for the request
+        format: Whether to format each non-string input as a markdown string
+        
+    Returns:
+        EmbeddingModelResponse: The embedding response
+    """
+    embedding_model = EmbeddingModel(model=model)
+    return embedding_model.run(
+        input=input,
+        dimensions=dimensions,
+        encoding_format=encoding_format,
+        timeout=timeout,
+        api_base=api_base,
+        api_version=api_version,
+        api_key=api_key,
+        api_type=api_type,
+        caching=caching,
+        user=user,
+        format=format,
+    )
+
+
+# Overloads for async_run_embedding_model
+@overload
+async def async_run_embedding_model(
+    input: List[Any] | Any,
+    *,
+    # Provider settings
+    model: EmbeddingModelName = "openai/text-embedding-3-small",
+    api_base: Optional[str] = None,
+    api_key: Optional[str] = None,
+    api_version: Optional[str] = None,
+    api_type: Optional[str] = None,
+    # Extended settings
+    dimensions: Optional[int] = None,
+    encoding_format: Optional[str] = None,
+    timeout: int = 600,
+    caching: bool = False,
+    user: Optional[str] = None,
+    format: bool = False,
+) -> EmbeddingModelResponse: ...
+
+
+async def async_run_embedding_model(
+    input: List[Any] | Any,
+    *,
+    # Provider settings
+    model: EmbeddingModelName = "openai/text-embedding-3-small",
+    api_base: Optional[str] = None,
+    api_key: Optional[str] = None,
+    api_version: Optional[str] = None,
+    api_type: Optional[str] = None,
+    # Extended settings
+    dimensions: Optional[int] = None,
+    encoding_format: Optional[str] = None,
+    timeout: int = 600,
+    caching: bool = False,
+    user: Optional[str] = None,
+    format: bool = False,
+) -> EmbeddingModelResponse:
+    """Asynchronously run an embedding model with the given input.
+    
+    Args:
+        input: The input text/content to generate embeddings for
+        model: The embedding model to use
+        api_base: The base URL for the API
+        api_key: The API key to use for the request
+        api_version: The version of the API
+        api_type: The API type to use for the request
+        dimensions: The number of dimensions for the embedding
+        encoding_format: The format to return the embeddings in
+        timeout: The timeout for the request
+        caching: Whether to cache the request
+        user: The user to use for the request
+        format: Whether to format each non-string input as a markdown string
+        
+    Returns:
+        EmbeddingModelResponse: The embedding response
+    """
+    embedding_model = EmbeddingModel(model=model)
+    return await embedding_model.async_run(
+        input=input,
+        dimensions=dimensions,
+        encoding_format=encoding_format,
+        timeout=timeout,
+        api_base=api_base,
+        api_version=api_version,
+        api_key=api_key,
+        api_type=api_type,
+        caching=caching,
+        user=user,
+        format=format,
+    )

hammad/genai/language_models/__init__.py

@@ -0,0 +1,35 @@
+"""hammad.genai.language_models"""
+
+from typing import TYPE_CHECKING
+from ..._internal import create_getattr_importer
+
+if TYPE_CHECKING:
+    from .language_model import LanguageModel
+    from .run import run_language_model, async_run_language_model
+    from .language_model_request import LanguageModelMessagesParam
+    from .language_model_response import LanguageModelResponse
+    from .language_model_response_chunk import LanguageModelResponseChunk
+    from .language_model_request import LanguageModelRequest
+
+__all__ = (
+    # hammad.genai.language_models.language_model
+    "LanguageModel",
+
+    # hammad.genai.language_models.run
+    "run_language_model",
+    "async_run_language_model",
+
+    # hammad.genai.language_models.language_model_request
+    "LanguageModelMessagesParam",
+    "LanguageModelRequest",
+
+    # hammad.genai.language_models.language_model_response
+    "LanguageModelResponse",
+    "LanguageModelResponseChunk",
+)
+
+__getattr__ = create_getattr_importer(__all__)
+
+def __dir__() -> list[str]:
+    """Get the attributes of the language_models module."""
+    return list(__all__)